diff --git a/Cargo.lock b/Cargo.lock index db3b3890..941b7c47 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3475,6 +3475,7 @@ dependencies = [ "prost-types", "serde", "serde_json", + "spur-client", "spur-core", "spur-net", "spur-proto", @@ -3487,6 +3488,17 @@ dependencies = [ "whoami", ] +[[package]] +name = "spur-client" +version = "0.3.0" +dependencies = [ + "tokio", + "tokio-stream", + "tonic", + "tonic-health", + "tracing", +] + [[package]] name = "spur-core" version = "0.3.0" @@ -3534,6 +3546,7 @@ name = "spur-ffi" version = "0.3.0" dependencies = [ "prost-types", + "spur-client", "spur-core", "spur-proto", "tokio", @@ -3712,6 +3725,7 @@ dependencies = [ "serde", "serde_json", "shlex", + "spur-client", "spur-core", "spur-devices", "spur-net", @@ -4022,7 +4036,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd" dependencies = [ "fastrand", - "getrandom 0.3.4", + "getrandom 0.4.3", "once_cell", "rustix", "windows-sys 0.61.2", @@ -4185,6 +4199,7 @@ dependencies = [ "futures-core", "pin-project-lite", "tokio", + "tokio-util", ] [[package]] @@ -4305,6 +4320,19 @@ dependencies = [ "syn 2.0.118", ] +[[package]] +name = "tonic-health" +version = "0.14.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fcfab99db777fba2802f0dfa861d1628d1ae916fb199d29819941f139ae85082" +dependencies = [ + "prost", + "tokio", + "tokio-stream", + "tonic", + "tonic-prost", +] + [[package]] name = "tonic-prost" version = "0.14.6" diff --git a/Cargo.toml b/Cargo.toml index a826fda3..0c368440 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -3,6 +3,7 @@ resolver = "2" members = [ "crates/spur-proto", "crates/spur-core", + "crates/spur-client", "crates/spur-net", "crates/spur-sched", "crates/spur-metrics", diff --git a/crates/spur-cli/Cargo.toml b/crates/spur-cli/Cargo.toml index 0e340d14..31aef687 100644 --- a/crates/spur-cli/Cargo.toml +++ b/crates/spur-cli/Cargo.toml @@ -11,6 +11,7 @@ path = "src/main.rs" [dependencies] spur-proto = { path = "../spur-proto" } +spur-client = { path = "../spur-client" } spur-core = { path = "../spur-core" } spur-update = { path = "../spur-update" } spur-net = { path = "../spur-net" } diff --git a/crates/spur-cli/src/exec.rs b/crates/spur-cli/src/exec.rs index 3a0f01fc..f447515a 100644 --- a/crates/spur-cli/src/exec.rs +++ b/crates/spur-cli/src/exec.rs @@ -38,9 +38,10 @@ pub async fn main() -> Result<()> { pub async fn main_with_args(args: Vec) -> Result<()> { let args = ExecArgs::try_parse_from(&args)?; - let mut client = SlurmControllerClient::connect(args.controller.clone()) + let channel = spur_client::connect_channel(&args.controller) .await .context("failed to connect to controller")?; + let mut client = SlurmControllerClient::new(channel); let resp = client .exec_in_job(ExecInJobRequest { diff --git a/crates/spur-cli/src/main.rs b/crates/spur-cli/src/main.rs index 1088dbbf..1c307c9f 100644 --- a/crates/spur-cli/src/main.rs +++ b/crates/spur-cli/src/main.rs @@ -49,21 +49,10 @@ fn load_controller_addr_from_config() { } if let Ok(config) = spur_core::config::SlurmConfig::load_from_file(&config_path) { - // Extract host and port from the config - let host = config - .controller - .hosts - .first() - .map(|h| h.as_str()) - .unwrap_or("localhost"); - let port = config - .controller - .listen_addr - .rsplit(':') - .next() - .unwrap_or("6817"); - let addr = format!("http://{}:{}", host, port); - std::env::set_var("SPUR_CONTROLLER_ADDR", &addr); + let endpoints = config.controller.endpoints(); + if !endpoints.is_empty() { + std::env::set_var("SPUR_CONTROLLER_ADDR", endpoints.join(",")); + } } } diff --git a/crates/spur-cli/src/node.rs b/crates/spur-cli/src/node.rs index efe3aaf0..7a6fed98 100644 --- a/crates/spur-cli/src/node.rs +++ b/crates/spur-cli/src/node.rs @@ -104,7 +104,7 @@ fn parse_label_args(label_args: &[String]) -> Result<(HashMap, V } async fn cmd_label(controller: &str, node: String, label_args: Vec) -> Result<()> { - let mut client = SlurmControllerClient::connect(controller.to_string()).await?; + let mut client = SlurmControllerClient::new(spur_client::connect_channel(controller).await?); let (set_labels, remove_labels) = parse_label_args(&label_args)?; client @@ -128,7 +128,7 @@ async fn cmd_label(controller: &str, node: String, label_args: Vec) -> R } async fn cmd_drain(controller: &str, node: String, reason: Option) -> Result<()> { - let mut client = SlurmControllerClient::connect(controller.to_string()).await?; + let mut client = SlurmControllerClient::new(spur_client::connect_channel(controller).await?); let resp = client .drain_node(spur_proto::proto::DrainNodeRequest { name: node.clone(), @@ -158,7 +158,7 @@ async fn cmd_remove( force: bool, reason: Option, ) -> Result<()> { - let mut client = SlurmControllerClient::connect(controller.to_string()).await?; + let mut client = SlurmControllerClient::new(spur_client::connect_channel(controller).await?); let resp = client .deregister_node(spur_proto::proto::DeregisterNodeRequest { name: node.clone(), diff --git a/crates/spur-cli/src/salloc.rs b/crates/spur-cli/src/salloc.rs index b0dd64dc..96707923 100644 --- a/crates/spur-cli/src/salloc.rs +++ b/crates/spur-cli/src/salloc.rs @@ -109,9 +109,10 @@ pub async fn main_with_args(args: Vec) -> Result<()> { let ntasks = args.ntasks; let cpus_per_task = args.cpus_per_task; - let mut client = SlurmControllerClient::connect(controller) + let channel = spur_client::connect_channel(&controller) .await .context("failed to connect to spurctld")?; + let mut client = SlurmControllerClient::new(channel); // Submit interactive allocation (sleep infinity holds the allocation) let job_spec = JobSpec { diff --git a/crates/spur-cli/src/sattach.rs b/crates/spur-cli/src/sattach.rs index 6d9efb01..3c93ab7b 100644 --- a/crates/spur-cli/src/sattach.rs +++ b/crates/spur-cli/src/sattach.rs @@ -47,9 +47,10 @@ pub async fn main_with_args(args: Vec) -> Result<()> { .and_then(|s| s.parse().ok()) .context("sattach: invalid job ID format (expected JOB_ID or JOB_ID.STEP_ID)")?; - let mut client = SlurmControllerClient::connect(args.controller) + let channel = spur_client::connect_channel(&args.controller) .await .context("failed to connect to spurctld")?; + let mut client = SlurmControllerClient::new(channel); // Look up the job to find which node it is running on let job = client diff --git a/crates/spur-cli/src/sbatch.rs b/crates/spur-cli/src/sbatch.rs index 09a77791..9829241e 100644 --- a/crates/spur-cli/src/sbatch.rs +++ b/crates/spur-cli/src/sbatch.rs @@ -646,9 +646,10 @@ pub async fn main_with_args(cli_args: Vec) -> Result<()> { }; // Submit to controller - let mut client = SlurmControllerClient::connect(args.controller) + let channel = spur_client::connect_channel(&args.controller) .await .context("failed to connect to spurctld")?; + let mut client = SlurmControllerClient::new(channel); let response = client .submit_job(SubmitJobRequest { diff --git a/crates/spur-cli/src/scancel.rs b/crates/spur-cli/src/scancel.rs index 6fef65c7..20d43e2e 100644 --- a/crates/spur-cli/src/scancel.rs +++ b/crates/spur-cli/src/scancel.rs @@ -75,9 +75,10 @@ pub async fn main_with_args(args: Vec) -> Result<()> { .user .unwrap_or_else(|| whoami::username().unwrap_or_else(|_| "unknown".into())); - let mut client = SlurmControllerClient::connect(args.controller) + let channel = spur_client::connect_channel(&args.controller) .await .context("failed to connect to spurctld")?; + let mut client = SlurmControllerClient::new(channel); if !args.job_ids.is_empty() { // Cancel specific jobs diff --git a/crates/spur-cli/src/scontrol.rs b/crates/spur-cli/src/scontrol.rs index 9b35c05a..1447575a 100644 --- a/crates/spur-cli/src/scontrol.rs +++ b/crates/spur-cli/src/scontrol.rs @@ -168,9 +168,10 @@ pub async fn main_with_args(args: Vec) -> Result<()> { } ScontrolCommand::Requeue { job_id } => { // Requeue = cancel + resubmit, simplified for now - let mut client = SlurmControllerClient::connect(args.controller.to_string()) + let channel = spur_client::connect_channel(&args.controller) .await .context("failed to connect to spurctld")?; + let mut client = SlurmControllerClient::new(channel); client .cancel_job(spur_proto::proto::CancelJobRequest { job_id, @@ -183,9 +184,10 @@ pub async fn main_with_args(args: Vec) -> Result<()> { Ok(()) } ScontrolCommand::Suspend { job_id } => { - let mut client = SlurmControllerClient::connect(args.controller.to_string()) + let channel = spur_client::connect_channel(&args.controller) .await .context("failed to connect to spurctld")?; + let mut client = SlurmControllerClient::new(channel); client .suspend_job(spur_proto::proto::SuspendJobRequest { job_id, @@ -197,9 +199,10 @@ pub async fn main_with_args(args: Vec) -> Result<()> { Ok(()) } ScontrolCommand::Resume { job_id } => { - let mut client = SlurmControllerClient::connect(args.controller.to_string()) + let channel = spur_client::connect_channel(&args.controller) .await .context("failed to connect to spurctld")?; + let mut client = SlurmControllerClient::new(channel); client .resume_job(spur_proto::proto::ResumeJobRequest { job_id, @@ -246,9 +249,10 @@ pub async fn main_with_args(args: Vec) -> Result<()> { .filter(|s| !s.is_empty()) .collect() }; - let mut client = SlurmControllerClient::connect(args.controller.to_string()) + let channel = spur_client::connect_channel(&args.controller) .await .context("failed to connect to spurctld")?; + let mut client = SlurmControllerClient::new(channel); client .update_reservation(spur_proto::proto::UpdateReservationRequest { name: name.clone(), @@ -272,9 +276,10 @@ pub async fn main_with_args(args: Vec) -> Result<()> { } async fn show(controller: &str, entity: &str, name: Option<&str>) -> Result<()> { - let mut client = SlurmControllerClient::connect(controller.to_string()) + let channel = spur_client::connect_channel(controller) .await .context("failed to connect to spurctld")?; + let mut client = SlurmControllerClient::new(channel); match entity.to_lowercase().as_str() { "job" | "jobs" => { @@ -496,9 +501,10 @@ async fn show(controller: &str, entity: &str, name: Option<&str>) -> Result<()> } async fn ping(controller: &str) -> Result<()> { - let mut client = SlurmControllerClient::connect(controller.to_string()) + let channel = spur_client::connect_channel(controller) .await .context("failed to connect to spurctld")?; + let mut client = SlurmControllerClient::new(channel); let resp = client.ping(()).await.context("ping failed")?; @@ -537,9 +543,10 @@ fn format_ts(ts: Option<&prost_types::Timestamp>) -> String { async fn send_job_update(controller: &str, req: spur_proto::proto::UpdateJobRequest) -> Result<()> { let hold = req.hold; let job_id = req.job_id; - let mut client = SlurmControllerClient::connect(controller.to_string()) + let channel = spur_client::connect_channel(controller) .await .context("failed to connect to spurctld")?; + let mut client = SlurmControllerClient::new(channel); client.update_job(req).await.context("update failed")?; @@ -624,9 +631,10 @@ async fn update_node( state: Option<&str>, reason: Option, ) -> Result<()> { - let mut client = SlurmControllerClient::connect(controller.to_string()) + let channel = spur_client::connect_channel(controller) .await .context("failed to connect to spurctld")?; + let mut client = SlurmControllerClient::new(channel); let proto_state = state.map(|s| match s.to_lowercase().as_str() { "idle" | "resume" => spur_proto::proto::NodeState::NodeIdle as i32, @@ -666,9 +674,10 @@ async fn create_reservation( accounts: &str, users: &str, ) -> Result<()> { - let mut client = SlurmControllerClient::connect(controller.to_string()) + let channel = spur_client::connect_channel(controller) .await .context("failed to connect to spurctld")?; + let mut client = SlurmControllerClient::new(channel); let node_list: Vec = nodes .split(',') @@ -704,9 +713,10 @@ async fn create_reservation( /// Delete a reservation via the controller. async fn delete_reservation(controller: &str, name: &str) -> Result<()> { - let mut client = SlurmControllerClient::connect(controller.to_string()) + let channel = spur_client::connect_channel(controller) .await .context("failed to connect to spurctld")?; + let mut client = SlurmControllerClient::new(channel); client .delete_reservation(spur_proto::proto::DeleteReservationRequest { diff --git a/crates/spur-cli/src/sdiag.rs b/crates/spur-cli/src/sdiag.rs index 4e38df95..fa82fe22 100644 --- a/crates/spur-cli/src/sdiag.rs +++ b/crates/spur-cli/src/sdiag.rs @@ -36,9 +36,10 @@ pub async fn main() -> Result<()> { pub async fn main_with_args(args: Vec) -> Result<()> { let args = SdiagArgs::try_parse_from(&args)?; - let mut client = SlurmControllerClient::connect(args.controller.clone()) + let channel = spur_client::connect_channel(&args.controller) .await .context("failed to connect to spurctld")?; + let mut client = SlurmControllerClient::new(channel); if args.reset { client diff --git a/crates/spur-cli/src/sinfo.rs b/crates/spur-cli/src/sinfo.rs index 0d28e297..d9ef0ca9 100644 --- a/crates/spur-cli/src/sinfo.rs +++ b/crates/spur-cli/src/sinfo.rs @@ -70,9 +70,10 @@ pub async fn main_with_args(args: Vec) -> Result<()> { let fields = format_engine::parse_format(&fmt, &format_engine::sinfo_header); - let mut client = SlurmControllerClient::connect(args.controller) + let channel = spur_client::connect_channel(&args.controller) .await .context("failed to connect to spurctld")?; + let mut client = SlurmControllerClient::new(channel); // Get partitions let partitions_resp = client diff --git a/crates/spur-cli/src/smd.rs b/crates/spur-cli/src/smd.rs index e217414a..ace211e5 100644 --- a/crates/spur-cli/src/smd.rs +++ b/crates/spur-cli/src/smd.rs @@ -42,9 +42,10 @@ pub async fn main_with_args(args: Vec) -> Result<()> { let args = SmdArgs::try_parse_from(&args)?; loop { - let mut client = SlurmControllerClient::connect(args.controller.clone()) + let channel = spur_client::connect_channel(&args.controller) .await .context("failed to connect to spurctld")?; + let mut client = SlurmControllerClient::new(channel); let nodes = client .get_nodes(GetNodesRequest { diff --git a/crates/spur-cli/src/sprio.rs b/crates/spur-cli/src/sprio.rs index e7229fab..6f14e7dd 100644 --- a/crates/spur-cli/src/sprio.rs +++ b/crates/spur-cli/src/sprio.rs @@ -53,9 +53,10 @@ pub async fn main_with_args(args: Vec) -> Result<()> { }) .unwrap_or_default(); - let mut client = SlurmControllerClient::connect(args.controller.clone()) + let channel = spur_client::connect_channel(&args.controller) .await .context("failed to connect to spurctld")?; + let mut client = SlurmControllerClient::new(channel); // Get pending jobs only (priority is relevant for pending jobs) let response = client diff --git a/crates/spur-cli/src/squeue.rs b/crates/spur-cli/src/squeue.rs index 420b3ce1..de81fc25 100644 --- a/crates/spur-cli/src/squeue.rs +++ b/crates/spur-cli/src/squeue.rs @@ -93,9 +93,10 @@ pub async fn main_with_args(args: Vec) -> Result<()> { .unwrap_or_default(); // Connect and fetch - let mut client = SlurmControllerClient::connect(args.controller) + let channel = spur_client::connect_channel(&args.controller) .await .context("failed to connect to spurctld")?; + let mut client = SlurmControllerClient::new(channel); let response = client .get_jobs(GetJobsRequest { diff --git a/crates/spur-cli/src/srun.rs b/crates/spur-cli/src/srun.rs index 703054e0..998a4691 100644 --- a/crates/spur-cli/src/srun.rs +++ b/crates/spur-cli/src/srun.rs @@ -218,9 +218,10 @@ pub async fn main_with_args(args: Vec) -> Result<()> { .unwrap_or(0); // Submit as a batch job - let mut client = SlurmControllerClient::connect(args.controller.clone()) + let channel = spur_client::connect_channel(&args.controller) .await .context("failed to connect to spurctld")?; + let mut client = SlurmControllerClient::new(channel); let job_spec = JobSpec { name, @@ -537,9 +538,10 @@ async fn run_as_step( ) -> Result<()> { use spur_proto::proto::RunStepRequest; - let mut client = SlurmControllerClient::connect(args.controller.clone()) + let channel = spur_client::connect_channel(&args.controller) .await .context("failed to connect to spurctld")?; + let mut client = SlurmControllerClient::new(channel); // Create a step on the controller for tracking; capture the assigned // step_id so the completion (and thus DerivedExitCode) records against it. diff --git a/crates/spur-cli/src/sstat.rs b/crates/spur-cli/src/sstat.rs index 4b846a0b..36905cad 100644 --- a/crates/spur-cli/src/sstat.rs +++ b/crates/spur-cli/src/sstat.rs @@ -53,9 +53,10 @@ pub async fn main_with_args(args: Vec) -> Result<()> { bail!("sstat: no valid job IDs specified"); } - let mut client = SlurmControllerClient::connect(args.controller.clone()) + let channel = spur_client::connect_channel(&args.controller) .await .context("failed to connect to spurctld")?; + let mut client = SlurmControllerClient::new(channel); // Determine which fields to show let fields = if let Some(ref fmt) = args.format { diff --git a/crates/spur-cli/src/token.rs b/crates/spur-cli/src/token.rs index f715a47c..3aab0862 100644 --- a/crates/spur-cli/src/token.rs +++ b/crates/spur-cli/src/token.rs @@ -73,7 +73,7 @@ fn parse_ttl(s: &str) -> Result { async fn cmd_create(controller: &str, ttl: Option) -> Result<()> { let ttl_secs = ttl.map(|t| parse_ttl(&t)).transpose()?; - let mut client = SlurmControllerClient::connect(controller.to_string()).await?; + let mut client = SlurmControllerClient::new(spur_client::connect_channel(controller).await?); let resp = client.create_token(CreateTokenRequest { ttl_secs }).await?; let inner = resp.into_inner(); @@ -83,7 +83,7 @@ async fn cmd_create(controller: &str, ttl: Option) -> Result<()> { } async fn cmd_list(controller: &str) -> Result<()> { - let mut client = SlurmControllerClient::connect(controller.to_string()).await?; + let mut client = SlurmControllerClient::new(spur_client::connect_channel(controller).await?); let resp = client.list_tokens(ListTokensRequest {}).await?; let tokens = resp.into_inner().tokens; @@ -109,7 +109,7 @@ async fn cmd_list(controller: &str) -> Result<()> { } async fn cmd_revoke(controller: &str, token_id: &str) -> Result<()> { - let mut client = SlurmControllerClient::connect(controller.to_string()).await?; + let mut client = SlurmControllerClient::new(spur_client::connect_channel(controller).await?); client .revoke_token(RevokeTokenRequest { token_id: token_id.to_string(), diff --git a/crates/spur-client/Cargo.toml b/crates/spur-client/Cargo.toml new file mode 100644 index 00000000..3685c498 --- /dev/null +++ b/crates/spur-client/Cargo.toml @@ -0,0 +1,18 @@ +[package] +name = "spur-client" +version.workspace = true +edition.workspace = true +license.workspace = true +publish.workspace = true + +[dependencies] +tonic = { workspace = true } +tracing = { workspace = true } + +[dev-dependencies] +tokio = { workspace = true } +tokio-stream = { version = "0.1", features = ["net"] } +tonic-health = "0.14" + +[lints] +workspace = true diff --git a/crates/spur-client/src/lib.rs b/crates/spur-client/src/lib.rs new file mode 100644 index 00000000..68e14138 --- /dev/null +++ b/crates/spur-client/src/lib.rs @@ -0,0 +1,168 @@ +// Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Transport dialing with client-side failover. +//! +//! Establishes a [`tonic::transport::Channel`] from a comma-separated list of +//! endpoints, rotating to the next endpoint when one is unreachable. The helper +//! is service-agnostic: callers wrap the returned channel in whatever generated +//! client they need (e.g. `SlurmControllerClient::new(channel)`). + +use tonic::transport::{Channel, Endpoint}; +use tracing::warn; + +/// Parse a comma-separated endpoint string into normalized `http(s)://host:port` +/// entries. Whitespace is trimmed and empty entries dropped. Entries without a +/// scheme are prefixed with `http://`. +pub fn parse_endpoints(raw: &str) -> Vec { + raw.split(',') + .map(str::trim) + .filter(|s| !s.is_empty()) + .map(normalize_endpoint) + .collect() +} + +fn normalize_endpoint(addr: &str) -> String { + if addr.starts_with("http://") || addr.starts_with("https://") { + addr.to_string() + } else { + format!("http://{addr}") + } +} + +/// Always yields at least one entry so callers get a real transport error +/// instead of a silent no-op when handed an empty string. +fn endpoint_list(raw: &str) -> Vec { + let parsed = parse_endpoints(raw); + if parsed.is_empty() { + vec![normalize_endpoint(raw.trim())] + } else { + parsed + } +} + +/// Connect to the first reachable controller endpoint. +/// +/// Endpoints are tried in the order given. On connection failure the next +/// endpoint is attempted; if every endpoint fails, the last error is returned. +/// Once connected to any running controller, server-side leader forwarding +/// routes writes to the current Raft leader, so a single reachable node is +/// sufficient regardless of which one is the leader. +pub async fn connect_channel(endpoints: &str) -> Result { + let list = endpoint_list(endpoints); + let last = list.len() - 1; + + for endpoint in &list[..last] { + match try_connect(endpoint).await { + Ok(channel) => return Ok(channel), + Err(e) => warn!( + %endpoint, + error = %e, + "controller endpoint unreachable, trying next" + ), + } + } + + try_connect(&list[last]).await +} + +async fn try_connect(endpoint: &str) -> Result { + Endpoint::from_shared(endpoint.to_string())?.connect().await +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn parse_single_endpoint() { + assert_eq!( + parse_endpoints("http://ctrl:6817"), + vec!["http://ctrl:6817"] + ); + } + + #[test] + fn parse_adds_scheme() { + assert_eq!(parse_endpoints("ctrl:6817"), vec!["http://ctrl:6817"]); + } + + #[test] + fn parse_preserves_https() { + assert_eq!( + parse_endpoints("https://ctrl:6817"), + vec!["https://ctrl:6817"] + ); + } + + #[test] + fn parse_comma_list_with_whitespace() { + assert_eq!( + parse_endpoints(" ctrl1:6817 , http://ctrl2:6817 ,ctrl3:6817"), + vec![ + "http://ctrl1:6817", + "http://ctrl2:6817", + "http://ctrl3:6817" + ] + ); + } + + #[test] + fn parse_drops_empty_entries() { + assert_eq!( + parse_endpoints("ctrl1:6817,,ctrl2:6817,"), + vec!["http://ctrl1:6817", "http://ctrl2:6817"] + ); + } + + #[test] + fn endpoint_list_never_empty() { + assert_eq!(endpoint_list(""), vec!["http://"]); + assert_eq!(endpoint_list(" "), vec!["http://"]); + } + + #[tokio::test] + async fn connect_all_down_returns_error() { + // Two ports that are bound then released, so connections are refused. + let down1 = free_addr().await; + let down2 = free_addr().await; + let endpoints = format!("http://{down1},http://{down2}"); + assert!(connect_channel(&endpoints).await.is_err()); + } + + #[tokio::test] + async fn connect_rotates_to_first_reachable() { + let down = free_addr().await; + let up = spawn_server().await; + let endpoints = format!("http://{down},http://{up}"); + assert!(connect_channel(&endpoints).await.is_ok()); + } + + /// Bind to an ephemeral port and release it, yielding an address that will + /// refuse connections. + async fn free_addr() -> std::net::SocketAddr { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0") + .await + .expect("bind ephemeral port"); + listener.local_addr().expect("local addr") + } + + /// Start a minimal h2 server on a bound listener and return its address. + /// The service is irrelevant: `connect_channel` only completes the h2 + /// handshake and returns a channel without issuing an RPC. + async fn spawn_server() -> std::net::SocketAddr { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0") + .await + .expect("bind ephemeral port"); + let addr = listener.local_addr().expect("local addr"); + let (_reporter, health) = tonic_health::server::health_reporter(); + let incoming = tokio_stream::wrappers::TcpListenerStream::new(listener); + tokio::spawn(async move { + let _ = tonic::transport::Server::builder() + .add_service(health) + .serve_with_incoming(incoming) + .await; + }); + addr + } +} diff --git a/crates/spur-core/src/config.rs b/crates/spur-core/src/config.rs index 25d27270..b6d7c78d 100644 --- a/crates/spur-core/src/config.rs +++ b/crates/spur-core/src/config.rs @@ -321,6 +321,19 @@ impl Default for ControllerConfig { } } +impl ControllerConfig { + /// Client-facing controller endpoints, one per `hosts` entry, each combined + /// with the port from `listen_addr`. Used for client-side failover across an + /// HA quorum. + pub fn endpoints(&self) -> Vec { + let port = self.listen_addr.rsplit(':').next().unwrap_or("6817"); + self.hosts + .iter() + .map(|host| format!("http://{host}:{port}")) + .collect() + } +} + #[derive(Debug, Clone, Serialize, Deserialize)] pub struct AccountingConfig { /// Address of the accounting daemon. @@ -947,6 +960,39 @@ pub fn format_time(total_minutes: Option) -> String { mod tests { use super::*; + #[test] + fn test_controller_endpoints_single_host() { + let cfg = ControllerConfig::default(); + assert_eq!(cfg.endpoints(), vec!["http://localhost:6817"]); + } + + #[test] + fn test_controller_endpoints_multi_host() { + let cfg = ControllerConfig { + listen_addr: "[::]:6817".into(), + hosts: vec!["ctrl1".into(), "ctrl2".into(), "ctrl3".into()], + ..Default::default() + }; + assert_eq!( + cfg.endpoints(), + vec![ + "http://ctrl1:6817", + "http://ctrl2:6817", + "http://ctrl3:6817" + ] + ); + } + + #[test] + fn test_controller_endpoints_custom_port() { + let cfg = ControllerConfig { + listen_addr: "0.0.0.0:7000".into(), + hosts: vec!["ctrl1".into()], + ..Default::default() + }; + assert_eq!(cfg.endpoints(), vec!["http://ctrl1:7000"]); + } + #[test] fn test_parse_time() { assert_eq!(parse_time_minutes("60"), Some(60)); diff --git a/crates/spur-ffi/Cargo.toml b/crates/spur-ffi/Cargo.toml index 735f4e67..45410e59 100644 --- a/crates/spur-ffi/Cargo.toml +++ b/crates/spur-ffi/Cargo.toml @@ -11,6 +11,7 @@ crate-type = ["cdylib", "staticlib"] [dependencies] spur-proto = { path = "../spur-proto" } +spur-client = { path = "../spur-client" } spur-core = { path = "../spur-core" } tokio = { workspace = true } tonic = { workspace = true } diff --git a/crates/spur-ffi/src/lib.rs b/crates/spur-ffi/src/lib.rs index c941c04c..0574b8ce 100644 --- a/crates/spur-ffi/src/lib.rs +++ b/crates/spur-ffi/src/lib.rs @@ -67,9 +67,8 @@ pub extern "C" fn slurm_submit_batch_job( use spur_proto::proto::slurm_controller_client::SlurmControllerClient; use spur_proto::proto::{JobSpec, SubmitJobRequest}; - let mut client = SlurmControllerClient::connect(controller_addr().to_string()) - .await - .ok()?; + let channel = spur_client::connect_channel(controller_addr()).await.ok()?; + let mut client = SlurmControllerClient::new(channel); let spec = JobSpec { name: c_str_to_string(desc.name), @@ -126,9 +125,8 @@ pub extern "C" fn slurm_load_jobs( use spur_proto::proto::slurm_controller_client::SlurmControllerClient; use spur_proto::proto::GetJobsRequest; - let mut client = SlurmControllerClient::connect(controller_addr().to_string()) - .await - .ok()?; + let channel = spur_client::connect_channel(controller_addr()).await.ok()?; + let mut client = SlurmControllerClient::new(channel); let resp = client.get_jobs(GetJobsRequest::default()).await.ok()?; @@ -204,9 +202,8 @@ pub extern "C" fn slurm_load_node( use spur_proto::proto::slurm_controller_client::SlurmControllerClient; use spur_proto::proto::GetNodesRequest; - let mut client = SlurmControllerClient::connect(controller_addr().to_string()) - .await - .ok()?; + let channel = spur_client::connect_channel(controller_addr()).await.ok()?; + let mut client = SlurmControllerClient::new(channel); let resp = client.get_nodes(GetNodesRequest::default()).await.ok()?; @@ -258,9 +255,8 @@ pub extern "C" fn slurm_load_partitions( use spur_proto::proto::slurm_controller_client::SlurmControllerClient; use spur_proto::proto::GetPartitionsRequest; - let mut client = SlurmControllerClient::connect(controller_addr().to_string()) - .await - .ok()?; + let channel = spur_client::connect_channel(controller_addr()).await.ok()?; + let mut client = SlurmControllerClient::new(channel); let resp = client .get_partitions(GetPartitionsRequest::default()) @@ -305,9 +301,8 @@ pub extern "C" fn slurm_kill_job(job_id: c_uint, signal: u16, _flags: u16) -> c_ use spur_proto::proto::slurm_controller_client::SlurmControllerClient; use spur_proto::proto::CancelJobRequest; - let mut client = SlurmControllerClient::connect(controller_addr().to_string()) - .await - .ok()?; + let channel = spur_client::connect_channel(controller_addr()).await.ok()?; + let mut client = SlurmControllerClient::new(channel); client .cancel_job(CancelJobRequest { diff --git a/crates/spurd/Cargo.toml b/crates/spurd/Cargo.toml index 141b1eb0..78e2c9fd 100644 --- a/crates/spurd/Cargo.toml +++ b/crates/spurd/Cargo.toml @@ -11,6 +11,7 @@ path = "src/main.rs" [dependencies] spur-proto = { path = "../spur-proto" } +spur-client = { path = "../spur-client" } spur-core = { path = "../spur-core" } spur-update = { path = "../spur-update" } spur-net = { path = "../spur-net" } diff --git a/crates/spurd/src/agent_server.rs b/crates/spurd/src/agent_server.rs index c381909b..18f52343 100644 --- a/crates/spurd/src/agent_server.rs +++ b/crates/spurd/src/agent_server.rs @@ -344,17 +344,12 @@ async fn report_completion( // RaisedSignal outcome from the reported `signal`. let state = spur_core::job::JobState::completion_state_for_exit_code(exit_code).to_proto_i32(); - let url = if controller_addr.starts_with("http") { - controller_addr.to_string() - } else { - format!("http://{}", controller_addr) - }; - // Retry up to 3 times with 1-second backoff — a single transient failure // must not permanently lose a job completion. for attempt in 1..=3 { - match SlurmControllerClient::connect(url.clone()).await { - Ok(mut client) => { + match spur_client::connect_channel(controller_addr).await { + Ok(channel) => { + let mut client = SlurmControllerClient::new(channel); let req = ReportJobStatusRequest { job_id, state, @@ -370,7 +365,7 @@ async fn report_completion( info!( job_id, exit_code, - controller = %url, + controller = %controller_addr, "reported completion to controller" ); return; diff --git a/crates/spurd/src/reporter.rs b/crates/spurd/src/reporter.rs index 9501ba15..364d1945 100644 --- a/crates/spurd/src/reporter.rs +++ b/crates/spurd/src/reporter.rs @@ -49,9 +49,10 @@ impl NodeReporter { /// Register with the controller. pub async fn register(&self) -> anyhow::Result<()> { - let mut client = SlurmControllerClient::connect(self.controller_addr.clone()) + let channel = spur_client::connect_channel(&self.controller_addr) .await .context("failed to connect to spurctld for registration")?; + let mut client = SlurmControllerClient::new(channel); let resp = client .register_agent(RegisterAgentRequest { @@ -83,9 +84,10 @@ impl NodeReporter { /// Notify the controller that this agent is shutting down. pub async fn deregister(&self, reason: &str) -> anyhow::Result<()> { let current_token = self.node_token.read().unwrap().clone(); - let mut client = SlurmControllerClient::connect(self.controller_addr.clone()) + let channel = spur_client::connect_channel(&self.controller_addr) .await .context("failed to connect to spurctld for deregistration")?; + let mut client = SlurmControllerClient::new(channel); client .deregister_agent(spur_proto::proto::DeregisterAgentRequest { @@ -112,8 +114,9 @@ impl NodeReporter { self.free_memory_mb.store(free_mem, Ordering::Relaxed); let current_token = self.node_token.read().unwrap().clone(); - match SlurmControllerClient::connect(self.controller_addr.clone()).await { - Ok(mut client) => { + match spur_client::connect_channel(&self.controller_addr).await { + Ok(channel) => { + let mut client = SlurmControllerClient::new(channel); match client .heartbeat(spur_proto::proto::HeartbeatRequest { hostname: self.hostname.clone(), diff --git a/docs/deployment/native-host.rst b/docs/deployment/native-host.rst index 5d514ebd..338a373c 100644 --- a/docs/deployment/native-host.rst +++ b/docs/deployment/native-host.rst @@ -131,6 +131,14 @@ Start the agent: The agent auto-detects CPUs, memory, and GPUs, then registers with the controller over the mesh. +For an HA quorum, pass every controller as a comma-separated list so the agent +and CLI fail over to a surviving node if one is unreachable. The same format +works for the ``SPUR_CONTROLLER_ADDR`` environment variable: + +.. code-block:: bash + + --controller http://10.44.0.1:6817,http://10.44.0.2:6817,http://10.44.0.3:6817 + Repeat for each worker, incrementing the WireGuard address. Verify: diff --git a/examples/spur.conf b/examples/spur.conf index 57dbc8fb..a19ef14e 100644 --- a/examples/spur.conf +++ b/examples/spur.conf @@ -3,6 +3,11 @@ cluster_name = "mi300x-cluster" [controller] listen_addr = "[::]:6817" state_dir = "/var/spool/spur" +# Controller hostname(s). For an HA quorum, list every controller so clients +# and agents fail over to a surviving node if one is unreachable. Each host is +# combined with the port from listen_addr. Equivalent to a comma-separated +# SPUR_CONTROLLER_ADDR (e.g. "http://ctrl1:6817,http://ctrl2:6817"). +# hosts = ["ctrl1", "ctrl2", "ctrl3"] [scheduler] plugin = "backfill" diff --git a/tests/native_host/e2e/cluster.py b/tests/native_host/e2e/cluster.py index 7f96e5ec..d66411a3 100644 --- a/tests/native_host/e2e/cluster.py +++ b/tests/native_host/e2e/cluster.py @@ -324,21 +324,28 @@ def teardown(self): # --- CLI wrappers --- - def cli(self, args: list[str]) -> str: - """Run a spur CLI command on the controller node.""" + def cli(self, args: list[str], controller_addr: str | None = None) -> str: + """Run a spur CLI command on the controller node. + + *controller_addr* overrides the endpoint(s) passed via + ``SPUR_CONTROLLER_ADDR`` (e.g. a comma-separated failover list). + """ cmd_parts = [ - f"SPUR_CONTROLLER_ADDR='{self.controller_addr}'", + f"SPUR_CONTROLLER_ADDR='{controller_addr or self.controller_addr}'", f"PATH='{self.bin_dir}':$PATH", f"'{self.bin_dir}/{args[0]}'", ] cmd_parts.extend(f"'{a}'" for a in args[1:]) return self.nodes[0].exec(" ".join(cmd_parts)) - def cli_allow_fail(self, args: list[str]) -> str: + def cli_allow_fail(self, args: list[str], controller_addr: str | None = None) -> str: """Run a spur CLI command, returning stdout+stderr regardless of exit - code. Use to assert on expected submission rejections.""" + code. Use to assert on expected submission rejections. + + *controller_addr* overrides ``SPUR_CONTROLLER_ADDR`` as in :meth:`cli`. + """ cmd_parts = [ - f"SPUR_CONTROLLER_ADDR='{self.controller_addr}'", + f"SPUR_CONTROLLER_ADDR='{controller_addr or self.controller_addr}'", f"PATH='{self.bin_dir}':$PATH", f"'{self.bin_dir}/{args[0]}'", ] diff --git a/tests/native_host/e2e/test_controller_failover.py b/tests/native_host/e2e/test_controller_failover.py new file mode 100644 index 00000000..067fa125 --- /dev/null +++ b/tests/native_host/e2e/test_controller_failover.py @@ -0,0 +1,67 @@ +# Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Client-side controller failover E2E tests. + +Exercises the comma-separated ``SPUR_CONTROLLER_ADDR`` endpoint list: the CLI +must rotate past an unreachable controller endpoint and reach a live one. These +run against the standard single-controller cluster; the dead endpoint is a +closed port on the same host, so no extra topology is needed. + +Full multi-controller Raft failover (kill the connected controller, survivor +takes over) is covered separately once the harness grows Raft support. +""" + +from cluster import parse_job_id, wait_job + +# Port with nothing listening — connections are refused immediately. +DEAD_PORT = 1 + + +def _dead(cluster) -> str: + return f"http://{cluster.nodes[0].host}:{DEAD_PORT}" + + +class TestControllerFailover: + def test_rotates_past_dead_first_endpoint(self, cluster): + endpoints = f"{_dead(cluster)},{cluster.controller_addr}" + out = cluster.cli(["sinfo"], controller_addr=endpoints) + for name in cluster.node_names: + assert name in out, ( + f"node {name} missing after rotating past dead endpoint:\n{out}" + ) + + def test_live_first_endpoint_still_works(self, cluster): + endpoints = f"{cluster.controller_addr},{_dead(cluster)}" + out = cluster.cli(["sinfo"], controller_addr=endpoints) + for name in cluster.node_names: + assert name in out, f"node {name} missing with live-first list:\n{out}" + + def test_all_endpoints_down_fails_cleanly(self, cluster): + endpoints = f"{_dead(cluster)},http://{cluster.nodes[0].host}:2" + out = cluster.cli_allow_fail(["sinfo"], controller_addr=endpoints) + assert not any(name in out for name in cluster.node_names), ( + f"expected failure with all endpoints down, got node output:\n{out}" + ) + assert "connect" in out.lower() or "transport" in out.lower(), ( + f"expected a connection error, got:\n{out}" + ) + + def test_job_submits_through_failover_list(self, cluster): + out_path = f"{cluster.remote_dir}/failover.out" + script = cluster.write_file( + "failover-job.sh", "#!/bin/bash\necho FAILOVER_JOB_OK\n" + ) + endpoints = f"{_dead(cluster)},{cluster.controller_addr}" + sb = cluster.cli( + ["sbatch", "-J", "failover", "-N", "1", "-o", out_path, script], + controller_addr=endpoints, + ) + job_id = parse_job_id(sb) + assert job_id is not None, f"submit through failover list failed:\n{sb}" + + 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 "FAILOVER_JOB_OK" in content, f"output:\n{content}"