Skip to content

Commit 36f3f93

Browse files
fix: Fix lint, and removed generated code
1 parent e8e7b13 commit 36f3f93

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
@@ -194,7 +194,6 @@ def _getRemotePilotOutput(self, pilotReference, pilotDict):
194194
# return res, correct or not
195195
return res
196196

197-
198197
##############################################################################
199198
types_selectPilots = [dict]
200199

@@ -274,7 +273,6 @@ def export_getGroupedPilotSummary(cls, columnList):
274273
"""
275274
return cls.pilotAgentsDB.getGroupedPilotSummary(columnList)
276275

277-
278276
types_countPilots = [dict]
279277

280278
@classmethod
@@ -283,7 +281,6 @@ def export_countPilots(cls, condDict, older=None, newer=None, timeStamp="Submiss
283281

284282
return cls.pilotAgentsDB.countPilots(condDict, older, newer, timeStamp)
285283

286-
287284
# --------------- Moved to DiracX ---------------
288285

289286
#############################################
@@ -296,13 +293,11 @@ def export_addPilotReferences(cls, pilotStamps, VO, gridType="DIRAC", pilotRefDi
296293
pilot_references = pilotRefDict.values()
297294
pilot_stamp_dict = dict(zip(pilotStamps, pilot_references))
298295

299-
300296
return cls.pilotAgentsDB.addPilotReferences(pilot_references, VO, gridType, pilot_stamp_dict)
301297

302-
303298
#############################################
304299
types_setJobForPilot = [[str, int], str]
305-
300+
306301
@classmethod
307302
def export_setJobForPilot(cls, jobID, pilotRef, destination=None):
308303
"""Report the DIRAC job ID which is executed by the given pilot job"""
@@ -334,7 +329,6 @@ def export_setAccountingFlag(cls, pilotRef, mark="True"):
334329
"""Set the pilot AccountingSent flag"""
335330
return cls.pilotAgentsDB.setAccountingFlag(pilotRef, mark)
336331

337-
338332
#############################################
339333
types_setPilotStatus = [str, str]
340334

@@ -348,12 +342,12 @@ def export_setPilotStatus(cls, pilotRef, status, destination=None, reason=None,
348342

349343
#############################################
350344
types_deletePilots = [[list, str, int]]
351-
345+
352346
@classmethod
353347
def export_deletePilots(cls, pilotIDs):
354348
if isinstance(pilotIDs, str):
355349
return cls.pilotAgentsDB.deletePilot(pilotIDs)
356-
350+
357351
# And list[str]????
358352
# pilot_id>>>S<<<
359353

@@ -374,8 +368,8 @@ def export_deletePilots(cls, pilotIDs):
374368
@classmethod
375369
def export_clearPilots(cls, interval=30, aborted_interval=7):
376370
return cls.pilotAgentsDB.clearPilots(interval, aborted_interval)
377-
378-
##############################################################################
371+
372+
#############################################
379373
types_getPilots = [[str, int]]
380374

381375
@classmethod
@@ -386,8 +380,8 @@ def export_getPilots(cls, jobID):
386380
return S_ERROR(f"Failed to get pilot for Job {int(jobID)}: {result.get('Message', '')}")
387381

388382
return cls.pilotAgentsDB.getPilotInfo(pilotID=result["Value"])
389-
390-
##############################################################################
383+
384+
#############################################
391385
types_getPilotInfo = [[list, str]]
392386

393387
@classmethod

0 commit comments

Comments
 (0)