refactor: extracted rescheduleJobs in utility, use it in agents

Author: fstagni

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,9 @@ 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])
+            cmd = f"{cmd} WHERE pilotID IN ({",".join(["%s" % x for x in pilotID])})"
         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/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,27 @@ 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)"""
+
+    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()