feat: use directly the utilities instead of going through the service

Author: fstagni

src/DIRAC/TransformationSystem/Agent/TransformationCleaningAgent.py

@@ -33,8 +33,12 @@
 from DIRAC.TransformationSystem.Client import TransformationStatus
 from DIRAC.TransformationSystem.Client.TransformationClient import TransformationClient
 from DIRAC.WorkloadManagementSystem.Client.JobMonitoringClient import JobMonitoringClient
-from DIRAC.WorkloadManagementSystem.Client.WMSClient import WMSClient
 from DIRAC.WorkloadManagementSystem.DB.JobDB import JobDB
+from DIRAC.WorkloadManagementSystem.Service.JobPolicy import (
+    RIGHT_DELETE,
+    RIGHT_KILL,
+)
+from DIRAC.WorkloadManagementSystem.Utilities.jobAdministration import kill_delete_jobs
 
 # # agent's name
 AGENT_NAME = "Transformation/TransformationCleaningAgent"
@@ -58,8 +62,6 @@ def __init__(self, *args, **kwargs):
 
         # # transformation client
         self.transClient = None
-        # # wms client
-        self.wmsClient = None
         # # request client
         self.reqClient = None
         # # file catalog client
@@ -120,8 +122,6 @@ def initialize(self):
 
         # # transformation client
         self.transClient = TransformationClient()
-        # # wms client
-        self.wmsClient = WMSClient()
         # # request client
         self.reqClient = ReqClient()
         # # file catalog client
@@ -613,8 +613,8 @@ def __removeWMSTasks(self, transJobIDs):
         # Prevent 0 job IDs
         jobIDs = [int(j) for j in transJobIDs if int(j)]
         allRemove = True
-        for jobList in breakListIntoChunks(jobIDs, 500):
-            res = self.wmsClient.killJob(jobList, force=True)
+        for jobList in breakListIntoChunks(jobIDs, 1000):
+            res = kill_delete_jobs(RIGHT_KILL, jobList, force=True)
             if res["OK"]:
                 self.log.info(f"Successfully killed {len(jobList)} jobs from WMS")
             elif ("InvalidJobIDs" in res) and ("NonauthorizedJobIDs" not in res) and ("FailedJobIDs" not in res):
@@ -626,7 +626,7 @@ def __removeWMSTasks(self, transJobIDs):
                 self.log.error("Failed to kill jobs", f"(n={len(res['FailedJobIDs'])})")
                 allRemove = False
 
-            res = self.wmsClient.deleteJob(jobList)
+            res = kill_delete_jobs(RIGHT_DELETE, jobList, force=True)
             if res["OK"]:
                 self.log.info("Successfully deleted jobs from WMS", f"(n={len(jobList)})")
             elif ("InvalidJobIDs" in res) and ("NonauthorizedJobIDs" not in res) and ("FailedJobIDs" not in res):

src/DIRAC/WorkloadManagementSystem/Agent/JobCleaningAgent.py

@@ -35,10 +35,12 @@
 from DIRAC.RequestManagementSystem.Client.ReqClient import ReqClient
 from DIRAC.RequestManagementSystem.Client.Request import Request
 from DIRAC.WorkloadManagementSystem.Client import JobStatus
-from DIRAC.WorkloadManagementSystem.Client.JobMonitoringClient import JobMonitoringClient
 from DIRAC.WorkloadManagementSystem.Client.WMSClient import WMSClient
 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.Utilities.JobParameters import getJobParameters
 
 
 class JobCleaningAgent(AgentModule):
@@ -230,11 +232,11 @@ def _deleteRemoveJobs(self, jobList, remove=False):
             if not res["OK"]:
                 self.log.error("No DN found", f"for {user}")
                 return res
-            wmsClient = WMSClient(useCertificates=True, delegatedDN=res["Value"][0], delegatedGroup=ownerGroup)
             if remove:
+                wmsClient = WMSClient(useCertificates=True, delegatedDN=res["Value"][0], delegatedGroup=ownerGroup)
                 result = wmsClient.removeJob(jobsList)
             else:
