Skip to content

Commit 4e66e2d

Browse files
fstagnifstagni
committed
feat: use directly the utilities instead of going through the service
1 parent 1a28820 commit 4e66e2d

7 files changed

Lines changed: 40 additions & 66 deletions

File tree

‎src/DIRAC/TransformationSystem/Agent/TransformationCleaningAgent.py‎

Lines changed: 8 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -16,14 +16,12 @@
1616

1717
# # from DIRAC
1818
from DIRAC import S_ERROR, S_OK
19-
from DIRAC.ConfigurationSystem.Client.ConfigurationData import gConfigurationData
2019
from DIRAC.ConfigurationSystem.Client.Helpers.Operations import Operations
2120
from DIRAC.Core.Base.AgentModule import AgentModule
2221
from DIRAC.Core.Utilities.DErrno import cmpError
2322
from DIRAC.Core.Utilities.List import breakListIntoChunks
2423
from DIRAC.Core.Utilities.Proxy import executeWithUserProxy
2524
from DIRAC.Core.Utilities.ReturnValues import returnSingleResult
26-
from DIRAC.DataManagementSystem.Client.DataManager import DataManager
2725
from DIRAC.RequestManagementSystem.Client.File import File
2826
from DIRAC.RequestManagementSystem.Client.Operation import Operation
2927
from DIRAC.RequestManagementSystem.Client.ReqClient import ReqClient
@@ -35,7 +33,11 @@
3533
from DIRAC.TransformationSystem.Client import TransformationStatus
3634
from DIRAC.TransformationSystem.Client.TransformationClient import TransformationClient
3735
from DIRAC.WorkloadManagementSystem.Client.JobMonitoringClient import JobMonitoringClient
38-
from DIRAC.WorkloadManagementSystem.Client.WMSClient import WMSClient
36+
from DIRAC.WorkloadManagementSystem.Service.JobPolicy import (
37+
RIGHT_DELETE,
38+
RIGHT_KILL,
39+
)
40+
from DIRAC.WorkloadManagementSystem.Utilities.jobAdministration import kill_delete_jobs
3941

4042
# # agent's name
4143
AGENT_NAME = "Transformation/TransformationCleaningAgent"
@@ -59,8 +61,6 @@ def __init__(self, *args, **kwargs):
5961

6062
# # transformation client
6163
self.transClient = None
62-
# # wms client
63-
self.wmsClient = None
6464
# # request client
6565
self.reqClient = None
6666
# # file catalog client
@@ -119,8 +119,6 @@ def initialize(self):
119119

120120
# # transformation client
121121
self.transClient = TransformationClient()
122-
# # wms client
123-
self.wmsClient = WMSClient()
124122
# # request client
125123
self.reqClient = ReqClient()
126124
# # file catalog client
@@ -610,8 +608,8 @@ def __removeWMSTasks(self, transJobIDs):
610608
# Prevent 0 job IDs
611609
jobIDs = [int(j) for j in transJobIDs if int(j)]
612610
allRemove = True
613-
for jobList in breakListIntoChunks(jobIDs, 500):
614-
res = self.wmsClient.killJob(jobList, force=True)
611+
for jobList in breakListIntoChunks(jobIDs, 1000):
612+
res = kill_delete_jobs(RIGHT_KILL, jobList, force=True)
615613
if res["OK"]:
616614
self.log.info(f"Successfully killed {len(jobList)} jobs from WMS")
617615
elif ("InvalidJobIDs" in res) and ("NonauthorizedJobIDs" not in res) and ("FailedJobIDs" not in res):
@@ -623,7 +621,7 @@ def __removeWMSTasks(self, transJobIDs):
623621
self.log.error("Failed to kill jobs", f"(n={len(res['FailedJobIDs'])})")
624622
allRemove = False
625623

626-
res = self.wmsClient.deleteJob(jobList)
624+
res = kill_delete_jobs(RIGHT_DELETE, jobList, force=True)
627625
if res["OK"]:
628626
self.log.info("Successfully deleted jobs from WMS", f"(n={len(jobList)})")
629627
elif ("InvalidJobIDs" in res) and ("NonauthorizedJobIDs" not in res) and ("FailedJobIDs" not in res):

