fix: getting cpu work left from a single source of truth

Author: aldbr

src/DIRAC/WorkloadManagementSystem/Agent/JobAgent.py

@@ -27,7 +27,6 @@
 from DIRAC.RequestManagementSystem.Client.ReqClient import ReqClient
 from DIRAC.RequestManagementSystem.Client.Request import Request
 from DIRAC.RequestManagementSystem.private.RequestValidator import RequestValidator
-from DIRAC.Resources.Computing.BatchSystems.TimeLeft.TimeLeft import TimeLeft
 from DIRAC.Resources.Computing.ComputingElementFactory import ComputingElementFactory
 from DIRAC.WorkloadManagementSystem.Client import JobStatus, PilotStatus
 from DIRAC.WorkloadManagementSystem.Client.JobManagerClient import JobManagerClient
@@ -66,7 +65,7 @@ def __init__(self, agentName, loadName, baseAgentName=False, properties=None):
 
         # Agent options
         # This is the factor to convert raw CPU to Normalized units (based on the CPU Model)
-        self.cpuFactor = 0.0
+        self.cpuPower = 0.0
         self.jobSubmissionDelay = 10
         self.fillingMode = True
         self.minimumTimeLeft = 5000
@@ -80,11 +79,10 @@ def __init__(self, agentName, loadName, baseAgentName=False, properties=None):
         self.logLevel = "INFO"
         self.defaultWrapperLocation = "DIRAC/WorkloadManagementSystem/JobWrapper/JobWrapperTemplate.py"
 
-        # Timeleft
-        self.initTimes = os.times()
-        self.initTimeLeft = 0.0
-        self.timeLeft = self.initTimeLeft
-        self.timeLeftUtil = None
+        # CPU work left (Wall-clock time * CPU Power)
+        self.initCPUWork = 0.0
+        self.timeLeft = self.initCPUWork
+        self.initTime = time.time()
         self.pilotInfoReportedFlag = False
 
         # Attributes related to the processed jobs, it should take the following form:
@@ -109,23 +107,18 @@ def initialize(self):
         if not result["OK"]:
             return result
 
-        result = self._getCEDict(self.computingElement)
-        if not result["OK"]:
-            return result
-        ceDict = result["Value"][0]
+        # Read initial CPU work left from config (seeded by pilot via dirac-wms-get-queue-cpu-time)
+        self.initCPUWork = gConfig.getValue("/LocalSite/CPUTimeLeft", self.initCPUWork)
+        self.timeLeft = self.initCPUWork
 
-        self.initTimeLeft = ceDict.get("CPUTime", self.initTimeLeft)
-        self.initTimeLeft = gConfig.getValue("/Resources/Computing/CEDefaults/MaxCPUTime", self.initTimeLeft)
-        self.timeLeft = self.initTimeLeft
-
-        self.initTimes = os.times()
+        self.initTime = time.time()
         # Localsite options
         self.siteName = siteName()
         self.pilotReference = gConfig.getValue("/LocalSite/PilotReference", self.pilotReference)
         self.defaultProxyLength = gConfig.getValue("/Registry/DefaultProxyLifeTime", self.defaultProxyLength)
         # Agent options
         # This is the factor to convert raw CPU to Normalized units (based on the CPU Model)
-        self.cpuFactor = gConfig.getValue("/LocalSite/CPUNormalizationFactor", self.cpuFactor)
+        self.cpuPower = gConfig.getValue("/LocalSite/CPUNormalizationFactor", self.cpuPower)
         self.jobSubmissionDelay = self.am_getOption("SubmissionDelay", self.jobSubmissionDelay)
         self.fillingMode = self.am_getOption("FillingModeFlag", self.fillingMode)
         self.minimumTimeLeft = self.am_getOption("MinimumTimeLeft", self.minimumTimeLeft)
@@ -136,9 +129,6 @@ def initialize(self):
         self.logLevel = self.am_getOption("DefaultLogLevel", self.logLevel)
         self.defaultWrapperLocation = self.am_getOption("JobWrapperTemplate", self.defaultWrapperLocation)
 
-        # Utilities
-        self.timeLeftUtil = TimeLeft()
-
         # Some innerCEs may want to make use of CGroup2 support, so we prepare it globally here
         res = CG2Manager().setUp()
         if res["OK"]:
@@ -180,15 +170,17 @@ def execute(self):
         if result["OK"] and result["Value"]:
             return result
 
