1
- from diracx .client ._generated .models import BodyPilotsAddJobsToPilot , BodyPilotsAddPilotStamps , ScalarSearchSpec , SearchParams
2
1
from DIRAC .Core .Utilities .ReturnValues import convertToReturnValue
3
2
4
3
from DIRAC .Core .Security .DiracX import DiracXClient
@@ -9,20 +8,14 @@ class PilotManagerClient:
9
8
def addPilotReferences (self , pilot_stamps , VO , gridType = "DIRAC" , pilot_references = {}):
10
9
with DiracXClient () as api :
11
10
# 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
+ )
19
14
20
15
def set_pilot_field (self , pilot_stamp , values_dict ):
21
16
with DiracXClient () as api :
22
17
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 )
26
19
27
20
@convertToReturnValue
28
21
def setPilotBenchmark (self , pilotStamp , mark ):
@@ -34,68 +27,52 @@ def setAccountingFlag(self, pilotStamp, flag):
34
27
35
28
@convertToReturnValue
36
29
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
+
45
41
@convertToReturnValue
46
42
def clearPilots (self , interval = 30 , aborted_interval = 7 ):
47
43
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 )
56
46
57
47
@convertToReturnValue
58
48
def deletePilots (self , pilot_stamps ):
59
49
with DiracXClient () as api :
60
- api .pilots .delete_pilots (
61
- pilot_stamps = pilot_stamps
62
- )
50
+ api .pilots .delete_pilots (pilot_stamps = pilot_stamps )
63
51
64
52
@convertToReturnValue
65
53
def setJobForPilot (self , job_id , pilot_stamp , destination = None ):
66
54
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
+ )
71
63
72
- self .set_pilot_field (pilot_stamp , {
73
- "DestinationSite" : destination ,
74
- })
75
-
76
64
@convertToReturnValue
77
65
def getPilots (self , job_id ):
78
66
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 )
82
68
83
69
query = [{"parameter" : "PilotID" , "operator" : "in" , "value" : pilot_ids }]
84
70
85
- return api .pilots .search (
86
- parameters = [],
87
- search = query ,
88
- sort = []
89
- )
71
+ return api .pilots .search (parameters = [], search = query , sort = [])
90
72
91
-
92
73
@convertToReturnValue
93
74
def getPilotInfo (self , pilot_stamp ):
94
75
with DiracXClient () as api :
95
76
query = [{"parameter" : "PilotStamp" , "operator" : "eq" , "value" : pilot_stamp }]
96
77
97
- return api .pilots .search (
98
- parameters = [],
99
- search = query ,
100
- sort = []
101
- )
78
+ return api .pilots .search (parameters = [], search = query , sort = [])
0 commit comments