-                result = wmsClient.deleteJob(jobsList)
+                result = kill_delete_jobs(RIGHT_DELETE, jobsList)
             if not result["OK"]:
                 self.log.error(
                     f"Could not {'remove' if remove else 'delete'} jobs",
@@ -294,7 +296,8 @@ def deleteJobOversizedSandbox(self, jobIDList):
         failed = {}
         successful = {}
 
-        result = JobMonitoringClient().getJobParameters(jobIDList, ["OutputSandboxLFN"])
+        jobIDs = [int(jobID) for jobID in jobIDList]
+        result = getJobParameters(jobIDs, "OutputSandboxLFN")
         if not result["OK"]:
             return result
         osLFNDict = result["Value"]

src/DIRAC/WorkloadManagementSystem/Agent/StalledJobAgent.py

@@ -14,16 +14,16 @@
 from DIRAC import S_ERROR, S_OK, gConfig
 from DIRAC.AccountingSystem.Client.Types.Job import Job
 from DIRAC.ConfigurationSystem.Client.Helpers import cfgPath
-from DIRAC.ConfigurationSystem.Client.Helpers.Registry import getDNForUsername
 from DIRAC.Core.Base.AgentModule import AgentModule
 from DIRAC.Core.Utilities import DErrno
 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.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.Service.JobPolicy import RIGHT_KILL
+from DIRAC.WorkloadManagementSystem.Utilities.jobAdministration import kill_delete_jobs
 from DIRAC.WorkloadManagementSystem.Utilities.JobParameters import getJobParameters
 from DIRAC.WorkloadManagementSystem.Utilities.Utils import rescheduleJobs
 
@@ -235,7 +235,7 @@ def _failStalledJobs(self, jobID):
         # Set the jobs Failed, send them a kill signal in case they are not really dead
         # and send accounting info
         if setFailed:
-            res = self._sendKillCommand(jobID)
+            res = kill_delete_jobs(RIGHT_KILL, [jobID], nonauthJobList=[], force=True)
             if not res["OK"]:
                 self.log.error("Failed to kill job", jobID)
 
@@ -574,26 +574,3 @@ def _failSubmittingJobs(self):
                 continue
 
         return S_OK()
-
-    def _sendKillCommand(self, job):
-        """Send a kill signal to the job such that it cannot continue running.
-
-        :param int job: ID of job to send kill command
-        """
-
-        res = self.jobDB.getJobAttribute(job, "Owner")
-        if not res["OK"]:
-            return res
-        owner = res["Value"]
-
-        res = self.jobDB.getJobAttribute(job, "OwnerGroup")
-        if not res["OK"]:
-            return res
-        ownerGroup = res["Value"]
-
-        wmsClient = WMSClient(
-            useCertificates=True,
-            delegatedDN=getDNForUsername(owner)["Value"][0] if owner else None,
-            delegatedGroup=ownerGroup,
-        )
-        return wmsClient.killJob(job)

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

@@ -1,10 +1,11 @@
 """ Test class for Job Cleaning Agent
 """
-import pytest
 from unittest.mock import MagicMock
 
+import pytest
+
 # DIRAC Components
-from DIRAC import gLogger, S_OK
+from DIRAC import S_OK, gLogger
 from DIRAC.WorkloadManagementSystem.Agent.JobCleaningAgent import JobCleaningAgent
 
 gLogger.setLevel("DEBUG")
@@ -97,7 +98,7 @@ def test_deleteJobsByStatus(jca, conditions, mockReplyInput, expected):
     "inputs, params, expected",
     [
         ([], {"OK": True, "Value": {}}, {"OK": True, "Value": {"Failed": {}, "Successful": {}}}),
-        (["a", "b"], {"OK": True, "Value": {}}, {"OK": True, "Value": {"Failed": {}, "Successful": {}}}),
+        (["123", "456"], {"OK": True, "Value": {}}, {"OK": True, "Value": {"Failed": {}, "Successful": {}}}),
         (
             [],
             {"OK": True, "Value": {1: {"OutputSandboxLFN": "/some/lfn/1.txt"}}},
@@ -112,11 +113,11 @@ def test_deleteJobsByStatus(jca, conditions, mockReplyInput, expected):
             {"OK": True, "Value": {"Failed": {}, "Successful": {1: "/some/lfn/1.txt", 2: "/some/other/lfn/2.txt"}}},
         ),
         (
-            ["a", "b"],
+            ["123", "456"],
             {"OK": True, "Value": {1: {"OutputSandboxLFN": "/some/lfn/1.txt"}}},
             {"OK": True, "Value": {"Failed": {}, "Successful": {1: "/some/lfn/1.txt"}}},
         ),
-        (["a", "b"], {"OK": False}, {"OK": False}),
+        (["123", "456"], {"OK": False}, {"OK": False}),
     ],
 )
 def test_deleteJobOversizedSandbox(mocker, inputs, params, expected):
