Withdraw stale split topologies on worker loss (#763)

Withdraw stale split topologies on worker loss

Validation
* Validation tier: Tier 3 - split-serving recovery/status hardening refreshed onto current main for PR #763; active stage refresh failures now mark cached stages failed and split fallback/withdraw events remove stale topology/status surfaces without protocol/schema/ABI changes.
* git fetch --no-tags origin main:refs/remotes/origin/main codex/split-worker-loss-withdraw:refs/remotes/origin/codex/split-worker-loss-withdraw: PASS, origin/main at d9d90969.
* git rebase origin/main: PASS, no conflicts.
* git diff --check origin/main...HEAD: PASS, no output.
* git diff --check: PASS, no output.
* git diff --cached --check: PASS, no output.
* cargo fmt --all -- --check: PASS.
* LLAMA_STAGE_BUILD_DIR=<stage-build-dir> cargo test -p mesh-llm-host-runtime stage_topology --lib -- --test-threads=1: PASS, 3 passed.
* LLAMA_STAGE_BUILD_DIR=<stage-build-dir> cargo test -p mesh-llm-host-runtime active_stage_refresh --lib -- --test-threads=1: PASS, 2 passed.
* LLAMA_STAGE_BUILD_DIR=<stage-build-dir> cargo test -p mesh-llm-host-runtime split_unavailable_active_stage_nodes --lib -- --test-threads=1: PASS, 2 passed.
* LLAMA_STAGE_BUILD_DIR=<stage-build-dir> cargo test -p mesh-llm-host-runtime split_loss_recovery --lib -- --test-threads=1: PASS, 7 passed.
* LLAMA_STAGE_BUILD_DIR=<stage-build-dir> cargo check -p mesh-llm: PASS.
* LLAMA_STAGE_BUILD_DIR=<stage-build-dir> cargo clippy -p mesh-llm-host-runtime --all-targets -- -D warnings: PASS.
* Remote CI: PASS on refreshed head 600d828b; PR Builds and PR Quality Checks completed successfully.
* Ledger: not applicable - not required for selected validation tier/change family.
* Version: not applicable - no release/version sync required for this non-release split recovery hardening.
* Not run: just build - not required for selected validation tier; no UI assets or release bundle changed.
* Not run: live two-node split worker-loss smoke - no local second split-capable node/runtime endpoint was available; targeted topology, refresh, coordinator recovery tests, and CI split smoke coverage cover the changed paths.

Rollback
* git revert <merge-commit-sha>
This commit is contained in:
Ivan Golovach 2026-05-31 14:17:41 -07:00 committed by GitHub
parent d9d9096955
commit 8219dc76ca
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 142 additions and 21 deletions

View file

@ -2616,6 +2616,15 @@ impl StageTopologyState {
self.record_topology(topology);
}
fn withdraw_topology(&mut self, topology_id: &str, run_id: &str) -> bool {
let topology_key = stage_topology_key(topology_id, run_id);
let removed_topology = self.topologies.remove(&topology_key).is_some();
let old_status_count = self.statuses.len();
self.statuses
.retain(|_, status| status.topology_id != topology_id || status.run_id != run_id);
removed_topology || self.statuses.len() != old_status_count
}
fn visible_topologies(&self) -> Vec<StageTopologyInstance> {
self.topologies
.values()
@ -2671,6 +2680,17 @@ impl StageTopologyState {
);
}
fn record_status_refresh_failure(&mut self, status: &StageRuntimeStatus, error: String) {
self.record_status(stage_runtime_status_from_snapshot(
status.node_id,
stage_snapshot_from_runtime_status(
status,
crate::inference::skippy::StageRuntimeState::Failed,
Some(error),
),
));
}
fn active_statuses(&self) -> Vec<StageRuntimeStatus> {
self.statuses
.values()
@ -3091,6 +3111,13 @@ impl Node {
.activate_topology(topology);
}
pub async fn withdraw_stage_topology(&self, topology_id: &str, run_id: &str) -> bool {
self.stage_topologies
.lock()
.await
.withdraw_topology(topology_id, run_id)
}
pub async fn stage_topologies(&self) -> Vec<StageTopologyInstance> {
self.stage_topologies.lock().await.visible_topologies()
}
@ -3129,15 +3156,13 @@ impl Node {
match tokio::time::timeout(timeout, refresh).await {
Ok(Ok(crate::inference::skippy::StageControlResponse::Status(statuses))) => {
if statuses.is_empty() {
self.record_stage_status(
Some(peer_id),
stage_snapshot_from_runtime_status(
self.stage_topologies
.lock()
.await
.record_status_refresh_failure(
&status,
crate::inference::skippy::StageRuntimeState::Failed,
Some("stage status missing from runtime".to_string()),
),
)
.await;
"stage status missing from runtime".to_string(),
);
} else {
for status in statuses {
self.record_stage_status(Some(peer_id), status).await;
@ -3149,23 +3174,25 @@ impl Node {
}
Ok(Ok(_)) => {}
Ok(Err(error)) => {
self.record_stage_status(
Some(peer_id),
stage_snapshot_from_runtime_status(
&status,
crate::inference::skippy::StageRuntimeState::Failed,
Some(error.to_string()),
),
)
.await;
self.stage_topologies
.lock()
.await
.record_status_refresh_failure(&status, error.to_string());
}
Err(_) => {
self.stage_topologies
.lock()
.await
.record_status_refresh_failure(
&status,
"stage status refresh timed out".to_string(),
);
tracing::debug!(
topology_id = %status.topology_id,
run_id = %status.run_id,
stage_id = %status.stage_id,
peer = %peer_id.fmt_short(),
"stage status refresh timed out; preserving last known status"
"stage status refresh timed out; marking stage failed"
);
}
}

View file

@ -7763,6 +7763,56 @@ fn active_stage_topology_replaces_previous_generation_for_model() {
assert!(state.runtime_statuses().is_empty());
}
#[test]
fn stage_topology_withdraw_removes_active_topology_and_statuses() {
let host_id = EndpointId::from(SecretKey::from_bytes(&[0x41; 32]).public());
let worker_id = EndpointId::from(SecretKey::from_bytes(&[0x42; 32]).public());
let mut state = StageTopologyState::default();
state.activate_topology(StageTopologyInstance {
topology_id: "topology-a".to_string(),
run_id: "run-a".to_string(),
model_id: "model-a".to_string(),
package_ref: "gguf:///model.gguf".to_string(),
manifest_sha256: "direct-gguf:1:model.gguf".to_string(),
stages: vec![
StageAssignment {
stage_id: "stage-0".to_string(),
stage_index: 0,
node_id: host_id,
layer_start: 0,
layer_end: 12,
endpoint: StageEndpoint {
bind_addr: "127.0.0.1:50000".to_string(),
},
},
StageAssignment {
stage_id: "stage-1".to_string(),
stage_index: 1,
node_id: worker_id,
layer_start: 12,
layer_end: 24,
endpoint: StageEndpoint {
bind_addr: "127.0.0.1:0".to_string(),
},
},
],
});
state.record_status(test_stage_status(
worker_id,
"stage-1",
1,
"127.0.0.1:51234",
crate::inference::skippy::StageRuntimeState::Ready,
));
assert_eq!(state.visible_topologies().len(), 1);
assert_eq!(state.runtime_statuses().len(), 1);
assert!(state.withdraw_topology("topology-a", "run-a"));
assert!(state.visible_topologies().is_empty());
assert!(state.runtime_statuses().is_empty());
assert!(!state.withdraw_topology("topology-a", "run-a"));
}
#[test]
fn empty_stage_status_snapshots_are_ignored() {
let node_id = EndpointId::from(SecretKey::from_bytes(&[0x39; 32]).public());
@ -7814,3 +7864,29 @@ fn active_stage_refresh_marks_missing_stage_failed() {
Some("stage status missing from runtime")
);
}
#[test]
fn active_stage_refresh_timeout_marks_cached_stage_failed() {
let node_id = EndpointId::from(SecretKey::from_bytes(&[0x43; 32]).public());
let mut state = StageTopologyState::default();
state.record_status(test_stage_status(
node_id,
"stage-1",
1,
"127.0.0.1:51234",
crate::inference::skippy::StageRuntimeState::Ready,
));
let cached = state.active_statuses().into_iter().next().unwrap();
state.record_status_refresh_failure(&cached, "stage status refresh timed out".to_string());
let status = state.runtime_statuses().into_iter().next().unwrap();
assert_eq!(
status.state,
crate::inference::skippy::StageRuntimeState::Failed
);
assert_eq!(
status.error.as_deref(),
Some("stage status refresh timed out")
);
}

View file

@ -245,6 +245,7 @@ pub(super) struct SplitCoordinatorLocalFallbackEvent {
pub(super) reason: &'static str,
pub(super) generation: u64,
pub(super) topology_id: String,
pub(super) run_id: String,
pub(super) unavailable_stage_nodes: Vec<iroh::EndpointId>,
pub(super) ack: tokio::sync::oneshot::Sender<SplitCoordinatorAck>,
}
@ -253,6 +254,7 @@ pub(super) struct SplitCoordinatorWithdrawEvent {
pub(super) reason: &'static str,
pub(super) generation: u64,
pub(super) topology_id: String,
pub(super) run_id: String,
pub(super) unavailable_stage_nodes: Vec<iroh::EndpointId>,
pub(super) ack: tokio::sync::oneshot::Sender<SplitCoordinatorAck>,
}
@ -2310,6 +2312,7 @@ impl SplitTopologyCoordinator {
reason,
generation: self.active.generation,
topology_id: self.active.topology_id.clone(),
run_id: self.active.run_id.clone(),
unavailable_stage_nodes,
ack: ack_tx,
});
@ -2332,6 +2335,7 @@ impl SplitTopologyCoordinator {
reason,
generation: self.active.generation,
topology_id: self.active.topology_id.clone(),
run_id: self.active.run_id.clone(),
unavailable_stage_nodes,
ack: ack_tx,
});

View file

@ -1711,6 +1711,10 @@ async fn startup_handle_local_fallback_event(
let unavailable_stage_nodes =
startup_split_unavailable_stage_nodes(&event.unavailable_stage_nodes);
let old_loaded_name = state.loaded_name.clone();
let withdrew_topology = ctx
.node
.withdraw_stage_topology(&event.topology_id, &event.run_id)
.await;
let Some(old_handle) = state.handle.take() else {
let _ = event.ack.send(SplitCoordinatorAck::Accepted);
return StartupLoopControl::Break;
@ -1834,9 +1838,11 @@ async fn startup_handle_local_fallback_event(
event.topology_id, state.loaded_name
),
context: Some(format!(
"reason={} generation={} unavailable_stage_nodes=[{}] previous_port={} new_port={} new_ctx={}",
"reason={} generation={} run_id={} topology_withdrawn={} unavailable_stage_nodes=[{}] previous_port={} new_port={} new_ctx={}",
event.reason,
event.generation,
event.run_id,
withdrew_topology,
unavailable_stage_nodes,
old_port,
new_port,
@ -1943,14 +1949,22 @@ async fn startup_handle_split_event(
SplitCoordinatorEvent::Withdraw(event) => {
let unavailable_stage_nodes =
startup_split_unavailable_stage_nodes(&event.unavailable_stage_nodes);
let withdrew_topology = ctx
.node
.withdraw_stage_topology(&event.topology_id, &event.run_id)
.await;
let _ = emit_event(OutputEvent::Warning {
message: format!(
"Split runtime topology '{}' lost required stage peer(s); withdrawing model '{}'",
event.topology_id, state.loaded_name
),
context: Some(format!(
"reason={} generation={} unavailable_stage_nodes=[{}]",
event.reason, event.generation, unavailable_stage_nodes
"reason={} generation={} run_id={} topology_withdrawn={} unavailable_stage_nodes=[{}]",
event.reason,
event.generation,
event.run_id,
withdrew_topology,
unavailable_stage_nodes
)),
});
let _ = event.ack.send(SplitCoordinatorAck::Accepted);