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/scontrol.rs b/crates/spur-cli/src/scontrol.rs index 9b35c05a..28906418 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, overlap) + #[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-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..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(), } } } @@ -381,6 +391,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 +421,7 @@ impl Default for SchedulerConfig { fairshare_halflife_days: 14, default_time_limit_minutes: 60, complete_wait_secs: 300, + resv_overrun_minutes: 0, } } } @@ -790,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(()) } @@ -1163,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 964ce8c1..0803ca21 100644 --- a/crates/spur-core/src/job.rs +++ b/crates/spur-core/src/job.rs @@ -251,9 +251,20 @@ pub enum PendingReason { QosGrpNodeLimit, BurstBufferResources, 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", @@ -294,6 +305,9 @@ impl PendingReason { Self::QosGrpNodeLimit => "QOSGrpNodeLimit", Self::BurstBufferResources => "BurstBufferResources", Self::BurstBufferStageIn => "BurstBufferStageIn", + Self::ReservedMaintenance => "ReqNodeNotAvail, Reserved for maintenance", + Self::ReservationDeleted => "ReservationDeleted", + Self::JobHoldMaxRequeue => "JobHoldMaxRequeue", } } } @@ -778,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, @@ -1170,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 763e35f4..e8368219 100644 --- a/crates/spur-core/src/reservation.rs +++ b/crates/spur-core/src/reservation.rs @@ -1,9 +1,73 @@ // 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, + #[serde(default)] + pub overlap: bool, +} + +impl ReservationFlags { + pub fn parse_list(values: &[String]) -> Result { + 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, + "overlap" => flags.overlap = true, + "" => {} + other => return Err(format!("unknown reservation flag '{other}'")), + } + } + } + Ok(flags) + } + + pub fn parse_csv(csv: &str) -> Result { + if csv.is_empty() { + return Ok(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"); + } + if self.overlap { + parts.push("OVERLAP"); + } + parts.join(",") + } +} + /// A resource reservation that blocks nodes for specific users/accounts /// during a time window. #[derive(Debug, Clone, Serialize, Deserialize)] @@ -14,23 +78,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; @@ -44,10 +125,174 @@ 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.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], + 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 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 { + return Some((job.job_id, node)); + } + } + None +} + +/// Whether a job without reservation access would overlap this reservation window. +pub fn prospective_overlap_at( + job: &Job, + reservation: &Reservation, + node: &str, + at: 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 = 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)] mod tests { use super::*; - use chrono::Duration; + use crate::job::JobSpec; fn make_reservation() -> Reservation { let now = Utc::now(); @@ -58,6 +303,7 @@ mod tests { nodes: vec!["node001".into(), "node002".into()], accounts: vec!["research".into()], users: vec!["alice".into()], + flags: ReservationFlags::default(), } } @@ -96,7 +342,131 @@ 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,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(); + 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] + 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(); + 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..03ad5c6e 100644 --- a/crates/spur-core/src/wal.rs +++ b/crates/spur-core/src/wal.rs @@ -4,8 +4,9 @@ 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; use crate::resource::{ResourceAllocations, ResourceSet}; @@ -26,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, @@ -57,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, @@ -118,6 +134,146 @@ 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, + }, +} + +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::*; + 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-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 7834b673..395c9590 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,20 @@ 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 { + reservations + .iter() + .any(|res| reservation::prospective_overlap_at(job, res, node, start, duration)) + } + /// Find nodes that satisfy a job's resource requirements. fn find_suitable_nodes( &self, @@ -126,6 +140,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 +149,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 +160,22 @@ impl BackfillScheduler { { return false; } + if let Some(res) = reservations.iter().find(|r| r.name == res_name) { + let earliest = now.max(res.start_time); + if earliest + job_duration > 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 +273,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 +1098,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 +1222,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 +1241,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..20e183ce 100644 --- a/crates/spur-tests/src/t55_format.rs +++ b/crates/spur-tests/src/t55_format.rs @@ -121,6 +121,12 @@ 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::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 e152e558..9ff20815 100644 --- a/crates/spurctld/src/cluster.rs +++ b/crates/spurctld/src/cluster.rs @@ -18,7 +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::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; @@ -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,16 +756,50 @@ 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(()) } + /// Return a preempted job to the pending queue. + pub fn requeue_preempted_job(&self, job_id: JobId, force: bool) -> anyhow::Result<()> { + 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 { + 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(()); + } + if !force && !job.spec.requeue { + return Ok(()); + } + job.state + }; + + self.propose(WalOperation::job_state_change( + job_id, + old_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 @@ -729,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 }; @@ -742,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(()) @@ -922,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(()) } @@ -985,6 +1141,9 @@ impl ClusterManager { job_id, old_priority: old, new_priority: p, + pending_reason: None, + reset_requeue_count: false, + clear_reservation: false, })?; } @@ -1252,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(); @@ -1347,6 +1506,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 @@ -1366,8 +1526,13 @@ 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 @@ -1529,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 }) @@ -1635,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; } @@ -1656,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; @@ -1750,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| { @@ -1778,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; @@ -1787,18 +1952,34 @@ 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) { - anyhow::bail!("reservation '{}' already exists", res.name); + /// Create a new reservation (validated, persisted via Raft). + pub fn create_reservation(&self, mut res: Reservation) -> Result<(), ReservationError> { + if self.reservations.read().iter().any(|r| r.name == 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(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 + ))); } - info!(name = %res.name, "reservation created"); - reservations.push(res); Ok(()) } - /// Update an existing reservation. + /// Update an existing reservation (validated, persisted via Raft). #[allow(clippy::too_many_arguments)] pub fn update_reservation( &self, @@ -1810,13 +1991,250 @@ impl ClusterManager { remove_users: &[String], add_accounts: &[String], remove_accounts: &[String], - ) -> anyhow::Result<()> { - let mut reservations = self.reservations.write(); - let res = reservations - .iter_mut() + ) -> Result<(), ReservationError> { + let mut preview = self + .reservations + .read() + .iter() .find(|r| r.name == name) - .ok_or_else(|| anyhow::anyhow!("reservation '{}' not found", name))?; + .cloned() + .ok_or_else(|| { + ReservationError::not_found(format!("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(ReservationError::invalid)?, + ); + } + 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)) + .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(), + 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(), + }) + .map_err(|e| ReservationError::raft(e.to_string()))?; + Ok(()) + } + + /// Delete a reservation by name (persisted via Raft). + 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!( + job.state, + JobState::Running | JobState::Completing | JobState::Suspended + ) { + continue; + } + if job.spec.reservation.as_deref() == Some(name) { + return Err(ReservationError::invalid(format!( + "reservation '{}' in use by running job {}", + name, job.job_id + ))); + } + } + + self.propose(WalOperation::ReservationDelete { + name: name.to_string(), + }) + .map_err(|e| ReservationError::raft(e.to_string()))?; + 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 (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() + .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 now > res.end_time + grace { + 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 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_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; + } + 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; + } + } + } + #[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 +2256,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. @@ -1886,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 @@ -1897,6 +2300,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(); @@ -2265,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) { @@ -2273,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; + } } } } @@ -2458,6 +2876,7 @@ impl ClusterManager { state: final_state, exit_code: final_exit, }], + ..Default::default() }; } } @@ -2567,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 { @@ -2727,6 +3158,62 @@ impl ClusterManager { t.revoked = true; } } + WalOperation::ReservationCreate { reservation } => { + 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, + 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 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 { + 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"); + } + } + } } self.next_job_id.store(next_id, Ordering::Relaxed); response @@ -2853,11 +3340,109 @@ impl StateMachineApply for ClusterManager { } } -/// `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. -fn reservation_block( - job: &Job, +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 all_blocked = true; + + for node_name in &candidates { + let mut node_blocked = false; + for res in cluster_state.reservations { + 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 { + Some(PendingReason::ReqNodeNotAvail) + } +} + +/// `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. +fn reservation_block( + job: &Job, reservations: &[Reservation], now: chrono::DateTime, ) -> Option { @@ -3386,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(); @@ -3434,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 { @@ -3472,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 }); @@ -3636,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. @@ -3673,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); @@ -3795,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(); @@ -3909,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 { @@ -3948,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 { @@ -3988,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 { @@ -4052,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, @@ -4111,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 { @@ -4139,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 { @@ -4187,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 { @@ -4224,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 { @@ -4279,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 { @@ -4314,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, @@ -4354,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, @@ -4387,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, @@ -4423,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 { @@ -4488,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 { @@ -4838,6 +5464,822 @@ 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 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 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(); @@ -5221,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"); @@ -5489,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(); @@ -5871,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 { @@ -5929,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(); @@ -5950,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"); @@ -6509,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/scheduler_loop.rs b/crates/spurctld/src/scheduler_loop.rs index e368a5a7..dd675344 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(). @@ -409,10 +411,23 @@ 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 cluster_nodes = cluster.get_nodes(); + + 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() @@ -420,30 +435,108 @@ 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 !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) + .unwrap_or(1); + if pending_tier <= candidate_tier { + continue; } - break; // One preemption per cycle, re-evaluate next cycle + } + + 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::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::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 { + warn!( + job_id = candidate.job_id, + error = %e, + "failed to preempt job" + ); + } else { + break; } } } } +/// 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 a2d77a95..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; @@ -1182,6 +1182,9 @@ 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) + .map_err(Status::invalid_argument)?; + let reservation = spur_core::reservation::Reservation { name: req.name, start_time, @@ -1189,11 +1192,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(reservation_rpc_status)?; Ok(Response::new(())) } @@ -1229,7 +1233,7 @@ impl SlurmController for ControllerService { &req.add_accounts, &req.remove_accounts, ) - .map_err(|e| Status::internal(e.to_string()))?; + .map_err(reservation_rpc_status)?; Ok(Response::new(())) } @@ -1255,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(())) } @@ -1273,6 +1277,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 +1287,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 { @@ -1762,6 +1769,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(), } } @@ -1783,6 +1791,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,15 +1922,30 @@ fn annotate_nodes_with_reservations( now: DateTime, ) { 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(); - break; + if node_info.active_reservation.is_empty() { + node_info.active_reservation = res.name.clone(); + } + if res.flags.maint { + node_info.reservation_maint = true; + } } } } } +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), + } +} + fn node_complete_to_status(err: NodeCompleteError) -> Status { let message = err.to_string(); let code = match err { @@ -1946,6 +1970,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 { @@ -1969,6 +1994,7 @@ mod tests { nodes: nodes.iter().map(|s| s.to_string()).collect(), accounts: Vec::new(), users: Vec::new(), + flags: Default::default(), } } @@ -2030,6 +2056,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/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 99837268..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,6 +186,7 @@ impl SlurmAccounting for AccountingService { qos: String::new(), array_job_id: 0, array_task_id: 0, + 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 6120e0f1..13a0bb81 100644 --- a/proto/slurm.proto +++ b/proto/slurm.proto @@ -196,6 +196,9 @@ message JobInfo { // Array jobs uint32 array_job_id = 50; uint32 array_task_id = 51; + + // Reservation + string reservation = 52; } message NodeInfo { @@ -220,6 +223,7 @@ message NodeInfo { string active_reservation = 50; map labels = 51; + bool reservation_maint = 52; } message PartitionInfo { @@ -698,6 +702,7 @@ message RecordJobStartRequest { string partition = 4; ResourceAllocations resources = 5; google.protobuf.Timestamp start_time = 6; + string reservation = 7; } message RecordJobEndRequest { @@ -913,6 +918,7 @@ message CreateReservationRequest { repeated string nodes = 4; repeated string accounts = 5; repeated string users = 6; + repeated string flags = 7; } message UpdateReservationRequest { @@ -943,6 +949,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/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 new file mode 100644 index 00000000..cfb9deb7 --- /dev/null +++ b/tests/native_host/e2e/test_reservations.py @@ -0,0 +1,197 @@ +# 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, wait_job_state + + +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_name}", + "--start-time=now", + "--duration=60", + f"--nodes={node}", + "--users=testuser", + ) + 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] + 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 + + 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())}" + 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_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") + sb = cluster.sbatch(["-N", "1", "-w", node, "-t", "10", long_script]) + job_id = parse_job_id(sb) + assert job_id is not None + + wait_job_state(cluster, job_id, "R", timeout=30) + + res_name = f"res-busy-{int(time.time())}" + out = cluster.cli_allow_fail( + [ + "scontrol", + "create-reservation", + f"--name={res_name}", + "--start-time=now", + "--duration=10", + f"--nodes={node}", + ] + ) + msg = out.lower() + assert "busy" in msg or "until after reservation start" in msg, f"unexpected: {out}"