diff --git a/CHANGELOG.md b/CHANGELOG.md index df972286..4f6dd4c6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -16,9 +16,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - New attributes for a Database Job: - extra - failed_node -- Now possible to initialize a pyslurm.db.Jobs collection with existing job +- Now possible to initialize a [pyslurm.db.Jobs][] collection with existing job ids or pyslurm.db.Job objects - Added `as_dict` function to all Collections +- Added a new Base Class [MultiClusterMap][pyslurm.xcollections.MultiClusterMap] that some Collections inherit from. ### Fixed @@ -28,9 +29,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - the Job was older than a day ### Changed - -- All Collections (like [pyslurm.Jobs](https://pyslurm.github.io/23.2/reference/job/#pyslurm.Jobs)) inherit from `list` now instead of `dict` + - `JobSearchFilter` has been renamed to `JobFilter` +- Renamed `as_dict` Function of some classes to `to_dict` ## [23.2.1](https://github.com/PySlurm/pyslurm/releases/tag/v23.2.1) - 2023-05-18 diff --git a/docs/reference/config.md b/docs/reference/config.md index 94b0438e..a461aba5 100644 --- a/docs/reference/config.md +++ b/docs/reference/config.md @@ -7,4 +7,3 @@ title: Config removed in the future when a replacement is introduced ::: pyslurm.config - handler: python diff --git a/docs/reference/constants.md b/docs/reference/constants.md index dd659b4c..65301afb 100644 --- a/docs/reference/constants.md +++ b/docs/reference/constants.md @@ -3,5 +3,3 @@ title: constants --- ::: pyslurm.constants - handler: python - diff --git a/docs/reference/db/cluster.md b/docs/reference/db/cluster.md index e6d0a900..219988d5 100644 --- a/docs/reference/db/cluster.md +++ b/docs/reference/db/cluster.md @@ -7,4 +7,3 @@ title: Cluster removed in the future when a replacement is introduced ::: pyslurm.slurmdb_clusters - handler: python diff --git a/docs/reference/db/connection.md b/docs/reference/db/connection.md index 27c904fc..7d77639e 100644 --- a/docs/reference/db/connection.md +++ b/docs/reference/db/connection.md @@ -3,4 +3,3 @@ title: Connection --- ::: pyslurm.db.Connection - handler: python diff --git a/docs/reference/db/event.md b/docs/reference/db/event.md index 020abcac..2816aaae 100644 --- a/docs/reference/db/event.md +++ b/docs/reference/db/event.md @@ -7,4 +7,3 @@ title: Event removed in the future when a replacement is introduced ::: pyslurm.slurmdb_events - handler: python diff --git a/docs/reference/db/job.md b/docs/reference/db/job.md index a2c7fadd..e806cc1f 100644 --- a/docs/reference/db/job.md +++ b/docs/reference/db/job.md @@ -7,7 +7,4 @@ title: Job will be removed in a future release ::: pyslurm.db.Job - handler: python - ::: pyslurm.db.Jobs - handler: python diff --git a/docs/reference/db/jobfilter.md b/docs/reference/db/jobfilter.md index 21aa55d1..523d7c9c 100644 --- a/docs/reference/db/jobfilter.md +++ b/docs/reference/db/jobfilter.md @@ -3,4 +3,3 @@ title: JobFilter --- ::: pyslurm.db.JobFilter - handler: python diff --git a/docs/reference/db/jobstats.md b/docs/reference/db/jobstats.md index 35f31ac6..1bc17d20 100644 --- a/docs/reference/db/jobstats.md +++ b/docs/reference/db/jobstats.md @@ -3,4 +3,3 @@ title: JobStatistics --- ::: pyslurm.db.JobStatistics - handler: python diff --git a/docs/reference/db/jobstep.md b/docs/reference/db/jobstep.md index 392fab65..a7bdc720 100644 --- a/docs/reference/db/jobstep.md +++ b/docs/reference/db/jobstep.md @@ -3,7 +3,4 @@ title: JobStep --- ::: pyslurm.db.JobStep - handler: python - ::: pyslurm.db.JobSteps - handler: python diff --git a/docs/reference/db/reservation.md b/docs/reference/db/reservation.md index 1a1af0c4..c1f110a3 100644 --- a/docs/reference/db/reservation.md +++ b/docs/reference/db/reservation.md @@ -7,4 +7,3 @@ title: Reservation removed in the future when a replacement is introduced ::: pyslurm.slurmdb_reservations - handler: python diff --git a/docs/reference/exceptions.md b/docs/reference/exceptions.md index 90876435..4abc0047 100644 --- a/docs/reference/exceptions.md +++ b/docs/reference/exceptions.md @@ -3,7 +3,4 @@ title: Exceptions --- ::: pyslurm.PyslurmError - handler: python - ::: pyslurm.RPCError - handler: python diff --git a/docs/reference/frontend.md b/docs/reference/frontend.md index 5247e540..f56a7ecd 100644 --- a/docs/reference/frontend.md +++ b/docs/reference/frontend.md @@ -7,4 +7,3 @@ title: Frontend removed in the future when a replacement is introduced ::: pyslurm.front_end - handler: python diff --git a/docs/reference/hostlist.md b/docs/reference/hostlist.md index dc2d81ee..33f8485d 100644 --- a/docs/reference/hostlist.md +++ b/docs/reference/hostlist.md @@ -7,4 +7,3 @@ title: Hostlist removed in the future when a replacement is introduced ::: pyslurm.hostlist - handler: python diff --git a/docs/reference/job.md b/docs/reference/job.md index 8e3d0c6e..cb1c19eb 100644 --- a/docs/reference/job.md +++ b/docs/reference/job.md @@ -7,7 +7,4 @@ title: Job removed in a future release ::: pyslurm.Job - handler: python - ::: pyslurm.Jobs - handler: python diff --git a/docs/reference/jobstep.md b/docs/reference/jobstep.md index 2ce6ef7f..b7b3e2b9 100644 --- a/docs/reference/jobstep.md +++ b/docs/reference/jobstep.md @@ -7,7 +7,4 @@ title: JobStep will be removed in a future release ::: pyslurm.JobStep - handler: python - ::: pyslurm.JobSteps - handler: python diff --git a/docs/reference/jobsubmitdescription.md b/docs/reference/jobsubmitdescription.md index bd31bac9..bf7eb6bd 100644 --- a/docs/reference/jobsubmitdescription.md +++ b/docs/reference/jobsubmitdescription.md @@ -3,4 +3,3 @@ title: JobSubmitDescription --- ::: pyslurm.JobSubmitDescription - handler: python diff --git a/docs/reference/node.md b/docs/reference/node.md index ccb16c54..e8e8d619 100644 --- a/docs/reference/node.md +++ b/docs/reference/node.md @@ -7,7 +7,4 @@ title: Node removed in a future release ::: pyslurm.Node - handler: python - ::: pyslurm.Nodes - handler: python diff --git a/docs/reference/partition.md b/docs/reference/partition.md index b9701f55..9181e10f 100644 --- a/docs/reference/partition.md +++ b/docs/reference/partition.md @@ -7,7 +7,4 @@ title: Partition will be removed in a future release ::: pyslurm.Partition - handler: python - ::: pyslurm.Partitions - handler: python diff --git a/docs/reference/reservation.md b/docs/reference/reservation.md index 563e29db..c5a3d891 100644 --- a/docs/reference/reservation.md +++ b/docs/reference/reservation.md @@ -7,4 +7,3 @@ title: Reservation removed in the future when a replacement is introduced ::: pyslurm.reservation - handler: python diff --git a/docs/reference/statistics.md b/docs/reference/statistics.md index 1f2b2e37..043461f8 100644 --- a/docs/reference/statistics.md +++ b/docs/reference/statistics.md @@ -7,4 +7,3 @@ title: Statistics removed in the future when a replacement is introduced ::: pyslurm.statistics - handler: python diff --git a/docs/reference/topology.md b/docs/reference/topology.md index 1cb107a1..c6b8f9cc 100644 --- a/docs/reference/topology.md +++ b/docs/reference/topology.md @@ -7,4 +7,3 @@ title: Topology removed in the future when a replacement is introduced ::: pyslurm.topology - handler: python diff --git a/docs/reference/trigger.md b/docs/reference/trigger.md index 308a3e3f..e6ea1e98 100644 --- a/docs/reference/trigger.md +++ b/docs/reference/trigger.md @@ -7,4 +7,3 @@ title: Trigger removed in the future when a replacement is introduced ::: pyslurm.trigger - handler: python diff --git a/docs/reference/utilities.md b/docs/reference/utilities.md index 63eb7bc0..dbf4a09e 100644 --- a/docs/reference/utilities.md +++ b/docs/reference/utilities.md @@ -3,37 +3,17 @@ title: utils --- ::: pyslurm.utils - handler: python + options: + members: [] ::: pyslurm.utils.timestr_to_secs - handler: python - ::: pyslurm.utils.timestr_to_mins - handler: python - ::: pyslurm.utils.secs_to_timestr - handler: python - ::: pyslurm.utils.mins_to_timestr - handler: python - ::: pyslurm.utils.date_to_timestamp - handler: python - ::: pyslurm.utils.timestamp_to_date - handler: python - ::: pyslurm.utils.expand_range_str - handler: python - ::: pyslurm.utils.humanize - handler: python - ::: pyslurm.utils.dehumanize - handler: python - ::: pyslurm.utils.nodelist_from_range_str - handler: python - ::: pyslurm.utils.nodelist_to_range_str - handler: python diff --git a/docs/reference/xcollections.md b/docs/reference/xcollections.md new file mode 100644 index 00000000..fd57ec09 --- /dev/null +++ b/docs/reference/xcollections.md @@ -0,0 +1,16 @@ +--- +title: xcollections +--- + +::: pyslurm.xcollections + handler: python + options: + members: + - MultiClusterMap + - BaseView + - KeysView + - MCKeysView + - ItemsView + - MCItemsView + - ValuesView + - ClustersView diff --git a/docs/stylesheets/extra.css b/docs/stylesheets/extra.css index 9562d9be..eab891415 100644 --- a/docs/stylesheets/extra.css +++ b/docs/stylesheets/extra.css @@ -2,3 +2,9 @@ .md-grid { max-width: 75%; } + +/* Indentation. */ +div.doc-contents:not(.first) { + padding-left: 25px; + border-left: .05rem solid var(--md-typeset-table-color); +} diff --git a/mkdocs.yml b/mkdocs.yml index daea3007..9d81f66b 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -56,6 +56,8 @@ plugins: docstring_style: google show_signature: true show_root_heading: true + show_symbol_type_toc: true + show_symbol_type_heading: true markdown_extensions: - admonition diff --git a/pyslurm/core/job/job.pxd b/pyslurm/core/job/job.pxd index bee4f9ec..4eb89bde 100644 --- a/pyslurm/core/job/job.pxd +++ b/pyslurm/core/job/job.pxd @@ -25,14 +25,12 @@ from pyslurm.utils cimport cstr, ctime from pyslurm.utils.uint cimport * from pyslurm.utils.ctime cimport time_t - from libc.string cimport memcpy, memset from libc.stdint cimport uint8_t, uint16_t, uint32_t, uint64_t, int64_t from libc.stdlib cimport free - from pyslurm.core.job.submission cimport JobSubmitDescription from pyslurm.core.job.step cimport JobSteps, JobStep - +from pyslurm.xcollections cimport MultiClusterMap from pyslurm cimport slurm from pyslurm.slurm cimport ( working_cluster_rec, @@ -67,8 +65,8 @@ from pyslurm.slurm cimport ( ) -cdef class Jobs(list): - """A collection of [pyslurm.Job][] objects. +cdef class Jobs(MultiClusterMap): + """A [`Multi Cluster`][pyslurm.xcollections.MultiClusterMap] collection of [pyslurm.Job][] objects. Args: jobs (Union[list, dict], optional=None): @@ -90,7 +88,7 @@ cdef class Jobs(list): This is the result of multiplying the run_time with the amount of cpus for each job. frozen (bool): - If this is set to True and the reload() method is called, then + If this is set to True and the `reload()` method is called, then *ONLY* Jobs that already exist in this collection will be reloaded. New Jobs that are discovered will not be added to this collection, but old Jobs which have already been purged from the @@ -115,15 +113,12 @@ cdef class Job: job_id (int): An Integer representing a Job-ID. - Raises: - MemoryError: If malloc fails to allocate memory. - Attributes: steps (JobSteps): Steps this Job has. Before you can access the Steps data for a Job, you have to call - the reload() method of a Job instance or the load_steps() method - of a Jobs collection. + the `reload()` method of a Job instance or the `load_steps()` + method of a Jobs collection. name (str): Name of the Job id (int): diff --git a/pyslurm/core/job/job.pyx b/pyslurm/core/job/job.pyx index 2c33d581..e2915608 100644 --- a/pyslurm/core/job/job.pyx +++ b/pyslurm/core/job/job.pyx @@ -34,7 +34,8 @@ from typing import Union from pyslurm.utils import cstr, ctime from pyslurm.utils.uint import * from pyslurm.core.job.util import * -from pyslurm.db.cluster import LOCAL_CLUSTER +from pyslurm.settings import LOCAL_CLUSTER +from pyslurm import xcollections from pyslurm.core.error import ( RPCError, verify_rpc, @@ -48,14 +49,11 @@ from pyslurm.utils.helpers import ( _getgrall_to_dict, _getpwall_to_dict, instance_to_dict, - collection_to_dict, - group_collection_by_cluster, - _sum_prop, _get_exit_code, ) -cdef class Jobs(list): +cdef class Jobs(MultiClusterMap): def __cinit__(self): self.info = NULL @@ -65,38 +63,11 @@ cdef class Jobs(list): def __init__(self, jobs=None, frozen=False): self.frozen = frozen - - if isinstance(jobs, list): - for job in jobs: - if isinstance(job, int): - self.append(Job(job)) - else: - self.append(job) - elif isinstance(jobs, str): - joblist = jobs.split(",") - self.extend([Job(int(job)) for job in joblist]) - elif isinstance(jobs, dict): - self.extend([job for job in jobs.values()]) - elif jobs is not None: - raise TypeError("Invalid Type: {type(jobs)}") - - def as_dict(self, recursive=False): - """Convert the collection data to a dict. - - Args: - recursive (bool, optional): - By default, the objects will not be converted to a dict. If - this is set to `True`, then additionally all objects are - converted to dicts. - - Returns: - (dict): Collection as a dict. - """ - col = collection_to_dict(self, identifier=Job.id, recursive=recursive) - return col.get(LOCAL_CLUSTER, {}) - - def group_by_cluster(self): - return group_collection_by_cluster(self) + super().__init__(data=jobs, + typ="Jobs", + val_type=Job, + id_attr=Job.id, + key_type=int) @staticmethod def load(preload_passwd_info=False, frozen=False): @@ -122,7 +93,7 @@ cdef class Jobs(list): cdef: dict passwd = {} dict groups = {} - Jobs jobs = Jobs.__new__(Jobs) + Jobs jobs = Jobs(frozen=frozen) int flags = slurm.SHOW_ALL | slurm.SHOW_DETAIL Job job @@ -150,16 +121,13 @@ cdef class Jobs(list): job.passwd = passwd job.groups = groups - jobs.append(job) + cluster = job.cluster + if cluster not in jobs.data: + jobs.data[cluster] = {} + jobs[cluster][job.id] = job - # At this point we memcpy'd all the memory for the Jobs. Setting this - # to 0 will prevent the slurm job free function to deallocate the - # memory for the individual jobs. This should be fine, because they - # are free'd automatically in __dealloc__ since the lifetime of each - # job-pointer is tied to the lifetime of its corresponding "Job" - # instance. + # We have extracted all pointers jobs.info.record_count = 0 - jobs.frozen = frozen return jobs @@ -169,29 +137,7 @@ cdef class Jobs(list): Raises: RPCError: When getting the Jobs from the slurmctld failed. """ - cdef: - Jobs reloaded_jobs - Jobs new_jobs = Jobs() - dict self_dict - - if not self: - return self - - reloaded_jobs = Jobs.load().as_dict() - for idx, jid in enumerate(self): - if jid in reloaded_jobs: - # Put the new data in. - new_jobs.append(reloaded_jobs[jid]) - - if not self.frozen: - self_dict = self.as_dict() - for jid in reloaded_jobs: - if jid not in self_dict: - new_jobs.append(reloaded_jobs[jid]) - - self.clear() - self.extend(new_jobs) - return self + return xcollections.multi_reload(self, frozen=self.frozen) def load_steps(self): """Load all Job steps for this collection of Jobs. @@ -207,32 +153,27 @@ cdef class Jobs(list): RPCError: When retrieving the Job information for all the Steps failed. """ - cdef dict steps = JobSteps.load().as_dict() - - for idx, job in enumerate(self): - # Ignore any Steps from Jobs which do not exist in this - # collection. + cdef dict steps = JobSteps.load_all() + for job in self.values(): jid = job.id if jid in steps: - job_steps = self[idx].steps - job_steps.clear() - job_steps.extend(steps[jid].values()) + job.steps = steps[jid] @property def memory(self): - return _sum_prop(self, Job.memory) + return xcollections.sum_property(self, Job.memory) @property def cpus(self): - return _sum_prop(self, Job.cpus) + return xcollections.sum_property(self, Job.cpus) @property def ntasks(self): - return _sum_prop(self, Job.ntasks) + return xcollections.sum_property(self, Job.ntasks) @property def cpu_time(self): - return _sum_prop(self, Job.cpu_time) + return xcollections.sum_property(self, Job.cpu_time) cdef class Job: @@ -246,7 +187,7 @@ cdef class Job: self.passwd = {} self.groups = {} cstr.fmalloc(&self.ptr.cluster, LOCAL_CLUSTER) - self.steps = JobSteps.__new__(JobSteps) + self.steps = JobSteps() def _alloc_impl(self): if not self.ptr: @@ -261,10 +202,8 @@ cdef class Job: def __dealloc__(self): self._dealloc_impl() - def __eq__(self, other): - if isinstance(other, Job): - return self.id == other.id and self.cluster == other.cluster - return NotImplemented + def __repr__(self): + return f'{self.__class__.__name__}({self.id})' @staticmethod def load(job_id): @@ -329,7 +268,6 @@ cdef class Job: wrap.groups = {} wrap.steps = JobSteps.__new__(JobSteps) memcpy(wrap.ptr, in_ptr, sizeof(slurm_job_info_t)) - return wrap cdef _swap_data(Job dst, Job src): @@ -340,6 +278,9 @@ cdef class Job: src.ptr = tmp def as_dict(self): + return self.to_dict() + + def to_dict(self): """Job information formatted as a dictionary. Returns: @@ -355,14 +296,14 @@ cdef class Job: Args: signal (Union[str, int]): Any valid signal which will be sent to the Job. Can be either - a str like 'SIGUSR1', or simply an int. + a str like `SIGUSR1`, or simply an [int][]. steps (str): Selects which steps should be signaled. Valid values for this - are: "all", "batch" and "children". The default value is - "children", where all steps except the batch-step will be + are: `all`, `batch` and `children`. The default value is + `children`, where all steps except the batch-step will be signaled. - The value "batch" in contrast means, that only the batch-step - will be signaled. With "all" every step is signaled. + The value `batch` in contrast means, that only the batch-step + will be signaled. With `all` every step is signaled. hurry (bool): If True, no burst buffer data will be staged out. The default value is False. @@ -480,9 +421,9 @@ cdef class Job: Args: mode (str): Determines in which mode the Job should be held. Possible - values are "user" or "admin". By default, the Job is held in - "admin" mode, meaning only an Administrator will be able to - release the Job again. If you specify the mode as "user", the + values are `user` or `admin`. By default, the Job is held in + `admin` mode, meaning only an Administrator will be able to + release the Job again. If you specify the mode as `user`, the User will also be able to release the job. Raises: @@ -524,7 +465,7 @@ cdef class Job: Args: hold (bool, optional): Controls whether the Job should be put in a held state or not. - Default for this is 'False', so it will not be held. + Default for this is `False`, so it will not be held. Raises: RPCError: When requeing the Job was not successful. @@ -1242,8 +1183,9 @@ cdef class Job: Return type may still be subject to change in the future Returns: - (dict): Resource layout, where the key is the name of the name and - its value another dict with the CPU-ids, memory and gres. + (dict): Resource layout, where the key is the name of the node and + the value another dict with the keys `cpu_ids`, `memory` and + `gres`. """ # The code for this function is a modified reimplementation from here: # https://github.com/SchedMD/slurm/blob/d525b6872a106d32916b33a8738f12510ec7cf04/src/api/job_info.c#L739 diff --git a/pyslurm/core/job/step.pxd b/pyslurm/core/job/step.pxd index 458ee506..489e9d64 100644 --- a/pyslurm/core/job/step.pxd +++ b/pyslurm/core/job/step.pxd @@ -49,16 +49,11 @@ from pyslurm.utils.ctime cimport time_t from pyslurm.core.job.task_dist cimport TaskDistribution -cdef class JobSteps(list): - """A collection of [pyslurm.JobStep][] objects for a given Job. - - Args: - job (Union[Job, int]): - A Job for which the Steps should be loaded. +cdef class JobSteps(dict): + """A [dict][] of [pyslurm.JobStep][] objects for a given Job. Raises: RPCError: When getting the Job steps from the slurmctld failed. - MemoryError: If malloc fails to allocate memory. """ cdef: @@ -68,8 +63,7 @@ cdef class JobSteps(list): @staticmethod cdef JobSteps _load_single(Job job) - - cdef _load_data(self, uint32_t job_id, int flags) + cdef dict _load_data(self, uint32_t job_id, int flags) cdef class JobStep: @@ -85,9 +79,6 @@ cdef class JobStep: time_limit (int): Time limit in Minutes for this step. - Raises: - MemoryError: If malloc fails to allocate memory. - Attributes: id (Union[str, int]): The id for this step. diff --git a/pyslurm/core/job/step.pyx b/pyslurm/core/job/step.pyx index d4038f54..54cb8f59 100644 --- a/pyslurm/core/job/step.pyx +++ b/pyslurm/core/job/step.pyx @@ -26,13 +26,12 @@ from typing import Union from pyslurm.utils import cstr, ctime from pyslurm.utils.uint import * from pyslurm.core.error import RPCError, verify_rpc -from pyslurm.db.cluster import LOCAL_CLUSTER +from pyslurm.settings import LOCAL_CLUSTER +from pyslurm import xcollections from pyslurm.utils.helpers import ( signal_to_num, instance_to_dict, uid_to_name, - collection_to_dict, - group_collection_by_cluster, humanize_step_id, dehumanize_step_id, ) @@ -46,7 +45,7 @@ from pyslurm.utils.ctime import ( ) -cdef class JobSteps(list): +cdef class JobSteps(dict): def __dealloc__(self): slurm_free_job_step_info_response_msg(self.info) @@ -55,73 +54,48 @@ cdef class JobSteps(list): self.info = NULL def __init__(self, steps=None): - if isinstance(steps, list): - self.extend(steps) + if isinstance(steps, dict): + self.update(steps) elif steps is not None: raise TypeError("Invalid Type: {type(steps)}") - def as_dict(self, recursive=False): - """Convert the collection data to a dict. - - Args: - recursive (bool, optional): - By default, the objects will not be converted to a dict. If - this is set to `True`, then additionally all objects are - converted to dicts. - - Returns: - (dict): Collection as a dict. - """ - col = collection_to_dict(self, identifier=JobStep.id, - recursive=recursive, group_id=JobStep.job_id) - col = col.get(LOCAL_CLUSTER, {}) - if self._job_id: - return col.get(self._job_id, {}) - - return col - - def group_by_cluster(self): - return group_collection_by_cluster(self) - @staticmethod - def load(job_id=0): + def load(job): """Load the Job Steps from the system. Args: - job_id (Union[Job, int]): + job (Union[Job, int]): The Job for which the Steps should be loaded. Returns: (pyslurm.JobSteps): JobSteps of the Job """ cdef: - Job job + Job _job JobSteps steps - if job_id: - job = Job.load(job_id.id if isinstance(job_id, Job) else job_id) - steps = JobSteps._load_single(job) - steps._job_id = job.id - return steps - else: - steps = JobSteps() - return steps._load_data(0, slurm.SHOW_ALL) + _job = Job.load(job.id if isinstance(job, Job) else job) + steps = JobSteps._load_single(_job) + steps._job_id = _job.id + return steps @staticmethod cdef JobSteps _load_single(Job job): cdef JobSteps steps = JobSteps() - steps._load_data(job.id, slurm.SHOW_ALL) - if not steps and not slurm.IS_JOB_PENDING(job.ptr): + data = steps._load_data(job.id, slurm.SHOW_ALL) + if not data and not slurm.IS_JOB_PENDING(job.ptr): msg = f"Failed to load step info for Job {job.id}." raise RPCError(msg=msg) + steps.update(data[job.id]) return steps - - cdef _load_data(self, uint32_t job_id, int flags): + + cdef dict _load_data(self, uint32_t job_id, int flags): cdef: JobStep step uint32_t cnt = 0 + dict steps = {} rc = slurm_get_job_steps(0, job_id, slurm.NO_VAL, &self.info, flags) @@ -133,21 +107,29 @@ cdef class JobSteps(list): # Put each job-step pointer into its own "JobStep" instance. for cnt in range(self.info.job_step_count): step = JobStep.from_ptr(&self.info.job_steps[cnt]) - # Prevent double free if xmalloc fails mid-loop and a MemoryError # is raised by replacing it with a zeroed-out job_step_info_t. self.info.job_steps[cnt] = self.tmp_info - self.append(step) - - # At this point we memcpy'd all the memory for the Steps. Setting this - # to 0 will prevent the slurm step free function to deallocate the - # memory for the individual steps. This should be fine, because they - # are free'd automatically in __dealloc__ since the lifetime of each - # step-pointer is tied to the lifetime of its corresponding JobStep - # instance. + + job_id = step.job_id + if not job_id in steps: + steps[job_id] = JobSteps() + steps[job_id][step.id] = step + + # We have extracted all pointers self.info.job_step_count = 0 + return steps - return self + @staticmethod + def load_all(): + """Loads all the steps in the system. + + Returns: + (dict): A dict where every JobID (key) is mapped with an instance + of its JobSteps (value). + """ + cdef JobSteps steps = JobSteps() + return steps._load_data(slurm.NO_VAL, slurm.SHOW_ALL) cdef class JobStep: @@ -160,6 +142,7 @@ cdef class JobStep: self._alloc_impl() self.job_id = job_id.id if isinstance(job_id, Job) else job_id self.id = step_id + cstr.fmalloc(&self.ptr.cluster, LOCAL_CLUSTER) # Initialize attributes, if any were provided for k, v in kwargs.items(): @@ -203,6 +186,9 @@ cdef class JobStep: # Call descriptors __set__ directly JobStep.__dict__[name].__set__(self, val) + def __repr__(self): + return f'{self.__class__.__name__}({self.id})' + @staticmethod def load(job_id, step_id): """Load information for a specific job step. @@ -221,7 +207,6 @@ cdef class JobStep: Raises: RPCError: When retrieving Step information from the slurmctld was not successful. - MemoryError: If malloc failed to allocate memory. Examples: >>> import pyslurm @@ -264,7 +249,7 @@ cdef class JobStep: Args: signal (Union[str, int]): Any valid signal which will be sent to the Job. Can be either - a str like 'SIGUSR1', or simply an int. + a str like `SIGUSR1`, or simply an [int][]. Raises: RPCError: When sending the signal was not successful. @@ -326,6 +311,9 @@ cdef class JobStep: verify_rpc(slurm_update_step(js.umsg)) def as_dict(self): + return self.to_dict() + + def to_dict(self): """JobStep information formatted as a dictionary. Returns: diff --git a/pyslurm/core/job/submission.pyx b/pyslurm/core/job/submission.pyx index bf47105b..df33992b 100644 --- a/pyslurm/core/job/submission.pyx +++ b/pyslurm/core/job/submission.pyx @@ -81,7 +81,6 @@ cdef class JobSubmitDescription: Raises: RPCError: When the job submission was not successful. - MemoryError: If malloc failed to allocate enough memory. Examples: >>> import pyslurm diff --git a/pyslurm/core/node.pxd b/pyslurm/core/node.pxd index ea59e6ff..5167de78 100644 --- a/pyslurm/core/node.pxd +++ b/pyslurm/core/node.pxd @@ -55,10 +55,11 @@ from pyslurm.utils cimport cstr from pyslurm.utils cimport ctime from pyslurm.utils.ctime cimport time_t from pyslurm.utils.uint cimport * +from pyslurm.xcollections cimport MultiClusterMap -cdef class Nodes(list): - """A collection of [pyslurm.Node][] objects. +cdef class Nodes(MultiClusterMap): + """A [`Multi Cluster`][pyslurm.xcollections.MultiClusterMap] collection of [pyslurm.Node][] objects. Args: nodes (Union[list, dict, str], optional=None): @@ -83,9 +84,6 @@ cdef class Nodes(list): Total amount of Watts consumed in this node collection. avg_watts (int): Amount of average watts consumed in this node collection. - - Raises: - MemoryError: If malloc fails to allocate memory. """ cdef: node_info_msg_t *info @@ -165,7 +163,7 @@ cdef class Node: memory_reserved_for_system (int): Raw Memory in Mebibytes reserved for the System not usable by Jobs. - temporary_disk_space_per_node (int): + temporary_disk (int): Amount of temporary disk space this node has, in Mebibytes. weight (int): Weight of the node in scheduling. @@ -223,9 +221,6 @@ cdef class Node: CPU Load on the Node. slurmd_port (int): Port the slurmd is listening on the node. - - Raises: - MemoryError: If malloc fails to allocate memory. """ cdef: node_info_t *info diff --git a/pyslurm/core/node.pyx b/pyslurm/core/node.pyx index 609016fe..eac1bfef 100644 --- a/pyslurm/core/node.pyx +++ b/pyslurm/core/node.pyx @@ -28,7 +28,8 @@ from pyslurm.utils import ctime from pyslurm.utils.uint import * from pyslurm.core.error import RPCError, verify_rpc from pyslurm.utils.ctime import timestamp_to_date, _raw_time -from pyslurm.db.cluster import LOCAL_CLUSTER +from pyslurm.settings import LOCAL_CLUSTER +from pyslurm import xcollections from pyslurm.utils.helpers import ( uid_to_name, gid_to_name, @@ -37,15 +38,12 @@ from pyslurm.utils.helpers import ( _getpwall_to_dict, cpubind_to_num, instance_to_dict, - collection_to_dict, - group_collection_by_cluster, - _sum_prop, nodelist_from_range_str, nodelist_to_range_str, ) -cdef class Nodes(list): +cdef class Nodes(MultiClusterMap): def __dealloc__(self): slurm_free_node_info_msg(self.info) @@ -56,38 +54,11 @@ cdef class Nodes(list): self.part_info = NULL def __init__(self, nodes=None): - if isinstance(nodes, list): - for node in nodes: - if isinstance(node, str): - self.append(Node(node)) - else: - self.append(node) - elif isinstance(nodes, str): - nodelist = nodes.split(",") - self.extend([Node(node) for node in nodelist]) - elif isinstance(nodes, dict): - self.extend([node for node in nodes.values()]) - elif nodes is not None: - raise TypeError("Invalid Type: {type(nodes)}") - - def as_dict(self, recursive=False): - """Convert the collection data to a dict. - - Args: - recursive (bool, optional): - By default, the objects will not be converted to a dict. If - this is set to `True`, then additionally all objects are - converted to dicts. - - Returns: - (dict): Collection as a dict. - """ - col = collection_to_dict(self, identifier=Node.name, - recursive=recursive) - return col.get(LOCAL_CLUSTER, {}) - - def group_by_cluster(self): - return group_collection_by_cluster(self) + super().__init__(data=nodes, + typ="Nodes", + val_type=Node, + id_attr=Node.name, + key_type=str) @staticmethod def load(preload_passwd_info=False): @@ -107,12 +78,11 @@ cdef class Nodes(list): Raises: RPCError: When getting all the Nodes from the slurmctld failed. - MemoryError: If malloc fails to allocate memory. """ cdef: dict passwd = {} dict groups = {} - Nodes nodes = Nodes.__new__(Nodes) + Nodes nodes = Nodes() int flags = slurm.SHOW_ALL Node node @@ -141,16 +111,13 @@ cdef class Nodes(list): node.passwd = passwd node.groups = groups - nodes.append(node) + cluster = node.cluster + if cluster not in nodes.data: + nodes.data[cluster] = {} + nodes.data[cluster][node.name] = node - # At this point we memcpy'd all the memory for the Nodes. Setting this - # to 0 will prevent the slurm node free function to deallocate the - # memory for the individual nodes. This should be fine, because they - # are free'd automatically in __dealloc__ since the lifetime of each - # node-pointer is tied to the lifetime of its corresponding "Node" - # instance. + # We have extracted all pointers nodes.info.record_count = 0 - return nodes def reload(self): @@ -164,19 +131,7 @@ cdef class Nodes(list): Raises: RPCError: When getting the Nodes from the slurmctld failed. """ - cdef Nodes reloaded_nodes - - if not self: - return self - - reloaded_nodes = Nodes.load().as_dict() - for idx, node in enumerate(self): - node_name = node.name - if node in reloaded_nodes: - # Put the new data in. - self[idx] = reloaded_nodes[node_name] - - return self + return xcollections.multi_reload(self) def modify(self, Node changes): """Modify all Nodes in a collection. @@ -199,50 +154,47 @@ cdef class Nodes(list): >>> # Apply the changes to all the nodes >>> nodes.modify(changes) """ - cdef: - Node n = changes - list node_names = [node.name for node in self] - - node_str = nodelist_to_range_str(node_names) + cdef Node n = changes + node_str = nodelist_to_range_str(list(self.keys())) n._alloc_umsg() cstr.fmalloc(&n.umsg.node_names, node_str) verify_rpc(slurm_update_node(n.umsg)) @property def free_memory(self): - return _sum_prop(self, Node.free_memory) + return xcollections.sum_property(self, Node.free_memory) @property def real_memory(self): - return _sum_prop(self, Node.real_memory) + return xcollections.sum_property(self, Node.real_memory) @property def allocated_memory(self): - return _sum_prop(self, Node.allocated_memory) + return xcollections.sum_property(self, Node.allocated_memory) @property def total_cpus(self): - return _sum_prop(self, Node.total_cpus) + return xcollections.sum_property(self, Node.total_cpus) @property def idle_cpus(self): - return _sum_prop(self, Node.idle_cpus) + return xcollections.sum_property(self, Node.idle_cpus) @property def allocated_cpus(self): - return _sum_prop(self, Node.allocated_cpus) + return xcollections.sum_property(self, Node.allocated_cpus) @property def effective_cpus(self): - return _sum_prop(self, Node.effective_cpus) + return xcollections.sum_property(self, Node.effective_cpus) @property def current_watts(self): - return _sum_prop(self, Node.current_watts) + return xcollections.sum_property(self, Node.current_watts) @property def avg_watts(self): - return _sum_prop(self, Node.avg_watts) + return xcollections.sum_property(self, Node.avg_watts) cdef class Node: @@ -293,8 +245,8 @@ cdef class Node: # Call descriptors __set__ directly Node.__dict__[name].__set__(self, val) - def __eq__(self, other): - return isinstance(other, Node) and self.name == other.name + def __repr__(self): + return f'{self.__class__.__name__}({self.name})' @staticmethod cdef Node from_ptr(node_info_t *in_ptr): @@ -325,7 +277,6 @@ cdef class Node: Raises: RPCError: If requesting the Node information from the slurmctld was not successful. - MemoryError: If malloc failed to allocate memory. Examples: >>> import pyslurm @@ -365,7 +316,7 @@ cdef class Node: Args: state (str, optional): An optional state the created Node should have. Allowed values - are "future" and "cloud". "future" is the default. + are `future` and `cloud`. `future` is the default. Returns: (pyslurm.Node): This function returns the current Node-instance @@ -373,7 +324,6 @@ cdef class Node: Raises: RPCError: If creating the Node was not successful. - MemoryError: If malloc failed to allocate memory. Examples: >>> import pyslurm @@ -424,7 +374,6 @@ cdef class Node: Raises: RPCError: If deleting the Node was not successful. - MemoryError: If malloc failed to allocate memory. Examples: >>> import pyslurm @@ -434,6 +383,9 @@ cdef class Node: verify_rpc(slurm_delete_node(self.umsg)) def as_dict(self): + return self.to_dict() + + def to_dict(self): """Node information formatted as a dictionary. Returns: @@ -442,7 +394,7 @@ cdef class Node: Examples: >>> import pyslurm >>> mynode = pyslurm.Node.load("mynode") - >>> mynode_dict = mynode.as_dict() + >>> mynode_dict = mynode.to_dict() """ return instance_to_dict(self) @@ -559,7 +511,7 @@ cdef class Node: return u64_parse(self.info.mem_spec_limit) @property - def temporary_disk_space(self): + def temporary_disk(self): return u32_parse(self.info.tmp_disk) @property diff --git a/pyslurm/core/partition.pxd b/pyslurm/core/partition.pxd index b10366b8..a5a638df 100644 --- a/pyslurm/core/partition.pxd +++ b/pyslurm/core/partition.pxd @@ -56,10 +56,11 @@ from pyslurm.utils cimport ctime from pyslurm.utils.ctime cimport time_t from pyslurm.utils.uint cimport * from pyslurm.core cimport slurmctld +from pyslurm.xcollections cimport MultiClusterMap -cdef class Partitions(list): - """A collection of [pyslurm.Partition][] objects. +cdef class Partitions(MultiClusterMap): + """A [`Multi Cluster`][pyslurm.xcollections.MultiClusterMap] collection of [pyslurm.Partition][] objects. Args: partitions (Union[list[str], dict[str, Partition], str], optional=None): @@ -167,7 +168,7 @@ cdef class Partition: This can also return [UNLIMITED][pyslurm.constants.UNLIMITED] min_nodes (int): Minimum number of Nodes that must be requested by Jobs - max_time_limit (int): + max_time (int): Max Time-Limit in minutes that Jobs can request This can also return [UNLIMITED][pyslurm.constants.UNLIMITED] diff --git a/pyslurm/core/partition.pyx b/pyslurm/core/partition.pyx index 56375d33..e1a1b6b1 100644 --- a/pyslurm/core/partition.pyx +++ b/pyslurm/core/partition.pyx @@ -30,7 +30,8 @@ from pyslurm.utils.uint import * from pyslurm.core.error import RPCError, verify_rpc from pyslurm.utils.ctime import timestamp_to_date, _raw_time from pyslurm.constants import UNLIMITED -from pyslurm.db.cluster import LOCAL_CLUSTER +from pyslurm.settings import LOCAL_CLUSTER +from pyslurm import xcollections from pyslurm.utils.helpers import ( uid_to_name, gid_to_name, @@ -38,9 +39,6 @@ from pyslurm.utils.helpers import ( _getpwall_to_dict, cpubind_to_num, instance_to_dict, - collection_to_dict, - group_collection_by_cluster, - _sum_prop, dehumanize, ) from pyslurm.utils.ctime import ( @@ -49,7 +47,7 @@ from pyslurm.utils.ctime import ( ) -cdef class Partitions(list): +cdef class Partitions(MultiClusterMap): def __dealloc__(self): slurm_free_partition_info_msg(self.info) @@ -58,38 +56,11 @@ cdef class Partitions(list): self.info = NULL def __init__(self, partitions=None): - if isinstance(partitions, list): - for part in partitions: - if isinstance(part, str): - self.append(Partition(part)) - else: - self.append(part) - elif isinstance(partitions, str): - partlist = partitions.split(",") - self.extend([Partition(part) for part in partlist]) - elif isinstance(partitions, dict): - self.extend([part for part in partitions.values()]) - elif partitions is not None: - raise TypeError("Invalid Type: {type(partitions)}") - - def as_dict(self, recursive=False): - """Convert the collection data to a dict. - - Args: - recursive (bool, optional): - By default, the objects will not be converted to a dict. If - this is set to `True`, then additionally all objects are - converted to dicts. - - Returns: - (dict): Collection as a dict. - """ - col = collection_to_dict(self, identifier=Partition.name, - recursive=recursive) - return col.get(LOCAL_CLUSTER, {}) - - def group_by_cluster(self): - return group_collection_by_cluster(self) + super().__init__(data=partitions, + typ="Partitions", + val_type=Partition, + id_attr=Partition.name, + key_type=str) @staticmethod def load(): @@ -103,7 +74,7 @@ cdef class Partitions(list): failed. """ cdef: - Partitions partitions = Partitions.__new__(Partitions) + Partitions partitions = Partitions() int flags = slurm.SHOW_ALL Partition partition slurmctld.Config slurm_conf @@ -126,18 +97,16 @@ cdef class Partitions(list): # is raised by replacing it with a zeroed-out partition_info_t. partitions.info.partition_array[cnt] = partitions.tmp_info + cluster = partition.cluster + if cluster not in partitions.data: + partitions.data[cluster] = {} + partition.power_save_enabled = power_save_enabled partition.slurm_conf = slurm_conf - partitions.append(partition) - - # At this point we memcpy'd all the memory for the Partitions. Setting - # this to 0 will prevent the slurm partition free function to - # deallocate the memory for the individual partitions. This should be - # fine, because they are free'd automatically in __dealloc__ since the - # lifetime of each partition-pointer is tied to the lifetime of its - # corresponding "Partition" instance. - partitions.info.record_count = 0 + partitions.data[cluster][partition.name] = partition + # We have extracted all pointers + partitions.info.record_count = 0 return partitions def reload(self): @@ -154,19 +123,7 @@ cdef class Partitions(list): Raises: RPCError: When getting the Partitions from the slurmctld failed. """ - cdef dict reloaded_parts - - if not self: - return self - - reloaded_parts = Partitions.load().as_dict() - for idx, part in enumerate(self): - part_name = part.name - if part_name in reloaded_parts: - # Put the new data in. - self[idx] = reloaded_parts[part_name] - - return self + return xcollections.multi_reload(self) def modify(self, changes): """Modify all Partitions in a Collection. @@ -189,16 +146,16 @@ cdef class Partitions(list): >>> # Apply the changes to all the partitions >>> parts.modify(changes) """ - for part in self: + for part in self.values(): part.modify(changes) @property def total_cpus(self): - return _sum_prop(self, Partition.total_cpus) + return xcollections.sum_property(self, Partition.total_cpus) @property def total_nodes(self): - return _sum_prop(self, Partition.total_nodes) + return xcollections.sum_property(self, Partition.total_nodes) cdef class Partition: @@ -228,6 +185,9 @@ cdef class Partition: def __dealloc__(self): self._dealloc_impl() + def __repr__(self): + return f'{self.__class__.__name__}({self.name})' + @staticmethod cdef Partition from_ptr(partition_info_t *in_ptr): cdef Partition wrap = Partition.__new__(Partition) @@ -243,6 +203,9 @@ cdef class Partition: return self.name def as_dict(self): + return self.to_dict() + + def to_dict(self): """Partition information formatted as a dictionary. Returns: @@ -251,7 +214,7 @@ cdef class Partition: Examples: >>> import pyslurm >>> mypart = pyslurm.Partition.load("mypart") - >>> mypart_dict = mypart.as_dict() + >>> mypart_dict = mypart.to_dict() """ return instance_to_dict(self) @@ -274,11 +237,11 @@ cdef class Partition: >>> import pyslurm >>> part = pyslurm.Partition.load("normal") """ - partitions = Partitions.load().as_dict() - if name not in partitions: + part = Partitions.load().get(name) + if not part: raise RPCError(msg=f"Partition '{name}' doesn't exist") - return partitions[name] + return part def create(self): """Create a Partition. @@ -341,7 +304,6 @@ cdef class Partition: """ cdef delete_part_msg_t del_part_msg memset(&del_part_msg, 0, sizeof(del_part_msg)) - del_part_msg.name = cstr.from_unicode(self._error_or_name()) verify_rpc(slurm_delete_partition(&del_part_msg)) @@ -357,6 +319,10 @@ cdef class Partition: def name(self): return cstr.to_unicode(self.ptr.name) + @property + def _id(self): + return self.name + @name.setter def name(self, val): cstr.fmalloc(&self.ptr.name, val) @@ -546,11 +512,11 @@ cdef class Partition: self.ptr.min_nodes = u32(val, zero_is_noval=False) @property - def max_time_limit(self): + def max_time(self): return _raw_time(self.ptr.max_time, on_inf=UNLIMITED) - @max_time_limit.setter - def max_time_limit(self, val): + @max_time.setter + def max_time(self, val): self.ptr.max_time = timestr_to_mins(val) @property diff --git a/pyslurm/db/__init__.py b/pyslurm/db/__init__.py index 0e78a734..acd36a40 100644 --- a/pyslurm/db/__init__.py +++ b/pyslurm/db/__init__.py @@ -42,4 +42,3 @@ Association, AssociationFilter, ) -from . import cluster diff --git a/pyslurm/db/assoc.pxd b/pyslurm/db/assoc.pxd index 12a0cde1..384dbb0a 100644 --- a/pyslurm/db/assoc.pxd +++ b/pyslurm/db/assoc.pxd @@ -49,12 +49,13 @@ from pyslurm.db.connection cimport Connection from pyslurm.utils cimport cstr from pyslurm.utils.uint cimport * from pyslurm.db.qos cimport QualitiesOfService, _set_qos_list +from pyslurm.xcollections cimport MultiClusterMap cdef _parse_assoc_ptr(Association ass) cdef _create_assoc_ptr(Association ass, conn=*) -cdef class Associations(list): +cdef class Associations(MultiClusterMap): pass @@ -69,8 +70,8 @@ cdef class AssociationFilter: cdef class Association: cdef: slurmdb_assoc_rec_t *ptr - dict qos_data - dict tres_data + QualitiesOfService qos_data + TrackableResources tres_data cdef public: group_tres diff --git a/pyslurm/db/assoc.pyx b/pyslurm/db/assoc.pyx index d1ac4789..4e535a46 100644 --- a/pyslurm/db/assoc.pyx +++ b/pyslurm/db/assoc.pyx @@ -25,47 +25,22 @@ from pyslurm.core.error import RPCError from pyslurm.utils.helpers import ( instance_to_dict, - collection_to_dict, - group_collection_by_cluster, user_to_uid, ) from pyslurm.utils.uint import * from pyslurm.db.connection import _open_conn_or_error -from pyslurm.db.cluster import LOCAL_CLUSTER +from pyslurm.settings import LOCAL_CLUSTER +from pyslurm import xcollections -cdef class Associations(list): +cdef class Associations(MultiClusterMap): - def __init__(self): - pass - - def as_dict(self, recursive=False, group_by_cluster=False): - """Convert the collection data to a dict. - - Args: - recursive (bool, optional): - By default, the objects will not be converted to a dict. If - this is set to `True`, then additionally all objects are - converted to dicts. - group_by_cluster (bool, optional): - By default, only the Jobs from your local Cluster are - returned. If this is set to `True`, then all the Jobs in the - collection will be grouped by the Cluster - with the name of - the cluster as the key and the value being the collection as - another dict. - - Returns: - (dict): Collection as a dict. - """ - col = collection_to_dict(self, identifier=Association.id, - recursive=recursive) - if not group_by_cluster: - return col.get(LOCAL_CLUSTER, {}) - - return col - - def group_by_cluster(self): - return group_collection_by_cluster(self) + def __init__(self, assocs=None): + super().__init__(data=assocs, + typ="Associations", + val_type=Association, + id_attr=Association.id, + key_type=int) @staticmethod def load(AssociationFilter db_filter=None, Connection db_connection=None): @@ -76,8 +51,8 @@ cdef class Associations(list): SlurmList assoc_data SlurmListItem assoc_ptr Connection conn - dict qos_data - dict tres_data + QualitiesOfService qos_data + TrackableResources tres_data # Prepare SQL Filter if not db_filter: @@ -96,10 +71,10 @@ cdef class Associations(list): # Fetch other necessary dependencies needed for translating some # attributes (i.e QoS IDs to its name) - qos_data = QualitiesOfService.load(db_connection=conn).as_dict( - name_is_key=False) - tres_data = TrackableResources.load(db_connection=conn).as_dict( - name_is_key=False) + qos_data = QualitiesOfService.load(db_connection=conn, + name_is_key=False) + tres_data = TrackableResources.load(db_connection=conn, + name_is_key=False) # Setup Association objects for assoc_ptr in SlurmList.iter_and_pop(assoc_data): @@ -107,7 +82,11 @@ cdef class Associations(list): assoc.qos_data = qos_data assoc.tres_data = tres_data _parse_assoc_ptr(assoc) - out.append(assoc) + + cluster = assoc.cluster + if cluster not in out.data: + out.data[cluster] = {} + out.data[cluster][assoc.id] = assoc return out @@ -226,13 +205,16 @@ cdef class Association: slurmdb_init_assoc_rec(self.ptr, 0) + def __repr__(self): + return f'{self.__class__.__name__}({self.id})' + @staticmethod cdef Association from_ptr(slurmdb_assoc_rec_t *in_ptr): cdef Association wrap = Association.__new__(Association) wrap.ptr = in_ptr return wrap - def as_dict(self): + def to_dict(self): """Database Association information formatted as a dictionary. Returns: @@ -408,8 +390,8 @@ cdef class Association: cdef _parse_assoc_ptr(Association ass): cdef: - dict tres = ass.tres_data - dict qos = ass.qos_data + TrackableResources tres = ass.tres_data + QualitiesOfService qos = ass.qos_data ass.group_tres = TrackableResourceLimits.from_ids( ass.ptr.grp_tres, tres) diff --git a/pyslurm/db/cluster.pxd b/pyslurm/db/cluster.pxd deleted file mode 100644 index 30acdbde..00000000 --- a/pyslurm/db/cluster.pxd +++ /dev/null @@ -1,27 +0,0 @@ -######################################################################### -# cluster.pxd - pyslurm slurmdbd cluster api -######################################################################### -# Copyright (C) 2023 Toni Harzendorf -# -# This file is part of PySlurm -# -# PySlurm is free software; you can redistribute it and/or modify -# it under the terms of the GNU General Public License as published by -# the Free Software Foundation; either version 2 of the License, or -# (at your option) any later version. - -# PySlurm is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU General Public License for more details. -# -# You should have received a copy of the GNU General Public License along -# with PySlurm; if not, write to the Free Software Foundation, Inc., -# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA. -# -# cython: c_string_type=unicode, c_string_encoding=default -# cython: language_level=3 - - -from pyslurm cimport slurm -from pyslurm.utils cimport cstr diff --git a/pyslurm/db/job.pxd b/pyslurm/db/job.pxd index fc395943..bf21c003 100644 --- a/pyslurm/db/job.pxd +++ b/pyslurm/db/job.pxd @@ -53,6 +53,7 @@ from pyslurm.db.connection cimport Connection from pyslurm.utils cimport cstr from pyslurm.db.qos cimport QualitiesOfService from pyslurm.db.tres cimport TrackableResources, TrackableResource +from pyslurm.xcollections cimport MultiClusterMap cdef class JobFilter: @@ -150,8 +151,8 @@ cdef class JobFilter: with_env -cdef class Jobs(list): - """A collection of [pyslurm.db.Job][] objects.""" +cdef class Jobs(MultiClusterMap): + """A [`Multi Cluster`][pyslurm.xcollections.MultiClusterMap] collection of [pyslurm.db.Job][] objects.""" pass @@ -161,6 +162,8 @@ cdef class Job: Args: job_id (int, optional=0): An Integer representing a Job-ID. + cluster (str, optional=None): + Name of the Cluster for this Job. Other Parameters: admin_comment (str): @@ -283,7 +286,7 @@ cdef class Job: """ cdef: slurmdb_job_rec_t *ptr - dict qos_data + QualitiesOfService qos_data cdef public: JobSteps steps diff --git a/pyslurm/db/job.pyx b/pyslurm/db/job.pyx index 636e1137..905f206a 100644 --- a/pyslurm/db/job.pyx +++ b/pyslurm/db/job.pyx @@ -27,7 +27,8 @@ from pyslurm.core.error import RPCError, PyslurmError from pyslurm.core import slurmctld from typing import Any from pyslurm.utils.uint import * -from pyslurm.db.cluster import LOCAL_CLUSTER +from pyslurm.settings import LOCAL_CLUSTER +from pyslurm import xcollections from pyslurm.utils.ctime import ( date_to_timestamp, timestr_to_mins, @@ -40,8 +41,6 @@ from pyslurm.utils.helpers import ( uid_to_name, nodelist_to_range_str, instance_to_dict, - collection_to_dict, - group_collection_by_cluster, _get_exit_code, ) from pyslurm.db.connection import _open_conn_or_error @@ -80,7 +79,7 @@ cdef class JobFilter: qos_data = QualitiesOfService.load() for user_input in self.qos: found = False - for qos in qos_data: + for qos in qos_data.values(): if (qos.id == user_input or qos.name == user_input or qos == user_input): @@ -189,54 +188,14 @@ cdef class JobFilter: JobSearchFilter = JobFilter -cdef class Jobs(list): +cdef class Jobs(MultiClusterMap): def __init__(self, jobs=None): - if isinstance(jobs, list): - for job in jobs: - if isinstance(job, int): - self.append(Job(job)) - else: - self.append(job) - elif isinstance(jobs, str): - joblist = jobs.split(",") - self.extend([Job(job) for job in joblist]) - elif isinstance(jobs, dict): - self.extend([job for job in jobs.values()]) - elif jobs is not None: - raise TypeError("Invalid Type: {type(jobs)}") - - def as_dict(self, recursive=False, group_by_cluster=False): - """Convert the collection data to a dict. - - Args: - recursive (bool, optional): - By default, the objects will not be converted to a dict. If - this is set to `True`, then additionally all objects are - converted to dicts. - group_by_cluster (bool, optional): - By default, only the Jobs from your local Cluster are - returned. If this is set to `True`, then all the Jobs in the - collection will be grouped by the Cluster - with the name of - the cluster as the key and the value being the collection as - another dict. - - Returns: - (dict): Collection as a dict. - """ - col = collection_to_dict(self, identifier=Job.id, recursive=recursive) - if not group_by_cluster: - return col.get(LOCAL_CLUSTER, {}) - - return col - - def group_by_cluster(self): - """Group Jobs by cluster name - - Returns: - (dict[str, Jobs]): Jobs grouped by cluster. - """ - return group_collection_by_cluster(self) + super().__init__(data=jobs, + typ="Jobs", + val_type=Job, + id_attr=Job.id, + key_type=int) @staticmethod def load(JobFilter db_filter=None, Connection db_connection=None): @@ -280,7 +239,7 @@ cdef class Jobs(list): SlurmList job_data SlurmListItem job_ptr Connection conn - dict qos_data + QualitiesOfService qos_data # Prepare SQL Filter if not db_filter: @@ -297,15 +256,14 @@ cdef class Jobs(list): # Fetch other necessary dependencies needed for translating some # attributes (i.e QoS IDs to its name) - qos_data = QualitiesOfService.load(db_connection=conn).as_dict( - name_is_key=False) + qos_data = QualitiesOfService.load(db_connection=conn, + name_is_key=False) # TODO: also get trackable resources with slurmdb_tres_get and store # it in each job instance. tres_alloc_str and tres_req_str only # contain the numeric tres ids, but it probably makes more sense to # convert them to its type name for the user in advance. - # TODO: For multi-cluster support, remove duplicate federation jobs # TODO: How to handle the possibility of duplicate job ids that could # appear if IDs on a cluster are resetted? for job_ptr in SlurmList.iter_and_pop(job_data): @@ -313,7 +271,11 @@ cdef class Jobs(list): job.qos_data = qos_data job._create_steps() JobStatistics._sum_step_stats_for_job(job, job.steps) - out.append(job) + + cluster = job.cluster + if cluster not in out.data: + out.data[cluster] = {} + out[cluster][job.id] = job return out @@ -363,7 +325,7 @@ cdef class Jobs(list): >>> changes = pyslurm.db.Job(comment="A comment for the job") >>> modified_jobs = pyslurm.db.Jobs.modify(db_filter, changes) >>> print(modified_jobs) - >>> [9999] + [9999] In the above example, the changes will be automatically committed if successful. @@ -380,7 +342,7 @@ cdef class Jobs(list): >>> >>> # Now you can first examine which Jobs have been modified >>> print(modified_jobs) - >>> [9999] + [9999] >>> # And then you can actually commit (or even rollback) the >>> # changes >>> db_conn.commit() @@ -420,7 +382,7 @@ cdef class Jobs(list): # # " submitted at " # - # We are just interest in the Job-ID, so extract it + # We are just interested in the Job-ID, so extract it job_id = response_str.split(" ")[0] if job_id and job_id.isdigit(): out.append(int(job_id)) @@ -444,10 +406,14 @@ cdef class Job: def __cinit__(self): self.ptr = NULL - def __init__(self, job_id=0, cluster=LOCAL_CLUSTER, **kwargs): + def __init__(self, job_id=0, cluster=None, **kwargs): self._alloc_impl() self.ptr.jobid = int(job_id) - cstr.fmalloc(&self.ptr.cluster, cluster) + cstr.fmalloc(&self.ptr.cluster, + LOCAL_CLUSTER if not cluster else cluster) + self.qos_data = QualitiesOfService() + self.steps = JobSteps() + self.stats = JobStatistics() for k, v in kwargs.items(): setattr(self, k, v) @@ -471,12 +437,20 @@ cdef class Job: return wrap @staticmethod - def load(job_id, cluster=LOCAL_CLUSTER, with_script=False, with_env=False): + def load(job_id, cluster=None, with_script=False, with_env=False): """Load the information for a specific Job from the Database. Args: job_id (int): ID of the Job to be loaded. + cluster (str): + Name of the Cluster to search in. + with_script (bool): + Whether the Job-Script should also be loaded. Mutually + exclusive with `with_env`. + with_env (bool): + Whether the Job Environment should also be loaded. Mutually + exclusive with `with_script`. Returns: (pyslurm.db.Job): Returns a new Database Job instance @@ -489,24 +463,24 @@ cdef class Job: >>> import pyslurm >>> db_job = pyslurm.db.Job.load(10000) - In the above example, attribute like "script" and "environment" + In the above example, attributes like `script` and `environment` are not populated. You must explicitly request one of them to be loaded: >>> import pyslurm >>> db_job = pyslurm.db.Job.load(10000, with_script=True) >>> print(db_job.script) - """ + cluster = LOCAL_CLUSTER if not cluster else cluster jfilter = JobFilter(ids=[int(job_id)], clusters=[cluster], with_script=with_script, with_env=with_env) - jobs = Jobs.load(jfilter) - if not jobs: + job = Jobs.load(jfilter).get((cluster, int(job_id))) + if not job: raise RPCError(msg=f"Job {job_id} does not exist on " f"Cluster {cluster}") # TODO: There might be multiple entries when job ids were reset. - return jobs[0] + return job def _create_steps(self): cdef: @@ -520,7 +494,10 @@ cdef class Job: self.steps[step.id] = step def as_dict(self): - """Database Job information formatted as a dictionary. + return self.to_dict() + + def to_dict(self): + """Convert Database Job information to a dictionary. Returns: (dict): Database Job information as dict @@ -528,20 +505,23 @@ cdef class Job: Examples: >>> import pyslurm >>> myjob = pyslurm.db.Job.load(10000) - >>> myjob_dict = myjob.as_dict() + >>> myjob_dict = myjob.to_dict() """ cdef dict out = instance_to_dict(self) if self.stats: - out["stats"] = self.stats.as_dict() + out["stats"] = self.stats.to_dict() steps = out.pop("steps", {}) out["steps"] = {} for step_id, step in steps.items(): - out["steps"][step_id] = step.as_dict() + out["steps"][step_id] = step.to_dict() return out + def __repr__(self): + return f'{self.__class__.__name__}({self.id})' + def modify(self, changes, db_connection=None): """Modify a Slurm database Job. diff --git a/pyslurm/db/qos.pxd b/pyslurm/db/qos.pxd index 9cb3df86..ea0fde2d 100644 --- a/pyslurm/db/qos.pxd +++ b/pyslurm/db/qos.pxd @@ -44,7 +44,7 @@ from pyslurm.utils cimport cstr cdef _set_qos_list(List *in_list, vals, QualitiesOfService data) -cdef class QualitiesOfService(list): +cdef class QualitiesOfService(dict): pass diff --git a/pyslurm/db/qos.pyx b/pyslurm/db/qos.pyx index a01ef9b0..299c0ed9 100644 --- a/pyslurm/db/qos.pyx +++ b/pyslurm/db/qos.pyx @@ -23,41 +23,26 @@ # cython: language_level=3 from pyslurm.core.error import RPCError -from pyslurm.utils.helpers import instance_to_dict, collection_to_dict_global +from pyslurm.utils.helpers import instance_to_dict from pyslurm.db.connection import _open_conn_or_error -cdef class QualitiesOfService(list): +cdef class QualitiesOfService(dict): def __init__(self): pass - def as_dict(self, recursive=False, name_is_key=True): - """Convert the collection data to a dict. + @staticmethod + def load(QualityOfServiceFilter db_filter=None, + Connection db_connection=None, name_is_key=True): + """Load QoS data from the Database Args: - recursive (bool, optional): - By default, the objects will not be converted to a dict. If - this is set to `True`, then additionally all objects are - converted to dicts. name_is_key (bool, optional): By default, the keys in this dict are the names of each QoS. If this is set to `False`, then the unique ID of the QoS will be used as dict keys. - - Returns: - (dict): Collection as a dict. """ - identifier = QualityOfService.name - if not name_is_key: - identifier = QualityOfService.id - - return collection_to_dict_global(self, identifier=identifier, - recursive=recursive) - - @staticmethod - def load(QualityOfServiceFilter db_filter=None, - Connection db_connection=None): cdef: QualitiesOfService out = QualitiesOfService() QualityOfService qos @@ -83,7 +68,8 @@ cdef class QualitiesOfService(list): # Setup QOS objects for qos_ptr in SlurmList.iter_and_pop(qos_data): qos = QualityOfService.from_ptr(qos_ptr.data) - out.append(qos) + _id = qos.name if name_is_key else qos.id + out[_id] = qos return out @@ -170,7 +156,10 @@ cdef class QualityOfService: wrap.ptr = in_ptr return wrap - def as_dict(self): + def __repr__(self): + return f'{self.__class__.__name__}({self.name})' + + def to_dict(self): """Database QualityOfService information formatted as a dictionary. Returns: @@ -195,11 +184,11 @@ cdef class QualityOfService: sucessful. """ qfilter = QualityOfServiceFilter(names=[name]) - qos_data = QualitiesOfService.load(qfilter) - if not qos_data: + qos = QualitiesOfService.load(qfilter).get(name) + if not qos: raise RPCError(msg=f"QualityOfService {name} does not exist") - return qos_data[0] + return qos @property def name(self): @@ -227,7 +216,7 @@ def _qos_names_to_ids(qos_list, QualitiesOfService data): def _validate_qos_single(qid, QualitiesOfService data): - for item in data: + for item in data.values(): if qid == item.id or qid == item.name: return item.id diff --git a/pyslurm/db/stats.pyx b/pyslurm/db/stats.pyx index 3ae0c8b5..7bbb2a8a 100644 --- a/pyslurm/db/stats.pyx +++ b/pyslurm/db/stats.pyx @@ -47,7 +47,7 @@ cdef class JobStatistics: self.min_cpu_time_node = None self.min_cpu_time_task = None - def as_dict(self): + def to_dict(self): return instance_to_dict(self) @staticmethod diff --git a/pyslurm/db/step.pxd b/pyslurm/db/step.pxd index aef7120b..ab0ff70c 100644 --- a/pyslurm/db/step.pxd +++ b/pyslurm/db/step.pxd @@ -44,7 +44,7 @@ from pyslurm.db.tres cimport TrackableResources, TrackableResource cdef class JobSteps(dict): - """A collection of [pyslurm.db.JobStep][] objects""" + """A [dict][] of [pyslurm.db.JobStep][] objects""" pass diff --git a/pyslurm/db/step.pyx b/pyslurm/db/step.pyx index fa4ab8bb..e39af066 100644 --- a/pyslurm/db/step.pyx +++ b/pyslurm/db/step.pyx @@ -57,9 +57,14 @@ cdef class JobStep: wrap.stats = JobStatistics.from_step(wrap) return wrap - def as_dict(self): + def to_dict(self): + """Convert Database JobStep information to a dictionary. + + Returns: + (dict): Database JobStep information as dict + """ cdef dict out = instance_to_dict(self) - out["stats"] = self.stats.as_dict() + out["stats"] = self.stats.to_dict() return out @property diff --git a/pyslurm/db/tres.pxd b/pyslurm/db/tres.pxd index 41ed1b4d..23b44ad2 100644 --- a/pyslurm/db/tres.pxd +++ b/pyslurm/db/tres.pxd @@ -42,7 +42,7 @@ from pyslurm.db.connection cimport Connection cdef find_tres_count(char *tres_str, typ, on_noval=*, on_inf=*) cdef find_tres_limit(char *tres_str, typ) cdef merge_tres_str(char **tres_str, typ, val) -cdef _tres_ids_to_names(char *tres_str, dict tres_data) +cdef _tres_ids_to_names(char *tres_str, TrackableResources tres_data) cdef _set_tres_limits(char **dest, TrackableResourceLimits src, TrackableResources tres_data) @@ -62,14 +62,14 @@ cdef class TrackableResourceLimits: license @staticmethod - cdef from_ids(char *tres_id_str, dict tres_data) + cdef from_ids(char *tres_id_str, TrackableResources tres_data) cdef class TrackableResourceFilter: cdef slurmdb_tres_cond_t *ptr -cdef class TrackableResources(list): +cdef class TrackableResources(dict): cdef public raw_str @staticmethod diff --git a/pyslurm/db/tres.pyx b/pyslurm/db/tres.pyx index df93dda0..78195654 100644 --- a/pyslurm/db/tres.pyx +++ b/pyslurm/db/tres.pyx @@ -25,7 +25,7 @@ from pyslurm.utils.uint import * from pyslurm.constants import UNLIMITED from pyslurm.core.error import RPCError -from pyslurm.utils.helpers import instance_to_dict, collection_to_dict_global +from pyslurm.utils.helpers import instance_to_dict from pyslurm.utils import cstr from pyslurm.db.connection import _open_conn_or_error import json @@ -56,7 +56,7 @@ cdef class TrackableResourceLimits: setattr(self, k, v) @staticmethod - cdef from_ids(char *tres_id_str, dict tres_data): + cdef from_ids(char *tres_id_str, TrackableResources tres_data): tres_list = _tres_ids_to_names(tres_id_str, tres_data) if not tres_list: return None @@ -76,7 +76,7 @@ cdef class TrackableResourceLimits: return out def _validate(self, TrackableResources tres_data): - id_dict = _tres_names_to_ids(self.as_dict(flatten_limits=True), + id_dict = _tres_names_to_ids(self.to_dict(flatten_limits=True), tres_data) return id_dict @@ -91,7 +91,7 @@ cdef class TrackableResourceLimits: return out - def as_dict(self, flatten_limits=False): + def to_dict(self, flatten_limits=False): cdef dict inst_dict = instance_to_dict(self) if flatten_limits: @@ -134,36 +134,21 @@ cdef class TrackableResourceFilter: self._alloc() -cdef class TrackableResources(list): +cdef class TrackableResources(dict): def __init__(self): pass - def as_dict(self, recursive=False, name_is_key=True): - """Convert the collection data to a dict. + @staticmethod + def load(Connection db_connection=None, name_is_key=True): + """Load Trackable Resources from the Database. Args: - recursive (bool, optional): - By default, the objects will not be converted to a dict. If - this is set to `True`, then additionally all objects are - converted to dicts. name_is_key (bool, optional): By default, the keys in this dict are the names of each TRES. If this is set to `False`, then the unique ID of the TRES will be used as dict keys. - - Returns: - (dict): Collection as a dict. """ - identifier = TrackableResource.type_and_name - if not name_is_key: - identifier = TrackableResource.id - - return collection_to_dict_global(self, identifier=identifier, - recursive=recursive) - - @staticmethod - def load(Connection db_connection=None): cdef: TrackableResources out = TrackableResources() TrackableResource tres @@ -188,7 +173,8 @@ cdef class TrackableResources(list): for tres_ptr in SlurmList.iter_and_pop(tres_data): tres = TrackableResource.from_ptr( tres_ptr.data) - out.append(tres) + _id = tres.type_and_name if name_is_key else tres.id + out[_id] = tres return out @@ -246,7 +232,7 @@ cdef class TrackableResource: wrap.ptr = in_ptr return wrap - def as_dict(self): + def to_dict(self): return instance_to_dict(self) @property @@ -307,7 +293,7 @@ cdef merge_tres_str(char **tres_str, typ, val): cstr.from_dict(tres_str, current) -cdef _tres_ids_to_names(char *tres_str, dict tres_data): +cdef _tres_ids_to_names(char *tres_str, TrackableResources tres_data): if not tres_str: return None @@ -342,7 +328,7 @@ def _tres_names_to_ids(dict tres_dict, TrackableResources tres_data): def _validate_tres_single(tid, TrackableResources tres_data): - for tres in tres_data: + for tres in tres_data.values(): if tid == tres.id or tid == tres.type_and_name: return tres.id diff --git a/pyslurm/db/cluster.pyx b/pyslurm/settings.pyx similarity index 92% rename from pyslurm/db/cluster.pyx rename to pyslurm/settings.pyx index 436183a8..5085a9f5 100644 --- a/pyslurm/db/cluster.pyx +++ b/pyslurm/settings.pyx @@ -1,5 +1,5 @@ ######################################################################### -# cluster.pyx - pyslurm slurmdbd cluster api +# settings.pyx - pyslurm global settings ######################################################################### # Copyright (C) 2023 Toni Harzendorf # @@ -23,6 +23,8 @@ # cython: language_level=3 from pyslurm.core import slurmctld +from pyslurm cimport slurm +from pyslurm.utils cimport cstr LOCAL_CLUSTER = cstr.to_unicode(slurm.slurm_conf.cluster_name) diff --git a/pyslurm/utils/helpers.pyx b/pyslurm/utils/helpers.pyx index fb1d2201..9fcd5896 100644 --- a/pyslurm/utils/helpers.pyx +++ b/pyslurm/utils/helpers.pyx @@ -341,60 +341,6 @@ def instance_to_dict(inst): return out -def collection_to_dict(collection, identifier, recursive=False, group_id=None): - cdef dict out = {} - - for item in collection: - cluster = item.cluster - if cluster not in out: - out[cluster] = {} - - _id = identifier.__get__(item) - data = item if not recursive else item.as_dict() - - if group_id: - grp_id = group_id.__get__(item) - if grp_id not in out[cluster]: - out[cluster][grp_id] = {} - out[cluster][grp_id].update({_id: data}) - else: - out[cluster][_id] = data - - return out - - -def collection_to_dict_global(collection, identifier, recursive=False): - cdef dict out = {} - for item in collection: - _id = identifier.__get__(item) - out[_id] = item if not recursive else item.as_dict() - return out - - -def group_collection_by_cluster(collection): - cdef dict out = {} - collection_type = type(collection) - - for item in collection: - cluster = item.cluster - if cluster not in out: - out[cluster] = collection_type() - - out[cluster].append(item) - - return out - - -def _sum_prop(obj, name, startval=0): - val = startval - for n in obj.values(): - v = name.__get__(n) - if v is not None: - val += v - - return val - - def _get_exit_code(exit_code): exit_state=sig = 0 if exit_code != slurm.NO_VAL: diff --git a/pyslurm/xcollections.pxd b/pyslurm/xcollections.pxd new file mode 100644 index 00000000..24007da7 --- /dev/null +++ b/pyslurm/xcollections.pxd @@ -0,0 +1,93 @@ +######################################################################### +# collections.pxd - pyslurm custom collections +######################################################################### +# Copyright (C) 2023 Toni Harzendorf +# +# This file is part of PySlurm +# +# PySlurm is free software; you can redistribute it and/or modify +# it under the terms of the GNU General Public License as published by +# the Free Software Foundation; either version 2 of the License, or +# (at your option) any later version. + +# PySlurm is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU General Public License for more details. +# +# You should have received a copy of the GNU General Public License along +# with PySlurm; if not, write to the Free Software Foundation, Inc., +# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA. +# +# cython: c_string_type=unicode, c_string_encoding=default +# cython: language_level=3 + + +cdef class MultiClusterMap: + """Mapping of Multi-Cluster Data for a Collection. + + !!! note "TL;DR" + + If you have no need to write Multi-Cluster capable code and just work + on a single Cluster, Collections inheriting from this Class behave + just like a normal `dict`. + + This class enables collections to hold data from multiple Clusters if + applicable. + For quite a few Entities in Slurm it is possible to gather data from + multiple Clusters. For example, with `squeue`, you can easily list Jobs + running on different Clusters - provided your Cluster is joined in a + Federation or simply part of a multi Cluster Setup. + + Collections like `pyslurm.Jobs` inherit from this Class to enable holding + such data from multiple Clusters. + Internally, the data is structured in a `dict` like this (with + `pyslurm.Jobs` as an example): + + ```python + data = { + "LOCAL_CLUSTER": + 1: pyslurm.Job, + 2: pyslurm.Job, + ... + "OTHER_REMOTE_CLUSTER": + 100: pyslurm.Job, + 101, pyslurm.Job + ... + ... + } + ``` + + When a collection inherits from this class, its functionality will + basically simulate a standard `dict` - with a few extensions to enable + multi-cluster code. + By default, even if your Collections contains Data from multiple Clusters, + any operation will be targeted on the local Cluster data, if available. + + For example, with the data from above: + + ```python + job = data[1] + ``` + + `job` would then hold the instance for Job 1 from the `LOCAL_CLUSTER` + data. + Alternatively, data can also be accessed like this: + + ```python + job = data["OTHER_REMOTE_CLUSTER"][100] + ``` + + Here, you are directly specifying which Cluster data you want to access. + + Similarly, every method (where applicable) from a standard dict is + extended with multi-cluster functionality (check out the examples on the + methods) + """ + cdef public dict data + + cdef: + _typ + _key_type + _val_type + _id_attr diff --git a/pyslurm/xcollections.pyx b/pyslurm/xcollections.pyx new file mode 100644 index 00000000..8be67d29 --- /dev/null +++ b/pyslurm/xcollections.pyx @@ -0,0 +1,581 @@ +######################################################################### +# collections.pyx - pyslurm custom collections +######################################################################### +# Copyright (C) 2023 Toni Harzendorf +# +# This file is part of PySlurm +# +# PySlurm is free software; you can redistribute it and/or modify +# it under the terms of the GNU General Public License as published by +# the Free Software Foundation; either version 2 of the License, or +# (at your option) any later version. + +# PySlurm is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU General Public License for more details. +# +# You should have received a copy of the GNU General Public License along +# with PySlurm; if not, write to the Free Software Foundation, Inc., +# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA. +# +# cython: c_string_type=unicode, c_string_encoding=default +# cython: language_level=3 +"""Custom Collection utilities""" + +from pyslurm.settings import LOCAL_CLUSTER +import json +from typing import Union, Any + + +class BaseView: + """Base View for all other Views""" + def __init__(self, mcm): + self._mcm = mcm + self._data = mcm.data + + def __len__(self): + return len(self._mcm) + + def __repr__(self): + data = ", ".join(map(repr, self)) + return f'{self.__class__.__name__}([{data}])' + + +class ValuesView(BaseView): + """A simple Value View + + When iterating over an instance of this View, this will yield all values + from all clusters. + """ + def __contains__(self, val): + try: + item = self._mcm.get( + key=self._mcm._item_id(val), + cluster=val.cluster + ) + return item is val or item == val + except AttributeError: + pass + + return False + + def __iter__(self): + for cluster in self._mcm.data.values(): + for item in cluster.values(): + yield item + + +class ClustersView(BaseView): + """A simple Cluster-Keys View + + When iterating over an instance of this View, it will yield all the + Cluster names of the collection. + """ + def __contains__(self, item): + return item in self._data + + def __len__(self): + return len(self._data) + + def __iter__(self): + yield from self._data + + +class MCKeysView(BaseView): + """A Multi-Cluster Keys View + + Unlike KeysView, when iterating over an MCKeysView instance, this will + yield a 2-tuple in the form `(cluster, key)`. + + Similarly, when checking whether this View contains a Key with the `in` + operator, a 2-tuple must be used in the form described above. + """ + def __contains__(self, item): + cluster, key, = item + return key in self._data[cluster] + + def __iter__(self): + for cluster, keys in self._data.items(): + for key in keys: + yield (cluster, key) + + +class KeysView(BaseView): + """A simple Keys View of a collection + + When iterating, this yields all the keys found from each Cluster in the + collection. Note that unlike the KeysView from a `dict`, the keys here + aren't unique and may appear multiple times. + + If you indeed have multiple Clusters in a collection and need to tell the + keys apart, use the `with_cluster()` function. + """ + def __contains__(self, item): + return item in self._mcm + + def __iter__(self): + for cluster, keys in self._data.items(): + yield from keys + + def with_cluster(self): + """Return a Multi-Cluster Keys View. + + Returns: + (MCKeysView): Multi-Cluster Keys View. + """ + return MCKeysView(self._mcm) + + +class ItemsView(BaseView): + """A simple Items View of a collection. + + Returns a 2-tuple in the form of `(key, value)` when iterating. + + Similarly, when checking whether this View contains an Item with the `in` + operator, a 2-tuple must be used. + """ + def __contains__(self, item): + key, val = item + + try: + out = self._mcm.data[item.cluster][key] + except (KeyError, AttributeError): + return False + else: + return out is val or out == val + + def __iter__(self): + for cluster, data in self._mcm.data.items(): + for key in data: + yield (key, data[key]) + + def with_cluster(self): + """Return a Multi-Cluster Items View. + + Returns: + (MCItemsView): Multi-Cluster Items View. + """ + return MCItemsView(self._mcm) + + +class MCItemsView(BaseView): + """A Multi-Cluster Items View. + + This differs from ItemsView in that it returns a 3-tuple in the form of + `(cluster, key, value)` when iterating. + + Similarly, when checking whether this View contains an Item with the `in` + operator, a 3-tuple must be used. + """ + def __contains__(self, item): + cluster, key, val = item + + try: + out = self._mcm.data[cluster][key] + except KeyError: + return False + else: + return out is val or out == val + + def __iter__(self): + for cluster, data in self._mcm.data.items(): + for key in data: + yield (cluster, key, data[key]) + + +cdef class MultiClusterMap: + + def __init__(self, data, typ=None, val_type=None, + key_type=None, id_attr=None, init_data=True): + self.data = {} if init_data else data + self._typ = typ + self._key_type = key_type + self._val_type = val_type + self._id_attr = id_attr + if init_data: + self._init_data(data) + + def _init_data(self, data): + if isinstance(data, list): + for item in data: + if isinstance(item, self._key_type): + item = self._val_type(item) + if LOCAL_CLUSTER not in self.data: + self.data[LOCAL_CLUSTER] = {} + + self.data[LOCAL_CLUSTER].update({self._item_id(item): item}) + elif isinstance(data, str): + itemlist = data.split(",") + items = {self._key_type(item):self._val_type(item) + for item in itemlist} + self.data[LOCAL_CLUSTER] = items + elif isinstance(data, dict): + self.update(data) + elif data is not None: + raise TypeError(f"Invalid Type: {type(data).__name__}") + + def _check_for_value(self, val_id, cluster): + cluster_data = self.data.get(cluster) + if cluster_data and val_id in cluster_data: + return True + return False + + def _get_cluster(self): + cluster = None + if not self.data or LOCAL_CLUSTER in self.data: + cluster = LOCAL_CLUSTER + else: + try: + cluster = next(iter(self.keys())) + except StopIteration: + raise KeyError("Collection is Empty") from None + + return cluster + + def _get_key_and_cluster(self, item): + if isinstance(item, self._val_type): + cluster, key = item.cluster, self._item_id(item) + elif isinstance(item, tuple) and len(item) == 2: + cluster, key = item + else: + cluster, key = self._get_cluster(), item + + return cluster, key + + def _check_val_type(self, item): + if not isinstance(item, self._val_type): + raise TypeError(f"Invalid Type: {type(item).__name__}. " + f"{self._val_type}.__name__ is required.") + + def _item_id(self, item): + return self._id_attr.__get__(item) + + def __eq__(self, other): + if isinstance(other, self.__class__): + return self.data == other.data + return NotImplemented + + def __getitem__(self, item): + if item in self.data: + return self.data[item] + + cluster, key = self._get_key_and_cluster(item) + return self.data[cluster][key] + + def __setitem__(self, where, item): + if where in self.data: + self.data[where] = item + else: + cluster, key = self._get_key_and_cluster(where) + self.data[cluster][key] = item + + def __delitem__(self, item): + if item in self.data: + del self.data[item] + else: + cluster, key = self._get_key_and_cluster(item) + del self.data[cluster][key] + + def __len__(self): + return sum(len(data) for data in self.data.values()) + + def __repr__(self): + return f'{self._typ}([{", ".join(map(repr, self.values()))}])' + + def __contains__(self, item): + if isinstance(item, self._val_type): + item = (item.cluster, self._item_id(item)) + return self.get(item, default=None) is not None + # return self._check_for_value(self._item_id(item), item.cluster) + elif isinstance(item, self._key_type): + found = False + for cluster, data in self.data.items(): + if item in data: + found = True + return found + elif isinstance(item, tuple): + return self.get(item, default=None) is not None + # return self._check_for_value(item, cluster) + + return False + + def __iter__(self): + return iter(self.keys()) + + def __bool__(self): + return bool(self.data) + + def __copy__(self): + return self.copy() + + def copy(self): + """Return a Copy of this instance.""" + out = self.__class__.__new__(self.__class__) + super(self.__class__, out).__init__( + data=self.data.copy(), + typ=self._typ, + key_type=self._key_type, + val_type=self._val_type, + init_data=False, + ) + return out + + def get(self, key, default=None): + """Get the specific value for a Key + + This behaves like `dict`'s `get` method, with the difference that you + can additionally pass in a 2-tuple in the form of `(cluster, key)` as + the key, which can be helpful if this collection contains data from + multiple Clusters. + + If just a key without notion of the Cluster is given, access to the + local cluster data is implied. If this collection does however not + contain data from the local cluster, the first cluster detected + according to `next(iter(self.keys()))` will be used. + + Examples: + Get a Job from the LOCAL_CLUSTER + + >>> job_id = 1 + >>> job = data.get(job_id) + + Get a Job from another Cluster in the Collection, by providing a + 2-tuple with the cluster identifier: + + >>> job_id = 1 + >>> job = data.get(("REMOTE_CLUSTER", job_id)) + """ + cluster, key = self._get_key_and_cluster(key) + return self.data.get(cluster, {}).get(key, default) + + def add(self, item): + """An Item to add to the collection + + Note that a collection can only hold its specific type. + For example, a collection of `pyslurm.Jobs` can only hold + `pyslurm.Job` objects. Trying to add anything other than the accepted + type will raise a TypeError. + + Args: + item (Any): + Item to add to the collection. + + Raises: + TypeError: When an item with an unexpected type not belonging to + the collection was added. + + Examples: + Add a `pyslurm.Job` instance to the `Jobs` collection. + + >>> data = pyslurm.Jobs() + >>> job = pyslurm.Job(1) + >>> data.add(job) + >>> print(data) + Jobs([Job(1)]) + """ + if item.cluster not in self.data: + self.data[item.cluster] = {} + + self._check_val_type(item) + self.data[item.cluster][self._item_id(item)] = item + + def to_json(self, multi_cluster=False): + """Convert the collection to JSON. + + Returns: + (str): JSON formatted string from `json.dumps()` + """ + data = multi_dict_recursive(self) + if multi_cluster: + return json.dumps(data) + else: + cluster = self._get_cluster() + return json.dumps(data[cluster]) + + def keys(self): + """Return a View of all the Keys in this collection + + Returns: + (KeysView): View of all Keys + + Examples: + Iterate over all Keys from all Clusters: + + >>> for key in collection.keys() + ... print(key) + + Iterate over all Keys from all Clusters with the name of the + Cluster additionally provided: + + >>> for cluster, key in collection.keys().with_cluster() + ... print(cluster, key) + """ + return KeysView(self) + + def items(self): + """Return a View of all the Values in this collection + + Returns: + (ItemsView): View of all Items + + Examples: + Iterate over all Items from all Clusters: + + >>> for key, value in collection.items() + ... print(key, value) + + Iterate over all Items from all Clusters with the name of the + Cluster additionally provided: + + >>> for cluster, key, value in collection.items().with_cluster() + ... print(cluster, key, value) + """ + return ItemsView(self) + + def values(self): + """Return a View of all the Values in this collection + + Returns: + (ValuesView): View of all Values + + Examples: + Iterate over all Values from all Clusters: + + >>> for value in collection.values() + ... print(value) + """ + return ValuesView(self) + + def clusters(self): + """Return a View of all the Clusters in this collection + + Returns: + (ClustersView): View of Cluster keys + + Examples: + Iterate over all Cluster-Names the Collection contains: + + >>> for cluster in collection.clusters() + ... print(cluster) + """ + return ClustersView(self) + + def popitem(self): + """Remove and return a `(key, value)` pair as a 2-tuple""" + try: + item = next(iter(self.values())) + except StopIteration: + raise KeyError from None + + key = self._item_id(item) + del self.data[item.cluster][key] + return (key, item) + + def clear(self): + """Clear the collection""" + self.data.clear() + + def pop(self, key, default=None): + """Remove key from the collection and return the value + + This behaves like `dict`'s `pop` method, with the difference that you + can additionally pass in a 2-tuple in the form of `(cluster, key)` as + the key, which can be helpful if this collection contains data from + multiple Clusters. + + If just a key without notion of the Cluster is given, access to the + local cluster data is implied. If this collection does however not + contain data from the local cluster, the first cluster detected + according to `next(iter(self.keys()))` will be used. + """ + item = self.get(key, default=default) + if item is default or item == default: + return default + + cluster = item.cluster + del self.data[cluster][key] + if not self.data[cluster]: + del self.data[cluster] + + return item + + def _update(self, data): + for key in data: + try: + iterator = iter(data[key]) + except TypeError as e: + cluster = self._get_cluster() + if not cluster in self.data: + self.data[cluster] = {} + self.data[cluster].update(data) + break + else: + cluster = key + if not cluster in self.data: + self.data[cluster] = {} + self.data[cluster].update(data[cluster]) +# col = data[cluster] +# if hasattr(col, "keys") and callable(col.keys): +# for k in col.keys(): + +# else: +# for item in col: +# k, v = item + + + def update(self, data={}, **kwargs): + """Update the collection. + + This functions like `dict`'s `update` method. + """ + self._update(data) + self._update(kwargs) + + +def multi_reload(cur, frozen=True): + if not cur: + return cur + + new = cur.__class__.load() + for cluster, item in list(cur.keys().with_cluster()): + if (cluster, item) in new.keys().with_cluster(): + cur[cluster][item] = new.pop(item, cluster) + elif not frozen: + del cur[cluster][item] + + if not frozen: + for cluster, item in new.keys().with_cluster(): + if (cluster, item) not in cur.keys().with_cluster(): + cur[cluster][item] = new[cluster][item] + + return cur + + +def dict_recursive(collection): + cdef dict out = {} + for item_id, item in collection.items(): + if hasattr(item, "to_dict"): + out[item_id] = item.to_dict() + return out + + +def to_json(collection): + return json.dumps(dict_recursive(collection)) + + +def multi_dict_recursive(collection): + cdef dict out = collection.data.copy() + for cluster, data in collection.data.items(): + out[cluster] = dict_recursive(data) + return out + + +def sum_property(collection, prop, startval=0): + out = startval + for item in collection.values(): + data = prop.__get__(item) + if data is not None: + out += data + + return out diff --git a/tests/integration/test_db_job.py b/tests/integration/test_db_job.py index 571ec0d2..1ea59690 100644 --- a/tests/integration/test_db_job.py +++ b/tests/integration/test_db_job.py @@ -49,7 +49,7 @@ def test_parse_all(submit_job): job = submit_job() util.wait() db_job = pyslurm.db.Job.load(job.id) - job_dict = db_job.as_dict() + job_dict = db_job.to_dict() assert job_dict["stats"] assert job_dict["steps"] diff --git a/tests/integration/test_db_qos.py b/tests/integration/test_db_qos.py index 11d9e870..e1cde024 100644 --- a/tests/integration/test_db_qos.py +++ b/tests/integration/test_db_qos.py @@ -38,7 +38,7 @@ def test_load_single(): def test_parse_all(submit_job): qos = pyslurm.db.QualityOfService.load("normal") - qos_dict = qos.as_dict() + qos_dict = qos.to_dict() assert qos_dict assert qos_dict["name"] == qos.name diff --git a/tests/integration/test_job.py b/tests/integration/test_job.py index cef42daf..9788af45 100644 --- a/tests/integration/test_job.py +++ b/tests/integration/test_job.py @@ -35,9 +35,7 @@ def test_parse_all(submit_job): job = submit_job() - # Use the as_dict() function to test if parsing works for all - # properties on a simple Job without error. - Job.load(job.id).as_dict() + Job.load(job.id).to_dict() def test_load(submit_job): @@ -150,7 +148,7 @@ def test_get_job_queue(submit_job): # Submit 10 jobs, gather the job_ids in a list job_list = [submit_job() for i in range(10)] - jobs = Jobs.load().as_dict() + jobs = Jobs.load() for job in job_list: # Check to see if all the Jobs we submitted exist assert job.id in jobs diff --git a/tests/integration/test_job_steps.py b/tests/integration/test_job_steps.py index b24409f5..8d13ba9f 100644 --- a/tests/integration/test_job_steps.py +++ b/tests/integration/test_job_steps.py @@ -102,7 +102,7 @@ def test_collection(submit_job): job = submit_job(script=create_job_script_multi_step()) time.sleep(util.WAIT_SECS_SLURMCTLD) - steps = JobSteps.load(job).as_dict() + steps = JobSteps.load(job) assert steps # We have 3 Steps: batch, 0 and 1 @@ -116,7 +116,7 @@ def test_cancel(submit_job): job = submit_job(script=create_job_script_multi_step()) time.sleep(util.WAIT_SECS_SLURMCTLD) - steps = JobSteps.load(job).as_dict() + steps = JobSteps.load(job) assert len(steps) == 3 assert ("batch" in steps and 0 in steps and @@ -125,7 +125,7 @@ def test_cancel(submit_job): steps[0].cancel() time.sleep(util.WAIT_SECS_SLURMCTLD) - steps = JobSteps.load(job).as_dict() + steps = JobSteps.load(job) assert len(steps) == 2 assert ("batch" in steps and 1 in steps) @@ -173,8 +173,5 @@ def test_load_with_wrong_step_id(submit_job): def test_parse_all(submit_job): job = submit_job() - - # Use the as_dict() function to test if parsing works for all - # properties on a simple JobStep without error. time.sleep(util.WAIT_SECS_SLURMCTLD) - JobStep.load(job, "batch").as_dict() + JobStep.load(job, "batch").to_dict() diff --git a/tests/integration/test_node.py b/tests/integration/test_node.py index 49a69db2..a1c9f6b6 100644 --- a/tests/integration/test_node.py +++ b/tests/integration/test_node.py @@ -29,7 +29,7 @@ def test_load(): - name = Nodes.load()[0].name + name, _ = Nodes.load().popitem() # Now load the node info node = Node.load(name) @@ -56,7 +56,7 @@ def test_create(): def test_modify(): - node = Node(Nodes.load()[0].name) + _, node = Nodes.load().popitem() node.modify(Node(weight=10000)) assert Node.load(node.name).weight == 10000 @@ -69,4 +69,5 @@ def test_modify(): def test_parse_all(): - Node.load(Nodes.load()[0].name).as_dict() + _, node = Nodes.load().popitem() + assert node.to_dict() diff --git a/tests/integration/test_partition.py b/tests/integration/test_partition.py index 8d7a4de4..712eeaff 100644 --- a/tests/integration/test_partition.py +++ b/tests/integration/test_partition.py @@ -28,7 +28,7 @@ def test_load(): - part = Partitions.load()[0] + name, part = Partitions.load().popitem() assert part.name assert part.state @@ -49,7 +49,7 @@ def test_create_delete(): def test_modify(): - part = Partitions.load()[0] + _, part = Partitions.load().popitem() part.modify(Partition(default_time=120)) assert Partition.load(part.name).default_time == 120 @@ -57,8 +57,8 @@ def test_modify(): part.modify(Partition(default_time="1-00:00:00")) assert Partition.load(part.name).default_time == 24*60 - part.modify(Partition(default_time="UNLIMITED")) - assert Partition.load(part.name).default_time == "UNLIMITED" + part.modify(Partition(max_time="UNLIMITED")) + assert Partition.load(part.name).max_time == "UNLIMITED" part.modify(Partition(state="DRAIN")) assert Partition.load(part.name).state == "DRAIN" @@ -68,23 +68,23 @@ def test_modify(): def test_parse_all(): - Partitions.load()[0].as_dict() + _, part = Partitions.load().popitem() + assert part.to_dict() def test_reload(): _partnames = [util.randstr() for i in range(3)] _tmp_parts = Partitions(_partnames) - for part in _tmp_parts: + for part in _tmp_parts.values(): part.create() all_parts = Partitions.load() assert len(all_parts) >= 3 my_parts = Partitions(_partnames[1:]).reload() - print(my_parts) assert len(my_parts) == 2 - for part in my_parts: + for part in my_parts.values(): assert part.state != "UNKNOWN" - for part in _tmp_parts: + for part in _tmp_parts.values(): part.delete() diff --git a/tests/unit/test_collection.py b/tests/unit/test_collection.py new file mode 100644 index 00000000..ccb27779 --- /dev/null +++ b/tests/unit/test_collection.py @@ -0,0 +1,328 @@ +######################################################################### +# test_collection.py - custom collection unit tests +######################################################################### +# Copyright (C) 2023 Toni Harzendorf +# +# This file is part of PySlurm +# +# PySlurm is free software; you can redistribute it and/or modify +# it under the terms of the GNU General Public License as published by +# the Free Software Foundation; either version 2 of the License, or +# (at your option) any later version. + +# PySlurm is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU General Public License for more details. +# +# You should have received a copy of the GNU General Public License along +# with PySlurm; if not, write to the Free Software Foundation, Inc., +# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA. +"""test_collection.py - Unit test custom collection functionality.""" + +import pytest +import pyslurm +from pyslurm.xcollections import sum_property + +LOCAL_CLUSTER = pyslurm.settings.LOCAL_CLUSTER +OTHER_CLUSTER = "other_cluster" + + +class TestMultiClusterMap: + + def _create_collection(self): + data = { + LOCAL_CLUSTER: { + 1: pyslurm.db.Job(1), + 2: pyslurm.db.Job(2), + }, + OTHER_CLUSTER: { + 1: pyslurm.db.Job(1, cluster="other_cluster"), + 10: pyslurm.db.Job(10, cluster="other_cluster"), + } + } + col = pyslurm.db.Jobs() + col.update(data) + return col + + def test_create(self): + jobs = pyslurm.db.Jobs("101,102") + assert len(jobs) == 2 + assert 101 in jobs + assert 102 in jobs + assert jobs[101].id == 101 + assert jobs[102].id == 102 + + jobs = pyslurm.db.Jobs([101, 102]) + assert len(jobs) == 2 + assert 101 in jobs + assert 102 in jobs + assert jobs[101].id == 101 + assert jobs[102].id == 102 + + jobs = pyslurm.db.Jobs( + { + 101: pyslurm.db.Job(101), + 102: pyslurm.db.Job(102), + } + ) + assert len(jobs) == 2 + assert 101 in jobs + assert 102 in jobs + assert jobs[101].id == 101 + assert jobs[102].id == 102 + assert True + + def test_add(self): + col = self._create_collection() + col_len = len(col) + + item = pyslurm.db.Job(20) + col.add(item) + + assert len(col[LOCAL_CLUSTER]) == 3 + assert len(col) == col_len+1 + + item = pyslurm.db.Job(20, cluster=OTHER_CLUSTER) + col.add(item) + + assert len(col[LOCAL_CLUSTER]) == 3 + assert len(col) == col_len+2 + + def test_get(self): + col = self._create_collection() + + item = col.get(1) + assert item is not None + assert isinstance(item, pyslurm.db.Job) + assert item.cluster == LOCAL_CLUSTER + + item = col.get((OTHER_CLUSTER, 1)) + assert item is not None + assert isinstance(item, pyslurm.db.Job) + assert item.cluster == OTHER_CLUSTER + + item = col.get(30) + assert item is None + + def test_keys(self): + col = self._create_collection() + + keys = col.keys() + keys_with_cluster = keys.with_cluster() + assert len(keys) == len(col) + + for k in keys: + assert k + + for cluster, k in keys_with_cluster: + assert cluster + assert cluster in col.data + assert k + + def test_values(self): + col = self._create_collection() + values = col.values() + + assert len(values) == len(col) + + for item in values: + assert item + print(item) + assert isinstance(item, pyslurm.db.Job) + assert item.cluster in col.data + + def test_getitem(self): + col = self._create_collection() + + item1 = col[LOCAL_CLUSTER][1] + item2 = col[1] + item3 = col[OTHER_CLUSTER][1] + + assert item1 + assert item2 + assert item3 + assert item1 == item2 + assert item1 != item3 + + with pytest.raises(KeyError): + item = col[30] + + with pytest.raises(KeyError): + item = col[OTHER_CLUSTER][30] + + def test_setitem(self): + col = self._create_collection() + col_len = len(col) + + item = pyslurm.db.Job(30) + col[item.id] = item + assert len(col[LOCAL_CLUSTER]) == 3 + assert len(col) == col_len+1 + + item = pyslurm.db.Job(50, cluster=OTHER_CLUSTER) + col[OTHER_CLUSTER][item.id] = item + assert len(col[OTHER_CLUSTER]) == 3 + assert len(col) == col_len+2 + + item = pyslurm.db.Job(100, cluster=OTHER_CLUSTER) + col[item] = item + assert len(col[OTHER_CLUSTER]) == 4 + assert len(col) == col_len+3 + + item = pyslurm.db.Job(101, cluster=OTHER_CLUSTER) + col[(item.cluster, item.id)] = item + assert len(col[OTHER_CLUSTER]) == 5 + assert len(col) == col_len+4 + + new_other_data = { + 1: pyslurm.db.Job(1), + 2: pyslurm.db.Job(2), + } + col[OTHER_CLUSTER] = new_other_data + assert len(col[OTHER_CLUSTER]) == 2 + assert len(col[LOCAL_CLUSTER]) == 3 + assert 1 in col[OTHER_CLUSTER] + assert 2 in col[OTHER_CLUSTER] + + def test_delitem(self): + col = self._create_collection() + col_len = len(col) + + del col[1] + assert len(col[LOCAL_CLUSTER]) == 1 + assert len(col) == col_len-1 + + del col[OTHER_CLUSTER][1] + assert len(col[OTHER_CLUSTER]) == 1 + assert len(col) == col_len-2 + + del col[OTHER_CLUSTER] + assert len(col) == 1 + assert OTHER_CLUSTER not in col.data + + def test_copy(self): + col = self._create_collection() + col_copy = col.copy() + assert col == col_copy + + def test_iter(self): + col = self._create_collection() + for k in col: + assert k + + def test_items(self): + col = self._create_collection() + for k, v in col.items(): + assert k + assert v + assert isinstance(v, pyslurm.db.Job) + + for c, k, v in col.items().with_cluster(): + assert c + assert k + assert v + assert isinstance(v, pyslurm.db.Job) + + def test_popitem(self): + col = self._create_collection() + col_len = len(col) + + key, item = col.popitem() + assert item + assert key + assert isinstance(item, pyslurm.db.Job) + assert len(col) == col_len-1 + + def test_update(self): + col = self._create_collection() + col_len = len(col) + + col_update = { + 30: pyslurm.db.Job(30), + 50: pyslurm.db.Job(50), + } + col.update(col_update) + assert len(col) == col_len+2 + assert len(col[LOCAL_CLUSTER]) == 4 + assert 30 in col + assert 50 in col + + col_update = { + "new_cluster": { + 80: pyslurm.db.Job(80, cluster="new_cluster"), + 50: pyslurm.db.Job(50, cluster="new_cluster"), + } + } + col.update(col_update) + assert len(col) == col_len+4 + assert len(col[LOCAL_CLUSTER]) == 4 + assert len(col["new_cluster"]) == 2 + assert 80 in col + assert 50 in col + + col_update = { + 200: pyslurm.db.Job(200, cluster=OTHER_CLUSTER), + 300: pyslurm.db.Job(300, cluster=OTHER_CLUSTER), + } + col.update({OTHER_CLUSTER: col_update}) + assert len(col) == col_len+6 + assert len(col[OTHER_CLUSTER]) == 4 + assert 200 in col + assert 300 in col + + empty_col = pyslurm.db.Jobs() + empty_col.update(col_update) + assert len(empty_col) == 2 + + def test_pop(self): + col = self._create_collection() + col_len = len(col) + + item = col.pop(1) + assert item + assert item.id == 1 + assert len(col) == col_len-1 + + item = col.pop(999, default="def") + assert item == "def" + + def test_contains(self): + col = self._create_collection() + item = pyslurm.db.Job(1) + assert item in col + + assert 10 in col + assert 20 not in col + + assert (OTHER_CLUSTER, 10) in col + assert (LOCAL_CLUSTER, 10) not in col + + def test_to_json(self): + col = self._create_collection() + data = col.to_json(multi_cluster=True) + assert data + + def test_cluster_view(self): + col = self._create_collection() + assert len(col.clusters()) == 2 + for c in col.clusters(): + assert c + + def test_sum_property(self): + class TestObject: + @property + def memory(self): + return 10240 + + @property + def cpus(self): + return None + + object_dict = {i: TestObject() for i in range(10)} + + expected = 10240 * 10 + assert sum_property(object_dict, TestObject.memory) == expected + + expected = 0 + assert sum_property(object_dict, TestObject.cpus) == expected diff --git a/tests/unit/test_common.py b/tests/unit/test_common.py index 1598d191..cf5353b1 100644 --- a/tests/unit/test_common.py +++ b/tests/unit/test_common.py @@ -55,12 +55,11 @@ nodelist_from_range_str, nodelist_to_range_str, instance_to_dict, - collection_to_dict, - collection_to_dict_global, - group_collection_by_cluster, - _sum_prop, ) from pyslurm.utils import cstr +from pyslurm.xcollections import ( + sum_property, +) class TestStrings: @@ -414,75 +413,3 @@ def test_nodelist_to_range_str(self): assert "node[001,007-009]" == nodelist_to_range_str(nodelist) assert "node[001,007-009]" == nodelist_to_range_str(nodelist_str) - def test_summarize_property(self): - class TestObject: - @property - def memory(self): - return 10240 - - @property - def cpus(self): - return None - - object_dict = {i: TestObject() for i in range(10)} - - expected = 10240 * 10 - assert _sum_prop(object_dict, TestObject.memory) == expected - - expected = 0 - assert _sum_prop(object_dict, TestObject.cpus) == expected - - def test_collection_to_dict(self): - class TestObject: - - def __init__(self, _id, _grp_id, cluster): - self._id = _id - self._grp_id = _grp_id - self.cluster = cluster - - @property - def id(self): - return self._id - - @property - def group_id(self): - return self._grp_id - - def as_dict(self): - return instance_to_dict(self) - - class TestCollection(list): - - def __init__(self, data): - super().__init__() - self.extend(data) - - OFFSET = 100 - RANGE = 10 - - data = [TestObject(x, x+OFFSET, "TestCluster") for x in range(RANGE)] - collection = TestCollection(data) - - coldict = collection_to_dict(collection, identifier=TestObject.id) - coldict = coldict.get("TestCluster", {}) - - assert len(coldict) == RANGE - for i in range(RANGE): - assert i in coldict - assert isinstance(coldict[i], TestObject) - - coldict = collection_to_dict(collection, identifier=TestObject.id, - group_id=TestObject.group_id) - coldict = coldict.get("TestCluster", {}) - - assert len(coldict) == RANGE - for i in range(RANGE): - assert i+OFFSET in coldict - assert i in coldict[i+OFFSET] - - coldict = collection_to_dict(collection, identifier=TestObject.id, - recursive=True) - coldict = coldict.get("TestCluster", {}) - - for item in coldict.values(): - assert isinstance(item, dict) diff --git a/tests/unit/test_db_job.py b/tests/unit/test_db_job.py index 7b77671f..c2ae8bb0 100644 --- a/tests/unit/test_db_job.py +++ b/tests/unit/test_db_job.py @@ -42,38 +42,11 @@ def test_filter(): job_filter._create() -def test_create_collection(): - jobs = pyslurm.db.Jobs("101,102") - assert len(jobs) == 2 - jobs = jobs.as_dict() - assert 101 in jobs - assert 102 in jobs - assert jobs[101].id == 101 - assert jobs[102].id == 102 - - jobs = pyslurm.db.Jobs([101, 102]) - assert len(jobs) == 2 - jobs = jobs.as_dict() - assert 101 in jobs - assert 102 in jobs - assert jobs[101].id == 101 - assert jobs[102].id == 102 - - jobs = pyslurm.db.Jobs( - { - 101: pyslurm.db.Job(101), - 102: pyslurm.db.Job(102), - } - ) - assert len(jobs) == 2 - jobs = jobs.as_dict() - assert 101 in jobs - assert 102 in jobs - assert jobs[101].id == 101 - assert jobs[102].id == 102 - assert True - - def test_create_instance(): job = pyslurm.db.Job(9999) assert job.id == 9999 + + +def test_parse_all(): + job = pyslurm.db.Job(9999) + assert job.to_dict() diff --git a/tests/unit/test_db_qos.py b/tests/unit/test_db_qos.py index 0d2fd538..5ee2db76 100644 --- a/tests/unit/test_db_qos.py +++ b/tests/unit/test_db_qos.py @@ -39,11 +39,6 @@ def test_search_filter(): qos_filter._create() -def test_create_collection_instance(): - # TODO - assert True - - def test_create_instance(): qos = pyslurm.db.QualityOfService("test") assert qos.name == "test" diff --git a/tests/unit/test_job.py b/tests/unit/test_job.py index edcf65d4..863fcfab 100644 --- a/tests/unit/test_job.py +++ b/tests/unit/test_job.py @@ -31,9 +31,7 @@ def test_create_instance(): def test_parse_all(): - # Use the as_dict() function to test if parsing works for all - # properties on a simple Job without error. - Job(9999).as_dict() + assert Job(9999).to_dict() def test_parse_dependencies_to_dict(): diff --git a/tests/unit/test_job_steps.py b/tests/unit/test_job_steps.py index fcd0d012..c8c52352 100644 --- a/tests/unit/test_job_steps.py +++ b/tests/unit/test_job_steps.py @@ -39,6 +39,4 @@ def test_create_instance(): def test_parse_all(): - # Use the as_dict() function to test if parsing works for all - # properties on a simple JobStep without error. - JobStep(9999, 1).as_dict() + assert JobStep(9999, 1).to_dict() diff --git a/tests/unit/test_node.py b/tests/unit/test_node.py index 755e85d9..c4dba73e 100644 --- a/tests/unit/test_node.py +++ b/tests/unit/test_node.py @@ -32,35 +32,7 @@ def test_create_instance(): def test_parse_all(): - Node("localhost").as_dict() - - -def test_create_nodes_collection(): - nodes = Nodes("node1,node2").as_dict() - assert len(nodes) == 2 - assert "node1" in nodes - assert "node2" in nodes - assert nodes["node1"].name == "node1" - assert nodes["node2"].name == "node2" - - nodes = Nodes(["node1", "node2"]).as_dict() - assert len(nodes) == 2 - assert "node1" in nodes - assert "node2" in nodes - assert nodes["node1"].name == "node1" - assert nodes["node2"].name == "node2" - - nodes = Nodes( - { - "node1": Node("node1"), - "node2": Node("node2"), - } - ).as_dict() - assert len(nodes) == 2 - assert "node1" in nodes - assert "node2" in nodes - assert nodes["node1"].name == "node1" - assert nodes["node2"].name == "node2" + assert Node("localhost").to_dict() def test_set_node_state(): diff --git a/tests/unit/test_partition.py b/tests/unit/test_partition.py index 89403ae2..b699893c 100644 --- a/tests/unit/test_partition.py +++ b/tests/unit/test_partition.py @@ -31,36 +31,8 @@ def test_create_instance(): assert part.name == "normal" -def test_create_collection(): - parts = Partitions("part1,part2").as_dict() - assert len(parts) == 2 - assert "part1" in parts - assert "part2" in parts - assert parts["part1"].name == "part1" - assert parts["part2"].name == "part2" - - parts = Partitions(["part1", "part2"]).as_dict() - assert len(parts) == 2 - assert "part1" in parts - assert "part2" in parts - assert parts["part1"].name == "part1" - assert parts["part2"].name == "part2" - - parts = Partitions( - { - "part1": Partition("part1"), - "part2": Partition("part2"), - } - ).as_dict() - assert len(parts) == 2 - assert "part1" in parts - assert "part2" in parts - assert parts["part1"].name == "part1" - assert parts["part2"].name == "part2" - - def test_parse_all(): - Partition("normal").as_dict() + assert Partition("normal").to_dict() def test_parse_memory():