Skip to content

Commit e98dcc8

Browse files
Robin-Van-de-MerghelRobin VAN DE MERGHEL
authored andcommitted
fix: Fix lint, and removed generated code
1 parent b6183bc commit e98dcc8

File tree

2 files changed

+36
-65
lines changed

2 files changed

+36
-65
lines changed
Lines changed: 29 additions & 52 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,3 @@
1-
from diracx.client._generated.models import BodyPilotsAddJobsToPilot, BodyPilotsAddPilotStamps, ScalarSearchSpec, SearchParams
21
from DIRAC.Core.Utilities.ReturnValues import convertToReturnValue
32

43
from DIRAC.Core.Security.DiracX import DiracXClient
@@ -9,20 +8,14 @@ class PilotManagerClient:
98
def addPilotReferences(self, pilot_stamps, VO, gridType="DIRAC", pilot_references={}):
109
with DiracXClient() as api:
1110
# We will move toward a stamp as identifier for the pilot
12-
return api.pilots.add_pilot_stamps(BodyPilotsAddPilotStamps(
13-
pilot_stamps=pilot_stamps,
14-
vo=VO,
15-
grid_type=gridType,
16-
pilot_references=pilot_references
17-
))
18-
11+
return api.pilots.add_pilot_stamps(
12+
{"pilot_stamps": pilot_stamps, "vo": VO, "grid_type": gridType, "pilot_references": pilot_references}
13+
)
1914

2015
def set_pilot_field(self, pilot_stamp, values_dict):
2116
with DiracXClient() as api:
2217
values_dict["PilotStamp"] = pilot_stamp
23-
return api.pilots.update_pilot_fields(
24-
values_dict
25-
)
18+
return api.pilots.update_pilot_fields(values_dict)
2619

2720
@convertToReturnValue
2821
def setPilotBenchmark(self, pilotStamp, mark):
@@ -34,68 +27,52 @@ def setAccountingFlag(self, pilotStamp, flag):
3427

3528
@convertToReturnValue
3629
def setPilotStatus(self, pilot_stamp, status, destination=None, reason=None, grid_site=None, queue=None):
37-
return self.set_pilot_field(pilot_stamp, {
38-
"Status": status,
39-
"DestinationSite": destination,
40-
"StatusReason": reason,
41-
"GridSite": grid_site,
42-
"Queue": queue
43-
})
44-
30+
return self.set_pilot_field(
31+
pilot_stamp,
32+
{
33+
"Status": status,
34+
"DestinationSite": destination,
35+
"StatusReason": reason,
36+
"GridSite": grid_site,
37+
"Queue": queue,
38+
},
39+
)
40+
4541
@convertToReturnValue
4642
def clearPilots(self, interval=30, aborted_interval=7):
4743
with DiracXClient() as api:
48-
api.pilots.delete_pilots(
49-
age_in_days=interval,
50-
delete_only_aborted=False
51-
)
52-
api.pilots.delete_pilots(
53-
age_in_days=aborted_interval,
54-
delete_only_aborted=True
55-
)
44+
api.pilots.delete_pilots(age_in_days=interval, delete_only_aborted=False)
45+
api.pilots.delete_pilots(age_in_days=aborted_interval, delete_only_aborted=True)
5646

5747
@convertToReturnValue
5848
def deletePilots(self, pilot_stamps):
5949
with DiracXClient() as api:
60-
api.pilots.delete_pilots(
61-
pilot_stamps=pilot_stamps
62-
)
50+
api.pilots.delete_pilots(pilot_stamps=pilot_stamps)
6351

6452
@convertToReturnValue
6553
def setJobForPilot(self, job_id, pilot_stamp, destination=None):
6654
with DiracXClient() as api:
67-
api.pilots.add_jobs_to_pilot(BodyPilotsAddJobsToPilot(
68-
pilot_stamp=pilot_stamp,
69-
job_ids=[job_id]
70-
))
55+
api.pilots.add_jobs_to_pilot({"pilot_stamp": pilot_stamp, "job_ids": [job_id]})
56+
57+
self.set_pilot_field(
58+
pilot_stamp,
59+
{
60+
"DestinationSite": destination,
61+
},
62+
)
7163

