Merge pull request #8192 from fstagni/90_wms_fixes

[9.0] StalledJobAgent improvements

Author: chrisburr

src/DIRAC/WorkloadManagementSystem/Agent/PushJobAgent.py

@@ -12,16 +12,14 @@
 import hashlib
 import json
 import os
-from pathlib import Path
 import random
 import shutil
 import sys
-from collections import defaultdict
 import time
+from collections import defaultdict
+from pathlib import Path
 
-from diraccfg import CFG
-
-from DIRAC import gConfig, S_OK, S_ERROR
+from DIRAC import S_ERROR, S_OK, gConfig
 from DIRAC.ConfigurationSystem.Client.Helpers.Operations import Operations
 from DIRAC.ConfigurationSystem.Client.Helpers.Resources import getQueues
 from DIRAC.Core.Utilities import DErrno
@@ -30,6 +28,7 @@
 from DIRAC.Core.Utilities.Version import getVersion
 from DIRAC.FrameworkSystem.Client.ProxyManagerClient import gProxyManager
 from DIRAC.Resources.Computing import ComputingElement
+from DIRAC.WorkloadManagementSystem.Agent.JobAgent import JobAgent
 from DIRAC.WorkloadManagementSystem.Client import JobMinorStatus, JobStatus, PilotStatus
 from DIRAC.WorkloadManagementSystem.Client.JobMonitoringClient import JobMonitoringClient
 from DIRAC.WorkloadManagementSystem.Client.JobReport import JobReport
@@ -41,11 +40,10 @@
     resolveInputData,
     transferInputSandbox,
 )
-from DIRAC.WorkloadManagementSystem.Utilities.QueueUtilities import getQueuesResolved
-from DIRAC.WorkloadManagementSystem.Agent.JobAgent import JobAgent
-from DIRAC.WorkloadManagementSystem.Utilities.Utils import createJobWrapper
 from DIRAC.WorkloadManagementSystem.private.ConfigHelper import findGenericPilotCredentials
+from DIRAC.WorkloadManagementSystem.Utilities.JobParameters import getJobParameters
 from DIRAC.WorkloadManagementSystem.Utilities.QueueUtilities import getQueuesResolved
+from DIRAC.WorkloadManagementSystem.Utilities.Utils import createJobWrapper
 
 MAX_JOBS_MANAGED = 100
 
@@ -740,7 +738,7 @@ def _checkSubmittedJobWrappers(self, ce: ComputingElement, site: str):
             return S_OK()
 
         # Get their parameters
-        if not (result := self.jobMonitoring.getJobParameters(jobs, ["GridCE", "TaskID", "Stamp"]))["OK"]:
+        if not (result := getJobParameters(jobs, ["GridCE", "TaskID", "Stamp"]))["OK"]:
             self.log.error("Failed to get the list of taskIDs", result["Message"])
             return result
 

src/DIRAC/WorkloadManagementSystem/Agent/StalledJobAgent.py

@@ -20,12 +20,12 @@
 from DIRAC.Core.Utilities.ClassAd.ClassAdLight import ClassAd
 from DIRAC.Core.Utilities.TimeUtilities import fromString, second, toEpoch
 from DIRAC.WorkloadManagementSystem.Client import JobMinorStatus, JobStatus
-from DIRAC.WorkloadManagementSystem.Client.JobManagerClient import JobManagerClient
-from DIRAC.WorkloadManagementSystem.Client.JobMonitoringClient import JobMonitoringClient
 from DIRAC.WorkloadManagementSystem.Client.WMSClient import WMSClient
 from DIRAC.WorkloadManagementSystem.DB.JobDB import JobDB
 from DIRAC.WorkloadManagementSystem.DB.JobLoggingDB import JobLoggingDB
 from DIRAC.WorkloadManagementSystem.DB.PilotAgentsDB import PilotAgentsDB
+from DIRAC.WorkloadManagementSystem.Utilities.JobParameters import getJobParameters
+from DIRAC.WorkloadManagementSystem.Utilities.JobStatusUtility import rescheduleJobs
 
 
 class StalledJobAgent(AgentModule):
@@ -254,11 +254,11 @@ def _failStalledJobs(self, jobID):
 
     def _getJobPilotStatus(self, jobID):
         """Get the job pilot status."""
-        result = JobMonitoringClient().getJobParameter(jobID, "Pilot_Reference")
+        result = getJobParameters(jobID, "Pilot_Reference")
         if not result["OK"]:
             return result
-        pilotReference = result["Value"].get("Pilot_Reference", "Unknown")
-        if pilotReference == "Unknown":
+        pilotReference = result["Value"].get("Pilot_Reference")
+        if not pilotReference:
             # There is no pilot reference, hence its status is unknown
             return S_OK("NoPilot")
 
