Merge pull request #44376 from valsdav/tensorflow-disable-gpu-tests_1…

…4_1_X Improve handling of TF CUDA tests for 14_1_X
cms-sw · Mar 18, 2024 · 28f76ff · 28f76ff
2 parents 23e94df + b9b41e4
commit 28f76ff
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 7 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");

diff --git a/PhysicsTools/TensorFlow/test/BuildFile.xml b/PhysicsTools/TensorFlow/test/BuildFile.xml
@@ -5,7 +5,7 @@
   <use name="PhysicsTools/TensorFlow"/>
 </bin>
 
-<iftool name="cuda">
+<iftool name="tf_cuda_support">
   <bin name="testTFHelloWorldCUDA" file="testRunner.cpp,testHelloWorldCUDA.cc">
     <use name="boost_filesystem"/>
     <use name="catch2"/>
@@ -29,7 +29,7 @@
   <use name="PhysicsTools/TensorFlow"/>
 </bin>
 
-<iftool name="cuda">
+<iftool name="tf_cuda_support">
   <bin name="testTFMetaGraphLoadingCUDA" file="testRunner.cpp,testMetaGraphLoadingCUDA.cc">
     <use name="boost_filesystem"/>
     <use name="catch2"/>
@@ -52,7 +52,7 @@
   <use name="PhysicsTools/TensorFlow"/>
 </bin>
 
-<iftool name="cuda">
+<iftool name="tf_cuda_support">
   <bin name="testTFGraphLoadingCUDA" file="testRunner.cpp,testGraphLoadingCUDA.cc">
     <use name="boost_filesystem"/>
     <use name="catch2"/>
@@ -75,7 +75,8 @@
   <use name="PhysicsTools/TensorFlow"/>
 </bin>
 
-<iftool name="cuda">
+
+<iftool name="tf_cuda_support">
   <bin name="testTFConstSessionCUDA" file="testRunner.cpp,testConstSessionCUDA.cc">
     <use name="boost_filesystem"/>
     <use name="catch2"/>
@@ -98,7 +99,7 @@
   <use name="PhysicsTools/TensorFlow"/>
 </bin>
 
-<iftool name="cuda">
+<iftool name="tf_cuda_support">
   <bin name="testTFSessionCacheCUDA" file="testRunner.cpp,testSessionCacheCUDA.cc">
     <use name="boost_filesystem"/>
     <use name="catch2"/>
@@ -121,7 +122,8 @@
   <use name="PhysicsTools/TensorFlow"/>
 </bin>
 
-<iftool name="cuda">
+
+<iftool name="tf_cuda_support">
   <bin name="testTFThreadPoolsCUDA" file="testRunner.cpp,testThreadPoolsCUDA.cc">
     <use name="boost_filesystem"/>
     <use name="catch2"/>
@@ -145,7 +147,7 @@
 </bin>
 
 
-<iftool name="cuda">
+<iftool name="tf_cuda_support">
   <bin name="testTFVisibleDevicesCUDA" file="testRunner.cpp,testVisibleDevicesCUDA.cc">
     <use name="boost_filesystem"/>
     <use name="catch2"/>