From 43be31e87d6a6c5ee0d80a6a3fd9b7ef07897923 Mon Sep 17 00:00:00 2001 From: yansun1996 Date: Tue, 28 Jul 2026 01:11:06 +0000 Subject: [PATCH] fix(spurd): redeliver completion reports across a controller outage MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The agent reported job completion with only 3 retries over ~3s, then dropped the report permanently. A controller unreachable longer than that window (Raft leader failover, restart, or a network blip) lost the completion, leaving the job stuck in COMPLETING until the complete_wait_secs force-finish timer (default 300s) — during which the node could be handed a new job it would reject. Retry delivery with capped exponential backoff (1s doubling to 30s) until the controller acks or returns a permanent error, and spawn the delivery so a slow controller no longer stalls the monitor loop's detection of other completions. Controller-side handling is idempotent (duplicate/late reports for a terminal or superseded run are acked), so unbounded retry is safe. A NotFound is now retryable to cover the brief cold-start window before a restarted controller replays the job's log entry. --- crates/spurd/src/agent_server.rs | 327 +++++++++++++----- .../e2e/test_completion_recovery.py | 55 +++ 2 files changed, 290 insertions(+), 92 deletions(-) create mode 100644 tests/native_host/e2e/test_completion_recovery.py diff --git a/crates/spurd/src/agent_server.rs b/crates/spurd/src/agent_server.rs index a44a7f3a..edc29ad6 100644 --- a/crates/spurd/src/agent_server.rs +++ b/crates/spurd/src/agent_server.rs @@ -318,16 +318,17 @@ impl AgentService { } else { None }; - report_completion( - &controller_addr, + // Spawn so a controller mid-failover/restart can't stall the + // monitor loop; delivery retries in the background until acked. + tokio::spawn(report_completion( + controller_addr.clone(), c.job_id, c.exit_code, c.signal, c.run_attempt, - &local_hostname, - drain.as_ref(), - ) - .await; + local_hostname.clone(), + drain, + )); } } }); @@ -437,21 +438,27 @@ impl Drop for LaunchReservationGuard { fn completion_report_retryable(status: &tonic::Status) -> bool { use tonic::Code; + // NotFound: a just-restarted controller may not have replayed this job's log + // entry yet; jobs are never purged, so retrying resolves once replay finishes. matches!( status.code(), - Code::Unavailable | Code::Internal | Code::DeadlineExceeded | Code::Unknown + Code::Unavailable + | Code::Internal + | Code::DeadlineExceeded + | Code::Unknown + | Code::NotFound ) } #[cfg(test)] mod completion_report_tests { - use super::completion_report_retryable; + use super::*; + use std::sync::atomic::{AtomicU32, Ordering}; use tonic::Status; #[test] fn permanent_errors_are_not_retryable() { assert!(!completion_report_retryable(&Status::invalid_argument("x"))); - assert!(!completion_report_retryable(&Status::not_found("x"))); } #[test] @@ -459,6 +466,109 @@ mod completion_report_tests { assert!(completion_report_retryable(&Status::unavailable("x"))); assert!(completion_report_retryable(&Status::internal("x"))); } + + #[test] + fn not_found_is_retryable_for_cold_start_replay() { + assert!(completion_report_retryable(&Status::not_found( + "job 7 not found" + ))); + } + + #[test] + fn classify_maps_result_to_delivery() { + assert_eq!(classify_report_result(&Ok(())), CompletionDelivery::Acked); + assert_eq!( + classify_report_result(&Err(Status::unavailable("x"))), + CompletionDelivery::Transient + ); + assert_eq!( + classify_report_result(&Err(Status::invalid_argument("x"))), + CompletionDelivery::Permanent + ); + } + + #[test] + fn backoff_doubles_then_caps() { + let max = std::time::Duration::from_secs(30); + let d = std::time::Duration::from_secs(1); + let d = next_backoff(d, max); + assert_eq!(d, std::time::Duration::from_secs(2)); + let d = next_backoff(d, max); + assert_eq!(d, std::time::Duration::from_secs(4)); + assert_eq!( + next_backoff(std::time::Duration::from_secs(20), max), + max, + "must saturate at the cap" + ); + } + + #[tokio::test(start_paused = true)] + async fn transient_failures_retry_until_acked() { + let calls = AtomicU32::new(0); + let attempts = drive_completion_report( + |_| { + let n = calls.fetch_add(1, Ordering::SeqCst) + 1; + async move { + if n < 4 { + Err(Status::unavailable("controller mid-failover")) + } else { + Ok(()) + } + } + }, + COMPLETION_RETRY_INITIAL, + COMPLETION_RETRY_MAX, + ) + .await; + assert_eq!(attempts, 4); + } + + #[tokio::test(start_paused = true)] + async fn permanent_failure_stops_immediately() { + let calls = AtomicU32::new(0); + let attempts = drive_completion_report( + |_| { + calls.fetch_add(1, Ordering::SeqCst); + async move { Err(Status::invalid_argument("node not allocated")) } + }, + COMPLETION_RETRY_INITIAL, + COMPLETION_RETRY_MAX, + ) + .await; + assert_eq!(attempts, 1); + assert_eq!(calls.load(Ordering::SeqCst), 1); + } + + #[tokio::test(start_paused = true)] + async fn success_on_first_attempt_makes_one_call() { + let attempts = drive_completion_report( + |_| async move { Ok(()) }, + COMPLETION_RETRY_INITIAL, + COMPLETION_RETRY_MAX, + ) + .await; + assert_eq!(attempts, 1); + } + + #[test] + fn request_carries_drain_fields_when_draining() { + let drain = DrainRequest { + reason: "prolog failed: boom".into(), + }; + let req = build_completion_request(7, -1, 0, 3, "node1".into(), Some(&drain)); + assert!(req.drain_node); + assert_eq!(req.drain_reason, "prolog failed: boom"); + assert_eq!(req.job_id, 7); + assert_eq!(req.reporting_node, "node1"); + assert_eq!(req.run_attempt, 3); + } + + #[test] + fn request_omits_drain_when_not_draining() { + let req = build_completion_request(7, 0, 0, 1, "node1".into(), None); + assert!(!req.drain_node); + assert!(req.drain_reason.is_empty()); + } } /// Reap an already-killed displaced run. Polls `try_wait` so both executor @@ -552,85 +662,122 @@ fn inject_script_args(script: &str, args: &[String]) -> Result { Ok(format!("{set_line}\n{script}")) } -async fn report_completion( - controller_addr: &str, +/// Initial and capped retry backoff for redelivering a completion report. +const COMPLETION_RETRY_INITIAL: std::time::Duration = std::time::Duration::from_secs(1); +const COMPLETION_RETRY_MAX: std::time::Duration = std::time::Duration::from_secs(30); + +/// Outcome of one `ReportJobStatus` attempt, driving the redelivery loop. +#[derive(Debug, PartialEq, Eq)] +enum CompletionDelivery { + Acked, + Permanent, + Transient, +} + +fn classify_report_result(result: &Result<(), tonic::Status>) -> CompletionDelivery { + match result { + Ok(()) => CompletionDelivery::Acked, + Err(s) if completion_report_retryable(s) => CompletionDelivery::Transient, + Err(_) => CompletionDelivery::Permanent, + } +} + +fn next_backoff(current: std::time::Duration, max: std::time::Duration) -> std::time::Duration { + current.saturating_mul(2).min(max) +} + +/// Redeliver a completion until acked or a permanent error; a transient failure +/// retries with capped exponential backoff. Returns the attempt count. +async fn drive_completion_report( + mut attempt: F, + initial: std::time::Duration, + max: std::time::Duration, +) -> u32 +where + F: FnMut(u32) -> Fut, + Fut: std::future::Future>, +{ + let mut delay = initial; + let mut n = 0; + loop { + n += 1; + match classify_report_result(&attempt(n).await) { + CompletionDelivery::Acked | CompletionDelivery::Permanent => return n, + CompletionDelivery::Transient => { + tokio::time::sleep(delay).await; + delay = next_backoff(delay, max); + } + } + } +} + +/// Build the completion RPC. `state` is advisory (derived from `exit_code` to +/// satisfy the validator); the controller rederives the true outcome from `signal`. +fn build_completion_request( job_id: u32, exit_code: i32, signal: i32, run_attempt: u32, - reporting_node: &str, + reporting_node: String, drain: Option<&DrainRequest>, +) -> ReportJobStatusRequest { + ReportJobStatusRequest { + job_id, + state: spur_core::job::JobState::completion_state_for_exit_code(exit_code).to_proto_i32(), + exit_code, + signal, + message: format!("exit_code={}", exit_code), + drain_node: drain.is_some(), + drain_reason: drain.map(|d| d.reason.clone()).unwrap_or_default(), + reporting_node, + run_attempt, + } +} + +async fn report_completion( + controller_addr: String, + job_id: u32, + exit_code: i32, + signal: i32, + run_attempt: u32, + reporting_node: String, + drain: Option, ) { - // Wire `state` is derived from `exit_code` alone (advisory): a signaled job - // reports Completed/0 because the controller's validator requires - // state<->exit_code agreement. The controller rederives the true Failed / - // RaisedSignal outcome from the reported `signal`. - let state = spur_core::job::JobState::completion_state_for_exit_code(exit_code).to_proto_i32(); - - // 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 spur_client::connect_channel(controller_addr).await { - Ok(channel) => { - let mut client = spur_proto::controller_client(channel); - let req = ReportJobStatusRequest { - job_id, - state, - exit_code, - signal, - message: format!("exit_code={}", exit_code), - drain_node: drain.is_some(), - drain_reason: drain.as_ref().map(|d| d.reason.clone()).unwrap_or_default(), - reporting_node: reporting_node.to_string(), - run_attempt, - }; - match client.report_job_status(req).await { - Ok(_) => { - info!( - job_id, - exit_code, - controller = %controller_addr, - "reported completion to controller" - ); - return; - } - Err(e) => { - if !completion_report_retryable(&e) { - error!( - job_id, - attempt, - code = ?e.code(), - error = %e, - "ReportJobStatus failed with non-retryable error" - ); - return; - } - warn!( - job_id, - attempt, - error = %e, - "ReportJobStatus RPC failed" - ); - } + let attempt = |n: u32| { + let controller_addr = controller_addr.clone(); + let req = build_completion_request( + job_id, + exit_code, + signal, + run_attempt, + reporting_node.clone(), + drain.as_ref(), + ); + async move { + let result = match spur_client::connect_channel(&controller_addr).await { + Ok(channel) => spur_proto::controller_client(channel) + .report_job_status(req) + .await + .map(|_| ()), + Err(e) => Err(tonic::Status::unavailable(e.to_string())), + }; + match &result { + Ok(()) => info!(job_id, exit_code, "reported completion to controller"), + Err(e) if completion_report_retryable(e) => { + warn!(job_id, attempt = n, error = %e, "ReportJobStatus failed; will retry") } - } - Err(e) => { - warn!( + Err(e) => error!( job_id, - attempt, + code = ?e.code(), error = %e, - "failed to connect to controller for completion report" - ); + "ReportJobStatus failed with non-retryable error" + ), } + result } - if attempt < 3 { - tokio::time::sleep(tokio::time::Duration::from_secs(1)).await; - } - } - error!( - job_id, - exit_code, "gave up reporting completion after 3 attempts" - ); + }; + + drive_completion_report(attempt, COMPLETION_RETRY_INITIAL, COMPLETION_RETRY_MAX).await; } #[tonic::async_trait] @@ -1114,21 +1261,17 @@ impl SlurmAgent for AgentService { let controller = self.reporter.controller_addr.clone(); let node_name = self.reporter.hostname.clone(); let drain_reason = format!("prolog failed: {}", err_msg); - tokio::spawn(async move { - let drain = DrainRequest { - reason: drain_reason.clone(), - }; - report_completion( - &controller, - job_id, - -1, - 0, - run_attempt, - &node_name, - Some(&drain), - ) - .await; - }); + tokio::spawn(report_completion( + controller, + job_id, + -1, + 0, + run_attempt, + node_name, + Some(DrainRequest { + reason: drain_reason, + }), + )); } Ok(Response::new(LaunchJobResponse { diff --git a/tests/native_host/e2e/test_completion_recovery.py b/tests/native_host/e2e/test_completion_recovery.py new file mode 100644 index 00000000..f1590974 --- /dev/null +++ b/tests/native_host/e2e/test_completion_recovery.py @@ -0,0 +1,55 @@ +# Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""E2E test for durable completion-report redelivery across a controller outage. + +If the controller is unreachable when a job finishes (a routine restart or +failover), the agent's completion report must not be permanently lost. The +agent keeps redelivering it, so once the controller returns the job leaves +COMPLETING in the normal sub-second path instead of waiting out the multi-minute +``complete_wait_secs`` force-finish timer. + +``complete_wait_secs`` is set high so the force-finish timer cannot be what +finalizes the job — a prompt terminal state after the controller returns proves +the redelivered report did it. +""" + +import time + +import pytest + +from cluster import parse_job_id, wait_job + + +class TestCompletionRecovery: + # Well above the time the job spends stopped-then-restarted, so a prompt + # finish can only come from the redelivered report, not the force-finish. + COMPLETE_WAIT = 300 + + @pytest.fixture + def cluster_config_overrides(self): + return {"scheduler": {"complete_wait_secs": self.COMPLETE_WAIT}} + + def test_completion_survives_controller_restart(self, cluster): + out_path = f"{cluster.remote_dir}/recover.out" + script = cluster.write_file( + "recover-job.sh", "#!/bin/bash\nsleep 3\necho RECOVER_OK\n" + ) + sb = cluster.sbatch(["-J", "recover", "-N", "1", "-o", out_path, script]) + job_id = parse_job_id(sb) + assert job_id is not None, f"submit failed:\n{sb}" + + # Stop the controller while the job is still running so the completion is + # generated during the outage and must be redelivered after restart. + cluster.stop_controller() + time.sleep(8) + cluster.restart_controller() + + # CD within 60s (not the 300s force-finish) proves redelivery finalized it. + state = wait_job(cluster, job_id, timeout=60) + assert state == "CD", ( + f"job did not complete promptly after controller restart (got {state}); " + "the completion report was not redelivered" + ) + content = cluster.read_output_on_any_node(out_path) + assert "RECOVER_OK" in content, f"job output missing marker:\n{content}"