-        # Check that we are allowed to continue and that time left is sufficient
+        # Update CPU work left: wall-clock is ticking whether a job is running or not
+        cpuWorkLeft = self._computeCPUWorkLeft()
+        result = self._setCPUWorkLeft(cpuWorkLeft)
+        if not result["OK"]:
+            return result
+
+        # After the first job, check filling mode eligibility
         if self.jobCount:
-            cpuWorkLeft = self._computeCPUWorkLeft()
             result = self._checkCPUWorkLeft(cpuWorkLeft)
             if not result["OK"]:
                 return result
-            result = self._setCPUWorkLeft(cpuWorkLeft)
-            if not result["OK"]:
-                return result
 
         # Get environment details and enhance them
         result = self._getCEDict(self.computingElement)
@@ -373,7 +365,7 @@ def _setCEDict(self, ceDict):
             ceDict["GridCE"] = gridCE
         if "PilotReference" not in ceDict:
             ceDict["PilotReference"] = str(self.pilotReference)
-        ceDict["PilotBenchmark"] = self.cpuFactor
+        ceDict["PilotBenchmark"] = self.cpuPower
         ceDict["PilotInfoReportedFlag"] = self.pilotInfoReportedFlag
 
         # Add possible job requirements
@@ -403,28 +395,22 @@ def _checkCEAvailability(self, computingElement):
         return S_OK()
 
     #############################################################################
-    def _computeCPUWorkLeft(self, processors=1):
+    def _computeCPUWorkLeft(self):
         """
-        Compute CPU Work Left in hepspec06 seconds
+        Compute CPU Work Left in hepspec06 seconds.
+
+        Uses a simple wall-clock countdown from the initial value (seeded by the pilot
+        via dirac-wms-get-queue-cpu-time). The elapsed wall-clock time is multiplied by
+        the CPU normalization factor to get the consumed CPU work.
 
-        :param int processors: number of processors available
         :return: cpu work left (cpu time left * cpu power of the cpus)
         """
-        # Sum all times but the last one (elapsed_time) and remove times at init (is this correct?)
-        cpuTimeConsumed = sum(os.times()[:-1]) - sum(self.initTimes[:-1])
-        result = self.timeLeftUtil.getTimeLeft(cpuTimeConsumed, processors)
-        if not result["OK"]:
-            self.log.warn("There were errors calculating time left using the Timeleft utility", result["Message"])
-            self.log.warn("The time left will be calculated using os.times() and the info in our possession")
-            self.log.info(f"Current raw CPU time consumed is {cpuTimeConsumed}")
-            if self.cpuFactor:
-                return self.initTimeLeft - cpuTimeConsumed * self.cpuFactor
-            return self.timeLeft
-        return result["Value"]
+        elapsed = time.time() - self.initTime
+        cpuWorkConsumed = elapsed * self.cpuPower
+        return self.initCPUWork - cpuWorkConsumed
 
     def _checkCPUWorkLeft(self, cpuWorkLeft):
         """Check that fillingMode is enabled and time left is sufficient to continue the execution"""
-        # Only call timeLeft utility after a job has been picked up
         self.log.info("Attempting to check CPU time left for filling mode")
         if not self.fillingMode:
             return self._finish("Filling Mode is Disabled")
@@ -435,7 +421,7 @@ def _checkCPUWorkLeft(self, cpuWorkLeft):
         return S_OK()
 
     def _setCPUWorkLeft(self, cpuWorkLeft):
-        """Update the TimeLeft within the CE and the configuration for next matching request"""
+        """Update the CPU work left within the CE and the configuration for next matching request"""
         self.timeLeft = cpuWorkLeft
 
         result = self.computingElement.setCPUTimeLeft(cpuTimeLeft=self.timeLeft)

src/DIRAC/WorkloadManagementSystem/Agent/test/Test_Agent_JobAgent.py

@@ -11,7 +11,6 @@
 from DIRAC.Core.Security.X509Chain import X509Chain  # pylint: disable=import-error
 
 from DIRAC import S_ERROR, S_OK, gLogger
-from DIRAC.Resources.Computing.BatchSystems.TimeLeft.TimeLeft import TimeLeft
 from DIRAC.Resources.Computing.ComputingElementFactory import ComputingElementFactory
 from DIRAC.Resources.Computing.test.Test_PoolComputingElement import badJobScript, jobScript
 from DIRAC.WorkloadManagementSystem.Agent.JobAgent import JobAgent
