[Bugfix] Add CUDA context availability check before setting curand seed (#4223)

include/dgl/runtime/device_api.h

@@ -44,6 +44,12 @@ class DeviceAPI {
  public:
   /*! \brief virtual destructor */
   virtual ~DeviceAPI() {}
+  /*!
+   * \brief Check whether the device is available.
+   */
+  virtual bool IsAvailable() {
+    return true;
+  }
   /*!
    * \brief Set the environment device id to ctx
    * \param ctx The context to be set.

src/random/random.cc

@@ -29,13 +29,15 @@ DGL_REGISTER_GLOBAL("rng._CAPI_SetSeed")
       }
     });
 #ifdef DGL_USE_CUDA
-    auto* thr_entry = CUDAThreadEntry::ThreadLocal();
-    if (!thr_entry->curand_gen) {
-      CURAND_CALL(curandCreateGenerator(&thr_entry->curand_gen, CURAND_RNG_PSEUDO_DEFAULT));
+    if (DeviceAPI::Get(kDLGPU)->IsAvailable()) {
+      auto* thr_entry = CUDAThreadEntry::ThreadLocal();
+      if (!thr_entry->curand_gen) {
+        CURAND_CALL(curandCreateGenerator(&thr_entry->curand_gen, CURAND_RNG_PSEUDO_DEFAULT));
+      }
+      CURAND_CALL(curandSetPseudoRandomGeneratorSeed(
+          thr_entry->curand_gen,
+          static_cast<uint64_t>(seed)));
     }
-    CURAND_CALL(curandSetPseudoRandomGeneratorSeed(
-        thr_entry->curand_gen,
-        static_cast<uint64_t>(seed)));
 #endif  // DGL_USE_CUDA
   });
 

src/runtime/cuda/cuda_device_api.cc

@@ -15,6 +15,23 @@ namespace runtime {
 
 class CUDADeviceAPI final : public DeviceAPI {
  public:
+  CUDADeviceAPI() {
+    int count;
+    auto err = cudaGetDeviceCount(&count);
+    switch (err) {
+      case cudaSuccess:
+        break;
+      default:
+        count = 0;
+        cudaGetLastError();
+    }
+    is_available_ = count > 0;
+  }
+
+  bool IsAvailable() final {
+    return is_available_;
+  }
+
   void SetDevice(DGLContext ctx) final {
     CUDA_CALL(cudaSetDevice(ctx.device_id));
   }