From b030b83925aa3aeb70467f9e01833957db87121e Mon Sep 17 00:00:00 2001 From: abbas1902 Date: Mon, 27 Jan 2025 19:29:47 +0000 Subject: [PATCH] Add provisioningModel for calendar future reservations --- .../modules/slurm_files/scripts/resume.py | 8 ++++---- .../modules/slurm_files/scripts/util.py | 5 +++++ 2 files changed, 9 insertions(+), 4 deletions(-) diff --git a/community/modules/scheduler/schedmd-slurm-gcp-v6-controller/modules/slurm_files/scripts/resume.py b/community/modules/scheduler/schedmd-slurm-gcp-v6-controller/modules/slurm_files/scripts/resume.py index 95e85aa8df..f352916cea 100755 --- a/community/modules/scheduler/schedmd-slurm-gcp-v6-controller/modules/slurm_files/scripts/resume.py +++ b/community/modules/scheduler/schedmd-slurm-gcp-v6-controller/modules/slurm_files/scripts/resume.py @@ -101,11 +101,11 @@ def instance_properties(nodeset:object, model:str, placement_group:Optional[str] props.resourcePolicies = [placement_group] if reservation := lookup().nodeset_reservation(nodeset): - update_reservation_props(reservation, props, placement_group) + update_reservation_props(reservation, props, placement_group, calendar_mode=False) if (fr := lookup().future_reservation(nodeset)) and fr.specific: assert fr.active_reservation - update_reservation_props(fr.active_reservation, props, placement_group) + update_reservation_props(fr.active_reservation, props, placement_group, fr.calendar) if props.resourcePolicies: props.scheduling.onHostMaintenance = "TERMINATE" @@ -120,14 +120,14 @@ def instance_properties(nodeset:object, model:str, placement_group:Optional[str] props.update(nodeset.get("instance_properties") or {}) return props -def update_reservation_props(reservation:ReservationDetails, props:NSDict, placement_group:Optional[str]) -> None: +def update_reservation_props(reservation:ReservationDetails, props:NSDict, placement_group:Optional[str], calendar_mode:bool) -> None: props.reservationAffinity = { "consumeReservationType": "SPECIFIC_RESERVATION", "key": f"compute.{util.universe_domain()}/reservation-name", "values": [reservation.bulk_insert_name], } - if reservation.dense: + if reservation.dense or calendar_mode: props.scheduling.provisioningModel = "RESERVATION_BOUND" # Figure out `resourcePolicies` diff --git a/community/modules/scheduler/schedmd-slurm-gcp-v6-controller/modules/slurm_files/scripts/util.py b/community/modules/scheduler/schedmd-slurm-gcp-v6-controller/modules/slurm_files/scripts/util.py index 48fc072ff1..7aa73f599a 100755 --- a/community/modules/scheduler/schedmd-slurm-gcp-v6-controller/modules/slurm_files/scripts/util.py +++ b/community/modules/scheduler/schedmd-slurm-gcp-v6-controller/modules/slurm_files/scripts/util.py @@ -1313,8 +1313,12 @@ class FutureReservation: specific: bool start_time: datetime end_time: datetime + reservation_mode: Optional[str] active_reservation: Optional[ReservationDetails] + @property + def calendar(self) -> bool: + return self.reservation_mode == "CALENDAR" @dataclass class Job: @@ -1681,6 +1685,7 @@ def future_reservation(self, nodeset:object) -> Optional[FutureReservation]: specific=fr["specificReservationRequired"], start_time=start_time, end_time=end_time, + reservation_mode=fr.get("reservationMode"), active_reservation=active_reservation )