@@ -389,7 +389,7 @@ def _sendAccounting(self, jobID):
             if lastHeartBeatTime is not None and lastHeartBeatTime > endTime:
                 endTime = lastHeartBeatTime
 
-            result = JobMonitoringClient().getJobParameter(jobID, "CPUNormalizationFactor")
+            result = getJobParameters(jobID, "CPUNormalizationFactor")
             if not result["OK"] or not result["Value"]:
                 self.log.error(
                     "Error getting Job Parameter CPUNormalizationFactor, setting 0",
@@ -518,8 +518,7 @@ def _checkLoggingInfo(self, jobID, jobDict):
         return startTime, endTime
 
     def _kickStuckJobs(self):
-        """Reschedule jobs stuck in initialization status Rescheduled,
-        Matched."""
+        """Reschedule jobs stuck in initialization status Rescheduled, Matched."""
 
         message = ""
 
@@ -530,17 +529,12 @@ def _kickStuckJobs(self):
             return result
 
         jobIDs = result["Value"]
-        jobManagerClient = JobManagerClient()
         if jobIDs:
             self.log.info(f"Rescheduling {len(jobIDs)} jobs stuck in {JobStatus.MATCHED} status")
-            result = jobManagerClient.rescheduleJob(jobIDs)
+            result = rescheduleJobs(jobIDs)
             if not result["OK"]:
                 message = f"Failed to reschedule jobs stuck in {JobStatus.MATCHED} status"
                 message += "\n" + result["Message"]
-                if "InvalidJobIDs" in result:
-                    message += "\n" + "\tInvalid job IDs: " + str(result["InvalidJobIDs"])
-                if "NonauthorizedJobIDs" in result:
-                    message += "\n" + "\tNon authorized job IDs: " + str(result["NonauthorizedJobIDs"])
 
         checkTime = datetime.datetime.utcnow() - self.rescheduledTime * second
         result = self.jobDB.selectJobs({"Status": JobStatus.RESCHEDULED}, older=checkTime)
@@ -550,18 +544,14 @@ def _kickStuckJobs(self):
 
         jobIDs = result["Value"]
         if jobIDs:
-            self.log.info(f"Rescheduling {len(jobIDs)} jobs stuck in Rescheduled status")
-            result = jobManagerClient.rescheduleJob(jobIDs)
+            self.log.info(f"Rescheduling {len(jobIDs)} jobs stuck in {JobStatus.RESCHEDULED} status")
+            result = rescheduleJobs(jobIDs)
             if not result["OK"]:
                 message = f"Failed to reschedule jobs stuck in {JobStatus.RESCHEDULED} status"
                 message += "\n" + result["Message"]
-                if "InvalidJobIDs" in result:
-                    message += "\n" + "\tInvalid job IDs: " + str(result["InvalidJobIDs"])
-                if "NonauthorizedJobIDs" in result:
-                    message += "\n" + "\tNon authorized job IDs: " + str(result["NonauthorizedJobIDs"])
 
         if message:
-            return S_ERROR(message)
+            self.log.error(message)
         return S_OK()
 
     def _failSubmittingJobs(self):

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

@@ -25,9 +25,9 @@ def sja(mocker):
     )
     mocker.patch("DIRAC.WorkloadManagementSystem.Agent.StalledJobAgent.JobDB")
     mocker.patch("DIRAC.WorkloadManagementSystem.Agent.StalledJobAgent.JobLoggingDB")
-    mocker.patch("DIRAC.WorkloadManagementSystem.Agent.StalledJobAgent.JobMonitoringClient", return_value=MagicMock())
-    mocker.patch("DIRAC.WorkloadManagementSystem.Agent.StalledJobAgent.JobManagerClient", return_value=MagicMock())
+    mocker.patch("DIRAC.WorkloadManagementSystem.Agent.StalledJobAgent.rescheduleJobs", return_value=MagicMock())
     mocker.patch("DIRAC.WorkloadManagementSystem.Agent.StalledJobAgent.PilotAgentsDB", return_value=MagicMock())
+    mocker.patch("DIRAC.WorkloadManagementSystem.Agent.StalledJobAgent.getJobParameters", return_value=MagicMock())
     mocker.patch("DIRAC.WorkloadManagementSystem.Agent.StalledJobAgent.WMSClient", return_value=MagicMock())
     mocker.patch("DIRAC.WorkloadManagementSystem.Agent.StalledJobAgent.getDNForUsername", return_value=MagicMock())
 

src/DIRAC/WorkloadManagementSystem/DB/PilotAgentsDB.py

@@ -447,9 +447,10 @@ def getJobsForPilot(self, pilotID):
         """Get IDs of Jobs that were executed by a pilot"""
         cmd = "SELECT pilotID,JobID FROM JobToPilotMapping "
         if isinstance(pilotID, list):
