Merge pull request #35263 from srimanob/121_runTheMatrixGPUWorkflow

Add GPU workflow to runTheMatrix
cms-sw · Sep 21, 2021 · 5d10e5e · 5d10e5e
2 parents 0ec8e62 + de6a246
commit 5d10e5e
Show file tree

Hide file tree

Showing 2 changed files with 291 additions and 214 deletions.
diff --git a/Configuration/PyReleaseValidation/python/MatrixInjector.py b/Configuration/PyReleaseValidation/python/MatrixInjector.py
@@ -65,6 +65,19 @@ def __init__(self,opt,mode='init',options=''):
         if(opt.batchName):
             self.batchName = '__'+opt.batchName+'-'+self.batchTime
 
+        ####################################
+        # Checking and setting up GPU attributes
+        ####################################
+        # Mendatory
+        self.RequiresGPU = opt.gpu
+        self.GPUMemoryMB = opt.GPUMemoryMB
+        self.CUDACapabilities = opt.CUDACapabilities
+        self.CUDARuntime = opt.CUDARuntime
+        # optional
+        self.GPUName = opt.GPUName
+        self.CUDADriverVersion = opt.CUDADriverVersion
+        self.CUDARuntimeVersion = opt.CUDARuntimeVersion
+
         # WMagent url
         if not self.wmagent:
             # Overwrite with env variable
@@ -180,8 +193,18 @@ def __init__(self,opt,mode='init',options=''):
             "nowmIO": {},
             "Multicore" : opt.nThreads,                       # this is the per-taskchain Multicore; it's the default assigned to a task if it has no value specified 
             "EventStreams": self.numberOfStreams,
-            "KeepOutput" : False
+            "KeepOutput" : False,
+            "RequiresGPU" : None,
+            "GPUParams": None
+            }
+        self.defaultGPUParams={
+            "GPUMemoryMB": self.GPUMemoryMB,
+            "CUDACapabilities": self.CUDACapabilities,
+            "CUDARuntime": self.CUDARuntime
             }
+        if self.GPUName: self.defaultGPUParams.update({"GPUName": self.GPUName})
+        if self.CUDADriverVersion: self.defaultGPUParams.update({"CUDADriverVersion": self.CUDADriverVersion})
+        if self.CUDARuntimeVersion: self.defaultGPUParams.update({"CUDARuntimeVersion": self.CUDARuntimeVersion})
 
         self.chainDicts={}
 
@@ -408,6 +431,9 @@ def prepare(self, mReader, directories, mode='init'):
                                     if setPrimaryDs:
                                         chainDict['nowmTasklist'][-1]['PrimaryDataset']=setPrimaryDs
                                 nextHasDSInput=None
+                                if 'GPU' in step and self.RequiresGPU != 'forbidden':
+                                    chainDict['nowmTasklist'][-1]['RequiresGPU'] = self.RequiresGPU
+                                    chainDict['nowmTasklist'][-1]['GPUParams']=json.dumps(self.defaultGPUParams)
                             else:
                                 #not first step and no inputDS
                                 chainDict['nowmTasklist'].append(copy.deepcopy(self.defaultTask))
@@ -420,6 +446,9 @@ def prepare(self, mReader, directories, mode='init'):
                                     chainDict['nowmTasklist'][-1]['LumisPerJob']=splitForThisWf
                                 if step in wmsplit:
                                     chainDict['nowmTasklist'][-1]['LumisPerJob']=wmsplit[step]
+                                if 'GPU' in step and self.RequiresGPU != 'forbidden':
+                                    chainDict['nowmTasklist'][-1]['RequiresGPU'] = self.RequiresGPU
+                                    chainDict['nowmTasklist'][-1]['GPUParams']=json.dumps(self.defaultGPUParams)
 
                             # change LumisPerJob for Hadronizer steps. 
                             if 'Hadronizer' in step: