feat: Computing Elements use MaxRAM

Author: fstagni

src/DIRAC/Core/Utilities/CGroups2.py

@@ -300,6 +300,7 @@ def systemCall(self, *args, **kwargs):
             if "ceParameters" in kwargs:
                 if cpuLimit := kwargs["ceParameters"].get("CPULimit", None):
                     cores = float(cpuLimit)
+                # MemoryLimitMB should be the job upper limit
                 if memoryMB := int(kwargs["ceParameters"].get("MemoryLimitMB", 0)):
                     memory = memoryMB * 1024 * 1024
                 if kwargs["ceParameters"].get("MemoryNoSwap", "no").lower() in ("yes", "true"):

src/DIRAC/Interfaces/API/Job.py

@@ -522,18 +522,49 @@ def setDestination(self, destination):
         return S_OK()
 
     #############################################################################
-    def setRAMRequirements(self, ramRequired: int = 0):
+    def setRAMRequirements(self, ramRequired: int = 0, maxRAM: int = 0):
         """Helper function.
-        Specify the RAM requirements for the job in GB. 0 (default) means no specific requirements.
+        Specify the RAM requirements for the job. 0 (default) means no specific requirements.
+
+        Example usage:
+
+        >>> job = Job()
+        >>> job.setRAMRequirements(ramRequired=2)
+        means that the job needs at least 2 GBs of RAM to work. This is taken into consideration at job's matching time.
+        The job definition does not specify an upper limit.
+        From a user's point of view this is fine (normally, not for admins).
+
+        >>> job.setRAMRequirements(ramRequired=2, maxRAM=4)
+        means that the job needs 2 GBs of RAM to work. 4 GBs will then be the upper limit for CG2 limits.
+
+        >>> job.setRAMRequirements(ramRequired=4, maxRAM=4)
+        means that we should match this job if there is at least 4 available GBs of run. At the same time, CG2 will not allow to use more than that.
+
+        >>> job.setRAMRequirements(maxRAM=4)
+        means that the job does not set a min amount of RAM (so can match--run "everywhere"), but the 4 GBs will then be the upper limit for CG2 limits.
+
+        >>> job.setRAMRequirements(ramRequired=8, maxRAM=4)
+        Makes no sense, an error will be raised
         """
+        if ramRequired and maxRAM and ramRequired > maxRAM:
+            return self._reportError("Invalid settings, ramRequired is higher than maxRAM")
+
         if ramRequired:
             self._addParameter(
                 self.workflow,
-                "MaxRAM",
+                "MinRAM",
                 "JDL",
                 ramRequired,
                 "GBs of RAM requested",
             )
+        if maxRAM:
+            self._addParameter(
+                self.workflow,
+                "MaxRAM",
+                "JDL",
+                maxRAM,
+                "Max GBs of RAM to be used",
+            )
 
     def setNumberOfProcessors(self, numberOfProcessors=None, minNumberOfProcessors=None, maxNumberOfProcessors=None):
         """Helper function.

src/DIRAC/Resources/Computing/ComputingElement.py

@@ -236,14 +236,14 @@ def setParameters(self, ceOptions):
                 generalCEDict.update(self.ceParameters)
                 self.ceParameters = generalCEDict
 
-        # If NumberOfProcessors/GPUs/RAM is present in the description but is equal to zero
+        # If NumberOfProcessors/GPUs/MaxRAM is present in the description but is equal to zero
         # interpret it as needing local evaluation
         if self.ceParameters.get("NumberOfProcessors", -1) == 0:
             self.ceParameters["NumberOfProcessors"] = getNumberOfProcessors()
         if self.ceParameters.get("NumberOfGPUs", -1) == 0:
             self.ceParameters["NumberOfGPUs"] = getNumberOfGPUs()
-        if self.ceParameters.get("RAM", -1) == 0:
-            self.ceParameters["RAM"] = getAvailableRAM()
+        if self.ceParameters.get("MaxRAM", 0) == 0:
+            self.ceParameters["MaxRAM"] = getAvailableRAM()
 
         for key in ceOptions:
             if key in INTEGER_PARAMETERS:

src/DIRAC/Resources/Computing/InProcessComputingElement.py

@@ -22,6 +22,7 @@ def __init__(self, ceUniqueID):
         self.runningJobs = 0
 
         self.processors = int(self.ceParameters.get("NumberOfProcessors", 1))
+        self.maxRAM = int(self.ceParameters.get("MaxRAM", 0))
         self.ceParameters["MaxTotalJobs"] = 1
 
     def submitJob(self, executableFile, proxy=None, inputs=None, **kwargs):
@@ -118,4 +119,6 @@ def getCEStatus(self):
         result["WaitingJobs"] = 0
         # processors
         result["AvailableProcessors"] = self.processors
+        # RAM
+        result["MaxRAM"] = self.maxRAM
         return result

src/DIRAC/Resources/Computing/PoolComputingElement.py

@@ -61,7 +61,9 @@ 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.ram = (
+            1024  # Available RAM for the node, in GB. The default value is an arbitrary large value in case of no limit
+        )
         self.ramPerTask = {}
 
         # This CE will effectively submit to another "Inner"CE
@@ -77,9 +79,8 @@ 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
+        if max_ram:
+            self.ram = self.ceParameters["MaxRAM"]
         # Indicates that the submission is done asynchronously
         # The result is not immediately available
         self.ceParameters["AsyncSubmission"] = True
@@ -190,14 +191,21 @@ def _getMemoryForJobs(self, kwargs):
         """
 
         # # job requirements