72-
self.set_pilot_field(pilot_stamp, {
73-
"DestinationSite": destination,
74-
})
75-
7664
@convertToReturnValue
7765
def getPilots(self, job_id):
7866
with DiracXClient() as api:
79-
pilot_ids = api.pilots.get_pilot_jobs(
80-
job_id=job_id
81-
)
67+
pilot_ids = api.pilots.get_pilot_jobs(job_id=job_id)
8268

8369
query = [{"parameter": "PilotID", "operator": "in", "value": pilot_ids}]
8470

85-
return api.pilots.search(
86-
parameters=[],
87-
search=query,
88-
sort=[]
89-
)
71+
return api.pilots.search(parameters=[], search=query, sort=[])
9072

91-
9273
@convertToReturnValue
9374
def getPilotInfo(self, pilot_stamp):
9475
with DiracXClient() as api:
9576
query = [{"parameter": "PilotStamp", "operator": "eq", "value": pilot_stamp}]
9677

97-
return api.pilots.search(
98-
parameters=[],
99-
search=query,
100-
sort=[]
101-
)
78+
return api.pilots.search(parameters=[], search=query, sort=[])

src/DIRAC/WorkloadManagementSystem/Service/PilotManagerHandler.py

Lines changed: 7 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -197,7 +197,6 @@ def _getRemotePilotOutput(self, pilotReference, pilotDict):
197197
# return res, correct or not
198198
return res
199199

200-
201200
##############################################################################
202201
types_selectPilots = [dict]
203202

@@ -277,7 +276,6 @@ def export_getGroupedPilotSummary(cls, columnList):
277276
"""
278277
return cls.pilotAgentsDB.getGroupedPilotSummary(columnList)
279278

280-
281279
types_countPilots = [dict]
282280

283281
@classmethod
@@ -286,7 +284,6 @@ def export_countPilots(cls, condDict, older=None, newer=None, timeStamp="Submiss
286284

287285
return cls.pilotAgentsDB.countPilots(condDict, older, newer, timeStamp)
288286

289-
290287
# --------------- Moved to DiracX ---------------
291288

292289
#############################################
@@ -299,13 +296,11 @@ def export_addPilotReferences(cls, pilotStamps, VO, gridType="DIRAC", pilotRefDi
299296
pilot_references = pilotRefDict.values()
300297
pilot_stamp_dict = dict(zip(pilotStamps, pilot_references))
301298

302-
303299
return cls.pilotAgentsDB.addPilotReferences(pilot_references, VO, gridType, pilot_stamp_dict)
304300

305-
306301
#############################################
307302
types_setJobForPilot = [[str, int], str]
308-
303+
309304
@classmethod
310305
def export_setJobForPilot(cls, jobID, pilotRef, destination=None):
311306
"""Report the DIRAC job ID which is executed by the given pilot job"""
@@ -337,7 +332,6 @@ def export_setAccountingFlag(cls, pilotRef, mark="True"):
337332
"""Set the pilot AccountingSent flag"""
338333
return cls.pilotAgentsDB.setAccountingFlag(pilotRef, mark)
339334

340-
341335
#############################################
342336
types_setPilotStatus = [str, str]
343337

@@ -351,12 +345,12 @@ def export_setPilotStatus(cls, pilotRef, status, destination=None, reason=None,
351345

352346
#############################################
353347
types_deletePilots = [[list, str, int]]
354-
348+
355349
@classmethod
356350
def export_deletePilots(cls, pilotIDs):
357351
if isinstance(pilotIDs, str):
358352
return cls.pilotAgentsDB.deletePilot(pilotIDs)
359-
353+
360354
# And list[str]????
361355
# pilot_id>>>S<<<
362356

@@ -377,8 +371,8 @@ def export_deletePilots(cls, pilotIDs):
377371
@classmethod
378372
def export_clearPilots(cls, interval=30, aborted_interval=7):
379373
return cls.pilotAgentsDB.clearPilots(interval, aborted_interval)
380-
381-
##############################################################################
374+
375+
#############################################
382376
types_getPilots = [[str, int]]
383377

384378
@classmethod
@@ -389,8 +383,8 @@ def export_getPilots(cls, jobID):
389383
return S_ERROR(f"Failed to get pilot for Job {int(jobID)}: {result.get('Message', '')}")
390384

391385
return cls.pilotAgentsDB.getPilotInfo(pilotID=result["Value"])
392-
393-
##############################################################################
386+
387+
#############################################
394388
types_getPilotInfo = [[list, str]]
395389

396390
@classmethod

0 commit comments

Comments
 (0)