-            cmd = cmd + " WHERE pilotID IN (%s)" % ",".join(["%s" % x for x in pilotID])
+            pilotIDs_string = ",".join(str(int(x)) for x in pilotID)
+            cmd = f"{cmd} WHERE pilotID IN ({pilotIDs_string})"
         else:
-            cmd = cmd + f" WHERE pilotID = {pilotID}"
+            cmd = f"{cmd} WHERE pilotID = {pilotID}"
 
         result = self._query(cmd)
         if not result["OK"]:

src/DIRAC/WorkloadManagementSystem/Service/JobManagerHandler.py

@@ -34,6 +34,7 @@
 )
 from DIRAC.WorkloadManagementSystem.Utilities.JobModel import JobDescriptionModel
 from DIRAC.WorkloadManagementSystem.Utilities.ParametricJob import generateParametricJobs, getParameterVectorLength
+from DIRAC.WorkloadManagementSystem.Utilities.JobStatusUtility import rescheduleJobs
 
 MAX_PARAMETRIC_JOBS = 20
 
@@ -345,8 +346,7 @@ def __getJobList(jobInput):
     types_rescheduleJob = []
 
     def export_rescheduleJob(self, jobIDs):
-        """Reschedule a single job. If the optional proxy parameter is given
-        it will be used to refresh the proxy in the Proxy Repository
+        """Reschedule a list of jobs.
 
         :param list jobIDs: list of job IDs
 
@@ -360,22 +360,12 @@ def export_rescheduleJob(self, jobIDs):
         validJobList, invalidJobList, nonauthJobList, ownerJobList = self.jobPolicy.evaluateJobRights(
             jobList, RIGHT_RESCHEDULE
         )
-        for jobID in validJobList:
-            self.taskQueueDB.deleteJob(jobID)
-            result = self.jobDB.rescheduleJob(jobID)
-            self.log.debug(str(result))
-            if not result["OK"]:
-                return result
-            self.jobLoggingDB.addLoggingRecord(
-                result["JobID"],
-                status=result["Status"],
-                minorStatus=result["MinorStatus"],
-                applicationStatus="Unknown",
-                source="JobManager",
-            )
+        res = rescheduleJobs(validJobList, source="JobManager")
+        if not res["OK"]:
+            self.log.error(res["Message"])
 
         if invalidJobList or nonauthJobList:
-            result = S_ERROR("Some jobs failed reschedule")
+            result = S_ERROR("Some jobs can not be rescheduled")
             if invalidJobList:
                 result["InvalidJobIDs"] = invalidJobList
             if nonauthJobList:

src/DIRAC/WorkloadManagementSystem/Service/JobMonitoringHandler.py

@@ -11,6 +11,7 @@
 from DIRAC.Core.Utilities.JEncode import strToIntDict
 from DIRAC.Core.Utilities.ObjectLoader import ObjectLoader
 from DIRAC.WorkloadManagementSystem.Client.PilotManagerClient import PilotManagerClient
+from DIRAC.WorkloadManagementSystem.Utilities.JobParameters import getJobParameters
 
 
 class JobMonitoringHandlerMixin:
@@ -270,27 +271,8 @@ def export_getJobParameters(self, jobIDs, parName=None):
         if not isinstance(jobIDs, list):
             jobIDs = [jobIDs]
         jobIDs = [int(jobID) for jobID in jobIDs]
-        res = self.elasticJobParametersDB.getJobParameters(jobIDs, self.vo, parName)
-        if not res["OK"]:
-            return res
-        parameters = res["Value"]
 
-        # Need anyway to get also from JobDB, for those jobs with parameters registered in MySQL or in both backends
-        res = self.jobDB.getJobParameters(jobIDs, parName)
-        if not res["OK"]:
-            return res
-        parametersM = res["Value"]
-
-        # and now combine
-        final = dict(parametersM)
-        # if job in JobDB, update with parameters from ES if any
-        for jobID in final:
-            final[jobID].update(parameters.get(jobID, {}))
-        # if job in ES and not in JobDB, take ES
-        for jobID in parameters:
-            if jobID not in final:
-                final[jobID] = parameters[jobID]
-        return S_OK(final)
+        return getJobParameters(jobIDs, parName, self.vo or "")
 
     ##############################################################################
     types_getAtticJobParameters = [int]

src/DIRAC/WorkloadManagementSystem/Utilities/JobParameters.py

