Merge pull request #8278 from fstagni/90_refactoring_WMS

[9.0] more WMS refactoring

Author: chrisburr

src/DIRAC/Interfaces/API/Dirac.py

@@ -46,6 +46,7 @@
 from DIRAC.WorkloadManagementSystem.Client.JobMonitoringClient import JobMonitoringClient
 from DIRAC.WorkloadManagementSystem.Client.SandboxStoreClient import SandboxStoreClient
 from DIRAC.WorkloadManagementSystem.Client.WMSClient import WMSClient
+from DIRAC.WorkloadManagementSystem.Utilities.jobAdministration import _filterJobStateTransition
 
 
 def parseArguments(args):
@@ -1450,10 +1451,13 @@ def deleteJob(self, jobID):
         # Remove any job IDs that can't change to the Killed or Deleted states
         filteredJobs = set()
         for filterState in (JobStatus.KILLED, JobStatus.DELETED):
-            filterRes = JobStatus.filterJobStateTransition(jobIDs, filterState)
-            if not filterRes["OK"]:
-                return filterRes
-            filteredJobs.update(filterRes["Value"])
+            # get a dictionary of jobID:status
+            res = JobMonitoringClient().getJobsStatus(jobIDs)
+            if not res["OK"]:
+                return res
+            js = {k: v["Status"] for k, v in res["Value"].items()}
+            # then filter
+            filteredJobs.update(_filterJobStateTransition(js, filterState))
 
         return WMSClient(useCertificates=self.useCertificates).deleteJob(list(filteredJobs))
 
@@ -1480,11 +1484,13 @@ def rescheduleJob(self, jobID):
             return ret
         jobIDs = ret["Value"]
 
-        # Remove any job IDs that can't change to the rescheduled state
-        filterRes = JobStatus.filterJobStateTransition(jobIDs, JobStatus.RESCHEDULED)
-        if not filterRes["OK"]:
-            return filterRes
-        jobIDsToReschedule = filterRes["Value"]
+        # get a dictionary of jobID:status
+        res = JobMonitoringClient().getJobsStatus(jobIDs)
+        if not res["OK"]:
+            return res
+        js = {k: v["Status"] for k, v in res["Value"].items()}
+        # then filter
+        jobIDsToReschedule = _filterJobStateTransition(js, JobStatus.RESCHEDULED)
 
         return WMSClient(useCertificates=self.useCertificates).rescheduleJob(jobIDsToReschedule)
 
@@ -1510,10 +1516,13 @@ def killJob(self, jobID):
         # Remove any job IDs that can't change to the Killed or Deleted states
         filteredJobs = set()
         for filterState in (JobStatus.KILLED, JobStatus.DELETED):
-            filterRes = JobStatus.filterJobStateTransition(jobIDs, filterState)
-            if not filterRes["OK"]:
-                return filterRes
-            filteredJobs.update(filterRes["Value"])
+            # get a dictionary of jobID:status
+            res = JobMonitoringClient().getJobsStatus(jobIDs)
+            if not res["OK"]:
+                return res
+            js = {k: v["Status"] for k, v in res["Value"].items()}
+            # then filter
+            filteredJobs.update(_filterJobStateTransition(js, filterState))
 
         return WMSClient(useCertificates=self.useCertificates).killJob(list(filteredJobs))
 

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

@@ -28,6 +28,7 @@ def sja(mocker):
     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.kill_delete_jobs", return_value=MagicMock())
 
     stalledJobAgent = StalledJobAgent()
     stalledJobAgent._AgentModule__configDefaults = mockAM

src/DIRAC/WorkloadManagementSystem/Client/JobStatus.py

