refactor: extracted rescheduleJobs in utility, use it in agents

Author: fstagni

src/DIRAC/WorkloadManagementSystem/Agent/PushJobAgent.py

@@ -41,10 +41,9 @@
     transferInputSandbox,
 )
 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
-from DIRAC.WorkloadManagementSystem.Utilities.JobParameters import getJobParameters
-
 
 MAX_JOBS_MANAGED = 100
 

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.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):
@@ -529,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,17 +545,13 @@ def _kickStuckJobs(self):
         jobIDs = result["Value"]
         if jobIDs:
             self.log.info(f"Rescheduling {len(jobIDs)} jobs stuck in {JobStatus.RESCHEDULED} status")
-            result = jobManagerClient.rescheduleJob(jobIDs)
+            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,7 +25,7 @@ def sja(mocker):
     )
     mocker.patch("DIRAC.WorkloadManagementSystem.Agent.StalledJobAgent.JobDB")
     mocker.patch("DIRAC.WorkloadManagementSystem.Agent.StalledJobAgent.JobLoggingDB")
-    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())

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/Utilities/JobParameters.py

@@ -146,8 +146,9 @@ def getNumberOfGPUs(siteName=None, gridCE=None, queue=None):
 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

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()