legion-ref/mesh-llm
1
0
Fork 0
mirror of https://github.com/Mesh-LLM/mesh-llm.git synced 2026-08-08 22:23:19 -04:00
Code Issues Projects Releases Packages Wiki Activity Actions

Add split coordinator fencing

Browse source
This commit is contained in:
James Dumay 2026-05-11 12:11:55 +10:00
parent 043319c767
commit 04986dfeb5
21 changed files with 1517 additions and 93 deletions
Download patch file Download diff file Expand all files Collapse all files

8
Cargo.lock generated
View file

@ -4101,6 +4101,7 @@ dependencies = [
"serde_json",
"serial_test",
"sha2 0.10.9",
"skippy-coordinator",
"skippy-protocol",
"skippy-runtime",
"skippy-server",
@ -7190,6 +7191,13 @@ dependencies = [
"skippy-protocol",
]
[[package]]
name = "skippy-coordinator"
version = "0.1.0"
dependencies = [
"thiserror 2.0.18",
]
[[package]]
name = "skippy-correctness"
version = "0.1.0"

1
Cargo.toml
View file

@ -19,6 +19,7 @@ members = [
"crates/model-hf",
"crates/model-resolver",
"crates/skippy-protocol",
"crates/skippy-coordinator",
"crates/skippy-topology",
"crates/skippy-cache",
"crates/skippy-metrics",

1
crates/mesh-llm-host-runtime/Cargo.toml
View file

@ -19,6 +19,7 @@ model-ref = { path = "../model-ref" }
model-resolver = { path = "../model-resolver" }
openai-frontend = { path = "../openai-frontend" }
skippy-protocol = { path = "../skippy-protocol" }
skippy-coordinator = { path = "../skippy-coordinator" }
skippy-runtime = { path = "../skippy-runtime" }
skippy-server = { path = "../skippy-server" }
skippy-topology = { path = "../skippy-topology" }

3
crates/mesh-llm-host-runtime/src/api/tests.rs
View file

@ -1536,6 +1536,9 @@ async fn runtime_data_api_routes_remain_payload_stable() {
flash_attn_type: skippy_protocol::FlashAttentionType::Enabled,
error: None,
shutdown_generation: 7,
coordinator_term: 11,
coordinator_id: Some(node.id()),
lease_until_unix_ms: 999_999,
},
)
.await;

4
crates/mesh-llm-host-runtime/src/inference/skippy/deployment.rs
View file

@ -58,6 +58,9 @@ pub(crate) fn remote_stage_load_request(
cache_type_v: context.kv_cache.cache_type_v().to_string(),
flash_attn_type: context.flash_attn_type,
shutdown_generation: 1,
coordinator_term: 0,
coordinator_id: None,
lease_until_unix_ms: 0,
load_mode: LoadMode::LayerPackage,
upstream: None,
downstream,
@ -127,6 +130,7 @@ pub(crate) fn stage_stop_request(
run_id: context.run_id.to_string(),
stage_id: stage.stage_id.clone(),
shutdown_generation,
coordinator_term: shutdown_generation,
}
}

6
crates/mesh-llm-host-runtime/src/inference/skippy/materialization.rs
View file

@ -1507,6 +1507,9 @@ mod tests {
cache_type_v: "f16".to_string(),
flash_attn_type: FlashAttentionType::Auto,
shutdown_generation: 1,
coordinator_term: 0,
coordinator_id: None,
lease_until_unix_ms: 0,
load_mode: LoadMode::LayerPackage,
upstream: None,
downstream: None,
@ -1950,6 +1953,9 @@ mod tests {
cache_type_v: "f16".to_string(),
flash_attn_type: skippy_protocol::FlashAttentionType::Auto,
shutdown_generation: 0,
coordinator_term: 0,
coordinator_id: None,
lease_until_unix_ms: 0,
load_mode: LoadMode::LayerPackage,
upstream: None,
downstream: None,

9
crates/mesh-llm-host-runtime/src/inference/skippy/mod.rs
View file

@ -48,10 +48,11 @@ pub(crate) use package::{
pub(crate) use stage::{
spawn_stage_control_loop, stage_load_timeout, LayerRange, SourceModelKind,
StageCancelPrepareRequest, StageControlCommand, StageControlRequest, StageControlResponse,
StageInventoryRequest, StageLayerInventory, StageLoadRequest, StagePackagePrefetcher,
StagePeerDescriptor, StagePreparationState, StagePreparationStatus,
StagePrepareAcceptedResponse, StagePrepareRequest, StageReadyResponse, StageRuntimeState,
StageStatusAck, StageStatusFilter, StageStatusSnapshot, StageStopRequest, StageWireDType,
StageCoordinatorClaim, StageCoordinatorClaimAck, StageInventoryRequest, StageLayerInventory,
StageLoadRequest, StagePackagePrefetcher, StagePeerDescriptor, StagePreparationState,
StagePreparationStatus, StagePrepareAcceptedResponse, StagePrepareRequest, StageReadyResponse,
StageRuntimeState, StageStatusAck, StageStatusFilter, StageStatusSnapshot, StageStopRequest,
StageWireDType,
};
pub(crate) use topology::{plan_package_identity_topology, StageTopologyParticipant};

12
crates/mesh-llm-host-runtime/src/inference/skippy/stage/inventory.rs
View file

@ -89,7 +89,7 @@ pub(super) async fn run_stage_prepare_task(
if !update_preparation(
&preparations,
&key,
preparation_status_from_load(&load, StagePreparationState::Resolving),
preparation_status_from_load(&load, StagePreparationState::Resolving, None),
)
.await
|| cancelled.load(Ordering::Acquire)
@ -107,7 +107,7 @@ pub(super) async fn run_stage_prepare_task(
&& !update_preparation(
&preparations,
&key,
preparation_status_from_load(&load, StagePreparationState::Downloading),
preparation_status_from_load(&load, StagePreparationState::Downloading, None),
)
.await
{
@ -119,13 +119,15 @@ pub(super) async fn run_stage_prepare_task(
}
let state = match result {
Ok(PrepareSourceResult { bytes_total }) => {
let mut status = preparation_status_from_load(&load, StagePreparationState::Available);
let mut status =
preparation_status_from_load(&load, StagePreparationState::Available, None);
status.bytes_done = bytes_total;
status.bytes_total = bytes_total;
status
}
Err(error) => {
let mut status = preparation_status_from_load(&load, StagePreparationState::Failed);
let mut status =
preparation_status_from_load(&load, StagePreparationState::Failed, None);
status.error = Some(match peer_prefetch_error {
Some(prefetch_error) => {
format!("{error}; peer artifact prefetch failed: {prefetch_error}")
@ -152,7 +154,7 @@ async fn prefetch_stage_package_if_needed(
let _ = update_preparation(
preparations,
key,
preparation_status_from_load(load, StagePreparationState::Downloading),
preparation_status_from_load(load, StagePreparationState::Downloading, None),
)
.await;
match prefetcher.prefetch_stage_package(request).await {

205
crates/mesh-llm-host-runtime/src/inference/skippy/stage/mod.rs
View file

@ -5,10 +5,11 @@ use std::{
atomic::{AtomicBool, Ordering},
Arc,
},
time::Duration,
time::{Duration, SystemTime, UNIX_EPOCH},
};
use anyhow::{anyhow, Context, Result};
use skippy_coordinator::{ClaimDecision, ClaimFence, LoadClaimRef};
use skippy_protocol::{FlashAttentionType, LoadMode, PeerConfig, StageConfig};
use skippy_server::{
binary_transport::{BinaryStageOptions, WireCondition},
@ -39,6 +40,7 @@ struct RunningStage {
#[derive(Default)]
struct StageControlState {
stages: HashMap<String, RunningStage>,
coordinator_claims: ClaimFence,
preparations: Arc<Mutex<HashMap<String, StagePreparationStatus>>>,
preparation_tasks: HashMap<String, StagePreparationTask>,
package_prefetcher: Option<Arc<dyn StagePackagePrefetcher>>,
@ -74,6 +76,10 @@ pub(crate) fn spawn_stage_control_loop(
impl StageControlState {
async fn handle(&mut self, request: StageControlRequest) -> Result<StageControlResponse> {
match request {
StageControlRequest::Claim(claim) => self
.claim(claim)
.await
.map(StageControlResponse::ClaimAccepted),
StageControlRequest::Load(load) => {
self.load(load).await.map(StageControlResponse::Ready)
}
@ -98,6 +104,46 @@ impl StageControlState {
}
}
async fn claim(&mut self, claim: StageCoordinatorClaim) -> Result<StageCoordinatorClaimAck> {
match self
.coordinator_claims
.accept_claim(claim, current_time_unix_ms())
{
ClaimDecision::Accepted {
supersedes_term: Some(_),
claim,
} => {
self.fence_stale_runtime_for_claim(&claim).await?;
Ok(StageCoordinatorClaimAck {
accepted: true,
claim,
error: None,
})
}
ClaimDecision::Accepted { claim, .. } => Ok(StageCoordinatorClaimAck {
accepted: true,
claim,
error: None,
}),
ClaimDecision::Rejected { current, reason } => Ok(StageCoordinatorClaimAck {
accepted: false,
claim: current.unwrap_or_else(|| StageCoordinatorClaim {
model_id: String::new(),
package_ref: String::new(),
manifest_sha256: String::new(),
topology_id: String::new(),
run_id: String::new(),
coordinator_id: String::new(),
coordinator_term: 0,
participant_set_hash: String::new(),
topology_hash: String::new(),
lease_until_unix_ms: 0,
}),
error: Some(reason.to_string()),
}),
}
}
async fn inventory(&self, request: StageInventoryRequest) -> StageLayerInventory {
let preparing_ranges = self
.preparations
@ -169,12 +215,24 @@ impl StageControlState {
&mut self,
request: StagePrepareRequest,
) -> Result<StagePrepareAcceptedResponse> {
if let Some(error) = self.validate_load_claim(&request.load) {
return Ok(StagePrepareAcceptedResponse {
accepted: false,
status: preparation_status_from_load(
&request.load,
StagePreparationState::Failed,
Some(error.clone()),
),
error: Some(error),
});
}
let key = stage_key(
&request.load.topology_id,
&request.load.run_id,
&request.load.stage_id,
);
let status = preparation_status_from_load(&request.load, StagePreparationState::Assigned);
let status =
preparation_status_from_load(&request.load, StagePreparationState::Assigned, None);
{
let mut preparations = self.preparations.lock().await;
if let Some(existing) = preparations.get(&key) {
@ -287,6 +345,13 @@ impl StageControlState {
"unsupported stage backend '{}'",
load.backend
);
if let Some(error) = self.validate_load_claim(&load) {
return Ok(StageReadyResponse {
accepted: false,
status: failed_status_from_load(&load, error.clone()),
error: Some(error),
});
}
let key = stage_key(&load.topology_id, &load.run_id, &load.stage_id);
if let Some(existing) = self.stages.remove(&key) {
existing.server.shutdown().await?;
@ -368,6 +433,19 @@ impl StageControlState {
error: None,
});
};
if stop.coordinator_term < existing.load.coordinator_term {
let current_term = existing.load.coordinator_term;
let status = status_from_running(&existing);
self.stages.insert(key, existing);
return Ok(StageReadyResponse {
accepted: false,
status,
error: Some(format!(
"stale coordinator term {} < {}",
stop.coordinator_term, current_term
)),
});
}
if stop.shutdown_generation < existing.load.shutdown_generation {
let status = status_from_running(&existing);
self.stages.insert(key, existing);
@ -396,6 +474,59 @@ impl StageControlState {
.map(status_from_running)
.collect()
}
fn validate_load_claim(&self, load: &StageLoadRequest) -> Option<String> {
if load.coordinator_term == 0 && load.coordinator_id.is_none() {
return None;
}
self.coordinator_claims
.validate_load(&load_claim_ref(load), current_time_unix_ms())
.err()
.map(|error| error.to_string())
}
async fn fence_stale_runtime_for_claim(&mut self, claim: &StageCoordinatorClaim) -> Result<()> {
let stale_keys = self
.stages
.iter()
.filter_map(|(key, stage)| {
(stage.load.model_id == claim.model_id
&& stage.load.package_ref == claim.package_ref
&& stage.load.manifest_sha256 == claim.manifest_sha256
&& stage.load.coordinator_term < claim.coordinator_term)
.then_some(key.clone())
})
.collect::<Vec<_>>();
for key in stale_keys {
if let Some(stage) = self.stages.remove(&key) {
stage.server.shutdown().await?;
}
}
let mut preparations = self.preparations.lock().await;
let stale_preparations = preparations
.iter()
.filter_map(|(key, status)| {
(status.model_id == claim.model_id
&& status.package_ref == claim.package_ref
&& status.manifest_sha256 == claim.manifest_sha256
&& status.coordinator_term < claim.coordinator_term)
.then_some(key.clone())
})
.collect::<Vec<_>>();
for key in stale_preparations {
if let Some(task) = self.preparation_tasks.remove(&key) {
task.cancelled.store(true, Ordering::Release);
task.handle.abort();
}
if let Some(status) = preparations.get_mut(&key) {
status.state = StagePreparationState::Cancelled;
status.error = Some("superseded by newer coordinator term".to_string());
}
}
Ok(())
}
}
impl StageStatusFilter {
@ -418,6 +549,25 @@ fn stage_key(topology_id: &str, run_id: &str, stage_id: &str) -> String {
format!("{topology_id}\n{run_id}\n{stage_id}")
}
fn current_time_unix_ms() -> u64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
.as_millis() as u64
}
fn load_claim_ref(load: &StageLoadRequest) -> LoadClaimRef {
LoadClaimRef {
model_id: load.model_id.clone(),
package_ref: load.package_ref.clone(),
manifest_sha256: load.manifest_sha256.clone(),
topology_id: load.topology_id.clone(),
run_id: load.run_id.clone(),
coordinator_id: load.coordinator_id.map(|id| id.to_string()),
coordinator_term: load.coordinator_term,
}
}
fn parse_bind_addr(bind_addr: &str) -> Result<SocketAddr> {
bind_addr
.parse()
@ -643,6 +793,9 @@ fn status_from_running(stage: &RunningStage) -> StageStatusSnapshot {
flash_attn_type: stage.load.flash_attn_type,
error: server.last_error.clone(),
shutdown_generation: stage.load.shutdown_generation,
coordinator_term: stage.load.coordinator_term,
coordinator_id: stage.load.coordinator_id,
lease_until_unix_ms: stage.load.lease_until_unix_ms,
}
}
@ -676,12 +829,52 @@ fn stopped_status(stop: &StageStopRequest) -> StageStatusSnapshot {
flash_attn_type: FlashAttentionType::Auto,
error: None,
shutdown_generation: stop.shutdown_generation,
coordinator_term: stop.coordinator_term,
coordinator_id: None,
lease_until_unix_ms: 0,
}
}
fn failed_status_from_load(load: &StageLoadRequest, error: String) -> StageStatusSnapshot {
StageStatusSnapshot {
topology_id: load.topology_id.clone(),
run_id: load.run_id.clone(),
model_id: load.model_id.clone(),
backend: load.backend.clone(),
package_ref: Some(load.package_ref.clone()),
manifest_sha256: Some(load.manifest_sha256.clone()),
source_model_path: load.model_path.clone(),
source_model_sha256: None,
source_model_bytes: load.source_model_bytes,
materialized_path: None,
materialized_pinned: false,
projector_path: load.projector_path.clone(),
stage_id: load.stage_id.clone(),
stage_index: load.stage_index,
layer_start: load.layer_start,
layer_end: load.layer_end,
state: StageRuntimeState::Failed,
bind_addr: load.bind_addr.clone(),
activation_width: load.activation_width.max(0) as u32,
wire_dtype: load.wire_dtype,
selected_device: load.selected_device.clone(),
ctx_size: load.ctx_size,
lane_count: load.lane_count,
n_batch: load.n_batch,
n_ubatch: load.n_ubatch,
flash_attn_type: load.flash_attn_type,
error: Some(error),
shutdown_generation: load.shutdown_generation,
coordinator_term: load.coordinator_term,
coordinator_id: load.coordinator_id,
lease_until_unix_ms: load.lease_until_unix_ms,
}
}
fn preparation_status_from_load(
load: &StageLoadRequest,
state: StagePreparationState,
error: Option<String>,
) -> StagePreparationStatus {
StagePreparationStatus {
topology_id: load.topology_id.clone(),
@ -698,8 +891,11 @@ fn preparation_status_from_load(
bytes_done: None,
bytes_total: None,
bind_addr: None,
error: None,
error,
shutdown_generation: load.shutdown_generation,
coordinator_term: load.coordinator_term,
coordinator_id: load.coordinator_id,
lease_until_unix_ms: load.lease_until_unix_ms,
}
}
@ -721,6 +917,9 @@ fn preparation_status_from_cancel(cancel: StageCancelPrepareRequest) -> StagePre
bind_addr: None,
error: None,
shutdown_generation: cancel.shutdown_generation,
coordinator_term: 0,
coordinator_id: None,
lease_until_unix_ms: 0,
}
}

78
crates/mesh-llm-host-runtime/src/inference/skippy/stage/tests.rs
View file

@ -42,6 +42,9 @@ fn load_request() -> StageLoadRequest {
cache_type_v: "q8_0".to_string(),
flash_attn_type: FlashAttentionType::Enabled,
shutdown_generation: 7,
coordinator_term: 0,
coordinator_id: None,
lease_until_unix_ms: 0,
load_mode: LoadMode::RuntimeSlice,
upstream: None,
downstream: Some(StagePeerDescriptor {
@ -53,6 +56,28 @@ fn load_request() -> StageLoadRequest {
}
}
fn coordinator_id() -> iroh::EndpointId {
iroh::EndpointId::from(iroh::SecretKey::from_bytes(&[0x5a; 32]).public())
}
fn coordinator_claim_from_load(
load: &StageLoadRequest,
coordinator_id: iroh::EndpointId,
) -> StageCoordinatorClaim {
StageCoordinatorClaim {
model_id: load.model_id.clone(),
package_ref: load.package_ref.clone(),
manifest_sha256: load.manifest_sha256.clone(),
topology_id: load.topology_id.clone(),
run_id: load.run_id.clone(),
coordinator_id: coordinator_id.to_string(),
coordinator_term: load.coordinator_term,
participant_set_hash: "participants".to_string(),
topology_hash: "topology".to_string(),
lease_until_unix_ms: u64::MAX,
}
}
struct BlockingPackagePrefetcher {
started: TokioMutex<Option<oneshot::Sender<()>>>,
release: TokioMutex<Option<oneshot::Receiver<Result<()>>>>,
@ -88,6 +113,53 @@ impl StagePackagePrefetcher for BlockingPackagePrefetcher {
}
}
#[tokio::test]
async fn fenced_prepare_requires_accepted_coordinator_claim() {
let mut load = load_request();
let coordinator_id = coordinator_id();
load.coordinator_term = 11;
load.coordinator_id = Some(coordinator_id);
load.lease_until_unix_ms = u64::MAX;
let mut state = StageControlState::default();
let response = state
.prepare(StagePrepareRequest {
load,
coordinator_id: None,
})
.await
.unwrap();
assert!(!response.accepted);
assert_eq!(response.error.as_deref(), Some("missing coordinator claim"));
assert_eq!(response.status.state, StagePreparationState::Failed);
}
#[tokio::test]
async fn accepted_coordinator_claim_allows_fenced_prepare() {
let mut load = load_request();
let coordinator_id = coordinator_id();
load.coordinator_term = 11;
load.coordinator_id = Some(coordinator_id);
load.lease_until_unix_ms = u64::MAX;
let claim = coordinator_claim_from_load(&load, coordinator_id);
let mut state = StageControlState::default();
let ack = state.claim(claim).await.unwrap();
assert!(ack.accepted);
let response = state
.prepare(StagePrepareRequest {
load,
coordinator_id: None,
})
.await
.unwrap();
assert!(response.accepted);
assert_eq!(response.status.state, StagePreparationState::Assigned);
}
#[test]
fn stage_config_preserves_backend_neutral_load_fields() {
let request = load_request();
@ -490,7 +562,7 @@ async fn stale_cancel_prepare_keeps_newer_prepare_status() {
let load = load_request();
let key = stage_key(&load.topology_id, &load.run_id, &load.stage_id);
let mut state = StageControlState::default();
let current = preparation_status_from_load(&load, StagePreparationState::Resolving);
let current = preparation_status_from_load(&load, StagePreparationState::Resolving, None);
state.preparations.lock().await.insert(key, current.clone());
let status = state
@ -511,7 +583,7 @@ async fn stale_cancel_prepare_keeps_newer_prepare_status() {
async fn status_update_upserts_preparation_status_and_rejects_stale_generation() {
let load = load_request();
let mut state = StageControlState::default();
let mut update = preparation_status_from_load(&load, StagePreparationState::Loading);
let mut update = preparation_status_from_load(&load, StagePreparationState::Loading, None);
update.bytes_done = Some(1024);
update.bytes_total = Some(4096);
@ -564,7 +636,7 @@ async fn inventory_retains_failed_prepare_status() {
let load = load_request();
let key = stage_key(&load.topology_id, &load.run_id, &load.stage_id);
let state = StageControlState::default();
let mut failed = preparation_status_from_load(&load, StagePreparationState::Failed);
let mut failed = preparation_status_from_load(&load, StagePreparationState::Failed, None);
failed.error = Some("source unavailable".to_string());
state.preparations.lock().await.insert(key, failed);

21
crates/mesh-llm-host-runtime/src/inference/skippy/stage/types.rs
View file

@ -11,6 +11,7 @@ pub(crate) struct StageControlCommand {
#[derive(Clone, Debug)]
#[allow(clippy::large_enum_variant)]
pub(crate) enum StageControlRequest {
Claim(StageCoordinatorClaim),
Load(StageLoadRequest),
Stop(StageStopRequest),
Status(StageStatusFilter),
@ -23,6 +24,7 @@ pub(crate) enum StageControlRequest {
#[derive(Clone, Debug)]
#[allow(clippy::large_enum_variant)]
pub(crate) enum StageControlResponse {
ClaimAccepted(StageCoordinatorClaimAck),
Ready(StageReadyResponse),
Status(Vec<StageStatusSnapshot>),
Inventory(StageLayerInventory),
@ -31,6 +33,15 @@ pub(crate) enum StageControlResponse {
StatusAck(StageStatusAck),
}
pub(crate) type StageCoordinatorClaim = skippy_coordinator::CoordinatorClaim;
#[derive(Clone, Debug)]
pub(crate) struct StageCoordinatorClaimAck {
pub(crate) accepted: bool,
pub(crate) claim: StageCoordinatorClaim,
pub(crate) error: Option<String>,
}
#[derive(Clone, Debug)]
pub(crate) struct StageLoadRequest {
pub(crate) topology_id: String,
@ -59,6 +70,9 @@ pub(crate) struct StageLoadRequest {
pub(crate) cache_type_v: String,
pub(crate) flash_attn_type: FlashAttentionType,
pub(crate) shutdown_generation: u64,
pub(crate) coordinator_term: u64,
pub(crate) coordinator_id: Option<iroh::EndpointId>,
pub(crate) lease_until_unix_ms: u64,
pub(crate) load_mode: LoadMode,
pub(crate) upstream: Option<StagePeerDescriptor>,
pub(crate) downstream: Option<StagePeerDescriptor>,
@ -70,6 +84,7 @@ pub(crate) struct StageStopRequest {
pub(crate) run_id: String,
pub(crate) stage_id: String,
pub(crate) shutdown_generation: u64,
pub(crate) coordinator_term: u64,
}
#[derive(Clone, Debug, Default)]
@ -202,6 +217,9 @@ pub(crate) struct StageStatusSnapshot {
pub(crate) flash_attn_type: FlashAttentionType,
pub(crate) error: Option<String>,
pub(crate) shutdown_generation: u64,
pub(crate) coordinator_term: u64,
pub(crate) coordinator_id: Option<iroh::EndpointId>,
pub(crate) lease_until_unix_ms: u64,
}
#[derive(Clone, Debug)]
@ -222,6 +240,9 @@ pub(crate) struct StagePreparationStatus {
pub(crate) bind_addr: Option<String>,
pub(crate) error: Option<String>,
pub(crate) shutdown_generation: u64,
pub(crate) coordinator_term: u64,
pub(crate) coordinator_id: Option<iroh::EndpointId>,
pub(crate) lease_until_unix_ms: u64,
}
#[derive(Clone, Debug)]

125
crates/mesh-llm-host-runtime/src/mesh/mod.rs
View file

@ -1969,16 +1969,17 @@ impl Node {
request: &crate::inference::skippy::StageControlRequest,
) -> std::time::Duration {
match request {
crate::inference::skippy::StageControlRequest::Load(load) => {
crate::inference::skippy::stage_load_timeout(load)
}
crate::inference::skippy::StageControlRequest::Stop(_)
crate::inference::skippy::StageControlRequest::Claim(_)
| crate::inference::skippy::StageControlRequest::Stop(_)
| crate::inference::skippy::StageControlRequest::Status(_)
| crate::inference::skippy::StageControlRequest::Inventory(_)
| crate::inference::skippy::StageControlRequest::CancelPrepare(_)
| crate::inference::skippy::StageControlRequest::StatusUpdate(_) => {
std::time::Duration::from_secs(30)
}
crate::inference::skippy::StageControlRequest::Load(load) => {
crate::inference::skippy::stage_load_timeout(load)
}
crate::inference::skippy::StageControlRequest::Prepare(prepare) => {
crate::inference::skippy::stage_load_timeout(&prepare.load)
}
@ -4585,6 +4586,7 @@ impl Node {
request: &mut crate::inference::skippy::StageControlRequest,
) -> anyhow::Result<()> {
match request {
crate::inference::skippy::StageControlRequest::Claim(_) => {}
crate::inference::skippy::StageControlRequest::Load(load) => {
if load.load_mode == skippy_protocol::LoadMode::RuntimeSlice
&& load
@ -5832,6 +5834,9 @@ fn stage_snapshot_from_runtime_status(
flash_attn_type: status.flash_attn_type,
error,
shutdown_generation: status.shutdown_generation,
coordinator_term: 0,
coordinator_id: None,
lease_until_unix_ms: 0,
}
}
@ -5865,6 +5870,9 @@ fn stage_control_request_to_proto(
use skippy_stage_proto::stage_control_request::Command;
let command = match request {
crate::inference::skippy::StageControlRequest::Claim(claim) => {
Command::ClaimCoordinator(stage_coordinator_claim_to_proto(claim))
}
crate::inference::skippy::StageControlRequest::Load(load) => {
Command::LoadStage(stage_load_to_proto(load))
}
@ -5874,6 +5882,7 @@ fn stage_control_request_to_proto(
run_id: stop.run_id,
stage_id: stop.stage_id,
shutdown_generation: stop.shutdown_generation,
coordinator_term: stop.coordinator_term,
})
}
crate::inference::skippy::StageControlRequest::Status(status) => {
@ -5948,6 +5957,9 @@ fn stage_load_to_proto(
cache_type_v: load.cache_type_v,
flash_attn_type: stage_flash_attn_type_to_proto(load.flash_attn_type) as i32,
shutdown_generation: load.shutdown_generation,
coordinator_term: load.coordinator_term,
coordinator_id: load.coordinator_id.map(|id| id.to_string()),
lease_until_unix_ms: load.lease_until_unix_ms,
load_mode: match load.load_mode {
skippy_protocol::LoadMode::RuntimeSlice => {
skippy_stage_proto::StageLoadMode::RuntimeSlice as i32
@ -5964,6 +5976,23 @@ fn stage_load_to_proto(
}
}
fn stage_coordinator_claim_to_proto(
claim: crate::inference::skippy::StageCoordinatorClaim,
) -> skippy_stage_proto::ClaimCoordinator {
skippy_stage_proto::ClaimCoordinator {
model_id: claim.model_id,
package_ref: claim.package_ref,
manifest_sha256: claim.manifest_sha256,
topology_id: claim.topology_id,
run_id: claim.run_id,
coordinator_id: claim.coordinator_id,
coordinator_term: claim.coordinator_term,
participant_set_hash: claim.participant_set_hash,
topology_hash: claim.topology_hash,
lease_until_unix_ms: claim.lease_until_unix_ms,
}
}
fn stage_peer_to_proto(
peer: crate::inference::skippy::StagePeerDescriptor,
) -> skippy_stage_proto::StagePeer {
@ -5993,6 +6022,11 @@ fn stage_control_request_from_proto(
.command
.ok_or_else(|| anyhow::anyhow!("missing stage control command"))?
{
Command::ClaimCoordinator(claim) => {
Ok(crate::inference::skippy::StageControlRequest::Claim(
stage_coordinator_claim_from_proto(claim)?,
))
}
Command::LoadStage(load) => Ok(crate::inference::skippy::StageControlRequest::Load(
stage_load_from_proto(load)?,
)),
@ -6002,6 +6036,7 @@ fn stage_control_request_from_proto(
run_id: stop.run_id,
stage_id: stop.stage_id,
shutdown_generation: stop.shutdown_generation,
coordinator_term: stop.coordinator_term,
},
)),
Command::GetStageStatus(status) => {
@ -6096,12 +6131,36 @@ fn stage_load_from_proto(
cache_type_v: load.cache_type_v,
flash_attn_type: stage_flash_attn_type_from_proto(load.flash_attn_type),
shutdown_generation: load.shutdown_generation,
coordinator_term: load.coordinator_term,
coordinator_id: load
.coordinator_id
.map(|id| id.parse())
.transpose()
.context("invalid stage load coordinator_id")?,
lease_until_unix_ms: load.lease_until_unix_ms,
load_mode: stage_load_mode_from_proto(load.load_mode),
upstream: load.upstream.map(stage_peer_from_proto).transpose()?,
downstream: load.downstream.map(stage_peer_from_proto).transpose()?,
})
}
fn stage_coordinator_claim_from_proto(
claim: skippy_stage_proto::ClaimCoordinator,
) -> anyhow::Result<crate::inference::skippy::StageCoordinatorClaim> {
Ok(crate::inference::skippy::StageCoordinatorClaim {
model_id: claim.model_id,
package_ref: claim.package_ref,
manifest_sha256: claim.manifest_sha256,
topology_id: claim.topology_id,
run_id: claim.run_id,
coordinator_id: claim.coordinator_id,
coordinator_term: claim.coordinator_term,
participant_set_hash: claim.participant_set_hash,
topology_hash: claim.topology_hash,
lease_until_unix_ms: claim.lease_until_unix_ms,
})
}
fn stage_device_from_proto(
device: skippy_stage_proto::StageDevice,
) -> anyhow::Result<skippy_protocol::StageDevice> {
@ -6168,6 +6227,15 @@ fn stage_control_unavailable_response(
request: crate::inference::skippy::StageControlRequest,
) -> crate::inference::skippy::StageControlResponse {
let status = match request {
crate::inference::skippy::StageControlRequest::Claim(claim) => {
return crate::inference::skippy::StageControlResponse::ClaimAccepted(
crate::inference::skippy::StageCoordinatorClaimAck {
accepted: false,
claim,
error: Some("stage control is not available".to_string()),
},
);
}
crate::inference::skippy::StageControlRequest::Load(load) => {
stage_status_from_load(&load, crate::inference::skippy::StageRuntimeState::Failed)
}
@ -6201,6 +6269,9 @@ fn stage_control_unavailable_response(
flash_attn_type: skippy_protocol::FlashAttentionType::Auto,
error: Some("stage control is not available".to_string()),
shutdown_generation: stop.shutdown_generation,
coordinator_term: stop.coordinator_term,
coordinator_id: None,
lease_until_unix_ms: 0,
}
}
crate::inference::skippy::StageControlRequest::Status(_) => {
@ -6296,6 +6367,9 @@ fn stage_status_from_load(
flash_attn_type: load.flash_attn_type,
error: Some("stage control is not available".to_string()),
shutdown_generation: load.shutdown_generation,
coordinator_term: load.coordinator_term,
coordinator_id: load.coordinator_id,
lease_until_unix_ms: load.lease_until_unix_ms,
}
}
@ -6321,6 +6395,9 @@ fn stage_preparation_status_from_load(
bind_addr: None,
error,
shutdown_generation: load.shutdown_generation,
coordinator_term: load.coordinator_term,
coordinator_id: load.coordinator_id,
lease_until_unix_ms: load.lease_until_unix_ms,
}
}
@ -6346,6 +6423,9 @@ fn stage_preparation_status_from_cancel(
bind_addr: None,
error,
shutdown_generation: cancel.shutdown_generation,
coordinator_term: 0,
coordinator_id: None,
lease_until_unix_ms: 0,
}
}
@ -6356,6 +6436,13 @@ fn stage_control_response_to_proto(
use skippy_stage_proto::stage_control_response::Response;
let response = match response {
crate::inference::skippy::StageControlResponse::ClaimAccepted(accepted) => {
Response::CoordinatorClaimAccepted(skippy_stage_proto::CoordinatorClaimAccepted {
accepted: accepted.accepted,
claim: Some(stage_coordinator_claim_to_proto(accepted.claim)),
error: accepted.error,
})
}
crate::inference::skippy::StageControlResponse::Ready(ready) => {
Response::StageReady(skippy_stage_proto::StageReady {
accepted: ready.accepted,
@ -6414,6 +6501,20 @@ fn stage_control_response_from_proto(
.response
.ok_or_else(|| anyhow::anyhow!("missing stage control response"))?
{
Response::CoordinatorClaimAccepted(accepted) => {
let claim = accepted
.claim
.ok_or_else(|| anyhow::anyhow!("coordinator claim accepted missing claim"))?;
Ok(
crate::inference::skippy::StageControlResponse::ClaimAccepted(
crate::inference::skippy::StageCoordinatorClaimAck {
accepted: accepted.accepted,
claim: stage_coordinator_claim_from_proto(claim)?,
error: accepted.error,
},
),
)
}
Response::StageReady(ready) => {
let status = ready
.status
@ -6619,6 +6720,9 @@ fn stage_preparation_status_to_proto(
bind_addr: status.bind_addr,
error: status.error,
shutdown_generation: status.shutdown_generation,
coordinator_term: status.coordinator_term,
coordinator_id: status.coordinator_id.map(|id| id.to_string()),
lease_until_unix_ms: status.lease_until_unix_ms,
}
}
@ -6642,6 +6746,9 @@ fn stage_preparation_status_from_proto(
bind_addr: status.bind_addr,
error: status.error,
shutdown_generation: status.shutdown_generation,
coordinator_term: status.coordinator_term,
coordinator_id: status.coordinator_id.and_then(|id| id.parse().ok()),
lease_until_unix_ms: status.lease_until_unix_ms,
}
}
@ -6677,6 +6784,9 @@ fn stage_status_to_proto(
materialized_pinned: Some(status.materialized_pinned),
projector_path: status.projector_path,
flash_attn_type: stage_flash_attn_type_to_proto(status.flash_attn_type) as i32,
coordinator_term: status.coordinator_term,
coordinator_id: status.coordinator_id.map(|id| id.to_string()),
lease_until_unix_ms: status.lease_until_unix_ms,
}
}
@ -6719,6 +6829,13 @@ fn stage_status_from_proto(
flash_attn_type: stage_flash_attn_type_from_proto(status.flash_attn_type),
error: status.error,
shutdown_generation: status.shutdown_generation,
coordinator_term: status.coordinator_term,
coordinator_id: status
.coordinator_id
.map(|id| id.parse())
.transpose()
.context("invalid stage status coordinator_id")?,
lease_until_unix_ms: status.lease_until_unix_ms,
})
}

41
crates/mesh-llm-host-runtime/src/mesh/tests.rs
View file

@ -61,6 +61,9 @@ fn stage_load_request() -> crate::inference::skippy::StageLoadRequest {
cache_type_v: "q8_0".to_string(),
flash_attn_type: skippy_protocol::FlashAttentionType::Auto,
shutdown_generation: 3,
coordinator_term: 11,
coordinator_id: None,
lease_until_unix_ms: 999_999,
load_mode: skippy_protocol::LoadMode::RuntimeSlice,
upstream: None,
downstream: None,
@ -1835,12 +1838,19 @@ fn gossip_frame_roundtrip_preserves_scanned_model_metadata() {
canonical_ref: Some("hf/bartowski/Qwen3-8B-GGUF/Qwen3-8B-Q4_K_M.gguf".into()),
repository: Some("bartowski/Qwen3-8B-GGUF".into()),
revision: Some("main".into()),
artifact: Some("Qwen3-8B-Q4_K_M.gguf".into()),
local_file_name: Some("Qwen3-8B-Q4_K_M.gguf".into()),
identity_hash: Some("identity-hash".into()),
},
format: ModelFormat::LlamaGGUF,
quantization: "Q4_K_M".to_string(),
size_bytes: 4_800_000_000,
capabilities: Default::default(),
topology: None,
}],
served_model_runtime: vec![ModelRuntimeDescriptor {
model_name: "Qwen3-8B-Q4_K_M".to_string(),
identity_hash: Some("identity-hash".to_string()),
context_length: Some(32768),
ready: true,
}],
served_model_runtime: vec![],
owner_attestation: None,
artifact_transfer_supported: false,
stage_status_list_supported: false,
@ -2078,7 +2088,7 @@ fn transitive_peer_update_refreshes_metadata_fields() {
latency_observer_id: None,
};
apply_transitive_ann(&mut existing, &addr, &ann, test_endpoint_id(0xee));
apply_transitive_ann(&mut existing, &addr, &ann, make_test_endpoint_id(0xee));
assert!(
existing.available_models.is_empty(),
@ -2162,7 +2172,7 @@ fn transitive_peer_merge_preserves_richer_direct_address() {
latency_observer_id: None,
};
apply_transitive_ann(&mut existing, &weak_addr, &ann, test_endpoint_id(0xee));
apply_transitive_ann(&mut existing, &weak_addr, &ann, make_test_endpoint_id(0xee));
assert_eq!(
existing.addr.addrs.len(),
@ -2214,8 +2224,17 @@ fn transitive_peer_merge_preserves_richer_direct_address() {
owner_attestation: None,
artifact_transfer_supported: true,
stage_status_list_supported: true,
latency_ms: None,
latency_source: None,
latency_age_ms: None,
latency_observer_id: None,
};
apply_transitive_ann(&mut existing, &richer_addr, &ann2, test_endpoint_id(0xee));
apply_transitive_ann(
&mut existing,
&richer_addr,
&ann2,
make_test_endpoint_id(0xee),
);
assert_eq!(
existing.addr.addrs.len(),
@ -2792,7 +2811,7 @@ fn transitive_peer_update_refreshes_last_mentioned() {
latency_observer_id: None,
};
apply_transitive_ann(&mut peer, &addr, &ann, test_endpoint_id(0xee));
apply_transitive_ann(&mut peer, &addr, &ann, make_test_endpoint_id(0xee));
// Before refreshing last_mentioned, verify the peer WOULD be pruned.
let prune_cutoff_pre =
@ -5586,6 +5605,9 @@ fn test_stage_load_request() -> crate::inference::skippy::StageLoadRequest {
cache_type_v: "f16".to_string(),
flash_attn_type: skippy_protocol::FlashAttentionType::Auto,
shutdown_generation: 7,
coordinator_term: 11,
coordinator_id: Some(make_test_endpoint_id(0x70)),
lease_until_unix_ms: 999_999,
load_mode: skippy_protocol::LoadMode::RuntimeSlice,
upstream: None,
downstream: Some(crate::inference::skippy::StagePeerDescriptor {
@ -5617,6 +5639,9 @@ fn test_preparation_status(
bind_addr: Some("127.0.0.1:51234".to_string()),
error: None,
shutdown_generation: 7,
coordinator_term: 11,
coordinator_id: Some(make_test_endpoint_id(0x70)),
lease_until_unix_ms: 999_999,
}
}

9
crates/mesh-llm-host-runtime/src/protocol/mod.rs
View file

@ -1265,6 +1265,7 @@ mod tests {
gpu_compute_tflops_fp16: None,
available_model_metadata: vec![],
experts_summary: None,
available_model_sizes: HashMap::new(),
served_model_descriptors: vec![],
served_model_runtime: vec![],
owner_attestation: Some(crate::crypto::SignedNodeOwnership {
@ -2091,6 +2092,10 @@ mod tests {
owner_attestation: None,
artifact_transfer_supported: true,
stage_status_list_supported: true,
latency_ms: None,
latency_source: None,
latency_age_ms: None,
latency_observer_id: None,
};
let proto_pa = local_ann_to_proto_ann(&ann_with_timestamp);
@ -2137,6 +2142,10 @@ mod tests {
owner_attestation: None,
artifact_transfer_supported: false,
stage_status_list_supported: false,
latency_ms: None,
latency_source: None,
latency_age_ms: None,
latency_observer_id: None,
};
let proto_pa = local_ann_to_proto_ann(&ann_without_timestamp);

139
crates/mesh-llm-host-runtime/src/runtime/local.rs
View file

@ -10,6 +10,7 @@ use crate::runtime_data::{
RuntimeLlamaEndpointStatus, RuntimeLlamaSlotSnapshot, RuntimeLlamaSlotsSnapshot,
};
use anyhow::{Context, Result};
use sha2::{Digest, Sha256};
use skippy_protocol::{FlashAttentionType, LoadMode, PeerConfig, StageConfig};
use std::collections::HashMap;
use std::path::{Path, PathBuf};
@ -19,6 +20,7 @@ const SPLIT_PARTICIPANT_POLL_INTERVAL: Duration = Duration::from_millis(500);
const SPLIT_PARTICIPANT_STABLE_FOR: Duration = Duration::from_secs(2);
const SPLIT_DEFAULT_MIN_PARTICIPANTS: usize = 2;
const SPLIT_INITIAL_SHUTDOWN_GENERATION: u64 = 1;
const SPLIT_COORDINATOR_LEASE_SECS: u64 = 4 * 60 * 60;
const RUNTIME_MODEL_FIT_HEADROOM_NUMERATOR: u64 = 11;
const RUNTIME_MODEL_FIT_HEADROOM_DENOMINATOR: u64 = 10;
@ -120,6 +122,10 @@ fn current_time_unix_ms() -> u64 {
.as_millis() as u64
}
fn split_coordinator_lease_until_unix_ms() -> u64 {
current_time_unix_ms().saturating_add(SPLIT_COORDINATOR_LEASE_SECS.saturating_mul(1000))
}
pub(super) struct ManagedModelController {
pub(super) model_name: String,
pub(super) stop_tx: tokio::sync::watch::Sender<bool>,
@ -787,6 +793,8 @@ async fn load_split_runtime_generation_inner(
spec.node.id().fmt_short()
);
claim_split_coordinator_lease(spec.node, spec.model_ref, spec.package, spec.generation).await?;
let mut ready_by_stage: HashMap<String, skippy::StageStatusSnapshot> = HashMap::new();
let mut downstream: Option<skippy::StagePeerDescriptor> = None;
let kv_cache = skippy::KvCachePolicy::for_model_size(spec.package.source_model_bytes);
@ -864,6 +872,9 @@ async fn load_split_runtime_generation_inner(
cache_type_v: effective_cache_type_v.clone(),
flash_attn_type: resolved_flash_attn_type,
shutdown_generation: spec.generation.generation,
coordinator_term: spec.generation.coordinator_term,
coordinator_id: Some(spec.node.id()),
lease_until_unix_ms: spec.generation.lease_until_unix_ms,
load_mode: load_mode.clone(),
upstream: None,
downstream: downstream.clone(),
@ -1006,6 +1017,79 @@ async fn load_split_runtime_generation_inner(
})
}
async fn claim_split_coordinator_lease(
node: &mesh::Node,
model_ref: &str,
package: &skippy::SkippyPackageIdentity,
generation: &SplitTopologyGeneration,
) -> Result<()> {
let claim = split_coordinator_claim(node.id(), model_ref, package, generation);
let required_accepts = skippy_coordinator::quorum_requirement(generation.stages.len());
let mut accepted = 0usize;
let mut errors = Vec::new();
for stage in &generation.stages {
let request = skippy::StageControlRequest::Claim(claim.clone());
let response = if stage.node_id == node.id() {
node.send_local_stage_control(request).await
} else {
node.send_stage_control(stage.node_id, request).await
};
match response {
Ok(skippy::StageControlResponse::ClaimAccepted(ack)) if ack.accepted => {
accepted += 1;
}
Ok(skippy::StageControlResponse::ClaimAccepted(ack)) => {
errors.push(format!(
"{} rejected claim: {}",
stage.node_id.fmt_short(),
ack.error.unwrap_or_else(|| "unknown rejection".to_string())
));
}
Ok(other) => {
errors.push(format!(
"{} returned unexpected claim response: {other:?}",
stage.node_id.fmt_short()
));
}
Err(err) => {
errors.push(format!(
"{} claim failed: {err:#}",
stage.node_id.fmt_short()
));
}
}
}
anyhow::ensure!(
accepted >= required_accepts,
"coordinator claim for {model_ref} accepted by {accepted}/{} planned stage(s), need {required_accepts}: {}",
generation.stages.len(),
errors.join("; ")
);
Ok(())
}
fn split_coordinator_claim(
coordinator_id: iroh::EndpointId,
model_ref: &str,
package: &skippy::SkippyPackageIdentity,
generation: &SplitTopologyGeneration,
) -> skippy::StageCoordinatorClaim {
skippy::StageCoordinatorClaim {
model_id: model_ref.to_string(),
package_ref: package.package_ref.clone(),
manifest_sha256: package.manifest_sha256.clone(),
topology_id: generation.topology_id.clone(),
run_id: generation.run_id.clone(),
coordinator_id: coordinator_id.to_string(),
coordinator_term: generation.coordinator_term,
participant_set_hash: split_participant_set_hash(&generation.participants),
topology_hash: split_topology_hash(&generation.stages),
lease_until_unix_ms: generation.lease_until_unix_ms,
}
}
fn stage_load_model_path(load_mode: LoadMode, package_ref: &str, model_path: &Path) -> String {
match load_mode {
LoadMode::LayerPackage => package_ref.to_string(),
@ -1030,6 +1114,8 @@ struct SplitTopologyGeneration {
topology_id: String,
run_id: String,
generation: u64,
coordinator_term: u64,
lease_until_unix_ms: u64,
participants: Vec<SplitParticipant>,
stages: Vec<RuntimeSliceStagePlan>,
}
@ -1046,6 +1132,8 @@ impl SplitTopologyGeneration {
topology_id,
run_id,
generation,
coordinator_term: now_unix_nanos().max(1) as u64,
lease_until_unix_ms: split_coordinator_lease_until_unix_ms(),
participants,
stages,
}
@ -1618,6 +1706,7 @@ async fn stop_split_generation(
run_id: generation.run_id.clone(),
stage_id: stage.stage_id.clone(),
shutdown_generation,
coordinator_term: generation.coordinator_term,
};
let result = if stage.node_id == node.id() {
node.send_local_stage_control(skippy::StageControlRequest::Stop(stop))
@ -1946,6 +2035,33 @@ fn split_participant_signature(participants: &[SplitParticipant]) -> SplitPartic
.collect()
}
fn split_participant_set_hash(participants: &[SplitParticipant]) -> String {
let mut hasher = Sha256::new();
for participant in split_participant_signature(participants) {
hasher.update(participant.0.as_bytes());
hasher.update(participant.1.to_le_bytes());
hasher.update(participant.2.to_le_bytes());
hasher.update(participant.3.to_le_bytes());
hasher.update(participant.4.unwrap_or_default().to_le_bytes());
hasher.update([u8::from(participant.5)]);
hasher.update(participant.6.to_le_bytes());
}
format!("{:x}", hasher.finalize())
}
fn split_topology_hash(stages: &[RuntimeSliceStagePlan]) -> String {
let mut hasher = Sha256::new();
for stage in stages {
hasher.update(stage.stage_id.as_bytes());
hasher.update(stage.stage_index.to_le_bytes());
hasher.update(stage.node_id.to_string().as_bytes());
hasher.update(stage.layer_start.to_le_bytes());
hasher.update(stage.layer_end.to_le_bytes());
hasher.update(stage.parameter_bytes.to_le_bytes());
}
format!("{:x}", hasher.finalize())
}
fn split_participant_labels(participants: &[SplitParticipant]) -> Vec<String> {
participants
.iter()
@ -2782,6 +2898,9 @@ mod tests {
flash_attn_type: load.flash_attn_type,
error: None,
shutdown_generation: load.shutdown_generation,
coordinator_term: load.coordinator_term,
coordinator_id: load.coordinator_id,
lease_until_unix_ms: load.lease_until_unix_ms,
}
}
@ -2815,6 +2934,9 @@ mod tests {
flash_attn_type: FlashAttentionType::Auto,
error: None,
shutdown_generation: stop.shutdown_generation,
coordinator_term: stop.coordinator_term,
coordinator_id: None,
lease_until_unix_ms: 0,
}
}
@ -2838,6 +2960,9 @@ mod tests {
bind_addr: None,
error: None,
shutdown_generation: load.shutdown_generation,
coordinator_term: load.coordinator_term,
coordinator_id: load.coordinator_id,
lease_until_unix_ms: load.lease_until_unix_ms,
}
}
@ -3386,6 +3511,15 @@ mod tests {
test_inventory_from_request(inventory),
))
}
skippy::StageControlRequest::Claim(claim) => {
Ok(skippy::StageControlResponse::ClaimAccepted(
skippy::StageCoordinatorClaimAck {
accepted: true,
claim: claim.clone(),
error: None,
},
))
}
skippy::StageControlRequest::Load(load) if load.stage_id == "stage-1" => {
Err(anyhow::anyhow!("injected stage load failure"))
}
@ -3456,6 +3590,11 @@ mod tests {
);
let requests = requests.lock().unwrap();
let claim_count = requests
.iter()
.filter(|request| matches!(request, skippy::StageControlRequest::Claim(_)))
.count();
assert_eq!(claim_count, generation.stages.len());
let load_stage_ids = requests
.iter()
.filter_map(|request| match request {

170
crates/mesh-llm-host-runtime/src/runtime/mod.rs
View file

@ -58,6 +58,7 @@ use zeroize::Zeroizing;
const PRETTY_DASHBOARD_INVENTORY_CACHE_TTL: Duration = Duration::from_secs(5);
const DASHBOARD_CONTEXT_USAGE_REFRESH_INTERVAL: Duration = Duration::from_millis(250);
const DASHBOARD_FIRST_PAINT_TIMEOUT: Duration = Duration::from_secs(2);
const SPLIT_STANDBY_RETRY_INTERVAL: Duration = Duration::from_secs(30);
type DashboardContextUsage =
Arc<tokio::sync::Mutex<HashMap<String, HashMap<DashboardContextUsageSource, u64>>>>;
@ -1090,21 +1091,6 @@ async fn startup_local_model_loop(params: StartupLocalModelTask) {
None
};
let startup_load_guard = startup_load_gate.lock().await;
let start_spec = LocalRuntimeModelStartSpec {
node: &node,
model_path: &model_path,
mmproj_override: mmproj_path.as_deref(),
ctx_size_override: ctx_size,
pinned_gpu: pinned_gpu.as_ref(),
cache_type_k_override: cache_type_k.as_deref(),
cache_type_v_override: cache_type_v.as_deref(),
n_batch_override: n_batch,
n_ubatch_override: n_ubatch,
flash_attention_override: flash_attention,
slots,
parallel_override,
};
let local_capacity = pinned_gpu
.as_ref()
.map(|gpu| gpu.vram_bytes)
@ -1140,7 +1126,21 @@ async fn startup_local_model_loop(params: StartupLocalModelTask) {
reason: SplitRuntimeReason::LocalCapacity,
} => survey::SurveyLaunchKind::MoeFallback,
};
let launch_started = Instant::now();
let make_start_spec = || LocalRuntimeModelStartSpec {
node: &node,
model_path: &model_path,
mmproj_override: mmproj_path.as_deref(),
ctx_size_override: ctx_size,
pinned_gpu: pinned_gpu.as_ref(),
cache_type_k_override: cache_type_k.as_deref(),
cache_type_v_override: cache_type_v.as_deref(),
n_batch_override: n_batch,
n_ubatch_override: n_ubatch,
flash_attention_override: flash_attention,
slots,
parallel_override,
};
let mut launch_started: Instant;
let (
mut loaded_name,
handle,
@ -1164,60 +1164,99 @@ async fn startup_local_model_loop(params: StartupLocalModelTask) {
)),
});
}
match start_runtime_split_model(start_spec, &model_ref).await {
Ok(SplitRuntimeStart::Started(loaded)) => {
let mut loaded = *loaded;
(
loaded.loaded_name,
loaded.handle,
loaded.death_rx,
loaded.cleanup.take(),
loaded.coordinator_rx.take(),
loaded.coordinator_task.take(),
)
}
Ok(SplitRuntimeStart::Standby { coordinator }) => {
drop(startup_load_guard);
let _ = emit_event(OutputEvent::Info {
message: format!(
"Split runtime coordinator is {}; standing by for stage assignment",
coordinator.fmt_short()
),
context: Some(format!("model={model_ref}")),
});
update_startup_target(&target_tx, &model_name, election::InferenceTarget::None);
if let Some(cs) = console_state {
cs.update(false, false).await;
let mut peer_rx = node.peer_change_rx.clone();
loop {
let startup_load_guard = startup_load_gate.lock().await;
launch_started = Instant::now();
match start_runtime_split_model(make_start_spec(), &model_ref).await {
Ok(SplitRuntimeStart::Started(loaded)) => {
drop(startup_load_guard);
let mut loaded = *loaded;
break (
loaded.loaded_name,
loaded.handle,
loaded.death_rx,
loaded.cleanup.take(),
loaded.coordinator_rx.take(),
loaded.coordinator_task.take(),
);
}
return;
}
Err(err) => {
survey_telemetry.record_launch_failure(
survey::SurveyModelSpec {
model: &model_name,
model_path: Some(&model_path),
launch_kind,
pinned_gpu: pinned_gpu.as_ref(),
backend: None,
context_length: ctx_size.map(u64::from),
},
launch_started.elapsed(),
survey::classify_launch_failure(&err),
);
let _ = emit_event(OutputEvent::Error {
message: format!("Failed to start model {model_name}: {err:#}"),
context: Some(format!("model={model_name}")),
});
update_startup_target(&target_tx, &model_name, election::InferenceTarget::None);
if let Some(cs) = console_state {
cs.update(false, false).await;
Ok(SplitRuntimeStart::Standby { coordinator }) => {
drop(startup_load_guard);
let _ = emit_event(OutputEvent::Info {
message: format!(
"Split runtime coordinator is {}; standing by for stage assignment",
coordinator.fmt_short()
),
context: Some(format!("model={model_ref}")),
});
update_startup_target(
&target_tx,
&model_name,
election::InferenceTarget::None,
);
if let Some(cs) = console_state.as_ref() {
cs.update(false, false).await;
}
}
Err(err) => {
survey_telemetry.record_launch_failure(
survey::SurveyModelSpec {
model: &model_name,
model_path: Some(&model_path),
launch_kind,
pinned_gpu: pinned_gpu.as_ref(),
backend: None,
context_length: ctx_size.map(u64::from),
},
launch_started.elapsed(),
survey::classify_launch_failure(&err),
);
let _ = emit_event(OutputEvent::Error {
message: format!("Failed to start model {model_name}: {err:#}"),
context: Some(format!("model={model_name}")),
});
update_startup_target(
&target_tx,
&model_name,
election::InferenceTarget::None,
);
if let Some(cs) = console_state.as_ref() {
cs.update(false, false).await;
}
return;
}
}
tokio::select! {
result = peer_rx.changed() => {
if result.is_err() {
return;
}
tokio::select! {
_ = tokio::time::sleep(Duration::from_secs(2)) => {}
result = stop_rx.changed() => {
if result.is_err() || *stop_rx.borrow() {
return;
}
}
}
}
_ = tokio::time::sleep(SPLIT_STANDBY_RETRY_INTERVAL) => {}
result = stop_rx.changed() => {
if result.is_err() || *stop_rx.borrow() {
return;
}
}
return;
}
}
}
StartupRuntimePlan::Local => {
match start_runtime_local_model(start_spec, &model_ref).await {
let startup_load_guard = startup_load_gate.lock().await;
launch_started = Instant::now();
let start_result = start_runtime_local_model(make_start_spec(), &model_ref).await;
drop(startup_load_guard);
match start_result {
Ok((loaded_name, handle, death_rx)) => {
(loaded_name, handle, death_rx, None, None, None)
}
@ -1239,7 +1278,7 @@ async fn startup_local_model_loop(params: StartupLocalModelTask) {
context: Some(format!("model={model_name}")),
});
update_startup_target(&target_tx, &model_name, election::InferenceTarget::None);
if let Some(cs) = console_state {
if let Some(cs) = console_state.as_ref() {
cs.update(false, false).await;
}
return;
@ -1247,7 +1286,6 @@ async fn startup_local_model_loop(params: StartupLocalModelTask) {
}
}
};
drop(startup_load_guard);
let mut survey_loaded_model = survey_telemetry.model(survey::SurveyModelSpec {
model: &loaded_name,

9
crates/skippy-coordinator/Cargo.toml Normal file
View file

@ -0,0 +1,9 @@
[package]
name = "skippy-coordinator"
edition.workspace = true
license.workspace = true
readme = "README.md"
version.workspace = true
[dependencies]
thiserror = "2"

348
crates/skippy-coordinator/README.md Normal file
View file

@ -0,0 +1,348 @@
# Skippy Coordinator
`skippy-coordinator` owns the pure coordinator lease and fencing rules for
Skippy split topologies. It deliberately does not know about QUIC, gossip,
protobuf, Tokio, iroh endpoint IDs, or process management. Those parts live in
`mesh-llm-host-runtime`; this crate only answers:
- Is this coordinator claim valid?
- Does this newer claim supersede an older coordinator?
- Is this load fenced by the current accepted claim?
- How many planned stages must accept before a split can start?
## Mental Model
A split runtime has one coordinator. The coordinator chooses a topology, asks
the planned stage runtimes to accept a lease, and only loads the split if a
majority of those planned stages accept.
Every accepted claim is keyed by:
- `model_id`
- `package_ref`
- `manifest_sha256`
That means a stage remembers the current coordinator claim for a specific model
artifact. A newer term for the same model/package/manifest can replace an older
one; stale terms cannot.
```mermaid
flowchart LR
C["Coordinator node"]
S1["Stage A"]
S2["Stage B"]
S3["Stage C"]
F1["Claim fence"]
F2["Claim fence"]
F3["Claim fence"]
C -->|"Claim term T, topology hash, participant hash, lease"| S1
C -->|"Claim term T, topology hash, participant hash, lease"| S2
C -->|"Claim term T, topology hash, participant hash, lease"| S3
S1 --> F1
S2 --> F2
S3 --> F3
F1 -->|"accept/reject"| C
F2 -->|"accept/reject"| C
F3 -->|"accept/reject"| C
```
## Terms, Leases, And Hashes
Each claim carries:
- `coordinator_id`: the node that owns the split generation.
- `coordinator_term`: a monotonic-ish term chosen by the runtime for the split generation.
- `topology_id` and `run_id`: the concrete split generation being claimed.
- `participant_set_hash`: hash of the planned participants and capacity inputs.
- `topology_hash`: hash of the planned stages, owners, and layer ranges.
- `lease_until_unix_ms`: wall-clock lease expiry.
The current host runtime uses a 4 hour split coordinator lease. The lease is a
fence, not a heartbeat protocol. If a load arrives after the lease expires, the
stage rejects it.
```mermaid
flowchart TD
Claim["CoordinatorClaim"]
Shape{"Shape valid?"}
Existing{"Existing claim?"}
Stale{"term < current term?"}
SameTerm{"same term?"}
SameEpoch{"same topology/run/coordinator/hashes?"}
Accept["Accept claim"]
Supersede["Accept and supersede older term"]
Reject["Reject claim"]
Claim --> Shape
Shape -- "missing fields, term 0, expired lease" --> Reject
Shape -- "ok" --> Existing
Existing -- "no" --> Accept
Existing -- "yes" --> Stale
Stale -- "yes" --> Reject
Stale -- "no" --> SameTerm
SameTerm -- "yes" --> SameEpoch
SameEpoch -- "yes" --> Accept
SameEpoch -- "no" --> Reject
SameTerm -- "no, newer term" --> Supersede
```
## Startup Flow
The coordinator must win a quorum before it loads stages. Quorum is a majority
of the planned stage count:
```text
quorum = planned_stage_count / 2 + 1
```
For a 3 stage split, 2 accepts are enough. For a 4 stage split, 3 accepts are
required.
```mermaid
sequenceDiagram
participant C as Coordinator
participant A as Stage A
participant B as Stage B
participant D as Stage C
C->>A: ClaimCoordinator(term T)
C->>B: ClaimCoordinator(term T)
C->>D: ClaimCoordinator(term T)
A-->>C: accepted
B-->>C: accepted
D-->>C: failed or unavailable
Note over C: 2/3 accepted, quorum reached
C->>A: LoadStage(term T, coordinator_id)
C->>B: LoadStage(term T, coordinator_id)
C->>D: LoadStage(term T, coordinator_id)
```
If quorum is not reached, the coordinator must not load the split.
## Load Fencing
Stages validate fenced loads against the accepted claim. A `LoadStage` is
accepted only when:
- it has a non-zero `coordinator_term`
- it has a `coordinator_id`
- the stage has an accepted claim for the same model/package/manifest
- the load term equals the claim term
- the load coordinator equals the claim coordinator
- the load topology/run equals the claim topology/run
- the accepted lease has not expired
Loads with `coordinator_term = 0` and no coordinator ID are treated as
unfenced legacy/local loads and bypass this coordinator policy.
```mermaid
flowchart TD
Load["LoadStage"]
Unfenced{"term 0 and no coordinator_id?"}
ClaimExists{"matching claim exists?"}
Match{"term, coordinator, topology, run match?"}
Lease{"lease still valid?"}
Start["Start stage runtime"]
Reject["Reject load"]
Load --> Unfenced
Unfenced -- "yes" --> Start
Unfenced -- "no" --> ClaimExists
ClaimExists -- "no" --> Reject
ClaimExists -- "yes" --> Match
Match -- "no" --> Reject
Match -- "yes" --> Lease
Lease -- "expired" --> Reject
Lease -- "valid" --> Start
```
## Superseding And Fencing Old Work
When a stage accepts a newer term for the same model/package/manifest, the host
runtime fences stale work:
- running stages with older coordinator terms are shut down
- preparations with older coordinator terms are cancelled
- stale stop requests cannot stop a newer running stage
```mermaid
sequenceDiagram
participant C1 as Old coordinator
participant S as Stage
participant C2 as New coordinator
C1->>S: ClaimCoordinator(term 10)
S-->>C1: accepted
C1->>S: LoadStage(term 10)
S-->>C1: ready
C2->>S: ClaimCoordinator(term 11)
Note over S: term 11 supersedes term 10
S->>S: shut down older running stage
S-->>C2: accepted
C2->>S: LoadStage(term 11)
S-->>C2: ready
C1->>S: StopStage(term 10)
S-->>C1: rejected as stale
```
## Scenarios
### 1. Fresh Split Startup
The coordinator plans a topology, sends `ClaimCoordinator` to every planned
stage, reaches majority, then loads stages with the same term and coordinator
ID.
Expected result: split starts.
### 2. One Planned Stage Is Down During Startup
For a 3 stage topology, if two stages accept and one is unreachable, quorum is
still reached.
Expected result: the coordinator may continue. The later load can still fail if
the missing stage is required for the concrete topology.
### 3. Too Many Planned Stages Are Down
For a 3 stage topology, if only one stage accepts, quorum is not reached.
Expected result: the split is not loaded. The node remains in standby/retry
behavior and can try again after peer changes or retry ticks.
```mermaid
flowchart LR
C["Coordinator"]
A["Stage A accepts"]
B["Stage B down"]
D["Stage C down"]
Q{"1/3 accepts >= 2?"}
Stop["Do not load split"]
C --> A
C -. "unreachable" .-> B
C -. "unreachable" .-> D
A --> Q
Q -- "no" --> Stop
```
### 4. Coordinator Disappears After Split Is Running
The stages do not elect a new coordinator by themselves. Other mesh nodes
observe peer/status changes through the runtime layer, plan a replacement split,
and attempt a newer coordinator claim.
Expected result: a reachable node can become coordinator by claiming a newer
term from a majority of the planned replacement stages.
### 5. Old Coordinator Comes Back
If the old coordinator tries to load or stop using an older term, stages reject
the stale request once they have accepted a newer claim.
Expected result: the old coordinator cannot corrupt the newer split.
```mermaid
flowchart TD
Old["Old coordinator term 10 rejoins"]
Stage["Stage has accepted term 11"]
Load["Old LoadStage or StopStage term 10"]
Reject["Reject stale request"]
Old --> Load
Load --> Stage
Stage --> Reject
```
### 6. Two Coordinators Race With The Same Term
The first valid claim for a term is accepted. A second claim with the same term
but different topology, coordinator, participant hash, or topology hash is
rejected as a conflicting same-term claim.
Expected result: same-term split brain is fenced at the stage.
### 7. Two Coordinators Race With Different Terms
The higher term wins at each stage. Accepting the higher term supersedes the
older claim and fences stale runtime work.
Expected result: convergence on the newer term for stages that receive it.
The coordinator still needs majority acceptance before it can load a split.
### 8. Network Partition
A partition means the mesh has split into groups that cannot talk to each
other. The coordinator protocol does not use gossip as consensus. Gossip helps
the runtime notice peer changes, but coordinator ownership is decided by direct
stage-control claims.
Expected result:
- the side that can claim a majority of the planned stages may load
- the side that cannot claim majority stays out
- when the partition heals, stale terms are rejected by stages that accepted a
newer term
```mermaid
flowchart LR
subgraph P1["Partition A"]
C1["Coordinator A"]
S1["Stage 1"]
S2["Stage 2"]
end
subgraph P2["Partition B"]
C2["Coordinator B"]
S3["Stage 3"]
end
C1 --> S1
C1 --> S2
C1 -. "cannot reach" .-> S3
C2 --> S3
C2 -. "cannot reach" .-> S1
C2 -. "cannot reach" .-> S2
S1 --> Q1{"A has 2/3"}
S2 --> Q1
S3 --> Q2{"B has 1/3"}
Q1 -->|"quorum"| LoadA["A may load"]
Q2 -->|"no quorum"| NoLoadB["B must not load"]
```
### 9. Lease Expiry
An expired claim cannot authorize new loads. A coordinator must claim again
with a valid lease before loading fenced stages.
Expected result: old or delayed load messages do not start stages after the
lease window.
### 10. Local Or Legacy Unfenced Load
Loads with term `0` and no coordinator ID bypass the split coordinator fence.
Expected result: non-split stage usage and tests that do not participate in
split coordination still work.
## What This Is Not
This is not Raft. There is no replicated log, no committed command sequence,
and no long-lived cluster membership stored in this crate. The runtime only
needs a fencing token for split ownership:
- majority claim before load
- monotonic term replacement
- stale request rejection
- lease expiry
That is enough to prevent old coordinators from continuing to mutate a split
after a newer coordinator has taken ownership, without turning every stage into
a consensus node.

389
crates/skippy-coordinator/src/lib.rs Normal file
View file

@ -0,0 +1,389 @@
//! Coordinator lease and fencing policy for skippy split topologies.
//!
//! This crate intentionally owns only the pure coordination rules. Mesh
//! transport, protobuf conversion, node identity types, and stage runtime
//! process management stay in the host runtime.
use std::collections::HashMap;
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct CoordinatorClaim {
pub model_id: String,
pub package_ref: String,
pub manifest_sha256: String,
pub topology_id: String,
pub run_id: String,
pub coordinator_id: String,
pub coordinator_term: u64,
pub participant_set_hash: String,
pub topology_hash: String,
pub lease_until_unix_ms: u64,
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct LoadClaimRef {
pub model_id: String,
pub package_ref: String,
pub manifest_sha256: String,
pub topology_id: String,
pub run_id: String,
pub coordinator_id: Option<String>,
pub coordinator_term: u64,
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum ClaimDecision {
Accepted {
supersedes_term: Option<u64>,
claim: CoordinatorClaim,
},
Rejected {
current: Option<CoordinatorClaim>,
reason: ClaimRejection,
},
}
#[derive(Clone, Debug, Eq, PartialEq, thiserror::Error)]
pub enum ClaimRejection {
#[error("coordinator claim requires model_id")]
MissingModelId,
#[error("coordinator claim requires package_ref")]
MissingPackageRef,
#[error("coordinator claim requires manifest_sha256")]
MissingManifestSha256,
#[error("coordinator claim requires topology_id and run_id")]
MissingTopologyRun,
#[error("coordinator claim requires coordinator_id")]
MissingCoordinatorId,
#[error("coordinator claim requires non-zero term")]
MissingTerm,
#[error("coordinator claim requires participant and topology hashes")]
MissingHashes,
#[error("coordinator claim lease is expired")]
ExpiredLease,
#[error("stale coordinator term {claim_term} < {current_term}")]
StaleTerm { claim_term: u64, current_term: u64 },
#[error("conflicting coordinator claim for existing term")]
ConflictingSameTerm,
}
#[derive(Clone, Debug, Eq, PartialEq, thiserror::Error)]
pub enum LoadRejection {
#[error("missing coordinator term")]
MissingTerm,
#[error("missing coordinator id")]
MissingCoordinatorId,
#[error("missing coordinator claim")]
MissingClaim,
#[error("coordinator term mismatch: load={load_term} claim={claim_term}")]
TermMismatch { load_term: u64, claim_term: u64 },
#[error("coordinator id mismatch")]
CoordinatorMismatch,
#[error("coordinator claim does not match topology/run")]
TopologyRunMismatch,
#[error("coordinator lease expired")]
ExpiredLease,
}
#[derive(Clone, Debug, Default)]
pub struct ClaimFence {
claims: HashMap<ClaimKey, CoordinatorClaim>,
}
impl ClaimFence {
pub fn accept_claim(&mut self, claim: CoordinatorClaim, now_unix_ms: u64) -> ClaimDecision {
if let Some(reason) = validate_claim_shape(&claim, now_unix_ms) {
return ClaimDecision::Rejected {
current: None,
reason,
};
}
let key = ClaimKey::from_claim(&claim);
if let Some(current) = self.claims.get(&key) {
if claim.coordinator_term < current.coordinator_term {
return ClaimDecision::Rejected {
current: Some(current.clone()),
reason: ClaimRejection::StaleTerm {
claim_term: claim.coordinator_term,
current_term: current.coordinator_term,
},
};
}
if claim.coordinator_term == current.coordinator_term
&& !same_claim_epoch(&claim, current)
{
return ClaimDecision::Rejected {
current: Some(current.clone()),
reason: ClaimRejection::ConflictingSameTerm,
};
}
}
let previous = self.claims.insert(key, claim.clone());
ClaimDecision::Accepted {
supersedes_term: previous.as_ref().and_then(|current| {
(claim.coordinator_term > current.coordinator_term)
.then_some(current.coordinator_term)
}),
claim,
}
}
pub fn validate_load(
&self,
load: &LoadClaimRef,
now_unix_ms: u64,
) -> Result<(), LoadRejection> {
if load.coordinator_term == 0 {
return Err(LoadRejection::MissingTerm);
}
let Some(coordinator_id) = load.coordinator_id.as_deref() else {
return Err(LoadRejection::MissingCoordinatorId);
};
let key = ClaimKey::from_load(load);
let Some(claim) = self.claims.get(&key) else {
return Err(LoadRejection::MissingClaim);
};
if claim.coordinator_term != load.coordinator_term {
return Err(LoadRejection::TermMismatch {
load_term: load.coordinator_term,
claim_term: claim.coordinator_term,
});
}
if claim.coordinator_id != coordinator_id {
return Err(LoadRejection::CoordinatorMismatch);
}
if claim.topology_id != load.topology_id || claim.run_id != load.run_id {
return Err(LoadRejection::TopologyRunMismatch);
}
if claim.lease_until_unix_ms < now_unix_ms {
return Err(LoadRejection::ExpiredLease);
}
Ok(())
}
pub fn current_claim_for(
&self,
model_id: &str,
package_ref: &str,
manifest_sha256: &str,
) -> Option<&CoordinatorClaim> {
self.claims.get(&ClaimKey {
model_id: model_id.to_string(),
package_ref: package_ref.to_string(),
manifest_sha256: manifest_sha256.to_string(),
})
}
}
pub fn quorum_requirement(planned_stage_count: usize) -> usize {
planned_stage_count / 2 + 1
}
pub fn same_claim_epoch(left: &CoordinatorClaim, right: &CoordinatorClaim) -> bool {
left.model_id == right.model_id
&& left.package_ref == right.package_ref
&& left.manifest_sha256 == right.manifest_sha256
&& left.topology_id == right.topology_id
&& left.run_id == right.run_id
&& left.coordinator_id == right.coordinator_id
&& left.participant_set_hash == right.participant_set_hash
&& left.topology_hash == right.topology_hash
}
fn validate_claim_shape(claim: &CoordinatorClaim, now_unix_ms: u64) -> Option<ClaimRejection> {
if claim.model_id.is_empty() {
return Some(ClaimRejection::MissingModelId);
}
if claim.package_ref.is_empty() {
return Some(ClaimRejection::MissingPackageRef);
}
if claim.manifest_sha256.is_empty() {
return Some(ClaimRejection::MissingManifestSha256);
}
if claim.topology_id.is_empty() || claim.run_id.is_empty() {
return Some(ClaimRejection::MissingTopologyRun);
}
if claim.coordinator_id.is_empty() {
return Some(ClaimRejection::MissingCoordinatorId);
}
if claim.coordinator_term == 0 {
return Some(ClaimRejection::MissingTerm);
}
if claim.participant_set_hash.is_empty() || claim.topology_hash.is_empty() {
return Some(ClaimRejection::MissingHashes);
}
if claim.lease_until_unix_ms <= now_unix_ms {
return Some(ClaimRejection::ExpiredLease);
}
None
}
#[derive(Clone, Debug, Eq, Hash, PartialEq)]
struct ClaimKey {
model_id: String,
package_ref: String,
manifest_sha256: String,
}
impl ClaimKey {
fn from_claim(claim: &CoordinatorClaim) -> Self {
Self {
model_id: claim.model_id.clone(),
package_ref: claim.package_ref.clone(),
manifest_sha256: claim.manifest_sha256.clone(),
}
}
fn from_load(load: &LoadClaimRef) -> Self {
Self {
model_id: load.model_id.clone(),
package_ref: load.package_ref.clone(),
manifest_sha256: load.manifest_sha256.clone(),
}
}
}
#[cfg(test)]
mod tests {
use super::*;
fn claim(term: u64) -> CoordinatorClaim {
CoordinatorClaim {
model_id: "model".to_string(),
package_ref: "hf://pkg".to_string(),
manifest_sha256: "manifest".to_string(),
topology_id: format!("topology-{term}"),
run_id: format!("run-{term}"),
coordinator_id: "node-a".to_string(),
coordinator_term: term,
participant_set_hash: "participants".to_string(),
topology_hash: format!("topology-hash-{term}"),
lease_until_unix_ms: 10_000,
}
}
fn load(term: u64) -> LoadClaimRef {
LoadClaimRef {
model_id: "model".to_string(),
package_ref: "hf://pkg".to_string(),
manifest_sha256: "manifest".to_string(),
topology_id: format!("topology-{term}"),
run_id: format!("run-{term}"),
coordinator_id: Some("node-a".to_string()),
coordinator_term: term,
}
}
#[test]
fn accepts_first_valid_claim() {
let mut fence = ClaimFence::default();
assert!(matches!(
fence.accept_claim(claim(1), 1_000),
ClaimDecision::Accepted {
supersedes_term: None,
..
}
));
}
#[test]
fn rejects_stale_claim_after_newer_term() {
let mut fence = ClaimFence::default();
fence.accept_claim(claim(2), 1_000);
assert!(matches!(
fence.accept_claim(claim(1), 1_000),
ClaimDecision::Rejected {
reason: ClaimRejection::StaleTerm {
claim_term: 1,
current_term: 2
},
..
}
));
}
#[test]
fn rejects_conflicting_claim_for_same_term() {
let mut fence = ClaimFence::default();
fence.accept_claim(claim(1), 1_000);
let mut conflicting = claim(1);
conflicting.coordinator_id = "node-b".to_string();
assert!(matches!(
fence.accept_claim(conflicting, 1_000),
ClaimDecision::Rejected {
reason: ClaimRejection::ConflictingSameTerm,
..
}
));
}
#[test]
fn newer_claim_supersedes_old_term() {
let mut fence = ClaimFence::default();
fence.accept_claim(claim(1), 1_000);
assert!(matches!(
fence.accept_claim(claim(2), 1_000),
ClaimDecision::Accepted {
supersedes_term: Some(1),
..
}
));
}
#[test]
fn validates_load_against_current_claim() {
let mut fence = ClaimFence::default();
fence.accept_claim(claim(3), 1_000);
assert_eq!(fence.validate_load(&load(3), 1_000), Ok(()));
}
#[test]
fn rejects_load_without_matching_claim() {
let fence = ClaimFence::default();
assert_eq!(
fence.validate_load(&load(1), 1_000),
Err(LoadRejection::MissingClaim)
);
}
#[test]
fn rejects_load_for_stale_term() {
let mut fence = ClaimFence::default();
fence.accept_claim(claim(2), 1_000);
assert_eq!(
fence.validate_load(&load(1), 1_000),
Err(LoadRejection::TermMismatch {
load_term: 1,
claim_term: 2,
})
);
}
#[test]
fn rejects_expired_claim_and_expired_load() {
let mut fence = ClaimFence::default();
assert!(matches!(
fence.accept_claim(claim(1), 10_001),
ClaimDecision::Rejected {
reason: ClaimRejection::ExpiredLease,
..
}
));
fence.accept_claim(claim(2), 1_000);
assert_eq!(
fence.validate_load(&load(2), 10_001),
Err(LoadRejection::ExpiredLease)
);
}
#[test]
fn quorum_is_majority_of_planned_stages() {
assert_eq!(quorum_requirement(1), 1);
assert_eq!(quorum_requirement(2), 2);
assert_eq!(quorum_requirement(3), 2);
assert_eq!(quorum_requirement(4), 3);
}
}

31
crates/skippy-protocol/proto/stage.proto
View file

@ -18,6 +18,7 @@ message StageControlRequest {
PrepareStage prepare_stage = 7;
CancelPrepareStage cancel_prepare_stage = 8;
StageStatusUpdate stage_status_update = 9;
ClaimCoordinator claim_coordinator = 10;
}
}
@ -31,6 +32,7 @@ message StageControlResponse {
StagePreparationStatus stage_preparation_status = 6;
StageStatusAck stage_status_ack = 7;
StageStatusList stage_statuses = 8;
CoordinatorClaimAccepted coordinator_claim_accepted = 9;
}
}
@ -86,6 +88,9 @@ message LoadStage {
optional uint32 n_ubatch = 27;
StageFlashAttnType flash_attn_type = 28;
optional uint64 source_model_bytes = 29;
uint64 coordinator_term = 30;
optional string coordinator_id = 31;
uint64 lease_until_unix_ms = 32;
}
message StopStage {
@ -93,6 +98,7 @@ message StopStage {
string run_id = 2;
string stage_id = 3;
uint64 shutdown_generation = 4;
uint64 coordinator_term = 5;
}
message GetStageStatus {
@ -123,6 +129,19 @@ message StageStatusUpdate {
StagePreparationStatus status = 1;
}
message ClaimCoordinator {
string model_id = 1;
string package_ref = 2;
string manifest_sha256 = 3;
string topology_id = 4;
string run_id = 5;
string coordinator_id = 6;
uint64 coordinator_term = 7;
string participant_set_hash = 8;
string topology_hash = 9;
uint64 lease_until_unix_ms = 10;
}
message StageReady {
bool accepted = 1;
StageStatus status = 2;
@ -140,6 +159,12 @@ message StageStatusAck {
optional string error = 2;
}
message CoordinatorClaimAccepted {
bool accepted = 1;
ClaimCoordinator claim = 2;
optional string error = 3;
}
message StageStatus {
string topology_id = 1;
string run_id = 2;
@ -169,6 +194,9 @@ message StageStatus {
optional uint32 n_batch = 26;
optional uint32 n_ubatch = 27;
StageFlashAttnType flash_attn_type = 28;
uint64 coordinator_term = 29;
optional string coordinator_id = 30;
uint64 lease_until_unix_ms = 31;
}
message StageStatusList {
@ -211,6 +239,9 @@ message StagePreparationStatus {
optional string bind_addr = 14;
optional string error = 15;
uint64 shutdown_generation = 16;
uint64 coordinator_term = 17;
optional string coordinator_id = 18;
uint64 lease_until_unix_ms = 19;
}
enum SourceModelKind {

1
crates/skippy-protocol/src/lib.rs
View file

@ -695,6 +695,7 @@ mod tests {
run_id: "run-a".to_string(),
stage_id: "stage-0".to_string(),
shutdown_generation: 7,
coordinator_term: 7,
})),
..frame.clone()
};