@@ -2,12 +2,7 @@
 This module contains constants and lists for the possible job states.
 """
 
-from DIRAC import gLogger, S_OK, S_ERROR
 from DIRAC.Core.Utilities.StateMachine import State, StateMachine
-from DIRAC.Core.Utilities.Decorators import deprecated
-
-from DIRAC.WorkloadManagementSystem.Client.JobMonitoringClient import JobMonitoringClient
-
 
 #:
 SUBMITTING = "Submitting"
@@ -98,57 +93,3 @@ def __init__(self, state):
             RECEIVED: State(1, [SCOUTING, CHECKING, STAGING, WAITING, FAILED, DELETED, KILLED], defState=RECEIVED),
             SUBMITTING: State(0, [RECEIVED, CHECKING, DELETED, KILLED], defState=SUBMITTING),  # initial state
         }
-
-
-@deprecated("Use filterJobStateTransition instead")
-def checkJobStateTransition(jobID, candidateState, currentStatus=None, jobMonitoringClient=None):
-    """Utility to check if a job state transition is allowed"""
-    if not currentStatus:
-        if not jobMonitoringClient:
-            from DIRAC.WorkloadManagementSystem.Client.JobMonitoringClient import JobMonitoringClient
-
-            jobMonitoringClient = JobMonitoringClient()
-
-        res = jobMonitoringClient.getJobsStatus(jobID)
-        if not res["OK"]:
-            return res
-        try:
-            currentStatus = res["Value"][jobID]["Status"]
-        except KeyError:
-            return S_ERROR("Job does not exist")
-
-    res = JobsStateMachine(currentStatus).getNextState(candidateState)
-    if not res["OK"]:
-        return res
-
-    # If the JobsStateMachine does not accept the candidate, return an ERROR
-    if candidateState != res["Value"]:
-        gLogger.error(
-            "Job Status Error",
-            f"{jobID} can't move from {currentStatus} to {candidateState}",
-        )
-        return S_ERROR("Job state transition not allowed")
-    return S_OK()
-
-
-def filterJobStateTransition(jobIDs, candidateState):
-    """Given a list of jobIDs, return a list that are allowed to transition
-    to the given candidate state.
-    """
-    allowedJobs = []
-
-    if not isinstance(jobIDs, list):
-        jobIDs = [jobIDs]
-
-    res = JobMonitoringClient().getJobsStatus(jobIDs)
-    if not res["OK"]:
-        return res
-
-    for jobID in jobIDs:
-        if jobID in res["Value"]:
-            curState = res["Value"][jobID]["Status"]
-            stateRes = JobsStateMachine(curState).getNextState(candidateState)
-            if stateRes["OK"]:
-                if stateRes["Value"] == candidateState:
-                    allowedJobs.append(jobID)
-    return S_OK(allowedJobs)

src/DIRAC/WorkloadManagementSystem/FutureClient/JobStateUpdateClient.py

@@ -83,12 +83,6 @@ def setJobAttribute(self, jobID: str | int, attribute: str, value: str):
             else:
                 return api.jobs.patch_metadata({jobID: {attribute: value}})
 
-    @stripValueIfOK
-    @convertToReturnValue
-    def setJobFlag(self, jobID: str | int, flag: str):
-        with DiracXClient() as api:
-            api.jobs.patch_metadata({jobID: {flag: True}})
-
     @stripValueIfOK
     @convertToReturnValue
     def setJobParameter(self, jobID: str | int, name: str, value: str):
@@ -151,12 +145,6 @@ def setJobsParameter(self, jobsParameterDict: dict):
             updates = {job_id: {k: v} for job_id, (k, v) in jobsParameterDict.items()}
             api.jobs.patch_metadata(updates)
 
-    @stripValueIfOK
-    @convertToReturnValue
-    def unsetJobFlag(self, jobID: str | int, flag: str):
-        with DiracXClient() as api:
-            api.jobs.patch_metadata({jobID: {flag: False}})
-
     @stripValueIfOK
     @convertToReturnValue
     def updateJobFromStager(self, jobID: str | int, status: str):

src/DIRAC/WorkloadManagementSystem/Service/JobStateUpdateHandler.py

@@ -128,22 +128,6 @@ def export_setJobSite(cls, jobID, site):
         """Allows the site attribute to be set for a job specified by its jobID."""
         return cls.jobDB.setJobAttribute(int(jobID), "Site", site)
 
-    ###########################################################################
-    types_setJobFlag = [[str, int], str]
-
-    @classmethod
-    def export_setJobFlag(cls, jobID, flag):
-        """Set job flag for job with jobID"""
-        return cls.jobDB.setJobAttribute(int(jobID), flag, "True")
-
-    ###########################################################################
-    types_unsetJobFlag = [[str, int], str]
-
-    @classmethod
-    def export_unsetJobFlag(cls, jobID, flag):
-        """Unset job flag for job with jobID"""
-        return cls.jobDB.setJobAttribute(int(jobID), flag, "False")
-
     ###########################################################################
     types_setJobApplicationStatus = [[str, int], str, str]
 

src/DIRAC/WorkloadManagementSystem/Utilities/jobAdministration.py

@@ -1,11 +1,10 @@
 from DIRAC import S_ERROR, S_OK, gLogger
 from DIRAC.StorageManagementSystem.DB.StorageManagementDB import StorageManagementDB
 from DIRAC.WorkloadManagementSystem.Client import JobStatus
-from DIRAC.WorkloadManagementSystem.Client.JobStatus import filterJobStateTransition
 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_KILL, RIGHT_DELETE
+from DIRAC.WorkloadManagementSystem.Service.JobPolicy import RIGHT_DELETE, RIGHT_KILL
 
 
 def _deleteJob(jobID, force=False):
@@ -79,18 +78,17 @@ def kill_delete_jobs(right, validJobList, nonauthJobList=[], force=False):
     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
-        filterRes = filterJobStateTransition(validJobList, JobStatus.KILLED)
-        if not filterRes["OK"]:
-            return filterRes
-        killJobList.extend(filterRes["Value"])
+        killJobList.extend(_filterJobStateTransition(jobStates, JobStatus.KILLED))
 
         if right == RIGHT_DELETE:
             # Get the jobs allowed to transition to the Deleted state
-            filterRes = filterJobStateTransition(validJobList, JobStatus.DELETED)
-            if not filterRes["OK"]:
-                return filterRes
-            deleteJobList.extend(filterRes["Value"])
+            deleteJobList.extend(_filterJobStateTransition(jobStates, JobStatus.DELETED))
 
         for jobID in killJobList:
             result = _killJob(jobID, force=force)
@@ -103,10 +101,7 @@ def kill_delete_jobs(right, validJobList, nonauthJobList=[], force=False):
                 badIDs.append(jobID)
 
         # Look for jobs that are in the Staging state to send kill signal to the stager
-        result = JobDB().getJobsAttributes(killJobList, ["Status"])
-        if not result["OK"]:
-            return result
-        stagingJobList = [jobID for jobID, sDict in result["Value"].items() if sDict["Status"] == JobStatus.STAGING]
+        stagingJobList = [jobID for jobID, sDict in jobStates.items() if sDict["Status"] == JobStatus.STAGING]
 
         if stagingJobList:
             stagerDB = StorageManagementDB()
@@ -127,3 +122,17 @@ def kill_delete_jobs(right, validJobList, nonauthJobList=[], force=False):
 
     jobsList = killJobList if right == RIGHT_KILL else deleteJobList
     return S_OK(jobsList)
+
+
+def _filterJobStateTransition(jobStates, candidateState):
+    """Given a dictionary of jobs states,
+    return a list of jobs that are allowed to transition to the given candidate state.
+    """
+    allowedJobs = []
+
+    for js in jobStates.items():
+        stateRes = JobStatus.JobsStateMachine(js[1]["Status"]).getNextState(candidateState)
+        if stateRes["OK"]:
+            if stateRes["Value"] == candidateState:
+                allowedJobs.append(js[0])
+    return allowedJobs

src/DIRAC/WorkloadManagementSystem/Utilities/test/Test_JobAdministration.py

@@ -10,31 +10,19 @@
 
 
 @pytest.mark.parametrize(
-    "jobIDs_list, right, filtered_jobs, expected_res, expected_value",
+    "jobIDs_list, right",
     [
-        ([], "Kill", [], True, []),
-        ([], "Delete", [], True, []),
-        (1, "Kill", [], True, []),
-        (1, "Kill", [1], True, [1]),
-        ([1, 2], "Kill", [], True, []),
-        ([1, 2], "Kill", [1], True, [1]),
-        (1, "Kill", [1], True, [1]),
-        ([1, 2], "Kill", [1], True, [1]),
-        ([1, 2], "Kill", [2], True, [2]),
-        ([1, 2], "Kill", [], True, []),
-        ([1, 2], "Kill", [1, 2], True, [1, 2]),
+        ([], "Kill"),
+        ([], "Delete"),
+        (1, "Kill"),
+        ([1, 2], "Kill"),
     ],
 )
-def test___kill_delete_jobs(mocker, jobIDs_list, right, filtered_jobs, expected_res, expected_value):
+def test___kill_delete_jobs(mocker, jobIDs_list, right):
     mocker.patch("DIRAC.WorkloadManagementSystem.Utilities.jobAdministration.JobDB", MagicMock())
     mocker.patch("DIRAC.WorkloadManagementSystem.Utilities.jobAdministration.TaskQueueDB", MagicMock())
     mocker.patch("DIRAC.WorkloadManagementSystem.Utilities.jobAdministration.PilotAgentsDB", MagicMock())
     mocker.patch("DIRAC.WorkloadManagementSystem.Utilities.jobAdministration.StorageManagementDB", MagicMock())
-    mocker.patch(
-        "DIRAC.WorkloadManagementSystem.Utilities.jobAdministration.filterJobStateTransition",
-        return_value={"OK": True, "Value": filtered_jobs},
-    )
 
     res = kill_delete_jobs(right, jobIDs_list)
-    assert res["OK"] == expected_res
-    assert res["Value"] == expected_value
+    assert res["OK"]

tests/Integration/FutureClient/WorkloadManagement/Test_JobStateUpdate.py

@@ -142,42 +142,6 @@ def test_setJobAttribute(monkeypatch, example_jobids, args):
     assert result[example_jobids[0]] == result[example_jobids[1]]
 
 
-def test_setJobFlag(monkeypatch, example_jobids):
-    # JobStateUpdateClient().setJobFlag(jobID: str | int, flag: str)
-    assert returnValueOrRaise(JobMonitoringClient().getJobAttribute(example_jobids[0], "AccountedFlag")) == "False"
-    assert returnValueOrRaise(JobMonitoringClient().getJobAttribute(example_jobids[1], "AccountedFlag")) == "False"
-
-    method = JobStateUpdateClient().setJobFlag
-    test_func1 = partial(method, example_jobids[0], "AccountedFlag")
-    test_func2 = partial(method, example_jobids[1], "AccountedFlag")
-    compare_results2(monkeypatch, test_func1, test_func2)
-
-    assert returnValueOrRaise(JobMonitoringClient().getJobAttribute(example_jobids[0], "AccountedFlag")) == "True"
-    assert returnValueOrRaise(JobMonitoringClient().getJobAttribute(example_jobids[1], "AccountedFlag")) == "True"
-
-
-def test_unsetJobFlag(monkeypatch, example_jobids):
-    # JobStateUpdateClient().unsetJobFlag(jobID: str | int, flag: str)
-    assert returnValueOrRaise(JobMonitoringClient().getJobAttribute(example_jobids[0], "AccountedFlag")) == "False"
-    assert returnValueOrRaise(JobMonitoringClient().getJobAttribute(example_jobids[1], "AccountedFlag")) == "False"
-
-    method = JobStateUpdateClient().setJobFlag
-    test_func1 = partial(method, example_jobids[0], "AccountedFlag")
-    test_func2 = partial(method, example_jobids[1], "AccountedFlag")
-    compare_results2(monkeypatch, test_func1, test_func2)
-
-    assert returnValueOrRaise(JobMonitoringClient().getJobAttribute(example_jobids[0], "AccountedFlag")) == "True"
-    assert returnValueOrRaise(JobMonitoringClient().getJobAttribute(example_jobids[1], "AccountedFlag")) == "True"
-
-    method = JobStateUpdateClient().unsetJobFlag
-    test_func1 = partial(method, example_jobids[0], "AccountedFlag")
-    test_func2 = partial(method, example_jobids[1], "AccountedFlag")
-    compare_results2(monkeypatch, test_func1, test_func2)
-
-    assert returnValueOrRaise(JobMonitoringClient().getJobAttribute(example_jobids[0], "AccountedFlag")) == "False"
-    assert returnValueOrRaise(JobMonitoringClient().getJobAttribute(example_jobids[1], "AccountedFlag")) == "False"
-
-
 @pytest.mark.parametrize(
     "args",
     [