Merge pull request #8283 from chaen/v9.0_fix_utilityAgain

fix (WMS): do not import DB in utilities

Author: fstagni

.github/workflows/basic.yml

@@ -126,7 +126,7 @@ jobs:
           fetch-depth: 0
           repository: DIRACGrid/diracx
           path: diracx
-      - uses: prefix-dev/setup-pixi@v0.8.11
+      - uses: prefix-dev/setup-pixi@v0.9.0
         with:
           run-install: false
           post-cleanup: false
@@ -143,7 +143,7 @@ jobs:
           pixi add --pypi --feature diracx-core 'DIRAC @ file://'$PWD'/../DIRAC'
           # Show any changes
           git diff
-      - uses: prefix-dev/setup-pixi@v0.8.11
+      - uses: prefix-dev/setup-pixi@v0.9.0
         with:
           cache: false
           manifest-path: diracx/pixi.toml

src/DIRAC/TransformationSystem/Agent/TransformationCleaningAgent.py

@@ -37,7 +37,7 @@
     RIGHT_DELETE,
     RIGHT_KILL,
 )
-from DIRAC.WorkloadManagementSystem.Utilities.jobAdministration import kill_delete_jobs
+from DIRAC.WorkloadManagementSystem.DB.StatusUtils import kill_delete_jobs
 
 # # agent's name
 AGENT_NAME = "Transformation/TransformationCleaningAgent"

src/DIRAC/WorkloadManagementSystem/Agent/JobCleaningAgent.py

@@ -39,7 +39,7 @@
 from DIRAC.WorkloadManagementSystem.DB.JobDB import JobDB
 from DIRAC.WorkloadManagementSystem.DB.SandboxMetadataDB import SandboxMetadataDB
 from DIRAC.WorkloadManagementSystem.Service.JobPolicy import RIGHT_DELETE
-from DIRAC.WorkloadManagementSystem.Utilities.jobAdministration import kill_delete_jobs
+from DIRAC.WorkloadManagementSystem.DB.StatusUtils import kill_delete_jobs
 from DIRAC.WorkloadManagementSystem.Utilities.JobParameters import getJobParameters
 
 

src/DIRAC/WorkloadManagementSystem/Agent/StalledJobAgent.py

@@ -23,7 +23,7 @@
 from DIRAC.WorkloadManagementSystem.DB.JobLoggingDB import JobLoggingDB
 from DIRAC.WorkloadManagementSystem.DB.PilotAgentsDB import PilotAgentsDB
 from DIRAC.WorkloadManagementSystem.Service.JobPolicy import RIGHT_KILL
-from DIRAC.WorkloadManagementSystem.Utilities.jobAdministration import kill_delete_jobs
+from DIRAC.WorkloadManagementSystem.DB.StatusUtils import kill_delete_jobs
 from DIRAC.WorkloadManagementSystem.Utilities.JobParameters import getJobParameters
 from DIRAC.WorkloadManagementSystem.Utilities.Utils import rescheduleJobs
 

src/DIRAC/WorkloadManagementSystem/DB/StatusUtils.py

