diff --git a/docs/reference/reservation.md b/docs/reference/reservation.md index 5826fcf9..8fc9f401 100644 --- a/docs/reference/reservation.md +++ b/docs/reference/reservation.md @@ -2,8 +2,7 @@ title: Reservation --- -!!! warning - This API is currently being completely reworked, and is subject to be - removed in the future when a replacement is introduced - -::: pyslurm.deprecated.reservation +::: pyslurm.Reservation +::: pyslurm.Reservations +::: pyslurm.ReservationFlags +::: pyslurm.ReservationReoccurrence diff --git a/pyslurm/__init__.py b/pyslurm/__init__.py index c6e41eb7..9bf54325 100644 --- a/pyslurm/__init__.py +++ b/pyslurm/__init__.py @@ -22,6 +22,12 @@ ) from pyslurm.core.node import Node, Nodes from pyslurm.core.partition import Partition, Partitions +from pyslurm.core.reservation import ( + Reservation, + Reservations, + ReservationFlags, + ReservationReoccurrence, +) from pyslurm.core import error from pyslurm.core.error import ( PyslurmError, diff --git a/pyslurm/core/reservation.pxd b/pyslurm/core/reservation.pxd new file mode 100644 index 00000000..e5b9e62a --- /dev/null +++ b/pyslurm/core/reservation.pxd @@ -0,0 +1,166 @@ +######################################################################### +# reservation.pxd - interface to work with reservations in slurm +######################################################################### +# Copyright (C) 2025 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 libc.string cimport memcpy, memset +from libc.stdint cimport uint8_t, uint16_t, uint32_t, uint64_t +from libc.stdlib cimport free +from pyslurm cimport slurm +from pyslurm.slurm cimport ( + reserve_info_t, + reserve_info_msg_t, + resv_desc_msg_t, + reservation_name_msg_t, + reserve_response_msg_t, + slurm_free_reservation_info_msg, + slurm_load_reservations, + slurm_delete_reservation, + slurm_update_reservation, + slurm_create_reservation, + slurm_init_resv_desc_msg, + xfree, + try_xmalloc, +) + +from pyslurm.utils cimport cstr +from pyslurm.utils cimport ctime +from pyslurm.utils.ctime cimport time_t +from pyslurm.utils.uint cimport ( + u32, + u32_parse, + u64_parse_bool_flag, + u64_set_bool_flag, +) +from pyslurm.xcollections cimport MultiClusterMap + +cdef extern void slurm_free_resv_desc_msg(resv_desc_msg_t *msg) +cdef extern void slurm_free_reserve_info_members(reserve_info_t *resv) + + +cdef class Reservations(MultiClusterMap): + """A [`Multi Cluster`][pyslurm.xcollections.MultiClusterMap] collection of [pyslurm.Reservation][] objects. + + Args: + reservations (Union[list[str], dict[str, pyslurm.Reservation], str], optional=None): + Reservations to initialize this collection with. + """ + cdef: + reserve_info_msg_t *info + reserve_info_t tmp_info + + +cdef class Reservation: + """A Slurm Reservation. + + Args: + name (str, optional=None): + Name for a Reservation. + **kwargs (Any, optional=None): + All Attributes of a Reservation are eligible to be set, except + `cpu_ids_by_node`. Although the `name` attribute can also be + changed on the instance, the change will not be taken into account + by `slurmctld` when calling `modify()`. + + Attributes: + accounts (list[str]): + List of account names that have access to the Reservation. + burst_buffer (str): + Burst Buffer specification. + comment (str): + Arbitrary comment for the Reservation. + cpus (int): + Amount of CPUs used by the Reservation + cpu_ids_by_node (dict[str, int]): + A Mapping where each key is the node-name, and the values are a + string of CPU-IDs reserved on the specific nodes. + end_time (int): + Unix Timestamp when the Reservation ends. + features (list[str]): + List of features required by the Reservation. + groups (list[str]): + List of Groups that can access the Reservation. + licenses (list[str]): + List of licenses to be reserved. + max_start_delay (int): + Maximum delay, in seconds, where Jobs are permitted to overlap with + the Reservation once Jobs are queued for it. + name (str): + Name of the Reservation. + node_count (int): + Count of Nodes required. + nodes (str): + Nodes to be reserved. + When creating or updating a Reservation, you can also pass the + string `ALL`, when you want all nodes to be included in the + Reservation. + partition (str): + Name of the partition to be used. + purge_time (int): + When the Reservation is idle for this amount of seconds, it will be + removed. + start_time (int): + When the Reservation starts. This is a Unix timestamp. + duration (int): + How long, in minutes, the reservation runs for. Unless a + `start_time` has already been specified, setting this will set the + `start_time` the current time, meaning the Reservation will start + immediately. + + For setting this attribute, instead of minutes you can also specify + a time-string like this: + + duration = "1-00:00:00" + + The above means that the Reservation will last for 1 day. + is_active (bool): + Whether the reservation is currently active or not. + tres (dict[str, int]): + TRES for the Reservation. + users (list[str]): + List of user names permitted to use the Reservation. + flags (pyslurm.ReservationFlags): + Optional Flags for the Reservation. + For convenience, instead of using [pyslurm.ReservationFlags][] in + combination with the logical Operators, you can set this attribute + via a [list][] of [str][]: + + flags = ["MAINTENANCE", "FLEX", "MAGNETIC"] + + When setting like this, the strings must match the names of members + in [pyslurm.ReservationFlags][]. + reoccurrence (pyslurm.ReservationReoccurrence): + Describes if and when this Reservation reoccurs. + Since [pyslurm.ReservationReoccurrence] members are also just + strings, you can conveniently also set the attribute like this: + + reoccurrence = "DAILY" + """ + cdef: + reserve_info_t *info + resv_desc_msg_t *umsg + _reoccurrence + + cdef readonly cluster + + @staticmethod + cdef Reservation from_ptr(reserve_info_t *in_ptr) diff --git a/pyslurm/core/reservation.pyx b/pyslurm/core/reservation.pyx new file mode 100644 index 00000000..e208c095 --- /dev/null +++ b/pyslurm/core/reservation.pyx @@ -0,0 +1,549 @@ +######################################################################### +# reservation.pyx - interface to work with reservations in slurm +######################################################################### +# Copyright (C) 2025 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 typing import Union, Any +from pyslurm.utils import cstr +from pyslurm.utils import ctime +from pyslurm.utils.uint import u32_parse +from pyslurm import settings +from pyslurm.core.slurmctld.config import _get_memory +from datetime import datetime +from pyslurm import xcollections +from pyslurm.utils.helpers import instance_to_dict +from pyslurm.utils.enums import SlurmEnum, SlurmFlag +from enum import auto, StrEnum +from pyslurm.utils.ctime import ( + _raw_time, + timestr_to_mins, + timestr_to_secs, + date_to_timestamp, +) +from pyslurm.core.error import ( + RPCError, + verify_rpc, + slurm_errno, +) + + +cdef class Reservations(MultiClusterMap): + + def __dealloc__(self): + slurm_free_reservation_info_msg(self.info) + self.info = NULL + + def __cinit__(self): + self.info = NULL + + def __init__(self, reservations=None): + super().__init__(data=reservations, + typ="Reservations", + val_type=Reservation, + id_attr=Reservation.name, + key_type=str) + + @staticmethod + def load(): + """Load all Reservations in the system. + + Returns: + (pyslurm.Reservations): Collection of [pyslurm.Reservation][] + objects. + + Raises: + (pyslurm.RPCError): When getting all the Reservations from the + slurmctld failed. + """ + cdef: + Reservations reservations = Reservations() + Reservation reservation + + verify_rpc(slurm_load_reservations(0, &reservations.info)) + + memset(&reservations.tmp_info, 0, sizeof(reserve_info_t)) + for cnt in range(reservations.info.record_count): + reservation = Reservation.from_ptr(&reservations.info.reservation_array[cnt]) + + # If we already parsed at least one Reservation, and if for some + # reason a MemoryError is raised after parsing subsequent + # reservations, invalid behaviour will be shown by Valgrind, since + # the Memory for the already parsed Reservation will be freed + # twice. So for all successfully parsed Reservations, replace it + # with a dummy struct that will be skipped in case of error. + reservations.info.reservation_array[cnt] = reservations.tmp_info + + cluster = reservation.cluster + if cluster not in reservations.data: + reservations.data[cluster] = {} + + reservations.data[cluster][reservation.name] = reservation + + reservations.info.record_count = 0 + return reservations + + +cdef class Reservation: + + def __cinit__(self): + self.info = NULL + self.umsg = NULL + + def __init__(self, name=None, **kwargs): + self._alloc_impl() + self.name = name + self._reoccurrence = ReservationReoccurrence.NO + self.cluster = settings.LOCAL_CLUSTER + for k, v in kwargs.items(): + setattr(self, k, v) + + def _alloc_impl(self): + self._alloc_info() + self._alloc_umsg() + + def _alloc_info(self): + if not self.info: + self.info = try_xmalloc(sizeof(reserve_info_t)) + if not self.info: + raise MemoryError("xmalloc failed for reserve_info_t") + + def _alloc_umsg(self): + if not self.umsg: + self.umsg = try_xmalloc(sizeof(resv_desc_msg_t)) + if not self.umsg: + raise MemoryError("xmalloc failed for resv_desc_msg_t") + slurm_init_resv_desc_msg(self.umsg) + self.umsg.flags = 0 + + def _dealloc_umsg(self): + slurm_free_resv_desc_msg(self.umsg) + self.umsg = NULL + + def _dealloc_impl(self): + self._dealloc_umsg() + slurm_free_reserve_info_members(self.info) + xfree(self.info) + self.info = NULL + + def __dealloc__(self): + self._dealloc_impl() + + def __setattr__(self, name, val): + # When a user wants to set attributes on a Reservation instance that + # was created by calling Reservations(), the "umsg" pointer is not yet + # allocated. We only allocate memory for it by the time the user + # actually wants to modify something. + self._alloc_umsg() + # Call descriptors __set__ directly + Reservation.__dict__[name].__set__(self, val) + + def __repr__(self): + return f'pyslurm.{self.__class__.__name__}({self.name})' + + @staticmethod + cdef Reservation from_ptr(reserve_info_t *in_ptr): + cdef Reservation wrap = Reservation.__new__(Reservation) + wrap._alloc_info() + wrap.cluster = settings.LOCAL_CLUSTER + memcpy(wrap.info, in_ptr, sizeof(reserve_info_t)) + wrap._reoccurrence = ReservationReoccurrence.from_flag(wrap.info.flags, + default=ReservationReoccurrence.NO) + wrap.info.flags &= ~wrap.reoccurrence._flag + return wrap + + def _error_or_name(self): + if not self.name: + raise RPCError(msg="No Reservation name was specified. " + "Did you set the `name` attribute on the " + "Reservation instance?") + return self.name + + def to_dict(self): + """Reservation information formatted as a dictionary. + + Returns: + (dict): Reservation information as dict + + Examples: + >>> import pyslurm + >>> resv = pyslurm.Reservation.load("maintenance") + >>> resv_dict = resv.to_dict() + >>> print(resv_dict) + """ + return instance_to_dict(self) + + @staticmethod + def load(name): + """Load information for a specific Reservation. + + Args: + name (str): + The name of the Reservation to load. + + Returns: + (pyslurm.Reservation): Returns a new Reservation instance. + + Raises: + (pyslurm.RPCError): If requesting the Reservation information from + the slurmctld was not successful. + + Examples: + >>> import pyslurm + >>> reservation = pyslurm.Reservation.load("maintenance") + """ + resv = Reservations.load().get(name) + if not resv: + raise RPCError(msg=f"Reservation '{name}' doesn't exist") + + return resv + + def create(self): + """Create a Reservation. + + If you did not specify at least a `start_time` and `duration` or + `end_time`, then by default the Reservation will start effective + immediately, with a duration of one year. + + Returns: + (pyslurm.Reservation): This function returns the current + Reservation instance object itself. + + Raises: + (pyslurm.RPCError): If creating the Reservation was not successful. + + Examples: + >>> import pyslurm + >>> from pyslurm import ReservationFlags, ReservationReoccurrence + >>> resv = pyslurm.Reservation( + ... name = "debug", + ... users = ["root"], + ... nodes = "node001", + ... duration = "1-00:00:00", + ... flags = ReservationFlags.MAINTENANCE | ReservationFlags.FLEX, + ... reoccurrence = ReservationReoccurrence.DAILY, + ... ) + >>> resv.create() + """ + cdef char* new_name = NULL + + if not self.start_time or not (self.duration and self.end_time): + raise RPCError(msg="You must at least specify a start_time, " + " combined with an end_time or a duration.") + + self.name = self._error_or_name() + new_name = slurm_create_reservation(self.umsg) + free(new_name) + verify_rpc(slurm_errno()) + return self + + def modify(self, Reservation changes=None): + """Modify a Reservation. + + Args: + changes (pyslurm.Reservation, optional=None): + Another Reservation object that contains all the changes to + apply. This is optional - you can also directly modify a + Reservation object and just call `modify()`, and the changes + will be sent to `slurmctld`. + + Raises: + (pyslurm.RPCError): When updating the Reservation was not + successful. + + Examples: + >>> import pyslurm + >>> + >>> resv = pyslurm.Reservation.load("maintenance") + >>> # Add 60 Minutes to the reservation + >>> resv.duration += 60 + >>> + >>> # You can also add a slurm timestring. + >>> # For example, extend the duration by another day: + >>> resv.duration += pyslurm.utils.timestr_to_mins("1-00:00:00") + >>> + >>> # Now send the changes to the Controller: + >>> resv.modify() + """ + cdef Reservation updates = changes if changes is not None else self + if not updates.umsg: + return + + self._error_or_name() + cstr.fmalloc(&updates.umsg.name, self.info.name) + verify_rpc(slurm_update_reservation(updates.umsg)) + + # Make sure we clean the object from any previous changes. + updates._dealloc_umsg() + + def delete(self): + """Delete a Reservation. + + Raises: + (pyslurm.RPCError): If deleting the Reservation was not successful. + + Examples: + >>> import pyslurm + >>> pyslurm.Reservation("maintenance").delete() + """ + cdef reservation_name_msg_t to_delete + memset(&to_delete, 0, sizeof(to_delete)) + to_delete.name = cstr.from_unicode(self._error_or_name()) + verify_rpc(slurm_delete_reservation(&to_delete)) + + @property + def accounts(self): + return cstr.to_list(self.info.accounts) + + @accounts.setter + def accounts(self, val): + cstr.from_list2(&self.info.accounts, &self.umsg.accounts, val) + + @property + def burst_buffer(self): + return cstr.to_unicode(self.info.burst_buffer) + + @burst_buffer.setter + def burst_buffer(self, val): + cstr.fmalloc2(&self.info.burst_buffer, &self.umsg.burst_buffer, val) + + @property + def comment(self): + return cstr.to_unicode(self.info.comment) + + @comment.setter + def comment(self, val): + cstr.fmalloc2(&self.info.comment, &self.umsg.comment, val) + + @property + def cpus(self): + return u32_parse(self.info.core_cnt, zero_is_noval=False) + + @cpus.setter + def cpus(self, val): + self.info.core_cnt = self.umsg.core_cnt = int(val) + + @property + def cpu_ids_by_node(self): + out = {} + for i in range(self.info.core_spec_cnt): + node = cstr.to_unicode(self.info.core_spec[i].node_name) + if node: + out[node] = cstr.to_unicode(self.info.core_spec[i].core_id) + + return out + + @property + def end_time(self): + return _raw_time(self.info.end_time) + + @end_time.setter + def end_time(self, val): + val = date_to_timestamp(val) + if self.start_time and val < self.info.start_time: + raise ValueError("end_time cannot be earlier then start_time.") + + self.info.end_time = self.umsg.end_time = val + + @property + def features(self): + return cstr.to_list(self.info.features) + + @features.setter + def features(self, val): + cstr.from_list2(&self.info.features, &self.umsg.features, val) + + @property + def groups(self): + return cstr.to_list(self.info.groups) + + @groups.setter + def groups(self, val): + cstr.from_list2(&self.info.groups, &self.umsg.groups, val) + + @property + def licenses(self): + return cstr.to_list(self.info.licenses) + + @licenses.setter + def licenses(self, val): + cstr.from_list2(&self.info.licenses, &self.umsg.licenses, val) + + @property + def max_start_delay(self): + return u32_parse(self.info.max_start_delay) + + @max_start_delay.setter + def max_start_delay(self, val): + self.info.max_start_delay = self.umsg.max_start_delay = int(val) + + @property + def name(self): + return cstr.to_unicode(self.info.name) + + @name.setter + def name(self, val): + cstr.fmalloc2(&self.info.name, &self.umsg.name, val) + + @property + def node_count(self): + return u32_parse(self.info.node_cnt, zero_is_noval=False) + + @node_count.setter + def node_count(self, val): + self.info.node_cnt = self.umsg.node_cnt = int(val) + + @property + def nodes(self): + return cstr.to_unicode(self.info.node_list) + + @nodes.setter + def nodes(self, val): + cstr.fmalloc2(&self.info.node_list, &self.umsg.node_list, val) + + @property + def partition(self): + return cstr.to_unicode(self.info.partition) + + @partition.setter + def partition(self, val): + cstr.fmalloc2(&self.info.partition, &self.umsg.partition, val) + + @property + def purge_time(self): + return u32_parse(self.info.purge_comp_time) + + @purge_time.setter + def purge_time(self, val): + self.info.purge_comp_time = self.umsg.purge_comp_time = timestr_to_secs(val) + if ReservationFlags.PURGE not in self.flags: + self.flags |= ReservationFlags.PURGE + + @property + def start_time(self): + return _raw_time(self.info.start_time) + + @start_time.setter + def start_time(self, val): + self.info.start_time = self.umsg.start_time = date_to_timestamp(val) + + @property + def duration(self): + cdef time_t duration = 0 + + if self.start_time and self.info.end_time >= self.info.start_time: + duration = ctime.difftime(self.info.end_time, + self.info.start_time) + + return int(duration / 60) + + @duration.setter + def duration(self, val): + val = timestr_to_mins(val) + if not self.start_time: + self.start_time = datetime.now() + self.end_time = self.start_time + (val * 60) + + @property + def is_active(self): + cdef time_t now = ctime.time(NULL) + if self.info.start_time <= now and self.info.end_time >= now: + return True + return False + + @property + def tres(self): + return cstr.to_dict(self.info.tres_str) + + @tres.setter + def tres(self, val): + cstr.fmalloc2(&self.info.tres_str, &self.umsg.tres_str, + cstr.dict_to_str(val)) + + @property + def users(self): + return cstr.to_list(self.info.users) + + @users.setter + def users(self, val): + cstr.from_list2(&self.info.users, &self.umsg.users, val) + + @property + def reoccurrence(self): + return self._reoccurrence + + @reoccurrence.setter + def reoccurrence(self, val): + v = ReservationReoccurrence(val) + current = self._reoccurrence + self._reoccurrence = v + if v == ReservationReoccurrence.NO: + self.umsg.flags |= current._clear_flag + else: + self.umsg.flags |= v._flag + + @property + def flags(self): + return ReservationFlags(self.info.flags) + + @flags.setter + def flags(self, val): + flag = val + if isinstance(val, list): + flag = ReservationFlags.from_list(val) + + self.info.flags = flag.value + self.umsg.flags = flag._get_flags_cleared() + + # TODO: RESERVE_FLAG_SKIP ? + + +class ReservationFlags(SlurmFlag): + """Flags for Reservations that can be set. + + See {scontrol#OPT_Flags} for more info. + """ + MAINTENANCE = slurm.RESERVE_FLAG_MAINT, slurm.RESERVE_FLAG_NO_MAINT + MAGNETIC = slurm.RESERVE_FLAG_MAGNETIC, slurm.RESERVE_FLAG_NO_MAGNETIC + FLEX = slurm.RESERVE_FLAG_FLEX, slurm.RESERVE_FLAG_NO_FLEX + IGNORE_RUNNING_JOBS = slurm.RESERVE_FLAG_IGN_JOBS, slurm.RESERVE_FLAG_NO_IGN_JOB + ANY_NODES = slurm.RESERVE_FLAG_ANY_NODES, slurm.RESERVE_FLAG_NO_ANY_NODES + STATIC_NODES = slurm.RESERVE_FLAG_STATIC, slurm.RESERVE_FLAG_NO_STATIC + PARTITION_NODES_ONLY = slurm.RESERVE_FLAG_PART_NODES, slurm.RESERVE_FLAG_NO_PART_NODES + USER_DELETION = slurm.RESERVE_FLAG_USER_DEL, slurm.RESERVE_FLAG_NO_USER_DEL + PURGE = slurm.RESERVE_FLAG_PURGE_COMP, slurm.RESERVE_FLAG_NO_PURGE_COMP + SPECIFIC_NODES = slurm.RESERVE_FLAG_SPEC_NODES + NO_JOB_HOLD_AFTER_END = slurm.RESERVE_FLAG_NO_HOLD_JOBS + OVERLAP = slurm.RESERVE_FLAG_OVERLAP + ALL_NODES = slurm.RESERVE_FLAG_ALL_NODES + + +class ReservationReoccurrence(SlurmEnum): + """Different reocurrences for a Reservation. + + See {scontrol#OPT_Flags} for more info. + """ + NO = auto() + DAILY = auto(), slurm.RESERVE_FLAG_DAILY, slurm.RESERVE_FLAG_NO_DAILY + HOURLY = auto(), slurm.RESERVE_FLAG_HOURLY, slurm.RESERVE_FLAG_NO_HOURLY + WEEKLY = auto(), slurm.RESERVE_FLAG_WEEKLY, slurm.RESERVE_FLAG_NO_WEEKLY + WEEKDAY = auto(), slurm.RESERVE_FLAG_WEEKDAY, slurm.RESERVE_FLAG_NO_WEEKDAY + WEEKEND = auto(), slurm.RESERVE_FLAG_WEEKEND, slurm.RESERVE_FLAG_NO_WEEKEND diff --git a/pyslurm/utils/enums.pyx b/pyslurm/utils/enums.pyx new file mode 100644 index 00000000..f453d2f8 --- /dev/null +++ b/pyslurm/utils/enums.pyx @@ -0,0 +1,109 @@ +######################################################################### +# utils/enums.pyx - pyslurm enum helpers +######################################################################### +# Copyright (C) 2025 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 enum import Enum, Flag +import inspect + +try: + from enum import EnumMeta as EnumType +except ImportError: + from enum import EnumType + + +class DocstringSupport(EnumType): + def __new__(metacls, clsname, bases, classdict): + cls = super().__new__(metacls, clsname, bases, classdict) + + # In the future, if we want to properly document enum members, + # implement this: + # source = inspect.getdoc(cls) + # docstrings = source.replace(" ", "").split("\n") + + for member in cls: + member.__doc__ = "" + + return cls + + +class SlurmEnum(str, Enum, metaclass=DocstringSupport): + + def __new__(cls, name, *args): + # https://docs.python.org/3/library/enum.html + # + # 1. + # Second argument to str is encoding, third is error. We don't really + # care for that, so no need to check. + # 2. + # Python Documentation recommends to not call super().__new__, but + # the corresponding types __new__ directly, so str here. + # 3. + # Docs recommend to set _value_ + v = str(name) + new_string = str.__new__(cls, v) + new_string._value_ = v + + new_string._flag = int(args[0]) if len(args) >= 1 else 0 + new_string._clear_flag = int(args[1]) if len(args) >= 2 else 0 + return new_string + + def __str__(self): + return str(self.value) + + @staticmethod + def _generate_next_value_(name, _start, _count, _last_values): + # We just care about the name of the member to be defined. + return name.upper() + + @classmethod + def from_flag(cls, flag, default): + out = cls(default) + for item in cls: + if item._flag & flag: + return item + return out + + +class SlurmFlag(Flag, metaclass=DocstringSupport): + + def __new__(cls, flag, *args): + obj = super()._new_member_(cls) + obj._value_ = int(flag) + obj._clear_flag = int(args[0]) if len(args) >= 1 else 0 + return obj + + @classmethod + def from_list(cls, inp): + out = cls(0) + for flag in cls: + if flag.name in inp: + out |= flag + + return out + + def _get_flags_cleared(self): + val = self.value + for flag in self.__class__: + if flag not in self: + val |= flag._clear_flag + return val diff --git a/pyslurm/utils/uint.pxd b/pyslurm/utils/uint.pxd index d886d6f3..44b2eb27 100644 --- a/pyslurm/utils/uint.pxd +++ b/pyslurm/utils/uint.pxd @@ -47,3 +47,5 @@ cdef u32_parse_bool_flag(uint32_t flags, flag) cdef u32_set_bool_flag(uint32_t *flags, boolean, true_flag, false_flag=*) cdef u16_parse_bool_flag(uint16_t flags, flag) cdef u16_set_bool_flag(uint16_t *flags, boolean, true_flag, false_flag=*) +cdef u8_parse_bool_flag(uint8_t flags, flag) +cdef u8_set_bool_flag(uint8_t *flags, boolean, true_flag, false_flag=*) diff --git a/pyslurm/utils/uint.pyx b/pyslurm/utils/uint.pyx index 23ddad48..1cdd2fb1 100644 --- a/pyslurm/utils/uint.pyx +++ b/pyslurm/utils/uint.pyx @@ -172,6 +172,10 @@ cdef u16_parse_bool(uint16_t val): return uint_parse_bool(val, slurm.NO_VAL16) +cdef u8_set_bool_flag(uint8_t *flags, boolean, true_flag, false_flag=0): + flags[0] = uint_set_bool_flag(flags[0], boolean, true_flag, false_flag) + + cdef u16_set_bool_flag(uint16_t *flags, boolean, true_flag, false_flag=0): flags[0] = uint_set_bool_flag(flags[0], boolean, true_flag, false_flag) @@ -188,6 +192,10 @@ cdef u16_parse_bool_flag(uint16_t flags, flag): return uint_parse_bool_flag(flags, flag, slurm.NO_VAL16) +cdef u8_parse_bool_flag(uint8_t flags, flag): + return uint_parse_bool_flag(flags, flag, slurm.NO_VAL8) + + cdef u32_parse_bool_flag(uint32_t flags, flag): return uint_parse_bool_flag(flags, flag, slurm.NO_VAL) diff --git a/scripts/griffe_exts.py b/scripts/griffe_exts.py index 905f8358..c55b9150 100644 --- a/scripts/griffe_exts.py +++ b/scripts/griffe_exts.py @@ -30,7 +30,11 @@ SLURM_DOCS_URL_BASE = "https://slurm.schedmd.com/archive" SLURM_DOCS_URL_VERSIONED = f"{SLURM_DOCS_URL_BASE}/slurm-{SLURM_VERSION}-latest" -config_files = ["acct_gather.conf", "slurm.conf", "cgroup.conf", "mpi.conf"] +config_files = ["acct_gather.conf", + "slurm.conf", + "cgroup.conf", + "mpi.conf", + "scontrol"] def replace_with_slurm_docs_url(match): @@ -94,6 +98,18 @@ def on_instance( logger.debug(f"Object {obj.path} does not have a __doc__ attribute") return + # Hack to improve generated docs for Enums. + if hasattr(obj.parent, "bases"): + for base in obj.parent.bases: + b = base.lower() + if "enum" in b and not obj.name.startswith("_"): + v = obj.value[:-1].split(" ")[-1] + obj.value = v + obj.labels = {} + + if "slurmflag" in b: + obj.value = None + if not docstring or not obj.docstring: return diff --git a/tests/integration/test_reservation.py b/tests/integration/test_reservation.py new file mode 100644 index 00000000..e1ee3482 --- /dev/null +++ b/tests/integration/test_reservation.py @@ -0,0 +1,95 @@ +######################################################################### +# test_reservation.py - reservation integration tests +######################################################################### +# Copyright (C) 2025 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_reservation.py - integration test reservation functionalities.""" + +import pyslurm +from pyslurm import ReservationFlags, ReservationReoccurrence +from datetime import datetime + + +def test_api_calls(): + start = datetime.now() + duration = "1-00:00:00" + resv = pyslurm.Reservation( + name="testing", + start_time=start, + duration=duration, + users=["root"], + node_count=1, + reoccurrence="DAILY" + ) + resv.create() + + reservations = pyslurm.Reservations.load() + resv = reservations["testing"] + assert len(reservations) == 1 + assert resv.name == "testing" + assert resv.to_dict() + + assert resv.start_time == int(start.timestamp()) + assert resv.duration == 60 * 24 + assert resv.end_time == resv.start_time + (60 * 60 * 24) + + resv.duration += 60 * 24 + resv.modify() + + resv = pyslurm.Reservation.load("testing") + assert resv.name == "testing" + assert resv.start_time == int(start.timestamp()) + assert resv.duration == 2 * 60 * 24 + assert resv.end_time == resv.start_time + (2 * 60 * 60 * 24) + + assert resv.reoccurrence == ReservationReoccurrence.DAILY + assert resv.reoccurrence == "DAILY" + # Can only remove this once the Reservation exists. Setting another + # reoccurrence doesn't work, probably a bug in slurmctld..., because it + # makes no sense why that shouldn't work. + resv.reoccurrence = ReservationReoccurrence.NO + resv.modify() + + resv = pyslurm.Reservation.load("testing") + assert resv.reoccurrence == "NO" + + resv.flags = ReservationFlags.MAINTENANCE | ReservationFlags.FLEX + resv.modify() + + resv = pyslurm.Reservation.load("testing") + assert resv.flags == ReservationFlags.MAINTENANCE | ReservationFlags.FLEX + + assert ReservationFlags.PURGE not in resv.flags + resv.purge_time = "2-00:00:00" + resv.modify() + + resv = pyslurm.Reservation.load("testing") + assert ReservationFlags.PURGE in resv.flags + assert resv.purge_time == 2 * 60 * 60 * 24 + + resv.purge_time = "3-00:00:00" + resv.modify() + + resv = pyslurm.Reservation.load("testing") + assert ReservationFlags.PURGE in resv.flags + assert resv.purge_time == 3 * 60 * 60 * 24 + + assert resv.to_dict() + resv.delete() + reservations = pyslurm.Reservations.load() + assert len(reservations) == 0 diff --git a/tests/unit/test_reservation.py b/tests/unit/test_reservation.py new file mode 100644 index 00000000..ec6eafde --- /dev/null +++ b/tests/unit/test_reservation.py @@ -0,0 +1,80 @@ +######################################################################### +# test_reservation.py - reservation unit tests +######################################################################### +# Copyright (C) 2025 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_reservation.py - Unit test basic reservation functionalities.""" + +import pyslurm +from pyslurm import ReservationFlags, ReservationReoccurrence +from datetime import datetime + + +def test_create_instance(): + resv = pyslurm.Reservation("test") + assert resv.name == "test" + assert resv.accounts == [] + assert resv.start_time == None + assert resv.end_time == None + assert resv.duration == 0 + assert resv.is_active is False + assert resv.cpu_ids_by_node == {} + assert resv.to_dict() + + start = datetime.now() + resv.start_time = start + resv.duration = "1-00:00:00" + + assert resv.start_time == int(start.timestamp()) + assert resv.duration == 60 * 24 + assert resv.end_time == resv.start_time + (60 * 60 * 24) + + resv.duration += pyslurm.utils.timestr_to_mins("1-00:00:00") + + assert resv.start_time == int(start.timestamp()) + assert resv.duration == 2 * 60 * 24 + assert resv.end_time == resv.start_time + (2 * 60 * 60 * 24) + + start = datetime.fromisoformat("2022-04-03T06:00:00") + end = resv.end_time + resv.start_time = int(start.timestamp()) + + assert resv.start_time == int(start.timestamp()) + assert resv.end_time == end + assert resv.duration == int((resv.end_time - resv.start_time) / 60) + + duration = resv.duration + resv.end_time += 60 * 60 * 24 + assert resv.start_time == int(start.timestamp()) + assert resv.end_time == end + (60 * 60 * 24) + assert resv.duration == duration + (60 * 24) + + assert resv.reoccurrence == ReservationReoccurrence.NO + assert resv.reoccurrence == "NO" + resv.reoccurrence = ReservationReoccurrence.DAILY + + assert resv.flags == resv.flags.__class__(0) + resv.flags = ReservationFlags.MAINTENANCE | ReservationFlags.FLEX + assert resv.flags == ReservationFlags.MAINTENANCE | ReservationFlags.FLEX + resv.flags |= ReservationFlags.MAGNETIC + assert resv.flags == ReservationFlags.MAINTENANCE | ReservationFlags.FLEX | ReservationFlags.MAGNETIC + + resv.flags = ["FLEX", "PURGE"] + assert resv.flags == ReservationFlags.FLEX | ReservationFlags.PURGE + +