Change how CUDA runtime and capabilities are defined in the task and Condor

src/python/WMCore/BossAir/Plugins/BasePlugin.py

@@ -7,6 +7,7 @@
 
 from builtins import object, str, bytes
 from future.utils import viewvalues
+from distutils.version import StrictVersion
 
 from Utils.Utilities import decodeBytesToUnicode
 from WMCore.WMException import WMException
@@ -181,3 +182,32 @@ def scramArchtoRequiredArch(scramArch=None):
             archs = defaultArch
 
         return archs
+
+    @staticmethod
+    def cudaCapabilityToSingleVersion(capabilities=None):
+        """
+        Given a list of CUDA capabilities (with strings in a version style),
+        finds the smallest version required and convert it to a single integer
+        for comparison/job matchmaking purposes.
+        Version conversion formula is: (1000 * major + 10 * medium + minor)
+        :param capabilities: a list of string versions
+        :return: an integer with the version value; None in case of failure
+
+        For further details:
+        https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART____VERSION.html
+        """
+        if not (isinstance(capabilities, list) and capabilities):
+            return None
+        # now order the list of string versions in place. Precedence of digits is from left to right
+        #    from: ["2.3.1", "1.2.3", "3.2.1", "1.3.2"]
+        #    to:   ["1.2.3", "1.3.2", "2.3.1", "3.2.1"]
+        capabilities.sort(key=StrictVersion)
+
+        smallestVersion = capabilities[0]
+        smallestVersion = smallestVersion.split(".")
+        # deal with versions like: "1", "1.2" and "1.2.3"
+        for _i in range(0, 3 - len(smallestVersion)):
+            smallestVersion.append(0)
+
+        intVersion = int(smallestVersion[0]) * 1000 + int(smallestVersion[1]) * 10 + int(smallestVersion[2])
+        return intVersion

src/python/WMCore/BossAir/Plugins/SimpleCondorPlugin.py

@@ -566,8 +566,14 @@ def getJobParameters(self, jobList):
             if job.get('gpuRequirements', None):
                 ad['My.GPUMemoryMB'] = str(job['gpuRequirements']['GPUMemoryMB'])
                 cudaCapabilities = ','.join(sorted(job['gpuRequirements']['CUDACapabilities']))
-                ad['My.CUDACapability'] = classad.quote(str(cudaCapabilities))
-                ad['My.CUDARuntime'] = classad.quote(job['gpuRequirements']['CUDARuntime'])
+                minimalCapability = self.cudaCapabilityToSingleVersion(job['gpuRequirements']['CUDACapabilities'])
+                if minimalCapability is None:  # this should never happen!!
+                    ad['My.CUDACapability'] = undefined
+                else:
+                    ad['My.CUDACapability'] = classad.quote(str(minimalCapability))
+                ad['My.OriginalCUDACapability'] = classad.quote(str(cudaCapabilities))
+                cudaRuntime = ','.join(sorted(job['gpuRequirements']['CUDARuntime']))
+                ad['My.CUDARuntime'] = classad.quote(str(cudaRuntime))
             else:
                 ad['My.GPUMemoryMB'] = undefined
                 ad['My.CUDACapability'] = undefined

src/python/WMCore/WMSpec/WMTask.py

@@ -1525,16 +1525,31 @@ def getRequiresGPU(self):
     def getGPURequirements(self):
         """
         Return the GPU requirements for this task.
-        If it's a multi-step task, the first step with a meaningful
-        dictionary value will be returned
+        For multi-step tasks, the following logic is applied:
+          * GPUMemoryMB: return the max of them
+          * CUDARuntime: returns a flat list of unique runtime versions
+          * CUDACapabilities: returns a flat list of unique capabilities
         :return: a dictionary with the GPU requirements for this task
         """
-        gpuRequirements = {}
+        gpuRequirements = []
         for stepName in sorted(self.listAllStepNames()):
             stepHelper = self.getStep(stepName)
             if stepHelper.stepType() == "CMSSW" and stepHelper.getGPURequirements():
