feat: PoolCE can subdivide RAM

Author: fstagni

src/DIRAC/Interfaces/API/Job.py

@@ -522,6 +522,19 @@ def setDestination(self, destination):
         return S_OK()
 
     #############################################################################
+    def setRAMRequirements(self, ramRequired: int = 0):
+        """Helper function.
+        Specify the RAM requirements for the job in GB. 0 (default) means no specific requirements.
+        """
+        if ramRequired:
+            self._addParameter(
+                self.workflow,
+                "MaxRAM",
+                "JDL",
+                ramRequired,
+                "GBs of RAM requested",
+            )
+
     def setNumberOfProcessors(self, numberOfProcessors=None, minNumberOfProcessors=None, maxNumberOfProcessors=None):
         """Helper function.
 

src/DIRAC/Resources/Computing/PoolComputingElement.py

@@ -61,6 +61,8 @@ def __init__(self, ceUniqueID):
         self.taskID = 0
         self.processorsPerTask = {}
         self.userNumberPerTask = {}
+        self.ram = 1024  # Default RAM in GB (this is an arbitrary large value in case of no limit)
+        self.ramPerTask = {}
 
         # This CE will effectively submit to another "Inner"CE
         # (by default to the InProcess CE)
@@ -74,22 +76,16 @@ def _reset(self):
 
         self.processors = int(self.ceParameters.get("NumberOfProcessors", self.processors))
         self.ceParameters["MaxTotalJobs"] = self.processors
+        max_ram = int(self.ceParameters.get("MaxRAM", 0))
+        if max_ram > 0:
+            self.ram = max_ram // 1024  # Convert from MB to GB
+        self.ceParameters["MaxRAM"] = self.ram
         # Indicates that the submission is done asynchronously
         # The result is not immediately available
         self.ceParameters["AsyncSubmission"] = True
         self.innerCESubmissionType = self.ceParameters.get("InnerCESubmissionType", self.innerCESubmissionType)
         return S_OK()
 
-    def getProcessorsInUse(self):
-        """Get the number of currently allocated processor cores
-
-        :return: number of processors in use
-        """
-        processorsInUse = 0
-        for future in self.processorsPerTask:
-            processorsInUse += self.processorsPerTask[future]
-        return processorsInUse
-
     #############################################################################
     def submitJob(self, executableFile, proxy=None, inputs=None, **kwargs):
         """Method to submit job.
@@ -112,15 +108,22 @@ def submitJob(self, executableFile, proxy=None, inputs=None, **kwargs):
             self.taskID += 1
             return S_OK(taskID)
 
-        # Now persisting the job limits for later use in pilot.cfg file (pilot 3 default)
+        memoryForJob = self._getMemoryForJobs(kwargs)
+        if memoryForJob is None:
+            self.taskResults[self.taskID] = S_ERROR("Not enough memory for the job")
+            taskID = self.taskID
+            self.taskID += 1
+            return S_OK(taskID)
+
+        # Now persisting the job limits for later use in pilot.cfg file
         cd = ConfigurationData(loadDefaultCFG=False)
         res = cd.loadFile("pilot.cfg")
         if not res["OK"]:
             self.log.error("Could not load pilot.cfg", res["Message"])
         else:
-            # only NumberOfProcessors for now, but RAM (or other stuff) can also be added
             jobID = int(kwargs.get("jobDesc", {}).get("jobID", 0))
             cd.setOptionInCFG("/Resources/Computing/JobLimits/%d/NumberOfProcessors" % jobID, processorsForJob)
+            cd.setOptionInCFG("/Resources/Computing/JobLimits/%d/MaxRAM" % jobID, memoryForJob)
             res = cd.dumpLocalCFGToFile("pilot.cfg")
             if not res["OK"]:
                 self.log.error("Could not dump cfg to pilot.cfg", res["Message"])
@@ -132,6 +135,7 @@ def submitJob(self, executableFile, proxy=None, inputs=None, **kwargs):
         # Submission
         future = self.pPool.submit(executeJob, executableFile, proxy, self.taskID, inputs, **taskKwargs)
         self.processorsPerTask[future] = processorsForJob
+        self.ramPerTask[future] = memoryForJob
         future.add_done_callback(functools.partial(self.finalizeJob, self.taskID))
 
         taskID = self.taskID
@@ -141,7 +145,7 @@ def submitJob(self, executableFile, proxy=None, inputs=None, **kwargs):
 
     def _getProcessorsForJobs(self, kwargs):
         """helper function"""
