feat: Computing Elements use MaxRAM

Author: fstagni

docs/source/UserGuide/Tutorials/JobManagementAdvanced/index.rst

@@ -335,16 +335,15 @@ Jobs that can (or should) run using more than 1 processor should be described as
 using the "setNumberOfProcessors" method of the API::
 
       j = Job()
-      j.setCPUTime(500)
-      j.setExecutable('echo',arguments='hello')
-      j.setExecutable('ls',arguments='-l')
-      j.setExecutable('echo', arguments='hello again')
-      j.setName('MP test')
+      ...
       j.setNumberOfProcessors(16)
 
 Calling ``Job().setNumberOfProcessors()``, with a value bigger than 1,
 will translate into adding also the "MultiProcessor" tag to the job description.
 
+``Job().setNumberOfProcessors()`` takes at most 3 arguments, in this order:
+``numberOfProcessors`` is the exact number of requested processors, ``minNumberOfProcessors`` is the minimum allowed, ``maxNumberOfProcessors`` the maximum.
+
 Users can specify in the job descriptions NumberOfProcessors and WholeNode parameters, e.g.::
 
    NumberOfProcessors = 16;
@@ -356,6 +355,20 @@ This will be translated internally into 16Processors and WholeNode tags.
 This would allow resources (WN's) to put flexibly requirements on jobs to be taken, for example, avoiding single-core jobs on a multi-core nodes.
 
 
+Setting memory limits
+@@@@@@@@@@@@@@@@@@@@@
+
+Jobs can (and probably should) set RAM limits.
+using the "setRAMRequrements" method of the API::
+
+      j = Job()
+      ...
+      j.setRAMRequrements(2, 4)
+
+Calling ``Job().setRAMRequirements()`` takes 2 values, where the first is the minimum required amount of RAM (in GB) that the job requests.
+The second value instead specifies the limit that should not be surpassed.
+
+
 Submitting jobs with specifc requirements (e.g. GPU)
 @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
 

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

@@ -20,6 +20,7 @@
 **Code Documentation**
 """
 import concurrent.futures
+import copy
 import functools
 
 from DIRAC import S_ERROR, S_OK
@@ -29,8 +30,8 @@
 from DIRAC.Resources.Computing.SingularityComputingElement import SingularityComputingElement
 
 
-def executeJob(executableFile, proxy, taskID, inputs, **kwargs):
-    """wrapper around ce.submitJob: decides which CE to use (InProcess or Singularity)
+def executeJob(executableFile, proxy, taskID, inputs, innerCEParameters, **kwargs):
+    """wrapper around ce.submitJob: decides which inner CE to use (InProcess or Singularity)
 
     :param str executableFile: location of the executable file
     :param str proxy: proxy file location to be used for job submission
@@ -46,6 +47,8 @@ def executeJob(executableFile, proxy, taskID, inputs, **kwargs):
     else:
         ce = InProcessComputingElement("Task-" + str(taskID))
 
+    ce.ceParameters["NumberOfProcessors"] = innerCEParameters["NumberOfProcessors"]
+    ce.ceParameters["MaxRAM"] = innerCEParameters["MaxRAM"]
     return ce.submitJob(executableFile, proxy, inputs=inputs, **kwargs)
 
 
@@ -60,8 +63,9 @@ def __init__(self, ceUniqueID):
         self.pPool = None
         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
@@ -76,10 +80,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 int(self.ceParameters.get("MaxRAM", 0)):
+            self.ram = self.ceParameters["MaxRAM"]
         # Indicates that the submission is done asynchronously
         # The result is not immediately available
         self.ceParameters["AsyncSubmission"] = True
@@ -128,12 +130,20 @@ def submitJob(self, executableFile, proxy=None, inputs=None, **kwargs):
             if not res["OK"]:
                 self.log.error("Could not dump cfg to pilot.cfg", res["Message"])
 
+        # Define the innerCEParameters
+        innerCEParameters = {}
+        innerCEParameters = copy.deepcopy(self.ceParameters)
+        innerCEParameters["NumberOfProcessors"] = processorsForJob
+        innerCEParameters["MaxRAM"] = memoryForJob
+
         # Here we define task kwargs: adding complex objects like thread.Lock can trigger errors in the task
         taskKwargs = {"InnerCESubmissionType": self.innerCESubmissionType}
         taskKwargs["jobDesc"] = kwargs.get("jobDesc", {})
 
         # Submission
-        future = self.pPool.submit(executeJob, executableFile, proxy, self.taskID, inputs, **taskKwargs)
+        future = self.pPool.submit(
+            executeJob, executableFile, proxy, self.taskID, inputs, innerCEParameters, **taskKwargs
+        )
         self.processorsPerTask[future] = processorsForJob
         self.ramPerTask[future] = memoryForJob
         future.add_done_callback(functools.partial(self.finalizeJob, self.taskID))
@@ -190,14 +200,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,23 +223,26 @@ 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
 
     def getCEStatus(self):
         """Method to return information on running and waiting jobs,
-            as well as the number of processors (used, and available).
+        as well as the number of processors (used, and available),
+        and the RAM (used, and available)
 
         :return: dictionary of numbers of jobs per status and processors (used, and available)
         """
 
         result = S_OK()
         nJobs = 0
+
         for _j, value in self.processorsPerTask.items():
             if value > 0:
                 nJobs += 1
@@ -234,6 +254,10 @@ def getCEStatus(self):
         processorsInUse = sum(self.processorsPerTask.values())
         result["UsedProcessors"] = processorsInUse
         result["AvailableProcessors"] = self.processors - processorsInUse
+        # dealing with RAM
+        result["RAMInUse"] = sum(self.ramPerTask.values())
+        result["RAMAvailable"] = self.ram - sum(self.ramPerTask.values())
+
         return result
 
     def getDescription(self):

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_InProcessComputingElement.py

@@ -52,6 +52,8 @@ def test_submitJob():
         maxNumberOfProcessors=8,
         wholeNode=False,
         mpTag=True,
+        MinRAM=2,
+        MaxRAM=4,
         jobDesc={"jobParams": jobParams, "resourceParams": resourceParams, "optimizerParams": optimizerParams},
     )
     assert res["OK"] is True

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

@@ -213,8 +213,8 @@ def getAvailableRAM(siteName=None, gridCE=None, queue=None):
     The siteName/gridCE/queue parameters are normally not necessary.
 
     Tries to find it in this order:
-    1) from the /Resources/Computing/CEDefaults/AvailableRAM (which is what the pilot might fill up)
-    2) if not present looks in CS for "AvailableRAM" Queue or CE option
+    1) from the /Resources/Computing/CEDefaults/MaxRAM (which is what the pilot might fill up)
+    2) if not present looks in CS for "MemoryLimitMB" Queue or CE or site option
     3) if not present but there's WholeNode tag, look what the WN provides using _getMemoryFromProc()
     4) return 0
     """
@@ -238,9 +238,9 @@ def getAvailableRAM(siteName=None, gridCE=None, queue=None):
 
     grid = siteName.split(".")[0]
     csPaths = [
-        f"/Resources/Sites/{grid}/{siteName}/CEs/{gridCE}/Queues/{queue}/MaxRAM",
-        f"/Resources/Sites/{grid}/{siteName}/CEs/{gridCE}/MaxRAM",
-        f"/Resources/Sites/{grid}/{siteName}/MaxRAM",
+        f"/Resources/Sites/{grid}/{siteName}/CEs/{gridCE}/Queues/{queue}/MemoryLimitMB",
+        f"/Resources/Sites/{grid}/{siteName}/CEs/{gridCE}/MemoryLimitMB",
+        f"/Resources/Sites/{grid}/{siteName}/MemoryLimitMB",
     ]
     for csPath in csPaths:
         gLogger.info("Looking in", csPath)
@@ -265,5 +265,22 @@ def getAvailableRAM(siteName=None, gridCE=None, queue=None):
         return _getMemoryFromProc()
 
     # 4) return 0
-    gLogger.info("AvailableRAM could not be found in CS, and WholeNode tag not found")
+    gLogger.info("RAM limits 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()