@@ -0,0 +1,125 @@
+from DIRAC import S_ERROR, S_OK, gLogger
+from DIRAC.StorageManagementSystem.DB.StorageManagementDB import StorageManagementDB
+from DIRAC.WorkloadManagementSystem.Client import JobStatus
+from DIRAC.WorkloadManagementSystem.DB.JobDB import JobDB
+from DIRAC.WorkloadManagementSystem.DB.PilotAgentsDB import PilotAgentsDB
+from DIRAC.WorkloadManagementSystem.DB.TaskQueueDB import TaskQueueDB
+from DIRAC.WorkloadManagementSystem.Service.JobPolicy import RIGHT_DELETE, RIGHT_KILL
+from DIRAC.WorkloadManagementSystem.Utilities.jobAdministration import _filterJobStateTransition
+
+
+def _deleteJob(jobID, force=False):
+    """Set the job status to "Deleted"
+    and remove the pilot that ran and its logging info if the pilot is finished.
+
+    :param int jobID: job ID
+    :return: S_OK()/S_ERROR()
+    """
+    if not (result := JobDB().setJobStatus(jobID, JobStatus.DELETED, "Checking accounting", force=force))["OK"]:
+        gLogger.warn("Failed to set job Deleted status", result["Message"])
+        return result
+
+    if not (result := TaskQueueDB().deleteJob(jobID))["OK"]:
+        gLogger.warn("Failed to delete job from the TaskQueue")
+
+    # if it was the last job for the pilot
+    result = PilotAgentsDB().getPilotsForJobID(jobID)
+    if not result["OK"]:
+        gLogger.error("Failed to get Pilots for JobID", result["Message"])
+        return result
+    for pilot in result["Value"]:
+        res = PilotAgentsDB().getJobsForPilot(pilot)
+        if not res["OK"]:
+            gLogger.error("Failed to get jobs for pilot", res["Message"])
+            return res
+        if not res["Value"]:  # if list of jobs for pilot is empty, delete pilot
+            result = PilotAgentsDB().getPilotInfo(pilotID=pilot)
+            if not result["OK"]:
+                gLogger.error("Failed to get pilot info", result["Message"])
+                return result
+            ret = PilotAgentsDB().deletePilot(result["Value"]["PilotJobReference"])
+            if not ret["OK"]:
+                gLogger.error("Failed to delete pilot from PilotAgentsDB", ret["Message"])
+                return ret
+
+    return S_OK()
+
+
+def _killJob(jobID, sendKillCommand=True, force=False):
+    """Kill one job
+
+    :param int jobID: job ID
+    :param bool sendKillCommand: send kill command
+
+    :return: S_OK()/S_ERROR()
+    """
+    if sendKillCommand:
+        if not (result := JobDB().setJobCommand(jobID, "Kill"))["OK"]:
+            gLogger.warn("Failed to set job Kill command", result["Message"])
+            return result
+
+    gLogger.info("Job marked for termination", jobID)
+    if not (result := JobDB().setJobStatus(jobID, JobStatus.KILLED, "Marked for termination", force=force))["OK"]:
+        gLogger.warn("Failed to set job Killed status", result["Message"])
+    if not (result := TaskQueueDB().deleteJob(jobID))["OK"]:
+        gLogger.warn("Failed to delete job from the TaskQueue", result["Message"])
+
+    return S_OK()
+
+
+def kill_delete_jobs(right, validJobList, nonauthJobList=[], force=False):
+    """Kill (== set the status to "KILLED") or delete (== set the status to "DELETED") jobs as necessary
+
+    :param str right: RIGHT_KILL or RIGHT_DELETE
+
+    :return: S_OK()/S_ERROR()
+    """
+    badIDs = []
+
+    killJobList = []
+    deleteJobList = []
+    if validJobList:
+        result = JobDB().getJobsAttributes(killJobList, ["Status"])
+        if not result["OK"]:
+            return result
+        jobStates = result["Value"]
+
+        # Get the jobs allowed to transition to the Killed state
+        killJobList.extend(_filterJobStateTransition(jobStates, JobStatus.KILLED))
+
+        if right == RIGHT_DELETE:
+            # Get the jobs allowed to transition to the Deleted state
+            deleteJobList.extend(_filterJobStateTransition(jobStates, JobStatus.DELETED))
+
+        for jobID in killJobList:
+            result = _killJob(jobID, force=force)
+            if not result["OK"]:
+                badIDs.append(jobID)
+
+        for jobID in deleteJobList:
+            result = _deleteJob(jobID, force=force)
+            if not result["OK"]:
+                badIDs.append(jobID)
+
+        # Look for jobs that are in the Staging state to send kill signal to the stager
+        stagingJobList = [jobID for jobID, sDict in jobStates.items() if sDict["Status"] == JobStatus.STAGING]
+
+        if stagingJobList:
+            stagerDB = StorageManagementDB()
+            gLogger.info("Going to send killing signal to stager as well!")
+            result = stagerDB.killTasksBySourceTaskID(stagingJobList)
+            if not result["OK"]:
+                gLogger.warn("Failed to kill some Stager tasks", result["Message"])
+
+    if nonauthJobList or badIDs:
+        result = S_ERROR("Some jobs failed deletion")
+        if nonauthJobList:
+            gLogger.warn("Non-authorized JobIDs won't be deleted", str(nonauthJobList))
+            result["NonauthorizedJobIDs"] = nonauthJobList
+        if badIDs:
+            gLogger.warn("JobIDs failed to be deleted", str(badIDs))
+            result["FailedJobIDs"] = badIDs
+        return result
+
+    jobsList = killJobList if right == RIGHT_KILL else deleteJobList
+    return S_OK(jobsList)