@@ -2,8 +2,10 @@
 """
 import multiprocessing
 
-from DIRAC import gConfig, gLogger
+from DIRAC import gConfig, gLogger, S_OK
 from DIRAC.Core.Utilities.List import fromChar
+from DIRAC.WorkloadManagementSystem.DB.JobDB import JobDB
+from DIRAC.WorkloadManagementSystem.DB.JobParametersDB import JobParametersDB
 
 
 def getMemoryFromProc():
@@ -139,3 +141,62 @@ def getNumberOfGPUs(siteName=None, gridCE=None, queue=None):
     # 3) return 0
     gLogger.info("NumberOfGPUs could not be found in CS")
     return 0
+
+
+def getJobParameters(jobIDs: list[int], parName: str | None, vo: str = "") -> dict:
+    """Utility to get a job parameter for a list of jobIDs pertaining to a VO.
+    If the jobID is not in the JobParametersDB, it will be looked up in the JobDB.
+
+    Requires direct access to the JobParametersDB and JobDB.
+
+    :param jobIDs: list of jobIDs
+    :param parName: name of the parameter to be retrieved
+    :param vo: VO of the jobIDs
+    :return: dictionary with jobID as key and the parameter as value
+    :rtype: dict
+    """
+    elasticJobParametersDB = JobParametersDB()
+    jobDB = JobDB()
+
+    if vo:  # a user is connecting, with a proxy
+        res = elasticJobParametersDB.getJobParameters(jobIDs, vo, parName)
+        if not res["OK"]:
+            return res
+        parameters = res["Value"]
+    else:  # a service is connecting, no proxy, e.g. StalledJobAgent
+        q = f"SELECT JobID, VO FROM Jobs WHERE JobID IN ({','.join([str(jobID) for jobID in jobIDs])})"
+        res = jobDB._query(q)
+        if not res["OK"]:
+            return res
+        if not res["Value"]:
+            return S_OK({})
+        # get the VO for each jobID
+        voDict = {}
+        for jobID, vo in res["Value"]:
+            if vo not in voDict:
+                voDict[vo] = []
+            voDict[vo].append(jobID)
+        # get the parameters for each VO
+        parameters = {}
+        for vo, jobIDs in voDict.items():
+            res = elasticJobParametersDB.getJobParameters(jobIDs, vo, parName)
+            if not res["OK"]:
+                return res
+            parameters.update(res["Value"])
+
+    # Need anyway to get also from JobDB, for those jobs with parameters registered in MySQL or in both backends
+    res = jobDB.getJobParameters(jobIDs, parName)
+    if not res["OK"]:
+        return res
+    parametersM = res["Value"]
+
+    # and now combine
+    final = dict(parametersM)
+    # if job in JobDB, update with parameters from ES if any
+    for jobID in final:
+        final[jobID].update(parameters.get(jobID, {}))
+    # if job in ES and not in JobDB, take ES
+    for jobID in parameters:
+        if jobID not in final:
+            final[jobID] = parameters[jobID]
+    return S_OK(final)

src/DIRAC/WorkloadManagementSystem/Utilities/JobStatusUtility.py

@@ -9,10 +9,9 @@
 from DIRAC.Core.Utilities import TimeUtilities
 from DIRAC.Core.Utilities.ObjectLoader import ObjectLoader
 from DIRAC.WorkloadManagementSystem.Client import JobStatus
-
-if TYPE_CHECKING:
-    from DIRAC.WorkloadManagementSystem.DB.JobDB import JobDB
-    from DIRAC.WorkloadManagementSystem.DB.JobLoggingDB import JobLoggingDB
+from DIRAC.WorkloadManagementSystem.DB.JobLoggingDB import JobLoggingDB
+from DIRAC.WorkloadManagementSystem.DB.JobDB import JobDB
+from DIRAC.WorkloadManagementSystem.DB.TaskQueueDB import TaskQueueDB
 
 
 class JobStatusUtility:
@@ -243,3 +242,35 @@ def getNewStatus(
         minor = sDict.get("MinorStatus", minor)
         application = sDict.get("ApplicationStatus", application)
     return S_OK((status, minor, application))
+
+
+def rescheduleJobs(jobIDs: list[int], source: str = "") -> dict:
+    """Utility to reschedule jobs (not atomic, nor bulk)
+    Requires direct access to the JobDB and TaskQueueDB
+
+    :param jobIDs: list of jobIDs
+    :param source: source of the reschedule
+    :return: S_OK/S_ERROR
+    :rtype: dict
+
+    """
+
+    failedJobs = []
+
+    for jobID in jobIDs:
+        result = JobDB().rescheduleJob(jobID)
+        if not result["OK"]:
+            failedJobs.append(jobID)
+            continue
+        TaskQueueDB().deleteJob(jobID)
+        JobLoggingDB().addLoggingRecord(
+            result["JobID"],
+            status=result["Status"],
+            minorStatus=result["MinorStatus"],
+            applicationStatus="Unknown",
+            source=source,
+        )
+
+    if failedJobs:
+        return S_ERROR(f"Failed to reschedule jobs {failedJobs}")
+    return S_OK()