fix: Move rescheduleJobs into Utils

Author: chrisburr

src/DIRAC/WorkloadManagementSystem/Agent/StalledJobAgent.py

@@ -25,7 +25,7 @@
 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
+from DIRAC.WorkloadManagementSystem.Utilities.Utils import rescheduleJobs
 
 
 class StalledJobAgent(AgentModule):

src/DIRAC/WorkloadManagementSystem/Service/JobManagerHandler.py

@@ -34,7 +34,7 @@
 )
 from DIRAC.WorkloadManagementSystem.Utilities.JobModel import JobDescriptionModel
 from DIRAC.WorkloadManagementSystem.Utilities.ParametricJob import generateParametricJobs, getParameterVectorLength
-from DIRAC.WorkloadManagementSystem.Utilities.JobStatusUtility import rescheduleJobs
+from DIRAC.WorkloadManagementSystem.Utilities.Utils import rescheduleJobs
 
 MAX_PARAMETRIC_JOBS = 20
 

src/DIRAC/WorkloadManagementSystem/Utilities/JobStatusUtility.py

@@ -9,9 +9,10 @@
 from DIRAC.Core.Utilities import TimeUtilities
 from DIRAC.Core.Utilities.ObjectLoader import ObjectLoader
 from DIRAC.WorkloadManagementSystem.Client import JobStatus
-from DIRAC.WorkloadManagementSystem.DB.JobLoggingDB import JobLoggingDB
-from DIRAC.WorkloadManagementSystem.DB.JobDB import JobDB
-from DIRAC.WorkloadManagementSystem.DB.TaskQueueDB import TaskQueueDB
+
+if TYPE_CHECKING:
+    from DIRAC.WorkloadManagementSystem.DB.JobLoggingDB import JobLoggingDB
+    from DIRAC.WorkloadManagementSystem.DB.JobDB import JobDB
 
 
 class JobStatusUtility:
@@ -242,35 +243,3 @@ 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()

src/DIRAC/WorkloadManagementSystem/Utilities/Utils.py

@@ -4,9 +4,12 @@
 import sys
 import json
 
-from DIRAC import gLogger, S_OK
+from DIRAC import gLogger, S_OK, S_ERROR
 from DIRAC.Core.Utilities.File import mkDir
 from DIRAC.FrameworkSystem.private.standardLogging.Logging import Logging
+from DIRAC.WorkloadManagementSystem.DB.JobLoggingDB import JobLoggingDB
+from DIRAC.WorkloadManagementSystem.DB.JobDB import JobDB
+from DIRAC.WorkloadManagementSystem.DB.TaskQueueDB import TaskQueueDB
 
 
 def createJobWrapper(
@@ -113,3 +116,35 @@ def createJobWrapper(
     if rootLocation != wrapperPath:
         generatedFiles["JobExecutableRelocatedPath"] = os.path.join(rootLocation, os.path.basename(jobExeFile))
     return S_OK(generatedFiles)
+
+
+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()