@@ -34,6 +34,7 @@ from typing import Union
3434from pyslurm.utils import cstr, ctime
3535from pyslurm.utils.uint import *
3636from pyslurm.core.job.util import *
37+ from pyslurm.db.cluster import LOCAL_CLUSTER
3738from pyslurm.core.error import (
3839 RPCError,
3940 verify_rpc,
@@ -47,12 +48,14 @@ from pyslurm.utils.helpers import (
4748 _getgrall_to_dict,
4849 _getpwall_to_dict,
4950 instance_to_dict,
51+ collection_to_dict,
52+ group_collection_by_cluster,
5053 _sum_prop,
5154 _get_exit_code,
5255)
5356
5457
55- cdef class Jobs(dict ):
58+ cdef class Jobs(list ):
5659
5760 def __cinit__ (self ):
5861 self .info = NULL
@@ -63,14 +66,37 @@ cdef class Jobs(dict):
6366 def __init__ (self , jobs = None , frozen = False ):
6467 self .frozen = frozen
6568
66- if isinstance (jobs, dict ):
67- self .update(jobs)
68- elif jobs is not None :
69+ if isinstance (jobs, list ):
6970 for job in jobs:
7071 if isinstance (job, int ):
71- self [job] = Job(job)
72+ self .append( Job(job) )
7273 else :
73- self [job.id] = job
74+ self .append(job)
75+ elif isinstance (jobs, str ):
76+ joblist = jobs.split(" ," )
77+ self .extend([Job(int (job)) for job in joblist])
78+ elif isinstance (jobs, dict ):
79+ self .extend([job for job in jobs.values()])
80+ elif jobs is not None :
81+ raise TypeError (" Invalid Type: {type(jobs)}" )
82+
83+ def as_dict (self , recursive = False ):
84+ """ Convert the collection data to a dict.
85+
86+ Args:
87+ recursive (bool, optional):
88+ By default, the objects will not be converted to a dict. If
89+ this is set to `True`, then additionally all objects are
90+ converted to dicts.
91+
92+ Returns:
93+ (dict): Collection as a dict.
94+ """
95+ col = collection_to_dict(self , identifier = Job.id, recursive = recursive)
96+ return col.get(LOCAL_CLUSTER, {})
97+
98+ def group_by_cluster (self ):
99+ return group_collection_by_cluster(self )
74100
75101 @staticmethod
76102 def load (preload_passwd_info = False , frozen = False ):
@@ -124,7 +150,7 @@ cdef class Jobs(dict):
124150 job.passwd = passwd
125151 job.groups = groups
126152
127- jobs[job.id] = job
153+ jobs.append( job)
128154
129155 # At this point we memcpy'd all the memory for the Jobs. Setting this
130156 # to 0 will prevent the slurm job free function to deallocate the
@@ -143,28 +169,34 @@ cdef class Jobs(dict):
143169 Raises:
144170 RPCError: When getting the Jobs from the slurmctld failed.
145171 """
146- cdef Jobs reloaded_jobs = Jobs.load()
172+ cdef:
173+ Jobs reloaded_jobs
174+ Jobs new_jobs = Jobs()
175+ dict self_dict
147176
148- for jid in list (self .keys()):
177+ if not self :
178+ return self
179+
180+ reloaded_jobs = Jobs.load().as_dict()
181+ for idx, jid in enumerate (self ):
149182 if jid in reloaded_jobs:
150183 # Put the new data in.
151- self [jid] = reloaded_jobs[jid]
152- elif not self .frozen:
153- # Remove this instance from the current collection, as the Job
154- # doesn't exist anymore.
155- del self [jid]
184+ new_jobs.append(reloaded_jobs[jid])
156185
157186 if not self .frozen:
187+ self_dict = self .as_dict()
158188 for jid in reloaded_jobs:
159- if jid not in self :
160- self [jid] = reloaded_jobs[jid]
189+ if jid not in self_dict :
190+ new_jobs.append( reloaded_jobs[jid])
161191
192+ self .clear()
193+ self .extend(new_jobs)
162194 return self
163195
164196 def load_steps (self ):
165197 """ Load all Job steps for this collection of Jobs.
166198
167- This function fills in the " steps" attribute for all Jobs in the
199+ This function fills in the ` steps` attribute for all Jobs in the
168200 collection.
169201
170202 !!! note
@@ -175,21 +207,16 @@ cdef class Jobs(dict):
175207 RPCError: When retrieving the Job information for all the Steps
176208 failed.
177209 """
178- cdef dict step_info = JobSteps.load_all ()
210+ cdef dict steps = JobSteps.load().as_dict ()
179211
180- for jid in self :
212+ for idx, job in enumerate ( self ) :
181213 # Ignore any Steps from Jobs which do not exist in this
182214 # collection.
183- if jid in step_info:
184- self [jid].steps = step_info[jid]
185-
186- def as_list (self ):
187- """ Format the information as list of Job objects.
188-
189- Returns:
190- (list[pyslurm.Job]): List of Job objects
191- """
192- return list (self .values())
215+ jid = job.id
216+ if jid in steps:
217+ job_steps = self [idx].steps
218+ job_steps.clear()
219+ job_steps.extend(steps[jid].values())
193220
194221 @property
195222 def memory (self ):
@@ -218,6 +245,7 @@ cdef class Job:
218245 self .ptr.job_id = job_id
219246 self .passwd = {}
220247 self .groups = {}
248+ cstr.fmalloc(& self .ptr.cluster, LOCAL_CLUSTER)
221249 self .steps = JobSteps.__new__ (JobSteps)
222250
223251 def _alloc_impl (self ):
@@ -234,7 +262,9 @@ cdef class Job:
234262 self ._dealloc_impl()
235263
236264 def __eq__ (self , other ):
237- return isinstance (other, Job) and self .id == other.id
265+ if isinstance (other, Job):
266+ return self .id == other.id and self .cluster == other.cluster
267+ return NotImplemented
238268
239269 @staticmethod
240270 def load (job_id ):
@@ -278,7 +308,7 @@ cdef class Job:
278308 if not slurm.IS_JOB_PENDING(wrap.ptr):
279309 # Just ignore if the steps couldn't be loaded here.
280310 try :
281- wrap.steps = JobSteps._load (wrap)
311+ wrap.steps = JobSteps._load_single (wrap)
282312 except RPCError:
283313 pass
284314 else :
0 commit comments