‎src/DIRAC/WorkloadManagementSystem/Agent/JobCleaningAgent.py‎

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -35,10 +35,12 @@
3535
from DIRAC.RequestManagementSystem.Client.ReqClient import ReqClient
3636
from DIRAC.RequestManagementSystem.Client.Request import Request
3737
from DIRAC.WorkloadManagementSystem.Client import JobStatus
38-
from DIRAC.WorkloadManagementSystem.Client.JobMonitoringClient import JobMonitoringClient
3938
from DIRAC.WorkloadManagementSystem.Client.SandboxStoreClient import SandboxStoreClient
4039
from DIRAC.WorkloadManagementSystem.Client.WMSClient import WMSClient
4140
from DIRAC.WorkloadManagementSystem.DB.JobDB import JobDB
41+
from DIRAC.WorkloadManagementSystem.Service.JobPolicy import RIGHT_DELETE
42+
from DIRAC.WorkloadManagementSystem.Utilities.jobAdministration import kill_delete_jobs
43+
from DIRAC.WorkloadManagementSystem.Utilities.JobParameters import getJobParameters
4244

4345

4446
class JobCleaningAgent(AgentModule):
@@ -229,11 +231,11 @@ def _deleteRemoveJobs(self, jobList, remove=False):
229231
if not res["OK"]:
230232
self.log.error("No DN found", f"for {user}")
231233
return res
232-
wmsClient = WMSClient(useCertificates=True, delegatedDN=res["Value"][0], delegatedGroup=ownerGroup)
233234
if remove:
235+
wmsClient = WMSClient(useCertificates=True, delegatedDN=res["Value"][0], delegatedGroup=ownerGroup)
234236
result = wmsClient.removeJob(jobsList)
235237
else:
236-
result = wmsClient.deleteJob(jobsList)
238+
result = kill_delete_jobs(RIGHT_DELETE, jobsList)
237239
if not result["OK"]:
238240
self.log.error(
239241
f"Could not {'remove' if remove else 'delete'} jobs",
@@ -293,7 +295,8 @@ def deleteJobOversizedSandbox(self, jobIDList):
293295
failed = {}
294296
successful = {}
295297

296-
result = JobMonitoringClient().getJobParameters(jobIDList, ["OutputSandboxLFN"])
298+
jobIDs = [int(jobID) for jobID in jobIDList]
299+
result = getJobParameters(jobIDs, "OutputSandboxLFN")
297300
if not result["OK"]:
298301
return result
299302
osLFNDict = result["Value"]

‎src/DIRAC/WorkloadManagementSystem/Agent/StalledJobAgent.py‎

Lines changed: 3 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -14,16 +14,16 @@
1414
from DIRAC import S_ERROR, S_OK, gConfig
1515
from DIRAC.AccountingSystem.Client.Types.Job import Job
1616
from DIRAC.ConfigurationSystem.Client.Helpers import cfgPath
17-
from DIRAC.ConfigurationSystem.Client.Helpers.Registry import getDNForUsername
1817
from DIRAC.Core.Base.AgentModule import AgentModule
1918
from DIRAC.Core.Utilities import DErrno
2019
from DIRAC.Core.Utilities.ClassAd.ClassAdLight import ClassAd
2120
from DIRAC.Core.Utilities.TimeUtilities import fromString, second, toEpoch
2221
from DIRAC.WorkloadManagementSystem.Client import JobMinorStatus, JobStatus
23-
from DIRAC.WorkloadManagementSystem.Client.WMSClient import WMSClient
2422
from DIRAC.WorkloadManagementSystem.DB.JobDB import JobDB
2523
from DIRAC.WorkloadManagementSystem.DB.JobLoggingDB import JobLoggingDB
2624
from DIRAC.WorkloadManagementSystem.DB.PilotAgentsDB import PilotAgentsDB
25+
from DIRAC.WorkloadManagementSystem.Service.JobPolicy import RIGHT_KILL
26+
from DIRAC.WorkloadManagementSystem.Utilities.jobAdministration import kill_delete_jobs
2727
from DIRAC.WorkloadManagementSystem.Utilities.JobParameters import getJobParameters
2828
from DIRAC.WorkloadManagementSystem.Utilities.Utils import rescheduleJobs
2929

@@ -235,7 +235,7 @@ def _failStalledJobs(self, jobID):
235235
# Set the jobs Failed, send them a kill signal in case they are not really dead
236236
# and send accounting info
237237
if setFailed:
238-
res = self._sendKillCommand(jobID)
238+
res = kill_delete_jobs(RIGHT_KILL, [jobID], nonauthJobList=[], force=True)
239239
if not res["OK"]:
240240
self.log.error("Failed to kill job", jobID)
241241

@@ -574,26 +574,3 @@ def _failSubmittingJobs(self):
574574
continue
575575

576576
return S_OK()
577-
578-
def _sendKillCommand(self, job):
579-
"""Send a kill signal to the job such that it cannot continue running.
580-
581-
:param int job: ID of job to send kill command
582-
"""
583-
584-
res = self.jobDB.getJobAttribute(job, "Owner")
585-
if not res["OK"]:
586-
return res
587-
owner = res["Value"]
588-
589-
res = self.jobDB.getJobAttribute(job, "OwnerGroup")
590-
if not res["OK"]:
591-
return res
592-
ownerGroup = res["Value"]
593-
594-
wmsClient = WMSClient(
595-
useCertificates=True,
596-
delegatedDN=getDNForUsername(owner)["Value"][0] if owner else None,
597-
delegatedGroup=ownerGroup,
598-
)
599-
return wmsClient.killJob(job)

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

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,11 @@
11
""" Test class for Job Cleaning Agent
22
"""
3-
import pytest
43
from unittest.mock import MagicMock
54

5+
import pytest
6+
67
# DIRAC Components
7-
from DIRAC import gLogger, S_OK
8+
from DIRAC import S_OK, gLogger
89
from DIRAC.WorkloadManagementSystem.Agent.JobCleaningAgent import JobCleaningAgent
910

1011
gLogger.setLevel("DEBUG")
@@ -97,7 +98,7 @@ def test_deleteJobsByStatus(jca, conditions, mockReplyInput, expected):
9798
"inputs, params, expected",
9899
[
99100
([], {"OK": True, "Value": {}}, {"OK": True, "Value": {"Failed": {}, "Successful": {}}}),
100-
(["a", "b"], {"OK": True, "Value": {}}, {"OK": True, "Value": {"Failed": {}, "Successful": {}}}),
101+
(["123", "456"], {"OK": True, "Value": {}}, {"OK": True, "Value": {"Failed": {}, "Successful": {}}}),
101102
(
102103
[],
103104
{"OK": True, "Value": {1: {"OutputSandboxLFN": "/some/lfn/1.txt"}}},
@@ -112,11 +113,11 @@ def test_deleteJobsByStatus(jca, conditions, mockReplyInput, expected):
112113
{"OK": True, "Value": {"Failed": {}, "Successful": {1: "/some/lfn/1.txt", 2: "/some/other/lfn/2.txt"}}},
113114
),
114115
(
115-
["a", "b"],
116+
["123", "456"],
116117
{"OK": True, "Value": {1: {"OutputSandboxLFN": "/some/lfn/1.txt"}}},
117118
{"OK": True, "Value": {"Failed": {}, "Successful": {1: "/some/lfn/1.txt"}}},
118119
),
119-
(["a", "b"], {"OK": False}, {"OK": False}),
120+
(["123", "456"], {"OK": False}, {"OK": False}),
120121
],
121122
)
122123
def test_deleteJobOversizedSandbox(mocker, inputs, params, expected):
@@ -129,15 +130,14 @@ def test_deleteJobOversizedSandbox(mocker, inputs, params, expected):
129130
mocker.patch(
130131
"DIRAC.WorkloadManagementSystem.Agent.JobCleaningAgent.getDNForUsername", return_value=S_OK(["/bih/boh/DN"])
131132
)
133+
mocker.patch("DIRAC.WorkloadManagementSystem.Agent.JobCleaningAgent.getJobParameters", return_value=params)
132134

