feat: created a utility function getAvailableRAM

Author: fstagni

src/DIRAC/Resources/Computing/ComputingElement.py

@@ -57,6 +57,7 @@
 from DIRAC.WorkloadManagementSystem.Utilities.JobParameters import (
     getNumberOfProcessors,
     getNumberOfGPUs,
+    getAvailableRAM
 )
 
 INTEGER_PARAMETERS = ["CPUTime", "NumberOfProcessors", "NumberOfPayloadProcessors", "MaxRAM"]
@@ -235,12 +236,14 @@ def setParameters(self, ceOptions):
                 generalCEDict.update(self.ceParameters)
                 self.ceParameters = generalCEDict
 
-        # If NumberOfProcessors/GPUs is present in the description but is equal to zero
+        # If NumberOfProcessors/GPUs/RAM 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()
 
         for key in ceOptions:
             if key in INTEGER_PARAMETERS:

src/DIRAC/WorkloadManagementSystem/Utilities/JobParameters.py

@@ -34,31 +34,7 @@ def getJobFeatures():
     return features
 
 
-def getProcessorFromMJF():
-    jobFeatures = getJobFeatures()
-    if jobFeatures:
-        try:
-            return int(jobFeatures["allocated_cpu"])
-        except KeyError:
-            gLogger.error(
-                "MJF is available but allocated_cpu is not an integer", repr(jobFeatures.get("allocated_cpu"))
-            )
-    return None
-
-
-def getMemoryFromMJF():
-    jobFeatures = getJobFeatures()
-    if jobFeatures:
-        try:
-            return int(jobFeatures["max_rss_bytes"])
-        except KeyError:
-            gLogger.error(
-                "MJF is available but max_rss_bytes is not an integer", repr(jobFeatures.get("max_rss_bytes"))
-            )
-    return None
-
-
-def getMemoryFromProc():
+def _getMemoryFromProc():
     meminfo = {i.split()[0].rstrip(":"): int(i.split()[1]) for i in open("/proc/meminfo").readlines()}
     maxRAM = meminfo["MemTotal"]
     if maxRAM:
@@ -72,7 +48,6 @@ def getNumberOfProcessors(siteName=None, gridCE=None, queue=None):
 
     Tries to find it in this order:
     1) from the /Resources/Computing/CEDefaults/NumberOfProcessors (which is what the pilot fills up)
-    2) if not present from JobFeatures
     3) if not present looks in CS for "NumberOfProcessors" Queue or CE option
     4) if not present but there's WholeNode tag, look what the WN provides using multiprocessing.cpu_count()
     5) return 1
@@ -84,14 +59,7 @@ def getNumberOfProcessors(siteName=None, gridCE=None, queue=None):
     if numberOfProcessors:
         return numberOfProcessors
 
-    # 2) from MJF
-    gLogger.info("Getting numberOfProcessors from MJF")
-    numberOfProcessors = getProcessorFromMJF()
-    if numberOfProcessors:
-        return numberOfProcessors
-    gLogger.info("NumberOfProcessors could not be found in MJF")
-
-    # 3) looks in CS for "NumberOfProcessors" Queue or CE or site option
+    # 2) looks in CS for "NumberOfProcessors" Queue or CE or site option
     if not siteName:
         siteName = gConfig.getValue("/LocalSite/Site", "")
     if not gridCE:
@@ -237,3 +205,65 @@ def getNumberOfGPUs(siteName=None, gridCE=None, queue=None):
     # 3) return 0
     gLogger.info("NumberOfGPUs could not be found in CS")
     return 0
+
+
+def getAvailableRAM(siteName=None, gridCE=None, queue=None):
+    """Gets the available RAM on a certain CE/queue/node (what the pilot administers)
+
+    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
+    3) if not present but there's WholeNode tag, look what the WN provides using _getMemoryFromProc()
+    4) return 0
+    """
+
+    # 1) from /Resources/Computing/CEDefaults/MaxRAM
+    gLogger.info("Getting MaxRAM from /Resources/Computing/CEDefaults/MaxRAM")
+    availableRAM = gConfig.getValue("/Resources/Computing/CEDefaults/MaxRAM", None)
+    if availableRAM:
+        return availableRAM
+
+    # 2) looks in CS for "MaxRAM" Queue or CE or site option
+    if not siteName:
+        siteName = gConfig.getValue("/LocalSite/Site", "")
+    if not gridCE:
+        gridCE = gConfig.getValue("/LocalSite/GridCE", "")
+    if not queue:
+        queue = gConfig.getValue("/LocalSite/CEQueue", "")
+    if not (siteName and gridCE and queue):
+        gLogger.error("Could not find AvailableRAM: missing siteName or gridCE or queue. Returning 0")
+        return 0
+
+    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",
+    ]
+    for csPath in csPaths:
+        gLogger.info("Looking in", csPath)
+        availableRAM = gConfig.getValue(csPath, None)
+        if availableRAM:
+            return availableRAM
+        
+    # 3) checks if 'WholeNode' is one of the used tags
+    # Tags of the CE
+    tags = fromChar(
+        gConfig.getValue(f"/Resources/Sites/{siteName.split('.')[0]}/{siteName}/CEs/{gridCE}/Tag", "")
+    ) + fromChar(gConfig.getValue(f"/Resources/Sites/{siteName.split('.')[0]}/{siteName}/Cloud/{gridCE}/Tag", ""))
+    # Tags of the Queue
+    tags += fromChar(
+        gConfig.getValue(f"/Resources/Sites/{siteName.split('.')[0]}/{siteName}/CEs/{gridCE}/Queues/{queue}/Tag", "")
+    ) + fromChar(
+        gConfig.getValue(f"/Resources/Sites/{siteName.split('.')[0]}/{siteName}/Cloud/{gridCE}/VMTypes/{queue}/Tag", "")
+    )
+
+    if "WholeNode" in tags:
+        gLogger.info("Found WholeNode tag, using _getMemoryFromProc()")
+        return _getMemoryFromProc()
+    
+    # 4) return 0
+    gLogger.info("AvailableRAM could not be found in CS, and WholeNode tag not found")
+    return 0

src/DIRAC/WorkloadManagementSystem/scripts/dirac_wms_get_wn_parameters.py

@@ -39,15 +39,12 @@ def main():
     gLogger.info("Getting number of processors")
     numberOfProcessor = JobParameters.getNumberOfProcessors(Site, ceName, Queue)
 
-    gLogger.info("Getting memory (RAM) from MJF")
-    maxRAM = JobParameters.getMemoryFromMJF()
-    if not maxRAM:
-        gLogger.info("maxRAM could not be found in MJF, using JobParameters.getMemoryFromProc()")
-        maxRAM = JobParameters.getMemoryFromProc()
-
     gLogger.info("Getting number of GPUs")
     numberOfGPUs = JobParameters.getNumberOfGPUs(Site, ceName, Queue)
 
+    gLogger.info("Getting maximum RAM")
+    maxRAM = JobParameters.getAvailableRAM(Site, ceName, Queue)
+
     # just communicating it back
     gLogger.notice(" ".join(str(wnPar) for wnPar in [numberOfProcessor, maxRAM, numberOfGPUs]))