@@ -150,28 +149,27 @@ def test__checkCEAvailability(mocker, ceType, mockCEReply, expectedResult):
 
 
 @pytest.mark.parametrize(
-    "initTimeLeft, timeLeft, cpuFactor, mockTimeLeftReply, expectedTimeLeft",
+    "initCPUWork, cpuPower, elapsedSeconds, expectedTimeLeft",
     [
-        (100000, 75000, None, {"OK": False, "Message": "Error"}, 75000),
-        (100000, 75000, 10, {"OK": False, "Message": "Error"}, 100000),
-        (100000, 75000, 10, {"OK": True, "Value": 25000}, 25000),
+        # No CPU power: no work consumed, time left equals initial
+        (100000, 0, 100, 100000),
+        # With CPU power: elapsed * cpuPower is subtracted from initCPUWork
+        (100000, 10, 100, 99000),
+        # Longer elapsed time
+        (100000, 10, 5000, 50000),
     ],
 )
-def test__computeCPUWorkLeft(mocker, initTimeLeft, timeLeft, cpuFactor, mockTimeLeftReply, expectedTimeLeft):
+def test__computeCPUWorkLeft(mocker, initCPUWork, cpuPower, elapsedSeconds, expectedTimeLeft):
     """Test JobAgent()._computeCPUWorkLeft()"""
     mocker.patch("DIRAC.WorkloadManagementSystem.Agent.JobAgent.AgentModule.__init__")
-    mocker.patch(
-        "DIRAC.Resources.Computing.BatchSystems.TimeLeft.TimeLeft.TimeLeft.getTimeLeft", return_value=mockTimeLeftReply
-    )
 
     jobAgent = JobAgent("Test", "Test1")
     jobAgent.log = gLogger
     jobAgent.log.setLevel("DEBUG")
-    jobAgent.timeLeftUtil = TimeLeft()
 
-    jobAgent.initTimeLeft = initTimeLeft
-    jobAgent.timeLeft = timeLeft
-    jobAgent.cpuFactor = cpuFactor
+    jobAgent.initCPUWork = initCPUWork
+    jobAgent.cpuPower = cpuPower
+    jobAgent.initTime = time.time() - elapsedSeconds
     result = jobAgent._computeCPUWorkLeft()
 
     assert abs(result - expectedTimeLeft) < 10

src/DIRAC/WorkloadManagementSystem/Client/CPUNormalization.py

@@ -14,15 +14,11 @@
 
 
 def getCPUTime(cpuNormalizationFactor):
-    """Trying to get CPUTime left for execution (in seconds).
+    """Compute the initial CPUTime left for execution (in seconds).
 
-    It will first look to get the work left looking for batch system information useing the TimeLeft utility.
-    If it succeeds, it will convert it in real second, and return it.
-
-    If it fails, it tries to get it from the static info found in CS.
-    If it fails, it returns the default, which is a large 9999999, that we may consider as "Infinite".
-
-    This is a generic method, independent from the middleware of the resource if TimeLeft doesn't return a value
+    This is called at pilot bootstrap (via dirac-wms-get-queue-cpu-time) to seed
+    the initial CPUTimeLeft value. It queries the batch system first, then falls
+    back to static CS configuration.
 
     args:
       cpuNormalizationFactor (float): the CPU power of the current Worker Node.