-                return stepHelper.getGPURequirements()
-        return gpuRequirements
+                gpuRequirements.append(stepHelper.getGPURequirements())
+        if not gpuRequirements:
+            return {}
+
+        # in this case, it requires GPUs and it can be multi-steps GPU
+        bestGPUParams = {"GPUMemoryMB": 0, "CUDARuntime": [], "CUDACapabilities": []}
+        for params in gpuRequirements:
+            if params["GPUMemoryMB"] > bestGPUParams["GPUMemoryMB"]:
+                bestGPUParams["GPUMemoryMB"] = params["GPUMemoryMB"]
+            bestGPUParams["CUDARuntime"].append(params["CUDARuntime"])
+            bestGPUParams["CUDACapabilities"].extend(params["CUDACapabilities"])
+        # make the flat list elements unique
+        bestGPUParams["CUDARuntime"] = list(set(bestGPUParams["CUDARuntime"]))
+        bestGPUParams["CUDACapabilities"] = list(set(bestGPUParams["CUDACapabilities"]))
+        return bestGPUParams
 
     def _getStepValue(self, keyDict, defaultValue):
         """

test/python/WMCore_t/BossAir_t/BasePlugin_t.py

@@ -78,5 +78,27 @@ def testScramArchtoRequiredArch(self):
 
         return
 
+    def testCudaCapabilityToSingleVersion(self):
+        """
+        Test conversion of a list of version strings to a single integer version
+        """
+        bp = BasePlugin(config=None)
+
+        # bad input
+        self.assertEqual(bp.cudaCapabilityToSingleVersion([]), None)
+        self.assertEqual(bp.cudaCapabilityToSingleVersion({}), None)
+        self.assertEqual(bp.cudaCapabilityToSingleVersion(None), None)
+        # good and expected input
+        self.assertEqual(bp.cudaCapabilityToSingleVersion(["5.0"]), 5000)
+        unorderedL = ["2.3.1", "1.2.3", "3.2.1", "1.3.2", "1.2"]
+        self.assertEqual(bp.cudaCapabilityToSingleVersion(unorderedL), 1020)
+        orderedL = ["1.2", "1.2.3", "1.3.2", "2.3.1", "3.2.1"]
+        self.assertEqual(bp.cudaCapabilityToSingleVersion(orderedL), 1020)
+        orderedL = ["1.2.3", "1.3.2", "2.3.1", "3.2.1"]
+        self.assertEqual(bp.cudaCapabilityToSingleVersion(orderedL), 1023)
+        orderedL = ["2.3.1", "3.2.1"]
+        self.assertEqual(bp.cudaCapabilityToSingleVersion(orderedL), 2031)
+
+
 if __name__ == '__main__':
     unittest.main()

test/python/WMCore_t/WMSpec_t/StdSpecs_t/StepChain_t.py

@@ -2489,9 +2489,10 @@ def testGPUStepChainsTasks(self):
             testArguments[s]['ConfigCacheID'] = configDocs[s]
         testArguments['Step2']['KeepOutput'] = False
 
-        gpuParams = {"GPUMemoryMB": 1234, "CUDARuntime": "11.2.3", "CUDACapabilities": ["7.5", "8.0"]}
-        testArguments['Step1'].update({"RequiresGPU": "optional", "GPUParams": json.dumps(gpuParams)})
-        testArguments['Step2'].update({"RequiresGPU": "required", "GPUParams": json.dumps(gpuParams)})
+        gpuParams1 = {"GPUMemoryMB": 1234, "CUDARuntime": "11.2.3", "CUDACapabilities": ["7.5", "8.0"]}
+        testArguments['Step1'].update({"RequiresGPU": "optional", "GPUParams": json.dumps(gpuParams1)})
+        gpuParams2 = {"GPUMemoryMB": 2345, "CUDARuntime": "9.6", "CUDACapabilities": ["7.4"]}
+        testArguments['Step2'].update({"RequiresGPU": "required", "GPUParams": json.dumps(gpuParams2)})
         factory = StepChainWorkloadFactory()
         testWorkload = factory.factoryWorkloadConstruction("TestWorkload", testArguments)
 
@@ -2503,8 +2504,8 @@ def testGPUStepChainsTasks(self):
 
         # validate GPU parameters
         self.assertEqual(testArguments['GPUParams'], json.dumps(None))
-        self.assertEqual(testArguments["Step1"]['GPUParams'], json.dumps(gpuParams))
-        self.assertEqual(testArguments["Step2"]['GPUParams'], json.dumps(gpuParams))
+        self.assertEqual(testArguments["Step1"]['GPUParams'], json.dumps(gpuParams1))
+        self.assertEqual(testArguments["Step2"]['GPUParams'], json.dumps(gpuParams2))
         self.assertTrue("GPUParams" not in testArguments["Step3"])
 
         for taskName in testWorkload.listAllTaskNames():
@@ -2520,10 +2521,10 @@ def testGPUStepChainsTasks(self):
                 elif stepHelper.stepType() == "CMSSW" and taskName == "GENSIM":
                     if stepHelper.name() == "cmsRun1":
                         self.assertEqual(stepHelper.data.application.gpu.gpuRequired, testArguments["Step1"]['RequiresGPU'])
-                        self.assertItemsEqual(stepHelper.data.application.gpu.gpuRequirements, gpuParams)
+                        self.assertItemsEqual(stepHelper.data.application.gpu.gpuRequirements, gpuParams1)
                     elif stepHelper.name() == "cmsRun2":
                         self.assertEqual(stepHelper.data.application.gpu.gpuRequired, testArguments["Step2"]['RequiresGPU'])
-                        self.assertItemsEqual(stepHelper.data.application.gpu.gpuRequirements, gpuParams)
+                        self.assertItemsEqual(stepHelper.data.application.gpu.gpuRequirements, gpuParams2)
                     elif stepHelper.name() == "cmsRun3":
                         self.assertEqual(stepHelper.data.application.gpu.gpuRequired, "forbidden")
                         self.assertIsNone(stepHelper.data.application.gpu.gpuRequirements)
@@ -2535,18 +2536,17 @@ def testGPUStepChainsTasks(self):
         prodTask = testWorkload.getTask('GENSIM')
         gpuRequired, gpuRequirements = prodTask.getStepHelper('cmsRun1').getGPUSettings()
         self.assertEqual(gpuRequired, testArguments["Step1"]['RequiresGPU'])
-        self.assertItemsEqual(gpuRequirements, gpuParams)
+        self.assertItemsEqual(gpuRequirements, gpuParams1)
 
         gpuRequired, gpuRequirements = prodTask.getStepHelper('cmsRun2').getGPUSettings()
         self.assertEqual(gpuRequired, testArguments["Step2"]['RequiresGPU'])
-        self.assertItemsEqual(gpuRequirements, gpuParams)
+        self.assertItemsEqual(gpuRequirements, gpuParams2)
 
         gpuRequired, gpuRequirements = prodTask.getStepHelper('cmsRun3').getGPUSettings()
         self.assertEqual(gpuRequired, testArguments["Step3"].get('RequiresGPU', "forbidden"))
         self.assertIsNone(gpuRequirements)
 
 
-
         # test assignment with wrong Trust flags
         assignDict = {"SiteWhitelist": ["T2_US_Nebraska"], "Team": "The-A-Team",
                       "RequestStatus": "assigned"}
@@ -2560,18 +2560,18 @@ def testGPUStepChainsTasks(self):
 
         # validate GPU parameters
         self.assertEqual(testArguments['GPUParams'], json.dumps(None))
-        self.assertEqual(testArguments["Step1"]['GPUParams'], json.dumps(gpuParams))
-        self.assertEqual(testArguments["Step2"]['GPUParams'], json.dumps(gpuParams))
+        self.assertEqual(testArguments["Step1"]['GPUParams'], json.dumps(gpuParams1))
+        self.assertEqual(testArguments["Step2"]['GPUParams'], json.dumps(gpuParams2))
         self.assertTrue("GPUParams" not in testArguments["Step3"])
 
         prodTask = testWorkload.getTask('GENSIM')
         gpuRequired, gpuRequirements = prodTask.getStepHelper('cmsRun1').getGPUSettings()
         self.assertEqual(gpuRequired, testArguments["Step1"]['RequiresGPU'])
-        self.assertItemsEqual(gpuRequirements, gpuParams)
+        self.assertItemsEqual(gpuRequirements, gpuParams1)
 
         gpuRequired, gpuRequirements = prodTask.getStepHelper('cmsRun2').getGPUSettings()
         self.assertEqual(gpuRequired, testArguments["Step2"]['RequiresGPU'])
-        self.assertItemsEqual(gpuRequirements, gpuParams)
+        self.assertItemsEqual(gpuRequirements, gpuParams2)
 
         gpuRequired, gpuRequirements = prodTask.getStepHelper('cmsRun3').getGPUSettings()
         self.assertEqual(gpuRequired, testArguments["Step3"].get('RequiresGPU', "forbidden"))

test/python/WMCore_t/WMSpec_t/WMTask_t.py

@@ -858,6 +858,8 @@ def testGPUTaskSettings(self):
         ### Now set a single value for both tasks
         gpuParams = {"GPUMemoryMB": 1234, "CUDARuntime": "11.2.3", "CUDACapabilities": ["7.5", "8.0"]}
         task1.setTaskGPUSettings("required", json.dumps(gpuParams))
+        # CUDARuntime returns as a list
+        gpuParams["CUDARuntime"] = [gpuParams["CUDARuntime"]]
         for taskObj in task1.taskIterator():
             # task level check
             self.assertEqual(taskObj.getRequiresGPU(), "required")