133135
jobCleaningAgent = JobCleaningAgent()
134136
jobCleaningAgent.log = gLogger
135137
jobCleaningAgent.log.setLevel("DEBUG")
136138
jobCleaningAgent._AgentModule__configDefaults = mockAM
137139
jobCleaningAgent.initialize()
138140

139-
mockJMC.getJobParameters.return_value = params
140-
141141
result = jobCleaningAgent.deleteJobOversizedSandbox(inputs)
142142

143143
assert result == expected

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

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -28,8 +28,6 @@ def sja(mocker):
2828
mocker.patch("DIRAC.WorkloadManagementSystem.Agent.StalledJobAgent.rescheduleJobs", return_value=MagicMock())
2929
mocker.patch("DIRAC.WorkloadManagementSystem.Agent.StalledJobAgent.PilotAgentsDB", return_value=MagicMock())
3030
mocker.patch("DIRAC.WorkloadManagementSystem.Agent.StalledJobAgent.getJobParameters", return_value=MagicMock())
31-
mocker.patch("DIRAC.WorkloadManagementSystem.Agent.StalledJobAgent.WMSClient", return_value=MagicMock())
32-
mocker.patch("DIRAC.WorkloadManagementSystem.Agent.StalledJobAgent.getDNForUsername", return_value=MagicMock())
3331

