Checking if Nvidia GPUs are visible to TF when visible to cmssw

cms-sw · Mar 13, 2024 · b9b41e4 · b9b41e4
1 parent 7447264
commit b9b41e4
Showing 1 changed file with 9 additions and 0 deletions.
diff --git a/PhysicsTools/TensorFlow/src/TensorFlow.cc b/PhysicsTools/TensorFlow/src/TensorFlow.cc
@@ -42,6 +42,15 @@ namespace tensorflow {
     // NVidia GPU
     else if (backend == Backend::cuda) {
       if (not ri->nvidiaDriverVersion().empty()) {
+        // Check if one GPU device is visible to TF
+        // If not, an exception is raised --> this can happen in case of driver version mismatch
+        // or missing CUDA support in TF compilation
+        if ((*_options.config.mutable_device_count())["GPU"] == 0) {
+          edm::Exception ex(edm::errors::UnavailableAccelerator);
+          ex << "Cuda backend requested, NVIDIA GPU visible to cmssw, but not visible to TensorFlow in the job";
+          ex.addContext("Calling tensorflow::setBackend()");
+          throw ex;
+        }
         // Take only the first GPU in the CUDA_VISIBLE_DEVICE list
         (*_options.config.mutable_device_count())["GPU"] = 1;
         _options.config.mutable_gpu_options()->set_visible_device_list("0");