src/DIRAC/WorkloadManagementSystem/DB/tests/Test_StatusUtils.py

@@ -0,0 +1,28 @@
+""" unit test (pytest) of JobAdministration module
+"""
+
+from unittest.mock import MagicMock
+
+import pytest
+
+# sut
+from DIRAC.WorkloadManagementSystem.DB.StatusUtils import kill_delete_jobs
+
+
+@pytest.mark.parametrize(
+    "jobIDs_list, right",
+    [
+        ([], "Kill"),
+        ([], "Delete"),
+        (1, "Kill"),
+        ([1, 2], "Kill"),
+    ],
+)
+def test___kill_delete_jobs(mocker, jobIDs_list, right):
+    mocker.patch("DIRAC.WorkloadManagementSystem.DB.StatusUtils.JobDB", MagicMock())
+    mocker.patch("DIRAC.WorkloadManagementSystem.DB.StatusUtils.TaskQueueDB", MagicMock())
+    mocker.patch("DIRAC.WorkloadManagementSystem.DB.StatusUtils.PilotAgentsDB", MagicMock())
+    mocker.patch("DIRAC.WorkloadManagementSystem.DB.StatusUtils.StorageManagementDB", MagicMock())
+
+    res = kill_delete_jobs(right, jobIDs_list)
+    assert res["OK"]

src/DIRAC/WorkloadManagementSystem/Service/JobManagerHandler.py

@@ -30,7 +30,7 @@
     RIGHT_SUBMIT,
     JobPolicy,
 )
-from DIRAC.WorkloadManagementSystem.Utilities.jobAdministration import kill_delete_jobs
+from DIRAC.WorkloadManagementSystem.DB.StatusUtils import kill_delete_jobs
 from DIRAC.WorkloadManagementSystem.Utilities.JobModel import JobDescriptionModel
 from DIRAC.WorkloadManagementSystem.Utilities.ParametricJob import generateParametricJobs, getParameterVectorLength
 from DIRAC.WorkloadManagementSystem.Utilities.Utils import rescheduleJobs

src/DIRAC/WorkloadManagementSystem/Utilities/jobAdministration.py

@@ -1,127 +1,4 @@
-from DIRAC import S_ERROR, S_OK, gLogger
-from DIRAC.StorageManagementSystem.DB.StorageManagementDB import StorageManagementDB
 from DIRAC.WorkloadManagementSystem.Client import JobStatus