@@ -129,15 +130,14 @@ def test_deleteJobOversizedSandbox(mocker, inputs, params, expected):
     mocker.patch(
         "DIRAC.WorkloadManagementSystem.Agent.JobCleaningAgent.getDNForUsername", return_value=S_OK(["/bih/boh/DN"])
     )
+    mocker.patch("DIRAC.WorkloadManagementSystem.Agent.JobCleaningAgent.getJobParameters", return_value=params)
 
     jobCleaningAgent = JobCleaningAgent()
     jobCleaningAgent.log = gLogger
     jobCleaningAgent.log.setLevel("DEBUG")
     jobCleaningAgent._AgentModule__configDefaults = mockAM
     jobCleaningAgent.initialize()
 
-    mockJMC.getJobParameters.return_value = params
-
     result = jobCleaningAgent.deleteJobOversizedSandbox(inputs)
 
     assert result == expected

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

@@ -28,8 +28,6 @@ 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.WMSClient", return_value=MagicMock())
-    mocker.patch("DIRAC.WorkloadManagementSystem.Agent.StalledJobAgent.getDNForUsername", return_value=MagicMock())
 
     stalledJobAgent = StalledJobAgent()
     stalledJobAgent._AgentModule__configDefaults = mockAM

src/DIRAC/WorkloadManagementSystem/Service/JobManagerHandler.py

@@ -378,8 +378,6 @@ def export_removeJob(self, jobIDs):
         :return: S_OK()/S_ERROR() -- confirmed job IDs
         """
 
-        # FIXME: extract the logic to a utility function
-
         jobList = self.__getJobList(jobIDs)
         if not jobList:
             return S_ERROR("Invalid job specification: " + str(jobIDs))

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

@@ -2,35 +2,35 @@
 """
 
 from unittest.mock import MagicMock
-import pytest
 
+import pytest
 
 # sut
 from DIRAC.WorkloadManagementSystem.Utilities.jobAdministration import kill_delete_jobs
 
 
 @pytest.mark.parametrize(
-    "jobIDs_list, right, expected_res, expected_value",
+    "jobIDs_list, right, filtered_jobs, expected_res, expected_value",
     [
-        ([], 'Kill', True, []),
-        ([], 'Delete', True, []),
-        (1, 'Kill', True, []),
-        (1, 'Kill', True, []),
-        ([1, 2], 'Kill', True, []),
-        ([1, 2], 'Kill', True, []),
-        (1, 'Kill', True, [1]),
-        ([1, 2], 'Kill', True, [1]),
-        ([1, 2], 'Kill', True, [1]),
-        ([1, 2], 'Kill', True, []),
-        ([1, 2], 'Kill', True, [1, 2]),
+        ([], "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]),
     ],
 )
-def test___kill_delete_jobs(mocker, jobIDs_list, right, expected_res, expected_value):
+def test___kill_delete_jobs(mocker, jobIDs_list, right, filtered_jobs, expected_res, expected_value):
     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", 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