From d3dabf26410fd0da83c978258081c73b4792db9b Mon Sep 17 00:00:00 2001 From: Sudheendra Gopinath Date: Sat, 4 Jul 2026 14:33:49 +0000 Subject: [PATCH 1/5] feat(reservations): persist CRUD via raft WAL and enforce scheduling parity Route reservation create/update/delete through WalOperation so changes survive log replay, not just periodic snapshots. Add MAINT, IGNORE_JOBS, and NO_HOLD_JOBS flags with overlap validation, prospective backfill blocking, lifecycle enforcement (purge, end-time cancel, hold on delete), and scontrol/sinfo display updates including maint-fenced pending reasons. --- crates/spur-cli/src/scontrol.rs | 22 ++ crates/spur-cli/src/sinfo.rs | 3 + crates/spur-cli/src/smd.rs | 3 + crates/spur-core/src/job.rs | 4 + crates/spur-core/src/reservation.rs | 233 +++++++++++++- crates/spur-core/src/wal.rs | 66 ++++ crates/spur-sched/src/backfill.rs | 102 ++++++- crates/spur-tests/src/t07_sched.rs | 2 + crates/spur-tests/src/t50_core.rs | 1 + crates/spur-tests/src/t55_format.rs | 5 + crates/spurctld/src/cluster.rs | 339 +++++++++++++++++++-- crates/spurctld/src/scheduler_loop.rs | 2 + crates/spurctld/src/server.rs | 14 +- proto/slurm.proto | 4 + tests/native_host/e2e/test_reservations.py | 118 +++++++ 15 files changed, 879 insertions(+), 39 deletions(-) create mode 100644 tests/native_host/e2e/test_reservations.py diff --git a/crates/spur-cli/src/scontrol.rs b/crates/spur-cli/src/scontrol.rs index 9b35c05a..53326fc6 100644 --- a/crates/spur-cli/src/scontrol.rs +++ b/crates/spur-cli/src/scontrol.rs @@ -87,6 +87,9 @@ pub enum ScontrolCommand { /// Comma-separated users (optional) #[arg(long, default_value = "")] users: String, + /// Comma-separated flags (maint, ignore_jobs, no_hold_jobs) + #[arg(long, default_value = "")] + flags: String, }, /// Update a reservation #[command(name = "update-reservation")] @@ -218,6 +221,7 @@ pub async fn main_with_args(args: Vec) -> Result<()> { nodes, accounts, users, + flags, } => { create_reservation( &args.controller, @@ -227,6 +231,7 @@ pub async fn main_with_args(args: Vec) -> Result<()> { &nodes, &accounts, &users, + &flags, ) .await } @@ -371,6 +376,9 @@ async fn show(controller: &str, entity: &str, name: Option<&str>) -> Result<()> label_str.sort(); println!(" Labels={}", label_str.join(",")); } + if !node.active_reservation.is_empty() { + println!(" ActiveReservation={}", node.active_reservation); + } println!(" CpuLoad={}", node.cpu_load as f64 / 100.0); println!(); } @@ -421,6 +429,12 @@ async fn show(controller: &str, entity: &str, name: Option<&str>) -> Result<()> println!(" StartTime={}", res.start_time); println!(" EndTime={}", res.end_time); println!(" Nodes={}", res.nodes); + if !res.state.is_empty() { + println!(" State={}", res.state); + } + if !res.flags.is_empty() { + println!(" Flags={}", res.flags); + } if !res.accounts.is_empty() { println!(" Accounts={}", res.accounts); } @@ -657,6 +671,7 @@ async fn update_node( } /// Create a reservation via the controller. +#[allow(clippy::too_many_arguments)] async fn create_reservation( controller: &str, name: &str, @@ -665,6 +680,7 @@ async fn create_reservation( nodes: &str, accounts: &str, users: &str, + flags: &str, ) -> Result<()> { let mut client = SlurmControllerClient::connect(controller.to_string()) .await @@ -685,6 +701,11 @@ async fn create_reservation( .map(|s| s.trim().to_string()) .filter(|s| !s.is_empty()) .collect(); + let flag_list: Vec = flags + .split(',') + .map(|s| s.trim().to_string()) + .filter(|s| !s.is_empty()) + .collect(); client .create_reservation(spur_proto::proto::CreateReservationRequest { @@ -694,6 +715,7 @@ async fn create_reservation( nodes: node_list, accounts: account_list, users: user_list, + flags: flag_list, }) .await .context("failed to create reservation")?; diff --git a/crates/spur-cli/src/sinfo.rs b/crates/spur-cli/src/sinfo.rs index 0d28e297..ef73f92a 100644 --- a/crates/spur-cli/src/sinfo.rs +++ b/crates/spur-cli/src/sinfo.rs @@ -288,6 +288,9 @@ fn effective_state_str(node: &NodeInfo) -> String { if !node.active_reservation.is_empty() && node.state == spur_proto::proto::NodeState::NodeIdle as i32 { + if node.reservation_maint { + return "maint".into(); + } return "resv".into(); } spur_core::node::NodeState::from_proto_i32(node.state) diff --git a/crates/spur-cli/src/smd.rs b/crates/spur-cli/src/smd.rs index e217414a..09453716 100644 --- a/crates/spur-cli/src/smd.rs +++ b/crates/spur-cli/src/smd.rs @@ -109,6 +109,9 @@ pub async fn main_with_args(args: Vec) -> Result<()> { fn node_state_str(node: &NodeInfo) -> &'static str { if !node.active_reservation.is_empty() && node.state == NodeState::NodeIdle as i32 { + if node.reservation_maint { + return "maint"; + } return "resv"; } spur_core::node::NodeState::from_proto_i32(node.state) diff --git a/crates/spur-core/src/job.rs b/crates/spur-core/src/job.rs index 964ce8c1..78656b35 100644 --- a/crates/spur-core/src/job.rs +++ b/crates/spur-core/src/job.rs @@ -251,6 +251,8 @@ pub enum PendingReason { QosGrpNodeLimit, BurstBufferResources, BurstBufferStageIn, + ReservedMaintenance, + ReservationDeleted, } impl PendingReason { @@ -294,6 +296,8 @@ impl PendingReason { Self::QosGrpNodeLimit => "QOSGrpNodeLimit", Self::BurstBufferResources => "BurstBufferResources", Self::BurstBufferStageIn => "BurstBufferStageIn", + Self::ReservedMaintenance => "ReqNodeNotAvail, Reserved for maintenance", + Self::ReservationDeleted => "ReservationDeleted", } } } diff --git a/crates/spur-core/src/reservation.rs b/crates/spur-core/src/reservation.rs index 763e35f4..2cacc584 100644 --- a/crates/spur-core/src/reservation.rs +++ b/crates/spur-core/src/reservation.rs @@ -1,9 +1,68 @@ // Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -use chrono::{DateTime, Utc}; +use chrono::{DateTime, Duration, Utc}; use serde::{Deserialize, Serialize}; +use crate::hostlist; +use crate::job::{Job, JobState}; + +/// Optional reservation behavior flags (comma-separated in admin CLI). +#[derive(Debug, Clone, Copy, Default, Serialize, Deserialize, PartialEq, Eq)] +pub struct ReservationFlags { + #[serde(default)] + pub maint: bool, + #[serde(default)] + pub ignore_jobs: bool, + #[serde(default)] + pub no_hold_jobs: bool, +} + +impl ReservationFlags { + pub fn parse_list(values: &[String]) -> Self { + let mut flags = Self::default(); + for v in values { + for part in v.split(',') { + match part.trim().to_ascii_lowercase().as_str() { + "maint" => flags.maint = true, + "ignore_jobs" => flags.ignore_jobs = true, + "no_hold_jobs" => flags.no_hold_jobs = true, + "" => {} + _ => {} + } + } + } + flags + } + + pub fn parse_csv(csv: &str) -> Self { + if csv.is_empty() { + return Self::default(); + } + Self::parse_list( + &csv + .split(',') + .map(|s| s.trim().to_string()) + .filter(|s| !s.is_empty()) + .collect::>(), + ) + } + + pub fn display_csv(&self) -> String { + let mut parts = Vec::new(); + if self.maint { + parts.push("MAINT"); + } + if self.ignore_jobs { + parts.push("IGNORE_JOBS"); + } + if self.no_hold_jobs { + parts.push("NO_HOLD_JOBS"); + } + parts.join(",") + } +} + /// A resource reservation that blocks nodes for specific users/accounts /// during a time window. #[derive(Debug, Clone, Serialize, Deserialize)] @@ -14,23 +73,40 @@ pub struct Reservation { pub nodes: Vec, pub accounts: Vec, pub users: Vec, + #[serde(default)] + pub flags: ReservationFlags, } impl Reservation { - /// Check if the reservation is active at the given time. pub fn is_active(&self, now: DateTime) -> bool { now >= self.start_time && now < self.end_time } - /// Check if the reservation covers a specific node. + pub fn is_future(&self, now: DateTime) -> bool { + now < self.start_time + } + + pub fn is_expired(&self, now: DateTime) -> bool { + now >= self.end_time + } + + pub fn state_label(&self, now: DateTime) -> &'static str { + if self.is_active(now) { + "ACTIVE" + } else if self.is_future(now) { + "INACTIVE" + } else { + "EXPIRED" + } + } + pub fn covers_node(&self, node: &str) -> bool { self.nodes.iter().any(|n| n == node) } - /// Check if a user (and optionally account) is allowed to use this reservation. pub fn allows_user(&self, user: &str, account: Option<&str>) -> bool { if self.users.is_empty() && self.accounts.is_empty() { - return true; // No restrictions + return true; } if self.users.iter().any(|u| u == user) { return true; @@ -42,12 +118,125 @@ impl Reservation { } false } + + pub fn job_targets(&self, job: &Job) -> bool { + job.spec + .reservation + .as_deref() + .is_some_and(|n| n == self.name) + } +} + +/// Expand hostlist patterns and verify each name exists in the cluster. +pub fn normalize_node_list( + patterns: &[String], + known_nodes: &std::collections::HashSet, +) -> Result, String> { + let mut out = Vec::new(); + for pattern in patterns { + let trimmed = pattern.trim(); + if trimmed.is_empty() { + continue; + } + let expanded = hostlist::expand(trimmed).map_err(|e| e.to_string())?; + for node in expanded { + if !known_nodes.contains(&node) { + return Err(format!("unknown node '{}'", node)); + } + if !out.contains(&node) { + out.push(node); + } + } + } + if out.is_empty() { + return Err("reservation requires at least one node".into()); + } + Ok(out) +} + +/// Estimated end time for a running job (start + time limit). +pub fn running_job_end(job: &Job, now: DateTime) -> Option> { + let _ = now; + if !matches!( + job.state, + JobState::Running | JobState::Completing | JobState::Suspended + ) { + return None; + } + let start = job.start_time?; + let limit = job.spec.time_limit?; + Some(start + limit) +} + +/// Returns the first running job that would still occupy `nodes` at `start_time`. +pub fn running_jobs_overlap_start( + jobs: &std::collections::HashMap, + nodes: &[String], + start_time: DateTime, + except_reservation: Option<&str>, +) -> Option<(crate::job::JobId, String)> { + let node_set: std::collections::HashSet<&str> = nodes.iter().map(String::as_str).collect(); + for job in jobs.values() { + if !matches!( + job.state, + JobState::Running | JobState::Completing | JobState::Suspended + ) { + continue; + } + if let Some(ex) = except_reservation { + if job.spec.reservation.as_deref() == Some(ex) { + continue; + } + } + let overlaps_node = job.allocated_nodes.iter().any(|n| node_set.contains(n.as_str())); + if !overlaps_node { + continue; + } + let end = running_job_end(job, Utc::now()).unwrap_or(start_time); + if end > start_time { + let node = job + .allocated_nodes + .iter() + .find(|n| node_set.contains(n.as_str())) + .cloned() + .unwrap_or_default(); + return Some((job.job_id, node)); + } + } + None +} + +/// Whether a job without reservation access would overlap this reservation window. +pub fn prospective_overlap( + job: &Job, + reservation: &Reservation, + node: &str, + now: DateTime, + job_duration: Duration, +) -> bool { + if !reservation.covers_node(node) { + return false; + } + let res_name = job.spec.reservation.as_deref().filter(|s| !s.is_empty()); + if res_name == Some(reservation.name.as_str()) + && reservation.allows_user(&job.spec.user, job.spec.account.as_deref()) + { + return false; + } + let job_end = now + job_duration; + if reservation.is_active(now) { + return job_end > reservation.end_time || now < reservation.start_time; + } + if reservation.is_future(now) { + return job_end > reservation.start_time; + } + false } #[cfg(test)] mod tests { use super::*; - use chrono::Duration; + use crate::job::JobSpec; fn make_reservation() -> Reservation { let now = Utc::now(); @@ -58,6 +247,7 @@ mod tests { nodes: vec!["node001".into(), "node002".into()], accounts: vec!["research".into()], users: vec!["alice".into()], + flags: ReservationFlags::default(), } } @@ -96,7 +286,38 @@ mod tests { nodes: vec!["node001".into()], accounts: Vec::new(), users: Vec::new(), + flags: ReservationFlags::default(), }; assert!(res.allows_user("anyone", None)); } + + #[test] + fn flags_parse_csv() { + let f = ReservationFlags::parse_csv("maint,ignore_jobs"); + assert!(f.maint); + assert!(f.ignore_jobs); + assert!(!f.no_hold_jobs); + } + + #[test] + fn prospective_overlap_blocks_long_job_before_reservation() { + let now = Utc::now(); + let res = Reservation { + name: "r1".into(), + start_time: now + Duration::hours(1), + end_time: now + Duration::hours(2), + nodes: vec!["node001".into()], + accounts: Vec::new(), + users: vec!["alice".into()], + flags: ReservationFlags::default(), + }; + let job = Job::new(1, JobSpec::default()); + assert!(prospective_overlap( + &job, + &res, + "node001", + now, + Duration::hours(3) + )); + } } diff --git a/crates/spur-core/src/wal.rs b/crates/spur-core/src/wal.rs index 085665fd..e9a4d968 100644 --- a/crates/spur-core/src/wal.rs +++ b/crates/spur-core/src/wal.rs @@ -6,6 +6,7 @@ use serde::{Deserialize, Serialize}; use crate::admission::AdmissionToken; use crate::job::{JobId, JobSpec, JobState}; use crate::node::NodeState; +use crate::reservation::Reservation; use std::collections::HashMap; use crate::resource::{ResourceAllocations, ResourceSet}; @@ -118,6 +119,71 @@ pub enum WalOperation { TokenRevoke { token_id: String, }, + + ReservationCreate { + reservation: Reservation, + }, + ReservationUpdate { + name: String, + duration_minutes: u32, + add_nodes: Vec, + remove_nodes: Vec, + add_users: Vec, + remove_users: Vec, + add_accounts: Vec, + remove_accounts: Vec, + }, + ReservationDelete { + name: String, + }, +} + +#[cfg(test)] +mod reservation_wal_tests { + use super::*; + use crate::reservation::{Reservation, ReservationFlags}; + use chrono::Utc; + + #[test] + fn reservation_create_round_trips() { + let now = Utc::now(); + let op = WalOperation::ReservationCreate { + reservation: Reservation { + name: "r1".into(), + start_time: now, + end_time: now + chrono::Duration::hours(1), + nodes: vec!["n1".into()], + accounts: Vec::new(), + users: vec!["alice".into()], + flags: ReservationFlags { + maint: true, + ..Default::default() + }, + }, + }; + let json = serde_json::to_string(&op).unwrap(); + let back: WalOperation = serde_json::from_str(&json).unwrap(); + match back { + WalOperation::ReservationCreate { reservation } => { + assert_eq!(reservation.name, "r1"); + assert!(reservation.flags.maint); + } + _ => panic!("wrong variant"), + } + } + + #[test] + fn reservation_delete_round_trips() { + let op = WalOperation::ReservationDelete { + name: "r1".into(), + }; + let json = serde_json::to_string(&op).unwrap(); + let back: WalOperation = serde_json::from_str(&json).unwrap(); + match back { + WalOperation::ReservationDelete { name } => assert_eq!(name, "r1"), + _ => panic!("wrong variant"), + } + } } #[cfg(test)] diff --git a/crates/spur-sched/src/backfill.rs b/crates/spur-sched/src/backfill.rs index 7834b673..ad0b658f 100644 --- a/crates/spur-sched/src/backfill.rs +++ b/crates/spur-sched/src/backfill.rs @@ -8,7 +8,7 @@ use tracing::debug; use spur_core::job::{Job, JobId}; use spur_core::node::Node; -use spur_core::reservation::Reservation; +use spur_core::reservation::{self, Reservation}; use spur_core::resource::{build_exclusive_allocation, build_node_allocation, ResourceSet}; use crate::timeline::NodeTimeline; @@ -45,6 +45,33 @@ impl BackfillScheduler { .collect(); } + /// Returns true when `[start, start+duration)` intersects a reservation on `node` + /// and the job is not authorized for that reservation. + fn start_overlaps_reservation( + job: &Job, + node: &str, + reservations: &[Reservation], + start: chrono::DateTime, + duration: chrono::Duration, + ) -> bool { + let end = start + duration; + for res in reservations { + if !res.covers_node(node) { + continue; + } + let res_name = job.spec.reservation.as_deref().filter(|s| !s.is_empty()); + if res_name == Some(res.name.as_str()) + && res.allows_user(&job.spec.user, job.spec.account.as_deref()) + { + continue; + } + if end > res.start_time && start < res.end_time { + return true; + } + } + false + } + /// Find nodes that satisfy a job's resource requirements. fn find_suitable_nodes( &self, @@ -126,6 +153,7 @@ impl BackfillScheduler { } // Reservation enforcement let now = Utc::now(); + let job_duration = job.spec.time_limit.unwrap_or(Duration::hours(1)); let job_reservation = job.spec.reservation.as_deref().filter(|s| !s.is_empty()); let active_reservations: Vec<&Reservation> = reservations @@ -134,11 +162,9 @@ impl BackfillScheduler { .collect(); if let Some(res_name) = job_reservation { - // Job targets a reservation -- only allow nodes in that reservation if !active_reservations.iter().any(|r| r.name == res_name) { return false; } - // Check user/account is allowed let user = &job.spec.user; let account = job.spec.account.as_deref(); if !active_reservations @@ -147,11 +173,27 @@ impl BackfillScheduler { { return false; } + if let Some(res) = reservations.iter().find(|r| r.name == res_name) { + let job_end = now + job_duration; + if now < res.start_time || job_end > res.end_time { + return false; + } + } } else { - // Job does NOT target a reservation -- skip reserved nodes if !active_reservations.is_empty() { return false; } + for res in reservations { + if reservation::prospective_overlap( + job, + res, + &node.name, + now, + job_duration, + ) { + return false; + } + } } // Check resource capacity (total, not current available) @@ -249,6 +291,16 @@ impl Scheduler for BackfillScheduler { }) .collect(); + node_starts.retain(|(ni, start)| { + !Self::start_overlaps_reservation( + job, + &cluster.nodes[*ni].name, + cluster.reservations, + *start, + duration, + ) + }); + // For --spread-job, sort by least-loaded (ascending alloc) so we // prefer nodes with the most available resources. For normal jobs, // sort by earliest start time. For weighted nodes, prefer higher @@ -1064,10 +1116,11 @@ mod tests { Reservation { name: name.into(), start_time: now - Duration::hours(1), - end_time: now + Duration::hours(1), + end_time: now + Duration::hours(2), nodes, accounts: Vec::new(), users, + flags: Default::default(), } } @@ -1187,6 +1240,7 @@ mod tests { nodes: vec!["node001".into()], accounts: Vec::new(), users: vec!["alice".into()], + flags: Default::default(), }; // Non-reservation job should be able to use node001 since reservation is inactive @@ -1205,6 +1259,44 @@ mod tests { // (scheduler picks by weight/time, node001 is a valid target) } + #[test] + fn test_prospective_reservation_blocks_long_job() { + let mut sched = BackfillScheduler::new(100); + let nodes = make_nodes(1); + let partitions = vec![Partition { + name: "default".into(), + ..Default::default() + }]; + + let now = Utc::now(); + let future_reservation = Reservation { + name: "upcoming".into(), + start_time: now + Duration::minutes(30), + end_time: now + Duration::hours(2), + nodes: vec!["node001".into()], + accounts: Vec::new(), + users: vec!["alice".into()], + flags: Default::default(), + }; + + let mut job = make_job(1, 1, 1); + job.spec.time_limit = Some(Duration::hours(2)); + + let pending = vec![job]; + let cluster = ClusterState { + nodes: &nodes, + partitions: &partitions, + reservations: &[future_reservation], + topology: None, + }; + + let assignments = sched.schedule(&pending, &cluster); + assert!( + assignments.is_empty(), + "job that would overlap upcoming reservation must not schedule" + ); + } + #[test] fn test_job_resource_request_includes_countable_gres() { let job = Job::new( diff --git a/crates/spur-tests/src/t07_sched.rs b/crates/spur-tests/src/t07_sched.rs index 68f32b83..c37a6e28 100644 --- a/crates/spur-tests/src/t07_sched.rs +++ b/crates/spur-tests/src/t07_sched.rs @@ -552,6 +552,7 @@ mod tests { nodes: vec!["node001".into()], users: vec!["alice".into()], accounts: Vec::new(), + flags: Default::default(), }; // A job with no reservation spec should NOT land on node001. @@ -585,6 +586,7 @@ mod tests { nodes: vec!["node001".into()], users: vec!["bob".into()], accounts: Vec::new(), + flags: Default::default(), }; let mut job = make_job_with_resources("reserved-job", 1, 1, 1, Some(60)); diff --git a/crates/spur-tests/src/t50_core.rs b/crates/spur-tests/src/t50_core.rs index 60de1ec7..e18f9dfa 100644 --- a/crates/spur-tests/src/t50_core.rs +++ b/crates/spur-tests/src/t50_core.rs @@ -784,6 +784,7 @@ mod tests { nodes: vec!["node001".into()], users: vec!["alice".into()], accounts: Vec::new(), + flags: Default::default(), }; // All fields must be mutable for an update to work. res.end_time = now + chrono::Duration::hours(8); diff --git a/crates/spur-tests/src/t55_format.rs b/crates/spur-tests/src/t55_format.rs index d945d2de..020deb05 100644 --- a/crates/spur-tests/src/t55_format.rs +++ b/crates/spur-tests/src/t55_format.rs @@ -121,6 +121,11 @@ mod tests { (PendingReason::Held, "JobHeldUser"), // Reason-code vocabulary parity (Slurm 25.11.6 job_reason_string). (PendingReason::Reservation, "Reservation"), + ( + PendingReason::ReservedMaintenance, + "ReqNodeNotAvail, Reserved for maintenance", + ), + (PendingReason::ReservationDeleted, "ReservationDeleted"), (PendingReason::QosMaxCpuPerJobLimit, "QOSMaxCpuPerJobLimit"), ( PendingReason::QosMaxWallDurationPerJobLimit, diff --git a/crates/spurctld/src/cluster.rs b/crates/spurctld/src/cluster.rs index e152e558..d791312e 100644 --- a/crates/spurctld/src/cluster.rs +++ b/crates/spurctld/src/cluster.rs @@ -18,7 +18,9 @@ use spur_core::job::{Job, JobId, JobSpec, JobState, NodeCompleteError, PendingRe use spur_core::node::{Node, NodeEvent, NodeSource, NodeState}; use spur_core::partition::Partition; use spur_core::qos::{check_qos_limits, QosCheckResult}; -use spur_core::reservation::Reservation; +use spur_core::reservation::{ + self, normalize_node_list, running_jobs_overlap_start, Reservation, +}; use spur_core::resource::{ResourceAllocations, ResourceSet}; use spur_core::step::{JobStep, StepState, STEP_BATCH, STEP_RESERVED_MIN}; use spur_core::wal::WalOperation; @@ -1787,18 +1789,26 @@ impl ClusterManager { } } - /// Create a new reservation. - pub fn create_reservation(&self, res: Reservation) -> anyhow::Result<()> { - let mut reservations = self.reservations.write(); - if reservations.iter().any(|r| r.name == res.name) { + /// Create a new reservation (validated, persisted via Raft). + pub fn create_reservation(&self, mut res: Reservation) -> anyhow::Result<()> { + if self + .reservations + .read() + .iter() + .any(|r| r.name == res.name) + { anyhow::bail!("reservation '{}' already exists", res.name); } - info!(name = %res.name, "reservation created"); - reservations.push(res); + let known: std::collections::HashSet = + self.nodes.read().keys().cloned().collect(); + res.nodes = normalize_node_list(&res.nodes, &known) + .map_err(|e| anyhow::anyhow!("{e}"))?; + self.validate_reservation_job_overlap(&res, None)?; + self.propose(WalOperation::ReservationCreate { reservation: res })?; Ok(()) } - /// Update an existing reservation. + /// Update an existing reservation (validated, persisted via Raft). #[allow(clippy::too_many_arguments)] pub fn update_reservation( &self, @@ -1811,12 +1821,219 @@ impl ClusterManager { add_accounts: &[String], remove_accounts: &[String], ) -> anyhow::Result<()> { - let mut reservations = self.reservations.write(); - let res = reservations - .iter_mut() + let mut preview = self + .reservations + .read() + .iter() + .find(|r| r.name == name) + .cloned() + .ok_or_else(|| anyhow::anyhow!("reservation '{}' not found", name))?; + + if duration_minutes > 0 { + preview.end_time = + preview.start_time + chrono::Duration::minutes(duration_minutes as i64); + } + let known: std::collections::HashSet = + self.nodes.read().keys().cloned().collect(); + let mut add_expanded = Vec::new(); + for n in add_nodes { + add_expanded.extend(normalize_node_list(std::slice::from_ref(n), &known).map_err( + |e| anyhow::anyhow!("{e}"), + )?); + } + for node in &add_expanded { + if !preview.nodes.contains(node) { + preview.nodes.push(node.clone()); + } + } + preview.nodes.retain(|n| !remove_nodes.contains(n)); + for user in add_users { + if !preview.users.contains(user) { + preview.users.push(user.clone()); + } + } + preview.users.retain(|u| !remove_users.contains(u)); + for account in add_accounts { + if !preview.accounts.contains(account) { + preview.accounts.push(account.clone()); + } + } + preview.accounts.retain(|a| !remove_accounts.contains(a)); + + self.validate_reservation_job_overlap(&preview, Some(name))?; + + self.propose(WalOperation::ReservationUpdate { + name: name.to_string(), + duration_minutes, + add_nodes: add_expanded, + remove_nodes: remove_nodes.to_vec(), + add_users: add_users.to_vec(), + remove_users: remove_users.to_vec(), + add_accounts: add_accounts.to_vec(), + remove_accounts: remove_accounts.to_vec(), + })?; + Ok(()) + } + + /// Delete a reservation by name (persisted via Raft). + pub fn delete_reservation(&self, name: &str) -> anyhow::Result<()> { + let res = self + .reservations + .read() + .iter() .find(|r| r.name == name) + .cloned() .ok_or_else(|| anyhow::anyhow!("reservation '{}' not found", name))?; + for job in self.jobs.read().values() { + if !matches!( + job.state, + JobState::Running | JobState::Completing | JobState::Suspended + ) { + continue; + } + if job.spec.reservation.as_deref() == Some(name) { + anyhow::bail!( + "reservation '{}' in use by running job {}", + name, + job.job_id + ); + } + } + + self.propose(WalOperation::ReservationDelete { + name: name.to_string(), + })?; + + if !res.flags.no_hold_jobs { + self.hold_jobs_for_deleted_reservation(name); + } + Ok(()) + } + + /// Remove reservations past their end time when no jobs still reference them. + pub fn purge_expired_reservations(&self) { + let now = Utc::now(); + let expired: Vec = self + .reservations + .read() + .iter() + .filter(|r| r.is_expired(now)) + .map(|r| r.name.clone()) + .collect(); + for name in expired { + let in_use = self.jobs.read().values().any(|job| { + matches!( + job.state, + JobState::Running | JobState::Completing | JobState::Suspended + ) && job.spec.reservation.as_deref() == Some(name.as_str()) + }); + if in_use { + continue; + } + if let Err(e) = self.propose(WalOperation::ReservationDelete { name: name.clone() }) { + warn!(name = %name, error = %e, "failed to purge expired reservation"); + } + } + } + + /// Cancel running jobs whose reservation window has ended. + pub fn enforce_reservation_end_times(&self) { + let now = Utc::now(); + let reservations: std::collections::HashMap = self + .get_reservations() + .into_iter() + .map(|r| (r.name.clone(), r)) + .collect(); + let to_cancel: Vec = self + .jobs + .read() + .values() + .filter_map(|job| { + if !matches!( + job.state, + JobState::Running | JobState::Completing | JobState::Suspended + ) { + return None; + } + let res_name = job.spec.reservation.as_ref()?; + let res = reservations.get(res_name)?; + if res.is_expired(now) { + Some(job.job_id) + } else { + None + } + }) + .collect(); + for job_id in to_cancel { + if let Err(e) = self.complete_job(job_id, -1, JobState::Cancelled) { + warn!(job_id, error = %e, "failed to cancel job after reservation ended"); + } + } + } + + fn validate_reservation_job_overlap( + &self, + res: &Reservation, + except_name: Option<&str>, + ) -> anyhow::Result<()> { + if res.flags.ignore_jobs { + return Ok(()); + } + let jobs = self.jobs.read(); + if let Some((job_id, node)) = running_jobs_overlap_start( + &jobs, + &res.nodes, + res.start_time, + except_name, + ) { + anyhow::bail!( + "requested nodes are busy (job {} on {} until after reservation start)", + job_id, + node + ); + } + Ok(()) + } + + fn hold_jobs_for_deleted_reservation(&self, name: &str) { + let pending: Vec = self + .jobs + .read() + .values() + .filter(|j| { + j.state == JobState::Pending + && j.spec.reservation.as_deref() == Some(name) + && j.pending_reason != PendingReason::Held + }) + .map(|j| j.job_id) + .collect(); + for job_id in pending { + if let Err(e) = self.hold_job(job_id) { + warn!(job_id, error = %e, "failed to hold job after reservation delete"); + continue; + } + if let Some(job) = self.jobs.write().get_mut(&job_id) { + job.pending_reason = PendingReason::ReservationDeleted; + } + } + } + + #[allow(clippy::too_many_arguments)] + fn apply_reservation_update_locked( + reservations: &mut [Reservation], + name: &str, + duration_minutes: u32, + add_nodes: &[String], + remove_nodes: &[String], + add_users: &[String], + remove_users: &[String], + add_accounts: &[String], + remove_accounts: &[String], + ) { + let Some(res) = reservations.iter_mut().find(|r| r.name == name) else { + return; + }; if duration_minutes > 0 { res.end_time = res.start_time + chrono::Duration::minutes(duration_minutes as i64); } @@ -1838,21 +2055,6 @@ impl ClusterManager { } } res.accounts.retain(|a| !remove_accounts.contains(a)); - - info!(name, "reservation updated"); - Ok(()) - } - - /// Delete a reservation by name. - pub fn delete_reservation(&self, name: &str) -> anyhow::Result<()> { - let mut reservations = self.reservations.write(); - let len_before = reservations.len(); - reservations.retain(|r| r.name != name); - if reservations.len() == len_before { - anyhow::bail!("reservation '{}' not found", name); - } - info!(name, "reservation deleted"); - Ok(()) } /// Get all reservations. @@ -1897,6 +2099,11 @@ impl ClusterManager { continue; } + if let Some(reason) = reservation_fence_reason(job, cluster_state) { + job_entry.pending_reason = reason; + continue; + } + // Determine the correct reason let partition_name = job.spec.partition.as_deref(); @@ -2727,6 +2934,42 @@ impl ClusterManager { t.revoked = true; } } + WalOperation::ReservationCreate { reservation } => { + self.reservations.write().push(reservation.clone()); + info!(name = %reservation.name, "reservation created"); + } + WalOperation::ReservationUpdate { + name, + duration_minutes, + add_nodes, + remove_nodes, + add_users, + remove_users, + add_accounts, + remove_accounts, + } => { + let mut reservations = self.reservations.write(); + Self::apply_reservation_update_locked( + reservations.as_mut(), + name, + *duration_minutes, + add_nodes, + remove_nodes, + add_users, + remove_users, + add_accounts, + remove_accounts, + ); + info!(name, "reservation updated"); + } + WalOperation::ReservationDelete { name } => { + let mut reservations = self.reservations.write(); + let len_before = reservations.len(); + reservations.retain(|r| r.name != *name); + if reservations.len() < len_before { + info!(name, "reservation deleted"); + } + } } self.next_job_id.store(next_id, Ordering::Relaxed); response @@ -2853,6 +3096,50 @@ impl StateMachineApply for ClusterManager { } } +fn reservation_fence_reason( + job: &Job, + cluster_state: &spur_sched::traits::ClusterState, +) -> Option { + let now = Utc::now(); + let duration = job + .spec + .time_limit + .unwrap_or(chrono::Duration::hours(1)); + let mut maint_block = false; + let mut any_block = false; + + for node in cluster_state.nodes { + if let Some(pname) = job.spec.partition.as_deref() { + let requested: Vec<&str> = pname.split(',').map(str::trim).collect(); + if !requested + .iter() + .any(|rp| node.partitions.iter().any(|np| np == rp)) + { + continue; + } + } + for res in cluster_state.reservations { + if !res.covers_node(&node.name) { + continue; + } + if reservation::prospective_overlap(job, res, &node.name, now, duration) { + any_block = true; + if res.flags.maint { + maint_block = true; + } + } + } + } + + if maint_block { + Some(PendingReason::ReservedMaintenance) + } else if any_block { + Some(PendingReason::ReqNodeNotAvail) + } else { + None + } +} + /// `Reservation` if the job's `--reservation` is absent/inactive/expired or /// denies it, else `None`. Shared by `pending_jobs()` (drop) and /// `tag_blocked_pending_reasons()` (displayed reason) so the two agree. diff --git a/crates/spurctld/src/scheduler_loop.rs b/crates/spurctld/src/scheduler_loop.rs index e368a5a7..1794c2b8 100644 --- a/crates/spurctld/src/scheduler_loop.rs +++ b/crates/spurctld/src/scheduler_loop.rs @@ -103,6 +103,8 @@ pub async fn run(cluster: Arc, raft: Arc) { // a follow-up; drive_bb_stage_in() is the controller-side seam only. cluster.drive_bb_stage_in(); cluster.advance_bb_staging(); + cluster.purge_expired_reservations(); + cluster.enforce_reservation_end_times(); // Tag jobs pending_jobs() will drop (QoS/license/reservation/BB) with // their real reason, since they never reach update_pending_reasons(). diff --git a/crates/spurctld/src/server.rs b/crates/spurctld/src/server.rs index a2d77a95..3608b54b 100644 --- a/crates/spurctld/src/server.rs +++ b/crates/spurctld/src/server.rs @@ -1182,6 +1182,8 @@ impl SlurmController for ControllerService { let end_time = start_time + chrono::Duration::minutes(req.duration_minutes as i64); + let flags = spur_core::reservation::ReservationFlags::parse_list(&req.flags); + let reservation = spur_core::reservation::Reservation { name: req.name, start_time, @@ -1189,11 +1191,12 @@ impl SlurmController for ControllerService { nodes: req.nodes, accounts: req.accounts, users: req.users, + flags, }; self.cluster .create_reservation(reservation) - .map_err(|e| Status::internal(e.to_string()))?; + .map_err(|e| Status::invalid_argument(e.to_string()))?; Ok(Response::new(())) } @@ -1229,7 +1232,7 @@ impl SlurmController for ControllerService { &req.add_accounts, &req.remove_accounts, ) - .map_err(|e| Status::internal(e.to_string()))?; + .map_err(|e| Status::invalid_argument(e.to_string()))?; Ok(Response::new(())) } @@ -1273,6 +1276,7 @@ impl SlurmController for ControllerService { } } let reservations = self.cluster.get_reservations(); + let now = Utc::now(); let infos: Vec = reservations .iter() .map(|r| ReservationInfo { @@ -1282,6 +1286,8 @@ impl SlurmController for ControllerService { nodes: r.nodes.join(","), accounts: r.accounts.join(","), users: r.users.join(","), + flags: r.flags.display_csv(), + state: r.state_label(now).into(), }) .collect(); Ok(Response::new(ListReservationsResponse { @@ -1783,6 +1789,7 @@ fn node_to_proto(node: &spur_core::node::Node) -> NodeInfo { switch_name: node.switch_name.clone().unwrap_or_default(), active_reservation: String::new(), labels: node.labels.clone(), + reservation_maint: false, } } @@ -1913,9 +1920,11 @@ fn annotate_nodes_with_reservations( now: DateTime, ) { for node_info in nodes.iter_mut() { + node_info.reservation_maint = false; for res in reservations { if res.is_active(now) && res.covers_node(&node_info.name) { node_info.active_reservation = res.name.clone(); + node_info.reservation_maint = res.flags.maint; break; } } @@ -1969,6 +1978,7 @@ mod tests { nodes: nodes.iter().map(|s| s.to_string()).collect(), accounts: Vec::new(), users: Vec::new(), + flags: Default::default(), } } diff --git a/proto/slurm.proto b/proto/slurm.proto index 6120e0f1..73052770 100644 --- a/proto/slurm.proto +++ b/proto/slurm.proto @@ -220,6 +220,7 @@ message NodeInfo { string active_reservation = 50; map labels = 51; + bool reservation_maint = 52; } message PartitionInfo { @@ -913,6 +914,7 @@ message CreateReservationRequest { repeated string nodes = 4; repeated string accounts = 5; repeated string users = 6; + repeated string flags = 7; } message UpdateReservationRequest { @@ -943,6 +945,8 @@ message ReservationInfo { string nodes = 4; string accounts = 5; string users = 6; + string flags = 7; + string state = 8; } // ============================================================ diff --git a/tests/native_host/e2e/test_reservations.py b/tests/native_host/e2e/test_reservations.py new file mode 100644 index 00000000..4c4e88e0 --- /dev/null +++ b/tests/native_host/e2e/test_reservations.py @@ -0,0 +1,118 @@ +# Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""E2E tests for node reservations.""" + +import time + +from cluster import parse_job_id, wait_job + + +class TestReservations: + def test_create_list_and_delete_reservation(self, cluster): + node = cluster.node_names[0] + create_out = cluster.scontrol( + "create-reservation", + f"--name=res-e2e-{int(time.time())}", + "--start-time=now", + "--duration=60", + f"--nodes={node}", + "--users=testuser", + ) + assert "created" in create_out.lower() + + show_out = cluster.scontrol("show", "reservation") + assert node in show_out + assert "ACTIVE" in show_out or "INACTIVE" in show_out + + def test_unauthorized_job_blocked_on_reserved_node(self, cluster): + res_name = f"res-block-{int(time.time())}" + node = cluster.node_names[0] + cluster.scontrol( + "create-reservation", + f"--name={res_name}", + "--start-time=now", + "--duration=30", + f"--nodes={node}", + "--users=resuser", + ) + + script = cluster.write_file("res-block.sh", "#!/bin/bash\nsleep 120\n") + sb = cluster.sbatch(["-N", "1", "-w", node, "-t", "1", script]) + job_id = parse_job_id(sb) + assert job_id is not None + + deadline = time.time() + 30 + blocked = False + while time.time() < deadline: + sq = cluster.squeue_all() + if str(job_id) in sq and "PD" in sq.split(str(job_id))[1][:40]: + blocked = True + break + time.sleep(2) + + assert blocked, f"job {job_id} should stay pending on reserved node:\n{sq}" + + def test_reservation_job_schedules_for_authorized_user(self, cluster): + res_name = f"res-auth-{int(time.time())}" + node = cluster.node_names[0] + cluster.scontrol( + "create-reservation", + f"--name={res_name}", + "--start-time=now", + "--duration=30", + f"--nodes={node}", + "--users=testuser", + ) + + script = cluster.write_file("res-auth.sh", "#!/bin/bash\necho RES_OK\n") + out_path = f"{cluster.remote_dir}/res-auth.out" + sb = cluster.sbatch( + [ + "-N", + "1", + f"--reservation={res_name}", + "-w", + node, + "-t", + "1", + "-o", + out_path, + script, + ] + ) + job_id = parse_job_id(sb) + assert job_id is not None + + state = wait_job(cluster, job_id, timeout=60) + assert state in ("CD", "GONE"), f"expected completed, got {state}" + + content = cluster.read_output_on_any_node(out_path) + assert "RES_OK" in content + + def test_create_rejects_busy_node_without_ignore_jobs(self, cluster): + node = cluster.node_names[0] + long_script = cluster.write_file("res-long.sh", "#!/bin/bash\nsleep 300\n") + sb = cluster.sbatch(["-N", "1", "-w", node, "-t", "10", long_script]) + job_id = parse_job_id(sb) + assert job_id is not None + + deadline = time.time() + 30 + while time.time() < deadline: + sq = cluster.squeue_all() + if str(job_id) in sq and " R " in sq: + break + time.sleep(2) + + res_name = f"res-busy-{int(time.time())}" + try: + cluster.scontrol( + "create-reservation", + f"--name={res_name}", + "--start-time=now", + "--duration=10", + f"--nodes={node}", + ) + assert False, "expected reservation create to fail on busy node" + except Exception as exc: + assert "busy" in str(exc).lower() or "error" in str(exc).lower() From aca1266db484c9bd6499fadd96f237ffe2ab7989 Mon Sep 17 00:00:00 2001 From: Sudheendra Gopinath Date: Sat, 4 Jul 2026 15:56:05 +0000 Subject: [PATCH 2/5] feat(reservations): add scheduling parity and preemption interaction Add OVERLAP flag with reservation overlap validation, resv_overrun_minutes grace for end-of-window job cancel, reservation-first pending job ordering, PreemptMode-aware preemption with reservation immunity unless priority_tier is higher, JobInfo.reservation proto field, and squeue %v column. --- crates/spur-cli/src/format_engine.rs | 1 + crates/spur-cli/src/squeue.rs | 1 + crates/spur-core/src/config.rs | 4 + crates/spur-core/src/reservation.rs | 113 ++++++++++++- crates/spurctld/src/cluster.rs | 228 +++++++++++++++++++++++++- crates/spurctld/src/scheduler_loop.rs | 84 +++++++--- crates/spurctld/src/server.rs | 1 + proto/slurm.proto | 1 + 8 files changed, 410 insertions(+), 23 deletions(-) diff --git a/crates/spur-cli/src/format_engine.rs b/crates/spur-cli/src/format_engine.rs index 3f646247..dd00d9f1 100644 --- a/crates/spur-cli/src/format_engine.rs +++ b/crates/spur-cli/src/format_engine.rs @@ -193,6 +193,7 @@ pub fn squeue_header(spec: char) -> &'static str { 'N' => "NODELIST", 'b' => "GRES", 'L' => "TIME_LEFT", + 'v' => "RESERVATION", _ => "?", } } diff --git a/crates/spur-cli/src/squeue.rs b/crates/spur-cli/src/squeue.rs index 420b3ce1..fec0d274 100644 --- a/crates/spur-cli/src/squeue.rs +++ b/crates/spur-cli/src/squeue.rs @@ -152,6 +152,7 @@ fn resolve_job_field(job: &spur_proto::proto::JobInfo, spec: char) -> String { 'o' => job.command.clone(), 'S' => format_timestamp(job.start_time.as_ref()), 'V' => format_timestamp(job.submit_time.as_ref()), + 'v' => job.reservation.clone(), 'e' => format_timestamp(job.end_time.as_ref()), _ => "?".into(), } diff --git a/crates/spur-core/src/config.rs b/crates/spur-core/src/config.rs index 25d27270..8b3b6469 100644 --- a/crates/spur-core/src/config.rs +++ b/crates/spur-core/src/config.rs @@ -381,6 +381,9 @@ pub struct SchedulerConfig { /// Max seconds to wait in COMPLETING before force-finishing the job. #[serde(default = "default_complete_wait")] pub complete_wait_secs: u32, + /// Grace minutes after a reservation ends before cancelling its running jobs. + #[serde(default)] + pub resv_overrun_minutes: u32, } fn default_scheduler_plugin() -> String { @@ -408,6 +411,7 @@ impl Default for SchedulerConfig { fairshare_halflife_days: 14, default_time_limit_minutes: 60, complete_wait_secs: 300, + resv_overrun_minutes: 0, } } } diff --git a/crates/spur-core/src/reservation.rs b/crates/spur-core/src/reservation.rs index 2cacc584..ef795e0e 100644 --- a/crates/spur-core/src/reservation.rs +++ b/crates/spur-core/src/reservation.rs @@ -16,6 +16,8 @@ pub struct ReservationFlags { pub ignore_jobs: bool, #[serde(default)] pub no_hold_jobs: bool, + #[serde(default)] + pub overlap: bool, } impl ReservationFlags { @@ -27,6 +29,7 @@ impl ReservationFlags { "maint" => flags.maint = true, "ignore_jobs" => flags.ignore_jobs = true, "no_hold_jobs" => flags.no_hold_jobs = true, + "overlap" => flags.overlap = true, "" => {} _ => {} } @@ -59,6 +62,9 @@ impl ReservationFlags { if self.no_hold_jobs { parts.push("NO_HOLD_JOBS"); } + if self.overlap { + parts.push("OVERLAP"); + } parts.join(",") } } @@ -127,6 +133,57 @@ impl Reservation { } } +/// Whether two reservations share a node and overlap in time. +pub fn reservations_overlap(a: &Reservation, b: &Reservation) -> bool { + if a.end_time <= b.start_time || b.end_time <= a.start_time { + return false; + } + a.nodes.iter().any(|n| b.covers_node(n)) +} + +/// Whether overlap between `candidate` and `existing` is permitted by flags. +pub fn overlap_allowed(candidate: &Reservation, existing: &Reservation) -> bool { + candidate.flags.overlap + || candidate.flags.maint + || existing.flags.overlap + || existing.flags.maint +} + +/// Running job is tied to an active reservation on its allocated nodes. +pub fn job_runs_in_active_reservation( + job: &Job, + reservations: &[Reservation], + now: DateTime, +) -> bool { + let Some(res_name) = job.spec.reservation.as_deref() else { + return false; + }; + reservations.iter().any(|r| { + r.name == res_name + && r.is_active(now) + && job + .allocated_nodes + .iter() + .any(|node| r.covers_node(node)) + }) +} + +/// Pending or running job has a matching active reservation (for priority ordering). +pub fn job_has_active_reservation( + job: &Job, + reservations: &[Reservation], + now: DateTime, +) -> bool { + let Some(res_name) = job.spec.reservation.as_deref() else { + return false; + }; + reservations.iter().any(|r| { + r.name == res_name + && (r.is_active(now) || r.is_future(now)) + && r.allows_user(&job.spec.user, job.spec.account.as_deref()) + }) +} + /// Expand hostlist patterns and verify each name exists in the cluster. pub fn normalize_node_list( patterns: &[String], @@ -293,10 +350,64 @@ mod tests { #[test] fn flags_parse_csv() { - let f = ReservationFlags::parse_csv("maint,ignore_jobs"); + let f = ReservationFlags::parse_csv("maint,ignore_jobs,overlap"); assert!(f.maint); assert!(f.ignore_jobs); assert!(!f.no_hold_jobs); + assert!(f.overlap); + } + + #[test] + fn reservations_overlap_detects_shared_node_and_time() { + let now = Utc::now(); + let a = Reservation { + name: "a".into(), + start_time: now, + end_time: now + Duration::hours(1), + nodes: vec!["node001".into()], + accounts: Vec::new(), + users: Vec::new(), + flags: ReservationFlags::default(), + }; + let b = Reservation { + name: "b".into(), + start_time: now + Duration::minutes(30), + end_time: now + Duration::hours(2), + nodes: vec!["node001".into()], + accounts: Vec::new(), + users: Vec::new(), + flags: ReservationFlags::default(), + }; + assert!(reservations_overlap(&a, &b)); + } + + #[test] + fn overlap_allowed_with_overlap_flag() { + let now = Utc::now(); + let mut a = Reservation { + name: "a".into(), + start_time: now, + end_time: now + Duration::hours(1), + nodes: vec!["node001".into()], + accounts: Vec::new(), + users: Vec::new(), + flags: ReservationFlags { + overlap: true, + ..Default::default() + }, + }; + let b = Reservation { + name: "b".into(), + start_time: now, + end_time: now + Duration::hours(1), + nodes: vec!["node001".into()], + accounts: Vec::new(), + users: Vec::new(), + flags: ReservationFlags::default(), + }; + assert!(overlap_allowed(&a, &b)); + a.flags.overlap = false; + assert!(!overlap_allowed(&a, &b)); } #[test] diff --git a/crates/spurctld/src/cluster.rs b/crates/spurctld/src/cluster.rs index d791312e..03fd6d8c 100644 --- a/crates/spurctld/src/cluster.rs +++ b/crates/spurctld/src/cluster.rs @@ -717,6 +717,36 @@ impl ClusterManager { Ok(()) } + /// Return a preempted job to the pending queue. + pub fn requeue_preempted_job(&self, job_id: JobId, force: bool) -> anyhow::Result<()> { + const MAX_REQUEUE: u32 = 3; + let old_state = { + let jobs = self.jobs.read(); + let Some(job) = jobs.get(&job_id) else { + return Ok(()); + }; + if job.requeue_count >= MAX_REQUEUE { + return Ok(()); + } + if job.state != JobState::Preempted { + return Ok(()); + } + if !force && !job.spec.requeue { + return Ok(()); + } + job.state + }; + + self.propose(WalOperation::JobStateChange { + job_id, + old_state, + new_state: JobState::Pending, + })?; + + info!(job_id, "preempted job requeued"); + Ok(()) + } + /// Requeue a job back to Pending after a dispatch failure. /// Unlike `maybe_requeue`, this is unconditional and doesn't require /// the requeue flag on the spec. Used when the agent rejects a job @@ -1349,6 +1379,7 @@ impl ClusterManager { // Recompute effective priority with age + partition tier let now = Utc::now(); let partitions = self.partitions.read(); + let reservations = self.get_reservations(); for job in &mut pending { let age_minutes = (now - job.submit_time).num_minutes().max(0); let partition_tier = job @@ -1368,8 +1399,15 @@ impl ClusterManager { partition_tier, ); } + drop(partitions); - pending.sort_by_key(|j| std::cmp::Reverse(j.priority)); + pending.sort_by(|a, b| { + let a_res = reservation::job_has_active_reservation(a, &reservations, now); + let b_res = reservation::job_has_active_reservation(b, &reservations, now); + b_res + .cmp(&a_res) + .then(b.priority.cmp(&a.priority)) + }); // License reservation, in priority order. `remaining` starts from current // availability (config total minus licenses held by running jobs) and each @@ -1804,6 +1842,7 @@ impl ClusterManager { res.nodes = normalize_node_list(&res.nodes, &known) .map_err(|e| anyhow::anyhow!("{e}"))?; self.validate_reservation_job_overlap(&res, None)?; + self.validate_reservation_storage_overlap(&res, None)?; self.propose(WalOperation::ReservationCreate { reservation: res })?; Ok(()) } @@ -1861,6 +1900,7 @@ impl ClusterManager { preview.accounts.retain(|a| !remove_accounts.contains(a)); self.validate_reservation_job_overlap(&preview, Some(name))?; + self.validate_reservation_storage_overlap(&preview, Some(name))?; self.propose(WalOperation::ReservationUpdate { name: name.to_string(), @@ -1937,9 +1977,10 @@ impl ClusterManager { } } - /// Cancel running jobs whose reservation window has ended. + /// Cancel running jobs whose reservation window has ended (after optional grace). pub fn enforce_reservation_end_times(&self) { let now = Utc::now(); + let grace = chrono::Duration::minutes(self.config.scheduler.resv_overrun_minutes as i64); let reservations: std::collections::HashMap = self .get_reservations() .into_iter() @@ -1958,7 +1999,7 @@ impl ClusterManager { } let res_name = job.spec.reservation.as_ref()?; let res = reservations.get(res_name)?; - if res.is_expired(now) { + if now > res.end_time + grace { Some(job.job_id) } else { None @@ -1996,6 +2037,29 @@ impl ClusterManager { Ok(()) } + fn validate_reservation_storage_overlap( + &self, + res: &Reservation, + except_name: Option<&str>, + ) -> anyhow::Result<()> { + for existing in self.reservations.read().iter() { + if except_name == Some(existing.name.as_str()) { + continue; + } + if !reservation::reservations_overlap(res, existing) { + continue; + } + if reservation::overlap_allowed(res, existing) { + continue; + } + anyhow::bail!( + "reservation overlaps with existing reservation '{}'", + existing.name + ); + } + Ok(()) + } + fn hold_jobs_for_deleted_reservation(&self, name: &str) { let pending: Vec = self .jobs @@ -5125,6 +5189,164 @@ mod tests { ); } + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn reservation_overlap_rejected_without_flag() { + let dir = TempDir::new().unwrap(); + let cm = test_cluster(&dir).await; + register_node(&cm, "n1", 8, 16000); + let now = chrono::Utc::now(); + let base = Reservation { + name: String::new(), + start_time: now, + end_time: now + chrono::Duration::hours(2), + nodes: vec!["n1".into()], + accounts: Vec::new(), + users: vec!["alice".into()], + flags: Default::default(), + }; + let mut r1 = base.clone(); + r1.name = "r1".into(); + cm.create_reservation(r1).unwrap(); + let mut r2 = base; + r2.name = "r2".into(); + r2.start_time = now + chrono::Duration::hours(1); + assert!(cm.create_reservation(r2).is_err()); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn reservation_overlap_allowed_with_overlap_flag() { + let dir = TempDir::new().unwrap(); + let cm = test_cluster(&dir).await; + register_node(&cm, "n1", 8, 16000); + let now = chrono::Utc::now(); + let base = Reservation { + name: String::new(), + start_time: now, + end_time: now + chrono::Duration::hours(2), + nodes: vec!["n1".into()], + accounts: Vec::new(), + users: vec!["alice".into()], + flags: Default::default(), + }; + let mut r1 = base.clone(); + r1.name = "r1".into(); + cm.create_reservation(r1).unwrap(); + let mut r2 = base; + r2.name = "r2".into(); + r2.start_time = now + chrono::Duration::hours(1); + r2.flags.overlap = true; + assert!(cm.create_reservation(r2).is_ok()); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn pending_jobs_prioritize_active_reservation_targets() { + let dir = TempDir::new().unwrap(); + let cm = test_cluster(&dir).await; + register_node(&cm, "n1", 8, 16000); + let now = chrono::Utc::now(); + cm.create_reservation(Reservation { + name: "r1".into(), + start_time: now - chrono::Duration::minutes(5), + end_time: now + chrono::Duration::hours(2), + nodes: vec!["n1".into()], + accounts: Vec::new(), + users: vec!["testuser".into()], + flags: Default::default(), + }) + .unwrap(); + + let mut plain = basic_spec("plain"); + plain.priority = Some(5000); + let plain_id = submit_and_wait(&cm, plain); + + let mut resv = basic_spec("resv"); + resv.priority = Some(1000); + resv.reservation = Some("r1".into()); + let resv_id = submit_and_wait(&cm, resv); + + let pending = cm.pending_jobs(); + assert_eq!(pending.first().map(|j| j.job_id), Some(resv_id)); + assert!(pending.iter().any(|j| j.job_id == plain_id)); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn resv_overrun_grace_delays_cancel() { + let dir = TempDir::new().unwrap(); + let mut config = test_config(); + config.scheduler.resv_overrun_minutes = 30; + let cm = test_cluster_with_config(&dir, config).await; + register_node(&cm, "n1", 8, 16000); + let now = chrono::Utc::now(); + cm.create_reservation(Reservation { + name: "r1".into(), + start_time: now - chrono::Duration::hours(1), + end_time: now - chrono::Duration::minutes(5), + nodes: vec!["n1".into()], + accounts: Vec::new(), + users: vec!["testuser".into()], + flags: Default::default(), + }) + .unwrap(); + + let mut spec = basic_spec("resv-run"); + spec.reservation = Some("r1".into()); + let job_id = submit_and_wait(&cm, spec); + let res = scalar_alloc(1, 1000); + cm.start_job( + job_id, + vec!["n1".into()], + res.clone(), + per_node_for(&["n1"], res), + ) + .unwrap(); + settle(&cm, job_id, JobState::Running); + + cm.enforce_reservation_end_times(); + assert_eq!(cm.get_job(job_id).unwrap().state, JobState::Running); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn preempt_skips_jobs_in_active_reservation_at_same_tier() { + let dir = TempDir::new().unwrap(); + let mut config = test_config(); + config.partitions[0].preempt_mode = "cancel".into(); + let cm = test_cluster_with_config(&dir, config).await; + register_node(&cm, "n1", 8, 16000); + let now = chrono::Utc::now(); + cm.create_reservation(Reservation { + name: "r1".into(), + start_time: now - chrono::Duration::minutes(5), + end_time: now + chrono::Duration::hours(2), + nodes: vec!["n1".into()], + accounts: Vec::new(), + users: vec!["testuser".into()], + flags: Default::default(), + }) + .unwrap(); + + let mut low = basic_spec("low"); + low.priority = Some(100); + low.reservation = Some("r1".into()); + let low_id = submit_and_wait(&cm, low); + let res = scalar_alloc(1, 1000); + cm.start_job( + low_id, + vec!["n1".into()], + res.clone(), + per_node_for(&["n1"], res), + ) + .unwrap(); + settle(&cm, low_id, JobState::Running); + + let mut high = basic_spec("high"); + high.priority = Some(10_000); + let high_id = submit_and_wait(&cm, high); + let high_job = cm.get_job(high_id).unwrap(); + + crate::scheduler_loop::try_preempt(&cm, &[&high_job]); + assert_eq!(cm.get_job(low_id).unwrap().state, JobState::Running); + } + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn tag_blocked_sets_licenses_reason() { let dir = TempDir::new().unwrap(); diff --git a/crates/spurctld/src/scheduler_loop.rs b/crates/spurctld/src/scheduler_loop.rs index 1794c2b8..cb06ddaa 100644 --- a/crates/spurctld/src/scheduler_loop.rs +++ b/crates/spurctld/src/scheduler_loop.rs @@ -411,10 +411,22 @@ pub(crate) fn compute_job_allocation( } /// Try to preempt lower-priority running jobs to make room for higher-priority pending jobs. -fn try_preempt(cluster: &Arc, unscheduled: &[&spur_core::job::Job]) { +pub(crate) fn try_preempt(cluster: &Arc, unscheduled: &[&spur_core::job::Job]) { use spur_core::job::JobState; + use spur_core::partition::{Partition, PreemptMode}; + use spur_core::reservation::job_runs_in_active_reservation; + + let now = chrono::Utc::now(); + let reservations = cluster.get_reservations(); + let partitions = cluster.get_partitions(); + + let partition_for = |job: &spur_core::job::Job| -> Option<&Partition> { + job.spec + .partition + .as_ref() + .and_then(|pname| partitions.iter().find(|p| p.name == *pname)) + }; - // Get running jobs sorted by priority (lowest first = best preemption candidates) let mut running: Vec = cluster .get_jobs(&[JobState::Running], None, None, None, &[]) .into_iter() @@ -422,25 +434,59 @@ fn try_preempt(cluster: &Arc, unscheduled: &[&spur_core::job::Jo running.sort_by_key(|j| j.priority); for pending in unscheduled { - // Only preempt if pending job has significantly higher priority + let Some(pending_part) = partition_for(pending) else { + continue; + }; + let preempt_mode = pending_part.preempt_mode; + if preempt_mode == PreemptMode::Off { + continue; + } + let pending_tier = pending_part.priority_tier; + for candidate in &running { - if candidate.priority < pending.priority / 2 { - // Preempt: cancel the lower-priority job - info!( - preempted_job = candidate.job_id, - preempted_priority = candidate.priority, - pending_job = pending.job_id, - pending_priority = pending.priority, - "preempting lower-priority job" - ); - if let Err(e) = cluster.complete_job(candidate.job_id, -1, JobState::Preempted) { - warn!( - job_id = candidate.job_id, - error = %e, - "failed to preempt job" - ); + if candidate.priority >= pending.priority / 2 { + continue; + } + + if job_runs_in_active_reservation(candidate, &reservations, now) { + let candidate_tier = partition_for(candidate) + .map(|p| p.priority_tier) + .unwrap_or(1); + if pending_tier <= candidate_tier { + continue; + } + } + + info!( + preempted_job = candidate.job_id, + preempted_priority = candidate.priority, + pending_job = pending.job_id, + pending_priority = pending.priority, + preempt_mode = ?preempt_mode, + "preempting lower-priority job" + ); + + let result = match preempt_mode { + PreemptMode::Off => Ok(()), + PreemptMode::Suspend => cluster.suspend_job(candidate.job_id, "scheduler"), + PreemptMode::Cancel => { + cluster + .complete_job(candidate.job_id, -1, JobState::Preempted) + .and_then(|_| cluster.requeue_preempted_job(candidate.job_id, false)) } - break; // One preemption per cycle, re-evaluate next cycle + PreemptMode::Requeue => cluster + .complete_job(candidate.job_id, -1, JobState::Preempted) + .and_then(|_| cluster.requeue_preempted_job(candidate.job_id, true)), + }; + + if let Err(e) = result { + warn!( + job_id = candidate.job_id, + error = %e, + "failed to preempt job" + ); + } else { + break; } } } diff --git a/crates/spurctld/src/server.rs b/crates/spurctld/src/server.rs index 3608b54b..98022df0 100644 --- a/crates/spurctld/src/server.rs +++ b/crates/spurctld/src/server.rs @@ -1768,6 +1768,7 @@ fn job_to_proto(job: &spur_core::job::Job) -> JobInfo { qos: job.spec.qos.clone().unwrap_or_default(), array_job_id: job.spec.array_job_id.unwrap_or(0), array_task_id: job.spec.array_task_id.unwrap_or(0), + reservation: job.spec.reservation.clone().unwrap_or_default(), } } diff --git a/proto/slurm.proto b/proto/slurm.proto index 73052770..d711ef58 100644 --- a/proto/slurm.proto +++ b/proto/slurm.proto @@ -196,6 +196,7 @@ message JobInfo { // Array jobs uint32 array_job_id = 50; uint32 array_task_id = 51; + string reservation = 52; } message NodeInfo { From b71198506e66c93ec470c9be98b025e299a072b6 Mon Sep 17 00:00:00 2001 From: Sudheendra Gopinath Date: Sun, 5 Jul 2026 07:23:39 +0000 Subject: [PATCH 3/5] fix(reservations): fix issues in scheduling and durability Reject unknown reservation flags, fix backfill window for future starts, gate preemption on node overlap, hold pending jobs on purge_expired, add WAL apply and e2e coverage, and fix spurdbd JobInfo.reservation. --- crates/spur-cli/src/scontrol.rs | 2 +- crates/spur-core/src/reservation.rs | 39 +++-- crates/spur-core/src/wal.rs | 4 +- crates/spur-sched/src/backfill.rs | 13 +- crates/spurctld/src/cluster.rs | 164 +++++++++++++++++---- crates/spurctld/src/scheduler_loop.rs | 57 ++++++- crates/spurctld/src/server.rs | 3 +- crates/spurdbd/src/server.rs | 1 + tests/native_host/e2e/test_reservations.py | 23 ++- 9 files changed, 227 insertions(+), 79 deletions(-) diff --git a/crates/spur-cli/src/scontrol.rs b/crates/spur-cli/src/scontrol.rs index 53326fc6..28906418 100644 --- a/crates/spur-cli/src/scontrol.rs +++ b/crates/spur-cli/src/scontrol.rs @@ -87,7 +87,7 @@ pub enum ScontrolCommand { /// Comma-separated users (optional) #[arg(long, default_value = "")] users: String, - /// Comma-separated flags (maint, ignore_jobs, no_hold_jobs) + /// Comma-separated flags (maint, ignore_jobs, no_hold_jobs, overlap) #[arg(long, default_value = "")] flags: String, }, diff --git a/crates/spur-core/src/reservation.rs b/crates/spur-core/src/reservation.rs index ef795e0e..d234ff7b 100644 --- a/crates/spur-core/src/reservation.rs +++ b/crates/spur-core/src/reservation.rs @@ -21,7 +21,7 @@ pub struct ReservationFlags { } impl ReservationFlags { - pub fn parse_list(values: &[String]) -> Self { + pub fn parse_list(values: &[String]) -> Result { let mut flags = Self::default(); for v in values { for part in v.split(',') { @@ -31,20 +31,19 @@ impl ReservationFlags { "no_hold_jobs" => flags.no_hold_jobs = true, "overlap" => flags.overlap = true, "" => {} - _ => {} + other => return Err(format!("unknown reservation flag '{other}'")), } } } - flags + Ok(flags) } - pub fn parse_csv(csv: &str) -> Self { + pub fn parse_csv(csv: &str) -> Result { if csv.is_empty() { - return Self::default(); + return Ok(Self::default()); } Self::parse_list( - &csv - .split(',') + &csv.split(',') .map(|s| s.trim().to_string()) .filter(|s| !s.is_empty()) .collect::>(), @@ -124,13 +123,6 @@ impl Reservation { } false } - - pub fn job_targets(&self, job: &Job) -> bool { - job.spec - .reservation - .as_deref() - .is_some_and(|n| n == self.name) - } } /// Whether two reservations share a node and overlap in time. @@ -161,10 +153,7 @@ pub fn job_runs_in_active_reservation( reservations.iter().any(|r| { r.name == res_name && r.is_active(now) - && job - .allocated_nodes - .iter() - .any(|node| r.covers_node(node)) + && job.allocated_nodes.iter().any(|node| r.covers_node(node)) }) } @@ -179,7 +168,7 @@ pub fn job_has_active_reservation( }; reservations.iter().any(|r| { r.name == res_name - && (r.is_active(now) || r.is_future(now)) + && r.is_active(now) && r.allows_user(&job.spec.user, job.spec.account.as_deref()) }) } @@ -245,7 +234,10 @@ pub fn running_jobs_overlap_start( continue; } } - let overlaps_node = job.allocated_nodes.iter().any(|n| node_set.contains(n.as_str())); + let overlaps_node = job + .allocated_nodes + .iter() + .any(|n| node_set.contains(n.as_str())); if !overlaps_node { continue; } @@ -350,13 +342,18 @@ mod tests { #[test] fn flags_parse_csv() { - let f = ReservationFlags::parse_csv("maint,ignore_jobs,overlap"); + let f = ReservationFlags::parse_csv("maint,ignore_jobs,overlap").unwrap(); assert!(f.maint); assert!(f.ignore_jobs); assert!(!f.no_hold_jobs); assert!(f.overlap); } + #[test] + fn flags_reject_unknown() { + assert!(ReservationFlags::parse_csv("not_a_flag").is_err()); + } + #[test] fn reservations_overlap_detects_shared_node_and_time() { let now = Utc::now(); diff --git a/crates/spur-core/src/wal.rs b/crates/spur-core/src/wal.rs index e9a4d968..1fc6c910 100644 --- a/crates/spur-core/src/wal.rs +++ b/crates/spur-core/src/wal.rs @@ -174,9 +174,7 @@ mod reservation_wal_tests { #[test] fn reservation_delete_round_trips() { - let op = WalOperation::ReservationDelete { - name: "r1".into(), - }; + let op = WalOperation::ReservationDelete { name: "r1".into() }; let json = serde_json::to_string(&op).unwrap(); let back: WalOperation = serde_json::from_str(&json).unwrap(); match back { diff --git a/crates/spur-sched/src/backfill.rs b/crates/spur-sched/src/backfill.rs index ad0b658f..976028b9 100644 --- a/crates/spur-sched/src/backfill.rs +++ b/crates/spur-sched/src/backfill.rs @@ -174,8 +174,8 @@ impl BackfillScheduler { return false; } if let Some(res) = reservations.iter().find(|r| r.name == res_name) { - let job_end = now + job_duration; - if now < res.start_time || job_end > res.end_time { + let earliest = now.max(res.start_time); + if earliest + job_duration > res.end_time { return false; } } @@ -184,13 +184,8 @@ impl BackfillScheduler { return false; } for res in reservations { - if reservation::prospective_overlap( - job, - res, - &node.name, - now, - job_duration, - ) { + if reservation::prospective_overlap(job, res, &node.name, now, job_duration) + { return false; } } diff --git a/crates/spurctld/src/cluster.rs b/crates/spurctld/src/cluster.rs index 03fd6d8c..9b0ac412 100644 --- a/crates/spurctld/src/cluster.rs +++ b/crates/spurctld/src/cluster.rs @@ -18,9 +18,7 @@ use spur_core::job::{Job, JobId, JobSpec, JobState, NodeCompleteError, PendingRe use spur_core::node::{Node, NodeEvent, NodeSource, NodeState}; use spur_core::partition::Partition; use spur_core::qos::{check_qos_limits, QosCheckResult}; -use spur_core::reservation::{ - self, normalize_node_list, running_jobs_overlap_start, Reservation, -}; +use spur_core::reservation::{self, normalize_node_list, running_jobs_overlap_start, Reservation}; use spur_core::resource::{ResourceAllocations, ResourceSet}; use spur_core::step::{JobStep, StepState, STEP_BATCH, STEP_RESERVED_MIN}; use spur_core::wal::WalOperation; @@ -1404,9 +1402,7 @@ impl ClusterManager { pending.sort_by(|a, b| { let a_res = reservation::job_has_active_reservation(a, &reservations, now); let b_res = reservation::job_has_active_reservation(b, &reservations, now); - b_res - .cmp(&a_res) - .then(b.priority.cmp(&a.priority)) + b_res.cmp(&a_res).then(b.priority.cmp(&a.priority)) }); // License reservation, in priority order. `remaining` starts from current @@ -1829,18 +1825,11 @@ impl ClusterManager { /// Create a new reservation (validated, persisted via Raft). pub fn create_reservation(&self, mut res: Reservation) -> anyhow::Result<()> { - if self - .reservations - .read() - .iter() - .any(|r| r.name == res.name) - { + if self.reservations.read().iter().any(|r| r.name == res.name) { anyhow::bail!("reservation '{}' already exists", res.name); } - let known: std::collections::HashSet = - self.nodes.read().keys().cloned().collect(); - res.nodes = normalize_node_list(&res.nodes, &known) - .map_err(|e| anyhow::anyhow!("{e}"))?; + let known: std::collections::HashSet = self.nodes.read().keys().cloned().collect(); + res.nodes = normalize_node_list(&res.nodes, &known).map_err(|e| anyhow::anyhow!("{e}"))?; self.validate_reservation_job_overlap(&res, None)?; self.validate_reservation_storage_overlap(&res, None)?; self.propose(WalOperation::ReservationCreate { reservation: res })?; @@ -1872,13 +1861,13 @@ impl ClusterManager { preview.end_time = preview.start_time + chrono::Duration::minutes(duration_minutes as i64); } - let known: std::collections::HashSet = - self.nodes.read().keys().cloned().collect(); + let known: std::collections::HashSet = self.nodes.read().keys().cloned().collect(); let mut add_expanded = Vec::new(); for n in add_nodes { - add_expanded.extend(normalize_node_list(std::slice::from_ref(n), &known).map_err( - |e| anyhow::anyhow!("{e}"), - )?); + add_expanded.extend( + normalize_node_list(std::slice::from_ref(n), &known) + .map_err(|e| anyhow::anyhow!("{e}"))?, + ); } for node in &add_expanded { if !preview.nodes.contains(node) { @@ -1971,8 +1960,18 @@ impl ClusterManager { if in_use { continue; } + let no_hold = self + .reservations + .read() + .iter() + .find(|r| r.name == name) + .is_some_and(|r| r.flags.no_hold_jobs); if let Err(e) = self.propose(WalOperation::ReservationDelete { name: name.clone() }) { warn!(name = %name, error = %e, "failed to purge expired reservation"); + continue; + } + if !no_hold { + self.hold_jobs_for_deleted_reservation(&name); } } } @@ -2022,12 +2021,9 @@ impl ClusterManager { return Ok(()); } let jobs = self.jobs.read(); - if let Some((job_id, node)) = running_jobs_overlap_start( - &jobs, - &res.nodes, - res.start_time, - except_name, - ) { + if let Some((job_id, node)) = + running_jobs_overlap_start(&jobs, &res.nodes, res.start_time, except_name) + { anyhow::bail!( "requested nodes are busy (job {} on {} until after reservation start)", job_id, @@ -3165,10 +3161,7 @@ fn reservation_fence_reason( cluster_state: &spur_sched::traits::ClusterState, ) -> Option { let now = Utc::now(); - let duration = job - .spec - .time_limit - .unwrap_or(chrono::Duration::hours(1)); + let duration = job.spec.time_limit.unwrap_or(chrono::Duration::hours(1)); let mut maint_block = false; let mut any_block = false; @@ -5347,6 +5340,115 @@ mod tests { assert_eq!(cm.get_job(low_id).unwrap().state, JobState::Running); } + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn preempt_skips_jobs_on_unrelated_nodes() { + let dir = TempDir::new().unwrap(); + let mut config = test_config(); + config.partitions[0].preempt_mode = "cancel".into(); + let cm = test_cluster_with_config(&dir, config).await; + register_node(&cm, "n1", 8, 16000); + register_node(&cm, "n2", 8, 16000); + + let mut low = basic_spec("low"); + low.priority = Some(100); + low.nodelist = Some("n2".into()); + let low_id = submit_and_wait(&cm, low); + let res = scalar_alloc(1, 1000); + cm.start_job( + low_id, + vec!["n2".into()], + res.clone(), + per_node_for(&["n2"], res), + ) + .unwrap(); + settle(&cm, low_id, JobState::Running); + + let mut high = basic_spec("high"); + high.priority = Some(10_000); + high.nodelist = Some("n1".into()); + let high_id = submit_and_wait(&cm, high); + let high_job = cm.get_job(high_id).unwrap(); + + crate::scheduler_loop::try_preempt(&cm, &[&high_job]); + assert_eq!(cm.get_job(low_id).unwrap().state, JobState::Running); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn purge_expired_holds_pending_reservation_jobs() { + let dir = TempDir::new().unwrap(); + let cm = test_cluster(&dir).await; + register_node(&cm, "n1", 8, 16000); + let now = chrono::Utc::now(); + cm.apply_operation(&WalOperation::ReservationCreate { + reservation: Reservation { + name: "r1".into(), + start_time: now - chrono::Duration::hours(2), + end_time: now - chrono::Duration::minutes(1), + nodes: vec!["n1".into()], + accounts: Vec::new(), + users: vec!["testuser".into()], + flags: Default::default(), + }, + }); + + let mut spec = basic_spec("resv-pending"); + spec.reservation = Some("r1".into()); + let job_id = submit_and_wait(&cm, spec); + + cm.purge_expired_reservations(); + let job = cm.get_job(job_id).unwrap(); + assert_eq!(job.state, JobState::Pending); + assert_eq!(job.pending_reason, PendingReason::ReservationDeleted); + assert_eq!(job.priority, 0); + assert!(cm.get_reservations().is_empty()); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn apply_reservation_create_update_delete() { + let dir = TempDir::new().unwrap(); + let cm = test_cluster(&dir).await; + register_node(&cm, "n1", 8, 16000); + let now = chrono::Utc::now(); + let res = Reservation { + name: "r1".into(), + start_time: now, + end_time: now + chrono::Duration::hours(1), + nodes: vec!["n1".into()], + accounts: Vec::new(), + users: vec!["alice".into()], + flags: Default::default(), + }; + cm.apply_operation(&WalOperation::ReservationCreate { + reservation: res.clone(), + }); + assert_eq!(cm.get_reservations().len(), 1); + assert_eq!(cm.get_reservations()[0].name, "r1"); + + cm.apply_operation(&WalOperation::ReservationUpdate { + name: "r1".into(), + duration_minutes: 120, + add_nodes: Vec::new(), + remove_nodes: Vec::new(), + add_users: vec!["bob".into()], + remove_users: Vec::new(), + add_accounts: Vec::new(), + remove_accounts: Vec::new(), + }); + let updated = cm + .get_reservations() + .into_iter() + .find(|r| r.name == "r1") + .unwrap(); + assert!(updated.users.contains(&"bob".into())); + assert_eq!( + updated.end_time, + updated.start_time + chrono::Duration::minutes(120) + ); + + cm.apply_operation(&WalOperation::ReservationDelete { name: "r1".into() }); + assert!(cm.get_reservations().is_empty()); + } + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn tag_blocked_sets_licenses_reason() { let dir = TempDir::new().unwrap(); diff --git a/crates/spurctld/src/scheduler_loop.rs b/crates/spurctld/src/scheduler_loop.rs index cb06ddaa..dd675344 100644 --- a/crates/spurctld/src/scheduler_loop.rs +++ b/crates/spurctld/src/scheduler_loop.rs @@ -419,6 +419,7 @@ pub(crate) fn try_preempt(cluster: &Arc, unscheduled: &[&spur_co let now = chrono::Utc::now(); let reservations = cluster.get_reservations(); let partitions = cluster.get_partitions(); + let cluster_nodes = cluster.get_nodes(); let partition_for = |job: &spur_core::job::Job| -> Option<&Partition> { job.spec @@ -448,6 +449,10 @@ pub(crate) fn try_preempt(cluster: &Arc, unscheduled: &[&spur_co continue; } + if !preempt_overlaps_pending_nodes(pending, candidate, &cluster_nodes) { + continue; + } + if job_runs_in_active_reservation(candidate, &reservations, now) { let candidate_tier = partition_for(candidate) .map(|p| p.priority_tier) @@ -467,16 +472,14 @@ pub(crate) fn try_preempt(cluster: &Arc, unscheduled: &[&spur_co ); let result = match preempt_mode { - PreemptMode::Off => Ok(()), PreemptMode::Suspend => cluster.suspend_job(candidate.job_id, "scheduler"), - PreemptMode::Cancel => { - cluster - .complete_job(candidate.job_id, -1, JobState::Preempted) - .and_then(|_| cluster.requeue_preempted_job(candidate.job_id, false)) - } + PreemptMode::Cancel => cluster + .complete_job(candidate.job_id, -1, JobState::Preempted) + .and_then(|_| cluster.requeue_preempted_job(candidate.job_id, false)), PreemptMode::Requeue => cluster .complete_job(candidate.job_id, -1, JobState::Preempted) .and_then(|_| cluster.requeue_preempted_job(candidate.job_id, true)), + PreemptMode::Off => unreachable!("filtered above"), }; if let Err(e) = result { @@ -492,6 +495,48 @@ pub(crate) fn try_preempt(cluster: &Arc, unscheduled: &[&spur_co } } +/// True when `candidate` occupies a node the pending job could target. +fn preempt_overlaps_pending_nodes( + pending: &spur_core::job::Job, + candidate: &spur_core::job::Job, + nodes: &[spur_core::node::Node], +) -> bool { + if candidate.allocated_nodes.is_empty() { + return false; + } + let occupied: HashSet<&str> = candidate + .allocated_nodes + .iter() + .map(String::as_str) + .collect(); + + if let Some(ref nodelist) = pending.spec.nodelist { + return nodelist + .split(',') + .map(str::trim) + .any(|n| occupied.contains(n)); + } + + let partitions: Vec<&str> = pending + .spec + .partition + .as_deref() + .map(|p| p.split(',').map(str::trim).collect()) + .unwrap_or_default(); + + nodes.iter().any(|node| { + if !occupied.contains(node.name.as_str()) { + return false; + } + if partitions.is_empty() { + return true; + } + partitions + .iter() + .any(|p| node.partitions.iter().any(|np| np == p)) + }) +} + /// Forward unschedulable jobs to federation peer clusters. /// /// Tries each peer in order; stops forwarding a job as soon as one peer accepts it. diff --git a/crates/spurctld/src/server.rs b/crates/spurctld/src/server.rs index 98022df0..c9dd0f53 100644 --- a/crates/spurctld/src/server.rs +++ b/crates/spurctld/src/server.rs @@ -1182,7 +1182,8 @@ impl SlurmController for ControllerService { let end_time = start_time + chrono::Duration::minutes(req.duration_minutes as i64); - let flags = spur_core::reservation::ReservationFlags::parse_list(&req.flags); + let flags = spur_core::reservation::ReservationFlags::parse_list(&req.flags) + .map_err(Status::invalid_argument)?; let reservation = spur_core::reservation::Reservation { name: req.name, diff --git a/crates/spurdbd/src/server.rs b/crates/spurdbd/src/server.rs index 99837268..6bc1af7f 100644 --- a/crates/spurdbd/src/server.rs +++ b/crates/spurdbd/src/server.rs @@ -185,6 +185,7 @@ impl SlurmAccounting for AccountingService { qos: String::new(), array_job_id: 0, array_task_id: 0, + reservation: String::new(), }) .collect(); diff --git a/tests/native_host/e2e/test_reservations.py b/tests/native_host/e2e/test_reservations.py index 4c4e88e0..4c08671a 100644 --- a/tests/native_host/e2e/test_reservations.py +++ b/tests/native_host/e2e/test_reservations.py @@ -10,10 +10,11 @@ class TestReservations: def test_create_list_and_delete_reservation(self, cluster): + res_name = f"res-e2e-{int(time.time())}" node = cluster.node_names[0] create_out = cluster.scontrol( "create-reservation", - f"--name=res-e2e-{int(time.time())}", + f"--name={res_name}", "--start-time=now", "--duration=60", f"--nodes={node}", @@ -22,9 +23,16 @@ def test_create_list_and_delete_reservation(self, cluster): assert "created" in create_out.lower() show_out = cluster.scontrol("show", "reservation") + assert res_name in show_out assert node in show_out assert "ACTIVE" in show_out or "INACTIVE" in show_out + delete_out = cluster.scontrol("delete-reservation", res_name) + assert "deleted" in delete_out.lower() + + show_after = cluster.scontrol("show", "reservation") + assert res_name not in show_after + def test_unauthorized_job_blocked_on_reserved_node(self, cluster): res_name = f"res-block-{int(time.time())}" node = cluster.node_names[0] @@ -105,14 +113,15 @@ def test_create_rejects_busy_node_without_ignore_jobs(self, cluster): time.sleep(2) res_name = f"res-busy-{int(time.time())}" - try: - cluster.scontrol( + out = cluster.cli_allow_fail( + [ + "scontrol", "create-reservation", f"--name={res_name}", "--start-time=now", "--duration=10", f"--nodes={node}", - ) - assert False, "expected reservation create to fail on busy node" - except Exception as exc: - assert "busy" in str(exc).lower() or "error" in str(exc).lower() + ] + ) + msg = out.lower() + assert "busy" in msg or "until after reservation start" in msg, f"unexpected: {out}" From acab2d4a039d687251abfdaa1d26066c9bf9b707 Mon Sep 17 00:00:00 2001 From: Sudheendra Gopinath Date: Sun, 5 Jul 2026 07:55:31 +0000 Subject: [PATCH 4/5] fix(reservations): address Copilot review on overlap, RPC, and e2e waits Block unauthorized jobs on active reservations in prospective_overlap, treat unknown running-job end times as occupied, split reservation RPC errors (internal vs not_found vs invalid_argument), OR maint flags across overlapping reservations, and wait on specific job states in e2e tests. --- crates/spur-core/src/reservation.rs | 52 +++++++++++++++--- crates/spurctld/src/server.rs | 64 ++++++++++++++++++++-- tests/native_host/e2e/cluster.py | 15 +++++ tests/native_host/e2e/test_reservations.py | 20 +------ 4 files changed, 120 insertions(+), 31 deletions(-) diff --git a/crates/spur-core/src/reservation.rs b/crates/spur-core/src/reservation.rs index d234ff7b..06a12105 100644 --- a/crates/spur-core/src/reservation.rs +++ b/crates/spur-core/src/reservation.rs @@ -241,14 +241,16 @@ pub fn running_jobs_overlap_start( if !overlaps_node { continue; } - let end = running_job_end(job, Utc::now()).unwrap_or(start_time); + let node = job + .allocated_nodes + .iter() + .find(|n| node_set.contains(n.as_str())) + .cloned() + .unwrap_or_default(); + let Some(end) = running_job_end(job, Utc::now()) else { + return Some((job.job_id, node)); + }; if end > start_time { - let node = job - .allocated_nodes - .iter() - .find(|n| node_set.contains(n.as_str())) - .cloned() - .unwrap_or_default(); return Some((job.job_id, node)); } } @@ -274,7 +276,7 @@ pub fn prospective_overlap( } let job_end = now + job_duration; if reservation.is_active(now) { - return job_end > reservation.end_time || now < reservation.start_time; + return true; } if reservation.is_future(now) { return job_end > reservation.start_time; @@ -407,6 +409,40 @@ mod tests { assert!(!overlap_allowed(&a, &b)); } + #[test] + fn prospective_overlap_blocks_unauthorized_job_on_active_reservation() { + let now = Utc::now(); + let res = Reservation { + name: "r1".into(), + start_time: now - Duration::minutes(30), + end_time: now + Duration::hours(1), + nodes: vec!["node001".into()], + accounts: Vec::new(), + users: vec!["alice".into()], + flags: ReservationFlags::default(), + }; + let job = Job::new(1, JobSpec::default()); + assert!(prospective_overlap( + &job, + &res, + "node001", + now, + Duration::minutes(1), + )); + } + + #[test] + fn running_jobs_overlap_start_treats_unknown_end_as_occupied() { + let now = Utc::now(); + let start_time = now + Duration::hours(1); + let mut job = Job::new(1, JobSpec::default()); + job.state = JobState::Running; + job.allocated_nodes = vec!["node001".into()]; + let jobs = std::collections::HashMap::from([(1, job)]); + let overlap = running_jobs_overlap_start(&jobs, &["node001".into()], start_time, None); + assert_eq!(overlap, Some((1, "node001".into()))); + } + #[test] fn prospective_overlap_blocks_long_job_before_reservation() { let now = Utc::now(); diff --git a/crates/spurctld/src/server.rs b/crates/spurctld/src/server.rs index c9dd0f53..d20297df 100644 --- a/crates/spurctld/src/server.rs +++ b/crates/spurctld/src/server.rs @@ -1197,7 +1197,7 @@ impl SlurmController for ControllerService { self.cluster .create_reservation(reservation) - .map_err(|e| Status::invalid_argument(e.to_string()))?; + .map_err(reservation_rpc_status)?; Ok(Response::new(())) } @@ -1233,7 +1233,7 @@ impl SlurmController for ControllerService { &req.add_accounts, &req.remove_accounts, ) - .map_err(|e| Status::invalid_argument(e.to_string()))?; + .map_err(reservation_rpc_status)?; Ok(Response::new(())) } @@ -1259,7 +1259,7 @@ impl SlurmController for ControllerService { let name = request.into_inner().name; self.cluster .delete_reservation(&name) - .map_err(|e| Status::not_found(e.to_string()))?; + .map_err(reservation_rpc_status)?; Ok(Response::new(())) } @@ -1923,16 +1923,32 @@ fn annotate_nodes_with_reservations( ) { for node_info in nodes.iter_mut() { node_info.reservation_maint = false; + node_info.active_reservation.clear(); for res in reservations { if res.is_active(now) && res.covers_node(&node_info.name) { - node_info.active_reservation = res.name.clone(); - node_info.reservation_maint = res.flags.maint; - break; + if node_info.active_reservation.is_empty() { + node_info.active_reservation = res.name.clone(); + } + if res.flags.maint { + node_info.reservation_maint = true; + } } } } } +#[allow(clippy::result_large_err)] +fn reservation_rpc_status(err: impl std::fmt::Display) -> Status { + let message = err.to_string(); + if message.contains("raft propose failed") { + Status::internal(message) + } else if message.contains("not found") { + Status::not_found(message) + } else { + Status::invalid_argument(message) + } +} + fn node_complete_to_status(err: NodeCompleteError) -> Status { let message = err.to_string(); let code = match err { @@ -1957,6 +1973,7 @@ mod tests { use super::*; use chrono::Duration; use spur_core::job::{JobState, NodeCompleteError}; + use spur_core::reservation::ReservationFlags; use tonic::Code; fn make_node_info(name: &str) -> NodeInfo { @@ -2042,6 +2059,41 @@ mod tests { assert_eq!(nodes[0].active_reservation, "first"); } + #[test] + fn test_annotate_maint_flag_with_overlapping_reservations() { + let mut nodes = vec![make_node_info("n1")]; + let now = Utc::now(); + let plain = Reservation { + name: "plain".into(), + start_time: now - Duration::hours(1), + end_time: now + Duration::hours(1), + nodes: vec!["n1".into()], + accounts: Vec::new(), + users: Vec::new(), + flags: ReservationFlags { + overlap: true, + ..Default::default() + }, + }; + let maint = Reservation { + name: "maint".into(), + start_time: now - Duration::hours(1), + end_time: now + Duration::hours(1), + nodes: vec!["n1".into()], + accounts: Vec::new(), + users: Vec::new(), + flags: ReservationFlags { + maint: true, + overlap: true, + ..Default::default() + }, + }; + let reservations = vec![plain, maint]; + annotate_nodes_with_reservations(&mut nodes, &reservations, Utc::now()); + assert_eq!(nodes[0].active_reservation, "plain"); + assert!(nodes[0].reservation_maint); + } + #[test] fn node_complete_error_status_mapping_covers_all_variants() { let cases: Vec<(NodeCompleteError, Code, bool)> = vec![ diff --git a/tests/native_host/e2e/cluster.py b/tests/native_host/e2e/cluster.py index 7f96e5ec..85270f70 100644 --- a/tests/native_host/e2e/cluster.py +++ b/tests/native_host/e2e/cluster.py @@ -918,6 +918,21 @@ def job_state(squeue_output: str, job_id: int) -> str | None: return None +def wait_job_state( + cluster: SpurCluster, job_id: int, state: str, timeout: int = 120 +) -> None: + """Wait until the job reaches the given squeue state (e.g. PD, R).""" + deadline = time.time() + timeout + while time.time() < deadline: + sq = cluster.squeue_all() + if job_state(sq, job_id) == state: + return + time.sleep(2) + raise TimeoutError( + f"Job {job_id} did not reach state {state} within {timeout}s" + ) + + def wait_job(cluster: SpurCluster, job_id: int, timeout: int = 120) -> str: """ Wait for a job to reach a terminal state. Returns the final state string. diff --git a/tests/native_host/e2e/test_reservations.py b/tests/native_host/e2e/test_reservations.py index 4c08671a..1fd2e284 100644 --- a/tests/native_host/e2e/test_reservations.py +++ b/tests/native_host/e2e/test_reservations.py @@ -5,7 +5,7 @@ import time -from cluster import parse_job_id, wait_job +from cluster import parse_job_id, wait_job, wait_job_state class TestReservations: @@ -50,16 +50,7 @@ def test_unauthorized_job_blocked_on_reserved_node(self, cluster): job_id = parse_job_id(sb) assert job_id is not None - deadline = time.time() + 30 - blocked = False - while time.time() < deadline: - sq = cluster.squeue_all() - if str(job_id) in sq and "PD" in sq.split(str(job_id))[1][:40]: - blocked = True - break - time.sleep(2) - - assert blocked, f"job {job_id} should stay pending on reserved node:\n{sq}" + wait_job_state(cluster, job_id, "PD", timeout=30) def test_reservation_job_schedules_for_authorized_user(self, cluster): res_name = f"res-auth-{int(time.time())}" @@ -105,12 +96,7 @@ def test_create_rejects_busy_node_without_ignore_jobs(self, cluster): job_id = parse_job_id(sb) assert job_id is not None - deadline = time.time() + 30 - while time.time() < deadline: - sq = cluster.squeue_all() - if str(job_id) in sq and " R " in sq: - break - time.sleep(2) + wait_job_state(cluster, job_id, "R", timeout=30) res_name = f"res-busy-{int(time.time())}" out = cluster.cli_allow_fail( From d3cd08cf8d1d992f6bf29760e27f09431449d013 Mon Sep 17 00:00:00 2001 From: Sudheendra Gopinath Date: Tue, 7 Jul 2026 13:47:19 +0000 Subject: [PATCH 5/5] fix(reservations): durable requeue holds and delete side effects Make max-requeue holds and reservation delete/purge job updates WAL-atomic via extended JobStateChange and JobPriorityChange. Add max_batch_requeue config, JobHoldMaxRequeue hold reason, overlap and no_hold_jobs flags, accounting reservation tracking, and e2e coverage. --- crates/spur-core/src/config.rs | 44 + crates/spur-core/src/job.rs | 12 + crates/spur-core/src/reservation.rs | 25 +- crates/spur-core/src/wal.rs | 94 +- crates/spur-metrics/src/job.rs | 23 +- crates/spur-sched/src/backfill.rs | 19 +- crates/spur-tests/src/t55_format.rs | 1 + crates/spurctld/src/accounting.rs | 2 + crates/spurctld/src/cluster.rs | 1077 +++++++++++++++++--- crates/spurctld/src/raft.rs | 2 + crates/spurctld/src/server.rs | 17 +- crates/spurdbd/src/db.rs | 59 +- crates/spurdbd/src/server.rs | 3 +- examples/spur.conf | 2 + proto/slurm.proto | 3 + tests/native_host/e2e/test_reservations.py | 84 ++ 16 files changed, 1302 insertions(+), 165 deletions(-) diff --git a/crates/spur-core/src/config.rs b/crates/spur-core/src/config.rs index 8b3b6469..86ac5c8c 100644 --- a/crates/spur-core/src/config.rs +++ b/crates/spur-core/src/config.rs @@ -283,6 +283,15 @@ pub struct ControllerConfig { /// Defaults to 90 when absent. #[serde(default)] pub heartbeat_timeout_secs: Option, + + /// Maximum automatic requeues before a job is held with `JobHoldMaxRequeue`. + /// Configured in TOML as `[controller] max_batch_requeue` (default: 5). + #[serde(default = "default_max_batch_requeue")] + pub max_batch_requeue: u32, +} + +fn default_max_batch_requeue() -> u32 { + 5 } fn default_listen_addr() -> String { @@ -317,6 +326,7 @@ impl Default for ControllerConfig { node_id: None, raft_listen_addr: "[::]:6821".into(), heartbeat_timeout_secs: None, + max_batch_requeue: default_max_batch_requeue(), } } } @@ -794,6 +804,12 @@ impl SlurmConfig { if self.cluster_name.is_empty() { return Err(ConfigError::MissingField("cluster_name".into())); } + if self.controller.max_batch_requeue == 0 { + return Err(ConfigError::InvalidValue { + field: "controller.max_batch_requeue".into(), + value: "0 (must be at least 1)".into(), + }); + } Ok(()) } @@ -1167,6 +1183,34 @@ memory_mb = 1024000 fn controller_config_defaults_for_new_fields() { let config = ControllerConfig::default(); assert_eq!(config.heartbeat_timeout_secs, None); + assert_eq!(config.max_batch_requeue, 5); + } + + #[test] + fn controller_config_parses_max_batch_requeue() { + let toml = r#" +cluster_name = "test" + +[controller] +max_batch_requeue = 7 +"#; + let config = SlurmConfig::load_from_str(toml).unwrap(); + assert_eq!(config.controller.max_batch_requeue, 7); + } + + #[test] + fn controller_config_rejects_zero_max_batch_requeue() { + let toml = r#" +cluster_name = "test" + +[controller] +max_batch_requeue = 0 +"#; + let err = SlurmConfig::load_from_str(toml).unwrap_err(); + assert!( + err.to_string().contains("max_batch_requeue"), + "unexpected error: {err}" + ); } #[test] diff --git a/crates/spur-core/src/job.rs b/crates/spur-core/src/job.rs index 78656b35..0803ca21 100644 --- a/crates/spur-core/src/job.rs +++ b/crates/spur-core/src/job.rs @@ -253,9 +253,18 @@ pub enum PendingReason { BurstBufferStageIn, ReservedMaintenance, ReservationDeleted, + JobHoldMaxRequeue, } impl PendingReason { + /// True when the job is held and must not be scheduled until released. + pub fn is_scheduling_hold(&self) -> bool { + matches!( + self, + Self::Held | Self::JobHeldAdmin | Self::JobHoldMaxRequeue | Self::ReservationDeleted + ) + } + pub fn display(&self) -> &'static str { match self { Self::None => "None", @@ -298,6 +307,7 @@ impl PendingReason { Self::BurstBufferStageIn => "BurstBufferStageIn", Self::ReservedMaintenance => "ReqNodeNotAvail, Reserved for maintenance", Self::ReservationDeleted => "ReservationDeleted", + Self::JobHoldMaxRequeue => "JobHoldMaxRequeue", } } } @@ -782,6 +792,7 @@ impl Job { // Requeue transitions: terminal → Pending (for --requeue jobs) (JobState::Timeout, JobState::Pending) => true, (JobState::Preempted, JobState::Pending) => true, + (JobState::Preempted, JobState::Cancelled) => true, (JobState::NodeFail, JobState::Pending) => true, (JobState::Failed, JobState::Pending) => true, _ => false, @@ -1174,6 +1185,7 @@ mod tests { (PendingReason::QosGrpNodeLimit, "QOSGrpNodeLimit"), (PendingReason::BurstBufferResources, "BurstBufferResources"), (PendingReason::BurstBufferStageIn, "BurstBufferStageIn"), + (PendingReason::JobHoldMaxRequeue, "JobHoldMaxRequeue"), ]; #[test] diff --git a/crates/spur-core/src/reservation.rs b/crates/spur-core/src/reservation.rs index 06a12105..e8368219 100644 --- a/crates/spur-core/src/reservation.rs +++ b/crates/spur-core/src/reservation.rs @@ -258,11 +258,11 @@ pub fn running_jobs_overlap_start( } /// Whether a job without reservation access would overlap this reservation window. -pub fn prospective_overlap( +pub fn prospective_overlap_at( job: &Job, reservation: &Reservation, node: &str, - now: DateTime, + at: DateTime, job_duration: Duration, ) -> bool { if !reservation.covers_node(node) { @@ -274,14 +274,19 @@ pub fn prospective_overlap( { return false; } - let job_end = now + job_duration; - if reservation.is_active(now) { - return true; - } - if reservation.is_future(now) { - return job_end > reservation.start_time; - } - false + let job_end = at + job_duration; + job_end > reservation.start_time && at < reservation.end_time +} + +/// Whether a job without reservation access would overlap this reservation window at `now`. +pub fn prospective_overlap( + job: &Job, + reservation: &Reservation, + node: &str, + now: DateTime, + job_duration: Duration, +) -> bool { + prospective_overlap_at(job, reservation, node, now, job_duration) } #[cfg(test)] diff --git a/crates/spur-core/src/wal.rs b/crates/spur-core/src/wal.rs index 1fc6c910..03ad5c6e 100644 --- a/crates/spur-core/src/wal.rs +++ b/crates/spur-core/src/wal.rs @@ -4,7 +4,7 @@ use serde::{Deserialize, Serialize}; use crate::admission::AdmissionToken; -use crate::job::{JobId, JobSpec, JobState}; +use crate::job::{JobId, JobSpec, JobState, PendingReason}; use crate::node::NodeState; use crate::reservation::Reservation; use std::collections::HashMap; @@ -27,6 +27,12 @@ pub enum WalOperation { job_id: JobId, old_state: JobState, new_state: JobState, + /// When set with `new_state == Pending`, applied atomically instead of clearing to `None`. + #[serde(default)] + pending_reason: Option, + /// When set with `new_state == Pending`, sets priority in the same apply step (e.g. hold at 0). + #[serde(default)] + pending_priority: Option, }, JobStart { job_id: JobId, @@ -58,6 +64,15 @@ pub enum WalOperation { job_id: JobId, old_priority: u32, new_priority: u32, + /// When set, applied on all replicas so pending reason survives replay. + #[serde(default)] + pending_reason: Option, + /// When true, clears automatic requeue counter (admin release after max requeue). + #[serde(default)] + reset_requeue_count: bool, + /// When true, clears `spec.reservation` (admin release after reservation delete hold). + #[serde(default)] + clear_reservation: bool, }, JobSuspend { job_id: JobId, @@ -138,6 +153,83 @@ pub enum WalOperation { }, } +impl WalOperation { + pub fn job_state_change(job_id: JobId, old_state: JobState, new_state: JobState) -> Self { + Self::JobStateChange { + job_id, + old_state, + new_state, + pending_reason: None, + pending_priority: None, + } + } + + /// Pending transition that applies a scheduling hold atomically (priority 0 + reason). + pub fn job_state_change_held_pending( + job_id: JobId, + old_state: JobState, + reason: PendingReason, + ) -> Self { + Self::JobStateChange { + job_id, + old_state, + new_state: JobState::Pending, + pending_reason: Some(reason), + pending_priority: Some(0), + } + } +} + +#[cfg(test)] +mod job_state_change_wal_tests { + use super::*; + + #[test] + fn job_state_change_held_pending_round_trips() { + let op = WalOperation::job_state_change_held_pending( + 1, + JobState::Preempted, + PendingReason::JobHoldMaxRequeue, + ); + let json = serde_json::to_string(&op).unwrap(); + let back: WalOperation = serde_json::from_str(&json).unwrap(); + match back { + WalOperation::JobStateChange { + job_id, + old_state, + new_state, + pending_reason, + pending_priority, + } => { + assert_eq!(job_id, 1); + assert_eq!(old_state, JobState::Preempted); + assert_eq!(new_state, JobState::Pending); + assert_eq!(pending_reason, Some(PendingReason::JobHoldMaxRequeue)); + assert_eq!(pending_priority, Some(0)); + } + _ => panic!("wrong variant"), + } + } + + #[test] + fn job_state_change_without_hold_fields_deserializes() { + let op = WalOperation::job_state_change(1, JobState::Pending, JobState::Running); + let json = serde_json::to_string(&op).unwrap(); + let back: WalOperation = serde_json::from_str(&json).unwrap(); + match back { + WalOperation::JobStateChange { + pending_reason, + pending_priority, + .. + } => { + assert_eq!(pending_reason, None); + assert_eq!(pending_priority, None); + } + _ => panic!("wrong variant"), + } + } +} + #[cfg(test)] mod reservation_wal_tests { use super::*; diff --git a/crates/spur-metrics/src/job.rs b/crates/spur-metrics/src/job.rs index 7cd5f2dd..d6f90bd4 100644 --- a/crates/spur-metrics/src/job.rs +++ b/crates/spur-metrics/src/job.rs @@ -1,7 +1,7 @@ // Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -use spur_core::job::{Job, JobState, PendingReason}; +use spur_core::job::{Job, JobState}; /// Number of [`JobState`] variants (index for `by_state`). pub const JOB_STATE_COUNT: usize = JobState::COUNT; @@ -16,7 +16,7 @@ pub struct JobMetricsSnapshot { pub total: u64, /// Count per [`JobState`]; index via [`job_state_index`]. pub by_state: [u64; JOB_STATE_COUNT], - /// Pending jobs with `pending_reason == Held`. + /// Pending jobs held from scheduling (`Held`, `JobHeldAdmin`, `JobHoldMaxRequeue`). pub held_pending: u64, /// Sum of allocated CPUs for jobs in Running or Completing. pub running_cpus: u64, @@ -46,7 +46,7 @@ impl JobMetricsSnapshot { snap.total += 1; snap.by_state[job_state_index(job.state)] += 1; - if job.state == JobState::Pending && job.pending_reason == PendingReason::Held { + if job.state == JobState::Pending && job.pending_reason.is_scheduling_hold() { snap.held_pending += 1; } @@ -121,6 +121,23 @@ mod tests { assert_eq!(snap.count_state(JobState::Completed), 1); } + #[test] + fn counts_reservation_deleted_as_held() { + let mut job = job_with_state(1, JobState::Pending, false); + job.pending_reason = PendingReason::ReservationDeleted; + job.priority = 0; + let snap = JobMetricsSnapshot::collect([&job]); + assert_eq!(snap.held_pending, 1); + } + + #[test] + fn counts_job_hold_max_requeue_as_held() { + let mut job = job_with_state(1, JobState::Pending, false); + job.pending_reason = PendingReason::JobHoldMaxRequeue; + let snap = JobMetricsSnapshot::collect([&job]); + assert_eq!(snap.held_pending, 1); + } + #[test] fn running_alloc_sums_running_and_completing() { let jobs = [ diff --git a/crates/spur-sched/src/backfill.rs b/crates/spur-sched/src/backfill.rs index 976028b9..395c9590 100644 --- a/crates/spur-sched/src/backfill.rs +++ b/crates/spur-sched/src/backfill.rs @@ -54,22 +54,9 @@ impl BackfillScheduler { start: chrono::DateTime, duration: chrono::Duration, ) -> bool { - let end = start + duration; - for res in reservations { - if !res.covers_node(node) { - continue; - } - let res_name = job.spec.reservation.as_deref().filter(|s| !s.is_empty()); - if res_name == Some(res.name.as_str()) - && res.allows_user(&job.spec.user, job.spec.account.as_deref()) - { - continue; - } - if end > res.start_time && start < res.end_time { - return true; - } - } - false + reservations + .iter() + .any(|res| reservation::prospective_overlap_at(job, res, node, start, duration)) } /// Find nodes that satisfy a job's resource requirements. diff --git a/crates/spur-tests/src/t55_format.rs b/crates/spur-tests/src/t55_format.rs index 020deb05..20e183ce 100644 --- a/crates/spur-tests/src/t55_format.rs +++ b/crates/spur-tests/src/t55_format.rs @@ -126,6 +126,7 @@ mod tests { "ReqNodeNotAvail, Reserved for maintenance", ), (PendingReason::ReservationDeleted, "ReservationDeleted"), + (PendingReason::JobHoldMaxRequeue, "JobHoldMaxRequeue"), (PendingReason::QosMaxCpuPerJobLimit, "QOSMaxCpuPerJobLimit"), ( PendingReason::QosMaxWallDurationPerJobLimit, diff --git a/crates/spurctld/src/accounting.rs b/crates/spurctld/src/accounting.rs index 0640668b..4e546173 100644 --- a/crates/spurctld/src/accounting.rs +++ b/crates/spurctld/src/accounting.rs @@ -35,6 +35,7 @@ impl AccountingNotifier { partition: String, resources: &ResourceAllocations, start_time: DateTime, + reservation: Option, ) { let req = RecordJobStartRequest { job_id, @@ -43,6 +44,7 @@ impl AccountingNotifier { partition, resources: Some(allocations_to_proto(resources)), start_time: Some(datetime_to_proto(start_time)), + reservation: reservation.unwrap_or_default(), }; let mut client = self.client.clone(); tokio::spawn(async move { diff --git a/crates/spurctld/src/cluster.rs b/crates/spurctld/src/cluster.rs index 9b0ac412..9ff20815 100644 --- a/crates/spurctld/src/cluster.rs +++ b/crates/spurctld/src/cluster.rs @@ -42,6 +42,46 @@ pub enum NodeCompleteResult { AlreadyTerminal, } +/// Reservation CRUD errors for the gRPC boundary. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum ReservationError { + InvalidArgument(String), + NotFound(String), + AlreadyExists(String), + Raft(String), +} + +impl std::fmt::Display for ReservationError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::InvalidArgument(m) + | Self::NotFound(m) + | Self::AlreadyExists(m) + | Self::Raft(m) => f.write_str(m), + } + } +} + +impl std::error::Error for ReservationError {} + +impl ReservationError { + pub fn invalid(msg: impl Into) -> Self { + Self::InvalidArgument(msg.into()) + } + + pub fn not_found(msg: impl Into) -> Self { + Self::NotFound(msg.into()) + } + + pub fn already_exists(msg: impl Into) -> Self { + Self::AlreadyExists(msg.into()) + } + + pub fn raft(msg: impl Into) -> Self { + Self::Raft(msg.into()) + } +} + /// Central cluster state manager. /// /// Thread-safe via RwLock. The scheduler and gRPC server both access this. @@ -453,11 +493,11 @@ impl ClusterManager { } // propose() handles: state transition, resource allocation, license subtraction - self.propose(WalOperation::JobStateChange { + self.propose(WalOperation::job_state_change( job_id, old_state, - new_state: JobState::Running, - })?; + JobState::Running, + ))?; self.propose(WalOperation::JobStart { job_id, nodes: node_names.clone(), @@ -507,6 +547,7 @@ impl ClusterManager { spec_for_notify.partition.clone().unwrap_or_default(), &resources, Utc::now(), + spec_for_notify.reservation.clone(), ); } @@ -690,13 +731,23 @@ impl ClusterManager { /// Requeue a job if spec.requeue is set and attempt limit not exceeded. fn maybe_requeue(&self, job_id: JobId) -> anyhow::Result<()> { - const MAX_REQUEUE: u32 = 3; + let max = self.config.controller.max_batch_requeue; let (should_requeue, old_state) = { let jobs = self.jobs.read(); let Some(job) = jobs.get(&job_id) else { return Ok(()); }; - if !job.spec.requeue || job.requeue_count >= MAX_REQUEUE { + if job.requeue_count >= max { + if matches!( + job.state, + JobState::Preempted | JobState::Timeout | JobState::NodeFail + ) { + drop(jobs); + return self.hold_job_at_max_requeue(job_id); + } + return Ok(()); + } + if !job.spec.requeue { return Ok(()); } (true, job.state) @@ -705,11 +756,11 @@ impl ClusterManager { return Ok(()); } - self.propose(WalOperation::JobStateChange { + self.propose(WalOperation::job_state_change( job_id, old_state, - new_state: JobState::Pending, - })?; + JobState::Pending, + ))?; info!(job_id, from = %old_state, "job requeued"); Ok(()) @@ -717,15 +768,19 @@ impl ClusterManager { /// Return a preempted job to the pending queue. pub fn requeue_preempted_job(&self, job_id: JobId, force: bool) -> anyhow::Result<()> { - const MAX_REQUEUE: u32 = 3; + let max = self.config.controller.max_batch_requeue; let old_state = { let jobs = self.jobs.read(); let Some(job) = jobs.get(&job_id) else { return Ok(()); }; - if job.requeue_count >= MAX_REQUEUE { + if job.requeue_count >= max { + if job.state == JobState::Preempted { + drop(jobs); + return self.hold_job_at_max_requeue(job_id); + } return Ok(()); - } + }; if job.state != JobState::Preempted { return Ok(()); } @@ -735,11 +790,11 @@ impl ClusterManager { job.state }; - self.propose(WalOperation::JobStateChange { + self.propose(WalOperation::job_state_change( job_id, old_state, - new_state: JobState::Pending, - })?; + JobState::Pending, + ))?; info!(job_id, "preempted job requeued"); Ok(()) @@ -759,6 +814,10 @@ impl ClusterManager { if job.state.is_terminal() { return Ok(()); } + if job.requeue_count >= self.config.controller.max_batch_requeue { + drop(jobs); + return self.hold_job_at_max_requeue(job_id); + } job.state }; @@ -772,11 +831,11 @@ impl ClusterManager { // Failed → Pending resets allocation fields and makes // the job schedulable again. - self.propose(WalOperation::JobStateChange { + self.propose(WalOperation::job_state_change( job_id, - old_state: JobState::Failed, - new_state: JobState::Pending, - })?; + JobState::Failed, + JobState::Pending, + ))?; info!(job_id, from = %old_state, "job requeued after dispatch failure"); Ok(()) @@ -952,35 +1011,102 @@ impl ClusterManager { job_id, old_priority, new_priority: 0, + pending_reason: Some(PendingReason::Held), + reset_requeue_count: false, + clear_reservation: false, })?; - // Set held reason (not WAL-tracked) - if let Some(job) = self.jobs.write().get_mut(&job_id) { - job.pending_reason = PendingReason::Held; - } info!(job_id, "job held"); Ok(()) } + /// Hold a job that exhausted automatic requeues (`JobHoldMaxRequeue`). + fn hold_job_at_max_requeue(&self, job_id: JobId) -> anyhow::Result<()> { + let mut state = { + let jobs = self.jobs.read(); + let job = jobs + .get(&job_id) + .ok_or_else(|| anyhow::anyhow!("job {} not found", job_id))?; + job.state + }; + + if state == JobState::Running { + self.propose(WalOperation::JobComplete { + job_id, + exit_code: -1, + state: JobState::Failed, + })?; + state = JobState::Failed; + } + + if matches!( + state, + JobState::Preempted | JobState::Timeout | JobState::NodeFail | JobState::Failed + ) { + self.propose(WalOperation::job_state_change_held_pending( + job_id, + state, + PendingReason::JobHoldMaxRequeue, + ))?; + info!(job_id, "job held at max requeue limit"); + return Ok(()); + } + + if state != JobState::Pending { + anyhow::bail!( + "cannot hold job {} at max requeue from state {:?}", + job_id, + state + ); + } + + let needs_hold = self.jobs.read().get(&job_id).is_some_and(|j| { + j.pending_reason != PendingReason::JobHoldMaxRequeue || j.priority != 0 + }); + if needs_hold { + let old_priority = self + .jobs + .read() + .get(&job_id) + .map(|j| j.priority) + .unwrap_or(0); + self.propose(WalOperation::JobPriorityChange { + job_id, + old_priority, + new_priority: 0, + pending_reason: Some(PendingReason::JobHoldMaxRequeue), + reset_requeue_count: false, + clear_reservation: false, + })?; + } + info!(job_id, "job held at max requeue limit"); + Ok(()) + } + /// Release a held job. pub fn release_job(&self, job_id: JobId) -> anyhow::Result<()> { - { + let (reset_requeue, clear_reservation, old_priority) = { let jobs = self.jobs.read(); let job = jobs .get(&job_id) .ok_or_else(|| anyhow::anyhow!("job {} not found", job_id))?; - if job.pending_reason != PendingReason::Held { + if !job.pending_reason.is_scheduling_hold() { anyhow::bail!("job {} is not held", job_id); } - } + ( + job.pending_reason == PendingReason::JobHoldMaxRequeue, + job.pending_reason == PendingReason::ReservationDeleted, + job.priority, + ) + }; self.propose(WalOperation::JobPriorityChange { job_id, - old_priority: 0, + old_priority, new_priority: 1000, + pending_reason: Some(PendingReason::Priority), + reset_requeue_count: reset_requeue, + clear_reservation, })?; - if let Some(job) = self.jobs.write().get_mut(&job_id) { - job.pending_reason = PendingReason::Priority; - } info!(job_id, "job released"); Ok(()) } @@ -1015,6 +1141,9 @@ impl ClusterManager { job_id, old_priority: old, new_priority: p, + pending_reason: None, + reset_requeue_count: false, + clear_reservation: false, })?; } @@ -1282,7 +1411,7 @@ impl ClusterManager { let jobs = self.jobs.read(); let mut pending: Vec = jobs .values() - .filter(|j| j.state == JobState::Pending && j.pending_reason != PendingReason::Held) + .filter(|j| j.state == JobState::Pending && !j.pending_reason.is_scheduling_hold()) .cloned() .collect(); @@ -1565,7 +1694,7 @@ impl ClusterManager { .filter(|j| { j.state == JobState::Pending && j.bb_stage_state == BbStageState::None - && j.pending_reason != PendingReason::Held + && !j.pending_reason.is_scheduling_hold() && j.pending_reason != PendingReason::DeadLine && extract_bb_requirement(&j.spec) > 0 }) @@ -1671,7 +1800,7 @@ impl ClusterManager { for job in jobs.values() { if job.state != JobState::Pending || job.spec.dependency.is_empty() - || job.pending_reason == PendingReason::Held + || job.pending_reason.is_scheduling_hold() { continue; } @@ -1692,7 +1821,7 @@ impl ClusterManager { // Don't clobber Held or DeadLine — matches // update_pending_reasons(). if j.state == JobState::Pending - && j.pending_reason != PendingReason::Held + && !j.pending_reason.is_scheduling_hold() && j.pending_reason != PendingReason::DeadLine { j.pending_reason = PendingReason::Dependency; @@ -1786,7 +1915,7 @@ impl ClusterManager { jobs.values() .filter(|job| { job.state == JobState::Pending - && job.pending_reason != PendingReason::Held + && !job.pending_reason.is_scheduling_hold() && job.pending_reason != PendingReason::DeadLine }) .filter_map(|job| { @@ -1814,7 +1943,7 @@ impl ClusterManager { // Re-check under the write lock: the read snapshot was released, // so the job may have started or been held/deadlined since. if j.state == JobState::Pending - && j.pending_reason != PendingReason::Held + && !j.pending_reason.is_scheduling_hold() && j.pending_reason != PendingReason::DeadLine { j.pending_reason = reason; @@ -1824,15 +1953,29 @@ impl ClusterManager { } /// Create a new reservation (validated, persisted via Raft). - pub fn create_reservation(&self, mut res: Reservation) -> anyhow::Result<()> { + pub fn create_reservation(&self, mut res: Reservation) -> Result<(), ReservationError> { if self.reservations.read().iter().any(|r| r.name == res.name) { - anyhow::bail!("reservation '{}' already exists", res.name); + return Err(ReservationError::already_exists(format!( + "reservation '{}' already exists", + res.name + ))); } let known: std::collections::HashSet = self.nodes.read().keys().cloned().collect(); - res.nodes = normalize_node_list(&res.nodes, &known).map_err(|e| anyhow::anyhow!("{e}"))?; - self.validate_reservation_job_overlap(&res, None)?; - self.validate_reservation_storage_overlap(&res, None)?; - self.propose(WalOperation::ReservationCreate { reservation: res })?; + res.nodes = normalize_node_list(&res.nodes, &known).map_err(ReservationError::invalid)?; + self.validate_reservation_job_overlap(&res, None) + .map_err(|e| ReservationError::invalid(e.to_string()))?; + self.validate_reservation_storage_overlap(&res, None) + .map_err(|e| ReservationError::invalid(e.to_string()))?; + let name = res.name.clone(); + let resp = self + .propose(WalOperation::ReservationCreate { reservation: res }) + .map_err(|e| ReservationError::raft(e.to_string()))?; + if !resp.reservation_created { + return Err(ReservationError::already_exists(format!( + "reservation '{}' already exists", + name + ))); + } Ok(()) } @@ -1848,14 +1991,16 @@ impl ClusterManager { remove_users: &[String], add_accounts: &[String], remove_accounts: &[String], - ) -> anyhow::Result<()> { + ) -> Result<(), ReservationError> { let mut preview = self .reservations .read() .iter() .find(|r| r.name == name) .cloned() - .ok_or_else(|| anyhow::anyhow!("reservation '{}' not found", name))?; + .ok_or_else(|| { + ReservationError::not_found(format!("reservation '{}' not found", name)) + })?; if duration_minutes > 0 { preview.end_time = @@ -1866,7 +2011,7 @@ impl ClusterManager { for n in add_nodes { add_expanded.extend( normalize_node_list(std::slice::from_ref(n), &known) - .map_err(|e| anyhow::anyhow!("{e}"))?, + .map_err(ReservationError::invalid)?, ); } for node in &add_expanded { @@ -1888,8 +2033,10 @@ impl ClusterManager { } preview.accounts.retain(|a| !remove_accounts.contains(a)); - self.validate_reservation_job_overlap(&preview, Some(name))?; - self.validate_reservation_storage_overlap(&preview, Some(name))?; + self.validate_reservation_job_overlap(&preview, Some(name)) + .map_err(|e| ReservationError::invalid(e.to_string()))?; + self.validate_reservation_storage_overlap(&preview, Some(name)) + .map_err(|e| ReservationError::invalid(e.to_string()))?; self.propose(WalOperation::ReservationUpdate { name: name.to_string(), @@ -1900,19 +2047,19 @@ impl ClusterManager { remove_users: remove_users.to_vec(), add_accounts: add_accounts.to_vec(), remove_accounts: remove_accounts.to_vec(), - })?; + }) + .map_err(|e| ReservationError::raft(e.to_string()))?; Ok(()) } /// Delete a reservation by name (persisted via Raft). - pub fn delete_reservation(&self, name: &str) -> anyhow::Result<()> { - let res = self - .reservations - .read() - .iter() - .find(|r| r.name == name) - .cloned() - .ok_or_else(|| anyhow::anyhow!("reservation '{}' not found", name))?; + pub fn delete_reservation(&self, name: &str) -> Result<(), ReservationError> { + if !self.reservations.read().iter().any(|r| r.name == name) { + return Err(ReservationError::not_found(format!( + "reservation '{}' not found", + name + ))); + } for job in self.jobs.read().values() { if !matches!( @@ -1922,21 +2069,17 @@ impl ClusterManager { continue; } if job.spec.reservation.as_deref() == Some(name) { - anyhow::bail!( + return Err(ReservationError::invalid(format!( "reservation '{}' in use by running job {}", - name, - job.job_id - ); + name, job.job_id + ))); } } self.propose(WalOperation::ReservationDelete { name: name.to_string(), - })?; - - if !res.flags.no_hold_jobs { - self.hold_jobs_for_deleted_reservation(name); - } + }) + .map_err(|e| ReservationError::raft(e.to_string()))?; Ok(()) } @@ -1960,18 +2103,8 @@ impl ClusterManager { if in_use { continue; } - let no_hold = self - .reservations - .read() - .iter() - .find(|r| r.name == name) - .is_some_and(|r| r.flags.no_hold_jobs); if let Err(e) = self.propose(WalOperation::ReservationDelete { name: name.clone() }) { warn!(name = %name, error = %e, "failed to purge expired reservation"); - continue; - } - if !no_hold { - self.hold_jobs_for_deleted_reservation(&name); } } } @@ -2056,25 +2189,33 @@ impl ClusterManager { Ok(()) } - fn hold_jobs_for_deleted_reservation(&self, name: &str) { - let pending: Vec = self - .jobs - .read() - .values() - .filter(|j| { - j.state == JobState::Pending - && j.spec.reservation.as_deref() == Some(name) - && j.pending_reason != PendingReason::Held - }) - .map(|j| j.job_id) - .collect(); - for job_id in pending { - if let Err(e) = self.hold_job(job_id) { - warn!(job_id, error = %e, "failed to hold job after reservation delete"); + fn hold_jobs_for_deleted_reservation_jobs(jobs: &mut HashMap, name: &str) { + for job in jobs.values_mut() { + if job.state != JobState::Pending { continue; } - if let Some(job) = self.jobs.write().get_mut(&job_id) { - job.pending_reason = PendingReason::ReservationDeleted; + if job.spec.reservation.as_deref() != Some(name) { + continue; + } + if job.pending_reason.is_scheduling_hold() { + continue; + } + job.priority = 0; + job.pending_reason = PendingReason::ReservationDeleted; + } + } + + fn detach_jobs_from_deleted_reservation_jobs(jobs: &mut HashMap, name: &str) { + for job in jobs.values_mut() { + if job.state != JobState::Pending { + continue; + } + if job.spec.reservation.as_deref() != Some(name) { + continue; + } + job.spec.reservation = None; + if job.pending_reason == PendingReason::ReservationDeleted { + job.pending_reason = PendingReason::None; } } } @@ -2148,7 +2289,7 @@ impl ClusterManager { }; // Don't overwrite held jobs - if job_entry.pending_reason == PendingReason::Held { + if job_entry.pending_reason.is_scheduling_hold() { continue; } // Don't overwrite a DeadLine reason set by the deadline-enforcement @@ -2532,7 +2673,11 @@ impl ClusterManager { next_id = next_id.max(job_id + 1); } WalOperation::JobStateChange { - job_id, new_state, .. + job_id, + new_state, + pending_reason, + pending_priority, + .. } => { if let Some(job) = jobs.get_mut(job_id) { if let Err(e) = job.transition(*new_state) { @@ -2540,13 +2685,19 @@ impl ClusterManager { } // Requeue: reset allocation fields when returning to Pending if *new_state == JobState::Pending { - job.requeue_count += 1; + let max = self.config.controller.max_batch_requeue; + if job.requeue_count < max { + job.requeue_count += 1; + } job.start_time = None; job.exit_code = None; job.allocated_nodes.clear(); job.allocated_resources = None; job.per_node_alloc.clear(); - job.pending_reason = PendingReason::None; + job.pending_reason = pending_reason.clone().unwrap_or(PendingReason::None); + if let Some(priority) = pending_priority { + job.priority = *priority; + } } } } @@ -2725,6 +2876,7 @@ impl ClusterManager { state: final_state, exit_code: final_exit, }], + ..Default::default() }; } } @@ -2834,10 +2986,22 @@ impl ClusterManager { WalOperation::JobPriorityChange { job_id, new_priority, + pending_reason, + reset_requeue_count, + clear_reservation, .. } => { if let Some(job) = jobs.get_mut(job_id) { job.priority = *new_priority; + if let Some(reason) = pending_reason { + job.pending_reason = reason.clone(); + } + if *reset_requeue_count { + job.requeue_count = 0; + } + if *clear_reservation { + job.spec.reservation = None; + } } } WalOperation::NodeRegister { @@ -2995,8 +3159,17 @@ impl ClusterManager { } } WalOperation::ReservationCreate { reservation } => { - self.reservations.write().push(reservation.clone()); - info!(name = %reservation.name, "reservation created"); + let mut reservations = self.reservations.write(); + if reservations.iter().any(|r| r.name == reservation.name) { + warn!( + name = %reservation.name, + "duplicate reservation create in WAL apply, ignoring" + ); + } else { + reservations.push(reservation.clone()); + response.reservation_created = true; + info!(name = %reservation.name, "reservation created"); + } } WalOperation::ReservationUpdate { name, @@ -3023,11 +3196,22 @@ impl ClusterManager { info!(name, "reservation updated"); } WalOperation::ReservationDelete { name } => { + let deleted = { + let reservations = self.reservations.read(); + reservations.iter().find(|r| r.name == *name).cloned() + }; let mut reservations = self.reservations.write(); let len_before = reservations.len(); reservations.retain(|r| r.name != *name); if reservations.len() < len_before { - info!(name, "reservation deleted"); + if let Some(res) = deleted { + if res.flags.no_hold_jobs { + Self::detach_jobs_from_deleted_reservation_jobs(&mut jobs, name); + } else { + Self::hold_jobs_for_deleted_reservation_jobs(&mut jobs, name); + } + info!(name, "reservation deleted"); + } } } } @@ -3156,44 +3340,101 @@ impl StateMachineApply for ClusterManager { } } +fn job_candidate_node_names(job: &Job, nodes: &[spur_core::node::Node]) -> Vec { + let partitions = self_partitions_for_job(job); + let nodelist: Option> = job.spec.nodelist.as_deref().map(|s| { + s.split(',') + .map(str::trim) + .filter(|p| !p.is_empty()) + .collect() + }); + let exclude: std::collections::HashSet<&str> = job + .spec + .exclude + .as_deref() + .map(|s| { + s.split(',') + .map(str::trim) + .filter(|p| !p.is_empty()) + .collect() + }) + .unwrap_or_default(); + + nodes + .iter() + .filter(|node| { + if exclude.contains(node.name.as_str()) { + return false; + } + if let Some(ref nl) = nodelist { + if !nl.iter().any(|n| *n == node.name) { + return false; + } + } + if !partitions.is_empty() + && !partitions + .iter() + .any(|p| node.partitions.iter().any(|np| np == p)) + { + return false; + } + true + }) + .map(|n| n.name.clone()) + .collect() +} + +fn self_partitions_for_job(job: &Job) -> Vec { + job.spec + .partition + .as_deref() + .map(|p| { + p.split(',') + .map(str::trim) + .filter(|s| !s.is_empty()) + .map(String::from) + .collect() + }) + .unwrap_or_default() +} + fn reservation_fence_reason( job: &Job, cluster_state: &spur_sched::traits::ClusterState, ) -> Option { + let candidates = job_candidate_node_names(job, cluster_state.nodes); + if candidates.is_empty() { + return None; + } + let now = Utc::now(); let duration = job.spec.time_limit.unwrap_or(chrono::Duration::hours(1)); let mut maint_block = false; - let mut any_block = false; + let mut all_blocked = true; - for node in cluster_state.nodes { - if let Some(pname) = job.spec.partition.as_deref() { - let requested: Vec<&str> = pname.split(',').map(str::trim).collect(); - if !requested - .iter() - .any(|rp| node.partitions.iter().any(|np| np == rp)) - { - continue; - } - } + for node_name in &candidates { + let mut node_blocked = false; for res in cluster_state.reservations { - if !res.covers_node(&node.name) { - continue; - } - if reservation::prospective_overlap(job, res, &node.name, now, duration) { - any_block = true; + if reservation::prospective_overlap(job, res, node_name, now, duration) { + node_blocked = true; if res.flags.maint { maint_block = true; } } } + if !node_blocked { + all_blocked = false; + break; + } } + if !all_blocked { + return None; + } if maint_block { Some(PendingReason::ReservedMaintenance) - } else if any_block { - Some(PendingReason::ReqNodeNotAvail) } else { - None + Some(PendingReason::ReqNodeNotAvail) } } @@ -3730,6 +3971,8 @@ mod tests { job_id: 1, old_state: JobState::Pending, new_state: JobState::Running, + pending_reason: None, + pending_priority: None, }); let job = cm.get_job(1).unwrap(); @@ -3778,6 +4021,8 @@ mod tests { job_id: 1, old_state: JobState::Pending, new_state: JobState::Running, + pending_reason: None, + pending_priority: None, }); let alloc = scalar_alloc(4, 8000); cm.apply_operation(&WalOperation::JobStart { @@ -3816,6 +4061,8 @@ mod tests { job_id: 1, old_state: JobState::Pending, new_state: JobState::Running, + pending_reason: None, + pending_priority: None, }); let t0 = chrono::Utc::now(); cm.apply_operation(&WalOperation::JobSuspend { job_id: 1, at: t0 }); @@ -3980,6 +4227,8 @@ mod tests { job_id: 1, old_state: JobState::Pending, new_state: JobState::Running, + pending_reason: None, + pending_priority: None, }); let t0 = chrono::Utc::now(); // Cycle 1: 10s suspended. @@ -4017,6 +4266,8 @@ mod tests { job_id: 1, old_state: JobState::Pending, new_state: JobState::Running, + pending_reason: None, + pending_priority: None, }); // Suspended 30s ago, then cancelled now (JobComplete stamps Utc::now()). let since = chrono::Utc::now() - chrono::Duration::seconds(30); @@ -4139,6 +4390,9 @@ mod tests { job_id: 1, old_priority: 1000, new_priority: 5000, + pending_reason: None, + reset_requeue_count: false, + clear_reservation: false, }); let job = cm.get_job(1).unwrap(); @@ -4253,6 +4507,8 @@ mod tests { job_id: 1, old_state: JobState::Pending, new_state: JobState::Running, + pending_reason: None, + pending_priority: None, }); let alloc = scalar_alloc(2, 4000); cm.apply_operation(&WalOperation::JobStart { @@ -4292,6 +4548,8 @@ mod tests { job_id: 1, old_state: JobState::Pending, new_state: JobState::Running, + pending_reason: None, + pending_priority: None, }); let alloc = scalar_alloc(2, 4000); cm.apply_operation(&WalOperation::JobStart { @@ -4332,6 +4590,8 @@ mod tests { job_id: 1, old_state: JobState::Pending, new_state: JobState::Running, + pending_reason: None, + pending_priority: None, }); let alloc = scalar_alloc(2, 4000); cm.apply_operation(&WalOperation::JobStart { @@ -4396,6 +4656,8 @@ mod tests { job_id: 1, old_state: JobState::Pending, new_state: JobState::Running, + pending_reason: None, + pending_priority: None, }); cm.apply_operation(&WalOperation::JobStart { job_id: 1, @@ -4455,6 +4717,8 @@ mod tests { job_id: 1, old_state: JobState::Pending, new_state: JobState::Running, + pending_reason: None, + pending_priority: None, }); cm.apply_operation(&WalOperation::JobStepComplete { @@ -4483,6 +4747,8 @@ mod tests { job_id: 1, old_state: JobState::Pending, new_state: JobState::Running, + pending_reason: None, + pending_priority: None, }); let alloc = scalar_alloc(2, 4000); cm.apply_operation(&WalOperation::JobStart { @@ -4531,6 +4797,8 @@ mod tests { job_id: 1, old_state: JobState::Pending, new_state: JobState::Running, + pending_reason: None, + pending_priority: None, }); let alloc = scalar_alloc(2, 4000); cm.apply_operation(&WalOperation::JobStart { @@ -4568,6 +4836,8 @@ mod tests { job_id: 1, old_state: JobState::Pending, new_state: JobState::Running, + pending_reason: None, + pending_priority: None, }); let alloc = scalar_alloc(2, 4000); cm.apply_operation(&WalOperation::JobStart { @@ -4623,6 +4893,8 @@ mod tests { job_id: 1, old_state: JobState::Pending, new_state: JobState::Running, + pending_reason: None, + pending_priority: None, }); let alloc = scalar_alloc(2, 4000); cm.apply_operation(&WalOperation::JobStart { @@ -4658,6 +4930,8 @@ mod tests { job_id: 1, old_state: JobState::Pending, new_state: JobState::Running, + pending_reason: None, + pending_priority: None, }); cm.apply_operation(&WalOperation::JobStart { job_id: 1, @@ -4698,6 +4972,8 @@ mod tests { job_id: 1, old_state: JobState::Pending, new_state: JobState::Running, + pending_reason: None, + pending_priority: None, }); cm.apply_operation(&WalOperation::JobStart { job_id: 1, @@ -4731,6 +5007,8 @@ mod tests { job_id: 1, old_state: JobState::Pending, new_state: JobState::Running, + pending_reason: None, + pending_priority: None, }); cm.apply_operation(&WalOperation::JobStart { job_id: 1, @@ -4767,6 +5045,8 @@ mod tests { job_id: 1, old_state: JobState::Pending, new_state: JobState::Running, + pending_reason: None, + pending_priority: None, }); let alloc = scalar_alloc(2, 4000); cm.apply_operation(&WalOperation::JobStart { @@ -4832,6 +5112,8 @@ mod tests { job_id: 1, old_state: JobState::Pending, new_state: JobState::Running, + pending_reason: None, + pending_priority: None, }); let alloc = scalar_alloc(2, 4000); cm.apply_operation(&WalOperation::JobStart { @@ -5449,6 +5731,555 @@ mod tests { assert!(cm.get_reservations().is_empty()); } + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn apply_reservation_create_idempotent() { + let dir = TempDir::new().unwrap(); + let cm = test_cluster(&dir).await; + register_node(&cm, "n1", 8, 16000); + let now = chrono::Utc::now(); + let res = Reservation { + name: "r1".into(), + start_time: now, + end_time: now + chrono::Duration::hours(1), + nodes: vec!["n1".into()], + accounts: Vec::new(), + users: vec!["alice".into()], + flags: Default::default(), + }; + cm.apply_operation(&WalOperation::ReservationCreate { + reservation: res.clone(), + }); + cm.apply_operation(&WalOperation::ReservationCreate { reservation: res }); + assert_eq!(cm.get_reservations().len(), 1); + assert_eq!(cm.get_reservations()[0].name, "r1"); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn concurrent_reservation_create_keeps_single_entry() { + let dir = TempDir::new().unwrap(); + let cm = test_cluster(&dir).await; + register_node(&cm, "n1", 8, 16000); + let now = chrono::Utc::now(); + let base = Reservation { + name: "r1".into(), + start_time: now, + end_time: now + chrono::Duration::hours(1), + nodes: vec!["n1".into()], + accounts: Vec::new(), + users: vec!["alice".into()], + flags: Default::default(), + }; + let cm1 = cm.clone(); + let cm2 = cm.clone(); + let r1 = base.clone(); + let r2 = base; + let (first, second) = tokio::join!( + tokio::task::spawn_blocking(move || cm1.create_reservation(r1)), + tokio::task::spawn_blocking(move || cm2.create_reservation(r2)), + ); + let outcomes = [first.unwrap(), second.unwrap()]; + assert_eq!( + outcomes.iter().filter(|r| r.is_ok()).count(), + 1, + "exactly one concurrent create must succeed" + ); + assert_eq!(cm.get_reservations().len(), 1); + assert_eq!(cm.get_reservations()[0].name, "r1"); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn reservation_survives_wal_replay() { + let dir = TempDir::new().unwrap(); + { + let cm = test_cluster(&dir).await; + register_node(&cm, "n1", 8, 16000); + let now = chrono::Utc::now(); + cm.create_reservation(Reservation { + name: "r1".into(), + start_time: now, + end_time: now + chrono::Duration::hours(1), + nodes: vec!["n1".into()], + accounts: Vec::new(), + users: vec!["alice".into()], + flags: Default::default(), + }) + .unwrap(); + assert_eq!(cm.get_reservations().len(), 1); + } + + let cm2 = test_cluster(&dir).await; + wait_for("reservation replayed from WAL", || { + cm2.get_reservations().iter().any(|r| r.name == "r1") + }); + assert_eq!(cm2.get_reservations().len(), 1); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn no_hold_jobs_reservation_delete_does_not_permanently_block_job() { + let dir = TempDir::new().unwrap(); + let cm = test_cluster(&dir).await; + register_node(&cm, "n1", 8, 16000); + let now = chrono::Utc::now(); + cm.create_reservation(Reservation { + name: "r1".into(), + start_time: now - chrono::Duration::minutes(5), + end_time: now + chrono::Duration::hours(2), + nodes: vec!["n1".into()], + accounts: Vec::new(), + users: vec!["testuser".into()], + flags: spur_core::reservation::ReservationFlags { + no_hold_jobs: true, + ..Default::default() + }, + }) + .unwrap(); + + let mut spec = basic_spec("resv-pending"); + spec.reservation = Some("r1".into()); + let job_id = submit_and_wait(&cm, spec); + + cm.delete_reservation("r1").unwrap(); + wait_for("reservation deleted", || cm.get_reservations().is_empty()); + + let job = cm.get_job(job_id).unwrap(); + assert_eq!(job.state, JobState::Pending); + assert_eq!(job.spec.reservation, None); + assert_ne!(job.pending_reason, PendingReason::Held); + assert_ne!(job.priority, 0); + assert!( + cm.pending_jobs().iter().any(|j| j.job_id == job_id), + "job must remain schedulable after no_hold_jobs delete" + ); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn no_hold_jobs_purge_does_not_permanently_block_job() { + let dir = TempDir::new().unwrap(); + let cm = test_cluster(&dir).await; + register_node(&cm, "n1", 8, 16000); + let now = chrono::Utc::now(); + cm.apply_operation(&WalOperation::ReservationCreate { + reservation: Reservation { + name: "r1".into(), + start_time: now - chrono::Duration::hours(2), + end_time: now - chrono::Duration::minutes(1), + nodes: vec!["n1".into()], + accounts: Vec::new(), + users: vec!["testuser".into()], + flags: spur_core::reservation::ReservationFlags { + no_hold_jobs: true, + ..Default::default() + }, + }, + }); + + let mut spec = basic_spec("resv-pending"); + spec.reservation = Some("r1".into()); + let job_id = submit_and_wait(&cm, spec); + + cm.purge_expired_reservations(); + wait_for("reservation purged", || cm.get_reservations().is_empty()); + + let job = cm.get_job(job_id).unwrap(); + assert_eq!(job.spec.reservation, None); + assert!( + cm.pending_jobs().iter().any(|j| j.job_id == job_id), + "job must remain schedulable after no_hold_jobs purge" + ); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn reservation_fence_reason_requires_all_candidates_blocked() { + let dir = TempDir::new().unwrap(); + let cm = test_cluster(&dir).await; + register_node(&cm, "n1", 8, 16000); + register_node(&cm, "n2", 8, 16000); + let now = chrono::Utc::now(); + cm.create_reservation(Reservation { + name: "r1".into(), + start_time: now - chrono::Duration::minutes(5), + end_time: now + chrono::Duration::hours(2), + nodes: vec!["n1".into()], + accounts: Vec::new(), + users: vec!["alice".into()], + flags: Default::default(), + }) + .unwrap(); + + let job_id = submit_and_wait(&cm, basic_spec("fence")); + cm.tag_blocked_pending_reasons(); + let job = cm.get_job(job_id).unwrap(); + assert_ne!( + job.pending_reason, + PendingReason::ReqNodeNotAvail, + "n2 is an unblocked candidate; must not fence the job" + ); + assert_ne!(job.pending_reason, PendingReason::ReservedMaintenance); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn job_state_change_held_pending_apply_is_atomic() { + let dir = TempDir::new().unwrap(); + let cm = test_cluster(&dir).await; + let job_id = submit_and_wait(&cm, basic_spec("hold-atomic")); + cm.apply_operation(&WalOperation::job_state_change( + job_id, + JobState::Pending, + JobState::Running, + )); + cm.apply_operation(&WalOperation::JobComplete { + job_id, + exit_code: -1, + state: JobState::Preempted, + }); + for _ in 0..5 { + cm.apply_operation(&WalOperation::job_state_change( + job_id, + JobState::Preempted, + JobState::Pending, + )); + cm.apply_operation(&WalOperation::job_state_change( + job_id, + JobState::Pending, + JobState::Running, + )); + cm.apply_operation(&WalOperation::JobComplete { + job_id, + exit_code: -1, + state: JobState::Preempted, + }); + } + assert_eq!(cm.get_job(job_id).unwrap().requeue_count, 5); + + cm.apply_operation(&WalOperation::job_state_change_held_pending( + job_id, + JobState::Preempted, + PendingReason::JobHoldMaxRequeue, + )); + let job = cm.get_job(job_id).unwrap(); + assert_eq!(job.state, JobState::Pending); + assert_eq!(job.priority, 0); + assert_eq!(job.pending_reason, PendingReason::JobHoldMaxRequeue); + assert!( + !cm.pending_jobs().iter().any(|j| j.job_id == job_id), + "hold must apply priority and reason in one WAL entry" + ); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn requeue_preempted_job_holds_at_max_requeue() { + let dir = TempDir::new().unwrap(); + let cm = test_cluster(&dir).await; + + let job_id = submit_and_wait(&cm, basic_spec("preempt-me")); + cm.apply_operation(&WalOperation::JobStateChange { + job_id, + old_state: JobState::Pending, + new_state: JobState::Running, + pending_reason: None, + pending_priority: None, + }); + cm.apply_operation(&WalOperation::JobComplete { + job_id, + exit_code: -1, + state: JobState::Preempted, + }); + for _ in 0..5 { + cm.apply_operation(&WalOperation::JobStateChange { + job_id, + old_state: JobState::Preempted, + new_state: JobState::Pending, + pending_reason: None, + pending_priority: None, + }); + cm.apply_operation(&WalOperation::JobStateChange { + job_id, + old_state: JobState::Pending, + new_state: JobState::Running, + pending_reason: None, + pending_priority: None, + }); + cm.apply_operation(&WalOperation::JobComplete { + job_id, + exit_code: -1, + state: JobState::Preempted, + }); + } + assert_eq!(cm.get_job(job_id).unwrap().requeue_count, 5); + + cm.requeue_preempted_job(job_id, true).unwrap(); + wait_for("job held at max requeue", || { + cm.get_job(job_id).is_some_and(|j| { + j.state == JobState::Pending && j.pending_reason == PendingReason::JobHoldMaxRequeue + }) + }); + let job = cm.get_job(job_id).unwrap(); + assert_eq!(job.priority, 0); + assert!( + !cm.pending_jobs().iter().any(|j| j.job_id == job_id), + "held job must not be schedulable" + ); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn reservation_delete_hold_survives_wal_replay() { + let dir = TempDir::new().unwrap(); + { + let cm = test_cluster(&dir).await; + register_node(&cm, "n1", 8, 16000); + let now = chrono::Utc::now(); + cm.create_reservation(Reservation { + name: "r1".into(), + start_time: now - chrono::Duration::minutes(5), + end_time: now + chrono::Duration::hours(2), + nodes: vec!["n1".into()], + accounts: Vec::new(), + users: vec!["testuser".into()], + flags: Default::default(), + }) + .unwrap(); + + let mut spec = basic_spec("resv-hold"); + spec.reservation = Some("r1".into()); + let job_id = submit_and_wait(&cm, spec); + cm.delete_reservation("r1").unwrap(); + wait_for("job held after delete", || { + cm.get_job(job_id).is_some_and(|j| { + j.pending_reason == PendingReason::ReservationDeleted && j.priority == 0 + }) + }); + } + + let cm2 = test_cluster(&dir).await; + wait_for("held job replayed from WAL", || { + cm2.jobs + .read() + .values() + .any(|j| j.pending_reason == PendingReason::ReservationDeleted && j.priority == 0) + }); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn reservation_delete_no_hold_survives_wal_replay() { + let dir = TempDir::new().unwrap(); + let job_id; + { + let cm = test_cluster(&dir).await; + register_node(&cm, "n1", 8, 16000); + let now = chrono::Utc::now(); + cm.create_reservation(Reservation { + name: "r1".into(), + start_time: now - chrono::Duration::minutes(5), + end_time: now + chrono::Duration::hours(2), + nodes: vec!["n1".into()], + accounts: Vec::new(), + users: vec!["testuser".into()], + flags: spur_core::reservation::ReservationFlags { + no_hold_jobs: true, + ..Default::default() + }, + }) + .unwrap(); + + let mut spec = basic_spec("resv-no-hold"); + spec.reservation = Some("r1".into()); + job_id = submit_and_wait(&cm, spec); + cm.delete_reservation("r1").unwrap(); + wait_for("reservation detached from job", || { + cm.get_job(job_id) + .is_some_and(|j| j.spec.reservation.is_none() && j.priority > 0) + }); + } + + let cm2 = test_cluster(&dir).await; + wait_for("detached job replayed from WAL", || { + cm2.get_job(job_id) + .is_some_and(|j| j.spec.reservation.is_none() && j.priority > 0) + }); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn release_job_after_reservation_delete_unblocks_job() { + let dir = TempDir::new().unwrap(); + let cm = test_cluster(&dir).await; + register_node(&cm, "n1", 8, 16000); + let now = chrono::Utc::now(); + cm.create_reservation(Reservation { + name: "r1".into(), + start_time: now - chrono::Duration::minutes(5), + end_time: now + chrono::Duration::hours(2), + nodes: vec!["n1".into()], + accounts: Vec::new(), + users: vec!["testuser".into()], + flags: Default::default(), + }) + .unwrap(); + + let mut spec = basic_spec("release-me"); + spec.reservation = Some("r1".into()); + let job_id = submit_and_wait(&cm, spec); + cm.delete_reservation("r1").unwrap(); + wait_for("job held after delete", || { + cm.get_job(job_id).is_some_and(|j| { + j.pending_reason == PendingReason::ReservationDeleted && j.priority == 0 + }) + }); + assert!( + !cm.pending_jobs().iter().any(|j| j.job_id == job_id), + "held job must not be schedulable before release" + ); + + cm.release_job(job_id).unwrap(); + wait_for("job released after reservation delete", || { + cm.get_job(job_id).is_some_and(|j| { + j.spec.reservation.is_none() + && j.priority > 0 + && j.pending_reason == PendingReason::Priority + }) + }); + assert!( + cm.pending_jobs().iter().any(|j| j.job_id == job_id), + "released job must be schedulable" + ); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn tag_blocked_preserves_reservation_deleted_reason() { + let dir = TempDir::new().unwrap(); + let cm = test_cluster(&dir).await; + let mut spec = basic_spec("resv-deleted"); + spec.reservation = Some("gone".into()); + let job_id = submit_and_wait(&cm, spec); + { + let mut jobs = cm.jobs.write(); + jobs.get_mut(&job_id).unwrap().pending_reason = PendingReason::ReservationDeleted; + jobs.get_mut(&job_id).unwrap().priority = 0; + } + + cm.tag_blocked_pending_reasons(); + assert_eq!( + cm.get_job(job_id).unwrap().pending_reason, + PendingReason::ReservationDeleted + ); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn maybe_requeue_holds_at_max_for_timeout() { + let dir = TempDir::new().unwrap(); + let cm = test_cluster(&dir).await; + register_node(&cm, "n1", 8, 16000); + + let mut spec = basic_spec("timeout-requeue"); + spec.requeue = true; + let job_id = submit_and_wait(&cm, spec); + let alloc = scalar_alloc(1, 1000); + let nodes = vec!["n1".into()]; + let per_node = per_node_for(&["n1"], alloc.clone()); + + for _ in 0..5 { + cm.start_job(job_id, nodes.clone(), alloc.clone(), per_node.clone()) + .unwrap(); + settle(&cm, job_id, JobState::Running); + cm.complete_job(job_id, -1, JobState::Timeout).unwrap(); + settle(&cm, job_id, JobState::Pending); + } + assert_eq!(cm.get_job(job_id).unwrap().requeue_count, 5); + + cm.start_job(job_id, nodes, alloc, per_node).unwrap(); + settle(&cm, job_id, JobState::Running); + cm.complete_job(job_id, -1, JobState::Timeout).unwrap(); + wait_for("job held at max requeue after timeout", || { + cm.get_job(job_id).is_some_and(|j| { + j.state == JobState::Pending && j.pending_reason == PendingReason::JobHoldMaxRequeue + }) + }); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn release_job_resets_requeue_count_after_max_hold() { + let dir = TempDir::new().unwrap(); + let cm = test_cluster(&dir).await; + + let job_id = submit_and_wait(&cm, basic_spec("release-me")); + cm.apply_operation(&WalOperation::JobStateChange { + job_id, + old_state: JobState::Pending, + new_state: JobState::Running, + pending_reason: None, + pending_priority: None, + }); + cm.apply_operation(&WalOperation::JobComplete { + job_id, + exit_code: -1, + state: JobState::Preempted, + }); + for _ in 0..5 { + cm.apply_operation(&WalOperation::JobStateChange { + job_id, + old_state: JobState::Preempted, + new_state: JobState::Pending, + pending_reason: None, + pending_priority: None, + }); + cm.apply_operation(&WalOperation::JobStateChange { + job_id, + old_state: JobState::Pending, + new_state: JobState::Running, + pending_reason: None, + pending_priority: None, + }); + cm.apply_operation(&WalOperation::JobComplete { + job_id, + exit_code: -1, + state: JobState::Preempted, + }); + } + cm.requeue_preempted_job(job_id, true).unwrap(); + wait_for("job held at max requeue", || { + cm.get_job(job_id) + .is_some_and(|j| j.pending_reason == PendingReason::JobHoldMaxRequeue) + }); + assert_eq!(cm.get_job(job_id).unwrap().requeue_count, 5); + + cm.release_job(job_id).unwrap(); + wait_for("release resets requeue budget", || { + cm.get_job(job_id).is_some_and(|j| { + j.requeue_count == 0 + && j.priority > 0 + && j.pending_reason == PendingReason::Priority + }) + }); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn cancel_preempted_job_uses_cancelled_state() { + let dir = TempDir::new().unwrap(); + let cm = test_cluster(&dir).await; + let job_id = submit_and_wait(&cm, basic_spec("cancel-preempted")); + + cm.apply_operation(&WalOperation::JobStateChange { + job_id, + old_state: JobState::Pending, + new_state: JobState::Running, + pending_reason: None, + pending_priority: None, + }); + cm.apply_operation(&WalOperation::JobComplete { + job_id, + exit_code: -1, + state: JobState::Preempted, + }); + assert_eq!(cm.get_job(job_id).unwrap().state, JobState::Preempted); + + cm.cancel_job(job_id, "testuser").unwrap(); + wait_for("preempted job cancelled", || { + cm.get_job(job_id) + .is_some_and(|j| j.state == JobState::Cancelled) + }); + let job = cm.get_job(job_id).unwrap(); + assert_ne!(job.pending_reason, PendingReason::JobHoldMaxRequeue); + } + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn tag_blocked_sets_licenses_reason() { let dir = TempDir::new().unwrap(); @@ -5832,6 +6663,8 @@ mod tests { job_id: 1, old_state: JobState::Pending, new_state: JobState::Running, + pending_reason: None, + pending_priority: None, }); let mut child = basic_spec("child"); @@ -6100,6 +6933,8 @@ mod tests { job_id: id, old_state: JobState::Timeout, new_state: JobState::Pending, + pending_reason: None, + pending_priority: None, }); let job = cm.get_job(id).unwrap(); @@ -6482,6 +7317,8 @@ mod tests { job_id: id, old_state: JobState::Pending, new_state: JobState::Running, + pending_reason: None, + pending_priority: None, }); } cm.apply_operation(&WalOperation::JobComplete { @@ -6540,6 +7377,8 @@ mod tests { job_id: 2, old_state: JobState::Pending, new_state: JobState::Running, + pending_reason: None, + pending_priority: None, }); let cancelled = cm.cancel_unsatisfiable_dependency_jobs(); @@ -6561,6 +7400,8 @@ mod tests { job_id: 1, old_state: JobState::Pending, new_state: JobState::Running, + pending_reason: None, + pending_priority: None, }); let mut child = basic_spec("child"); @@ -7120,6 +7961,8 @@ mod tests { job_id: id, old_state: JobState::Pending, new_state: JobState::Running, + pending_reason: None, + pending_priority: None, }); cm.apply_operation(&WalOperation::JobStart { job_id: id, diff --git a/crates/spurctld/src/raft.rs b/crates/spurctld/src/raft.rs index 2ec45d7d..dc84004d 100644 --- a/crates/spurctld/src/raft.rs +++ b/crates/spurctld/src/raft.rs @@ -48,6 +48,8 @@ pub struct JobFinalized { pub struct ClientResponse { #[serde(default, skip_serializing_if = "Vec::is_empty")] pub jobs_finalized: Vec, + #[serde(default)] + pub reservation_created: bool, } /// Trait for applying committed Raft entries to the cluster state. diff --git a/crates/spurctld/src/server.rs b/crates/spurctld/src/server.rs index d20297df..34ab5f97 100644 --- a/crates/spurctld/src/server.rs +++ b/crates/spurctld/src/server.rs @@ -17,7 +17,7 @@ use spur_proto::proto::slurm_controller_client::SlurmControllerClient; use spur_proto::proto::slurm_controller_server::{SlurmController, SlurmControllerServer}; use spur_proto::proto::*; -use crate::cluster::ClusterManager; +use crate::cluster::{ClusterManager, ReservationError}; use crate::raft::RaftHandle; use crate::rpc_middleware::RpcStatsLayer; use crate::rpc_stats::RpcStatsCollector; @@ -1937,15 +1937,12 @@ fn annotate_nodes_with_reservations( } } -#[allow(clippy::result_large_err)] -fn reservation_rpc_status(err: impl std::fmt::Display) -> Status { - let message = err.to_string(); - if message.contains("raft propose failed") { - Status::internal(message) - } else if message.contains("not found") { - Status::not_found(message) - } else { - Status::invalid_argument(message) +fn reservation_rpc_status(err: ReservationError) -> Status { + match err { + ReservationError::InvalidArgument(m) => Status::invalid_argument(m), + ReservationError::NotFound(m) => Status::not_found(m), + ReservationError::AlreadyExists(m) => Status::already_exists(m), + ReservationError::Raft(m) => Status::internal(m), } } diff --git a/crates/spurdbd/src/db.rs b/crates/spurdbd/src/db.rs index 3a8d0fb3..f7ea6b8f 100644 --- a/crates/spurdbd/src/db.rs +++ b/crates/spurdbd/src/db.rs @@ -129,6 +129,7 @@ CREATE INDEX IF NOT EXISTS idx_assoc_account ON associations(account); ALTER TABLE jobs ADD COLUMN IF NOT EXISTS exit_signal INTEGER NOT NULL DEFAULT 0; ALTER TABLE jobs ADD COLUMN IF NOT EXISTS derived_exit_code INTEGER NOT NULL DEFAULT 0; +ALTER TABLE jobs ADD COLUMN IF NOT EXISTS reservation TEXT NOT NULL DEFAULT ''; "#; /// Record a job start in the database. @@ -144,14 +145,16 @@ pub async fn record_job_start( cpus_per_task: i32, memory_mb: i64, start_time: DateTime, + reservation: &str, ) -> anyhow::Result<()> { sqlx::query( r#" - INSERT INTO jobs (job_id, user_name, account, partition_name, num_nodes, num_tasks, cpus_per_task, memory_mb, start_time, state) - VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, 'RUNNING') + INSERT INTO jobs (job_id, user_name, account, partition_name, num_nodes, num_tasks, cpus_per_task, memory_mb, start_time, state, reservation) + VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, 'RUNNING', $10) ON CONFLICT (job_id) DO UPDATE SET start_time = $9, - state = 'RUNNING' + state = 'RUNNING', + reservation = $10 "#, ) .bind(job_id) @@ -163,6 +166,7 @@ pub async fn record_job_start( .bind(cpus_per_task) .bind(memory_mb) .bind(start_time) + .bind(reservation) .execute(pool) .await?; Ok(()) @@ -276,6 +280,7 @@ pub struct JobRecord { pub submit_time: DateTime, pub start_time: Option>, pub end_time: Option>, + pub reservation: String, } /// Query job history. @@ -291,7 +296,7 @@ pub async fn get_job_history( let mut qb = QueryBuilder::::new( "SELECT job_id, name, user_name, account, partition_name, state, exit_code, \ exit_signal, derived_exit_code, num_nodes, num_tasks, nodelist, \ - submit_time, start_time, end_time \ + submit_time, start_time, end_time, reservation \ FROM jobs WHERE 1=1", ); @@ -340,6 +345,7 @@ pub async fn get_job_history( submit_time: row.get("submit_time"), start_time: row.get("start_time"), end_time: row.get("end_time"), + reservation: row.get("reservation"), }) .collect(); @@ -694,13 +700,52 @@ mod job_history_tests { delete_jobs(&pool, &ids).await.ok(); - record_job_start(&pool, id0, &user_a, &account_one, "debug", 1, 1, 1, 0, t1).await?; + record_job_start( + &pool, + id0, + &user_a, + &account_one, + "debug", + 1, + 1, + 1, + 0, + t1, + "", + ) + .await?; record_job_end(&pool, id0, "COMPLETED", 0, t1 + Duration::minutes(5), 0, 0).await?; - record_job_start(&pool, id1, &user_b, &account_one, "debug", 1, 1, 1, 0, t1).await?; + record_job_start( + &pool, + id1, + &user_b, + &account_one, + "debug", + 1, + 1, + 1, + 0, + t1, + "", + ) + .await?; record_job_end(&pool, id1, "FAILED", 137, t1 + Duration::minutes(5), 9, 137).await?; - record_job_start(&pool, id2, &user_a, &account_two, "debug", 1, 1, 1, 0, t2).await?; + record_job_start( + &pool, + id2, + &user_a, + &account_two, + "debug", + 1, + 1, + 1, + 0, + t2, + "", + ) + .await?; record_job_end(&pool, id2, "COMPLETED", 0, t2 + Duration::minutes(5), 0, 0).await?; let by_user = get_job_history(&pool, Some(&user_a), None, None, None, &[], 100).await?; diff --git a/crates/spurdbd/src/server.rs b/crates/spurdbd/src/server.rs index 6bc1af7f..f972fdff 100644 --- a/crates/spurdbd/src/server.rs +++ b/crates/spurdbd/src/server.rs @@ -47,6 +47,7 @@ impl SlurmAccounting for AccountingService { 1, memory_mb, start_time, + &req.reservation, ) .await .map_err(|e| Status::internal(e.to_string()))?; @@ -185,7 +186,7 @@ impl SlurmAccounting for AccountingService { qos: String::new(), array_job_id: 0, array_task_id: 0, - reservation: String::new(), + reservation: r.reservation.clone(), }) .collect(); diff --git a/examples/spur.conf b/examples/spur.conf index 57dbc8fb..c2338a44 100644 --- a/examples/spur.conf +++ b/examples/spur.conf @@ -3,6 +3,8 @@ cluster_name = "mi300x-cluster" [controller] listen_addr = "[::]:6817" state_dir = "/var/spool/spur" +# Maximum automatic requeues before a job is held (default 5). +max_batch_requeue = 5 [scheduler] plugin = "backfill" diff --git a/proto/slurm.proto b/proto/slurm.proto index d711ef58..13a0bb81 100644 --- a/proto/slurm.proto +++ b/proto/slurm.proto @@ -196,6 +196,8 @@ message JobInfo { // Array jobs uint32 array_job_id = 50; uint32 array_task_id = 51; + + // Reservation string reservation = 52; } @@ -700,6 +702,7 @@ message RecordJobStartRequest { string partition = 4; ResourceAllocations resources = 5; google.protobuf.Timestamp start_time = 6; + string reservation = 7; } message RecordJobEndRequest { diff --git a/tests/native_host/e2e/test_reservations.py b/tests/native_host/e2e/test_reservations.py index 1fd2e284..cfb9deb7 100644 --- a/tests/native_host/e2e/test_reservations.py +++ b/tests/native_host/e2e/test_reservations.py @@ -89,6 +89,90 @@ def test_reservation_job_schedules_for_authorized_user(self, cluster): content = cluster.read_output_on_any_node(out_path) assert "RES_OK" in content + def test_hold_on_delete_and_release(self, cluster): + res_name = f"res-hold-{int(time.time())}" + node = cluster.node_names[0] + cluster.scontrol( + "create-reservation", + f"--name={res_name}", + "--start-time=now", + "--duration=60", + f"--nodes={node}", + "--users=testuser", + ) + + script = cluster.write_file("res-hold.sh", "#!/bin/bash\necho HOLD_RELEASE_OK\n") + out_path = f"{cluster.remote_dir}/res-hold.out" + sb = cluster.sbatch( + [ + "-N", + "1", + f"--reservation={res_name}", + "-w", + node, + "-t", + "1", + "-o", + out_path, + script, + ] + ) + job_id = parse_job_id(sb) + assert job_id is not None + wait_job_state(cluster, job_id, "PD", timeout=30) + + cluster.scontrol("delete-reservation", res_name) + + wait_job_state(cluster, job_id, "PD", timeout=30) + held = cluster.squeue(["-j", str(job_id), "-o", "%T %r"]) + assert "PD" in held + assert "ReservationDeleted" in held + + cluster.scontrol("release", str(job_id)) + + state = wait_job(cluster, job_id, timeout=60) + assert state in ("CD", "GONE"), f"expected completed after release, got {state}" + content = cluster.read_output_on_any_node(out_path) + assert "HOLD_RELEASE_OK" in content + + def test_no_hold_jobs_delete(self, cluster): + res_name = f"res-nohold-{int(time.time())}" + node = cluster.node_names[0] + cluster.scontrol( + "create-reservation", + f"--name={res_name}", + "--start-time=now", + "--duration=60", + f"--nodes={node}", + "--users=testuser", + "--flags=no_hold_jobs", + ) + + script = cluster.write_file("res-nohold.sh", "#!/bin/bash\nsleep 120\n") + sb = cluster.sbatch( + [ + "-N", + "1", + f"--reservation={res_name}", + "-w", + node, + "-t", + "1", + script, + ] + ) + job_id = parse_job_id(sb) + assert job_id is not None + wait_job_state(cluster, job_id, "PD", timeout=30) + + cluster.scontrol("delete-reservation", res_name) + + wait_job_state(cluster, job_id, "PD", timeout=30) + show = cluster.squeue(["-j", str(job_id), "-o", "%T %r %v"]) + assert "PD" in show + assert "Held" not in show + assert res_name not in show + def test_create_rejects_busy_node_without_ignore_jobs(self, cluster): node = cluster.node_names[0] long_script = cluster.write_file("res-long.sh", "#!/bin/bash\nsleep 300\n")