@@ -31,55 +27,58 @@ def getCPUTime(cpuNormalizationFactor):
     returns:
       cpuTimeLeft (int): the CPU time left, in seconds
     """
-    cpuTimeLeft = 0.0
-    cpuWorkLeft = gConfig.getValue("/LocalSite/CPUTimeLeft", 0)
-
-    if not cpuWorkLeft:
-        # Try and get the information from the CPU left utility
-        result = TimeLeft().getTimeLeft()
-        if result["OK"]:
-            cpuWorkLeft = result["Value"]
-
-    if cpuWorkLeft > 0:
-        # This is in HS06sseconds
-        # We need to convert in real seconds
-        if not cpuNormalizationFactor:  # if cpuNormalizationFactor passed in is 0, try get it from the local cfg
+
+    # 1. Try to compute time left from the batch system (sacct, qstat, etc.)
+    result = TimeLeft().getTimeLeft()
+    if result["OK"]:
+        cpuWorkLeft = result["Value"]
+        # Batch system answered — trust it, even if 0
+        if not cpuNormalizationFactor:
             cpuNormalizationFactor = gConfig.getValue("/LocalSite/CPUNormalizationFactor", 0.0)
         if cpuNormalizationFactor:
-            cpuTimeLeft = cpuWorkLeft / cpuNormalizationFactor
+            return int(cpuWorkLeft / cpuNormalizationFactor)
+        return 0
 
-    if not cpuTimeLeft:
-        # now we know that we have to find the CPUTimeLeft by looking in the CS
-        # this is not granted to be correct as the CS units may not be real seconds
-        gridCE = gConfig.getValue("/LocalSite/GridCE")
-        ceQueue = gConfig.getValue("/LocalSite/CEQueue")
-        if not ceQueue:
-            # we have to look for a ceQueue in the CS
-            # A bit hacky. We should better profit from something generic
-            gLogger.warn("No CEQueue in local configuration, looking to find one in CS")
-            siteName = DIRAC.siteName()
-            queueSection = f"/Resources/Sites/{siteName.split('.')[0]}/{siteName}/CEs/{gridCE}/Queues"
-            res = gConfig.getSections(queueSection)
-            if not res["OK"]:
-                raise RuntimeError(res["Message"])
-            queues = res["Value"]
-            cpuTimes = [gConfig.getValue(queueSection + "/" + queue + "/maxCPUTime", 9999999.0) for queue in queues]
-            # These are (real, wall clock) minutes - damn BDII!
+    cpuTimeLeft = 0.0
+
+    # 2. Fall back to queue configuration in the CS.
+    # These values are wall-clock minutes from BDII, so we convert to seconds.
+    gridCE = gConfig.getValue("/LocalSite/GridCE")
+    ceQueue = gConfig.getValue("/LocalSite/CEQueue")
+    if not ceQueue:
+        # we have to look for a ceQueue in the CS
+        # A bit hacky. We should better profit from something generic
+        gLogger.warn("No CEQueue in local configuration, looking to find one in CS")
+        siteName = DIRAC.siteName()
+        queueSection = f"/Resources/Sites/{siteName.split('.')[0]}/{siteName}/CEs/{gridCE}/Queues"
+        res = gConfig.getSections(queueSection)
+        if not res["OK"]:
+            raise RuntimeError(res["Message"])
+        queues = res["Value"]
+        cpuTimes = [gConfig.getValue(queueSection + "/" + queue + "/maxCPUTime", 0.0) for queue in queues]
+        cpuTimes = [t for t in cpuTimes if t > 0]
+        if cpuTimes:
             cpuTimeLeft = min(cpuTimes) * 60
+    else:
+        queueInfo = getQueueInfo(f"{gridCE}/{ceQueue}")
+        if not queueInfo["OK"] or not queueInfo["Value"]:
+            gLogger.warn("Can't find a CE/queue in CS")
         else:
-            queueInfo = getQueueInfo(f"{gridCE}/{ceQueue}")
-            cpuTimeLeft = 9999999.0
-            if not queueInfo["OK"] or not queueInfo["Value"]:
-                gLogger.warn("Can't find a CE/queue, defaulting CPUTime to %d" % cpuTimeLeft)
+            queueCSSection = queueInfo["Value"]["QueueCSSection"]
+            cpuTimeInMinutes = gConfig.getValue(f"{queueCSSection}/maxCPUTime", 0.0)
+            if cpuTimeInMinutes:
+                cpuTimeLeft = cpuTimeInMinutes * 60.0
+                gLogger.info(f"CPUTime for {queueCSSection}: {cpuTimeLeft:f}")
             else:
-                queueCSSection = queueInfo["Value"]["QueueCSSection"]
-                # These are (real, wall clock) minutes - damn BDII!
-                cpuTimeInMinutes = gConfig.getValue(f"{queueCSSection}/maxCPUTime", 0.0)
-                if cpuTimeInMinutes:
-                    cpuTimeLeft = cpuTimeInMinutes * 60.0
-                    gLogger.info(f"CPUTime for {queueCSSection}: {cpuTimeLeft:f}")
-                else:
-                    gLogger.warn(f"Can't find maxCPUTime for {queueCSSection}, defaulting CPUTime to {cpuTimeLeft:f}")
+                gLogger.warn(f"Can't find maxCPUTime for {queueCSSection}")
+
+    if not cpuTimeLeft:
+        # 3. Last resort: global default from CS, or 0 (fail safe: match no more jobs)
+        cpuTimeLeft = gConfig.getValue("/Resources/Computing/CEDefaults/MaxCPUTime", 0)
+        if cpuTimeLeft:
+            gLogger.warn(f"Using fallback MaxCPUTime: {cpuTimeLeft}")
+        else:
+            gLogger.warn("Could not determine CPUTime left")
 
     return int(cpuTimeLeft)