feat: PoolCE can subdivide RAM

Author: fstagni

src/DIRAC/Interfaces/API/Job.py

@@ -740,7 +740,7 @@ def setTag(self, tags):
         Example usage:
 
         >>> job = Job()
-        >>> job.setTag( ['WholeNode','8GB'] )
+        >>> job.setTag( ['WholeNode','8GB', 16GB_MAX] )
 
         :param tags: single tag string or a list of tags
         :type tags: str or python:list

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,6 +76,10 @@ 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
@@ -90,6 +96,16 @@ def getProcessorsInUse(self):
             processorsInUse += self.processorsPerTask[future]
         return processorsInUse
 
+    def getRAMInUse(self):
+        """Get the amount of RAM in use by the currently running jobs
+
+        :return: amount of RAM in use
+        """
+        ramInUse = 0
+        for future in self.ramPerTask:
+            ramInUse += self.ramPerTask[future]
+        return ramInUse
+
     #############################################################################
     def submitJob(self, executableFile, proxy=None, inputs=None, **kwargs):
         """Method to submit job.
@@ -112,15 +128,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 +155,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
@@ -178,6 +202,31 @@ 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 = 0
+
+        # If there's a memory tag, this should be considered as the requested memory by the job
+        for job_args in kwargs:
+            if job_args.endswith("GB"):
+                requestedMemory = int(job_args.replace("GB", ""))
+            if job_args.endswith("GB_MAX"):
+                requestedMemory = max(int(job_args.replace("GB_MAX", "")), requestedMemory)
+
+        # # now check what the slot can provide
+        # Do we have enough memory?
+        availableMemory = self.ram - self.getRAMInUse()
+        if availableMemory < requestedMemory:
+            return None
+
+        return requestedMemory
+
     def finalizeJob(self, taskID, future):
         """Finalize the job by updating the process utilisation counters
 

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,38 @@ 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, "8GB": True}, 1, 8),
+        (None, None, {"mpTag": True, "wholeNode": True}, 16, 0),
+        (None, None, {"mpTag": True, "wholeNode": False}, 1, 0),
+        (None, None, {"mpTag": True, "numberOfProcessors": 4, "4GB": True}, 4, 4),
+        (None, None, {"mpTag": True, "numberOfProcessors": 4, "maxNumberOfProcessors": 8}, 8, 0),
+        (None, None, {"mpTag": True, "4GB": True, "MaxRAM": 8}, 1, 8),
+        (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, "8GB": True}, 2, 8),
+        ({1: 4}, {1: 4}, {"mpTag": True, "numberOfProcessors": 16, "12GB": True}, 0, 12),
+        ({1: 4}, {1: 4}, {"mpTag": True, "maxNumberOfProcessors": 2, "16GB": True}, 2, 16),
+        ({1: 4}, {1: 4}, {"mpTag": True, "maxNumberOfProcessors": 16, "32GB": True}, 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