-        requestedMemory = kwargs.get("MaxRAM", 0)
+        requestedMemory = kwargs.get("MinRAM", 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
 
+        # if there's a MaxRAM requested, we allocate it all (if it fits)
+        if kwargs.get("MaxRAM", 0):
+            if availableMemory >= kwargs.get("MaxRAM"):
+                requestedMemory = kwargs.get("MaxRAM")
+            else:
+                return None
+
         return requestedMemory
 
     def finalizeJob(self, taskID, future):
@@ -206,10 +214,11 @@ def finalizeJob(self, taskID, future):
         :param future: evaluating the future result
         """
         nProc = self.processorsPerTask.pop(future)
+        ram = self.ramPerTask.pop(future, None)
 
         result = future.result()  # This would be the result of the e.g. InProcess.submitJob()
         if result["OK"]:
-            self.log.info("Task finished successfully:", f"{taskID}; {nProc} processor(s) freed")
+            self.log.info("Task finished successfully:", f"{taskID}; {nProc} processor(s) and {ram}GB freed")
         else:
             self.log.error("Task failed submission:", f"{taskID}; message: {result['Message']}")
         self.taskResults[taskID] = result

src/DIRAC/Resources/Computing/SingularityComputingElement.py

@@ -115,6 +115,7 @@ def __init__(self, ceUniqueID):
             self.__installDIRACInContainer = False
 
         self.processors = int(self.ceParameters.get("NumberOfProcessors", 1))
+        self.maxRAM = int(self.ceParameters.get("MaxRAM", 0))
 
     @staticmethod
     def __findInstallBaseDir():
@@ -413,6 +414,12 @@ def submitJob(self, executableFile, proxy=None, **kwargs):
 
         self.log.debug(f"Execute singularity command: {cmd}")
         self.log.debug(f"Execute singularity env: {self.__getEnv()}")
+        # systemCall below uses ceParameters["MemoryLimitMB"] as CG2 upper memory limit
+        # if there's a max RAM available to the job, use that
+        if self.maxRAM:
+            self.ceParameters["MemoryLimitMB"] = min(
+                self.maxRAM * 1024, self.ceParameters.get("MemoryLimitMB", 1024 * 1024)
+            )  # 1024 * 1024 is an arbitrary large number
         result = CG2Manager().systemCall(
             0, cmd, callbackFunction=self.sendOutput, env=self.__getEnv(), ceParameters=self.ceParameters
         )
@@ -449,4 +456,6 @@ def getCEStatus(self):
         result["WaitingJobs"] = 0
         # processors
         result["AvailableProcessors"] = self.processors
+        # RAM
+        result["MaxRAM"] = self.maxRAM
         return result

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

@@ -378,14 +378,18 @@ def test_executeJob_WholeNodeJobs(createAndDelete):
         (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, "MinRAM": 2}, 4, 2),
         (None, None, {"mpTag": True, "numberOfProcessors": 4, "MaxRAM": 4}, 4, 4),
+        (None, None, {"mpTag": True, "numberOfProcessors": 4, "MaxRAM": 36}, 4, None),
+        (None, None, {"mpTag": True, "numberOfProcessors": 4, "MinRAM": 2, "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, "numberOfProcessors": 2, "MinRAM": 8}, 2, 8),
+        ({1: 4}, {1: 4}, {"mpTag": True, "numberOfProcessors": 16, "MinRAM": 8, "MaxRAM": 12}, 0, 12),
         ({1: 4}, {1: 4}, {"mpTag": True, "maxNumberOfProcessors": 2, "MaxRAM": 16}, 2, 16),
+        ({1: 4}, {1: 4}, {"mpTag": True, "numberOfProcessors": 2, "MaxRAM": 8}, 2, 8),
         ({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),

src/DIRAC/WorkloadManagementSystem/Utilities/JobParameters.py

@@ -267,3 +267,20 @@ def getAvailableRAM(siteName=None, gridCE=None, queue=None):
     # 4) return 0
     gLogger.info("AvailableRAM could not be found in CS, and WholeNode tag not found")
     return 0
+
+
+def getRAMForJob(jobID):
+    """Gets the RAM allowed for the job.
+    This can be used to communicate to your job payload the RAM it's allowed to use,
+    so this function should be called from your extension.
+
+    If the JobAgent is using "InProcess" CE (which is the default),
+    then what's returned will basically be the same of what's returned by the getAvailableRAM() function above
+    """
+
+    # from /Resources/Computing/JobLimits/jobID/MaxRAM (set by PoolComputingElement)
+    ram = gConfig.getValue(f"Resources/Computing/JobLimits/{jobID}/MaxRAM")
+    if ram:
+        return ram
+
+    return getAvailableRAM()