-from DIRAC.WorkloadManagementSystem.DB.JobDB import JobDB
-from DIRAC.WorkloadManagementSystem.DB.PilotAgentsDB import PilotAgentsDB
-from DIRAC.WorkloadManagementSystem.DB.TaskQueueDB import TaskQueueDB
-from DIRAC.WorkloadManagementSystem.Service.JobPolicy import RIGHT_DELETE, RIGHT_KILL
-
-
-def _deleteJob(jobID, force=False):
-    """Set the job status to "Deleted"
-    and remove the pilot that ran and its logging info if the pilot is finished.
-
-    :param int jobID: job ID
-    :return: S_OK()/S_ERROR()
-    """
-    if not (result := JobDB().setJobStatus(jobID, JobStatus.DELETED, "Checking accounting", force=force))["OK"]:
-        gLogger.warn("Failed to set job Deleted status", result["Message"])
-        return result
-
-    if not (result := TaskQueueDB().deleteJob(jobID))["OK"]:
-        gLogger.warn("Failed to delete job from the TaskQueue")
-
-    # if it was the last job for the pilot
-    result = PilotAgentsDB().getPilotsForJobID(jobID)
-    if not result["OK"]:
-        gLogger.error("Failed to get Pilots for JobID", result["Message"])
-        return result
-    for pilot in result["Value"]:
-        res = PilotAgentsDB().getJobsForPilot(pilot)
-        if not res["OK"]:
-            gLogger.error("Failed to get jobs for pilot", res["Message"])
-            return res
-        if not res["Value"]:  # if list of jobs for pilot is empty, delete pilot
-            result = PilotAgentsDB().getPilotInfo(pilotID=pilot)
-            if not result["OK"]:
-                gLogger.error("Failed to get pilot info", result["Message"])
-                return result
-            ret = PilotAgentsDB().deletePilot(result["Value"]["PilotJobReference"])
-            if not ret["OK"]:
-                gLogger.error("Failed to delete pilot from PilotAgentsDB", ret["Message"])
-                return ret
-
-    return S_OK()
-
-
-def _killJob(jobID, sendKillCommand=True, force=False):
-    """Kill one job
-
-    :param int jobID: job ID
-    :param bool sendKillCommand: send kill command
-
-    :return: S_OK()/S_ERROR()
-    """
-    if sendKillCommand:
-        if not (result := JobDB().setJobCommand(jobID, "Kill"))["OK"]:
-            gLogger.warn("Failed to set job Kill command", result["Message"])
-            return result
-
-    gLogger.info("Job marked for termination", jobID)
-    if not (result := JobDB().setJobStatus(jobID, JobStatus.KILLED, "Marked for termination", force=force))["OK"]:
-        gLogger.warn("Failed to set job Killed status", result["Message"])
-    if not (result := TaskQueueDB().deleteJob(jobID))["OK"]:
-        gLogger.warn("Failed to delete job from the TaskQueue", result["Message"])
-
-    return S_OK()
-
-
-def kill_delete_jobs(right, validJobList, nonauthJobList=[], force=False):
-    """Kill (== set the status to "KILLED") or delete (== set the status to "DELETED") jobs as necessary
-
-    :param str right: RIGHT_KILL or RIGHT_DELETE
-
-    :return: S_OK()/S_ERROR()
-    """
-    badIDs = []
-
-    killJobList = []
-    deleteJobList = []
-    if validJobList:
-        result = JobDB().getJobsAttributes(killJobList, ["Status"])
-        if not result["OK"]:
-            return result
-        jobStates = result["Value"]
-
-        # Get the jobs allowed to transition to the Killed state
-        killJobList.extend(_filterJobStateTransition(jobStates, JobStatus.KILLED))
-
-        if right == RIGHT_DELETE:
-            # Get the jobs allowed to transition to the Deleted state
-            deleteJobList.extend(_filterJobStateTransition(jobStates, JobStatus.DELETED))
-
-        for jobID in killJobList:
-            result = _killJob(jobID, force=force)
-            if not result["OK"]:
-                badIDs.append(jobID)
-
-        for jobID in deleteJobList:
-            result = _deleteJob(jobID, force=force)
-            if not result["OK"]:
-                badIDs.append(jobID)
-
-        # Look for jobs that are in the Staging state to send kill signal to the stager
-        stagingJobList = [jobID for jobID, sDict in jobStates.items() if sDict["Status"] == JobStatus.STAGING]
-
-        if stagingJobList:
-            stagerDB = StorageManagementDB()
-            gLogger.info("Going to send killing signal to stager as well!")
-            result = stagerDB.killTasksBySourceTaskID(stagingJobList)
-            if not result["OK"]:
-                gLogger.warn("Failed to kill some Stager tasks", result["Message"])
-
-    if nonauthJobList or badIDs:
-        result = S_ERROR("Some jobs failed deletion")
-        if nonauthJobList:
-            gLogger.warn("Non-authorized JobIDs won't be deleted", str(nonauthJobList))
-            result["NonauthorizedJobIDs"] = nonauthJobList
-        if badIDs:
-            gLogger.warn("JobIDs failed to be deleted", str(badIDs))
-            result["FailedJobIDs"] = badIDs
-        return result
-
-    jobsList = killJobList if right == RIGHT_KILL else deleteJobList
-    return S_OK(jobsList)
 
 
 def _filterJobStateTransition(jobStates, candidateState):