-        processorsInUse = self.getProcessorsInUse()
+        processorsInUse = sum(self.processorsPerTask.values())
         availableProcessors = self.processors - processorsInUse
 
         self.log.verbose(
@@ -178,6 +182,24 @@ def _getProcessorsForJobs(self, kwargs):
 
         return requestedProcessors
 
+    def _getMemoryForJobs(self, kwargs):
+        """helper function to get the memory that will be allocated for the job
+
+        :param kwargs: job parameters
+        :return: memory in GB or None if not enough memory
+        """
+
+        # # job requirements
+        requestedMemory = kwargs.get("MaxRAM", 0)
+
+        # # now check what the slot can provide
+        # Do we have enough memory?
+        availableMemory = self.ram - sum(self.ramPerTask.values())
+        if availableMemory < requestedMemory:
+            return None
+
+        return requestedMemory
+
     def finalizeJob(self, taskID, future):
         """Finalize the job by updating the process utilisation counters
 
@@ -209,7 +231,7 @@ def getCEStatus(self):
         result["WaitingJobs"] = 0
 
         # dealing with processors
-        processorsInUse = self.getProcessorsInUse()
+        processorsInUse = sum(self.processorsPerTask.values())
         result["UsedProcessors"] = processorsInUse
         result["AvailableProcessors"] = self.processors - processorsInUse
         return result

src/DIRAC/Resources/Computing/test/Test_PoolComputingElement.py

@@ -83,7 +83,7 @@ def createAndDelete():
 def test_submit_and_shutdown(createAndDelete):
     time.sleep(0.5)
 
-    ceParameters = {"WholeNode": True, "NumberOfProcessors": 4}
+    ceParameters = {"WholeNode": True, "NumberOfProcessors": 4, "MaxRAM": 4}
     ce = PoolComputingElement("TestPoolCE")
     ce.setParameters(ceParameters)
 
@@ -371,28 +371,37 @@ def test_executeJob_WholeNodeJobs(createAndDelete):
 
 
 @pytest.mark.parametrize(
-    "processorsPerTask, kwargs, expected",
+    "processorsPerTask, ramPerTask, kwargs, expected_processors, expected_memory",
     [
-        (None, {}, 1),
-        (None, {"mpTag": False}, 1),
-        (None, {"mpTag": True}, 1),
-        (None, {"mpTag": True, "wholeNode": True}, 16),
-        (None, {"mpTag": True, "wholeNode": False}, 1),
-        (None, {"mpTag": True, "numberOfProcessors": 4}, 4),
-        (None, {"mpTag": True, "numberOfProcessors": 4, "maxNumberOfProcessors": 8}, 8),
-        (None, {"mpTag": True, "numberOfProcessors": 4, "maxNumberOfProcessors": 32}, 16),
-        ({1: 4}, {"mpTag": True, "wholeNode": True}, 0),
-        ({1: 4}, {"mpTag": True, "wholeNode": False}, 1),
-        ({1: 4}, {"mpTag": True, "numberOfProcessors": 2}, 2),
-        ({1: 4}, {"mpTag": True, "maxNumberOfProcessors": 2}, 2),
-        ({1: 4}, {"mpTag": True, "maxNumberOfProcessors": 16}, 12),
+        (None, None, {}, 1, 0),
+        (None, None, {"mpTag": False}, 1, 0),
+        (None, None, {"mpTag": True, "MaxRAM": 8}, 1, 8),
+        (None, None, {"mpTag": True, "wholeNode": True}, 16, 0),
+        (None, None, {"mpTag": True, "wholeNode": False}, 1, 0),
+        (None, None, {"mpTag": True, "numberOfProcessors": 4, "MaxRAM": 4}, 4, 4),
+        (None, None, {"mpTag": True, "numberOfProcessors": 4, "maxNumberOfProcessors": 8}, 8, 0),
+        (None, None, {"mpTag": True, "numberOfProcessors": 4, "maxNumberOfProcessors": 32}, 16, 0),
+        ({1: 4}, {1: 4}, {"mpTag": True, "wholeNode": True}, 0, 0),
+        ({1: 4}, {1: 4}, {"mpTag": True, "wholeNode": False}, 1, 0),
+        ({1: 4}, {1: 4}, {"mpTag": True, "numberOfProcessors": 2, "MaxRAM": 8}, 2, 8),
+        ({1: 4}, {1: 4}, {"mpTag": True, "numberOfProcessors": 16, "MaxRAM": 12}, 0, 12),
+        ({1: 4}, {1: 4}, {"mpTag": True, "maxNumberOfProcessors": 2, "MaxRAM": 16}, 2, 16),
+        ({1: 4}, {1: 4}, {"mpTag": True, "maxNumberOfProcessors": 16, "MaxRAM": 32}, 12, None),
+        ({1: 4, 2: 8}, {1: 4}, {"mpTag": True, "numberOfProcessors": 2}, 2, 0),
+        ({1: 4, 2: 8}, {1: 4}, {"mpTag": True, "numberOfProcessors": 4}, 4, 0),
+        ({1: 4, 2: 8, 3: 8}, {1: 4}, {"mpTag": True, "numberOfProcessors": 4}, 0, 0),
     ],
 )
-def test__getProcessorsForJobs(processorsPerTask, kwargs, expected):
+def test__getLimitsForJobs(processorsPerTask, ramPerTask, kwargs, expected_processors, expected_memory):
     ce = PoolComputingElement("TestPoolCE")
     ce.processors = 16
+    ce.ram = 32
 
     if processorsPerTask:
         ce.processorsPerTask = processorsPerTask
+    if ramPerTask:
+        ce.ramPerTask = ramPerTask
     res = ce._getProcessorsForJobs(kwargs)
-    assert res == expected
+    assert res == expected_processors
+    res = ce._getMemoryForJobs(kwargs)
+    assert res == expected_memory