3432
stalledJobAgent = StalledJobAgent()
3533
stalledJobAgent._AgentModule__configDefaults = mockAM

‎src/DIRAC/WorkloadManagementSystem/Service/JobManagerHandler.py‎

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -388,8 +388,6 @@ def export_removeJob(self, jobIDs):
388388
:return: S_OK()/S_ERROR() -- confirmed job IDs
389389
"""
390390

391-
# FIXME: extract the logic to a utility function
392-
393391
jobList = self.__getJobList(jobIDs)
394392
if not jobList:
395393
return S_ERROR("Invalid job specification: " + str(jobIDs))

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

Lines changed: 15 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -2,35 +2,35 @@
22
"""
33

44
from unittest.mock import MagicMock
5-
import pytest
65

6+
import pytest
77

88
# sut
99
from DIRAC.WorkloadManagementSystem.Utilities.jobAdministration import kill_delete_jobs
1010

1111

1212
@pytest.mark.parametrize(
13-
"jobIDs_list, right, expected_res, expected_value",
13+
"jobIDs_list, right, filtered_jobs, expected_res, expected_value",
1414
[
15-
([], 'Kill', True, []),
16-
([], 'Delete', True, []),
17-
(1, 'Kill', True, []),
18-
(1, 'Kill', True, []),
19-
([1, 2], 'Kill', True, []),
20-
([1, 2], 'Kill', True, []),
21-
(1, 'Kill', True, [1]),
22-
([1, 2], 'Kill', True, [1]),
23-
([1, 2], 'Kill', True, [1]),
24-
([1, 2], 'Kill', True, []),
25-
([1, 2], 'Kill', True, [1, 2]),
15+
([], "Kill", [], True, []),
16+
([], "Delete", [], True, []),
17+
(1, "Kill", [], True, []),
18+
(1, "Kill", [1], True, [1]),
19+
([1, 2], "Kill", [], True, []),
20+
([1, 2], "Kill", [1], True, [1]),
21+
(1, "Kill", [1], True, [1]),
22+
([1, 2], "Kill", [1], True, [1]),
23+
([1, 2], "Kill", [2], True, [2]),
24+
([1, 2], "Kill", [], True, []),
25+
([1, 2], "Kill", [1,2], True, [1, 2]),
2626
],
2727
)
28-
def test___kill_delete_jobs(mocker, jobIDs_list, right, expected_res, expected_value):
28+
def test___kill_delete_jobs(mocker, jobIDs_list, right, filtered_jobs, expected_res, expected_value):
2929
mocker.patch("DIRAC.WorkloadManagementSystem.Utilities.jobAdministration.JobDB", MagicMock())
3030
mocker.patch("DIRAC.WorkloadManagementSystem.Utilities.jobAdministration.TaskQueueDB", MagicMock())
3131
mocker.patch("DIRAC.WorkloadManagementSystem.Utilities.jobAdministration.PilotAgentsDB", MagicMock())
3232
mocker.patch("DIRAC.WorkloadManagementSystem.Utilities.jobAdministration.StorageManagementDB", MagicMock())
33-
mocker.patch("DIRAC.WorkloadManagementSystem.Utilities.jobAdministration.filterJobStateTransition", MagicMock())
33+
mocker.patch("DIRAC.WorkloadManagementSystem.Utilities.jobAdministration.filterJobStateTransition", return_value={"OK": True, "Value": filtered_jobs})
3434

3535
res = kill_delete_jobs(right, jobIDs_list)
3636
assert res["OK"] == expected_res

0 commit comments

Comments
 (0)