diff --git a/Cargo.lock b/Cargo.lock index 7b3b751..8d3ac29 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2542,6 +2542,7 @@ dependencies = [ "tokio", "toml 0.8.23", "ureq", + "wiremock", ] [[package]] diff --git a/apps/myelin/Cargo.toml b/apps/myelin/Cargo.toml index bb1c0b3..4a2a98c 100644 --- a/apps/myelin/Cargo.toml +++ b/apps/myelin/Cargo.toml @@ -31,6 +31,9 @@ blake3 = "1" toml = "0.8" ureq = "2" +[dev-dependencies] +wiremock = "0.6" + [target.'cfg(target_os = "linux")'.dependencies] libc = "0.2" signal-hook = "0.3" diff --git a/apps/myelin/src/chat/node_image.rs b/apps/myelin/src/chat/node_image.rs index e34c3ff..5ef1b88 100644 --- a/apps/myelin/src/chat/node_image.rs +++ b/apps/myelin/src/chat/node_image.rs @@ -106,14 +106,7 @@ fn prepare_node_image_inner( run_status_command( &root, "cargo", - &[ - "build", - "--quiet", - "-p", - "myelin", - "--bin", - "myelin-worker", - ], + &["build", "--quiet", "-p", "myelin", "--bin", "myelin-worker"], "build myelin-worker", None, progress, @@ -365,7 +358,12 @@ fn collect_hash_inputs(root: &Path, path: &Path, out: &mut Vec) -> Resu } return Ok(()); } - if !metadata.is_dir() || matches!(path.file_name().and_then(|name| name.to_str()), Some(".git" | "target" | "__pycache__")) { + if !metadata.is_dir() + || matches!( + path.file_name().and_then(|name| name.to_str()), + Some(".git" | "target" | "__pycache__") + ) + { return Ok(()); } let entries = fs::read_dir(path).map_err(|e| format!("read dir {display}: {e}"))?; @@ -703,7 +701,7 @@ fn run_status_command( ); let mut child = match Command::new(program) .current_dir(root) - .args(&args) + .args(&args) .stdin(Stdio::null()) .stdout(Stdio::piped()) .stderr(Stdio::piped()) diff --git a/apps/myelin/src/chat/runtime.rs b/apps/myelin/src/chat/runtime.rs index d278421..6139979 100644 --- a/apps/myelin/src/chat/runtime.rs +++ b/apps/myelin/src/chat/runtime.rs @@ -889,7 +889,10 @@ impl Config { /// cached GGUF model without explicit flags. Explicit `--cached-model` is /// always honored (and stays strict); the default degrades gracefully to /// the normal download path when no cached model is present. - fn cached_model_source(args: &ParsedArgs, provider: &ProviderKind) -> Option { + fn cached_model_source( + args: &ParsedArgs, + provider: &ProviderKind, + ) -> Option { match &args.cached_model { Some(source) => Some(source.clone()), None if provider == &provider_kind::process() => Some(CachedModelSource::Discover), @@ -1388,8 +1391,8 @@ fn wait_for_rpc_ready( } impl InProcessOrch { -// spawns the orchestrator process; process control, out of scope (ENGINE_SPEC.md §2) -#[allow(clippy::disallowed_methods)] + // spawns the orchestrator process; process control, out of scope (ENGINE_SPEC.md §2) + #[allow(clippy::disallowed_methods)] fn spawn(config: &Config, image_ref: &str) -> Result { let args = config.orchestrator_cli_args(image_ref); let (stop_tx, stop_rx) = mpsc::channel(); @@ -1418,8 +1421,8 @@ impl InProcessOrch { }) } -// synchronous process-control readiness sequencing; the engine drives all background work (ENGINE_SPEC.md §2) -#[allow(clippy::disallowed_methods)] + // synchronous process-control readiness sequencing; the engine drives all background work (ENGINE_SPEC.md §2) + #[allow(clippy::disallowed_methods)] fn shutdown(&mut self) { if self.cleaned { return; @@ -1517,8 +1520,8 @@ impl OrchChild { // The orchestrator shutdown spec is still pending. Replace this with the approved // shutdown contract when it is finalized; do not add private stdin commands here. -// synchronous process-control readiness sequencing; the engine drives all background work (ENGINE_SPEC.md §2) -#[allow(clippy::disallowed_methods)] + // synchronous process-control readiness sequencing; the engine drives all background work (ENGINE_SPEC.md §2) + #[allow(clippy::disallowed_methods)] fn shutdown(&mut self) { if self.cleaned { return; @@ -1688,14 +1691,7 @@ where config.skip_rebuild, &config.worker_bin, "myelin-worker", - &[ - "build", - "--quiet", - "-p", - "myelin", - "--bin", - "myelin-worker", - ], + &["build", "--quiet", "-p", "myelin", "--bin", "myelin-worker"], )?; emit_chat_progress( &mut progress, @@ -1733,14 +1729,7 @@ where config.skip_rebuild, &config.worker_bin, "myelin-worker", - &[ - "build", - "--quiet", - "-p", - "myelin", - "--bin", - "myelin-worker", - ], + &["build", "--quiet", "-p", "myelin", "--bin", "myelin-worker"], )?; emit_chat_progress( &mut progress, @@ -1895,7 +1884,15 @@ fn run_chat_session_with_output_and_progress( let mut progress = progress; let mut next_request_id = 1_u64; let mut next_prompt_index = 1_u64; - let prompt_exited = |progress: &mut Option<&mut ChatDatastream>, reason: &str| emit_chat_progress(progress, CHAT_PROMPT_CHANNEL, "prompt_loop", "exited", json!({"reason": reason})); + let prompt_exited = |progress: &mut Option<&mut ChatDatastream>, reason: &str| { + emit_chat_progress( + progress, + CHAT_PROMPT_CHANNEL, + "prompt_loop", + "exited", + json!({"reason": reason}), + ) + }; loop { if STOP_REQUESTED.load(Ordering::SeqCst) { diff --git a/apps/myelin/src/node/worker_node_runtime.rs b/apps/myelin/src/node/worker_node_runtime.rs index 09f39c0..db426c5 100644 --- a/apps/myelin/src/node/worker_node_runtime.rs +++ b/apps/myelin/src/node/worker_node_runtime.rs @@ -50,8 +50,8 @@ use iroh_driver::{EndpointAddrMask, MVP_IROH_ENDPOINT_ADDR_MASK_ENV, advertised_ use parking_lot::Mutex; use serde_json::{Value, json}; use swactor::actor::{ActorAddress, ActorInterface}; -use swactor_engine::{Engine, EngineHandle, TokioBackend, TokioConfig}; use swactor::runtime::{Ctx, ExternalSender}; +use swactor_engine::{Engine, EngineHandle, TokioBackend, TokioConfig}; use tokio::io::{AsyncBufReadExt, AsyncWriteExt}; const DEFAULT_WORKER_SCRIPT: &str = "/usr/local/share/myelin/tinygrad_worker.py"; @@ -329,14 +329,18 @@ fn spawn_debug_join_listener( let listener = match tokio::net::UnixListener::bind(&path) { Ok(l) => l, Err(e) => { - let _ = ready_tx - .send(Err(format!("bind debug join socket {}: {e}", path.display()))); + let _ = ready_tx.send(Err(format!( + "bind debug join socket {}: {e}", + path.display() + ))); return; } }; if let Err(e) = fs::set_permissions(&path, fs::Permissions::from_mode(0o600)) { - let _ = ready_tx - .send(Err(format!("chmod debug join socket {}: {e}", path.display()))); + let _ = ready_tx.send(Err(format!( + "chmod debug join socket {}: {e}", + path.display() + ))); return; } let _ = ready_tx.send(Ok(())); @@ -577,10 +581,7 @@ fn submit_sampler_sample_health( "failed", json!({"state":"error","sample_seq":seq,"error":error}), ), - None => ( - "ready", - json!({"state":"sample_observed","sample_seq":seq}), - ), + None => ("ready", json!({"state":"sample_observed","sample_seq":seq})), }; submit_sampler_health( producer, @@ -683,7 +684,7 @@ fn spawn_host_cpu_sampler( let engine_inner = engine.clone(); engine.spawn(async move { use datastream::hardware::cpu::{ - CpuSampler, CPU_SAMPLE_INTERVAL, HOST_CPU_CHANNEL, HostCpuSample, + CPU_SAMPLE_INTERVAL, CpuSampler, HOST_CPU_CHANNEL, HostCpuSample, }; submit_sampler_started( &producer, @@ -1092,9 +1093,12 @@ impl WorkerEdgeRuntime { let record = match record { Ok(record) => record, Err(e) => { - let _ = stack - .runtime - .send_to(node_actor, NodeAgentMsg::OutputFault { edge_id: outbound.edge_id }); + let _ = stack.runtime.send_to( + node_actor, + NodeAgentMsg::OutputFault { + edge_id: outbound.edge_id, + }, + ); return Err(format!("read egress ring: {e}")); } }; @@ -1127,9 +1131,12 @@ impl WorkerEdgeRuntime { .ok_or_else(|| "outbound edge sender missing".to_owned())?; let edge_send_started = Instant::now(); if let Err(e) = sender.send(record) { - let _ = stack - .runtime - .send_to(node_actor, NodeAgentMsg::OutputFault { edge_id: outbound.edge_id }); + let _ = stack.runtime.send_to( + node_actor, + NodeAgentMsg::OutputFault { + edge_id: outbound.edge_id, + }, + ); return Err(e); } let edge_send_ms = duration_ms_u64(edge_send_started.elapsed()); @@ -1199,9 +1206,13 @@ impl WorkerEdgeRuntime { Ok(Some(record)) => records.push(record), Ok(None) => break, Err(e) => { - let _ = stack - .runtime - .send_to(node_actor, NodeAgentMsg::ObjectFailed { edge_id, object_id: None }); + let _ = stack.runtime.send_to( + node_actor, + NodeAgentMsg::ObjectFailed { + edge_id, + object_id: None, + }, + ); return Err(e); } } @@ -1252,9 +1263,13 @@ impl WorkerEdgeRuntime { ) { Ok(loaded) => loaded, Err(e) => { - let _ = stack - .runtime - .send_to(node_actor, NodeAgentMsg::ObjectFailed { edge_id, object_id: Some(record.object_id) }); + let _ = stack.runtime.send_to( + node_actor, + NodeAgentMsg::ObjectFailed { + edge_id, + object_id: Some(record.object_id), + }, + ); return Err(e); } }; @@ -1585,10 +1600,7 @@ impl WorkerEdgeRuntime { edge::EdgeLifecycleEvent::EdgeFaulted { edge_id, reason } => { stack .runtime - .send_to( - node_actor, - NodeAgentMsg::EdgeFault { edge_id: edge_id.0 }, - ) + .send_to(node_actor, NodeAgentMsg::EdgeFault { edge_id: edge_id.0 }) .map_err(|e| format!("report edge fault: {e}"))?; return Err(format!("edge {} faulted: {reason:?}", edge_id.0)); } @@ -1736,7 +1748,11 @@ fn run() -> Result<(), String> { .and_then(|backend| Engine::new(parts, backend)) { Ok(engine) => { - boot("engine", "ready", json!({"backend":"tokio","owns":"core+substrate"}))?; + boot( + "engine", + "ready", + json!({"backend":"tokio","owns":"core+substrate"}), + )?; engine } Err(error) => { @@ -1928,28 +1944,26 @@ fn run() -> Result<(), String> { ); } let mut debug_join_rx = match &config.debug_join_socket { - Some(path) => { - match spawn_debug_join_listener(engine.handle(), PathBuf::from(path)) { - Ok(rx) => { - node_runtime( - &mut datastream, - "debug_join_socket", - "ready", - json!({"socket":path}), - ); - Some(rx) - } - Err(error) => { - node_runtime( - &mut datastream, - "debug_join_socket", - "failed", - json!({"socket":path,"error":error}), - ); - return Err(format!("bind debug join socket {}: {error}", path)); - } + Some(path) => match spawn_debug_join_listener(engine.handle(), PathBuf::from(path)) { + Ok(rx) => { + node_runtime( + &mut datastream, + "debug_join_socket", + "ready", + json!({"socket":path}), + ); + Some(rx) } - } + Err(error) => { + node_runtime( + &mut datastream, + "debug_join_socket", + "failed", + json!({"socket":path,"error":error}), + ); + return Err(format!("bind debug join socket {}: {error}", path)); + } + }, None => { node_runtime( &mut datastream, @@ -2547,7 +2561,7 @@ impl PendingRuntimeReady { .coordinator_endpoint .as_ref() .map(|endpoint| DistNodeId(*endpoint.id.as_bytes())), - readiness_id: 1, + readiness_id: config.attempt_id, attempts: 0, next_attempt_at: Instant::now(), backoff: RUNTIME_READY_RETRY_INITIAL, @@ -2748,9 +2762,7 @@ fn handle_prompt_request( "started", json!({"request_id":request_id,"command":"InferPrompt","max_tokens":max_tokens}), ); - match worker.infer_prompt( - request_id, &prompt, max_tokens, config, datastream, - ) { + match worker.infer_prompt(request_id, &prompt, max_tokens, config, datastream) { Ok(result) => { let text = result .get("text") @@ -3722,13 +3734,7 @@ fn run_self_test( config, datastream, )?; - let result = worker.infer_prompt( - 0, - prompt, - config.self_test_max_tokens, - config, - datastream, - )?; + let result = worker.infer_prompt(0, prompt, config.self_test_max_tokens, config, datastream)?; let record = json!({"type":"self_test_completed","prompt_bytes":prompt.len(),"result":result}); datastream.submit_text(datastream.channels.node_self_test, record.to_string()); emit_node_event( @@ -3753,6 +3759,7 @@ fn env_optional(name: &str) -> Option { struct DeploymentConfig { run_id: u64, logical_node_id: u64, + attempt_id: u64, stage_index: u32, coordinator_endpoint: Option, orchestrator_actor: Option, @@ -3786,6 +3793,7 @@ impl DeploymentConfig { } let run_id = env_parse!("MYELIN_RUN_ID", 1)?; let logical_node_id = env_parse!("MYELIN_LOGICAL_NODE_ID", 1)?; + let attempt_id = env_parse!("MYELIN_NODE_ATTEMPT_ID", 1)?; let relay = relay_runtime_config_from_env(run_id)?; let debug_join_socket = match env_optional("MYELIN_DEBUG_JOIN_SOCKET").as_deref() { Some("disabled") => None, @@ -3808,6 +3816,7 @@ impl DeploymentConfig { Ok(Self { run_id, logical_node_id, + attempt_id, stage_index: env_parse!("MYELIN_STAGE_INDEX", 0)?, coordinator_endpoint: env_optional("MYELIN_COORDINATOR_ENDPOINT") .map(|value| { @@ -3832,7 +3841,8 @@ impl DeploymentConfig { worker_script: env_optional("MYELIN_TINYGRAD_WORKER") .unwrap_or_else(|| DEFAULT_WORKER_SCRIPT.to_owned()), device: env_optional("DEV").unwrap_or_else(|| default_device.to_owned()), - model_id: env_optional("MYELIN_MODEL_ID").unwrap_or_else(|| DEFAULT_MODEL_ID.to_owned()), + model_id: env_optional("MYELIN_MODEL_ID") + .unwrap_or_else(|| DEFAULT_MODEL_ID.to_owned()), gguf_source: if let Some(path) = env_optional("MYELIN_GGUF_LOCAL_PATH") { GgufSource::LocalPath(path) } else { @@ -4097,7 +4107,11 @@ struct TinygradWorker { } impl TinygradWorker { - fn spawn(config: &DeploymentConfig, arena_fd: std::os::fd::RawFd, engine: EngineHandle) -> Result { + fn spawn( + config: &DeploymentConfig, + arena_fd: std::os::fd::RawFd, + engine: EngineHandle, + ) -> Result { let mut child = Command::new("python3") .arg(&config.worker_script) .env("DEV", &config.device) diff --git a/apps/myelin/src/observability/lifecycle.rs b/apps/myelin/src/observability/lifecycle.rs index 3c9f477..653567c 100644 --- a/apps/myelin/src/observability/lifecycle.rs +++ b/apps/myelin/src/observability/lifecycle.rs @@ -1,4 +1,3 @@ - use serde::{Deserialize, Serialize}; #[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)] diff --git a/apps/myelin/src/observability/provisioning_logs.rs b/apps/myelin/src/observability/provisioning_logs.rs index 5df5891..051b677 100644 --- a/apps/myelin/src/observability/provisioning_logs.rs +++ b/apps/myelin/src/observability/provisioning_logs.rs @@ -1,4 +1,3 @@ - use std::io::{BufRead, BufReader, Read}; use std::thread::{self, JoinHandle}; @@ -179,4 +178,3 @@ pub(crate) fn parse_stdio_datastream_frame( payload: frame.payload.to_string(), }) } - diff --git a/apps/myelin/src/observability/telemetry.rs b/apps/myelin/src/observability/telemetry.rs index 5573396..c7b7180 100644 --- a/apps/myelin/src/observability/telemetry.rs +++ b/apps/myelin/src/observability/telemetry.rs @@ -1,4 +1,3 @@ - //! Myelin-system-owned datastream channel records. use datastream::Record; @@ -64,4 +63,3 @@ pub(crate) fn myelin_provision_log_channel(node_id: u64, stream: ProvisionLogStr impl Record for MyelinProvisionLogRecord { const CHANNEL: &'static str = MYELIN_PROVISIONING_LOGS; } - diff --git a/apps/myelin/src/orchestration/actor.rs b/apps/myelin/src/orchestration/actor.rs index bf6719b..53ead05 100644 --- a/apps/myelin/src/orchestration/actor.rs +++ b/apps/myelin/src/orchestration/actor.rs @@ -234,7 +234,10 @@ impl OrchestratorActor { run_id: core::RunId(run_id), stage_index, }), - OrchestratorMsg::ObserveNodeRuntimeReady { .. } | OrchestratorMsg::ObserveNodeRuntimeReadyAck { .. } | OrchestratorMsg::ObserveWeightsReady { .. } | OrchestratorMsg::Snapshot { .. } => {} + OrchestratorMsg::ObserveNodeRuntimeReady { .. } + | OrchestratorMsg::ObserveNodeRuntimeReadyAck { .. } + | OrchestratorMsg::ObserveWeightsReady { .. } + | OrchestratorMsg::Snapshot { .. } => {} OrchestratorMsg::ObserveTokenInEndpointReady => { self.core.observe(core::RunEvent::TokenInEndpointReady) } diff --git a/apps/myelin/src/orchestration/app.rs b/apps/myelin/src/orchestration/app.rs index 4454fef..5059ff6 100644 --- a/apps/myelin/src/orchestration/app.rs +++ b/apps/myelin/src/orchestration/app.rs @@ -8,7 +8,7 @@ use std::path::{Path, PathBuf}; use std::process::Command; use std::sync::{Arc, mpsc}; use std::thread; -use std::time::{Duration, Instant}; +use std::time::{Duration, Instant, SystemTime}; use crate::DEFAULT_PIPELINE_CACHED_MODEL_FILE; use crate::codecs::register_myelin_actor_codecs; @@ -21,7 +21,6 @@ use crate::observability::dashboard_view::MyelinClusterDashboardView; use crate::observability::{benchmark, frame_archive::FrameArchive}; use crate::orchestration::actor::{OrchestratorActor, OrchestratorMsg, OrchestratorReport}; use crate::orchestration::config::{DEFAULT_CONFIG_PATH, TomlConfigOverlay}; -const PROVIDER_START_MAX_ATTEMPTS: usize = 4; use crate::gguf_shard::{StageShardPlan, plan_stage_shard}; use crate::node_provisioning::{ProviderKind, provider_kind}; @@ -29,10 +28,11 @@ use crate::observability::telemetry::{ MYELIN_PROVISIONING_EVENTS, MyelinProvisionEventRecord, MyelinProvisionLogRecord, myelin_provision_log_channel, }; +use crate::orchestration::cluster_reconciler::{ProvisionedClusterGuard, ReconcilerNodeBinding}; use crate::orchestration::distribution_stack::{DistributionRuntimeStack, duration_ms_u64}; use crate::orchestration::provider_adapters::relay::{ - MYELIN_IROH_RELAY_URL_ENV, RelayRuntimeConfig, SWACTOR_IROH_RELAY_URL_ENV, relay_mode_env_value, - relay_runtime_config_from_settings, + MYELIN_IROH_RELAY_URL_ENV, RelayRuntimeConfig, SWACTOR_IROH_RELAY_URL_ENV, + relay_mode_env_value, relay_runtime_config_from_settings, }; use crate::orchestration::provider_adapters::vastai::{ SshCommandBootstrapLauncher, ToolsVastAiLeaseClient, VastAiProvisioningConfig, @@ -48,6 +48,11 @@ use crate::provisioning::{ }; use crate::run_fsm::{RunConfig, RunId}; use crate::run_plan::{self, GgufSource, TokenizerSource}; +use ::provisioning::{ + BootSpec, ClusterShape, DesiredNodeShape, LogicalNodeId as ReconcilerLogicalNodeId, + NodeGroupId, ProviderKind as ReconcilerProviderKind, RetryPolicy, RoleId, + RunId as ClusterRunId, RunNodeGroupSpec, SwactorId, SwarmJoinTemplate, +}; use data_plane::object_record as ingress; use datastream::{ ChannelContent, ChannelId, ChannelRef, DatastreamEndpoint, DatastreamEvent, DatastreamProducer, @@ -55,8 +60,8 @@ use datastream::{ StreamId, StreamOrigin, SubscriptionRequest, }; use distribution::node::DistributedNodeConfig; -use distribution::telemetry::{MembershipTransition, SwimProbeEvent}; use distribution::swim::telemetry::ObservedTransition; +use distribution::telemetry::{MembershipTransition, SwimProbeEvent}; use distribution::types::{MemberState, NodeId as DistNodeId}; use iroh::EndpointAddr; use iroh_driver::{ @@ -514,7 +519,12 @@ where }), ); - let rpc_addr = match spawn_prompt_rpc(&engine.handle(), config.rpc_bind, work_tx, config.default_max_tokens) { + let rpc_addr = match spawn_prompt_rpc( + &engine.handle(), + config.rpc_bind, + work_tx, + config.default_max_tokens, + ) { Ok(addr) => { bootstrap( &mut orch_datastream, @@ -1135,7 +1145,9 @@ impl ConfigBuilder { self.config_profile = RuntimeConfigProfile::parse(&profile)?; }); env_parse!("MYELIN_RUN_ID", |run_id| { self.run_id = run_id }); - env_parse!("MYELIN_LOGICAL_NODE_ID", |node_id| { self.node_id = node_id }); + env_parse!("MYELIN_LOGICAL_NODE_ID", |node_id| { + self.node_id = node_id + }); env_parse!("MYELIN_STAGE_INDEX", |stage_index| { self.stage_index = stage_index }); @@ -1821,6 +1833,7 @@ impl Config { let mut keys: Vec = vec![ "MYELIN_RUN_ID", "MYELIN_LOGICAL_NODE_ID", + "MYELIN_NODE_ATTEMPT_ID", "MYELIN_NODE_PROVIDER", "MYELIN_STAGE_INDEX", "MYELIN_COORDINATOR_ENDPOINT", @@ -1956,6 +1969,7 @@ impl Config { Ok(NodeProvisionSpec { run_id: self.run_id, node_id: logical_node_id, + attempt_id: 0, stage_index: Some(stage_index), image: self.image.clone(), env, @@ -2058,7 +2072,8 @@ fn wait_for_runtime_ready_acks( ctx: RuntimeReadyAckLoop<'_>, targets: &[RuntimeReadyAckTarget], collector_endpoint: &EndpointAddr, -) -> Result<(), String> { + cluster: &mut ProvisionedClusterGuard, +) -> Result { let RuntimeReadyAckLoop { driver, stack, @@ -2103,6 +2118,15 @@ fn wait_for_runtime_ready_acks( let mut last_send = None::; while !pending.is_empty() { + cluster + .poll(SystemTime::now()) + .map_err(|error| format!("cluster reconcile while awaiting ready ack: {error}"))?; + if targets.iter().any(|target| { + cluster.current_attempt(target.node_id) + != Some(::provisioning::NodeAttemptId(target.ready.readiness_id)) + }) { + return Ok(false); + } pump(driver, frame_tx); drain_orch_stdio_capture( orch_stdio_rx, @@ -2116,13 +2140,9 @@ fn wait_for_runtime_ready_acks( "shutdown requested while waiting for runtime-ready acknowledgements".to_owned(), ); } - drain_observations_with_exit( - obs_rx, - dashboard, - orch_datastream, - provider, - |node_id, status| format!("node {node_id} exited before ready: {status:?}"), - )?; + while let Ok(observation) = obs_rx.try_recv() { + emit_plugin_observation(orch_datastream, dashboard, provider, &observation); + } drain_frames(frame_rx, dashboard, orch_datastream); while let Some(report) = orchestrator_reports.try_recv() { let OrchestratorReport::NodeRuntimeReadyAck { @@ -2154,7 +2174,7 @@ fn wait_for_runtime_ready_acks( ); } if pending.is_empty() { - return Ok(()); + return Ok(true); } if last_send.is_none_or(|sent_at| sent_at.elapsed() >= RUNTIME_READY_ACK_RETRY_INTERVAL) { for (key, target) in &pending { @@ -2195,66 +2215,112 @@ fn wait_for_runtime_ready_acks( } thread::sleep(PUMP_INTERVAL); } - Ok(()) + Ok(true) } -struct ProvisionedClusterGuard { +fn reconciler_group(config: &Config, spec: &NodeProvisionSpec) -> RunNodeGroupSpec { + let group_id = NodeGroupId(format!("node-{}", spec.node_id)); + let ssh_user = config + .vastai + .as_ref() + .map(|vastai| vastai.provisioning.ssh_user.clone()) + .unwrap_or_else(|| "root".to_owned()); + let disk_gb = config + .vastai + .as_ref() + .map(|vastai| vastai.provisioning.disk_gb) + .unwrap_or_default(); + let orchestrator = spec + .env + .iter() + .find(|(name, _)| name == "MYELIN_ORCHESTRATOR_ACTOR") + .map(|(_, value)| value.clone()) + .unwrap_or_default(); + RunNodeGroupSpec { + run_id: ClusterRunId(config.run_id), + group_id, + role: RoleId(format!( + "stage-{}", + spec.stage_index.unwrap_or(config.stage_index) + )), + count: 1, + provider: ReconcilerProviderKind::new(config.provider.as_str()), + shape: DesiredNodeShape { + image: spec.image.clone(), + disk_gb, + gpu_name: None, + min_gpu_ram_mb: None, + min_down_mbps: None, + min_up_mbps: None, + min_reliability: None, + require_verified: false, + provider_labels: BTreeMap::from([( + "myelin.provider_config".to_owned(), + config.provider_datastream_detail().to_string(), + )]), + }, + boot: BootSpec { + ssh_user, + verify_commands: Vec::new(), + start_swactor_command: spec.args.join(" "), + stdout_sources: Vec::new(), + stderr_sources: Vec::new(), + env: spec.env.clone(), + args: spec.args.clone(), + mounts: spec.mounts.clone(), + }, + swarm_join: SwarmJoinTemplate { + orch_swactor_addr: orchestrator, + join_token_ref: "myelin-runtime-ready".to_owned(), + }, + } +} + +fn build_reconciled_cluster( provisioner: Box, - handles: Vec, + config: &Config, + stage_specs: &[NodeProvisionSpec], + runtime: swactor::runtime::Runtime, + engine: EngineHandle, + sink: PluginSink, +) -> Result { + let groups = stage_specs + .iter() + .map(|spec| reconciler_group(config, spec)) + .collect::>(); + let desired = ClusterShape { + run_id: ClusterRunId(config.run_id), + generation: 1, + groups, + }; + let expanded = desired.expand().map_err(|error| error.to_string())?; + let mut first = Some(provisioner); + let mut bindings = Vec::with_capacity(stage_specs.len()); + for spec in stage_specs { + let logical_node_id = ReconcilerLogicalNodeId(format!("node-{}-0", spec.node_id)); + if !expanded.contains_key(&logical_node_id) { + return Err(format!( + "reconciler shape did not expand node {}", + logical_node_id.0 + )); + } + let plugin = match first.take() { + Some(plugin) => plugin, + None => config.build_provisioner(runtime.clone())?, + }; + bindings.push(ReconcilerNodeBinding { + logical_node_id, + provision: spec.clone(), + plugin, + }); + } + ProvisionedClusterGuard::new(desired, bindings, RetryPolicy::default(), engine, sink) } -impl ProvisionedClusterGuard { - fn new( - provisioner: Box, - handles: Vec, - ) -> Self { - Self { - provisioner, - handles, - } - } - - fn stop(&mut self) -> Result<(), String> { - let mut first_error = None; - while let Some(handle) = self.handles.pop() { - if let Err(error) = self.provisioner.stop_node(&handle) - && first_error.is_none() - { - first_error = Some(error); - } - } - match first_error { - Some(error) => Err(error), - None => Ok(()), - } - } - - fn complete_bootstrap(&mut self) -> Result<(), String> { - let mut first_error = None; - for handle in &self.handles { - if let Err(error) = self.provisioner.complete_bootstrap(handle) - && first_error.is_none() - { - first_error = Some(error); - } - } - match first_error { - Some(error) => Err(error), - None => Ok(()), - } - } -} - -impl Drop for ProvisionedClusterGuard { - fn drop(&mut self) { - let _ = self.stop(); - } -} - -// provider lifecycle/provisioning is out of scope (ENGINE_SPEC.md §2) +// Synchronous orchestration sequencing; provider work and timers are engine-hosted. #[allow(clippy::disallowed_methods)] fn start_and_provision_workers( - mut provisioner: Box, + provisioner: Box, config: &Config, pipeline_plan: Option<&run_plan::RunPlan>, ctx: RuntimeReadyAckLoop<'_>, @@ -2262,7 +2328,14 @@ fn start_and_provision_workers( coordinator: EndpointAddr, pipeline_coordinator: EndpointAddr, orchestrator_actor: ActorAddress, -) -> Result<(ProvisionedClusterGuard, PromptRuntimeReady, BTreeMap), String> { +) -> Result< + ( + ProvisionedClusterGuard, + PromptRuntimeReady, + BTreeMap, + ), + String, +> { let RuntimeReadyAckLoop { driver, stack, @@ -2321,124 +2394,79 @@ fn start_and_provision_workers( } else { BTreeMap::new() }; - let mut handles = Vec::with_capacity(stage_specs.len()); - let mut pending_specs = stage_specs; - for attempt in 1..=PROVIDER_START_MAX_ATTEMPTS { - for node_spec in &pending_specs { - orch_datastream.emit_event( - dashboard, - ProvisionEvent { - run_id: config.run_id, - node_id: node_spec.node_id, - kind: ProvisionEventKind::ProvisionStart, - provider: Some(config.provider.as_str().to_owned()), - message: Some(format!( - "starting {} image {}", - config.provider.as_str(), - config.image - )), - }, - ); - bootstrap( - orch_datastream, - "provider_start", - "started", - json!({ - "provider":config.provider.as_str(), - "image":&config.image, - "node_id":node_spec.node_id, - "stage_index":node_spec.stage_index, - "attempt":attempt, - }), - ); - } - let (tx, rx) = mpsc::channel(); - thread::spawn({ - let sink = sink.clone(); - move || { - let mut provisioner = provisioner; - let results = provisioner.start_nodes(pending_specs, sink); - let _ = tx.send((provisioner, results)); - } - }); - let start_results = loop { - match rx.recv_timeout(Duration::from_millis(100)) { - Ok((returned_provisioner, start_results)) => { - provisioner = returned_provisioner; - break start_results; - } - Err(mpsc::RecvTimeoutError::Timeout) => { - drain_orch_stdio_capture( - orch_stdio_rx, - orch_datastream, - dashboard, - config.run_id, - config.node_id, - ); - } - Err(mpsc::RecvTimeoutError::Disconnected) => { - return Err("provider start worker disconnected".to_owned()); - } - } - }; - let mut failed_specs = Vec::new(); - let mut first_error = None; - for (node_spec, handle_result) in start_results { - match handle_result { - Ok(handle) => { - bootstrap( - orch_datastream, - "provider_start", - "ready", - json!({ - "provider":config.provider.as_str(), - "node_id":node_spec.node_id, - "stage_index":node_spec.stage_index, - "attempt":attempt, - }), - ); - handles.push(handle); - } - Err(error) => { - bootstrap( - orch_datastream, - "provider_start", - "failed", - json!({ - "provider":config.provider.as_str(), - "node_id":node_spec.node_id, - "stage_index":node_spec.stage_index, - "attempt":attempt, - "error":error, - }), - ); - failed_specs.push(node_spec); - if first_error.is_none() { - first_error = Some(error); - } - } - } - } - if first_error.is_none() { + for node_spec in &stage_specs { + orch_datastream.emit_event( + dashboard, + ProvisionEvent { + run_id: config.run_id, + node_id: node_spec.node_id, + kind: ProvisionEventKind::ProvisionStart, + provider: Some(config.provider.as_str().to_owned()), + message: Some(format!( + "reconciling {} image {}", + config.provider.as_str(), + config.image + )), + }, + ); + bootstrap( + orch_datastream, + "provider_start", + "started", + json!({ + "provider":config.provider.as_str(), + "image":&config.image, + "node_id":node_spec.node_id, + "stage_index":node_spec.stage_index, + "generation":1, + }), + ); + } + let mut provisioned_nodes = build_reconciled_cluster( + provisioner, + config, + &stage_specs, + stack.runtime.clone(), + stack.engine.clone(), + sink, + )?; + loop { + provisioned_nodes + .poll(SystemTime::now()) + .map_err(|error| format!("cluster reconcile: {error}"))?; + if provisioned_nodes.awaiting_runtime() || provisioned_nodes.is_converged() { break; } - if attempt == PROVIDER_START_MAX_ATTEMPTS { - let error = first_error.expect("checked provider-start failure"); - while let Some(handle) = handles.pop() { - let _ = provisioner.stop_node(&handle); - } - drain_orch_stdio_capture( - orch_stdio_rx, - orch_datastream, - dashboard, - config.run_id, - config.node_id, - ); - return Err(error); + pump(driver, frame_tx); + drain_frames(frame_rx, dashboard, orch_datastream); + drain_orch_stdio_capture( + orch_stdio_rx, + orch_datastream, + dashboard, + config.run_id, + config.node_id, + ); + while let Ok(observation) = obs_rx.try_recv() { + emit_plugin_observation(orch_datastream, dashboard, &config.provider, &observation); } - pending_specs = failed_specs; + if stop_requested(stop_rx) { + return Err("shutdown requested while reconciling worker nodes".to_owned()); + } + thread::sleep(PUMP_INTERVAL); + } + for node_spec in &stage_specs { + bootstrap( + orch_datastream, + "provider_start", + "ready", + json!({ + "provider":config.provider.as_str(), + "node_id":node_spec.node_id, + "stage_index":node_spec.stage_index, + "attempt":provisioned_nodes.current_attempt(node_spec.node_id).map(|attempt| attempt.0), + }), + ); } - let mut provisioned_nodes = ProvisionedClusterGuard::new(provisioner, handles); drain_orch_stdio_capture( orch_stdio_rx, orch_datastream, @@ -2453,8 +2481,8 @@ fn start_and_provision_workers( "started", json!({"worker_count":expected_node_ids.len(),"node_ids":expected_node_ids}), ); - let readies = if pipeline_plan.is_some() { - match wait_for_runtime_readies( + let readies = loop { + let readies = match wait_for_runtime_readies( RuntimeReadyAckLoop { driver, stack, @@ -2472,6 +2500,7 @@ fn start_and_provision_workers( orchestrator_actor, }, &expected_node_ids, + &mut provisioned_nodes, ) { Ok(readies) => readies, Err(error) => { @@ -2483,75 +2512,80 @@ fn start_and_provision_workers( ); return Err(error); } - } - } else { - let ready = match wait_for_runtime_ready(RuntimeReadyAckLoop { - driver, - stack, - obs_rx, - frame_rx, - frame_tx, - orchestrator_reports, - stop_rx, - dashboard, - orch_datastream, - orch_stdio_rx, - run_id: config.run_id, - orchestrator_node_id: config.node_id, - provider: &config.provider, - orchestrator_actor, - }) { - Ok(ready) => ready, - Err(error) => { - bootstrap( - orch_datastream, - "node_runtime_ready", - "failed", - json!({"error":error}), - ); - return Err(error); - } }; - BTreeMap::from([(config.node_id, ready)]) - }; - for (node_id, ready) in &readies { + for (node_id, ready) in &readies { + bootstrap( + orch_datastream, + "node_runtime_ready", + "ready", + json!({"endpoint":&ready.endpoint,"node_actor":ready.node_actor,"node_id":node_id,"stage_index":ready.stage_index,"attempt":ready.readiness_id}), + ); + } + let ack_targets = readies + .iter() + .map(|(node_id, ready)| RuntimeReadyAckTarget { + node_id: *node_id, + ready: ready.clone(), + }) + .collect::>(); + if wait_for_runtime_ready_acks( + RuntimeReadyAckLoop { + driver, + stack, + obs_rx, + frame_rx, + frame_tx, + orchestrator_reports, + stop_rx, + dashboard, + orch_datastream, + orch_stdio_rx, + run_id: config.run_id, + orchestrator_node_id: config.node_id, + provider: &config.provider, + orchestrator_actor, + }, + &ack_targets, + &pipeline_coordinator, + &mut provisioned_nodes, + )? { + break readies; + } bootstrap( orch_datastream, "node_runtime_ready", - "ready", - json!({"endpoint":&ready.endpoint,"node_actor":ready.node_actor,"node_id":node_id,"stage_index":ready.stage_index}), + "retry", + json!({"reason":"node attempt changed before ready acknowledgement"}), ); + }; + for (node_id, ready) in &readies { + let attempt = ::provisioning::NodeAttemptId(ready.readiness_id); + if !provisioned_nodes.observe_runtime_ready( + *node_id, + attempt, + SwactorId(format!("{:?}", ready.node_actor)), + SystemTime::now(), + ) { + return Err(format!( + "stale runtime-ready observation for node {node_id} attempt {}", + attempt.0 + )); + } + } + while !provisioned_nodes.is_converged() { + provisioned_nodes + .poll(SystemTime::now()) + .map_err(|error| format!("cluster convergence: {error}"))?; + pump(driver, frame_tx); + drain_frames(frame_rx, dashboard, orch_datastream); + while let Ok(observation) = obs_rx.try_recv() { + emit_plugin_observation(orch_datastream, dashboard, &config.provider, &observation); + } + if stop_requested(stop_rx) { + return Err("shutdown requested while converging worker nodes".to_owned()); + } + thread::sleep(PUMP_INTERVAL); } - let ack_targets = readies - .iter() - .map(|(node_id, ready)| RuntimeReadyAckTarget { - node_id: *node_id, - ready: ready.clone(), - }) - .collect::>(); - wait_for_runtime_ready_acks( - RuntimeReadyAckLoop { - driver, - stack, - obs_rx, - frame_rx, - frame_tx, - orchestrator_reports, - stop_rx, - dashboard, - orch_datastream, - orch_stdio_rx, - run_id: config.run_id, - orchestrator_node_id: config.node_id, - provider: &config.provider, - orchestrator_actor, - }, - &ack_targets, - &pipeline_coordinator, - )?; - provisioned_nodes - .complete_bootstrap() - .map_err(|e| format!("complete provider bootstrap after runtime-ready: {e}"))?; bootstrap( orch_datastream, @@ -2927,6 +2961,7 @@ fn stage_ring_spec_wire(spec: run_plan::RingSpec) -> StageRingSpecWire { fn wait_for_runtime_readies( ctx: RuntimeReadyAckLoop<'_>, expected_node_ids: &[u64], + cluster: &mut ProvisionedClusterGuard, ) -> Result, String> { let RuntimeReadyAckLoop { driver, @@ -2946,6 +2981,13 @@ fn wait_for_runtime_readies( let expected = expected_node_ids.iter().copied().collect::>(); let mut pending = BTreeMap::::new(); loop { + cluster + .poll(SystemTime::now()) + .map_err(|error| format!("cluster reconcile while awaiting runtime: {error}"))?; + pending.retain(|node_id, ready| { + cluster.current_attempt(*node_id) + == Some(::provisioning::NodeAttemptId(ready.readiness_id)) + }); pump(driver, frame_tx); emit_swim_transitions( orch_datastream, @@ -2972,13 +3014,9 @@ fn wait_for_runtime_readies( PluginObservation::DatastreamFrame { .. } | PluginObservation::ProviderLine { .. } | PluginObservation::StdoutLine { .. } - | PluginObservation::StderrLine { .. } => {} - PluginObservation::Failed { reason, .. } => return Err(reason), - PluginObservation::Exited { - status, node_id, .. - } => { - return Err(format!("node {node_id} exited before ready: {status:?}")); - } + | PluginObservation::StderrLine { .. } + | PluginObservation::Failed { .. } + | PluginObservation::Exited { .. } => {} } } while let Some(report) = orchestrator_reports.try_recv() { @@ -2993,6 +3031,8 @@ fn wait_for_runtime_readies( } = report && report_run_id == run_id && expected.contains(&node_id) + && cluster.current_attempt(node_id) + == Some(::provisioning::NodeAttemptId(readiness_id)) { pending.insert( node_id, @@ -3739,7 +3779,9 @@ impl OrchDatastream { producer, channels: BTreeMap::new(), channel_names: BTreeMap::new(), - archive: frame_log.map(|p| FrameArchive::open_with_label(p, "datastream frame log")).transpose()? + archive: frame_log + .map(|p| FrameArchive::open_with_label(p, "datastream frame log")) + .transpose()?, }; for name in [ MYELIN_PROVISIONING_EVENTS, @@ -3788,8 +3830,8 @@ impl OrchDatastream { fn emit_log(&mut self, dashboard: Option<&DashboardSupport>, line: ProvisionLogLine) { let channel = myelin_provision_log_channel(line.node_id, line.stream); - let payload = - serde_json::to_vec(&MyelinProvisionLogRecord::new(line)).expect("serialize provision log"); + let payload = serde_json::to_vec(&MyelinProvisionLogRecord::new(line)) + .expect("serialize provision log"); self.emit_bytes(dashboard, &channel, payload); } @@ -4068,7 +4110,8 @@ impl DashboardSupport { fn start(enabled: bool, _engine: &EngineHandle) -> Result, String> { if enabled { return Err( - "MYELIN_DASHBOARD requires building myelin-system with feature dashboard".to_owned(), + "MYELIN_DASHBOARD requires building myelin-system with feature dashboard" + .to_owned(), ); } Ok(None) @@ -4077,7 +4120,6 @@ impl DashboardSupport { fn publish_frame(&self, _stream: &StreamId, _channel: &str, _frame: &Frame) {} } - struct ChannelObservationSink { tx: Mutex>, } @@ -4100,8 +4142,8 @@ fn spawn_prompt_rpc( // reusing the synchronous request/response parser unchanged. There is no // listener thread and no per-connection std thread (ENGINE_SPEC.md); // no raw Tokio handle or second runtime is introduced. - let std_listener = TcpListener::bind(bind) - .map_err(|e| format!("bind prompt RPC {bind}: {e}"))?; + let std_listener = + TcpListener::bind(bind).map_err(|e| format!("bind prompt RPC {bind}: {e}"))?; let addr = std_listener .local_addr() .map_err(|e| format!("read prompt RPC addr: {e}"))?; @@ -4123,8 +4165,7 @@ fn spawn_prompt_rpc( if let Ok(std_stream) = stream.into_std() { // The synchronous parser uses blocking I/O. let _ = std_stream.set_nonblocking(false); - let _ = - handle_prompt_connection(std_stream, tx, default_max_tokens); + let _ = handle_prompt_connection(std_stream, tx, default_max_tokens); } }); } @@ -4171,119 +4212,6 @@ fn handle_prompt_connection( Ok(()) } -// synchronous process-control/orchestration sequencing; the engine drives all background work (ENGINE_SPEC.md §2) -#[allow(clippy::disallowed_methods)] -fn wait_for_runtime_ready(ctx: RuntimeReadyAckLoop<'_>) -> Result { - let RuntimeReadyAckLoop { - driver, - stack, - obs_rx, - frame_rx, - frame_tx, - orchestrator_reports, - stop_rx, - dashboard, - orch_datastream, - orch_stdio_rx, - run_id, - orchestrator_node_id: node_id, - provider, - .. - } = ctx; - let mut pending_ready: Option = None; - let mut node_swim_started = false; - let mut node_swim_ready = false; - let mut node_route_started = false; - loop { - pump(driver, frame_tx); - drain_frames(frame_rx, dashboard, orch_datastream); - drain_orch_stdio_capture(orch_stdio_rx, orch_datastream, dashboard, run_id, node_id); - if stop_requested(stop_rx) { - return Err("shutdown requested while waiting for node ready".to_owned()); - } - drain_observations_with_exit(obs_rx, dashboard, orch_datastream, provider, |_, status| { - format!("node exited before ready: {status:?}") - })?; - while let Some(report) = orchestrator_reports.try_recv() { - if let OrchestratorReport::NodeRuntimeReady { - run_id: report_run_id, - node_id: report_node_id, - stage_index, - endpoint, - node_actor, - datastream_publisher, - readiness_id, - } = report - { - if report_run_id == run_id && report_node_id == node_id { - let reset_progress = pending_ready - .as_ref() - .map(|ready| ready.readiness_id != readiness_id) - .unwrap_or(true); - if reset_progress { - node_swim_started = false; - node_swim_ready = false; - node_route_started = false; - } - let swim_node_id = DistNodeId(*endpoint.id.as_bytes()); - pending_ready = Some(RuntimeReady { - endpoint, - node_actor, - datastream_publisher, - stage_index, - readiness_id, - swim_node_id, - }); - } - } - } - if let Some(ready) = pending_ready.as_ref() { - if runtime_ready_barrier_met(stack, ready) { - return Ok(ready.clone()); - } - let swim_ready = stack.member_state(ready.swim_node_id) == Some(MemberState::Alive); - let route_ready = stack.route_owner(ready.node_actor) == Some(ready.swim_node_id); - if !swim_ready { - if !node_swim_started { - orch_datastream.emit_bootstrap( - dashboard, - run_id, - node_id, - "node_swim", - "started", - json!({"node":format!("{:?}", ready.swim_node_id),"readiness_id":ready.readiness_id}), - ); - node_swim_started = true; - } - } else if !route_ready { - if !node_swim_ready { - orch_datastream.emit_bootstrap( - dashboard, - run_id, - node_id, - "node_swim", - "ready", - json!({"node":format!("{:?}", ready.swim_node_id),"readiness_id":ready.readiness_id}), - ); - node_swim_ready = true; - } - if !node_route_started { - orch_datastream.emit_bootstrap( - dashboard, - run_id, - node_id, - "node_route", - "started", - json!({"node_actor":ready.node_actor,"node":format!("{:?}", ready.swim_node_id),"readiness_id":ready.readiness_id}), - ); - node_route_started = true; - } - } - } - thread::sleep(PUMP_INTERVAL); - } -} - fn provision_stage( stack: &DistributionRuntimeStack, node_actor: ActorAddress, @@ -5065,9 +4993,10 @@ fn serve_prompts( "started", json!({"source":"stdin"}), ); - let _ = stack - .runtime - .send_to(orchestrator_actor, OrchestratorMsg::ObserveOperatorStop { run_id }); + let _ = stack.runtime.send_to( + orchestrator_actor, + OrchestratorMsg::ObserveOperatorStop { run_id }, + ); pump(driver, frame_tx); return Ok(()); } @@ -5479,10 +5408,7 @@ fn emit_swim_probe_events( /// progression and protocol tick injection are owned by the engine (see /// `spawn_protocol_ticker`); this only drains integration-owned queues /// (ENGINE_SPEC.md). -fn pump( - driver: &IrohDriver, - frame_tx: &mpsc::Sender, -) { +fn pump(driver: &IrohDriver, frame_tx: &mpsc::Sender) { drain_datastream_connections(driver, frame_tx); } @@ -5589,7 +5515,8 @@ fn derive_ssh_public_key(identity: &Path) -> Result { fn ssh_public_key_fingerprint(public_key: &str) -> String { const UNAVAILABLE: &str = "unavailable"; - let path = std::env::temp_dir().join(format!("myelin-vastai-ssh-key-{}.pub", std::process::id())); + let path = + std::env::temp_dir().join(format!("myelin-vastai-ssh-key-{}.pub", std::process::id())); if std::fs::write(&path, format!("{public_key}\n")).is_err() { return UNAVAILABLE.to_owned(); } diff --git a/apps/myelin/src/orchestration/cluster_reconciler.rs b/apps/myelin/src/orchestration/cluster_reconciler.rs new file mode 100644 index 0000000..e29d363 --- /dev/null +++ b/apps/myelin/src/orchestration/cluster_reconciler.rs @@ -0,0 +1,1618 @@ +//! Myelin integration for the provider-neutral cluster reconciler. + +use std::collections::{BTreeMap, BTreeSet, VecDeque}; +use std::sync::mpsc::{self, Receiver, Sender, TryRecvError}; +use std::sync::{Arc, Mutex, MutexGuard, RwLock, RwLockReadGuard, RwLockWriteGuard}; +use std::time::{Duration, SystemTime}; + +use provisioning::{ + BlockingEffectSpawner, BootstrapSessionId, ClusterDriver, ClusterShape, CreateLeaseResult, + DestroyHandle, DriverError, EffectBackend, EffectError, ExecutorOperationStatus, + IdempotentEffectExecutor, LeaseFacts, LogicalNodeId, NodeAttemptId, NodeIntent, + NodeManagerCommand, NodeObservation, NodeStage, OperationOutcome, PlannedEffect, + ProviderLeaseId, RetryPolicy, SshEndpoint, SwactorId, +}; +use swactor_engine::EngineHandle; + +use crate::provisioning::{ + NodeProvisionSpec, PluginNodeHandle, PluginObservation, PluginObservationSink, PluginSink, + ProvisionPlugin, +}; + +const PERIODIC_RECONCILE: Duration = Duration::from_secs(30); +const ATTEMPT_ENV: &str = "MYELIN_NODE_ATTEMPT_ID"; + +pub(crate) struct ReconcilerNodeBinding { + pub logical_node_id: LogicalNodeId, + pub provision: NodeProvisionSpec, + pub plugin: Box, +} + +struct LiveNode { + attempt: NodeAttemptId, + handle: PluginNodeHandle, + lease: LeaseFacts, + endpoint: SshEndpoint, + failure_sink: Arc, + bootstrap_started: bool, +} + +struct StagedNodeEffects { + template: NodeProvisionSpec, + plugin: Box, +} + +struct NodeEffects { + template: NodeProvisionSpec, + plugin: Box, + live: Option, + staged: Option, +} + +#[derive(Clone, Debug)] +struct TaggedFailure { + node: LogicalNodeId, + attempt: NodeAttemptId, + reason: String, +} + +enum FailureGate { + Pending(Vec), + Armed, + Discarded, +} + +struct AttemptObservationSink { + downstream: PluginSink, + failures: Sender, + node: LogicalNodeId, + attempt: NodeAttemptId, + gate: Mutex, +} + +impl AttemptObservationSink { + fn arm(&self) { + let mut gate = self + .gate + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let FailureGate::Pending(reasons) = std::mem::replace(&mut *gate, FailureGate::Armed) + else { + return; + }; + for reason in reasons { + self.send_failure(reason); + } + } + + fn discard(&self) { + *self + .gate + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) = FailureGate::Discarded; + } + + fn record_failure(&self, reason: String) { + let mut gate = self + .gate + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + match &mut *gate { + FailureGate::Pending(reasons) => reasons.push(reason), + FailureGate::Armed => self.send_failure(reason), + FailureGate::Discarded => {} + } + } + + fn send_failure(&self, reason: String) { + let _ = self.failures.send(TaggedFailure { + node: self.node.clone(), + attempt: self.attempt, + reason, + }); + } +} + +impl PluginObservationSink for AttemptObservationSink { + fn observe(&self, observation: PluginObservation) { + let reason = match &observation { + PluginObservation::Failed { reason, .. } => Some(reason.clone()), + PluginObservation::Exited { + node_id, status, .. + } => Some(format!( + "node {node_id} exited during bootstrap: {status:?}" + )), + PluginObservation::StdoutLine { .. } + | PluginObservation::StderrLine { .. } + | PluginObservation::DatastreamFrame { .. } + | PluginObservation::ProviderLine { .. } => None, + }; + self.downstream.observe(observation); + if let Some(reason) = reason { + self.record_failure(reason); + } + } +} + +pub(crate) struct MyelinEffectBackend { + nodes: RwLock>>>, + sink: PluginSink, + failure_tx: Sender, +} + +impl MyelinEffectBackend { + fn new( + bindings: Vec, + sink: PluginSink, + failure_tx: Sender, + ) -> Result<(Self, BTreeMap), String> { + let mut nodes = BTreeMap::new(); + let mut by_external_id = BTreeMap::new(); + for binding in bindings { + if nodes.contains_key(&binding.logical_node_id) { + return Err(format!( + "duplicate reconciler binding {}", + binding.logical_node_id.0 + )); + } + if by_external_id + .insert(binding.provision.node_id, binding.logical_node_id.clone()) + .is_some() + { + return Err(format!( + "duplicate provision node id {}", + binding.provision.node_id + )); + } + nodes.insert( + binding.logical_node_id, + Arc::new(Mutex::new(NodeEffects { + template: binding.provision, + plugin: binding.plugin, + live: None, + staged: None, + })), + ); + } + Ok(( + Self { + nodes: RwLock::new(nodes), + sink, + failure_tx, + }, + by_external_id, + )) + } + + fn node(&self, id: &LogicalNodeId) -> Result>, EffectError> { + lock_nodes_read(&self.nodes) + .get(id) + .cloned() + .ok_or_else(|| EffectError::definite(format!("no effect binding for node {}", id.0))) + } + + fn has_node(&self, id: &LogicalNodeId) -> bool { + lock_nodes_read(&self.nodes).contains_key(id) + } + + fn register_binding(&self, binding: ReconcilerNodeBinding) -> (u64, LogicalNodeId) { + let external_id = binding.provision.node_id; + let logical_id = binding.logical_node_id; + let staged = StagedNodeEffects { + template: binding.provision, + plugin: binding.plugin, + }; + let mut nodes = lock_nodes_write(&self.nodes); + if let Some(effects) = nodes.get(&logical_id) { + lock_node(effects).staged = Some(staged); + } else { + nodes.insert( + logical_id.clone(), + Arc::new(Mutex::new(NodeEffects { + template: staged.template, + plugin: staged.plugin, + live: None, + staged: None, + })), + ); + } + (external_id, logical_id) + } + + fn stop_all(&self) -> Result<(), String> { + let mut first_error = None; + let nodes = lock_nodes_read(&self.nodes) + .values() + .cloned() + .collect::>(); + for effects in nodes { + let mut effects = lock_node(&effects); + let Some(live) = effects.live.take() else { + continue; + }; + live.failure_sink.discard(); + if let Err(error) = effects.plugin.stop_node(&live.handle) { + effects.live = Some(live); + if first_error.is_none() { + first_error = Some(error); + } + } + } + match first_error { + Some(error) => Err(error), + None => Ok(()), + } + } + + fn create_or_adopt( + &self, + effect: &PlannedEffect, + request: &provisioning::CreateLeaseRequest, + ) -> Result { + let effects = self.node(&effect.node)?; + let mut effects = lock_node(&effects); + if let Some(live) = &effects.live { + if live.attempt == effect.operation.attempt { + return Ok(OperationOutcome::LeaseCreated(CreateLeaseResult { + lease: live.lease.clone(), + endpoint: Some(live.endpoint.clone()), + })); + } + return Err(EffectError::ambiguous(format!( + "node {} still owns attempt {} while creating attempt {}", + effect.node.0, live.attempt.0, effect.operation.attempt.0 + ))); + } + + if let Some(staged) = effects.staged.take() { + effects.template = staged.template; + effects.plugin = staged.plugin; + } + + let mut spec = effects.template.clone(); + spec.run_id = request.spec.run_id.0; + spec.attempt_id = effect.operation.attempt.0; + spec.image = request.spec.shape.image.clone(); + spec.env = request.spec.boot.env.clone(); + spec.args = request.spec.boot.args.clone(); + spec.mounts = request.spec.boot.mounts.clone(); + spec.env.retain(|(name, _)| name != ATTEMPT_ENV); + spec.env + .push((ATTEMPT_ENV.to_owned(), spec.attempt_id.to_string())); + let failure_sink = Arc::new(AttemptObservationSink { + downstream: self.sink.clone(), + failures: self.failure_tx.clone(), + node: effect.node.clone(), + attempt: effect.operation.attempt, + gate: Mutex::new(FailureGate::Pending(Vec::new())), + }); + let attempt_sink = PluginSink::new(failure_sink.clone()); + let handle = match effects.plugin.create_node(spec, attempt_sink) { + Ok(handle) => handle, + Err(error) => { + failure_sink.discard(); + return Err(EffectError::ambiguous(error)); + } + }; + + let provider = request.spec.provider.clone(); + let lease_id = ProviderLeaseId(format!( + "run-{}-node-{}-attempt-{}", + request.spec.run_id.0, effect.node.0, effect.operation.attempt.0 + )); + let provider_contract_id = lease_id.0.clone(); + let lease = LeaseFacts { + provider: provider.clone(), + lease_id: lease_id.clone(), + provider_contract_id: provider_contract_id.clone(), + offer_id: None, + destroy_handle: DestroyHandle { + provider, + lease_id, + provider_contract_id, + }, + provider_metadata: BTreeMap::from([ + ("logical_node_id".to_owned(), effect.node.0.clone()), + ("attempt".to_owned(), effect.operation.attempt.0.to_string()), + ]), + }; + let endpoint = SshEndpoint { + host: "managed-by-myelin-plugin".to_owned(), + port: 0, + user: request.spec.boot.ssh_user.clone(), + auth_ref: "managed-by-myelin-plugin".to_owned(), + }; + effects.live = Some(LiveNode { + attempt: effect.operation.attempt, + handle, + lease: lease.clone(), + endpoint: endpoint.clone(), + failure_sink: Arc::clone(&failure_sink), + bootstrap_started: false, + }); + failure_sink.arm(); + Ok(OperationOutcome::LeaseCreated(CreateLeaseResult { + lease, + endpoint: Some(endpoint), + })) + } +} + +impl EffectBackend for MyelinEffectBackend { + fn execute(&self, effect: &PlannedEffect) -> Result { + match &effect.command { + NodeManagerCommand::CreateLease(request) => self.create_or_adopt(effect, request), + NodeManagerCommand::LookupEndpoint(_) => { + let node = self.node(&effect.node)?; + let effects = lock_node(&node); + let endpoint = effects + .live + .as_ref() + .filter(|live| live.attempt == effect.operation.attempt) + .map(|live| live.endpoint.clone()); + Ok(OperationOutcome::EndpointLookup(endpoint)) + } + NodeManagerCommand::StartBootstrap(_) => { + let effects = self.node(&effect.node)?; + let mut effects = lock_node(&effects); + let (handle, bootstrap_started) = { + let live = effects.live.as_ref().ok_or_else(|| { + EffectError::definite(format!("node {} has no live lease", effect.node.0)) + })?; + if live.attempt != effect.operation.attempt { + return Err(EffectError::definite( + "bootstrap attempt does not own lease", + )); + } + (live.handle.clone(), live.bootstrap_started) + }; + if !bootstrap_started { + effects + .plugin + .start_bootstrap(&handle) + .map_err(EffectError::definite)?; + let live = effects + .live + .as_mut() + .expect("live lease retained while bootstrap starts"); + live.bootstrap_started = true; + } + Ok(OperationOutcome::BootstrapStarted { + session_id: BootstrapSessionId(effect.operation.attempt.0), + }) + } + NodeManagerCommand::BootstrapConvergenceObserved { session_id, .. } => { + if *session_id != BootstrapSessionId(effect.operation.attempt.0) { + return Err(EffectError::definite("bootstrap session/attempt mismatch")); + } + let effects = self.node(&effect.node)?; + let mut effects = lock_node(&effects); + let handle = effects + .live + .as_ref() + .filter(|live| live.attempt == effect.operation.attempt) + .map(|live| live.handle.clone()) + .ok_or_else(|| EffectError::definite("bootstrap lease is absent"))?; + effects + .plugin + .complete_bootstrap(&handle) + .map_err(EffectError::definite)?; + if let Some(live) = effects.live.as_ref() { + live.failure_sink.discard(); + } + Ok(OperationOutcome::BootstrapConvergenceAccepted) + } + NodeManagerCommand::CancelBootstrap { session_id } => { + if *session_id != BootstrapSessionId(effect.operation.attempt.0) { + return Err(EffectError::definite("bootstrap session/attempt mismatch")); + } + let effects = self.node(&effect.node)?; + let mut effects = lock_node(&effects); + if let Some(handle) = effects + .live + .as_ref() + .filter(|live| live.attempt == effect.operation.attempt) + .map(|live| live.handle.clone()) + { + effects + .plugin + .cancel_bootstrap(&handle) + .map_err(EffectError::definite)?; + if let Some(live) = effects.live.as_ref() { + live.failure_sink.discard(); + } + } + Ok(OperationOutcome::BootstrapCancelled) + } + NodeManagerCommand::DestroyLease(_) => { + let effects = self.node(&effect.node)?; + let mut effects = lock_node(&effects); + let Some(live) = effects.live.take() else { + return Ok(OperationOutcome::LeaseDestroyed); + }; + if live.attempt != effect.operation.attempt { + effects.live = Some(live); + return Err(EffectError::definite("destroy attempt does not own lease")); + } + live.failure_sink.discard(); + if let Err(error) = effects.plugin.stop_node(&live.handle) { + effects.live = Some(live); + return Err(EffectError::definite(error)); + } + Ok(OperationOutcome::LeaseDestroyed) + } + } + } +} + +#[derive(Clone)] +pub(crate) struct EngineEffectSpawner { + engine: EngineHandle, +} + +impl EngineEffectSpawner { + fn new(engine: EngineHandle) -> Self { + Self { engine } + } +} + +impl BlockingEffectSpawner for EngineEffectSpawner { + type SpawnError = String; + + fn spawn_blocking( + &self, + work: provisioning::BlockingEffectWork, + ) -> Result<(), Self::SpawnError> { + if !self.engine.capabilities().blocking { + return Err("engine blocking work capability is unavailable".to_owned()); + } + self.engine.spawn_blocking(work); + Ok(()) + } +} + +enum ControllerWake { + Deadline(SystemTime), + Periodic, +} + +pub(crate) struct ProvisionedClusterGuard { + driver: ClusterDriver, + executor: IdempotentEffectExecutor, + external_nodes: BTreeMap, + failure_rx: Receiver, + deferred_failures: VecDeque, + engine: EngineHandle, + wake_tx: Sender, + wake_rx: Receiver, + scheduled_deadline: Option, + stopped: bool, +} + +impl ProvisionedClusterGuard { + pub(crate) fn new( + desired: ClusterShape, + bindings: Vec, + retry: RetryPolicy, + engine: EngineHandle, + sink: PluginSink, + ) -> Result { + let expanded = desired.expand().map_err(|error| error.to_string())?; + let binding_ids = bindings + .iter() + .map(|binding| binding.logical_node_id.clone()) + .collect::>(); + for logical_id in expanded.keys() { + if !binding_ids.contains(logical_id) { + return Err(format!( + "no effect binding for desired node {}", + logical_id.0 + )); + } + } + for logical_id in &binding_ids { + if !expanded.contains_key(logical_id) { + return Err(format!( + "binding {} is absent from desired shape", + logical_id.0 + )); + } + } + let (failure_tx, failure_rx) = mpsc::channel(); + let (backend, external_nodes) = MyelinEffectBackend::new(bindings, sink, failure_tx)?; + let executor = + IdempotentEffectExecutor::new(backend, EngineEffectSpawner::new(engine.clone())); + let driver = ClusterDriver::new(desired, retry).map_err(|error| error.to_string())?; + let (wake_tx, wake_rx) = mpsc::channel(); + spawn_periodic_wake(&engine, wake_tx.clone()); + Ok(Self { + driver, + executor, + external_nodes, + failure_rx, + deferred_failures: VecDeque::new(), + engine, + wake_tx, + wake_rx, + scheduled_deadline: None, + stopped: false, + }) + } + + pub(crate) fn update_desired( + &mut self, + desired: ClusterShape, + bindings: Vec, + ) -> Result<(), String> { + let expanded = desired.expand().map_err(|error| error.to_string())?; + let mut new_external_ids = BTreeSet::new(); + let mut new_logical_ids = BTreeSet::new(); + for binding in &bindings { + if !expanded.contains_key(&binding.logical_node_id) { + return Err(format!( + "binding {} is absent from desired shape", + binding.logical_node_id.0 + )); + } + if self + .external_nodes + .get(&binding.provision.node_id) + .is_some_and(|logical_id| logical_id != &binding.logical_node_id) + || !new_external_ids.insert(binding.provision.node_id) + { + return Err(format!( + "duplicate provision node id {}", + binding.provision.node_id + )); + } + if !new_logical_ids.insert(binding.logical_node_id.clone()) { + return Err(format!( + "duplicate reconciler binding {}", + binding.logical_node_id.0 + )); + } + } + for (logical_id, desired_node) in &expanded { + let needs_binding = match self.driver.state().nodes.get(logical_id) { + Some(current) => current.record.desired != *desired_node, + None => true, + } || !self.executor.backend().has_node(logical_id); + if needs_binding && !new_logical_ids.contains(logical_id) { + return Err(format!( + "desired node {} requires a replacement effect binding", + logical_id.0 + )); + } + } + self.driver + .update_desired(desired) + .map_err(|error| error.to_string())?; + for binding in bindings { + let (external_id, logical_id) = self.executor.backend().register_binding(binding); + self.external_nodes + .retain(|_, existing| existing != &logical_id); + self.external_nodes.insert(external_id, logical_id); + } + Ok(()) + } + + pub(crate) fn is_converged(&self) -> bool { + self.driver.is_converged() + } + + pub(crate) fn current_attempt(&self, external_node_id: u64) -> Option { + let logical = self.external_nodes.get(&external_node_id)?; + self.driver + .state() + .nodes + .get(logical) + .map(|node| node.attempt) + } + + pub(crate) fn awaiting_runtime(&self) -> bool { + !self.driver.state().nodes.is_empty() + && self.driver.state().nodes.values().all(|node| { + (node.intent == NodeIntent::Active && node.record.ready && node.pending.is_none()) + || (node.intent == NodeIntent::Active + && node.record.stage == NodeStage::BootstrapRunning + && node.active_bootstrap.is_some() + && node.pending.is_none()) + }) + } + + pub(crate) fn observe_runtime_ready( + &mut self, + external_node_id: u64, + attempt: NodeAttemptId, + swactor_id: SwactorId, + now: SystemTime, + ) -> bool { + let Some(logical) = self.external_nodes.get(&external_node_id) else { + return false; + }; + self.driver.apply_observation( + logical, + attempt, + NodeObservation::SwactorJoined { + session_id: BootstrapSessionId(attempt.0), + swactor_id, + }, + now, + ) + } + + pub(crate) fn poll(&mut self, now: SystemTime) -> Result { + self.drain_wakes(now); + self.drain_executor_results(now); + self.drain_failures(now); + if self.classify_due_operations(now) { + self.drain_executor_results(now); + } + self.driver.trigger_if_due(now); + let submitted = self.driver.drive_until_blocked(now, &mut self.executor)?; + self.schedule_deadline(now); + Ok(submitted) + } + + fn begin_shutdown(&mut self) -> Result<(), String> { + if self.stopped { + return Ok(()); + } + let desired = ClusterShape { + run_id: self.driver.desired().run_id.clone(), + generation: self.driver.desired().generation.saturating_add(1), + groups: Vec::new(), + }; + self.update_desired(desired, Vec::new()) + } + + pub(crate) fn is_stopped(&self) -> bool { + self.driver.state().nodes.is_empty() + } + + // Synchronous orchestration waits while all provider work remains engine-hosted. + #[allow(clippy::disallowed_methods)] + pub(crate) fn stop(&mut self) -> Result<(), String> { + self.begin_shutdown()?; + while !self.is_stopped() { + self.poll(SystemTime::now()) + .map_err(|error| error.to_string())?; + std::thread::sleep(Duration::from_millis(10)); + } + self.executor.backend().stop_all()?; + self.stopped = true; + Ok(()) + } + + fn drain_wakes(&mut self, now: SystemTime) { + loop { + match self.wake_rx.try_recv() { + Ok(ControllerWake::Periodic) => self.driver.trigger(), + Ok(ControllerWake::Deadline(deadline)) => { + if self.scheduled_deadline == Some(deadline) { + self.scheduled_deadline = None; + } + self.driver.trigger_if_due(now); + } + Err(TryRecvError::Empty | TryRecvError::Disconnected) => return, + } + } + } + + fn drain_executor_results(&mut self, now: SystemTime) { + for result in self.executor.drain_results() { + let close = matches!( + result.result, + Ok(OperationOutcome::BootstrapConvergenceAccepted) + ) + .then_some((result.node.clone(), result.operation.attempt)); + if self.driver.apply_executor_result(result, now) + && let Some((node, attempt)) = close + { + self.driver.apply_observation( + &node, + attempt, + NodeObservation::BootstrapClosed { + session_id: BootstrapSessionId(attempt.0), + }, + now, + ); + } + } + } + + fn drain_failures(&mut self, now: SystemTime) { + while let Ok(failure) = self.failure_rx.try_recv() { + self.deferred_failures.push_back(failure); + } + let mut remaining = VecDeque::new(); + while let Some(failure) = self.deferred_failures.pop_front() { + let Some(node) = self.driver.state().nodes.get(&failure.node) else { + continue; + }; + if node.attempt != failure.attempt || node.intent == NodeIntent::Deleting { + continue; + } + let Some(session_id) = node.active_bootstrap else { + if node.pending.is_some() + || matches!( + node.record.stage, + NodeStage::LeaseCreated + | NodeStage::EndpointKnown + | NodeStage::BootstrapRunning + ) + { + remaining.push_back(failure); + } + continue; + }; + self.driver.apply_observation( + &failure.node, + failure.attempt, + NodeObservation::BootstrapFailed { + session_id, + reason: failure.reason, + }, + now, + ); + } + self.deferred_failures = remaining; + } + + fn classify_due_operations(&mut self, now: SystemTime) -> bool { + let mut expired = false; + for operation in self.driver.pending_operations_due(now) { + match self.executor.operation_status(operation.operation) { + ExecutorOperationStatus::Unknown => { + self.driver.operation_timed_out( + &operation, + "executor lost pending operation", + now, + ); + } + ExecutorOperationStatus::InFlight => { + expired |= self.executor.expire( + operation.operation, + "executor operation timed out with an ambiguous outcome", + ); + } + ExecutorOperationStatus::Completed => {} + } + } + expired + } + + fn schedule_deadline(&mut self, now: SystemTime) { + let deadline = self.driver.requeue_at(); + if deadline.is_some_and(|deadline| deadline <= now) + && self + .driver + .pending_operations_due(now) + .iter() + .any(|operation| { + self.executor.operation_status(operation.operation) + == ExecutorOperationStatus::InFlight + }) + { + self.scheduled_deadline = None; + return; + } + if deadline.is_none() || deadline == self.scheduled_deadline { + return; + } + let deadline = deadline.expect("checked deadline"); + self.scheduled_deadline = Some(deadline); + let delay = deadline.duration_since(now).unwrap_or(Duration::ZERO); + let timer = self.engine.timer(delay); + let wake = self.wake_tx.clone(); + self.engine.spawn(async move { + timer.await; + let _ = wake.send(ControllerWake::Deadline(deadline)); + }); + } +} + +impl Drop for ProvisionedClusterGuard { + fn drop(&mut self) { + if !self.stopped { + let _ = self.executor.backend().stop_all(); + } + } +} + +fn spawn_periodic_wake(engine: &EngineHandle, wake: Sender) { + let mut interval = engine.interval(PERIODIC_RECONCILE); + engine.spawn(async move { + loop { + (&mut interval).await; + if wake.send(ControllerWake::Periodic).is_err() { + return; + } + } + }); +} + +fn lock_node(node: &Mutex) -> MutexGuard<'_, NodeEffects> { + node.lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) +} + +fn lock_nodes_read( + nodes: &RwLock>>>, +) -> RwLockReadGuard<'_, BTreeMap>>> { + nodes + .read() + .unwrap_or_else(std::sync::PoisonError::into_inner) +} + +fn lock_nodes_write( + nodes: &RwLock>>>, +) -> RwLockWriteGuard<'_, BTreeMap>>> { + nodes + .write() + .unwrap_or_else(std::sync::PoisonError::into_inner) +} + +#[cfg(test)] +mod tests { + use std::sync::atomic::{AtomicUsize, Ordering}; + + use crate::provisioning::ProviderMount; + use ::provisioning::{ + BootSpec, BootstrapSessionSpec, ClusterShape, CreateLeaseRequest, DatastreamStreamId, + DesiredNodeShape, LogicalNodeSpec, NodeGroupId, OperationId, ProviderKind, RetryPolicy, + RoleId, RunId, RunNodeGroupSpec, SwarmJoinSpec, SwarmJoinTemplate, + }; + + use super::*; + + #[derive(Default)] + struct NullSink; + + impl PluginObservationSink for NullSink { + fn observe(&self, _observation: PluginObservation) {} + } + + #[test] + fn attempt_failures_are_published_only_while_the_start_is_live() { + let (failure_tx, failure_rx) = mpsc::channel(); + let sink = AttemptObservationSink { + downstream: PluginSink::new(Arc::new(NullSink)), + failures: failure_tx, + node: LogicalNodeId("node-7-0".to_owned()), + attempt: NodeAttemptId(9), + gate: Mutex::new(FailureGate::Pending(Vec::new())), + }; + let failed = |reason: &str| PluginObservation::Failed { + run_id: 5, + node_id: 7, + reason: reason.to_owned(), + }; + + sink.observe(failed("before start returned")); + assert!(matches!(failure_rx.try_recv(), Err(TryRecvError::Empty))); + sink.arm(); + assert_eq!(failure_rx.recv().unwrap().reason, "before start returned"); + sink.observe(failed("while live")); + assert_eq!(failure_rx.recv().unwrap().reason, "while live"); + sink.discard(); + sink.observe(failed("after cleanup")); + assert!(matches!(failure_rx.try_recv(), Err(TryRecvError::Empty))); + } + + #[derive(Default)] + struct PluginStats { + creates: AtomicUsize, + starts: AtomicUsize, + start_failures: AtomicUsize, + completes: AtomicUsize, + cancels: AtomicUsize, + stops: AtomicUsize, + stop_failures: AtomicUsize, + specs: Mutex>, + } + + struct FakePlugin { + stats: Arc, + } + + impl ProvisionPlugin for FakePlugin { + fn create_node( + &mut self, + spec: NodeProvisionSpec, + sink: PluginSink, + ) -> Result { + self.stats.creates.fetch_add(1, Ordering::SeqCst); + self.stats.specs.lock().unwrap().push(spec.clone()); + sink.observe(PluginObservation::ProviderLine { + run_id: spec.run_id, + node_id: spec.node_id, + line: "created".to_owned(), + }); + Ok(PluginNodeHandle { + id: spec.attempt_id, + provider_process_id: None, + }) + } + + fn start_bootstrap(&mut self, _handle: &PluginNodeHandle) -> Result<(), String> { + self.stats.starts.fetch_add(1, Ordering::SeqCst); + if self + .stats + .start_failures + .fetch_update(Ordering::SeqCst, Ordering::SeqCst, |remaining| { + (remaining > 0).then(|| remaining - 1) + }) + .is_ok() + { + Err("bootstrap start failed".to_owned()) + } else { + Ok(()) + } + } + + fn cancel_bootstrap(&mut self, _handle: &PluginNodeHandle) -> Result<(), String> { + self.stats.cancels.fetch_add(1, Ordering::SeqCst); + Ok(()) + } + + fn complete_bootstrap(&mut self, _handle: &PluginNodeHandle) -> Result<(), String> { + self.stats.completes.fetch_add(1, Ordering::SeqCst); + Ok(()) + } + + fn stop_node(&mut self, _handle: &PluginNodeHandle) -> Result<(), String> { + self.stats.stops.fetch_add(1, Ordering::SeqCst); + if self + .stats + .stop_failures + .fetch_update(Ordering::SeqCst, Ordering::SeqCst, |remaining| { + (remaining > 0).then(|| remaining - 1) + }) + .is_ok() + { + Err("node cleanup failed".to_owned()) + } else { + Ok(()) + } + } + } + + struct BlockingCreatePlugin { + stats: Arc, + entered: Sender<()>, + release: Receiver<()>, + } + + impl ProvisionPlugin for BlockingCreatePlugin { + fn create_node( + &mut self, + spec: NodeProvisionSpec, + _sink: PluginSink, + ) -> Result { + self.stats.creates.fetch_add(1, Ordering::SeqCst); + let _ = self.entered.send(()); + self.release + .recv() + .map_err(|_| "blocking test release closed".to_owned())?; + Ok(PluginNodeHandle { + id: spec.attempt_id, + provider_process_id: None, + }) + } + + fn start_bootstrap(&mut self, _handle: &PluginNodeHandle) -> Result<(), String> { + self.stats.starts.fetch_add(1, Ordering::SeqCst); + Ok(()) + } + + fn cancel_bootstrap(&mut self, _handle: &PluginNodeHandle) -> Result<(), String> { + self.stats.cancels.fetch_add(1, Ordering::SeqCst); + Ok(()) + } + + fn complete_bootstrap(&mut self, _handle: &PluginNodeHandle) -> Result<(), String> { + self.stats.completes.fetch_add(1, Ordering::SeqCst); + Ok(()) + } + + fn stop_node(&mut self, _handle: &PluginNodeHandle) -> Result<(), String> { + self.stats.stops.fetch_add(1, Ordering::SeqCst); + Ok(()) + } + } + + fn desired() -> LogicalNodeSpec { + let node = LogicalNodeId("node-7-0".to_owned()); + LogicalNodeSpec { + run_id: RunId(5), + logical_node_id: node.clone(), + group_id: NodeGroupId("node-7".to_owned()), + role: RoleId("worker".to_owned()), + provider: ProviderKind::new("fake"), + shape: DesiredNodeShape { + image: "node:v2".to_owned(), + disk_gb: 1, + gpu_name: None, + min_gpu_ram_mb: None, + min_down_mbps: None, + min_up_mbps: None, + min_reliability: None, + require_verified: false, + provider_labels: BTreeMap::new(), + }, + boot: BootSpec { + ssh_user: "root".to_owned(), + verify_commands: Vec::new(), + start_swactor_command: "swactor".to_owned(), + stdout_sources: Vec::new(), + stderr_sources: Vec::new(), + env: vec![("MYELIN_RUN_ID".to_owned(), "5".to_owned())], + args: vec!["--worker".to_owned()], + mounts: vec![ProviderMount { + host_path: "/model".to_owned(), + container_path: "/model".to_owned(), + readonly: true, + }], + }, + swarm_join: SwarmJoinSpec { + orch_swactor_addr: "orchestrator".to_owned(), + join_token_ref: "token".to_owned(), + expected_logical_node_id: node, + }, + } + } + + fn effect(sequence: u64, command: NodeManagerCommand) -> PlannedEffect { + PlannedEffect { + node: LogicalNodeId("node-7-0".to_owned()), + operation: OperationId { + attempt: NodeAttemptId(9), + sequence, + }, + command, + } + } + + #[test] + fn adapter_creates_and_starts_once_and_adopts_repeated_effects() { + let stats = Arc::new(PluginStats::default()); + let binding = ReconcilerNodeBinding { + logical_node_id: LogicalNodeId("node-7-0".to_owned()), + provision: NodeProvisionSpec { + run_id: 5, + node_id: 7, + attempt_id: 0, + stage_index: Some(0), + image: "node:v1".to_owned(), + env: Vec::new(), + args: Vec::new(), + mounts: Vec::new(), + }, + plugin: Box::new(FakePlugin { + stats: Arc::clone(&stats), + }), + }; + let downstream = PluginSink::new(Arc::new(NullSink)); + let (failure_tx, _failure_rx) = mpsc::channel(); + let (backend, _) = MyelinEffectBackend::new(vec![binding], downstream, failure_tx).unwrap(); + let desired = desired(); + + let create = effect( + 1, + NodeManagerCommand::CreateLease(CreateLeaseRequest { + spec: desired.clone(), + }), + ); + let created = backend.execute(&create).unwrap(); + assert!(matches!(created, OperationOutcome::LeaseCreated(_))); + let adopted = backend.execute(&create).unwrap(); + assert_eq!(adopted, created); + assert_eq!(stats.creates.load(Ordering::SeqCst), 1); + assert_eq!(stats.starts.load(Ordering::SeqCst), 0); + + let start = effect( + 2, + NodeManagerCommand::StartBootstrap(BootstrapSessionSpec { + run_id: desired.run_id.clone(), + logical_node_id: desired.logical_node_id.clone(), + lease_id: ProviderLeaseId("lease".to_owned()), + ssh: SshEndpoint { + host: "host".to_owned(), + port: 22, + user: "root".to_owned(), + auth_ref: "key".to_owned(), + }, + boot: desired.boot.clone(), + swarm_join: desired.swarm_join.clone(), + datastream: DatastreamStreamId("bootstrap".to_owned()), + }), + ); + assert!(matches!( + backend.execute(&start).unwrap(), + OperationOutcome::BootstrapStarted { + session_id: BootstrapSessionId(9) + } + )); + assert!(matches!( + backend.execute(&start).unwrap(), + OperationOutcome::BootstrapStarted { + session_id: BootstrapSessionId(9) + } + )); + assert_eq!(stats.starts.load(Ordering::SeqCst), 1); + let launched = stats.specs.lock().unwrap(); + assert_eq!(launched[0].attempt_id, 9); + assert_eq!(launched[0].image, "node:v2"); + assert!( + launched[0] + .env + .contains(&(ATTEMPT_ENV.to_owned(), "9".to_owned())) + ); + drop(launched); + + let convergence = effect( + 3, + NodeManagerCommand::BootstrapConvergenceObserved { + session_id: BootstrapSessionId(9), + swactor_id: SwactorId("joined".to_owned()), + }, + ); + backend.execute(&convergence).unwrap(); + assert_eq!(stats.completes.load(Ordering::SeqCst), 1); + + let cancel = effect( + 4, + NodeManagerCommand::CancelBootstrap { + session_id: BootstrapSessionId(9), + }, + ); + backend.execute(&cancel).unwrap(); + assert_eq!(stats.cancels.load(Ordering::SeqCst), 1); + + let destroy = effect( + 5, + NodeManagerCommand::DestroyLease(match created { + OperationOutcome::LeaseCreated(result) => result.lease.destroy_handle, + _ => unreachable!(), + }), + ); + backend.execute(&destroy).unwrap(); + backend.execute(&destroy).unwrap(); + assert_eq!(stats.stops.load(Ordering::SeqCst), 1); + } + + #[test] + fn returned_bootstrap_start_failure_is_definite_and_phase_correlated() { + let stats = Arc::new(PluginStats::default()); + stats.start_failures.store(1, Ordering::SeqCst); + let binding = ReconcilerNodeBinding { + logical_node_id: LogicalNodeId("node-7-0".to_owned()), + provision: NodeProvisionSpec { + run_id: 5, + node_id: 7, + attempt_id: 0, + stage_index: Some(0), + image: "node:v1".to_owned(), + env: Vec::new(), + args: Vec::new(), + mounts: Vec::new(), + }, + plugin: Box::new(FakePlugin { + stats: Arc::clone(&stats), + }), + }; + let (failure_tx, _failure_rx) = mpsc::channel(); + let (backend, _) = MyelinEffectBackend::new( + vec![binding], + PluginSink::new(Arc::new(NullSink)), + failure_tx, + ) + .unwrap(); + let desired = desired(); + let create = effect( + 1, + NodeManagerCommand::CreateLease(CreateLeaseRequest { + spec: desired.clone(), + }), + ); + backend.execute(&create).unwrap(); + let start = effect( + 2, + NodeManagerCommand::StartBootstrap(BootstrapSessionSpec { + run_id: desired.run_id.clone(), + logical_node_id: desired.logical_node_id.clone(), + lease_id: ProviderLeaseId("lease".to_owned()), + ssh: SshEndpoint { + host: "host".to_owned(), + port: 22, + user: "root".to_owned(), + auth_ref: "key".to_owned(), + }, + boot: desired.boot, + swarm_join: desired.swarm_join, + datastream: DatastreamStreamId("bootstrap".to_owned()), + }), + ); + + let error = backend.execute(&start).unwrap_err(); + assert_eq!( + error.disposition, + ::provisioning::EffectFailureDisposition::Definite + ); + assert_eq!(error.reason, "bootstrap start failed"); + assert_eq!(stats.creates.load(Ordering::SeqCst), 1); + assert_eq!(stats.starts.load(Ordering::SeqCst), 1); + } + + #[test] + fn replacement_binding_activates_only_after_old_lease_cleanup() { + let old_stats = Arc::new(PluginStats::default()); + let new_stats = Arc::new(PluginStats::default()); + let binding = |stats: &Arc, node_id| ReconcilerNodeBinding { + logical_node_id: LogicalNodeId("node-7-0".to_owned()), + provision: NodeProvisionSpec { + run_id: 5, + node_id, + attempt_id: 0, + stage_index: Some(0), + image: "node:v1".to_owned(), + env: Vec::new(), + args: Vec::new(), + mounts: Vec::new(), + }, + plugin: Box::new(FakePlugin { + stats: Arc::clone(stats), + }), + }; + let (failure_tx, _failure_rx) = mpsc::channel(); + let (backend, _) = MyelinEffectBackend::new( + vec![binding(&old_stats, 7)], + PluginSink::new(Arc::new(NullSink)), + failure_tx, + ) + .unwrap(); + let request = CreateLeaseRequest { spec: desired() }; + let create = effect(1, NodeManagerCommand::CreateLease(request.clone())); + let created = backend.execute(&create).unwrap(); + backend.register_binding(binding(&new_stats, 8)); + + assert_eq!(backend.execute(&create).unwrap(), created); + assert_eq!(old_stats.creates.load(Ordering::SeqCst), 1); + assert_eq!(new_stats.creates.load(Ordering::SeqCst), 0); + + let destroy = effect( + 2, + NodeManagerCommand::DestroyLease(match created { + OperationOutcome::LeaseCreated(result) => result.lease.destroy_handle, + _ => unreachable!(), + }), + ); + backend.execute(&destroy).unwrap(); + let mut replacement = effect(3, NodeManagerCommand::CreateLease(request)); + replacement.operation.attempt = NodeAttemptId(10); + backend.execute(&replacement).unwrap(); + + assert_eq!(old_stats.stops.load(Ordering::SeqCst), 1); + assert_eq!(old_stats.creates.load(Ordering::SeqCst), 1); + assert_eq!(new_stats.creates.load(Ordering::SeqCst), 1); + } + + #[test] + fn failed_node_cleanup_keeps_lease_owned_for_retry() { + let stats = Arc::new(PluginStats::default()); + stats.stop_failures.store(2, Ordering::SeqCst); + let binding = ReconcilerNodeBinding { + logical_node_id: LogicalNodeId("node-7-0".to_owned()), + provision: NodeProvisionSpec { + run_id: 5, + node_id: 7, + attempt_id: 0, + stage_index: Some(0), + image: "node:v1".to_owned(), + env: Vec::new(), + args: Vec::new(), + mounts: Vec::new(), + }, + plugin: Box::new(FakePlugin { + stats: Arc::clone(&stats), + }), + }; + let (failure_tx, _failure_rx) = mpsc::channel(); + let (backend, _) = MyelinEffectBackend::new( + vec![binding], + PluginSink::new(Arc::new(NullSink)), + failure_tx, + ) + .unwrap(); + let create = effect( + 1, + NodeManagerCommand::CreateLease(CreateLeaseRequest { spec: desired() }), + ); + let created = backend.execute(&create).unwrap(); + let destroy = effect( + 2, + NodeManagerCommand::DestroyLease(match created { + OperationOutcome::LeaseCreated(result) => result.lease.destroy_handle, + _ => unreachable!(), + }), + ); + + let error = backend.execute(&destroy).unwrap_err(); + assert_eq!( + error.disposition, + ::provisioning::EffectFailureDisposition::Definite + ); + assert_eq!(error.reason, "node cleanup failed"); + assert_eq!(backend.stop_all().unwrap_err(), "node cleanup failed"); + backend.stop_all().unwrap(); + assert_eq!(stats.stops.load(Ordering::SeqCst), 3); + } + + #[test] + #[allow(clippy::disallowed_methods)] + fn engine_hosted_controller_converges_and_cleans_up_end_to_end() { + let stats = Arc::new(PluginStats::default()); + let desired_node = desired(); + let shape = ClusterShape { + run_id: desired_node.run_id.clone(), + generation: 1, + groups: vec![RunNodeGroupSpec { + run_id: desired_node.run_id.clone(), + group_id: desired_node.group_id.clone(), + role: desired_node.role.clone(), + count: 1, + provider: desired_node.provider.clone(), + shape: desired_node.shape.clone(), + boot: desired_node.boot.clone(), + swarm_join: SwarmJoinTemplate { + orch_swactor_addr: desired_node.swarm_join.orch_swactor_addr.clone(), + join_token_ref: desired_node.swarm_join.join_token_ref.clone(), + }, + }], + }; + let binding = ReconcilerNodeBinding { + logical_node_id: desired_node.logical_node_id, + provision: NodeProvisionSpec { + run_id: 5, + node_id: 7, + attempt_id: 0, + stage_index: Some(0), + image: "node:v1".to_owned(), + env: Vec::new(), + args: Vec::new(), + mounts: Vec::new(), + }, + plugin: Box::new(FakePlugin { + stats: Arc::clone(&stats), + }), + }; + let parts = swactor::runtime::RuntimeParts::new(swactor::config::RuntimeConfig::default()); + let backend = + swactor_engine::TokioBackend::new(swactor_engine::TokioConfig::default()).unwrap(); + let engine = swactor_engine::Engine::new(parts, backend).unwrap(); + let sink = PluginSink::new(Arc::new(NullSink)); + let mut cluster = ProvisionedClusterGuard::new( + shape, + vec![binding], + RetryPolicy::default(), + engine.handle(), + sink, + ) + .unwrap(); + + for _ in 0..200 { + cluster.poll(SystemTime::now()).unwrap(); + if cluster.awaiting_runtime() { + break; + } + std::thread::sleep(Duration::from_millis(1)); + } + assert!(cluster.awaiting_runtime()); + let attempt = cluster.current_attempt(7).unwrap(); + assert!(cluster.observe_runtime_ready( + 7, + attempt, + SwactorId("worker-7".to_owned()), + SystemTime::now(), + )); + + for _ in 0..200 { + cluster.poll(SystemTime::now()).unwrap(); + if cluster.is_converged() { + break; + } + std::thread::sleep(Duration::from_millis(1)); + } + assert!(cluster.is_converged()); + assert_eq!(stats.starts.load(Ordering::SeqCst), 1); + assert_eq!(stats.creates.load(Ordering::SeqCst), 1); + assert_eq!(stats.completes.load(Ordering::SeqCst), 1); + let scaled_stats = Arc::new(PluginStats::default()); + let scaled = desired(); + cluster + .update_desired( + ClusterShape { + run_id: scaled.run_id.clone(), + generation: 2, + groups: vec![RunNodeGroupSpec { + run_id: scaled.run_id, + group_id: scaled.group_id, + role: scaled.role, + count: 2, + provider: scaled.provider, + shape: scaled.shape, + boot: scaled.boot, + swarm_join: SwarmJoinTemplate { + orch_swactor_addr: scaled.swarm_join.orch_swactor_addr, + join_token_ref: scaled.swarm_join.join_token_ref, + }, + }], + }, + vec![ReconcilerNodeBinding { + logical_node_id: LogicalNodeId("node-7-1".to_owned()), + provision: NodeProvisionSpec { + run_id: 5, + node_id: 8, + attempt_id: 0, + stage_index: Some(1), + image: "node:v1".to_owned(), + env: Vec::new(), + args: Vec::new(), + mounts: Vec::new(), + }, + plugin: Box::new(FakePlugin { + stats: Arc::clone(&scaled_stats), + }), + }], + ) + .unwrap(); + for _ in 0..200 { + cluster.poll(SystemTime::now()).unwrap(); + if cluster.awaiting_runtime() && cluster.current_attempt(8).is_some() { + break; + } + std::thread::sleep(Duration::from_millis(1)); + } + let scaled_attempt = cluster.current_attempt(8).unwrap(); + assert!(cluster.observe_runtime_ready( + 8, + scaled_attempt, + SwactorId("worker-8".to_owned()), + SystemTime::now(), + )); + for _ in 0..200 { + cluster.poll(SystemTime::now()).unwrap(); + + if cluster.is_converged() { + break; + } + std::thread::sleep(Duration::from_millis(1)); + } + assert!(cluster.is_converged()); + assert_eq!(scaled_stats.starts.load(Ordering::SeqCst), 1); + assert_eq!(scaled_stats.creates.load(Ordering::SeqCst), 1); + assert_eq!(scaled_stats.completes.load(Ordering::SeqCst), 1); + + cluster.stop().unwrap(); + assert!(cluster.is_stopped()); + assert_eq!(stats.stops.load(Ordering::SeqCst), 1); + assert_eq!(scaled_stats.stops.load(Ordering::SeqCst), 1); + } + #[test] + #[allow(clippy::disallowed_methods)] + fn ambiguous_create_timeout_retries_by_adoption_and_discards_late_success() { + let stats = Arc::new(PluginStats::default()); + let (entered_tx, entered_rx) = mpsc::channel(); + let (release_tx, release_rx) = mpsc::channel(); + let desired_node = desired(); + let shape = ClusterShape { + run_id: desired_node.run_id.clone(), + generation: 1, + groups: vec![RunNodeGroupSpec { + run_id: desired_node.run_id.clone(), + group_id: desired_node.group_id.clone(), + role: desired_node.role.clone(), + count: 1, + provider: desired_node.provider.clone(), + shape: desired_node.shape.clone(), + boot: desired_node.boot.clone(), + swarm_join: SwarmJoinTemplate { + orch_swactor_addr: desired_node.swarm_join.orch_swactor_addr.clone(), + join_token_ref: desired_node.swarm_join.join_token_ref.clone(), + }, + }], + }; + let binding = ReconcilerNodeBinding { + logical_node_id: desired_node.logical_node_id, + provision: NodeProvisionSpec { + run_id: 5, + node_id: 7, + attempt_id: 0, + stage_index: Some(0), + image: "node:v1".to_owned(), + env: Vec::new(), + args: Vec::new(), + mounts: Vec::new(), + }, + plugin: Box::new(BlockingCreatePlugin { + stats: Arc::clone(&stats), + entered: entered_tx, + release: release_rx, + }), + }; + let parts = swactor::runtime::RuntimeParts::new(swactor::config::RuntimeConfig::default()); + let backend = + swactor_engine::TokioBackend::new(swactor_engine::TokioConfig::default()).unwrap(); + let engine = swactor_engine::Engine::new(parts, backend).unwrap(); + let mut cluster = ProvisionedClusterGuard::new( + shape, + vec![binding], + RetryPolicy { + operation_timeout: Duration::from_secs(1), + ..RetryPolicy::default() + }, + engine.handle(), + PluginSink::new(Arc::new(NullSink)), + ) + .unwrap(); + + let mut entered = false; + for _ in 0..200 { + cluster.poll(SystemTime::now()).unwrap(); + if entered_rx.try_recv().is_ok() { + entered = true; + break; + } + std::thread::sleep(Duration::from_millis(1)); + } + assert!(entered, "lease creation did not enter the backend"); + + cluster + .poll(SystemTime::now() + Duration::from_secs(2)) + .unwrap(); + let managed = cluster.driver.state().nodes.values().next().unwrap(); + assert_eq!(managed.intent, NodeIntent::Active); + assert!(managed.pending.is_none()); + assert!( + managed + .retry + .last_error + .as_deref() + .is_some_and(|reason| reason.contains("ambiguous outcome")) + ); + + release_tx.send(()).unwrap(); + let retry_now = SystemTime::now() + Duration::from_secs(5); + for _ in 0..200 { + cluster.poll(retry_now).unwrap(); + if cluster.awaiting_runtime() { + break; + } + std::thread::sleep(Duration::from_millis(1)); + } + assert!(cluster.awaiting_runtime()); + assert_eq!(stats.creates.load(Ordering::SeqCst), 1); + assert_eq!(stats.starts.load(Ordering::SeqCst), 1); + + let attempt = cluster.current_attempt(7).unwrap(); + assert!(cluster.observe_runtime_ready( + 7, + attempt, + SwactorId("worker-7".to_owned()), + retry_now, + )); + for _ in 0..200 { + cluster.poll(retry_now).unwrap(); + if cluster.is_converged() { + break; + } + std::thread::sleep(Duration::from_millis(1)); + } + assert!(cluster.is_converged()); + assert_eq!(stats.completes.load(Ordering::SeqCst), 1); + + cluster.stop().unwrap(); + assert_eq!(stats.stops.load(Ordering::SeqCst), 1); + } +} diff --git a/apps/myelin/src/orchestration/distribution_stack.rs b/apps/myelin/src/orchestration/distribution_stack.rs index bfbca19..63ccee6 100644 --- a/apps/myelin/src/orchestration/distribution_stack.rs +++ b/apps/myelin/src/orchestration/distribution_stack.rs @@ -1,4 +1,3 @@ - //! Myelin-system swactor distribution runtime wiring. //! //! This is the production version of the actor-stack setup that integration @@ -16,8 +15,8 @@ use swactor::actor::{ActorAddress, ActorInterface}; use swactor::config::RuntimeConfig; use swactor::runtime::{Ctx, Runtime, RuntimeParts}; use swactor::stats::StatsHook; -use swactor_engine::EngineHandle; use swactor::std::StdExtension; +use swactor_engine::EngineHandle; use swactor_transport::{CodecRegistry, CodecRemoteSink, NetworkMessage, TransportRouter}; use distribution::directory_actor::{DirectoryActor, DirectoryIn}; @@ -77,7 +76,12 @@ impl DistributionRuntimeStack { pub(crate) fn build_runtime( extend_codecs: impl FnOnce(&mut CodecRegistry), stats_hook: Option>, - ) -> (RuntimeParts, Runtime, Arc, Arc) { + ) -> ( + RuntimeParts, + Runtime, + Arc, + Arc, + ) { let mut parts = RuntimeParts::new(RuntimeConfig::default()) .with_extension(Arc::new(StdExtension::new())); let mut codec = actor_codec_registry(); diff --git a/apps/myelin/src/orchestration/engine_builder/mod.rs b/apps/myelin/src/orchestration/engine_builder/mod.rs index f85d250..142437a 100644 --- a/apps/myelin/src/orchestration/engine_builder/mod.rs +++ b/apps/myelin/src/orchestration/engine_builder/mod.rs @@ -1,4 +1,3 @@ - //! Pool-based engine/node builder primitives. //! //! This module owns topology construction: acquire a role-neutral node pool, diff --git a/apps/myelin/src/orchestration/mod.rs b/apps/myelin/src/orchestration/mod.rs index 184fc00..334d3b8 100644 --- a/apps/myelin/src/orchestration/mod.rs +++ b/apps/myelin/src/orchestration/mod.rs @@ -7,6 +7,7 @@ pub(crate) mod actor; pub(crate) mod app; +pub(crate) mod cluster_reconciler; pub(crate) mod config; pub(crate) mod distribution_stack; #[cfg(test)] diff --git a/apps/myelin/src/orchestration/provider_adapters/relay.rs b/apps/myelin/src/orchestration/provider_adapters/relay.rs index de224cc..fefbee1 100644 --- a/apps/myelin/src/orchestration/provider_adapters/relay.rs +++ b/apps/myelin/src/orchestration/provider_adapters/relay.rs @@ -1,4 +1,3 @@ - //! Relay provisioning shims for Myelin runtimes. //! //! The current implementations are deliberately small: local tests get the same @@ -57,7 +56,6 @@ pub(crate) trait RelayProvider: Send { fn relay_mode(&self, lease: &RelayLease) -> Result; } - #[derive(Clone, Debug)] pub(crate) struct StaticRelayProvider { url: RelayUrl, diff --git a/apps/myelin/src/orchestration/provider_adapters/vastai/mod.rs b/apps/myelin/src/orchestration/provider_adapters/vastai/mod.rs index cd55189..0e2356d 100644 --- a/apps/myelin/src/orchestration/provider_adapters/vastai/mod.rs +++ b/apps/myelin/src/orchestration/provider_adapters/vastai/mod.rs @@ -3,12 +3,11 @@ // is driven by an explicit SingleThreadRuntime owned by that thread; the main // orchestration engine owns all bootstrap actors spawned on its runtime handle. #![allow(clippy::disallowed_methods)] -use parking_lot::Mutex; use std::collections::{BTreeMap, BTreeSet, HashSet}; use std::io::{BufRead, BufReader, Read}; use std::path::{Path, PathBuf}; use std::process::{Child, Command, Stdio}; -use std::sync::{Arc, mpsc}; +use std::sync::Arc; use std::sync::atomic::{AtomicBool, Ordering}; use std::thread::{self, JoinHandle}; use std::time::{Duration, Instant}; @@ -20,7 +19,7 @@ use swactor::runtime::{ }; use swactor_vastai::{ CreateInstanceRequest, LifecyclePolicy, Offer, ProvisionRequest, ProvisionedInstance, - SelectionPolicy, classify_vastai_error, create_instance, + SelectionPolicy, classify_vastai_error, }; use crate::observability::provisioning_logs::{BootstrapDatastreamBridge, node_stream_id}; @@ -112,12 +111,6 @@ impl Drop for VastAiProviderMonitor { pub(crate) trait VastAiLeaseClient: Send { fn provision_one(&mut self, request: ProvisionRequest) -> Result; - fn plan_first_wave_offers( - &mut self, - requests: &[ProvisionRequest], - ) -> Result>, String> { - Ok(vec![None; requests.len()]) - } fn ssh_endpoint( &mut self, @@ -143,8 +136,6 @@ pub(crate) trait VastAiLeaseClient: Send { pub(crate) struct ToolsVastAiLeaseClient { client: swactor_vastai::VastClient, runtime: tokio::runtime::Runtime, - planned_offer_pool: Arc>>, - planned_offer_ids: Arc>>, } impl ToolsVastAiLeaseClient { @@ -153,12 +144,7 @@ impl ToolsVastAiLeaseClient { .enable_all() .build() .map_err(|e| format!("vastai tokio runtime: {e}"))?; - Ok(Self { - client, - runtime, - planned_offer_pool: Arc::new(Mutex::new(Vec::new())), - planned_offer_ids: Arc::new(Mutex::new(HashSet::new())), - }) + Ok(Self { client, runtime }) } pub(crate) fn from_api_key(api_key: impl Into) -> Result { @@ -188,13 +174,6 @@ impl ToolsVastAiLeaseClient { } fn candidate_pool(&mut self, request: &ProvisionRequest) -> Result, String> { - let cached = self.planned_offer_pool.lock().clone(); - if request - .preferred_offer_id - .is_some_and(|offer_id| cached.iter().any(|offer| offer.id == offer_id)) - { - return Ok(cached); - } self.runtime .block_on(self.client.search_offers(&request.selection, 1)) } @@ -205,12 +184,9 @@ impl ToolsVastAiLeaseClient { offer: &Offer, ) -> Result { let create = Self::create_request_for_offer(request, offer.id); - let info = self.runtime.block_on(create_instance( - self.client.http(), - self.client.base_url(), - self.client.api_key(), - &create, - ))?; + let info = self + .runtime + .block_on(self.client.create_instance(&create))?; Ok(ProvisionedInstance { index: 0, contract_id: info.contract_id, @@ -221,6 +197,22 @@ impl ToolsVastAiLeaseClient { dph_total: offer.dph_total, }) } + + fn contract_by_label(&mut self, label: &str) -> Result, String> { + let instances = self.runtime.block_on(self.client.list_by_label(label))?; + match instances.as_slice() { + [] => Ok(None), + [instance] => Ok(Some(instance.contract_id)), + _ => Err(format!( + "multiple VastAI contracts share stable label {label}: {}", + instances + .iter() + .map(|instance| instance.contract_id.to_string()) + .collect::>() + .join(",") + )), + } + } } #[derive(Clone)] @@ -433,12 +425,22 @@ impl Clone for ToolsVastAiLeaseClient { .enable_all() .build() .expect("clone VastAI lease client runtime"), - planned_offer_pool: Arc::clone(&self.planned_offer_pool), - planned_offer_ids: Arc::clone(&self.planned_offer_ids), } } } +fn adopted_instance(contract_id: u64) -> ProvisionedInstance { + ProvisionedInstance { + index: 0, + contract_id, + offer_id: 0, + host_id: None, + gpu_name: "adopted".to_owned(), + gpu_ram: None, + dph_total: 0.0, + } +} + impl VastAiLeaseClient for ToolsVastAiLeaseClient { fn provision_one(&mut self, request: ProvisionRequest) -> Result { if request.count != 1 { @@ -447,9 +449,13 @@ impl VastAiLeaseClient for ToolsVastAiLeaseClient { request.count )); } + if let Some(label) = request.label.as_deref() + && let Some(contract_id) = self.contract_by_label(label)? + { + return Ok(adopted_instance(contract_id)); + } let pool = self.candidate_pool(&request)?; - let planned_offer_ids = self.planned_offer_ids.lock().clone(); let blocked_hosts = request .selection .blacklist_hosts @@ -458,23 +464,11 @@ impl VastAiLeaseClient for ToolsVastAiLeaseClient { .collect::>(); let mut failed_hosts = HashSet::new(); let mut tried_offer_ids = HashSet::new(); - let mut ordered = Vec::with_capacity(pool.len()); - if let Some(preferred_offer_id) = request.preferred_offer_id - && let Some(offer) = pool.iter().find(|offer| offer.id == preferred_offer_id) - { - ordered.push(offer.clone()); - } - ordered.extend(pool.into_iter()); - let mut last_error = None; - for offer in ordered { + for offer in pool { if !tried_offer_ids.insert(offer.id) { continue; } - if request.preferred_offer_id != Some(offer.id) && planned_offer_ids.contains(&offer.id) - { - continue; - } if offer.host_id.is_some_and(|host_id| { blocked_hosts.contains(&host_id) || failed_hosts.contains(&host_id) }) { @@ -483,6 +477,16 @@ impl VastAiLeaseClient for ToolsVastAiLeaseClient { match self.create_from_offer(&request, &offer) { Ok(instance) => return Ok(instance), Err(error) => { + if let Some(label) = request.label.as_deref() + && let Some(contract_id) = self.contract_by_label(label)? + { + return Ok(adopted_instance(contract_id)); + } + if classify_vastai_error(&error) + != swactor_vastai::VastAiFailureClass::VanishedOffer + { + return Err(format!("offer {}: {error}", offer.id)); + } if let Some(host_id) = offer.host_id { failed_hosts.insert(host_id); } @@ -499,36 +503,6 @@ impl VastAiLeaseClient for ToolsVastAiLeaseClient { )) } - fn plan_first_wave_offers( - &mut self, - requests: &[ProvisionRequest], - ) -> Result>, String> { - let Some(first) = requests.first() else { - self.planned_offer_pool.lock().clear(); - self.planned_offer_ids.lock().clear(); - return Ok(Vec::new()); - }; - let pool = self.runtime.block_on( - self.client - .search_offers(&first.selection, requests.len() as u32), - )?; - let planned = swactor_vastai::plan_distinct_host_first_wave( - &pool, - requests.len() as u32, - &first.selection.blacklist_hosts, - &[], - ); - let planned_ids = planned.iter().map(|offer| offer.id).collect::>(); - *self.planned_offer_pool.lock() = pool; - *self.planned_offer_ids.lock() = planned_ids; - let mut out = planned - .into_iter() - .map(|offer| Some(offer.id)) - .collect::>(); - out.resize(requests.len(), None); - Ok(out) - } - fn ssh_endpoint( &mut self, contract_id: u64, @@ -587,12 +561,7 @@ impl VastAiLeaseClient for ToolsVastAiLeaseClient { fn destroy_contract(&mut self, contract_id: u64) -> Result<(), String> { self.runtime - .block_on(swactor_vastai::destroy_instance_with_retry( - self.client.http(), - self.client.base_url(), - self.client.api_key(), - contract_id, - )) + .block_on(self.client.destroy_instance_with_retry(contract_id)) } } @@ -1165,6 +1134,8 @@ struct VastAiNode { node_id: u64, label: String, sink: PluginSink, + spec: NodeProvisionSpec, + endpoint: VastAiSshEndpoint, } impl VastAiProvisioningPlugin @@ -1187,8 +1158,8 @@ where fn label_for(&self, spec: &NodeProvisionSpec) -> String { format!( - "{}-{}-{}", - self.config.label_prefix, spec.run_id, spec.node_id + "{}-{}-{}-attempt-{}", + self.config.label_prefix, spec.run_id, spec.node_id, spec.attempt_id ) } @@ -1241,23 +1212,12 @@ fn classified_start_error(reason: String) -> String { format!("{reason} [class={class}]") } -struct VastAiStartedLease { - label: String, - instance: ProvisionedInstance, - endpoint: VastAiSshEndpoint, -} - -struct VastAiBatchStartError { - reason: String, - failed_host_id: Option, -} - impl ProvisionPlugin for VastAiProvisioningPlugin where C: VastAiLeaseClient + Clone + 'static, B: VastAiBootstrapLauncher + 'static, { - fn start_node( + fn create_node( &mut self, spec: NodeProvisionSpec, sink: PluginSink, @@ -1267,33 +1227,58 @@ where } let stream_id = node_stream_id(spec.run_id, spec.node_id); let label = self.label_for(&spec); - emit_node_line(&sink, spec.run_id, spec.node_id, format!("vastai provisioning label={label} stream={stream_id}")); + emit_node_line( + &sink, + spec.run_id, + spec.node_id, + format!("vastai provisioning label={label} stream={stream_id}"), + ); let request = self.build_request(&spec, label.clone()); - let instance = self.client.provision_one(request).map_err(|e| { - classified_start_error(format!("vastai provision node {}: {e}", spec.node_id)) + let instance = self.client.provision_one(request).map_err(|error| { + classified_start_error(format!("vastai provision node {}: {error}", spec.node_id)) })?; - emit_node_line(&sink, spec.run_id, spec.node_id, format!("vastai contract {} ready for SSH lookup", instance.contract_id)); - emit_node_line(&sink, spec.run_id, spec.node_id, serde_json::json!({ - "type": "VastAiLeaseReady", - "run_id": spec.run_id, - "node_id": spec.node_id, - "label": &label, - "image": &spec.image, - "contract_id": instance.contract_id, - "offer_id": instance.offer_id, - "host_id": instance.host_id, - "gpu_name": &instance.gpu_name, - "gpu_ram": instance.gpu_ram, - "dph_total": instance.dph_total, - }).to_string()); - emit_node_line(&sink, spec.run_id, spec.node_id, serde_json::json!({ - "type": "VastAiSshEndpointDiscoveryStarted", - "run_id": spec.run_id, - "node_id": spec.node_id, - "contract_id": instance.contract_id, - "label": &label, - }).to_string()); + emit_node_line( + &sink, + spec.run_id, + spec.node_id, + format!( + "vastai contract {} ready for SSH lookup", + instance.contract_id + ), + ); + emit_node_line( + &sink, + spec.run_id, + spec.node_id, + serde_json::json!({ + "type": "VastAiLeaseReady", + "run_id": spec.run_id, + "node_id": spec.node_id, + "label": &label, + "image": &spec.image, + "contract_id": instance.contract_id, + "offer_id": instance.offer_id, + "host_id": instance.host_id, + "gpu_name": &instance.gpu_name, + "gpu_ram": instance.gpu_ram, + "dph_total": instance.dph_total, + }) + .to_string(), + ); + emit_node_line( + &sink, + spec.run_id, + spec.node_id, + serde_json::json!({ + "type": "VastAiSshEndpointDiscoveryStarted", + "run_id": spec.run_id, + "node_id": spec.node_id, + "contract_id": instance.contract_id, + "label": &label, + }) + .to_string(), + ); let endpoint = match self.client.ssh_endpoint( instance.contract_id, @@ -1315,53 +1300,26 @@ where )); } }; - emit_node_line(&sink, spec.run_id, spec.node_id, serde_json::json!({ - "type": "VastAiSshEndpointReady", - "run_id": spec.run_id, - "node_id": spec.node_id, - "contract_id": instance.contract_id, - "host": &endpoint.host, - "port": endpoint.port, - "user": &endpoint.user, - }).to_string()); - - emit_node_line(&sink, spec.run_id, spec.node_id, serde_json::json!({ - "type": "VastAiBootstrapObservationStarted", - "run_id": spec.run_id, - "node_id": spec.node_id, - "contract_id": instance.contract_id, - "host": &endpoint.host, - "port": endpoint.port, - "user": &endpoint.user, - }).to_string()); - - let bootstrap = match self.bootstrap.start_bootstrap( - spec.clone(), - endpoint, - sink.clone(), - self.bootstrap_producer.clone(), - self.config.lifecycle.clone(), - ) { - Ok(handle) => handle, - Err(error) => { - if let Some(host_id) = instance.host_id { - self.failed_host_ids.insert(host_id); - } - return Err(self.cleanup_contract_after_start_error( - instance.contract_id, - classified_start_error(format!( - "vastai bootstrap node {}: {error}", - spec.node_id - )), - )); - } - }; + emit_node_line( + &sink, + spec.run_id, + spec.node_id, + serde_json::json!({ + "type": "VastAiSshEndpointReady", + "run_id": spec.run_id, + "node_id": spec.node_id, + "contract_id": instance.contract_id, + "host": &endpoint.host, + "port": endpoint.port, + "user": &endpoint.user, + }) + .to_string(), + ); let host_id = instance.host_id; if let Some(host_id) = host_id { self.leased_host_ids.insert(host_id); } - let handle = PluginNodeHandle { id: self.next_handle_id, provider_process_id: None, @@ -1371,7 +1329,7 @@ where handle.id, VastAiNode { contract_id: instance.contract_id, - bootstrap: Some(bootstrap), + bootstrap: None, provider_monitor: self.client.spawn_provider_monitor( instance.contract_id, label.clone(), @@ -1384,268 +1342,87 @@ where node_id: spec.node_id, label, sink, + spec, + endpoint, }, ); Ok(handle) } - fn start_nodes( - &mut self, - specs: Vec, - sink: PluginSink, - ) -> Vec<(NodeProvisionSpec, Result)> { - if specs.len() <= 1 { - return specs - .into_iter() - .map(|spec| { - let result = self.start_node(spec.clone(), sink.clone()); - (spec, result) - }) - .collect(); + fn start_bootstrap(&mut self, handle: &PluginNodeHandle) -> Result<(), String> { + let node = self + .nodes + .get_mut(&handle.id) + .ok_or_else(|| format!("vastai node handle {} is absent", handle.id))?; + if node.bootstrap.is_some() { + return Ok(()); } - - let mut results = (0..specs.len()).map(|_| None).collect::>(); - let mut start_inputs = Vec::new(); - for (index, spec) in specs.into_iter().enumerate() { - if !spec.mounts.is_empty() { - results[index] = Some(( - spec, - Err("vastai provider does not support host file mounts".to_owned()), - )); - continue; - } - let stream_id = node_stream_id(spec.run_id, spec.node_id); - let label = self.label_for(&spec); - emit_node_line(&sink, spec.run_id, spec.node_id, format!("vastai provisioning label={label} stream={stream_id}")); - let request = self.build_request(&spec, label.clone()); - start_inputs.push((index, spec, label, request)); - } - - let request_plan = start_inputs - .iter() - .map(|(_, _, _, request)| request.clone()) - .collect::>(); - let offer_plan = match self.client.plan_first_wave_offers(&request_plan) { - Ok(plan) => plan, - Err(error) => { - for (_, spec, _, _) in &start_inputs { - emit_node_line(&sink, spec.run_id, spec.node_id, format!("vastai first-wave offer planning failed; falling back to per-node selection: {error}")); - } - vec![None; start_inputs.len()] - } - }; - - let (completion_tx, completion_rx) = mpsc::channel(); - for (plan_index, (index, spec, label, mut request)) in start_inputs.into_iter().enumerate() - { - request.preferred_offer_id = offer_plan.get(plan_index).copied().flatten(); - if let Some(offer_id) = request.preferred_offer_id { - emit_node_line(&sink, spec.run_id, spec.node_id, serde_json::json!({ - "type": "VastAiFirstWaveOfferPlanned", - "run_id": spec.run_id, - "node_id": spec.node_id, - "label": &label, - "offer_id": offer_id, - }).to_string()); - } else { - emit_node_line(&sink, spec.run_id, spec.node_id, serde_json::json!({ - "type": "VastAiFirstWaveOfferPlanUnavailable", - "run_id": spec.run_id, - "node_id": spec.node_id, - "label": &label, - }).to_string()); - } - let mut client = self.client.clone(); - let config = self.config.clone(); - let worker_tx = completion_tx.clone(); - let worker_sink = sink.clone(); - std::thread::spawn(move || { - let started = match client.provision_one(request) { - Ok(instance) => { - emit_node_line(&worker_sink, spec.run_id, spec.node_id, serde_json::json!({ - "type": "VastAiLeaseReady", - "run_id": spec.run_id, - "node_id": spec.node_id, - "label": &label, - "image": &spec.image, - "contract_id": instance.contract_id, - "offer_id": instance.offer_id, - "host_id": instance.host_id, - "gpu_name": &instance.gpu_name, - "gpu_ram": instance.gpu_ram, - "dph_total": instance.dph_total, - }).to_string()); - emit_node_line(&worker_sink, spec.run_id, spec.node_id, serde_json::json!({ - "type": "VastAiSshEndpointDiscoveryStarted", - "run_id": spec.run_id, - "node_id": spec.node_id, - "contract_id": instance.contract_id, - "label": &label, - }).to_string()); - match client.ssh_endpoint( - instance.contract_id, - &label, - &config.lifecycle, - &config.ssh_user, - ) { - Ok(endpoint) => Ok(VastAiStartedLease { - label, - instance, - endpoint, - }), - Err(error) => { - let failed_host_id = instance.host_id; - let reason = match client.destroy_contract(instance.contract_id) { - Ok(()) => classified_start_error(format!( - "vastai SSH endpoint node {}: {error}", - spec.node_id - )), - Err(cleanup) => classified_start_error(format!( - "vastai SSH endpoint node {}: {error}; cleanup destroy {} failed: {cleanup}", - spec.node_id, instance.contract_id - )), - }; - Err(VastAiBatchStartError { - reason, - failed_host_id, - }) - } - } - } - Err(error) => Err(VastAiBatchStartError { - reason: classified_start_error(format!( - "vastai provision node {}: {error}", - spec.node_id - )), - failed_host_id: None, - }), - }; - let _ = worker_tx.send((index, spec, started)); - }); - } - drop(completion_tx); - - for (index, spec, started) in completion_rx { - match started { - Ok(started) => { - emit_node_line(&sink, spec.run_id, spec.node_id, serde_json::json!({ - "type": "VastAiSshEndpointReady", - "run_id": spec.run_id, - "node_id": spec.node_id, - "contract_id": started.instance.contract_id, - "host": &started.endpoint.host, - "port": started.endpoint.port, - "user": &started.endpoint.user, - }).to_string()); - emit_node_line(&sink, spec.run_id, spec.node_id, serde_json::json!({ - "type": "VastAiBootstrapObservationStarted", - "run_id": spec.run_id, - "node_id": spec.node_id, - "contract_id": started.instance.contract_id, - "host": &started.endpoint.host, - "port": started.endpoint.port, - "user": &started.endpoint.user, - }).to_string()); - - let bootstrap = match self.bootstrap.start_bootstrap( - spec.clone(), - started.endpoint, - sink.clone(), - self.bootstrap_producer.clone(), - self.config.lifecycle.clone(), - ) { - Ok(handle) => handle, - Err(error) => { - if let Some(host_id) = started.instance.host_id { - self.failed_host_ids.insert(host_id); - } - let node_id = spec.node_id; - results[index] = Some(( - spec, - Err(self.cleanup_contract_after_start_error( - started.instance.contract_id, - classified_start_error(format!( - "vastai bootstrap node {node_id}: {error}" - )), - )), - )); - continue; - } - }; - - let host_id = started.instance.host_id; - if let Some(host_id) = host_id { - self.leased_host_ids.insert(host_id); - } - let handle = PluginNodeHandle { - id: self.next_handle_id, - provider_process_id: None, - }; - self.next_handle_id = self.next_handle_id.wrapping_add(1).max(1); - self.nodes.insert( - handle.id, - VastAiNode { - contract_id: started.instance.contract_id, - bootstrap: Some(bootstrap), - provider_monitor: self.client.spawn_provider_monitor( - started.instance.contract_id, - started.label.clone(), - self.config.lifecycle.clone(), - spec.clone(), - sink.clone(), - ), - host_id, - run_id: spec.run_id, - node_id: spec.node_id, - label: started.label, - sink: sink.clone(), - }, - ); - results[index] = Some((spec, Ok(handle))); - } - Err(error) => { - if let Some(host_id) = error.failed_host_id { - self.failed_host_ids.insert(host_id); - } - results[index] = Some((spec, Err(error.reason))); - } - } - } - - results - .into_iter() - .enumerate() - .map(|(index, result)| { - result.unwrap_or_else(|| { - ( - NodeProvisionSpec { - run_id: 0, - node_id: u64::try_from(index).unwrap_or(u64::MAX), - stage_index: None, - image: String::new(), - env: Vec::new(), - args: Vec::new(), - mounts: Vec::new(), - }, - Err("vastai provision worker panicked".to_owned()), - ) - }) + emit_node_line( + &node.sink, + node.run_id, + node.node_id, + serde_json::json!({ + "type": "VastAiBootstrapObservationStarted", + "run_id": node.run_id, + "node_id": node.node_id, + "contract_id": node.contract_id, + "host": &node.endpoint.host, + "port": node.endpoint.port, + "user": &node.endpoint.user, }) - .collect() + .to_string(), + ); + match self.bootstrap.start_bootstrap( + node.spec.clone(), + node.endpoint.clone(), + node.sink.clone(), + self.bootstrap_producer.clone(), + self.config.lifecycle.clone(), + ) { + Ok(bootstrap) => { + node.bootstrap = Some(bootstrap); + Ok(()) + } + Err(error) => { + if let Some(host_id) = node.host_id { + self.failed_host_ids.insert(host_id); + } + Err(classified_start_error(format!( + "vastai bootstrap node {}: {error}", + node.node_id + ))) + } + } + } + + fn cancel_bootstrap(&mut self, handle: &PluginNodeHandle) -> Result<(), String> { + let Some(node) = self.nodes.get_mut(&handle.id) else { + return Ok(()); + }; + if let Some(mut bootstrap) = node.bootstrap.take() { + self.bootstrap.stop_bootstrap(&mut bootstrap); + } + Ok(()) } fn complete_bootstrap(&mut self, handle: &PluginNodeHandle) -> Result<(), String> { let Some(node) = self.nodes.get_mut(&handle.id) else { return Ok(()); }; - emit_node_line(&node.sink, node.run_id, node.node_id, serde_json::json!({ - "type": "VastAiRuntimeReadyAccepted", - "run_id": node.run_id, - "node_id": node.node_id, - "label": &node.label, - "contract_id": node.contract_id, - "classification": "runtime_ready_over_provider_staleness", - }).to_string()); + emit_node_line( + &node.sink, + node.run_id, + node.node_id, + serde_json::json!({ + "type": "VastAiRuntimeReadyAccepted", + "run_id": node.run_id, + "node_id": node.node_id, + "label": &node.label, + "contract_id": node.contract_id, + "classification": "runtime_ready_over_provider_staleness", + }) + .to_string(), + ); Ok(()) } @@ -1656,22 +1433,228 @@ where if let Some(mut monitor) = node.provider_monitor.take() { monitor.stop(); } - if let Some(host_id) = node.host_id { - self.leased_host_ids.remove(&host_id); - } if let Some(mut bootstrap) = node.bootstrap.take() { self.bootstrap.stop_bootstrap(&mut bootstrap); } let result = self.client.destroy_contract(node.contract_id); - emit_node_line(&node.sink, node.run_id, node.node_id, serde_json::json!({ - "type": "VastAiContractCleanup", - "run_id": node.run_id, - "node_id": node.node_id, - "label": &node.label, - "contract_id": node.contract_id, - "result": if result.is_ok() { "ok" } else { "failed" }, - "error": result.as_ref().err(), - }).to_string()); - result + emit_node_line( + &node.sink, + node.run_id, + node.node_id, + serde_json::json!({ + "type": "VastAiContractCleanup", + "run_id": node.run_id, + "node_id": node.node_id, + "label": &node.label, + "contract_id": node.contract_id, + "result": if result.is_ok() { "ok" } else { "failed" }, + "error": result.as_ref().err(), + }) + .to_string(), + ); + match result { + Ok(()) => { + if let Some(host_id) = node.host_id { + self.leased_host_ids.remove(&host_id); + } + Ok(()) + } + Err(error) => { + self.nodes.insert(handle.id, node); + Err(error) + } + } + } +} + +#[cfg(test)] +mod tests { + use serde_json::json; + use wiremock::matchers::{method, path}; + use wiremock::{Mock, MockServer, ResponseTemplate}; + + use super::*; + + #[test] + fn provision_one_adopts_stable_label_before_offer_search() { + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + let server = runtime.block_on(MockServer::start()); + runtime.block_on(async { + Mock::given(method("GET")) + .and(path("/api/v0/instances/")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "instances": [{ + "id": 73, + "label": "run-5-node-7-attempt-9", + "actual_status": "loading", + "ssh_host": "", + "ssh_port": 0, + "public_ipaddr": "" + }] + }))) + .mount(&server) + .await; + }); + let client = swactor_vastai::VastClient::with_base_url(server.uri(), "secret"); + let mut client = ToolsVastAiLeaseClient::new(client).unwrap(); + let adopted = client + .provision_one(ProvisionRequest { + count: 1, + image: "node:v1".to_owned(), + label: Some("run-5-node-7-attempt-9".to_owned()), + disk_gb: 10, + env: BTreeMap::new(), + per_instance_env: vec![BTreeMap::new()], + preferred_offer_id: None, + onstart: None, + selection: SelectionPolicy::default(), + lifecycle: LifecyclePolicy::default(), + confirm_lease: false, + }) + .unwrap(); + + assert_eq!(adopted.contract_id, 73); + assert_eq!(adopted.offer_id, 0); + let requests = runtime.block_on(server.received_requests()).unwrap(); + assert_eq!(requests.len(), 1); + assert_eq!(requests[0].url.path(), "/api/v0/instances/"); + } + + use std::sync::atomic::AtomicUsize; + + use crate::provisioning::PluginObservationSink; + + struct NullSink; + + impl PluginObservationSink for NullSink { + fn observe(&self, _observation: PluginObservation) {} + } + + #[derive(Clone)] + struct RetryDestroyClient { + destroy_calls: Arc, + destroy_failures: Arc, + } + + impl VastAiLeaseClient for RetryDestroyClient { + fn provision_one( + &mut self, + _request: ProvisionRequest, + ) -> Result { + Ok(ProvisionedInstance { + index: 0, + contract_id: 73, + offer_id: 11, + host_id: Some(44), + gpu_name: "test".to_owned(), + gpu_ram: None, + dph_total: 0.0, + }) + } + + fn ssh_endpoint( + &mut self, + _contract_id: u64, + _label: &str, + _lifecycle: &LifecyclePolicy, + ssh_user: &str, + ) -> Result { + Ok(VastAiSshEndpoint { + host: "host".to_owned(), + port: 22, + user: ssh_user.to_owned(), + }) + } + + fn destroy_contract(&mut self, _contract_id: u64) -> Result<(), String> { + self.destroy_calls.fetch_add(1, Ordering::SeqCst); + if self + .destroy_failures + .fetch_update(Ordering::SeqCst, Ordering::SeqCst, |remaining| { + (remaining > 0).then(|| remaining - 1) + }) + .is_ok() + { + Err("transient destroy failure".to_owned()) + } else { + Ok(()) + } + } + } + + struct CountingBootstrap { + starts: Arc, + stops: Arc, + } + + impl VastAiBootstrapLauncher for CountingBootstrap { + type Handle = (); + + fn start_bootstrap( + &mut self, + _spec: NodeProvisionSpec, + _endpoint: VastAiSshEndpoint, + _sink: PluginSink, + _producer: Option, + _lifecycle: LifecyclePolicy, + ) -> Result { + self.starts.fetch_add(1, Ordering::SeqCst); + Ok(()) + } + + fn stop_bootstrap(&mut self, _handle: &mut Self::Handle) { + self.stops.fetch_add(1, Ordering::SeqCst); + } + } + + #[test] + fn node_creation_defers_bootstrap_and_failed_destroy_remains_retryable() { + let destroy_calls = Arc::new(AtomicUsize::new(0)); + let destroy_failures = Arc::new(AtomicUsize::new(1)); + let starts = Arc::new(AtomicUsize::new(0)); + let stops = Arc::new(AtomicUsize::new(0)); + let mut plugin = VastAiProvisioningPlugin::new( + RetryDestroyClient { + destroy_calls: Arc::clone(&destroy_calls), + destroy_failures, + }, + CountingBootstrap { + starts: Arc::clone(&starts), + stops: Arc::clone(&stops), + }, + VastAiProvisioningConfig::default(), + ); + let spec = NodeProvisionSpec { + run_id: 5, + node_id: 7, + attempt_id: 9, + stage_index: Some(0), + image: "node:v1".to_owned(), + env: Vec::new(), + args: Vec::new(), + mounts: Vec::new(), + }; + let handle = plugin + .create_node(spec, PluginSink::new(Arc::new(NullSink))) + .unwrap(); + assert_eq!(starts.load(Ordering::SeqCst), 0); + + plugin.start_bootstrap(&handle).unwrap(); + assert_eq!(starts.load(Ordering::SeqCst), 1); + assert_eq!( + plugin.stop_node(&handle).unwrap_err(), + "transient destroy failure" + ); + assert!(plugin.nodes.contains_key(&handle.id)); + assert!(plugin.leased_host_ids.contains(&44)); + + plugin.stop_node(&handle).unwrap(); + assert!(!plugin.nodes.contains_key(&handle.id)); + assert!(!plugin.leased_host_ids.contains(&44)); + assert_eq!(destroy_calls.load(Ordering::SeqCst), 2); + assert_eq!(stops.load(Ordering::SeqCst), 1); } } diff --git a/apps/myelin/src/orchestration/provisioning.rs b/apps/myelin/src/orchestration/provisioning.rs index a4418c0..7ebc942 100644 --- a/apps/myelin/src/orchestration/provisioning.rs +++ b/apps/myelin/src/orchestration/provisioning.rs @@ -31,8 +31,10 @@ pub(crate) struct LocalDockerPlugin { } struct LocalDockerNode { + spec: NodeProvisionSpec, + sink: PluginSink, container_name: String, - stdin: ChildStdin, + stdin: Option, } pub(crate) struct LocalProcessPlugin { @@ -42,10 +44,14 @@ pub(crate) struct LocalProcessPlugin { } struct LocalProcessNode { - stdin: ChildStdin, - child: Arc>>, spec: NodeProvisionSpec, sink: PluginSink, + runtime: Option, +} + +struct LocalProcessRuntime { + stdin: ChildStdin, + child: Arc>>, } impl LocalProcessPlugin { @@ -68,6 +74,35 @@ impl LocalDockerPlugin { } } +fn docker_container_name(prefix: &str, spec: &NodeProvisionSpec) -> String { + format!( + "{prefix}-{}-{}-attempt-{}", + spec.run_id, spec.node_id, spec.attempt_id + ) +} + +#[allow(clippy::disallowed_methods)] +fn docker_container_is_absent(name: &str) -> Result { + let output = Command::new("docker") + .arg("inspect") + .arg(name) + .output() + .map_err(|error| format!("inspect Docker container {name}: {error}"))?; + if output.status.success() { + return Ok(false); + } + let stderr = String::from_utf8_lossy(&output.stderr); + if stderr.contains("No such object") || stderr.contains("No such container") { + Ok(true) + } else { + Err(format!( + "inspect Docker container {name} exited with {}: {}", + output.status, + stderr.trim() + )) + } +} + fn docker_mount_arg(mount: &ProviderMount) -> String { let mut arg = format!( "type=bind,src={},dst={}", @@ -235,13 +270,39 @@ fn lock_process_child( } impl ProvisionPlugin for LocalProcessPlugin { - // provider process supervision/lifecycle is out of scope (ENGINE_SPEC.md §2) - #[allow(clippy::disallowed_methods)] - fn start_node( + fn create_node( &mut self, spec: NodeProvisionSpec, sink: PluginSink, ) -> Result { + let handle = PluginNodeHandle { + id: self.next_handle_id, + provider_process_id: None, + }; + self.next_handle_id = self.next_handle_id.wrapping_add(1).max(1); + self.nodes.insert( + handle.id, + LocalProcessNode { + spec, + sink, + runtime: None, + }, + ); + Ok(handle) + } + + // provider process supervision/lifecycle is out of scope (ENGINE_SPEC.md §2) + #[allow(clippy::disallowed_methods)] + fn start_bootstrap(&mut self, handle: &PluginNodeHandle) -> Result<(), String> { + let node = self + .nodes + .get_mut(&handle.id) + .ok_or_else(|| format!("local process node handle {} is absent", handle.id))?; + if node.runtime.is_some() { + return Ok(()); + } + let spec = node.spec.clone(); + let sink = node.sink.clone(); let mut command = Command::new(&self.program); for (key, value) in &spec.env { command.env(key, value); @@ -269,36 +330,22 @@ impl ProvisionPlugin for LocalProcessPlugin { self.program.display() ) })?; - let provider_process_id = child.id(); - let stdin = child - .stdin - .take() - .ok_or_else(|| format!("local process node {} stdin missing", spec.node_id))?; - let stdout = child - .stdout - .take() - .ok_or_else(|| format!("local process node {} stdout missing", spec.node_id))?; - let stderr = child - .stderr - .take() - .ok_or_else(|| format!("local process node {} stderr missing", spec.node_id))?; + let (Some(stdin), Some(stdout), Some(stderr)) = + (child.stdin.take(), child.stdout.take(), child.stderr.take()) + else { + let _ = child.kill(); + let _ = child.wait(); + return Err(format!( + "local process node {} did not expose piped stdio", + spec.node_id + )); + }; let child = Arc::new(Mutex::new(Some(child))); - let handle = PluginNodeHandle { - id: self.next_handle_id, - provider_process_id: Some(provider_process_id), - }; - self.next_handle_id = self.next_handle_id.wrapping_add(1).max(1); - self.nodes.insert( - handle.id, - LocalProcessNode { - stdin, - child: Arc::clone(&child), - spec: spec.clone(), - sink: sink.clone(), - }, - ); - + node.runtime = Some(LocalProcessRuntime { + stdin, + child: Arc::clone(&child), + }); spawn_stdout_reader(spec.clone(), sink.clone(), stdout); spawn_stderr_reader(spec.clone(), sink.clone(), stderr); thread::spawn(move || { @@ -318,14 +365,11 @@ impl ProvisionPlugin for LocalProcessPlugin { }) } Ok(None) => None, - Err(error) => { - *slot = None; - Some(PluginObservation::Failed { - run_id: spec.run_id, - node_id: spec.node_id, - reason: format!("wait local process node: {error}"), - }) - } + Err(error) => Some(PluginObservation::Failed { + run_id: spec.run_id, + node_id: spec.node_id, + reason: format!("wait local process node: {error}"), + }), } }; if let Some(observation) = observation { @@ -335,8 +379,7 @@ impl ProvisionPlugin for LocalProcessPlugin { thread::sleep(Duration::from_millis(100)); } }); - - Ok(handle) + Ok(()) } fn complete_bootstrap(&mut self, _handle: &PluginNodeHandle) -> Result<(), String> { @@ -349,45 +392,53 @@ impl ProvisionPlugin for LocalProcessPlugin { let Some(mut node) = self.nodes.remove(&handle.id) else { return Ok(()); }; - let _ = node.stdin.write_all(b"shutdown\n"); - let _ = node.stdin.flush(); + let Some(mut runtime) = node.runtime.take() else { + return Ok(()); + }; + let _ = runtime.stdin.write_all(b"shutdown\n"); + let _ = runtime.stdin.flush(); let deadline = Instant::now() + Duration::from_secs(2); while Instant::now() < deadline { - let observation = { - let mut slot = lock_process_child(&node.child); + let status = { + let mut slot = lock_process_child(&runtime.child); let Some(child) = slot.as_mut() else { return Ok(()); }; match child.try_wait() { Ok(Some(status)) => { *slot = None; - Some(PluginObservation::Exited { - run_id: node.spec.run_id, - node_id: node.spec.node_id, - status: status.code(), - }) - } - Ok(None) => None, - Err(error) => { - *slot = None; - Some(PluginObservation::Failed { - run_id: node.spec.run_id, - node_id: node.spec.node_id, - reason: format!("wait local process node: {error}"), - }) + Ok(Some(status.code())) } + Ok(None) => Ok(None), + Err(error) => Err(format!("wait local process node: {error}")), } }; - if let Some(observation) = observation { - node.sink.observe(observation); - return Ok(()); + match status { + Ok(Some(status)) => { + node.sink.observe(PluginObservation::Exited { + run_id: node.spec.run_id, + node_id: node.spec.node_id, + status, + }); + return Ok(()); + } + Ok(None) => thread::sleep(Duration::from_millis(50)), + Err(reason) => { + node.sink.observe(PluginObservation::Failed { + run_id: node.spec.run_id, + node_id: node.spec.node_id, + reason: reason.clone(), + }); + node.runtime = Some(runtime); + self.nodes.insert(handle.id, node); + return Err(reason); + } } - thread::sleep(Duration::from_millis(50)); } - let observation = { - let mut slot = lock_process_child(&node.child); + let status = { + let mut slot = lock_process_child(&runtime.child); let Some(child) = slot.as_mut() else { return Ok(()); }; @@ -395,28 +446,40 @@ impl ProvisionPlugin for LocalProcessPlugin { unsafe { let _ = libc::kill(-(child.id() as i32), libc::SIGKILL); } - let _ = child.kill(); + let kill_error = child.kill().err(); match child.wait() { Ok(status) => { *slot = None; - PluginObservation::Exited { - run_id: node.spec.run_id, - node_id: node.spec.node_id, - status: status.code(), - } + Ok(status.code()) } - Err(error) => { - *slot = None; - PluginObservation::Failed { - run_id: node.spec.run_id, - node_id: node.spec.node_id, - reason: format!("kill local process node: {error}"), + Err(error) => Err(match kill_error { + Some(kill_error) => { + format!("kill local process node: {kill_error}; wait failed: {error}") } - } + None => format!("wait for killed local process node: {error}"), + }), } }; - node.sink.observe(observation); - Ok(()) + match status { + Ok(status) => { + node.sink.observe(PluginObservation::Exited { + run_id: node.spec.run_id, + node_id: node.spec.node_id, + status, + }); + Ok(()) + } + Err(reason) => { + node.sink.observe(PluginObservation::Failed { + run_id: node.spec.run_id, + node_id: node.spec.node_id, + reason: reason.clone(), + }); + node.runtime = Some(runtime); + self.nodes.insert(handle.id, node); + Err(reason) + } + } } } @@ -438,17 +501,41 @@ impl Drop for LocalProcessPlugin { } impl ProvisionPlugin for LocalDockerPlugin { - // provider process supervision/lifecycle is out of scope (ENGINE_SPEC.md §2) - #[allow(clippy::disallowed_methods)] - fn start_node( + fn create_node( &mut self, spec: NodeProvisionSpec, sink: PluginSink, ) -> Result { - let container_name = format!( - "{}-{}-{}", - self.container_name_prefix, spec.run_id, spec.node_id + let handle = PluginNodeHandle { + id: self.next_handle_id, + provider_process_id: None, + }; + self.next_handle_id = self.next_handle_id.wrapping_add(1).max(1); + self.nodes.insert( + handle.id, + LocalDockerNode { + container_name: docker_container_name(&self.container_name_prefix, &spec), + spec, + sink, + stdin: None, + }, ); + Ok(handle) + } + + // provider process supervision/lifecycle is out of scope (ENGINE_SPEC.md §2) + #[allow(clippy::disallowed_methods)] + fn start_bootstrap(&mut self, handle: &PluginNodeHandle) -> Result<(), String> { + let node = self + .nodes + .get_mut(&handle.id) + .ok_or_else(|| format!("Docker node handle {} is absent", handle.id))?; + if node.stdin.is_some() { + return Ok(()); + } + let spec = node.spec.clone(); + let sink = node.sink.clone(); + let container_name = node.container_name.clone(); let mut command = Command::new("docker"); command .arg("run") @@ -499,33 +586,24 @@ impl ProvisionPlugin for LocalDockerPlugin { .spawn() .map_err(|e| format!("spawn Docker node {}: {e}", spec.node_id))?; - let provider_process_id = child.id(); - let stdin = child - .stdin - .take() - .ok_or_else(|| format!("Docker node {} stdin missing", spec.node_id))?; - let stdout = child - .stdout - .take() - .ok_or_else(|| format!("Docker node {} stdout missing", spec.node_id))?; - let stderr = child - .stderr - .take() - .ok_or_else(|| format!("Docker node {} stderr missing", spec.node_id))?; - - let handle = PluginNodeHandle { - id: self.next_handle_id, - provider_process_id: Some(provider_process_id), + let (Some(stdin), Some(stdout), Some(stderr)) = + (child.stdin.take(), child.stdout.take(), child.stderr.take()) + else { + let _ = child.kill(); + let _ = child.wait(); + let _ = Command::new("docker") + .arg("rm") + .arg("-f") + .arg(&container_name) + .stdout(Stdio::null()) + .stderr(Stdio::null()) + .status(); + return Err(format!( + "Docker node {} did not expose piped stdio", + spec.node_id + )); }; - self.next_handle_id = self.next_handle_id.wrapping_add(1).max(1); - self.nodes.insert( - handle.id, - LocalDockerNode { - container_name: container_name.clone(), - stdin, - }, - ); - + node.stdin = Some(stdin); spawn_stdout_reader(spec.clone(), sink.clone(), stdout); spawn_stderr_reader(spec.clone(), sink.clone(), stderr); thread::spawn(move || match child.wait() { @@ -540,8 +618,7 @@ impl ProvisionPlugin for LocalDockerPlugin { reason: format!("wait Docker node: {error}"), }), }); - - Ok(handle) + Ok(()) } fn complete_bootstrap(&mut self, _handle: &PluginNodeHandle) -> Result<(), String> { @@ -552,9 +629,12 @@ impl ProvisionPlugin for LocalDockerPlugin { let Some(mut node) = self.nodes.remove(&handle.id) else { return Ok(()); }; - let _ = writeln!(node.stdin, "shutdown"); - let _ = node.stdin.flush(); - let status = Command::new("docker") + let Some(stdin) = node.stdin.as_mut() else { + return Ok(()); + }; + let _ = writeln!(stdin, "shutdown"); + let _ = stdin.flush(); + let result = match Command::new("docker") .arg("stop") .arg("-t") .arg("2") @@ -562,15 +642,32 @@ impl ProvisionPlugin for LocalDockerPlugin { .stdout(Stdio::null()) .stderr(Stdio::null()) .status() - .map_err(|e| format!("docker stop {}: {e}", node.container_name))?; - if status.success() { - Ok(()) - } else { - Err(format!( - "docker stop {} exited with {status}", - node.container_name - )) + { + Ok(status) if status.success() => Ok(()), + Ok(status) => match docker_container_is_absent(&node.container_name) { + Ok(true) => Ok(()), + Ok(false) => Err(format!( + "docker stop {} exited with {status}", + node.container_name + )), + Err(inspect_error) => Err(format!( + "docker stop {} exited with {status}; {inspect_error}", + node.container_name + )), + }, + Err(error) => match docker_container_is_absent(&node.container_name) { + Ok(true) => Ok(()), + Ok(false) => Err(format!("docker stop {}: {error}", node.container_name)), + Err(inspect_error) => Err(format!( + "docker stop {}: {error}; {inspect_error}", + node.container_name + )), + }, + }; + if result.is_err() { + self.nodes.insert(handle.id, node); } + result } } @@ -589,3 +686,73 @@ fn spawn_stderr_reader( ) { BootstrapDatastreamBridge::new(spec, sink, None).spawn_stderr_reader(stderr); } + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::mpsc; + + struct ChannelSink(mpsc::Sender); + + impl PluginObservationSink for ChannelSink { + fn observe(&self, observation: PluginObservation) { + let _ = self.0.send(observation); + } + } + + fn test_spec() -> NodeProvisionSpec { + NodeProvisionSpec { + run_id: 5, + node_id: 7, + attempt_id: 11, + stage_index: Some(0), + image: "node:v1".to_owned(), + env: Vec::new(), + args: Vec::new(), + mounts: Vec::new(), + } + } + + #[cfg(target_os = "linux")] + #[test] + fn local_process_creation_does_not_start_bootstrap() { + let (tx, rx) = mpsc::channel(); + let sink = PluginSink::new(Arc::new(ChannelSink(tx))); + let mut spec = test_spec(); + spec.args = vec![ + "-c".to_owned(), + "printf 'started\\n'; IFS= read -r line".to_owned(), + ]; + let mut plugin = LocalProcessPlugin::new("/bin/sh"); + + let handle = plugin.create_node(spec, sink).unwrap(); + assert!(matches!(rx.try_recv(), Err(mpsc::TryRecvError::Empty))); + + plugin.start_bootstrap(&handle).unwrap(); + let deadline = Instant::now() + Duration::from_secs(2); + let line = loop { + let remaining = deadline.saturating_duration_since(Instant::now()); + let observation = rx.recv_timeout(remaining).unwrap(); + if let PluginObservation::StdoutLine { line, .. } = observation { + break line; + } + }; + assert_eq!(line, "started"); + plugin.stop_node(&handle).unwrap(); + } + + #[test] + fn docker_container_identity_distinguishes_node_attempts() { + let mut spec = test_spec(); + + assert_eq!( + docker_container_name("myelin", &spec), + "myelin-5-7-attempt-11" + ); + spec.attempt_id = 12; + assert_eq!( + docker_container_name("myelin", &spec), + "myelin-5-7-attempt-12" + ); + } +} diff --git a/apps/myelin/src/orchestration/run_fsm.rs b/apps/myelin/src/orchestration/run_fsm.rs index 125aaa0..d7a217f 100644 --- a/apps/myelin/src/orchestration/run_fsm.rs +++ b/apps/myelin/src/orchestration/run_fsm.rs @@ -1,4 +1,3 @@ - #[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] pub(crate) struct RunId(pub(crate) u64); #[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] @@ -288,7 +287,9 @@ impl OrchestratorRun { sequence: sequence + 1, payload: TokenObjectPayload::Decode { token_id, - sampling: SamplingData { source_sequence: sequence }, + sampling: SamplingData { + source_sequence: sequence, + }, }, }); } else { diff --git a/apps/myelin/src/orchestration/run_plan.rs b/apps/myelin/src/orchestration/run_plan.rs index a7287db..9abd6e3 100644 --- a/apps/myelin/src/orchestration/run_plan.rs +++ b/apps/myelin/src/orchestration/run_plan.rs @@ -1,4 +1,3 @@ - pub(crate) const MO01_HEADER_BYTES: u64 = 40; const TOKEN_ID_WIDTH_BYTES: u32 = 4; @@ -431,8 +430,10 @@ pub(crate) fn plan_run(input: PlannerInput) -> Result { let mut stages = Vec::with_capacity(input.stage_count as usize); for placement in &placements { let stage_index = placement.stage_index; - let start = (u64::from(input.model.num_layers) * u64::from(stage_index) / u64::from(input.stage_count)) as u32; - let end = (u64::from(input.model.num_layers) * u64::from(stage_index + 1) / u64::from(input.stage_count)) as u32; + let start = (u64::from(input.model.num_layers) * u64::from(stage_index) + / u64::from(input.stage_count)) as u32; + let end = (u64::from(input.model.num_layers) * u64::from(stage_index + 1) + / u64::from(input.stage_count)) as u32; let inbound_edge = if stage_index == 0 { token_in_edge } else { diff --git a/apps/myelin/src/staging/control.rs b/apps/myelin/src/staging/control.rs index d024f0a..821623c 100644 --- a/apps/myelin/src/staging/control.rs +++ b/apps/myelin/src/staging/control.rs @@ -1,4 +1,3 @@ - use crate::gguf_shard::StageShardPlan; use crate::run_plan::{GgufSource, TokenizerSource}; diff --git a/apps/myelin/src/staging/gguf_metadata.rs b/apps/myelin/src/staging/gguf_metadata.rs index 52725b5..875d101 100644 --- a/apps/myelin/src/staging/gguf_metadata.rs +++ b/apps/myelin/src/staging/gguf_metadata.rs @@ -161,7 +161,10 @@ fn required_u32(map: &BTreeMap, key: &str, label: &str) -> Result(reader: &mut R, value_type: GgufValueType) -> Result<(), String> { +pub(crate) fn skip_scalar( + reader: &mut R, + value_type: GgufValueType, +) -> Result<(), String> { match value_type { GgufValueType::String => skip_gguf_string(reader), GgufValueType::Array => skip_array(reader), @@ -195,7 +198,10 @@ pub(crate) fn skip_array(reader: &mut R) -> Result<(), String> { } } -pub(crate) fn read_gguf_string(reader: &mut R, max_len: u64) -> Result { +pub(crate) fn read_gguf_string( + reader: &mut R, + max_len: u64, +) -> Result { let len = read_u64(reader)?; if len > max_len { return Err(format!( diff --git a/apps/myelin/src/staging/gguf_shard.rs b/apps/myelin/src/staging/gguf_shard.rs index a5f4bdd..413c86c 100644 --- a/apps/myelin/src/staging/gguf_shard.rs +++ b/apps/myelin/src/staging/gguf_shard.rs @@ -6,7 +6,9 @@ use crate::gguf_common::{GgufValueType, read_integer_value, read_u32, read_u64}; use serde::{Deserialize, Serialize}; use crate::run_plan::GgufSource; -use crate::staging::gguf_metadata::{skip_scalar as skip_value, read_gguf_string, GGUF_MAGIC, SUPPORTED_GGUF_VERSION}; +use crate::staging::gguf_metadata::{ + GGUF_MAGIC, SUPPORTED_GGUF_VERSION, read_gguf_string, skip_scalar as skip_value, +}; const DEFAULT_ALIGNMENT: u64 = 32; const MAX_STRING_BYTES: u64 = 64 * 1024 * 1024; @@ -204,7 +206,10 @@ pub(crate) fn source_url(source: &GgufSource) -> Result { } => Ok(format!( "https://huggingface.co/{repo}/resolve/{}/{}", revision.as_deref().unwrap_or("main"), - file.split('/').map(percent_encode_path_segment).collect::>().join("/") + file.split('/') + .map(percent_encode_path_segment) + .collect::>() + .join("/") )), GgufSource::LocalPath(path) => Err(format!( "stage shard range fetching requires a remote Hugging Face source; got local path {path:?}" diff --git a/apps/myelin/src/staging/mod.rs b/apps/myelin/src/staging/mod.rs index d50257c..ccbcf8b 100644 --- a/apps/myelin/src/staging/mod.rs +++ b/apps/myelin/src/staging/mod.rs @@ -1,4 +1,3 @@ - //! Myelin stage control, shard planning, and weight lifecycle public surface. pub(crate) mod control; diff --git a/apps/myelin/src/tests/local_mock/mock_node.rs b/apps/myelin/src/tests/local_mock/mock_node.rs index 0a39de8..a0cbb6e 100644 --- a/apps/myelin/src/tests/local_mock/mock_node.rs +++ b/apps/myelin/src/tests/local_mock/mock_node.rs @@ -1,7 +1,7 @@ use crate::run_plan as plan; +use crate::tests::harness::StageControllerHarness; use data_plane::edge_actor; use myelin::staging as stage; -use crate::tests::harness::StageControllerHarness; use super::mock_transport::MockObject; use super::mock_worker::MockWorker; diff --git a/apps/myelin/src/tests/orchestration_guarantees.rs b/apps/myelin/src/tests/orchestration_guarantees.rs index a81f6ae..7c18e56 100644 --- a/apps/myelin/src/tests/orchestration_guarantees.rs +++ b/apps/myelin/src/tests/orchestration_guarantees.rs @@ -679,9 +679,12 @@ mod run_fsm { run_id: fsm::RunId(7), stage_index: 99, }); - assert!(invalid.events().iter().any(|event| { - matches!(event, fsm::LifecycleEvent::RunFaulted { .. }) - })); + assert!( + invalid + .events() + .iter() + .any(|event| { matches!(event, fsm::LifecycleEvent::RunFaulted { .. }) }) + ); } // This proves execution has one start signal and advances by the token feedback diff --git a/apps/myelin/src/tests/staging_guarantees.rs b/apps/myelin/src/tests/staging_guarantees.rs index a7efea2..a4a2ab7 100644 --- a/apps/myelin/src/tests/staging_guarantees.rs +++ b/apps/myelin/src/tests/staging_guarantees.rs @@ -11,8 +11,8 @@ mod stage_controller { //! They assert the guarantees in //! `specs/BEHAVIOR_GUARANTEES.md`. - use myelin::staging as stage; use crate::tests::harness::StageControllerHarness; + use myelin::staging as stage; // This provision fixture represents a single middle stage. It has one inbound // and one outbound edge so tests can prove the controller uses assigned edges diff --git a/crates/provisioning/src/executor.rs b/crates/provisioning/src/executor.rs new file mode 100644 index 0000000..5779f40 --- /dev/null +++ b/crates/provisioning/src/executor.rs @@ -0,0 +1,511 @@ +//! Identity-aware asynchronous effect execution for the cluster reconciler. + +use std::collections::BTreeMap; +use std::fmt; +use std::sync::mpsc::{self, Receiver, Sender, TryRecvError}; +use std::sync::{Arc, Mutex, MutexGuard}; + +use crate::reconciler::{ + ExecutorResult, NodeAttemptId, OperationId, OperationKind, OperationOutcome, PlannedEffect, +}; + +/// Blocking work accepted by an execution substrate. +pub type BlockingEffectWork = Box; + +/// Substrate seam used to keep provider work outside reconcile transitions. +pub trait BlockingEffectSpawner: Clone + Send + Sync + 'static { + type SpawnError: fmt::Display; + + fn spawn_blocking(&self, work: BlockingEffectWork) -> Result<(), Self::SpawnError>; +} + +/// Provider/bootstrap implementation behind the identity-aware executor. +/// +/// Implementations must use `(run_id, logical_node_id, attempt)` from the effect +/// as the external request identity. Create and bootstrap-start operations must +/// look up and adopt an existing resource for that identity before creating a +/// new one. Cancel and destroy must treat an already-absent resource as success. +pub trait EffectBackend: Send + Sync + 'static { + fn execute(&self, effect: &PlannedEffect) -> Result; +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum EffectFailureDisposition { + Definite, + Ambiguous, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct EffectError { + pub reason: String, + pub disposition: EffectFailureDisposition, +} + +impl EffectError { + pub fn definite(reason: impl Into) -> Self { + Self { + reason: reason.into(), + disposition: EffectFailureDisposition::Definite, + } + } + + pub fn ambiguous(reason: impl Into) -> Self { + Self { + reason: reason.into(), + disposition: EffectFailureDisposition::Ambiguous, + } + } +} + +impl fmt::Display for EffectError { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str(&self.reason) + } +} + +impl std::error::Error for EffectError {} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct ExecutorSubmitError { + pub reason: String, +} + +impl ExecutorSubmitError { + fn new(reason: impl Into) -> Self { + Self { + reason: reason.into(), + } + } +} + +impl fmt::Display for ExecutorSubmitError { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str(&self.reason) + } +} + +impl std::error::Error for ExecutorSubmitError {} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum ExecutorOperationStatus { + Unknown, + InFlight, + Completed, +} + +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub struct AttemptResources { + pub lease_live: bool, + pub bootstrap_live: bool, +} + +// Ledger transitions retain effects inline to avoid a heap allocation per operation. +#[allow(clippy::large_enum_variant)] +#[derive(Clone, Debug)] +enum LedgerEntry { + Running(PlannedEffect), + Queued(PlannedEffect), + Completed { + effect: PlannedEffect, + result: ExecutorResult, + }, +} + +impl LedgerEntry { + fn effect(&self) -> &PlannedEffect { + match self { + Self::Running(effect) | Self::Queued(effect) | Self::Completed { effect, .. } => effect, + } + } +} + +#[derive(Default)] +struct ExecutorLedger { + operations: BTreeMap, + in_flight_by_attempt: BTreeMap, + running_by_attempt: BTreeMap, + queued_by_attempt: BTreeMap, + resources: BTreeMap, +} + +/// Process-local executor that deduplicates operation submissions and runs each +/// accepted operation as substrate-hosted blocking work. +/// +/// One instance belongs to one `ClusterDriver` run. Completed outcomes remain +/// cached for that lifetime, so resubmitting an `OperationId` returns the +/// recorded result without executing the backend again. +pub struct IdempotentEffectExecutor +where + B: EffectBackend, + S: BlockingEffectSpawner, +{ + backend: Arc, + spawner: S, + ledger: Arc>, + result_tx: Sender, + result_rx: Receiver, +} + +impl IdempotentEffectExecutor +where + B: EffectBackend, + S: BlockingEffectSpawner, +{ + pub fn new(backend: B, spawner: S) -> Self { + let (result_tx, result_rx) = mpsc::channel(); + Self { + backend: Arc::new(backend), + spawner, + ledger: Arc::new(Mutex::new(ExecutorLedger::default())), + result_tx, + result_rx, + } + } + + pub fn backend(&self) -> &B { + &self.backend + } + + pub fn drain_results(&mut self) -> Vec { + let mut results = Vec::new(); + loop { + match self.result_rx.try_recv() { + Ok(result) => results.push(result), + Err(TryRecvError::Empty | TryRecvError::Disconnected) => return results, + } + } + } + + pub fn operation_status(&self, operation: OperationId) -> ExecutorOperationStatus { + match lock_ledger(&self.ledger).operations.get(&operation) { + None => ExecutorOperationStatus::Unknown, + Some(LedgerEntry::Running(_) | LedgerEntry::Queued(_)) => { + ExecutorOperationStatus::InFlight + } + Some(LedgerEntry::Completed { .. }) => ExecutorOperationStatus::Completed, + } + } + + /// Classifies a still-running operation as ambiguous after its stored + /// deadline. The backend work is not forcibly cancelled; a late completion + /// is discarded, and any retry must adopt by the stable attempt identity. + pub fn expire(&self, operation: OperationId, reason: impl Into) -> bool { + let result = { + let mut ledger = lock_ledger(&self.ledger); + let effect = match ledger.operations.get(&operation).cloned() { + Some(LedgerEntry::Running(effect) | LedgerEntry::Queued(effect)) => effect, + None | Some(LedgerEntry::Completed { .. }) => return false, + }; + let result = ExecutorResult { + node: effect.node.clone(), + operation, + result: Err(EffectError::ambiguous(reason)), + }; + ledger.operations.insert( + operation, + LedgerEntry::Completed { + effect, + result: result.clone(), + }, + ); + if ledger.in_flight_by_attempt.get(&operation.attempt) == Some(&operation) { + ledger.in_flight_by_attempt.remove(&operation.attempt); + } + if ledger.queued_by_attempt.get(&operation.attempt) == Some(&operation) { + ledger.queued_by_attempt.remove(&operation.attempt); + } + result + }; + let _ = self.result_tx.send(result); + true + } + + pub fn attempt_resources(&self, attempt: NodeAttemptId) -> AttemptResources { + lock_ledger(&self.ledger) + .resources + .get(&attempt) + .copied() + .unwrap_or_default() + } + + fn submit_inner(&mut self, effect: &PlannedEffect) -> Result<(), ExecutorSubmitError> { + let should_spawn = { + let mut ledger = lock_ledger(&self.ledger); + if let Some(entry) = ledger.operations.get(&effect.operation).cloned() { + if entry.effect() != effect { + return Err(ExecutorSubmitError::new(format!( + "operation identity for attempt {} sequence {} was reused with different input", + effect.operation.attempt.0, effect.operation.sequence + ))); + } + return match entry { + LedgerEntry::Completed { result, .. } => { + self.result_tx.send(result).map_err(|_| { + ExecutorSubmitError::new("executor result receiver is closed") + }) + } + LedgerEntry::Running(_) | LedgerEntry::Queued(_) => Ok(()), + }; + } + if let Some(in_flight) = ledger.in_flight_by_attempt.get(&effect.operation.attempt) { + return Err(ExecutorSubmitError::new(format!( + "attempt {} already has operation {} in flight", + effect.operation.attempt.0, in_flight.sequence + ))); + } + + let attempt = effect.operation.attempt; + let physically_running = ledger.running_by_attempt.contains_key(&attempt); + let entry = if physically_running { + if ledger.queued_by_attempt.contains_key(&attempt) { + return Err(ExecutorSubmitError::new(format!( + "attempt {} already has an adoption operation queued", + attempt.0 + ))); + } + ledger.queued_by_attempt.insert(attempt, effect.operation); + LedgerEntry::Queued(effect.clone()) + } else { + ledger.running_by_attempt.insert(attempt, effect.operation); + LedgerEntry::Running(effect.clone()) + }; + ledger.operations.insert(effect.operation, entry); + ledger + .in_flight_by_attempt + .insert(attempt, effect.operation); + !physically_running + }; + + if !should_spawn { + return Ok(()); + } + if let Err(error) = spawn_effect( + self.spawner.clone(), + Arc::clone(&self.backend), + Arc::clone(&self.ledger), + self.result_tx.clone(), + effect.clone(), + ) { + let mut ledger = lock_ledger(&self.ledger); + if matches!( + ledger.operations.get(&effect.operation), + Some(LedgerEntry::Running(running)) if running == effect + ) { + ledger.operations.remove(&effect.operation); + } + if ledger.in_flight_by_attempt.get(&effect.operation.attempt) == Some(&effect.operation) + { + ledger + .in_flight_by_attempt + .remove(&effect.operation.attempt); + } + if ledger.running_by_attempt.get(&effect.operation.attempt) == Some(&effect.operation) { + ledger.running_by_attempt.remove(&effect.operation.attempt); + } + return Err(ExecutorSubmitError::new(format!( + "blocking effect submission failed: {error}" + ))); + } + Ok(()) + } +} + +fn spawn_effect( + spawner: S, + backend: Arc, + ledger: Arc>, + result_tx: Sender, + submitted: PlannedEffect, +) -> Result<(), S::SpawnError> +where + B: EffectBackend, + S: BlockingEffectSpawner, +{ + let nested_spawner = spawner.clone(); + spawner.spawn_blocking(Box::new(move || { + let operation = submitted.operation; + let expected_kind = OperationKind::for_command(&submitted.command); + let result = + std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| backend.execute(&submitted))) + .map_err(|panic| { + EffectError::definite(format!( + "effect backend panicked: {}", + panic_reason(panic) + )) + }) + .and_then(|result| result) + .and_then(|outcome| { + if outcome.kind() == expected_kind { + Ok(outcome) + } else { + Err(EffectError::definite(format!( + "executor returned {:?} for {:?}", + outcome.kind(), + expected_kind + ))) + } + }); + let executor_result = ExecutorResult { + node: submitted.node.clone(), + operation, + result, + }; + let (publish, next) = { + let mut ledger = lock_ledger(&ledger); + let publish = matches!( + ledger.operations.get(&operation), + Some(LedgerEntry::Running(effect)) if effect == &submitted + ); + if publish { + ledger.operations.insert( + operation, + LedgerEntry::Completed { + effect: submitted, + result: executor_result.clone(), + }, + ); + if ledger.in_flight_by_attempt.get(&operation.attempt) == Some(&operation) { + ledger.in_flight_by_attempt.remove(&operation.attempt); + } + if let Ok(outcome) = &executor_result.result { + update_resources(&mut ledger, operation.attempt, outcome); + } + } + if ledger.running_by_attempt.get(&operation.attempt) == Some(&operation) { + ledger.running_by_attempt.remove(&operation.attempt); + } + let next = promote_queued(&mut ledger, operation.attempt); + (publish, next) + }; + if publish { + let _ = result_tx.send(executor_result); + } + if let Some(next) = next { + spawn_promoted(nested_spawner, backend, ledger, result_tx, next); + } + })) +} + +fn spawn_promoted( + spawner: S, + backend: Arc, + ledger: Arc>, + result_tx: Sender, + effect: PlannedEffect, +) where + B: EffectBackend, + S: BlockingEffectSpawner, +{ + if let Err(error) = spawn_effect( + spawner.clone(), + Arc::clone(&backend), + Arc::clone(&ledger), + result_tx.clone(), + effect.clone(), + ) { + let (result, next) = record_spawn_failure( + &ledger, + &effect, + format!("blocking effect submission failed: {error}"), + ); + if let Some(result) = result { + let _ = result_tx.send(result); + } + if let Some(next) = next { + spawn_promoted(spawner, backend, ledger, result_tx, next); + } + } +} + +fn record_spawn_failure( + ledger: &Mutex, + effect: &PlannedEffect, + reason: String, +) -> (Option, Option) { + let mut ledger = lock_ledger(ledger); + let operation = effect.operation; + let publish = matches!( + ledger.operations.get(&operation), + Some(LedgerEntry::Running(running)) if running == effect + ); + let result = publish.then(|| ExecutorResult { + node: effect.node.clone(), + operation, + result: Err(EffectError::definite(reason)), + }); + if let Some(result) = result.as_ref() { + ledger.operations.insert( + operation, + LedgerEntry::Completed { + effect: effect.clone(), + result: result.clone(), + }, + ); + } + if ledger.in_flight_by_attempt.get(&operation.attempt) == Some(&operation) { + ledger.in_flight_by_attempt.remove(&operation.attempt); + } + if ledger.running_by_attempt.get(&operation.attempt) == Some(&operation) { + ledger.running_by_attempt.remove(&operation.attempt); + } + let next = promote_queued(&mut ledger, operation.attempt); + (result, next) +} + +fn promote_queued(ledger: &mut ExecutorLedger, attempt: NodeAttemptId) -> Option { + let operation = ledger.queued_by_attempt.remove(&attempt)?; + let Some(LedgerEntry::Queued(effect)) = ledger.operations.get(&operation).cloned() else { + return None; + }; + ledger + .operations + .insert(operation, LedgerEntry::Running(effect.clone())); + ledger.running_by_attempt.insert(attempt, operation); + Some(effect) +} + +impl crate::reconciler::EffectExecutor for IdempotentEffectExecutor +where + B: EffectBackend, + S: BlockingEffectSpawner, +{ + type SubmitError = ExecutorSubmitError; + + fn submit(&mut self, effect: &PlannedEffect) -> Result<(), Self::SubmitError> { + self.submit_inner(effect) + } +} + +fn lock_ledger(ledger: &Mutex) -> MutexGuard<'_, ExecutorLedger> { + ledger + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) +} + +fn panic_reason(panic: Box) -> String { + if let Some(reason) = panic.downcast_ref::<&str>() { + (*reason).to_owned() + } else if let Some(reason) = panic.downcast_ref::() { + reason.clone() + } else { + "non-string panic payload".to_owned() + } +} + +fn update_resources( + ledger: &mut ExecutorLedger, + attempt: NodeAttemptId, + outcome: &OperationOutcome, +) { + let resources = ledger.resources.entry(attempt).or_default(); + match outcome { + OperationOutcome::LeaseCreated(_) => resources.lease_live = true, + OperationOutcome::BootstrapStarted { .. } => resources.bootstrap_live = true, + OperationOutcome::BootstrapCancelled => resources.bootstrap_live = false, + OperationOutcome::LeaseDestroyed => { + resources.bootstrap_live = false; + resources.lease_live = false; + } + OperationOutcome::EndpointLookup(_) | OperationOutcome::BootstrapConvergenceAccepted => {} + } +} diff --git a/crates/provisioning/src/lib.rs b/crates/provisioning/src/lib.rs index 098c509..6b065a7 100644 --- a/crates/provisioning/src/lib.rs +++ b/crates/provisioning/src/lib.rs @@ -1,11 +1,15 @@ //! Reusable provider-neutral provisioning contracts. //! -//! This crate owns lease, boot, destroy, bootstrap-session, and provider plugin -//! state machines that do not depend on GGUF, prompts, stage execution, Docker, -//! VastAI, or MVP runtime policy. Concrete provider adapters live in application -//! crates and implement these contracts. +//! This crate owns provider-neutral lifecycle facts and commands together with +//! the level-triggered reconciler and identity-aware executor contracts. Concrete +//! provider adapters live in application crates and implement the executor +//! backend contract. +pub mod executor; pub mod node; pub mod plugin; +pub mod reconciler; +pub use executor::*; pub use node::*; +pub use reconciler::*; diff --git a/crates/provisioning/src/node.rs b/crates/provisioning/src/node.rs index 60c06d8..d277cca 100644 --- a/crates/provisioning/src/node.rs +++ b/crates/provisioning/src/node.rs @@ -1,13 +1,12 @@ -//! Provider-neutral node lease, bootstrap, and destroy state machines. +//! Provider-neutral node lease, bootstrap, and destroy contracts. //! -//! The module is intentionally in-process and deterministic. `NodeManager` is the -//! actor core: it owns one node record and emits commands for provider and -//! bootstrap effects. `BootstrapSession` is the transient pre-swactor SSH core. -//! Tests can drive both without Vast.ai, Docker, or real SSH. +//! `reconciler` owns lifecycle observation and effect selection. This module +//! retains the fact, command, provider, and bootstrap-session vocabulary. -use std::collections::{BTreeMap, VecDeque}; +use std::collections::BTreeMap; use std::time::SystemTime; +use crate::plugin::ProviderMount; use serde::{Deserialize, Serialize}; #[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)] @@ -67,6 +66,12 @@ pub struct BootSpec { pub start_swactor_command: String, pub stdout_sources: Vec, pub stderr_sources: Vec, + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub env: Vec<(String, String)>, + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub args: Vec, + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub mounts: Vec, } #[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] @@ -216,6 +221,7 @@ pub struct NodeRecord { pub bootstrap: Option, pub swactor: Option, pub failed_reason: Option, + pub failed_at: Option, pub destroyed_at: Option, } @@ -234,6 +240,7 @@ impl NodeRecord { bootstrap: None, swactor: None, failed_reason: None, + failed_at: None, destroyed_at: None, } } @@ -250,31 +257,6 @@ pub struct CreateLeaseResult { pub endpoint: Option, } -#[derive(Clone, Debug, PartialEq, Eq)] -pub struct ProviderError { - pub reason: String, -} - -impl ProviderError { - pub fn new(reason: impl Into) -> Self { - Self { - reason: reason.into(), - } - } -} - -pub trait ProviderPlugin { - fn create_lease( - &mut self, - request: CreateLeaseRequest, - ) -> Result; - - fn lookup_endpoint(&mut self, lease: &LeaseFacts) - -> Result, ProviderError>; - - fn destroy_lease(&mut self, handle: &DestroyHandle) -> Result<(), ProviderError>; -} - #[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] pub struct BootstrapObservation { pub stage: BootstrapStage, @@ -294,25 +276,6 @@ impl BootstrapObservation { } } -#[derive(Clone, Debug, PartialEq)] -pub enum NodeManagerMsg { - Start(LogicalNodeSpec), - LeaseCreated(CreateLeaseResult), - LeaseFailed(String), - EndpointKnown(SshEndpoint), - EndpointFailed(String), - BootstrapObserved(BootstrapObservation), - BootstrapFailed(String), - BootstrapClosed, - SwactorJoined { - logical_node_id: LogicalNodeId, - swactor_id: SwactorId, - }, - Destroy, - LeaseDestroyed, - DestroyFailed(String), -} - #[derive(Clone, Debug, PartialEq)] pub enum NodeManagerCommand { CreateLease(CreateLeaseRequest), @@ -328,361 +291,6 @@ pub enum NodeManagerCommand { DestroyLease(DestroyHandle), } -#[derive(Clone, Debug, PartialEq, Eq)] -pub struct NodeManagerError { - pub reason: String, -} - -impl NodeManagerError { - fn new(reason: impl Into) -> Self { - Self { - reason: reason.into(), - } - } -} - -#[derive(Clone, Debug)] -pub struct NodeManager { - record: Option, - active_bootstrap: Option, - next_bootstrap_session_id: u64, - lease_destroyed: bool, -} - -impl Default for NodeManager { - fn default() -> Self { - Self { - record: None, - active_bootstrap: None, - next_bootstrap_session_id: 1, - lease_destroyed: false, - } - } -} - -impl NodeManager { - pub fn new() -> Self { - Self::default() - } - - pub fn record(&self) -> Option<&NodeRecord> { - self.record.as_ref() - } - - pub fn is_ready(&self) -> bool { - self.record.as_ref().is_some_and(|record| record.ready) - } - - pub fn active_bootstrap(&self) -> Option { - self.active_bootstrap - } - - pub fn handle( - &mut self, - msg: NodeManagerMsg, - ) -> Result, NodeManagerError> { - match msg { - NodeManagerMsg::Start(spec) => self.start(spec), - NodeManagerMsg::LeaseCreated(result) => self.lease_created(result), - NodeManagerMsg::LeaseFailed(reason) => self.fail(reason), - NodeManagerMsg::EndpointKnown(endpoint) => self.endpoint_known(endpoint), - NodeManagerMsg::EndpointFailed(reason) => self.fail(reason), - NodeManagerMsg::BootstrapObserved(observation) => self.bootstrap_observed(observation), - NodeManagerMsg::BootstrapFailed(reason) => self.fail(reason), - NodeManagerMsg::BootstrapClosed => self.bootstrap_closed(), - NodeManagerMsg::SwactorJoined { - logical_node_id, - swactor_id, - } => self.swactor_joined(logical_node_id, swactor_id), - NodeManagerMsg::Destroy => self.destroy(), - NodeManagerMsg::LeaseDestroyed => self.lease_destroyed(), - NodeManagerMsg::DestroyFailed(reason) => self.destroy_failed(reason), - } - } - - fn start( - &mut self, - spec: LogicalNodeSpec, - ) -> Result, NodeManagerError> { - if self.record.is_some() { - return Err(NodeManagerError::new("node manager already started")); - } - let mut record = NodeRecord::from_spec(spec.clone()); - record.stage = NodeStage::LeaseRequested; - self.record = Some(record); - Ok(vec![NodeManagerCommand::CreateLease(CreateLeaseRequest { - spec, - })]) - } - - fn lease_created( - &mut self, - result: CreateLeaseResult, - ) -> Result, NodeManagerError> { - let record = self.record_mut()?; - record.lease = Some(result.lease.clone()); - record.stage = NodeStage::LeaseCreated; - if let Some(endpoint) = result.endpoint { - self.begin_bootstrap(endpoint) - } else { - Ok(vec![NodeManagerCommand::LookupEndpoint(result.lease)]) - } - } - - fn endpoint_known( - &mut self, - endpoint: SshEndpoint, - ) -> Result, NodeManagerError> { - if self.require_record()?.lease.is_none() { - return Err(NodeManagerError::new("endpoint cannot arrive before lease")); - } - self.begin_bootstrap(endpoint) - } - - fn begin_bootstrap( - &mut self, - endpoint: SshEndpoint, - ) -> Result, NodeManagerError> { - if self.active_bootstrap.is_some() { - return Err(NodeManagerError::new("bootstrap already active")); - } - let session_id = BootstrapSessionId(self.next_bootstrap_session_id); - self.next_bootstrap_session_id = self.next_bootstrap_session_id.wrapping_add(1).max(1); - let (run_id, logical_node_id, lease_id, boot, swarm_join) = { - let record = self.record_mut()?; - let lease_id = record - .lease - .as_ref() - .ok_or_else(|| NodeManagerError::new("bootstrap requires known lease"))? - .lease_id - .clone(); - record.connection = Some(endpoint.clone()); - record.stage = NodeStage::EndpointKnown; - record.bootstrap = Some(BootstrapFacts { - session_id, - last_stage: BootstrapStage::Created, - last_stdout_seq: None, - last_stderr_seq: None, - last_observed_at: SystemTime::now(), - }); - record.stage = NodeStage::BootstrapRunning; - ( - record.run_id.clone(), - record.logical_node_id.clone(), - lease_id, - record.desired.boot.clone(), - record.desired.swarm_join.clone(), - ) - }; - self.active_bootstrap = Some(session_id); - Ok(vec![NodeManagerCommand::StartBootstrap( - BootstrapSessionSpec { - datastream: DatastreamStreamId(format!( - "run/{}/node/{}/bootstrap", - run_id.0, logical_node_id.0 - )), - run_id, - logical_node_id, - lease_id, - ssh: endpoint, - boot, - swarm_join, - }, - )]) - } - - fn bootstrap_observed( - &mut self, - observation: BootstrapObservation, - ) -> Result, NodeManagerError> { - if self.active_bootstrap.is_none() { - return Err(NodeManagerError::new( - "bootstrap observation without active session", - )); - } - let record = self.record_mut()?; - let facts = record - .bootstrap - .as_mut() - .ok_or_else(|| NodeManagerError::new("missing bootstrap facts"))?; - facts.last_stage = observation.stage; - facts.last_observed_at = SystemTime::now(); - if observation.last_stdout_seq.is_some() { - facts.last_stdout_seq = observation.last_stdout_seq; - } - if observation.last_stderr_seq.is_some() { - facts.last_stderr_seq = observation.last_stderr_seq; - } - Ok(Vec::new()) - } - - fn swactor_joined( - &mut self, - logical_node_id: LogicalNodeId, - swactor_id: SwactorId, - ) -> Result, NodeManagerError> { - let expected = self.require_record()?.logical_node_id.clone(); - if logical_node_id != expected { - return Err(NodeManagerError::new(format!( - "swactor join for {}, expected {}", - logical_node_id.0, expected.0 - ))); - } - let session_id = self - .active_bootstrap - .ok_or_else(|| NodeManagerError::new("swactor join without active bootstrap"))?; - let record = self.record_mut()?; - record.swactor = Some(SwactorFacts { - swactor_id: swactor_id.clone(), - joined_at: SystemTime::now(), - handed_off_at: None, - }); - record.stage = NodeStage::SwactorJoined; - Ok(vec![NodeManagerCommand::BootstrapConvergenceObserved { - session_id, - swactor_id, - }]) - } - - fn bootstrap_closed(&mut self) -> Result, NodeManagerError> { - let record = self.record_mut()?; - if record.stage != NodeStage::SwactorJoined { - return Err(NodeManagerError::new( - "bootstrap closed before swactor convergence", - )); - } - let swactor = record - .swactor - .as_mut() - .ok_or_else(|| NodeManagerError::new("handoff requires swactor facts"))?; - swactor.handed_off_at = Some(SystemTime::now()); - record.stage = NodeStage::HandedOff; - record.ready = true; - record.stage = NodeStage::Dormant; - self.active_bootstrap = None; - Ok(Vec::new()) - } - - fn destroy(&mut self) -> Result, NodeManagerError> { - let mut commands = Vec::new(); - let active_bootstrap = self.active_bootstrap.take(); - let lease_already_destroyed = self.lease_destroyed; - let lease_command = { - let record = self.record_mut()?; - record.ready = false; - if let Some(session_id) = active_bootstrap { - commands.push(NodeManagerCommand::CancelBootstrap { session_id }); - } - if lease_already_destroyed { - None - } else { - record - .lease - .as_ref() - .map(|lease| NodeManagerCommand::DestroyLease(lease.destroy_handle.clone())) - } - }; - if let Some(command) = lease_command { - commands.push(command); - } else { - let record = self.record_mut()?; - record.stage = NodeStage::Destroyed; - record.destroyed_at = Some(SystemTime::now()); - } - Ok(commands) - } - - fn lease_destroyed(&mut self) -> Result, NodeManagerError> { - let record = self.record_mut()?; - record.stage = NodeStage::Destroyed; - record.ready = false; - record.destroyed_at = Some(SystemTime::now()); - self.lease_destroyed = true; - Ok(Vec::new()) - } - - fn destroy_failed( - &mut self, - reason: String, - ) -> Result, NodeManagerError> { - self.fail(format!("destroy: {reason}")) - } - - fn fail(&mut self, reason: String) -> Result, NodeManagerError> { - let record = self.record_mut()?; - record.stage = NodeStage::Failed; - record.ready = false; - record.failed_reason = Some(reason); - Ok(Vec::new()) - } - - fn require_record(&self) -> Result<&NodeRecord, NodeManagerError> { - self.record - .as_ref() - .ok_or_else(|| NodeManagerError::new("node manager not started")) - } - - fn record_mut(&mut self) -> Result<&mut NodeRecord, NodeManagerError> { - self.record - .as_mut() - .ok_or_else(|| NodeManagerError::new("node manager not started")) - } -} - -#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] -pub enum BootstrapLogSource { - SshBootstrap, -} - -#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] -pub enum BootstrapLogStream { - Stdout, - Stderr, -} - -#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] -pub struct BootstrapLogRecord { - pub run_id: RunId, - pub logical_node_id: LogicalNodeId, - pub lease_id: ProviderLeaseId, - pub source: BootstrapLogSource, - pub stream: BootstrapLogStream, - pub seq: u64, - pub timestamp: SystemTime, - pub line: String, -} - -pub trait BootstrapDatastreamSink { - fn record(&mut self, record: BootstrapLogRecord); - fn flush(&mut self); -} - -#[derive(Default, Debug, Clone, PartialEq, Eq)] -pub struct InMemoryBootstrapDatastream { - records: Vec, - flush_count: usize, -} - -impl InMemoryBootstrapDatastream { - pub fn records(&self) -> &[BootstrapLogRecord] { - &self.records - } - - pub fn flush_count(&self) -> usize { - self.flush_count - } -} - -impl BootstrapDatastreamSink for InMemoryBootstrapDatastream { - fn record(&mut self, record: BootstrapLogRecord) { - self.records.push(record); - } - - fn flush(&mut self) { - self.flush_count += 1; - } -} - #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] pub struct BootstrapSessionSpec { pub run_id: RunId, @@ -694,269 +302,6 @@ pub struct BootstrapSessionSpec { pub datastream: DatastreamStreamId, } -#[derive(Clone, Debug, PartialEq, Eq)] -pub enum BootstrapSessionEvent { - Observed(BootstrapObservation), - Failed(String), - Closed, -} - -#[derive(Clone, Debug, PartialEq, Eq)] -pub struct MockBootstrapScript { - pub ssh_ok: bool, - pub verify_ok: bool, - pub start_ok: bool, - pub records: Vec<(BootstrapLogStream, String)>, -} - -impl MockBootstrapScript { - pub fn successful(records: Vec<(BootstrapLogStream, String)>) -> Self { - Self { - ssh_ok: true, - verify_ok: true, - start_ok: true, - records, - } - } -} - -#[derive(Clone, Debug)] -pub struct BootstrapSession { - spec: BootstrapSessionSpec, - stage: BootstrapStage, - next_seq: u64, - last_stdout_seq: Option, - last_stderr_seq: Option, - closed: bool, -} - -impl BootstrapSession { - pub fn new(spec: BootstrapSessionSpec) -> Self { - Self { - spec, - stage: BootstrapStage::Created, - next_seq: 1, - last_stdout_seq: None, - last_stderr_seq: None, - closed: false, - } - } - - pub fn stage(&self) -> BootstrapStage { - self.stage - } - - pub fn is_closed(&self) -> bool { - self.closed - } - - pub fn start( - &mut self, - script: &MockBootstrapScript, - sink: &mut dyn BootstrapDatastreamSink, - ) -> Vec { - let mut events = Vec::new(); - self.stage = BootstrapStage::SshConnecting; - if !script.ssh_ok { - self.stage = BootstrapStage::SshConnectFailed; - return vec![BootstrapSessionEvent::Failed("ssh connect failed".into())]; - } - self.stage = BootstrapStage::SshReady; - events.push(BootstrapSessionEvent::Observed( - BootstrapObservation::stage(BootstrapStage::SshReady), - )); - - self.stage = BootstrapStage::StdoutStreaming; - for (stream, line) in &script.records { - let seq = self.next_seq; - self.next_seq += 1; - match stream { - BootstrapLogStream::Stdout => self.last_stdout_seq = Some(seq), - BootstrapLogStream::Stderr => self.last_stderr_seq = Some(seq), - } - sink.record(BootstrapLogRecord { - run_id: self.spec.run_id.clone(), - logical_node_id: self.spec.logical_node_id.clone(), - lease_id: self.spec.lease_id.clone(), - source: BootstrapLogSource::SshBootstrap, - stream: *stream, - seq, - timestamp: SystemTime::now(), - line: line.clone(), - }); - } - events.push(BootstrapSessionEvent::Observed(BootstrapObservation { - stage: BootstrapStage::StdoutStreaming, - last_stdout_seq: self.last_stdout_seq, - last_stderr_seq: self.last_stderr_seq, - marker: None, - })); - - self.stage = BootstrapStage::BootChecking; - events.push(BootstrapSessionEvent::Observed( - BootstrapObservation::stage(BootstrapStage::BootChecking), - )); - if !script.verify_ok { - self.stage = BootstrapStage::BootCheckFailed; - events.push(BootstrapSessionEvent::Failed("boot check failed".into())); - return events; - } - - self.stage = BootstrapStage::SwactorStarting; - events.push(BootstrapSessionEvent::Observed( - BootstrapObservation::stage(BootstrapStage::SwactorStarting), - )); - if !script.start_ok { - self.stage = BootstrapStage::StartFailed; - events.push(BootstrapSessionEvent::Failed("swactor start failed".into())); - return events; - } - - self.stage = BootstrapStage::WaitingForSwactorJoin; - events.push(BootstrapSessionEvent::Observed( - BootstrapObservation::stage(BootstrapStage::WaitingForSwactorJoin), - )); - events - } - - pub fn convergence_observed( - &mut self, - _swactor_id: SwactorId, - sink: &mut dyn BootstrapDatastreamSink, - ) -> Vec { - self.stage = BootstrapStage::Converged; - sink.flush(); - self.closed = true; - self.stage = BootstrapStage::Closed; - vec![ - BootstrapSessionEvent::Observed(BootstrapObservation::stage(BootstrapStage::Converged)), - BootstrapSessionEvent::Closed, - ] - } - - pub fn join_failed(&mut self, reason: impl Into) -> Vec { - self.stage = BootstrapStage::SwactorJoinFailed; - vec![BootstrapSessionEvent::Failed(reason.into())] - } - - pub fn cancel(&mut self) -> Vec { - self.stage = BootstrapStage::Cancelled; - self.closed = true; - vec![BootstrapSessionEvent::Closed] - } -} - -#[derive(Default, Debug, Clone)] -pub struct MockProviderPlugin { - next_contract_id: u64, - create_results: VecDeque>, - endpoint_results: VecDeque, ProviderError>>, - create_requests: Vec, - lookup_requests: Vec, - destroyed_handles: Vec, - destroy_failures: BTreeMap, -} - -impl MockProviderPlugin { - pub fn new() -> Self { - Self { - next_contract_id: 1, - ..Self::default() - } - } - - pub fn queue_create_result(&mut self, result: Result) { - self.create_results.push_back(result); - } - - pub fn queue_endpoint_result(&mut self, result: Result, ProviderError>) { - self.endpoint_results.push_back(result); - } - - pub fn fail_destroy(&mut self, lease_id: ProviderLeaseId, reason: impl Into) { - self.destroy_failures.insert(lease_id, reason.into()); - } - - pub fn create_requests(&self) -> &[CreateLeaseRequest] { - &self.create_requests - } - - pub fn lookup_requests(&self) -> &[ProviderLeaseId] { - &self.lookup_requests - } - - pub fn destroyed_handles(&self) -> &[DestroyHandle] { - &self.destroyed_handles - } - - pub fn result_with_endpoint(id: u64, endpoint: Option) -> CreateLeaseResult { - let lease_id = ProviderLeaseId(format!("mock:{id}")); - CreateLeaseResult { - lease: LeaseFacts { - provider: ProviderKind::new("mock"), - lease_id: lease_id.clone(), - provider_contract_id: id.to_string(), - offer_id: Some(format!("offer-{id}")), - destroy_handle: DestroyHandle { - provider: ProviderKind::new("mock"), - lease_id, - provider_contract_id: id.to_string(), - }, - provider_metadata: BTreeMap::new(), - }, - endpoint, - } - } - - fn default_endpoint(id: u64) -> SshEndpoint { - SshEndpoint { - host: "127.0.0.1".into(), - port: 22000 + id as u16, - user: "root".into(), - auth_ref: format!("mock-auth-{id}"), - } - } -} - -impl ProviderPlugin for MockProviderPlugin { - fn create_lease( - &mut self, - request: CreateLeaseRequest, - ) -> Result { - self.create_requests.push(request); - if let Some(result) = self.create_results.pop_front() { - return result; - } - let id = self.next_contract_id; - self.next_contract_id = self.next_contract_id.wrapping_add(1).max(1); - Ok(Self::result_with_endpoint( - id, - Some(Self::default_endpoint(id)), - )) - } - - fn lookup_endpoint( - &mut self, - lease: &LeaseFacts, - ) -> Result, ProviderError> { - self.lookup_requests.push(lease.lease_id.clone()); - if let Some(result) = self.endpoint_results.pop_front() { - return result; - } - Ok(Some(Self::default_endpoint( - lease.provider_contract_id.parse().unwrap_or(1), - ))) - } - - fn destroy_lease(&mut self, handle: &DestroyHandle) -> Result<(), ProviderError> { - if let Some(reason) = self.destroy_failures.get(&handle.lease_id) { - return Err(ProviderError::new(reason.clone())); - } - self.destroyed_handles.push(handle.clone()); - Ok(()) - } -} - #[cfg(test)] mod tests { use super::*; diff --git a/crates/provisioning/src/plugin.rs b/crates/provisioning/src/plugin.rs index fa738dd..727d58c 100644 --- a/crates/provisioning/src/plugin.rs +++ b/crates/provisioning/src/plugin.rs @@ -6,6 +6,9 @@ use serde::{Deserialize, Serialize}; pub struct NodeProvisionSpec { pub run_id: u64, pub node_id: u64, + /// Concrete attempt identity. Zero is reserved for an unbound template. + #[serde(default)] + pub attempt_id: u64, pub stage_index: Option, pub image: String, pub env: Vec<(String, String)>, @@ -115,24 +118,18 @@ impl PluginSink { } pub trait ProvisionPlugin: Send { - fn start_node( + /// Acquires or adopts the provider resource for one concrete node attempt. + fn create_node( &mut self, spec: NodeProvisionSpec, sink: PluginSink, ) -> Result; - fn start_nodes( - &mut self, - specs: Vec, - sink: PluginSink, - ) -> Vec<(NodeProvisionSpec, Result)> { - specs - .into_iter() - .map(|spec| { - let result = self.start_node(spec.clone(), sink.clone()); - (spec, result) - }) - .collect() + /// Starts or adopts bootstrap work on an already-created provider resource. + fn start_bootstrap(&mut self, handle: &PluginNodeHandle) -> Result<(), String>; + + fn cancel_bootstrap(&mut self, _handle: &PluginNodeHandle) -> Result<(), String> { + Ok(()) } fn complete_bootstrap(&mut self, handle: &PluginNodeHandle) -> Result<(), String>; @@ -161,6 +158,7 @@ mod tests { let spec = NodeProvisionSpec { run_id: 17, node_id: 23, + attempt_id: 7, stage_index: Some(2), image: "runtime:latest".to_owned(), env: vec![("A".to_owned(), "B".to_owned())], @@ -191,6 +189,7 @@ mod tests { .expect("legacy spec decodes"); assert!(decoded.mounts.is_empty()); + assert_eq!(decoded.attempt_id, 0); } #[test] diff --git a/crates/provisioning/src/reconciler.rs b/crates/provisioning/src/reconciler.rs new file mode 100644 index 0000000..141e7e3 --- /dev/null +++ b/crates/provisioning/src/reconciler.rs @@ -0,0 +1,1347 @@ +//! Level-triggered cluster reconciliation over the provider-neutral node lifecycle. + +use std::collections::{BTreeMap, BTreeSet}; +use std::fmt; +use std::time::{Duration, SystemTime}; + +use crate::executor::{EffectError, EffectFailureDisposition}; +use serde::{Deserialize, Serialize}; + +use crate::node::{ + BootstrapFacts, BootstrapObservation, BootstrapSessionId, BootstrapStage, CreateLeaseResult, + LogicalNodeId, LogicalNodeSpec, NodeManagerCommand, NodeRecord, NodeStage, RunId, + RunNodeGroupSpec, SshEndpoint, SwactorFacts, SwactorId, expand_node_group, +}; + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct ClusterShape { + pub run_id: RunId, + pub generation: u64, + pub groups: Vec, +} + +impl ClusterShape { + pub fn expand(&self) -> Result, ShapeError> { + let mut group_ids = BTreeSet::new(); + let mut nodes = BTreeMap::new(); + + for group in &self.groups { + if group.run_id != self.run_id { + return Err(ShapeError::new(format!( + "group {} belongs to run {}, expected {}", + group.group_id.0, group.run_id.0, self.run_id.0 + ))); + } + if !group_ids.insert(group.group_id.clone()) { + return Err(ShapeError::new(format!( + "duplicate node group {}", + group.group_id.0 + ))); + } + for (field, value) in [ + ("min_down_mbps", group.shape.min_down_mbps), + ("min_up_mbps", group.shape.min_up_mbps), + ("min_reliability", group.shape.min_reliability), + ] { + if value.is_some_and(|value| !value.is_finite()) { + return Err(ShapeError::new(format!( + "node group {} has non-finite {field}", + group.group_id.0 + ))); + } + } + for node in expand_node_group(group) { + let id = node.logical_node_id.clone(); + if nodes.insert(id.clone(), node).is_some() { + return Err(ShapeError::new(format!("duplicate logical node {}", id.0))); + } + } + } + + Ok(nodes) + } +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct ShapeError { + pub reason: String, +} + +impl ShapeError { + fn new(reason: impl Into) -> Self { + Self { + reason: reason.into(), + } + } +} + +impl fmt::Display for ShapeError { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str(&self.reason) + } +} + +impl std::error::Error for ShapeError {} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)] +pub struct NodeAttemptId(pub u64); + +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)] +pub struct OperationId { + pub attempt: NodeAttemptId, + pub sequence: u64, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub enum NodeIntent { + Active, + Deleting, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub enum OperationKind { + CreateLease, + LookupEndpoint, + StartBootstrap, + BootstrapConvergenceObserved, + CancelBootstrap, + DestroyLease, +} + +impl OperationKind { + pub(crate) fn for_command(command: &NodeManagerCommand) -> Self { + match command { + NodeManagerCommand::CreateLease(_) => Self::CreateLease, + NodeManagerCommand::LookupEndpoint(_) => Self::LookupEndpoint, + NodeManagerCommand::StartBootstrap(_) => Self::StartBootstrap, + NodeManagerCommand::BootstrapConvergenceObserved { .. } => { + Self::BootstrapConvergenceObserved + } + NodeManagerCommand::CancelBootstrap { .. } => Self::CancelBootstrap, + NodeManagerCommand::DestroyLease(_) => Self::DestroyLease, + } + } +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct PendingOperation { + pub id: OperationId, + pub kind: OperationKind, + pub deadline: SystemTime, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct PlannedOperation { + pub node: LogicalNodeId, + pub operation: OperationId, + pub kind: OperationKind, + pub deadline: SystemTime, +} + +#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)] +pub struct RetryState { + pub consecutive_failures: u32, + pub next_effect_at: Option, + pub restart_at: Option, + pub last_error: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub ambiguous_operation: Option, +} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct ManagedNode { + pub attempt: NodeAttemptId, + pub intent: NodeIntent, + pub record: NodeRecord, + pub active_bootstrap: Option, + pub pending: Option, + pub next_operation_sequence: u64, + pub retry: RetryState, +} + +impl ManagedNode { + pub fn new(attempt: NodeAttemptId, desired: LogicalNodeSpec) -> Self { + Self { + attempt, + intent: NodeIntent::Active, + record: NodeRecord::from_spec(desired), + active_bootstrap: None, + pending: None, + next_operation_sequence: 1, + retry: RetryState::default(), + } + } +} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct ClusterState { + pub observed_generation: u64, + pub next_attempt_id: u64, + pub nodes: BTreeMap, +} + +impl Default for ClusterState { + fn default() -> Self { + Self { + observed_generation: 0, + next_attempt_id: 1, + nodes: BTreeMap::new(), + } + } +} + +#[derive(Clone, Debug, PartialEq)] +pub struct ReconcilePlan { + pub actions: Vec, + pub observed_generation: u64, + pub requeue_at: Option, +} + +#[derive(Clone, Debug, PartialEq)] +pub enum NodeAction { + Insert { + attempt: NodeAttemptId, + desired: LogicalNodeSpec, + }, + BeginDelete { + node: LogicalNodeId, + expected_attempt: NodeAttemptId, + }, + MarkDestroyed { + node: LogicalNodeId, + expected_attempt: NodeAttemptId, + }, + Restart { + node: LogicalNodeId, + expected_attempt: NodeAttemptId, + new_attempt: NodeAttemptId, + desired: LogicalNodeSpec, + }, + Reap { + node: LogicalNodeId, + expected_attempt: NodeAttemptId, + }, + Dispatch(PlannedEffect), +} + +impl NodeAction { + pub fn node(&self) -> &LogicalNodeId { + match self { + Self::Insert { desired, .. } => &desired.logical_node_id, + Self::BeginDelete { node, .. } + | Self::MarkDestroyed { node, .. } + | Self::Restart { node, .. } + | Self::Reap { node, .. } => node, + Self::Dispatch(effect) => &effect.node, + } + } +} + +#[derive(Clone, Debug, PartialEq)] +pub struct PlannedEffect { + pub node: LogicalNodeId, + pub operation: OperationId, + pub command: NodeManagerCommand, +} + +/// Non-blocking submission boundary. Implementations queue engine-hosted work; +/// provider I/O must not execute inline in `submit`. +pub trait EffectExecutor { + type SubmitError: fmt::Display; + + fn submit(&mut self, effect: &PlannedEffect) -> Result<(), Self::SubmitError>; +} + +#[derive(Clone, Debug, PartialEq)] +pub struct NodeDecision { + pub action: Option, + pub requeue_at: Option, +} + +impl NodeDecision { + fn action(action: NodeDecisionAction) -> Self { + Self { + action: Some(action), + requeue_at: None, + } + } + + fn wait(requeue_at: Option) -> Self { + Self { + action: None, + requeue_at, + } + } +} + +#[derive(Clone, Debug, PartialEq)] +pub enum NodeDecisionAction { + BeginDelete, + MarkDestroyed, + Restart(LogicalNodeSpec), + Reap, + Dispatch(NodeManagerCommand), +} + +pub fn reconcile( + observed: &ClusterState, + desired: &ClusterShape, + now: SystemTime, +) -> Result { + let desired_nodes = desired.expand()?; + let keys: BTreeSet<_> = observed + .nodes + .keys() + .chain(desired_nodes.keys()) + .cloned() + .collect(); + let mut next_attempt = observed.next_attempt_id; + let mut actions = Vec::new(); + let mut requeue_at = None; + + for id in keys { + match (observed.nodes.get(&id), desired_nodes.get(&id)) { + (None, Some(spec)) => { + let attempt = allocate_attempt(&mut next_attempt)?; + actions.push(NodeAction::Insert { + attempt, + desired: spec.clone(), + }); + } + (Some(node), desired) => { + let decision = reconcile_node(node, desired, now); + requeue_at = earliest(requeue_at, decision.requeue_at); + let Some(action) = decision.action else { + continue; + }; + let action = match action { + NodeDecisionAction::BeginDelete => NodeAction::BeginDelete { + node: id, + expected_attempt: node.attempt, + }, + NodeDecisionAction::MarkDestroyed => NodeAction::MarkDestroyed { + node: id, + expected_attempt: node.attempt, + }, + NodeDecisionAction::Restart(spec) => NodeAction::Restart { + node: id, + expected_attempt: node.attempt, + new_attempt: allocate_attempt(&mut next_attempt)?, + desired: spec, + }, + NodeDecisionAction::Reap => NodeAction::Reap { + node: id, + expected_attempt: node.attempt, + }, + NodeDecisionAction::Dispatch(command) => NodeAction::Dispatch(PlannedEffect { + node: id, + operation: OperationId { + attempt: node.attempt, + sequence: node.next_operation_sequence, + }, + command, + }), + }; + actions.push(action); + } + (None, None) => unreachable!("union key must exist in one map"), + } + } + + Ok(ReconcilePlan { + actions, + observed_generation: desired.generation, + requeue_at, + }) +} + +fn allocate_attempt(next: &mut u64) -> Result { + let attempt = NodeAttemptId(*next); + *next = next + .checked_add(1) + .ok_or_else(|| ShapeError::new("node attempt allocator exhausted"))?; + Ok(attempt) +} + +pub fn reconcile_node( + node: &ManagedNode, + desired: Option<&LogicalNodeSpec>, + now: SystemTime, +) -> NodeDecision { + if node.record.stage == NodeStage::Destroyed { + return match desired { + None => NodeDecision::action(NodeDecisionAction::Reap), + Some(spec) => { + if let Some(restart_at) = node.retry.restart_at + && now < restart_at + { + NodeDecision::wait(Some(restart_at)) + } else { + NodeDecision::action(NodeDecisionAction::Restart(spec.clone())) + } + } + }; + } + + if node.intent == NodeIntent::Active + && (desired.is_none() + || desired.is_some_and(|spec| spec != &node.record.desired) + || node.record.stage == NodeStage::Failed) + { + return NodeDecision::action(NodeDecisionAction::BeginDelete); + } + + if let Some(pending) = &node.pending { + return NodeDecision::wait(Some(pending.deadline)); + } + + if node.intent == NodeIntent::Deleting { + if let Some(next_effect_at) = node.retry.next_effect_at + && now < next_effect_at + { + return NodeDecision::wait(Some(next_effect_at)); + } + match node.retry.ambiguous_operation { + Some(OperationKind::CreateLease) => { + return NodeDecision::action(NodeDecisionAction::Dispatch( + NodeManagerCommand::CreateLease(crate::node::CreateLeaseRequest { + spec: node.record.desired.clone(), + }), + )); + } + Some(OperationKind::StartBootstrap) => { + if let Some(command) = start_bootstrap_command(&node.record) { + return NodeDecision::action(NodeDecisionAction::Dispatch(command)); + } + } + _ => {} + } + if node.active_bootstrap.is_none() && node.record.lease.is_none() { + return NodeDecision::action(NodeDecisionAction::MarkDestroyed); + } + if let Some(session_id) = node.active_bootstrap { + return NodeDecision::action(NodeDecisionAction::Dispatch( + NodeManagerCommand::CancelBootstrap { session_id }, + )); + } + if let Some(lease) = &node.record.lease { + return NodeDecision::action(NodeDecisionAction::Dispatch( + NodeManagerCommand::DestroyLease(lease.destroy_handle.clone()), + )); + } + return NodeDecision::action(NodeDecisionAction::MarkDestroyed); + } + + if let Some(next_effect_at) = node.retry.next_effect_at + && now < next_effect_at + { + return NodeDecision::wait(Some(next_effect_at)); + } + + if node.record.ready && matches!(node.record.stage, NodeStage::HandedOff | NodeStage::Dormant) { + return NodeDecision::wait(None); + } + + if node.record.lease.is_none() { + return NodeDecision::action(NodeDecisionAction::Dispatch( + NodeManagerCommand::CreateLease(crate::node::CreateLeaseRequest { + spec: node.record.desired.clone(), + }), + )); + } + + if node.record.connection.is_none() { + return NodeDecision::action(NodeDecisionAction::Dispatch( + NodeManagerCommand::LookupEndpoint( + node.record.lease.clone().expect("lease checked above"), + ), + )); + } + + if node.record.stage == NodeStage::SwactorJoined { + return match (node.active_bootstrap, node.record.swactor.as_ref()) { + (Some(session_id), Some(swactor)) => NodeDecision::action( + NodeDecisionAction::Dispatch(NodeManagerCommand::BootstrapConvergenceObserved { + session_id, + swactor_id: swactor.swactor_id.clone(), + }), + ), + _ => NodeDecision::wait(None), + }; + } + + if node.active_bootstrap.is_some() || node.record.stage == NodeStage::BootstrapRunning { + return NodeDecision::wait(None); + } + + let command = + start_bootstrap_command(&node.record).expect("lease and connection checked above"); + NodeDecision::action(NodeDecisionAction::Dispatch(command)) +} + +fn start_bootstrap_command(record: &NodeRecord) -> Option { + let lease_id = record.lease.as_ref()?.lease_id.clone(); + let ssh = record.connection.clone()?; + Some(NodeManagerCommand::StartBootstrap( + crate::node::BootstrapSessionSpec { + datastream: crate::node::DatastreamStreamId(format!( + "run/{}/node/{}/bootstrap", + record.run_id.0, record.logical_node_id.0 + )), + run_id: record.run_id.clone(), + logical_node_id: record.logical_node_id.clone(), + lease_id, + ssh, + boot: record.desired.boot.clone(), + swarm_join: record.desired.swarm_join.clone(), + }, + )) +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct RetryPolicy { + pub initial_delay: Duration, + pub max_delay: Duration, + pub jitter: Duration, + pub operation_timeout: Duration, + pub endpoint_probe_interval: Duration, +} + +impl RetryPolicy { + pub fn delay_for_failure(&self, consecutive_failures: u32) -> Duration { + let exponent = consecutive_failures.saturating_sub(1).min(31); + let multiplier = 1_u32 << exponent; + self.initial_delay + .checked_mul(multiplier) + .unwrap_or(self.max_delay) + .min(self.max_delay) + } +} + +fn sampled_retry_delay( + policy: &RetryPolicy, + attempt: NodeAttemptId, + consecutive_failures: u32, +) -> Duration { + let base = policy.delay_for_failure(consecutive_failures); + let jitter_nanos = policy.jitter.as_nanos(); + if jitter_nanos == 0 { + return base; + } + let mixed = attempt + .0 + .wrapping_mul(0x9e37_79b9_7f4a_7c15) + .wrapping_add(u64::from(consecutive_failures)) + .rotate_left(17); + let sampled_nanos = u128::from(mixed) % (jitter_nanos + 1); + let sampled = Duration::new( + (sampled_nanos / 1_000_000_000) as u64, + (sampled_nanos % 1_000_000_000) as u32, + ); + base.checked_add(sampled).unwrap_or(Duration::MAX) +} + +impl Default for RetryPolicy { + fn default() -> Self { + Self { + initial_delay: Duration::from_secs(1), + max_delay: Duration::from_secs(60), + jitter: Duration::ZERO, + operation_timeout: Duration::from_secs(120), + endpoint_probe_interval: Duration::from_secs(2), + } + } +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum OperationOutcome { + LeaseCreated(CreateLeaseResult), + EndpointLookup(Option), + BootstrapStarted { session_id: BootstrapSessionId }, + BootstrapConvergenceAccepted, + BootstrapCancelled, + LeaseDestroyed, +} + +impl OperationOutcome { + pub(crate) fn kind(&self) -> OperationKind { + match self { + Self::LeaseCreated(_) => OperationKind::CreateLease, + Self::EndpointLookup(_) => OperationKind::LookupEndpoint, + Self::BootstrapStarted { .. } => OperationKind::StartBootstrap, + Self::BootstrapConvergenceAccepted => OperationKind::BootstrapConvergenceObserved, + Self::BootstrapCancelled => OperationKind::CancelBootstrap, + Self::LeaseDestroyed => OperationKind::DestroyLease, + } + } +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct ExecutorResult { + pub node: LogicalNodeId, + pub operation: OperationId, + pub result: Result, +} + +// Keeping outcomes inline avoids allocating on every executor result. +#[allow(clippy::large_enum_variant)] +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum NodeObservation { + OperationSucceeded { + operation: OperationId, + outcome: OperationOutcome, + }, + OperationFailed { + operation: OperationId, + error: EffectError, + }, + BootstrapObserved { + session_id: BootstrapSessionId, + observation: BootstrapObservation, + }, + BootstrapFailed { + session_id: BootstrapSessionId, + reason: String, + }, + SwactorJoined { + session_id: BootstrapSessionId, + swactor_id: SwactorId, + }, + BootstrapClosed { + session_id: BootstrapSessionId, + }, +} + +pub fn observe( + node: &mut ManagedNode, + observation: NodeObservation, + now: SystemTime, + retry: &RetryPolicy, +) { + if node.record.stage == NodeStage::Failed + && matches!( + &observation, + NodeObservation::BootstrapObserved { .. } + | NodeObservation::BootstrapFailed { .. } + | NodeObservation::SwactorJoined { .. } + | NodeObservation::BootstrapClosed { .. } + ) + { + return; + } + match observation { + NodeObservation::OperationSucceeded { operation, outcome } => { + let Some(pending) = node.pending.as_ref() else { + return; + }; + if pending.id != operation || pending.kind != outcome.kind() { + return; + } + node.pending = None; + node.retry.next_effect_at = None; + node.retry.ambiguous_operation = None; + match outcome { + OperationOutcome::LeaseCreated(result) => { + node.record.lease = Some(result.lease); + if let Some(endpoint) = result.endpoint { + node.record.connection = Some(endpoint); + node.record.stage = NodeStage::EndpointKnown; + } else { + node.record.stage = NodeStage::LeaseCreated; + } + } + OperationOutcome::EndpointLookup(Some(endpoint)) => { + node.record.connection = Some(endpoint); + node.record.stage = NodeStage::EndpointKnown; + } + OperationOutcome::EndpointLookup(None) => { + node.retry.next_effect_at = now.checked_add(retry.endpoint_probe_interval); + } + OperationOutcome::BootstrapStarted { session_id } => { + node.active_bootstrap = Some(session_id); + node.record.bootstrap = Some(BootstrapFacts { + session_id, + last_stage: BootstrapStage::Created, + last_stdout_seq: None, + last_stderr_seq: None, + last_observed_at: now, + }); + node.record.stage = NodeStage::BootstrapRunning; + } + OperationOutcome::BootstrapConvergenceAccepted => {} + OperationOutcome::BootstrapCancelled => { + if let Some(facts) = node.record.bootstrap.as_mut() { + facts.last_stage = BootstrapStage::Cancelled; + facts.last_observed_at = now; + } + node.active_bootstrap = None; + } + OperationOutcome::LeaseDestroyed => { + node.record.lease = None; + node.record.connection = None; + } + } + } + NodeObservation::OperationFailed { operation, error } => { + let Some(pending) = node.pending.as_ref() else { + return; + }; + if pending.id != operation { + return; + } + let kind = pending.kind; + node.pending = None; + record_failure(node, kind, error, now, retry); + } + NodeObservation::BootstrapObserved { + session_id, + observation, + } => { + if node.active_bootstrap != Some(session_id) { + return; + } + let stage = observation.stage; + let facts = node.record.bootstrap.get_or_insert(BootstrapFacts { + session_id, + last_stage: stage, + last_stdout_seq: None, + last_stderr_seq: None, + last_observed_at: now, + }); + facts.last_stage = stage; + facts.last_observed_at = now; + if observation.last_stdout_seq.is_some() { + facts.last_stdout_seq = observation.last_stdout_seq; + } + if observation.last_stderr_seq.is_some() { + facts.last_stderr_seq = observation.last_stderr_seq; + } + if is_bootstrap_failure(stage) { + mark_attempt_failed(node, format!("bootstrap stage {stage:?}"), now, retry); + } else if matches!(stage, BootstrapStage::Converged | BootstrapStage::Closed) { + finish_bootstrap(node, now, retry); + } else if stage == BootstrapStage::Cancelled { + node.active_bootstrap = None; + } + } + NodeObservation::BootstrapFailed { session_id, reason } => { + if node.active_bootstrap == Some(session_id) { + mark_attempt_failed(node, reason, now, retry); + } + } + NodeObservation::SwactorJoined { + session_id, + swactor_id, + } => { + if node.active_bootstrap != Some(session_id) { + return; + } + node.record.swactor = Some(SwactorFacts { + swactor_id, + joined_at: now, + handed_off_at: None, + }); + node.record.stage = NodeStage::SwactorJoined; + } + NodeObservation::BootstrapClosed { session_id } => { + if node.active_bootstrap == Some(session_id) { + finish_bootstrap(node, now, retry); + } + } + } +} + +fn observation_matches(node: &ManagedNode, observation: &NodeObservation) -> bool { + match observation { + NodeObservation::OperationSucceeded { operation, outcome } => node + .pending + .as_ref() + .is_some_and(|pending| pending.id == *operation && pending.kind == outcome.kind()), + NodeObservation::OperationFailed { operation, .. } => node + .pending + .as_ref() + .is_some_and(|pending| pending.id == *operation), + NodeObservation::BootstrapObserved { session_id, .. } + | NodeObservation::BootstrapFailed { session_id, .. } + | NodeObservation::SwactorJoined { session_id, .. } + | NodeObservation::BootstrapClosed { session_id } => { + node.record.stage != NodeStage::Failed && node.active_bootstrap == Some(*session_id) + } + } +} + +fn record_failure( + node: &mut ManagedNode, + kind: OperationKind, + error: EffectError, + now: SystemTime, + retry: &RetryPolicy, +) { + let EffectError { + reason, + disposition, + } = error; + let ambiguous_operation = (disposition == EffectFailureDisposition::Ambiguous + && matches!( + kind, + OperationKind::CreateLease | OperationKind::StartBootstrap + )) + .then_some(kind); + node.retry.ambiguous_operation = ambiguous_operation; + if kind == OperationKind::BootstrapConvergenceObserved + || (kind == OperationKind::StartBootstrap + && disposition == EffectFailureDisposition::Definite) + { + mark_attempt_failed(node, reason, now, retry); + return; + } + node.retry.consecutive_failures = node.retry.consecutive_failures.saturating_add(1); + node.retry.last_error = Some(reason); + if node.intent == NodeIntent::Deleting + && ambiguous_operation.is_none() + && !matches!( + kind, + OperationKind::CancelBootstrap | OperationKind::DestroyLease + ) + { + node.retry.next_effect_at = None; + return; + } + node.retry.next_effect_at = now.checked_add(sampled_retry_delay( + retry, + node.attempt, + node.retry.consecutive_failures, + )); +} + +fn mark_attempt_failed( + node: &mut ManagedNode, + reason: String, + now: SystemTime, + retry: &RetryPolicy, +) { + node.retry.consecutive_failures = node.retry.consecutive_failures.saturating_add(1); + node.retry.last_error = Some(reason.clone()); + node.retry.next_effect_at = None; + node.retry.ambiguous_operation = None; + node.retry.restart_at = now.checked_add(sampled_retry_delay( + retry, + node.attempt, + node.retry.consecutive_failures, + )); + node.record.stage = NodeStage::Failed; + node.record.ready = false; + node.record.failed_reason = Some(reason); + node.record.failed_at = Some(now); +} + +fn finish_bootstrap(node: &mut ManagedNode, now: SystemTime, retry: &RetryPolicy) { + if node.record.swactor.is_none() { + mark_attempt_failed( + node, + "bootstrap closed before swactor convergence".to_owned(), + now, + retry, + ); + return; + } + if let Some(swactor) = node.record.swactor.as_mut() { + swactor.handed_off_at = Some(now); + } + node.active_bootstrap = None; + node.record.stage = NodeStage::Dormant; + node.record.ready = node.intent == NodeIntent::Active; + node.retry.consecutive_failures = 0; + node.retry.next_effect_at = None; + node.retry.restart_at = None; + node.retry.last_error = None; + node.retry.ambiguous_operation = None; +} + +fn is_bootstrap_failure(stage: BootstrapStage) -> bool { + matches!( + stage, + BootstrapStage::SshConnectFailed + | BootstrapStage::BootCheckFailed + | BootstrapStage::StartFailed + | BootstrapStage::SwactorJoinFailed + | BootstrapStage::StreamError + ) +} + +fn earliest(left: Option, right: Option) -> Option { + match (left, right) { + (Some(left), Some(right)) => Some(left.min(right)), + (Some(value), None) | (None, Some(value)) => Some(value), + (None, None) => None, + } +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum DriverError { + Shape(ShapeError), + RunChanged { expected: RunId, supplied: RunId }, + GenerationRegressed { current: u64, supplied: u64 }, + ShapeChangedWithoutGeneration { generation: u64 }, + ReentrantPass, + AttemptAllocatorMismatch, + OperationSequenceExhausted, +} + +impl fmt::Display for DriverError { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::Shape(error) => error.fmt(formatter), + Self::RunChanged { expected, supplied } => write!( + formatter, + "driver run changed from {} to {}", + expected.0, supplied.0 + ), + Self::GenerationRegressed { current, supplied } => write!( + formatter, + "desired generation regressed from {current} to {supplied}" + ), + Self::ShapeChangedWithoutGeneration { generation } => write!( + formatter, + "desired shape changed without advancing generation {generation}" + ), + Self::ReentrantPass => formatter.write_str("reentrant reconcile pass"), + Self::AttemptAllocatorMismatch => formatter.write_str("attempt allocator mismatch"), + Self::OperationSequenceExhausted => formatter.write_str("operation sequence exhausted"), + } + } +} + +impl std::error::Error for DriverError {} + +impl From for DriverError { + fn from(value: ShapeError) -> Self { + Self::Shape(value) + } +} + +#[derive(Clone, Debug)] +pub struct ClusterDriver { + state: ClusterState, + desired: ClusterShape, + expanded_desired: BTreeMap, + retry: RetryPolicy, + queued: bool, + processing: bool, + dirty: bool, + requeue_at: Option, +} + +impl ClusterDriver { + pub fn new(desired: ClusterShape, retry: RetryPolicy) -> Result { + let expanded_desired = desired.expand()?; + Ok(Self { + state: ClusterState::default(), + desired, + expanded_desired, + retry, + queued: true, + processing: false, + dirty: false, + requeue_at: None, + }) + } + + pub fn state(&self) -> &ClusterState { + &self.state + } + + pub fn desired(&self) -> &ClusterShape { + &self.desired + } + + pub fn retry_policy(&self) -> &RetryPolicy { + &self.retry + } + + pub fn is_queued(&self) -> bool { + self.queued + } + + pub fn requeue_at(&self) -> Option { + self.requeue_at + } + + pub fn trigger(&mut self) { + if self.processing { + self.dirty = true; + } else { + self.queued = true; + } + } + + pub fn update_desired(&mut self, desired: ClusterShape) -> Result<(), DriverError> { + if desired.run_id != self.desired.run_id { + return Err(DriverError::RunChanged { + expected: self.desired.run_id.clone(), + supplied: desired.run_id, + }); + } + if desired.generation < self.desired.generation { + return Err(DriverError::GenerationRegressed { + current: self.desired.generation, + supplied: desired.generation, + }); + } + let expanded = desired.expand()?; + if desired.generation == self.desired.generation && desired != self.desired { + return Err(DriverError::ShapeChangedWithoutGeneration { + generation: desired.generation, + }); + } + self.desired = desired; + self.expanded_desired = expanded; + self.trigger(); + Ok(()) + } + + pub fn trigger_if_due(&mut self, now: SystemTime) -> bool { + if self.requeue_at.is_some_and(|deadline| deadline <= now) { + self.requeue_at = None; + self.trigger(); + true + } else { + false + } + } + + pub fn pending_operations_due(&self, now: SystemTime) -> Vec { + self.state + .nodes + .iter() + .filter_map(|(node, managed)| { + let pending = managed.pending.as_ref()?; + (pending.deadline <= now).then(|| PlannedOperation { + node: node.clone(), + operation: pending.id, + kind: pending.kind, + deadline: pending.deadline, + }) + }) + .collect() + } + + /// Folds a timeout only after the executor has stopped the operation or + /// classified an ambiguous outcome according to its idempotency contract. + pub fn operation_timed_out( + &mut self, + operation: &PlannedOperation, + reason: impl Into, + now: SystemTime, + ) -> bool { + self.apply_observation( + &operation.node, + operation.operation.attempt, + NodeObservation::OperationFailed { + operation: operation.operation, + error: EffectError::ambiguous(reason), + }, + now, + ) + } + + pub fn apply_executor_result(&mut self, result: ExecutorResult, now: SystemTime) -> bool { + let observation = match result.result { + Ok(outcome) => NodeObservation::OperationSucceeded { + operation: result.operation, + outcome, + }, + Err(error) => NodeObservation::OperationFailed { + operation: result.operation, + error, + }, + }; + self.apply_observation(&result.node, result.operation.attempt, observation, now) + } + pub fn apply_observation( + &mut self, + node: &LogicalNodeId, + attempt: NodeAttemptId, + observation: NodeObservation, + now: SystemTime, + ) -> bool { + let Some(managed) = self.state.nodes.get_mut(node) else { + return false; + }; + if managed.attempt != attempt || !observation_matches(managed, &observation) { + return false; + } + observe(managed, observation, now, &self.retry); + self.trigger(); + true + } + + pub fn submission_failed( + &mut self, + effect: &PlannedEffect, + reason: impl Into, + now: SystemTime, + ) -> bool { + self.apply_observation( + &effect.node, + effect.operation.attempt, + NodeObservation::OperationFailed { + operation: effect.operation, + error: EffectError::definite(reason), + }, + now, + ) + } + + pub fn drive_next(&mut self, now: SystemTime, executor: &mut E) -> Result + where + E: EffectExecutor, + { + if self.processing { + return Err(DriverError::ReentrantPass); + } + if !self.queued { + return Ok(0); + } + + self.queued = false; + self.processing = true; + let pass = self.run_pass(now); + let result = match pass { + Ok(effects) => { + let mut submitted = 0_usize; + for effect in effects { + match executor.submit(&effect) { + Ok(()) => submitted = submitted.saturating_add(1), + Err(error) => { + self.submission_failed( + &effect, + format!("executor submission failed: {error}"), + now, + ); + } + } + } + Ok(submitted) + } + Err(error) => Err(error), + }; + self.processing = false; + if self.dirty { + self.dirty = false; + self.queued = true; + } + result + } + + pub fn drive_until_blocked( + &mut self, + now: SystemTime, + executor: &mut E, + ) -> Result + where + E: EffectExecutor, + { + let mut submitted = 0_usize; + while self.queued { + submitted = submitted.saturating_add(self.drive_next(now, executor)?); + } + Ok(submitted) + } + + pub fn is_converged(&self) -> bool { + if self.state.observed_generation != self.desired.generation + || self.state.nodes.len() != self.expanded_desired.len() + { + return false; + } + self.expanded_desired.iter().all(|(id, desired)| { + self.state.nodes.get(id).is_some_and(|node| { + node.intent == NodeIntent::Active + && node.record.desired == *desired + && node.record.ready + && node.pending.is_none() + }) + }) + } + + fn run_pass(&mut self, now: SystemTime) -> Result, DriverError> { + let plan = reconcile(&self.state, &self.desired, now)?; + let mut effects = Vec::new(); + self.requeue_at = plan.requeue_at; + + let mut allocated_actions_valid = true; + for action in plan.actions { + if let Some(allocated_attempt) = allocated_attempt(&action) + && (!allocated_actions_valid || allocated_attempt.0 != self.state.next_attempt_id) + { + allocated_actions_valid = false; + self.dirty = true; + continue; + } + if let Some(effect) = self.apply_action(action, now)? { + effects.push(effect); + } + } + self.state.observed_generation = plan.observed_generation; + Ok(effects) + } + + fn apply_action( + &mut self, + action: NodeAction, + now: SystemTime, + ) -> Result, DriverError> { + match action { + NodeAction::Insert { attempt, desired } => { + if attempt.0 != self.state.next_attempt_id + || self.state.nodes.contains_key(&desired.logical_node_id) + { + self.dirty = true; + return Ok(None); + } + self.state.next_attempt_id = self + .state + .next_attempt_id + .checked_add(1) + .ok_or(DriverError::AttemptAllocatorMismatch)?; + self.state.nodes.insert( + desired.logical_node_id.clone(), + ManagedNode::new(attempt, desired), + ); + self.dirty = true; + Ok(None) + } + NodeAction::BeginDelete { + node, + expected_attempt, + } => { + let Some(managed) = self.current_node_mut(&node, expected_attempt) else { + self.dirty = true; + return Ok(None); + }; + managed.intent = NodeIntent::Deleting; + managed.record.ready = false; + if managed.retry.ambiguous_operation.is_none() { + managed.retry.next_effect_at = None; + } + self.dirty = true; + Ok(None) + } + NodeAction::MarkDestroyed { + node, + expected_attempt, + } => { + let Some(managed) = self.current_node_mut(&node, expected_attempt) else { + self.dirty = true; + return Ok(None); + }; + if managed.pending.is_some() + || managed.active_bootstrap.is_some() + || managed.record.lease.is_some() + { + self.dirty = true; + return Ok(None); + } + managed.record.stage = NodeStage::Destroyed; + managed.record.ready = false; + managed.record.destroyed_at = Some(now); + self.dirty = true; + Ok(None) + } + NodeAction::Restart { + node, + expected_attempt, + new_attempt, + desired, + } => { + if new_attempt.0 != self.state.next_attempt_id { + self.dirty = true; + return Ok(None); + } + let Some(old) = self.state.nodes.get(&node) else { + self.dirty = true; + return Ok(None); + }; + if old.attempt != expected_attempt || old.record.stage != NodeStage::Destroyed { + self.dirty = true; + return Ok(None); + } + let failures = old.retry.consecutive_failures; + self.state.next_attempt_id = self + .state + .next_attempt_id + .checked_add(1) + .ok_or(DriverError::AttemptAllocatorMismatch)?; + let mut replacement = ManagedNode::new(new_attempt, desired); + replacement.retry.consecutive_failures = failures; + self.state.nodes.insert(node, replacement); + self.dirty = true; + Ok(None) + } + NodeAction::Reap { + node, + expected_attempt, + } => { + if self.state.nodes.get(&node).is_some_and(|managed| { + managed.attempt == expected_attempt + && managed.record.stage == NodeStage::Destroyed + }) { + self.state.nodes.remove(&node); + } else { + self.dirty = true; + } + Ok(None) + } + NodeAction::Dispatch(effect) => { + let timeout = self.retry.operation_timeout; + let Some(managed) = self.current_node_mut(&effect.node, effect.operation.attempt) + else { + self.dirty = true; + return Ok(None); + }; + if managed.pending.is_some() + || managed.next_operation_sequence != effect.operation.sequence + { + self.dirty = true; + return Ok(None); + } + let deadline = now.checked_add(timeout).unwrap_or(now); + let next_sequence = managed + .next_operation_sequence + .checked_add(1) + .ok_or(DriverError::OperationSequenceExhausted)?; + managed.pending = Some(PendingOperation { + id: effect.operation, + kind: OperationKind::for_command(&effect.command), + deadline, + }); + managed.next_operation_sequence = next_sequence; + if matches!(effect.command, NodeManagerCommand::CreateLease(_)) { + managed.record.stage = NodeStage::LeaseRequested; + } + self.requeue_at = earliest(self.requeue_at, Some(deadline)); + Ok(Some(effect)) + } + } + } + + fn current_node_mut( + &mut self, + id: &LogicalNodeId, + attempt: NodeAttemptId, + ) -> Option<&mut ManagedNode> { + self.state + .nodes + .get_mut(id) + .filter(|node| node.attempt == attempt) + } +} + +fn allocated_attempt(action: &NodeAction) -> Option { + match action { + NodeAction::Insert { attempt, .. } => Some(*attempt), + NodeAction::Restart { new_attempt, .. } => Some(*new_attempt), + _ => None, + } +} diff --git a/crates/provisioning/tests/executor.rs b/crates/provisioning/tests/executor.rs new file mode 100644 index 0000000..06b13c4 --- /dev/null +++ b/crates/provisioning/tests/executor.rs @@ -0,0 +1,399 @@ +use std::collections::BTreeMap; +use std::convert::Infallible; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::{Arc, Mutex}; + +use provisioning::*; + +#[derive(Clone, Default)] +struct InlineSpawner; + +impl BlockingEffectSpawner for InlineSpawner { + type SpawnError = Infallible; + + fn spawn_blocking(&self, work: BlockingEffectWork) -> Result<(), Self::SpawnError> { + work(); + Ok(()) + } +} + +#[derive(Clone, Default)] +struct QueuedSpawner { + work: Arc>>, +} + +impl QueuedSpawner { + fn run_all(&self) { + loop { + let work = std::mem::take(&mut *self.work.lock().unwrap()); + if work.is_empty() { + return; + } + for operation in work { + operation(); + } + } + } + + fn len(&self) -> usize { + self.work.lock().unwrap().len() + } +} + +impl BlockingEffectSpawner for QueuedSpawner { + type SpawnError = Infallible; + + fn spawn_blocking(&self, work: BlockingEffectWork) -> Result<(), Self::SpawnError> { + self.work.lock().unwrap().push(work); + Ok(()) + } +} + +#[derive(Default)] +struct CountingBackend { + calls: AtomicUsize, +} + +impl EffectBackend for CountingBackend { + fn execute(&self, effect: &PlannedEffect) -> Result { + self.calls.fetch_add(1, Ordering::SeqCst); + Ok(outcome_for(effect)) + } +} + +struct PanickingBackend; + +impl EffectBackend for PanickingBackend { + fn execute(&self, _effect: &PlannedEffect) -> Result { + panic!("backend panic") + } +} + +#[derive(Default)] +struct AdoptingBackend { + creates: AtomicUsize, + lease: Mutex>, +} + +impl EffectBackend for AdoptingBackend { + fn execute(&self, effect: &PlannedEffect) -> Result { + if !matches!(effect.command, NodeManagerCommand::CreateLease(_)) { + return Ok(outcome_for(effect)); + } + let mut live = self.lease.lock().unwrap(); + if let Some(lease) = live.as_ref() { + return Ok(OperationOutcome::LeaseCreated(lease.clone())); + } + self.creates.fetch_add(1, Ordering::SeqCst); + *live = Some(lease_result(effect.operation.attempt)); + Err(EffectError::ambiguous( + "provider accepted create before transport failed", + )) + } +} + +fn logical_spec(node: &str) -> LogicalNodeSpec { + LogicalNodeSpec { + run_id: RunId(3), + logical_node_id: LogicalNodeId(node.to_owned()), + group_id: NodeGroupId("workers".to_owned()), + role: RoleId("worker".to_owned()), + provider: ProviderKind::new("mock"), + shape: DesiredNodeShape { + image: "node:v1".to_owned(), + disk_gb: 10, + gpu_name: None, + min_gpu_ram_mb: None, + min_down_mbps: None, + min_up_mbps: None, + min_reliability: None, + require_verified: false, + provider_labels: BTreeMap::new(), + }, + boot: BootSpec { + ssh_user: "root".to_owned(), + verify_commands: Vec::new(), + start_swactor_command: "swactor".to_owned(), + stdout_sources: Vec::new(), + stderr_sources: Vec::new(), + env: Vec::new(), + args: Vec::new(), + mounts: Vec::new(), + }, + swarm_join: SwarmJoinSpec { + orch_swactor_addr: "orchestrator".to_owned(), + join_token_ref: "token".to_owned(), + expected_logical_node_id: LogicalNodeId(node.to_owned()), + }, + } +} + +fn create_effect(node: &str, attempt: u64, sequence: u64) -> PlannedEffect { + PlannedEffect { + node: LogicalNodeId(node.to_owned()), + operation: OperationId { + attempt: NodeAttemptId(attempt), + sequence, + }, + command: NodeManagerCommand::CreateLease(CreateLeaseRequest { + spec: logical_spec(node), + }), + } +} + +fn lease_result(attempt: NodeAttemptId) -> CreateLeaseResult { + let provider = ProviderKind::new("mock"); + let lease_id = ProviderLeaseId(format!("lease-{}", attempt.0)); + CreateLeaseResult { + lease: LeaseFacts { + provider: provider.clone(), + lease_id: lease_id.clone(), + provider_contract_id: format!("contract-{}", attempt.0), + offer_id: None, + destroy_handle: DestroyHandle { + provider, + lease_id, + provider_contract_id: format!("contract-{}", attempt.0), + }, + provider_metadata: BTreeMap::new(), + }, + endpoint: Some(SshEndpoint { + host: "127.0.0.1".to_owned(), + port: 22, + user: "root".to_owned(), + auth_ref: "key".to_owned(), + }), + } +} + +fn outcome_for(effect: &PlannedEffect) -> OperationOutcome { + match &effect.command { + NodeManagerCommand::CreateLease(_) => { + OperationOutcome::LeaseCreated(lease_result(effect.operation.attempt)) + } + NodeManagerCommand::LookupEndpoint(_) => { + OperationOutcome::EndpointLookup(Some(SshEndpoint { + host: "127.0.0.1".to_owned(), + port: 22, + user: "root".to_owned(), + auth_ref: "key".to_owned(), + })) + } + NodeManagerCommand::StartBootstrap(_) => OperationOutcome::BootstrapStarted { + session_id: BootstrapSessionId(effect.operation.attempt.0), + }, + NodeManagerCommand::BootstrapConvergenceObserved { .. } => { + OperationOutcome::BootstrapConvergenceAccepted + } + NodeManagerCommand::CancelBootstrap { .. } => OperationOutcome::BootstrapCancelled, + NodeManagerCommand::DestroyLease(_) => OperationOutcome::LeaseDestroyed, + } +} + +#[test] +fn repeated_operation_id_returns_cached_outcome_without_reexecution() { + let effect = create_effect("worker-0", 1, 1); + let mut executor = IdempotentEffectExecutor::new(CountingBackend::default(), InlineSpawner); + + executor.submit(&effect).unwrap(); + executor.submit(&effect).unwrap(); + + let results = executor.drain_results(); + assert_eq!(results.len(), 2); + assert!( + results + .iter() + .all(|result| result.operation == effect.operation) + ); + assert_eq!(executor.backend().calls.load(Ordering::SeqCst), 1); +} + +#[test] +fn operation_identity_reuse_with_different_input_is_rejected() { + let mut executor = IdempotentEffectExecutor::new(CountingBackend::default(), InlineSpawner); + let first = create_effect("worker-0", 1, 1); + executor.submit(&first).unwrap(); + executor.drain_results(); + + let mut conflicting = first.clone(); + let NodeManagerCommand::CreateLease(request) = &mut conflicting.command else { + unreachable!(); + }; + request.spec.shape.image = "different:v2".to_owned(); + let error = executor.submit(&conflicting).unwrap_err(); + + assert!(error.reason.contains("reused with different input")); + assert_eq!(executor.backend().calls.load(Ordering::SeqCst), 1); + assert!(executor.drain_results().is_empty()); +} + +#[test] +fn backend_panic_becomes_a_correlated_error_result() { + let effect = create_effect("worker-0", 1, 1); + let mut executor = IdempotentEffectExecutor::new(PanickingBackend, InlineSpawner); + + executor.submit(&effect).unwrap(); + let results = executor.drain_results(); + + assert_eq!(results.len(), 1); + assert_eq!(results[0].node, effect.node); + assert_eq!(results[0].operation, effect.operation); + assert!( + results[0] + .result + .as_ref() + .unwrap_err() + .reason + .contains("effect backend panicked: backend panic") + ); + assert_eq!( + executor.operation_status(effect.operation), + ExecutorOperationStatus::Completed + ); +} + +#[test] +fn expired_ambiguous_create_is_adopted_and_late_result_is_discarded() { + let spawner = QueuedSpawner::default(); + let mut executor = IdempotentEffectExecutor::new(AdoptingBackend::default(), spawner.clone()); + let effect = create_effect("worker-0", 1, 1); + executor.submit(&effect).unwrap(); + + assert!(executor.expire(effect.operation, "ambiguous timeout")); + let timeout = executor.drain_results(); + assert_eq!(timeout.len(), 1); + assert_eq!( + timeout[0].result.as_ref().unwrap_err(), + &EffectError::ambiguous("ambiguous timeout") + ); + + let retry = create_effect("worker-0", 1, 2); + executor.submit(&retry).unwrap(); + assert_eq!( + spawner.len(), + 1, + "retry must wait for the ambiguous physical operation to finish" + ); + spawner.run_all(); + let retry_result = executor.drain_results(); + assert_eq!(retry_result.len(), 1); + assert_eq!(retry_result[0].operation, retry.operation); + assert!(matches!( + retry_result[0].result, + Ok(OperationOutcome::LeaseCreated(_)) + )); + assert_eq!(executor.backend().creates.load(Ordering::SeqCst), 1); + + executor.submit(&effect).unwrap(); + assert_eq!( + executor.drain_results()[0].result.as_ref().unwrap_err(), + &EffectError::ambiguous("ambiguous timeout") + ); +} +#[test] +fn different_attempts_are_accepted_as_independent_blocking_work() { + let spawner = QueuedSpawner::default(); + let mut executor = IdempotentEffectExecutor::new(CountingBackend::default(), spawner.clone()); + let first = create_effect("worker-0", 1, 1); + let second = create_effect("worker-1", 2, 1); + + executor.submit(&first).unwrap(); + executor.submit(&second).unwrap(); + + assert_eq!(spawner.len(), 2); + assert_eq!( + executor.operation_status(first.operation), + ExecutorOperationStatus::InFlight + ); + assert_eq!( + executor.operation_status(second.operation), + ExecutorOperationStatus::InFlight + ); + spawner.run_all(); + assert_eq!(executor.drain_results().len(), 2); +} + +#[test] +fn one_attempt_cannot_run_two_distinct_operations_concurrently() { + let spawner = QueuedSpawner::default(); + let mut executor = IdempotentEffectExecutor::new(CountingBackend::default(), spawner.clone()); + let first = create_effect("worker-0", 1, 1); + let second = create_effect("worker-0", 1, 2); + + executor.submit(&first).unwrap(); + let error = executor.submit(&second).unwrap_err(); + + assert!(error.reason.contains("already has operation")); + assert_eq!(spawner.len(), 1); +} + +#[test] +fn newer_create_operation_adopts_an_ambiguous_live_resource() { + let mut executor = IdempotentEffectExecutor::new(AdoptingBackend::default(), InlineSpawner); + let ambiguous = create_effect("worker-0", 1, 1); + let retry = create_effect("worker-0", 1, 2); + + executor.submit(&ambiguous).unwrap(); + let first = executor.drain_results(); + assert_eq!(first.len(), 1); + assert!(first[0].result.is_err()); + + executor.submit(&retry).unwrap(); + let second = executor.drain_results(); + assert_eq!(second.len(), 1); + assert!(matches!( + second[0].result, + Ok(OperationOutcome::LeaseCreated(_)) + )); + assert_eq!(executor.backend().creates.load(Ordering::SeqCst), 1); +} + +#[test] +fn successful_outcomes_track_one_live_lease_and_bootstrap_per_attempt() { + let attempt = NodeAttemptId(4); + let mut executor = IdempotentEffectExecutor::new(CountingBackend::default(), InlineSpawner); + let create = create_effect("worker-0", attempt.0, 1); + executor.submit(&create).unwrap(); + executor.drain_results(); + assert_eq!( + executor.attempt_resources(attempt), + AttemptResources { + lease_live: true, + bootstrap_live: false + } + ); + + let mut start = create.clone(); + start.operation.sequence = 2; + start.command = NodeManagerCommand::StartBootstrap(BootstrapSessionSpec { + run_id: RunId(3), + logical_node_id: LogicalNodeId("worker-0".to_owned()), + lease_id: ProviderLeaseId("lease-4".to_owned()), + ssh: lease_result(attempt).endpoint.unwrap(), + boot: logical_spec("worker-0").boot, + swarm_join: logical_spec("worker-0").swarm_join, + datastream: DatastreamStreamId("bootstrap".to_owned()), + }); + executor.submit(&start).unwrap(); + executor.drain_results(); + assert!(executor.attempt_resources(attempt).bootstrap_live); + + let mut cancel = start.clone(); + cancel.operation.sequence = 3; + cancel.command = NodeManagerCommand::CancelBootstrap { + session_id: BootstrapSessionId(attempt.0), + }; + executor.submit(&cancel).unwrap(); + executor.drain_results(); + assert!(!executor.attempt_resources(attempt).bootstrap_live); + + let mut destroy = cancel; + destroy.operation.sequence = 4; + destroy.command = NodeManagerCommand::DestroyLease(lease_result(attempt).lease.destroy_handle); + executor.submit(&destroy).unwrap(); + executor.drain_results(); + assert_eq!( + executor.attempt_resources(attempt), + AttemptResources::default() + ); +} diff --git a/crates/provisioning/tests/reconciler.rs b/crates/provisioning/tests/reconciler.rs new file mode 100644 index 0000000..70ff810 --- /dev/null +++ b/crates/provisioning/tests/reconciler.rs @@ -0,0 +1,962 @@ +use std::collections::BTreeMap; +use std::time::{Duration, SystemTime}; + +use provisioning::*; +use std::convert::Infallible; + +fn group(id: &str, count: u32) -> RunNodeGroupSpec { + RunNodeGroupSpec { + run_id: RunId(7), + group_id: NodeGroupId(id.to_owned()), + role: RoleId("worker".to_owned()), + count, + provider: ProviderKind::new("mock"), + shape: DesiredNodeShape { + image: "node:v1".to_owned(), + disk_gb: 20, + gpu_name: None, + min_gpu_ram_mb: None, + min_down_mbps: None, + min_up_mbps: None, + min_reliability: None, + require_verified: false, + provider_labels: BTreeMap::new(), + }, + boot: BootSpec { + ssh_user: "root".to_owned(), + verify_commands: vec!["true".to_owned()], + start_swactor_command: "swactor".to_owned(), + stdout_sources: Vec::new(), + stderr_sources: Vec::new(), + env: Vec::new(), + args: Vec::new(), + mounts: Vec::new(), + }, + swarm_join: SwarmJoinTemplate { + orch_swactor_addr: "127.0.0.1:9000".to_owned(), + join_token_ref: "token".to_owned(), + }, + } +} + +fn shape(generation: u64, groups: Vec) -> ClusterShape { + ClusterShape { + run_id: RunId(7), + generation, + groups, + } +} + +fn lease(id: &str, endpoint: bool) -> CreateLeaseResult { + let provider = ProviderKind::new("mock"); + let lease_id = ProviderLeaseId(id.to_owned()); + CreateLeaseResult { + lease: LeaseFacts { + provider: provider.clone(), + lease_id: lease_id.clone(), + provider_contract_id: format!("contract-{id}"), + offer_id: None, + destroy_handle: DestroyHandle { + provider, + lease_id, + provider_contract_id: format!("contract-{id}"), + }, + provider_metadata: BTreeMap::new(), + }, + endpoint: endpoint.then(|| SshEndpoint { + host: "127.0.0.1".to_owned(), + port: 22, + user: "root".to_owned(), + auth_ref: "test-key".to_owned(), + }), + } +} + +#[derive(Default)] +struct RecordingExecutor { + effects: Vec, +} + +impl EffectExecutor for RecordingExecutor { + type SubmitError = Infallible; + + fn submit(&mut self, effect: &PlannedEffect) -> Result<(), Self::SubmitError> { + self.effects.push(effect.clone()); + Ok(()) + } +} + +fn drive(driver: &mut ClusterDriver, now: SystemTime) -> Vec { + let mut executor = RecordingExecutor::default(); + driver.drive_until_blocked(now, &mut executor).unwrap(); + executor.effects +} + +fn succeed( + driver: &mut ClusterDriver, + effect: &PlannedEffect, + outcome: OperationOutcome, + now: SystemTime, +) { + assert!(driver.apply_observation( + &effect.node, + effect.operation.attempt, + NodeObservation::OperationSucceeded { + operation: effect.operation, + outcome, + }, + now, + )); +} + +fn only_effect(driver: &mut ClusterDriver, now: SystemTime) -> PlannedEffect { + let effects = drive(driver, now); + assert_eq!(effects.len(), 1, "expected one effect, got {effects:?}"); + effects.into_iter().next().unwrap() +} + +#[test] +fn shape_expansion_validates_run_and_group_identity() { + let duplicate = shape(1, vec![group("gpu", 1), group("gpu", 2)]); + assert!( + duplicate + .expand() + .unwrap_err() + .reason + .contains("duplicate node group") + ); + + let mut wrong_run = group("cpu", 1); + wrong_run.run_id = RunId(8); + assert!( + shape(1, vec![wrong_run]) + .expand() + .unwrap_err() + .reason + .contains("expected 7") + ); + + let mut non_finite = group("invalid-metrics", 1); + non_finite.shape.min_down_mbps = Some(f64::NAN); + assert!( + shape(1, vec![non_finite]) + .expand() + .unwrap_err() + .reason + .contains("non-finite min_down_mbps") + ); +} + +#[test] +fn reconcile_is_sorted_pure_and_allocates_distinct_attempts() { + let desired = shape(4, vec![group("z", 1), group("a", 2)]); + let observed = ClusterState::default(); + let now = SystemTime::UNIX_EPOCH + Duration::from_secs(10); + + let first = reconcile(&observed, &desired, now).unwrap(); + let second = reconcile(&observed, &desired, now).unwrap(); + assert_eq!(first, second); + assert_eq!(observed, ClusterState::default()); + assert_eq!(first.observed_generation, 4); + assert_eq!( + first + .actions + .iter() + .map(|action| action.node().0.as_str()) + .collect::>(), + vec!["a-0", "a-1", "z-0"] + ); + assert_eq!( + first + .actions + .iter() + .map(|action| match action { + NodeAction::Insert { attempt, .. } => attempt.0, + other => panic!("unexpected action {other:?}"), + }) + .collect::>(), + vec![1, 2, 3] + ); +} + +#[test] +fn scale_down_marks_only_highest_logical_slots_for_deletion() { + let now = SystemTime::UNIX_EPOCH + Duration::from_secs(50); + let mut driver = + ClusterDriver::new(shape(1, vec![group("gpu", 3)]), RetryPolicy::default()).unwrap(); + assert_eq!(drive(&mut driver, now).len(), 3); + + driver + .update_desired(shape(2, vec![group("gpu", 1)])) + .unwrap(); + assert!(drive(&mut driver, now).is_empty()); + + let intent = |id: &str| { + driver + .state() + .nodes + .get(&LogicalNodeId(id.to_owned())) + .unwrap() + .intent + }; + assert_eq!( + intent("gpu-0"), + NodeIntent::Active, + "the lowest stable slot remains desired" + ); + assert_eq!(intent("gpu-1"), NodeIntent::Deleting); + assert_eq!(intent("gpu-2"), NodeIntent::Deleting); +} + +#[test] +fn driver_records_operations_before_returning_effects_and_reaches_ready() { + let now = SystemTime::UNIX_EPOCH + Duration::from_secs(100); + let mut driver = + ClusterDriver::new(shape(1, vec![group("gpu", 1)]), RetryPolicy::default()).unwrap(); + + let create = only_effect(&mut driver, now); + let node = driver.state().nodes.get(&create.node).unwrap(); + assert_eq!(node.pending.as_ref().unwrap().id, create.operation); + assert_eq!(node.record.stage, NodeStage::LeaseRequested); + assert!(!driver.is_converged()); + + succeed( + &mut driver, + &create, + OperationOutcome::LeaseCreated(lease("lease-1", true)), + now, + ); + let bootstrap = only_effect(&mut driver, now); + assert!(matches!( + bootstrap.command, + NodeManagerCommand::StartBootstrap(_) + )); + succeed( + &mut driver, + &bootstrap, + OperationOutcome::BootstrapStarted { + session_id: BootstrapSessionId(41), + }, + now, + ); + assert!(drive(&mut driver, now).is_empty()); + + let id = bootstrap.node.clone(); + let attempt = bootstrap.operation.attempt; + assert!(driver.apply_observation( + &id, + attempt, + NodeObservation::SwactorJoined { + session_id: BootstrapSessionId(41), + swactor_id: SwactorId("swactor-1".to_owned()), + }, + now, + )); + let convergence = only_effect(&mut driver, now); + assert!(matches!( + convergence.command, + NodeManagerCommand::BootstrapConvergenceObserved { .. } + )); + succeed( + &mut driver, + &convergence, + OperationOutcome::BootstrapConvergenceAccepted, + now, + ); + assert!(driver.apply_observation( + &id, + attempt, + NodeObservation::BootstrapClosed { + session_id: BootstrapSessionId(41), + }, + now, + )); + assert!(drive(&mut driver, now).is_empty()); + assert!(driver.is_converged()); +} + +#[test] +fn deletion_cancels_bootstrap_then_destroys_lease_then_reaps() { + let now = SystemTime::UNIX_EPOCH + Duration::from_secs(200); + let mut driver = + ClusterDriver::new(shape(1, vec![group("gpu", 1)]), RetryPolicy::default()).unwrap(); + let create = only_effect(&mut driver, now); + succeed( + &mut driver, + &create, + OperationOutcome::LeaseCreated(lease("lease-2", true)), + now, + ); + let bootstrap = only_effect(&mut driver, now); + succeed( + &mut driver, + &bootstrap, + OperationOutcome::BootstrapStarted { + session_id: BootstrapSessionId(9), + }, + now, + ); + drive(&mut driver, now); + + driver.update_desired(shape(2, Vec::new())).unwrap(); + let cancel = only_effect(&mut driver, now); + assert!(matches!( + cancel.command, + NodeManagerCommand::CancelBootstrap { + session_id: BootstrapSessionId(9) + } + )); + assert!( + driver + .state() + .nodes + .get(&cancel.node) + .unwrap() + .record + .lease + .is_some() + ); + + succeed( + &mut driver, + &cancel, + OperationOutcome::BootstrapCancelled, + now, + ); + let destroy = only_effect(&mut driver, now); + assert!(matches!( + destroy.command, + NodeManagerCommand::DestroyLease(_) + )); + succeed(&mut driver, &destroy, OperationOutcome::LeaseDestroyed, now); + assert!(drive(&mut driver, now).is_empty()); + assert!(driver.state().nodes.is_empty()); + assert_eq!(driver.state().observed_generation, 2); +} + +#[test] +fn failed_bootstrap_attempt_rejects_late_join_and_close() { + let now = SystemTime::UNIX_EPOCH + Duration::from_secs(250); + let mut driver = + ClusterDriver::new(shape(1, vec![group("gpu", 1)]), RetryPolicy::default()).unwrap(); + let create = only_effect(&mut driver, now); + succeed( + &mut driver, + &create, + OperationOutcome::LeaseCreated(lease("lease-failed-bootstrap", true)), + now, + ); + let bootstrap = only_effect(&mut driver, now); + succeed( + &mut driver, + &bootstrap, + OperationOutcome::BootstrapStarted { + session_id: BootstrapSessionId(91), + }, + now, + ); + + assert!(driver.apply_observation( + &bootstrap.node, + bootstrap.operation.attempt, + NodeObservation::BootstrapFailed { + session_id: BootstrapSessionId(91), + reason: "remote start failed".to_owned(), + }, + now, + )); + assert!(!driver.apply_observation( + &bootstrap.node, + bootstrap.operation.attempt, + NodeObservation::SwactorJoined { + session_id: BootstrapSessionId(91), + swactor_id: SwactorId("late-swactor".to_owned()), + }, + now, + )); + assert!(!driver.apply_observation( + &bootstrap.node, + bootstrap.operation.attempt, + NodeObservation::BootstrapClosed { + session_id: BootstrapSessionId(91), + }, + now, + )); + + let node = driver.state().nodes.get(&bootstrap.node).unwrap(); + assert_eq!(node.record.stage, NodeStage::Failed); + assert!(!node.record.ready); + assert!(node.record.swactor.is_none()); +} + +#[test] +fn replacement_waits_for_cleanup_and_uses_a_fresh_attempt() { + let now = SystemTime::UNIX_EPOCH + Duration::from_secs(300); + let mut initial_group = group("gpu", 1); + let mut driver = ClusterDriver::new( + shape(1, vec![initial_group.clone()]), + RetryPolicy::default(), + ) + .unwrap(); + let create = only_effect(&mut driver, now); + let old_attempt = create.operation.attempt; + succeed( + &mut driver, + &create, + OperationOutcome::LeaseCreated(lease("lease-3", true)), + now, + ); + let bootstrap = only_effect(&mut driver, now); + succeed( + &mut driver, + &bootstrap, + OperationOutcome::BootstrapStarted { + session_id: BootstrapSessionId(10), + }, + now, + ); + drive(&mut driver, now); + + initial_group.shape.image = "node:v2".to_owned(); + driver + .update_desired(shape(2, vec![initial_group.clone()])) + .unwrap(); + initial_group.shape.image = "node:v3".to_owned(); + driver + .update_desired(shape(3, vec![initial_group])) + .unwrap(); + let cancel = only_effect(&mut driver, now); + succeed( + &mut driver, + &cancel, + OperationOutcome::BootstrapCancelled, + now, + ); + let destroy = only_effect(&mut driver, now); + succeed(&mut driver, &destroy, OperationOutcome::LeaseDestroyed, now); + let replacement_create = only_effect(&mut driver, now); + assert_ne!(replacement_create.operation.attempt, old_attempt); + let replacement = driver.state().nodes.get(&replacement_create.node).unwrap(); + assert_eq!(replacement.record.desired.shape.image, "node:v3"); + + assert!(!driver.apply_observation( + &replacement_create.node, + old_attempt, + NodeObservation::OperationFailed { + operation: create.operation, + error: EffectError::definite("late old result"), + }, + now, + )); + assert_eq!( + driver + .state() + .nodes + .get(&replacement_create.node) + .unwrap() + .pending + .as_ref() + .unwrap() + .id, + replacement_create.operation + ); +} + +#[test] +fn retry_backoff_is_per_node_and_stores_one_deadline() { + let now = SystemTime::UNIX_EPOCH + Duration::from_secs(400); + let retry = RetryPolicy { + initial_delay: Duration::from_secs(5), + max_delay: Duration::from_secs(30), + jitter: Duration::ZERO, + operation_timeout: Duration::from_secs(60), + endpoint_probe_interval: Duration::from_secs(2), + }; + let mut driver = ClusterDriver::new(shape(1, vec![group("gpu", 2)]), retry).unwrap(); + let effects = drive(&mut driver, now); + assert_eq!(effects.len(), 2); + let failed = &effects[0]; + let progressing = &effects[1]; + assert!(driver.submission_failed(failed, "provider unavailable", now)); + succeed( + &mut driver, + progressing, + OperationOutcome::LeaseCreated(lease("lease-4", true)), + now, + ); + + let next = drive(&mut driver, now); + assert_eq!(next.len(), 1); + assert_eq!(next[0].node, progressing.node); + assert!(matches!( + next[0].command, + NodeManagerCommand::StartBootstrap(_) + )); + let failed_node = driver.state().nodes.get(&failed.node).unwrap(); + assert_eq!( + failed_node.retry.next_effect_at, + now.checked_add(Duration::from_secs(5)) + ); + assert_eq!(driver.requeue_at(), failed_node.retry.next_effect_at); +} + +#[test] +fn ambiguous_bootstrap_start_retries_adoption_without_replacing_the_attempt() { + let now = SystemTime::UNIX_EPOCH + Duration::from_secs(450); + let retry_delay = Duration::from_secs(3); + let retry = RetryPolicy { + initial_delay: retry_delay, + max_delay: retry_delay, + jitter: Duration::ZERO, + ..RetryPolicy::default() + }; + let mut driver = ClusterDriver::new(shape(1, vec![group("gpu", 1)]), retry).unwrap(); + let create = only_effect(&mut driver, now); + succeed( + &mut driver, + &create, + OperationOutcome::LeaseCreated(lease("lease-ambiguous-start", true)), + now, + ); + let start = only_effect(&mut driver, now); + + assert!(driver.apply_executor_result( + ExecutorResult { + node: start.node.clone(), + operation: start.operation, + result: Err(EffectError::ambiguous("bootstrap start timed out")), + }, + now, + )); + assert!(drive(&mut driver, now).is_empty()); + let node = driver.state().nodes.get(&start.node).unwrap(); + assert_eq!(node.attempt, start.operation.attempt); + assert_eq!(node.intent, NodeIntent::Active); + assert_eq!(node.record.stage, NodeStage::EndpointKnown); + assert!(node.active_bootstrap.is_none()); + assert!(node.retry.restart_at.is_none()); + assert_eq!(node.retry.next_effect_at, now.checked_add(retry_delay)); + + assert!(driver.trigger_if_due(now + retry_delay)); + let retry = only_effect(&mut driver, now + retry_delay); + assert_eq!(retry.operation.attempt, start.operation.attempt); + assert_ne!(retry.operation, start.operation); + assert!(matches!( + retry.command, + NodeManagerCommand::StartBootstrap(_) + )); +} + +#[test] +fn deleting_node_resolves_ambiguous_create_before_marking_destroyed() { + let now = SystemTime::UNIX_EPOCH + Duration::from_secs(475); + let mut driver = + ClusterDriver::new(shape(1, vec![group("gpu", 1)]), RetryPolicy::default()).unwrap(); + let create = only_effect(&mut driver, now); + let timeout_at = driver + .state() + .nodes + .get(&create.node) + .unwrap() + .pending + .as_ref() + .unwrap() + .deadline; + + driver.update_desired(shape(2, Vec::new())).unwrap(); + assert!(drive(&mut driver, now).is_empty()); + assert_eq!( + driver.state().nodes.get(&create.node).unwrap().intent, + NodeIntent::Deleting + ); + let due = driver.pending_operations_due(timeout_at).pop().unwrap(); + assert!(driver.operation_timed_out(&due, "ambiguous create", timeout_at)); + assert!(drive(&mut driver, timeout_at).is_empty()); + + let retry_at = timeout_at + Duration::from_secs(1); + assert!(driver.trigger_if_due(retry_at)); + let adopt = only_effect(&mut driver, retry_at); + assert!(matches!(adopt.command, NodeManagerCommand::CreateLease(_))); + succeed( + &mut driver, + &adopt, + OperationOutcome::LeaseCreated(lease("adopted-before-delete", true)), + retry_at, + ); + let destroy = only_effect(&mut driver, retry_at); + assert!(matches!( + destroy.command, + NodeManagerCommand::DestroyLease(_) + )); + succeed( + &mut driver, + &destroy, + OperationOutcome::LeaseDestroyed, + retry_at, + ); + assert!(drive(&mut driver, retry_at).is_empty()); + assert!(driver.state().nodes.is_empty()); +} + +#[test] +fn deleting_node_resolves_ambiguous_bootstrap_before_cleanup() { + let now = SystemTime::UNIX_EPOCH + Duration::from_secs(500); + let mut driver = + ClusterDriver::new(shape(1, vec![group("gpu", 1)]), RetryPolicy::default()).unwrap(); + let create = only_effect(&mut driver, now); + succeed( + &mut driver, + &create, + OperationOutcome::LeaseCreated(lease("lease-ambiguous-delete", true)), + now, + ); + let start = only_effect(&mut driver, now); + let timeout_at = driver + .state() + .nodes + .get(&start.node) + .unwrap() + .pending + .as_ref() + .unwrap() + .deadline; + + driver.update_desired(shape(2, Vec::new())).unwrap(); + assert!(drive(&mut driver, now).is_empty()); + let due = driver.pending_operations_due(timeout_at).pop().unwrap(); + assert!(driver.operation_timed_out(&due, "ambiguous bootstrap", timeout_at)); + assert!(drive(&mut driver, timeout_at).is_empty()); + + let retry_at = timeout_at + Duration::from_secs(1); + assert!(driver.trigger_if_due(retry_at)); + let adopt = only_effect(&mut driver, retry_at); + assert!(matches!( + adopt.command, + NodeManagerCommand::StartBootstrap(_) + )); + succeed( + &mut driver, + &adopt, + OperationOutcome::BootstrapStarted { + session_id: BootstrapSessionId(92), + }, + retry_at, + ); + let cancel = only_effect(&mut driver, retry_at); + assert!(matches!( + cancel.command, + NodeManagerCommand::CancelBootstrap { + session_id: BootstrapSessionId(92) + } + )); + succeed( + &mut driver, + &cancel, + OperationOutcome::BootstrapCancelled, + retry_at, + ); + let destroy = only_effect(&mut driver, retry_at); + succeed( + &mut driver, + &destroy, + OperationOutcome::LeaseDestroyed, + retry_at, + ); + assert!(drive(&mut driver, retry_at).is_empty()); + assert!(driver.state().nodes.is_empty()); +} + +#[test] +fn driver_rejects_generation_and_run_contract_violations() { + let mut driver = + ClusterDriver::new(shape(3, vec![group("gpu", 1)]), RetryPolicy::default()).unwrap(); + let mut changed = group("gpu", 1); + changed.shape.image = "changed".to_owned(); + assert!(matches!( + driver.update_desired(shape(3, vec![changed])), + Err(DriverError::ShapeChangedWithoutGeneration { generation: 3 }) + )); + assert!(matches!( + driver.update_desired(shape(2, vec![group("gpu", 1)])), + Err(DriverError::GenerationRegressed { + current: 3, + supplied: 2 + }) + )); + let mut other_run = shape(4, vec![group("gpu", 1)]); + other_run.run_id = RunId(9); + assert!(matches!( + driver.update_desired(other_run), + Err(DriverError::RunChanged { .. }) + )); +} + +#[test] +fn attempt_failure_cleans_up_immediately_but_delays_replacement() { + let now = SystemTime::UNIX_EPOCH + Duration::from_secs(500); + let retry = RetryPolicy { + initial_delay: Duration::from_secs(10), + max_delay: Duration::from_secs(10), + jitter: Duration::ZERO, + operation_timeout: Duration::from_secs(60), + endpoint_probe_interval: Duration::from_secs(2), + }; + let mut driver = ClusterDriver::new(shape(1, vec![group("gpu", 1)]), retry).unwrap(); + let create = only_effect(&mut driver, now); + succeed( + &mut driver, + &create, + OperationOutcome::LeaseCreated(lease("lease-5", true)), + now, + ); + let bootstrap = only_effect(&mut driver, now); + assert!(driver.submission_failed(&bootstrap, "bootstrap refused", now)); + + let destroy = only_effect(&mut driver, now); + assert!(matches!( + destroy.command, + NodeManagerCommand::DestroyLease(_) + )); + let failed = driver.state().nodes.get(&destroy.node).unwrap(); + assert_eq!(failed.record.failed_at, Some(now)); + assert_eq!( + failed.retry.restart_at, + now.checked_add(Duration::from_secs(10)) + ); + succeed(&mut driver, &destroy, OperationOutcome::LeaseDestroyed, now); + assert!(drive(&mut driver, now).is_empty()); + let destroyed = driver.state().nodes.get(&destroy.node).unwrap(); + assert_eq!(destroyed.record.stage, NodeStage::Destroyed); + assert_eq!( + driver.requeue_at(), + now.checked_add(Duration::from_secs(10)) + ); + + let restart_at = now + Duration::from_secs(10); + assert!(driver.trigger_if_due(restart_at)); + let replacement = only_effect(&mut driver, restart_at); + assert_ne!(replacement.operation.attempt, create.operation.attempt); +} + +#[test] +fn endpoint_not_ready_uses_probe_deadline_without_counting_a_failure() { + let now = SystemTime::UNIX_EPOCH + Duration::from_secs(600); + let mut driver = + ClusterDriver::new(shape(1, vec![group("gpu", 1)]), RetryPolicy::default()).unwrap(); + let create = only_effect(&mut driver, now); + succeed( + &mut driver, + &create, + OperationOutcome::LeaseCreated(lease("lease-6", false)), + now, + ); + let lookup = only_effect(&mut driver, now); + assert!(matches!( + lookup.command, + NodeManagerCommand::LookupEndpoint(_) + )); + succeed( + &mut driver, + &lookup, + OperationOutcome::EndpointLookup(None), + now, + ); + assert!(drive(&mut driver, now).is_empty()); + let node = driver.state().nodes.get(&lookup.node).unwrap(); + assert_eq!(node.retry.consecutive_failures, 0); + let probe_at = now + Duration::from_secs(2); + assert_eq!(node.retry.next_effect_at, Some(probe_at)); + + assert!(driver.trigger_if_due(probe_at)); + let retried_lookup = only_effect(&mut driver, probe_at); + assert!(matches!( + retried_lookup.command, + NodeManagerCommand::LookupEndpoint(_) + )); +} + +#[test] +fn executor_deadline_is_exposed_and_timeout_is_correlated() { + let now = SystemTime::UNIX_EPOCH + Duration::from_secs(700); + let retry = RetryPolicy { + initial_delay: Duration::from_secs(3), + max_delay: Duration::from_secs(30), + jitter: Duration::ZERO, + operation_timeout: Duration::from_secs(10), + endpoint_probe_interval: Duration::from_secs(2), + }; + let mut driver = ClusterDriver::new(shape(1, vec![group("gpu", 1)]), retry).unwrap(); + let create = only_effect(&mut driver, now); + assert!( + driver + .pending_operations_due(now + Duration::from_secs(9)) + .is_empty() + ); + let due = driver.pending_operations_due(now + Duration::from_secs(10)); + assert_eq!(due.len(), 1); + assert_eq!(due[0].operation, create.operation); + + let timeout_at = now + Duration::from_secs(10); + assert!(driver.operation_timed_out(&due[0], "executor timeout", timeout_at)); + let node = driver.state().nodes.get(&create.node).unwrap(); + assert!(node.pending.is_none()); + assert_eq!( + node.retry.next_effect_at, + Some(timeout_at + Duration::from_secs(3)) + ); + assert!(!driver.operation_timed_out(&due[0], "duplicate timeout", timeout_at)); +} + +#[test] +fn deleting_node_never_becomes_ready_from_late_bootstrap_completion() { + let now = SystemTime::UNIX_EPOCH + Duration::from_secs(800); + let desired = shape(1, vec![group("gpu", 1)]) + .expand() + .unwrap() + .into_values() + .next() + .unwrap(); + let mut node = ManagedNode::new(NodeAttemptId(1), desired); + node.intent = NodeIntent::Deleting; + node.active_bootstrap = Some(BootstrapSessionId(8)); + node.record.swactor = Some(SwactorFacts { + swactor_id: SwactorId("joined".to_owned()), + joined_at: now, + handed_off_at: None, + }); + + observe( + &mut node, + NodeObservation::BootstrapClosed { + session_id: BootstrapSessionId(8), + }, + now, + &RetryPolicy::default(), + ); + + assert_eq!(node.intent, NodeIntent::Deleting); + assert!(!node.record.ready); + assert_eq!(node.record.stage, NodeStage::Dormant); +} + +#[test] +fn same_generation_rejects_any_changed_shape_content() { + let mut driver = ClusterDriver::new(shape(1, Vec::new()), RetryPolicy::default()).unwrap(); + let changed = shape(1, vec![group("empty", 0)]); + + assert!(matches!( + driver.update_desired(changed), + Err(DriverError::ShapeChangedWithoutGeneration { generation: 1 }) + )); +} + +#[derive(Default)] +struct RejectingExecutor { + seen: Vec, +} + +impl EffectExecutor for RejectingExecutor { + type SubmitError = String; + + fn submit(&mut self, effect: &PlannedEffect) -> Result<(), Self::SubmitError> { + self.seen.push(effect.clone()); + Err("queue closed".to_owned()) + } +} + +#[test] +fn submission_failure_is_folded_after_pending_is_recorded() { + let now = SystemTime::UNIX_EPOCH + Duration::from_secs(900); + let mut driver = + ClusterDriver::new(shape(1, vec![group("gpu", 1)]), RetryPolicy::default()).unwrap(); + let mut executor = RejectingExecutor::default(); + + assert_eq!( + driver + .drive_until_blocked(now, &mut executor) + .expect("driver pass"), + 0 + ); + assert_eq!(executor.seen.len(), 1); + let node = driver.state().nodes.values().next().unwrap(); + assert!(node.pending.is_none()); + assert_eq!(node.record.stage, NodeStage::LeaseRequested); + assert!(node.retry.next_effect_at > Some(now)); + assert!( + node.retry + .last_error + .as_deref() + .is_some_and(|error| error.contains("queue closed")) + ); +} + +#[test] +fn repeated_queued_triggers_coalesce_to_one_dispatch() { + let now = SystemTime::UNIX_EPOCH + Duration::from_secs(1_000); + let mut driver = + ClusterDriver::new(shape(1, vec![group("gpu", 1)]), RetryPolicy::default()).unwrap(); + driver.trigger(); + driver.trigger(); + driver.trigger(); + + let effects = drive(&mut driver, now); + assert_eq!(effects.len(), 1); + assert!( + driver + .state() + .nodes + .values() + .next() + .unwrap() + .pending + .is_some() + ); + assert!(!driver.is_queued()); +} + +#[test] +fn cleanup_failure_retries_without_releasing_live_facts() { + let now = SystemTime::UNIX_EPOCH + Duration::from_secs(1_100); + let retry = RetryPolicy { + initial_delay: Duration::from_secs(4), + max_delay: Duration::from_secs(4), + jitter: Duration::ZERO, + operation_timeout: Duration::from_secs(60), + endpoint_probe_interval: Duration::from_secs(2), + }; + let mut driver = ClusterDriver::new(shape(1, vec![group("gpu", 1)]), retry).unwrap(); + let create = only_effect(&mut driver, now); + succeed( + &mut driver, + &create, + OperationOutcome::LeaseCreated(lease("cleanup", true)), + now, + ); + let bootstrap = only_effect(&mut driver, now); + succeed( + &mut driver, + &bootstrap, + OperationOutcome::BootstrapStarted { + session_id: BootstrapSessionId(77), + }, + now, + ); + drive(&mut driver, now); + driver.update_desired(shape(2, Vec::new())).unwrap(); + let cancel = only_effect(&mut driver, now); + assert!(driver.submission_failed(&cancel, "cancel busy", now)); + + let node = driver.state().nodes.get(&cancel.node).unwrap(); + assert_eq!(node.active_bootstrap, Some(BootstrapSessionId(77))); + assert!(node.record.lease.is_some()); + assert!(drive(&mut driver, now).is_empty()); + + let retry_at = now + Duration::from_secs(4); + assert!(driver.trigger_if_due(retry_at)); + let retried = only_effect(&mut driver, retry_at); + assert!(matches!( + retried.command, + NodeManagerCommand::CancelBootstrap { .. } + )); + assert_ne!(retried.operation, cancel.operation); +} diff --git a/docs/specs/archive/RECONCILER_SPEC.md b/docs/specs/archive/RECONCILER_SPEC.md new file mode 100644 index 0000000..f555df6 --- /dev/null +++ b/docs/specs/archive/RECONCILER_SPEC.md @@ -0,0 +1,588 @@ +# cluster reconciler — specification + +Id: 3 +Last modified: af49ba5c2cbcf7d742a69c6e213597ce5664fb43 +Last reviewed: +> Any edit to this spec must update `Last modified` above to the current `git HEAD` commit. + +**Scope:** a level-triggered reconciler that drives a declared cluster shape toward +convergence over the existing node lifecycle, living in `crates/provisioning` +alongside the node lifecycle types. + +## 1. Purpose + +Today, node lifecycle is edge-triggered and imperative: lifecycle handlers react +to discrete observations and emit commands, while `apps/myelin` explicitly +sequences node acquisition, bootstrap, readiness, retry, and teardown. There is +no object that owns "the cluster should look like this." + +This spec introduces a cluster reconciler. A pure decider compares the latest +desired shape with driver-owned observed state and returns the next actions that +move observed toward desired. A stateful driver coalesces triggers, applies state +transitions, dispatches effects, folds results back into observed state, and runs +the decider again. The system converges from whatever state is currently known; +it does not resume an imperative script from an event-specific continuation. + +The design borrows Kubernetes controller mechanics—level-based decisions, +spec/status separation, keyed and coalesced triggers, durable deletion intent, +retries outside business transitions, and distinct logical versus concrete +identity—without requiring a Kubernetes API server. Swactor supplies serialized +state transitions; the engine supplies effect execution and timers. + +## 2. Scope + +**In scope** + +- Desired cluster shape and driver-owned observed cluster state. +- A pure `reconcile(observed, desired, now) -> plan` contract. +- Observation folding that is separate from effect selection. +- Scale up/down and immutable-spec replacement across node groups. +- Stable logical-node identity and per-incarnation attempt identity. +- Coalesced event and periodic driving, non-reentrancy, and timed requeue. +- Pending-operation identity, effect-result correlation, and stale-result + rejection. +- Per-node retry/backoff and cleanup-before-restart semantics. +- Mapping these mechanics onto `NodeRecord`, `NodeStage`, and + `NodeManagerCommand` without a second node lifecycle. + +**Out of scope** + +- Actor workload placement on reconciled nodes. +- Data-plane and connectivity reconciliation as part of cluster shape. +- Provider-specific implementation details. +- Cluster-wide admission and backpressure policy. +- Persistence, process-crash recovery, and leader election. v1 assumes one + process-local driver. The state model must not preclude later persistence. +- Configurable rollout availability budgets. v1 may replace multiple stale + nodes concurrently; replacement is phased per logical node, not advertised as + an availability-preserving rolling update. + +## 3. Roles and ownership + +Three roles uphold one state-ownership rule. + +- **Reconciler** — a pure deterministic decider. It reads snapshots of desired + and observed state plus injected time and returns driver actions. It performs + no I/O and mutates no input. Purity is a Myelin testing seam, not a claim that + Kubernetes reconcilers themselves are pure. +- **Driver** — the sole writer of observed state. It owns the node map, folds + observations, coalesces triggers, calls the reconciler, applies driver-state + actions, records operations as pending before dispatch, and schedules timed + requeues. +- **Executor** — applies provider and bootstrap effects outside the driver + transition. It reports accepted results or failures tagged with the operation + and node-attempt identity. Blocking provider calls run as engine-hosted + blocking work and never block a reconcile pass. + +**Invariant — there is one authoritative observed state.** +`ClusterState.nodes` is authoritative. Observation reducers and the reconciler +operate on that state. A refactored `NodeManager` must not retain a second copy +of the same `NodeRecord` beside `ClusterState`. + +## 4. Kubernetes-derived controller mechanics + +The following mechanics are normative for this spec. + +1. **Triggers carry identity, not decision input.** An event means only that the + cluster may be dirty. Reconciliation rereads the latest complete state; it + never branches on which event caused the pass. +2. **Latest desired state wins.** If desired shape changes A -> B -> C before a + pass, convergence may proceed directly toward C. There is no obligation to + touch B. +3. **Triggers coalesce.** Repeated triggers while a pass is queued collapse into + one pass. A trigger arriving during a pass marks the cluster dirty and causes + one further pass after the current pass completes. +4. **Desired and observed revisions are distinct.** `generation` identifies a + desired-shape revision. `observed_generation` says only that the driver has + evaluated that revision; readiness separately reports convergence. +5. **Deletion is state, not absence plus a one-shot command.** Once cleanup has + begun it runs to completion. Reintroducing the same logical node while its + old attempt is deleting does not resurrect the old attempt; the latest + desired spec starts a fresh attempt after cleanup. +6. **Logical identity differs from concrete identity.** A stable logical slot + may have many sequential attempts. Results from an old attempt cannot mutate + the current attempt. +7. **Actuation is interruption-safe within the v1 process lifetime.** An effect + is recorded as pending before dispatch, tagged with a stable operation ID, + and correlated on completion. Ambiguous create outcomes use lookup/adoption + rather than blind duplicate creation. +8. **Retry scheduling is controller state.** Backoff and timed requeue do not + masquerade as node lifecycle stages. A retry deadline is sampled once, + stored, and read by the pure reconciler. + +## 5. Desired and observed state + +### 5.1 Desired state + +`RunNodeGroupSpec`, `LogicalNodeSpec`, and `expand_node_group` remain the desired +node vocabulary. `ClusterShape` adds a caller-controlled generation: + +```rust +pub struct ClusterShape { + pub run_id: RunId, + /// Strictly increases whenever the supplied desired shape changes. + pub generation: u64, + pub groups: Vec, +} + +impl ClusterShape { + pub fn expand( + &self, + ) -> Result, ShapeError>; +} +``` + +Shape expansion validates before reconciliation: + +- every group has `run_id == ClusterShape::run_id`; +- group IDs are unique; and +- expanded logical-node IDs are unique. + +Separately, the driver requires `run_id` to remain fixed for its lifetime, +rejects a generation lower than the last accepted desired generation, and +rejects changed shape content at the same generation. `expand` does not depend +on driver history. + +The map, rather than an unvalidated `Vec`, is the desired set. Expansion still +uses the existing `{group_id}-{index}` identity convention. Scaling up adds +higher indices; scaling down makes higher indices absent first. + +For v1, a `LogicalNodeSpec` is an immutable attempt template. Any inequality +between the current attempt's `record.desired` and the latest desired spec—role, +provider, shape, boot, or swarm-join data—requires replacement. In-place node +mutation can be introduced later only with an explicit field policy and +transition contract. + +### 5.2 Observed state + +`NodeRecord` remains the provider-neutral lifecycle fact record. It gains +`failed_at: Option`; its existing `desired` field is the immutable +spec snapshot implemented by that concrete attempt. + +Cluster-level mechanics wrap, rather than duplicate, the node lifecycle: + +```rust +pub struct NodeAttemptId(pub u64); + +pub struct OperationId { + pub attempt: NodeAttemptId, + pub sequence: u64, +} + +pub enum NodeIntent { + Active, + Deleting, +} + +pub struct PendingOperation { + pub id: OperationId, + pub kind: OperationKind, + /// Executor timeout sampled and stored before dispatch. + pub deadline: SystemTime, +} + +pub struct RetryState { + pub consecutive_failures: u32, + /// Earliest time another external effect may be dispatched. + pub next_effect_at: Option, + /// Earliest time a destroyed failed attempt may be replaced. + pub restart_at: Option, + pub last_error: Option, + /// Timed-out create/start operation that must be adopted before cleanup. + pub ambiguous_operation: Option, +} + +pub struct ManagedNode { + pub attempt: NodeAttemptId, + pub intent: NodeIntent, + pub record: NodeRecord, + /// Currently addressable bootstrap session; facts may outlive this handle. + pub active_bootstrap: Option, + pub pending: Option, + pub next_operation_sequence: u64, + pub retry: RetryState, +} + +pub struct ClusterState { + /// Latest desired generation evaluated by a completed pass. + pub observed_generation: u64, + /// Cluster-wide monotonic allocator; attempt IDs are never reused. + pub next_attempt_id: u64, + pub nodes: BTreeMap, +} +``` + +`LogicalNodeId` identifies the stable slot. `NodeAttemptId` is allocated from +`ClusterState::next_attempt_id`, is unique within the run, and is never reused +after reaping a slot. `OperationId` combines that attempt with a monotonically +increasing per-attempt sequence. + +A `NodeRecord` is replaced only when a new attempt starts. Old records may be +emitted to observability before replacement, but they are not simultaneously +live under the same map key. + +`BeginDelete` immediately makes the record non-ready but preserves all live +resource facts and any pending operation so its eventual result can still be +folded. `Restart` requires `Destroyed`, installs a fresh globally allocated +attempt and `NodeRecord` from the latest desired spec, clears +pending/live-session state and deadlines, resets the per-attempt operation +sequence, and preserves the consecutive-failure count until the new attempt +becomes ready. `Reap` requires `Destroyed`. + +## 6. Reconciler contract + +The plan contains driver transitions as well as external effects. This is +necessary because a pure decider cannot itself mark deletion, install a fresh +attempt, or record an operation as pending. + +```rust +pub struct ReconcilePlan { + /// At most one action per logical node, sorted by LogicalNodeId. + pub actions: Vec, + /// The desired generation evaluated by this plan. + pub observed_generation: u64, + /// Earliest known deadline requiring another pass without an event. + pub requeue_at: Option, +} + +pub enum NodeAction { + Insert { + attempt: NodeAttemptId, + desired: LogicalNodeSpec, + }, + BeginDelete { + node: LogicalNodeId, + expected_attempt: NodeAttemptId, + }, + MarkDestroyed { + node: LogicalNodeId, + expected_attempt: NodeAttemptId, + }, + Restart { + node: LogicalNodeId, + expected_attempt: NodeAttemptId, + new_attempt: NodeAttemptId, + desired: LogicalNodeSpec, + }, + Reap { + node: LogicalNodeId, + expected_attempt: NodeAttemptId, + }, + Dispatch(PlannedEffect), +} + +pub struct PlannedEffect { + pub node: LogicalNodeId, + pub operation: OperationId, + pub command: NodeManagerCommand, +} + +/// Pure and deterministic for identical inputs, including `now`. +pub fn reconcile( + observed: &ClusterState, + desired: &ClusterShape, + now: SystemTime, +) -> Result; +``` + +Every action carries enough identity or precondition to be safe if the driver +has changed since the snapshot. A stale action is discarded and the cluster is +marked dirty; it is never applied to a different attempt. + +While constructing the sorted plan, `reconcile` assigns distinct sequential +attempt IDs from the snapshotted `next_attempt_id`. The driver applies +`Insert`/`Restart` only when each assigned ID equals the current allocator, then +advances it with checked arithmetic. An allocator mismatch invalidates that and +all later allocated-attempt actions in the plan and marks the cluster dirty. + +A pass chooses at most one action per logical node. Different nodes can advance +in the same pass. A driver transition such as `Insert`, `BeginDelete`, +`MarkDestroyed`, `Restart`, or `Reap` completes that node's step for the pass; +its resulting external effect is considered only in a later pass. + +Before submitting `Dispatch`, the driver atomically: + +1. verifies the attempt, operation sequence, and absence of another pending + operation; +2. samples and stores the executor deadline in `PendingOperation`, schedules + that deadline, advances `next_operation_sequence`, and applies any + command-requested status such as `LeaseRequested`; and +3. submits the effect to the executor. + +If submission itself fails, the driver folds that as an operation failure. No +pass can observe an unrecorded in-flight effect. + +## 7. Observation folding and per-node progression + +Observation folding and effect selection are separate operations: + +```rust +pub fn observe( + node: &mut ManagedNode, + observation: NodeObservation, + now: SystemTime, + retry: &RetryPolicy, +); + +pub fn reconcile_node( + node: &ManagedNode, + desired: Option<&LogicalNodeSpec>, + now: SystemTime, +) -> NodeDecision; +``` + +`observe` mutates facts and emits no command. `reconcile_node` reads facts and +returns no more than one action. Time and retry policy enter mutation only +through the driver-provided arguments; neither function reads a global clock or +random source. + +Executor results and asynchronous observations carry `LogicalNodeId`, +`NodeAttemptId`, and, for command results, `OperationId`. Results for a stale +attempt or non-current operation are ignored after observability is recorded. + +### 7.1 Effect-result folding + +- `CreateLease` success stores `LeaseFacts`; an included endpoint also stores + `connection` and yields `EndpointKnown`, otherwise the stage is + `LeaseCreated`. +- `LookupEndpoint` with an endpoint stores it and yields `EndpointKnown`. + "Not available yet" retains the lease and stores a future probe deadline; it + is not an attempt-ending failure. +- `StartBootstrap` success returns a `BootstrapSessionId`, stores it as + `active_bootstrap`, stores `BootstrapFacts`, and yields `BootstrapRunning`. +- Bootstrap observations update stage and sequence facts only. +- A swactor-join observation stores `SwactorFacts` and yields + `SwactorJoined`; it does not itself emit convergence commands. +- Bootstrap convergence/closure clears `active_bootstrap`, marks handoff + complete, and yields the existing ready `Dormant` state. +- Bootstrap cancellation clears `active_bootstrap` while retaining terminal + bootstrap facts for observability. +- Lease destruction clears the live lease facts. A later `MarkDestroyed` + transition yields `Destroyed` and records `destroyed_at`. + +A successful command result clears the matching pending operation before the +next decision. An operation failure also clears it, records retry state, and +follows §11. + +### 7.2 Level-to-effect table + +`pending.is_some()` always means wait for its result or stored executor +deadline. Every dispatch row below also requires `next_effect_at` to be absent +or due; otherwise the node waits and contributes that deadline to +`requeue_at`. With no pending operation, progression is: + +| intent / observed facts | next action | +|---|---| +| active, `New` or `LeaseRequested`, no lease, retry due | `CreateLease` | +| active, lease known, no connection, probe due | `LookupEndpoint` | +| active, connection known, no bootstrap session | `StartBootstrap` | +| active, `BootstrapRunning`, no swactor | none; await observation | +| active, `SwactorJoined`, active bootstrap | `BootstrapConvergenceObserved` | +| active, `HandedOff` / `Dormant`, ready | none; steady state | +| active, attempt-ending `Failed` | `BeginDelete` | +| deleting, active bootstrap | `CancelBootstrap` | +| deleting, no active bootstrap, lease present | `DestroyLease` | +| deleting, no active bootstrap or lease | `MarkDestroyed` | +| `Destroyed`, desired present, restart deadline due | `Restart` with latest desired spec | +| `Destroyed`, desired absent | `Reap` | + +Cleanup ordering is deliberately sequential: cancel bootstrap, then destroy the +lease, then mark/reap or restart. The earlier draft's simultaneous cancel and +destroy effects violated one-step progression and made partial success +ambiguous. + +## 8. Topology and replacement + +Top-level reconciliation compares the validated desired map with observed nodes: + +- **desired only** — `Insert` a `ManagedNode` with the next globally allocated + attempt ID, `Active` intent, and a fresh `NodeRecord`. A later pass emits + `CreateLease`. +- **observed only** — if active, `BeginDelete`; if already deleting, continue + cleanup; if destroyed, `Reap`. +- **both, same spec** — run per-node progression. +- **both, different spec** — if active, `BeginDelete`. Once the old attempt is + destroyed and any restart deadline has elapsed, `Restart` installs the latest + desired spec under the next globally allocated attempt ID. + +Replacement never places one logical ID in simultaneous start and destroy +lists. Once `Deleting` begins it is not cancelled, even if the old spec becomes +desired again; cleanup finishes and the latest desired spec starts as a new +attempt. This is the process-local equivalent of a Kubernetes object name having +successive concrete UIDs. + +Scale-down order follows identity expansion: higher indices become absent +first. Multiple independent topology actions may occur in one pass. v1 defines +no availability budget across replacements; adding one is a group-policy +extension over this per-node lifecycle. + +## 9. Driver and workqueue semantics + +v1 uses one cluster reconcile key and one non-reentrant driver. Triggers come +from: + +- desired-shape generation changes; +- executor results and bootstrap/swactor observations; +- stored retry or probe deadlines; and +- a periodic safety tick. + +The driver maintains queued, processing, and dirty state equivalent to a +single-key Kubernetes workqueue: + +- adding an already queued key is a no-op; +- adding the key while it is processing marks it dirty; and +- completing a dirty pass immediately queues one further pass. + +Each pass: + +1. snapshots `ClusterState` and the latest `ClusterShape`; +2. calls `reconcile`; +3. applies each still-valid driver transition or records-and-submits each + `Dispatch` without waiting for provider I/O, scheduling every newly stored + pending-operation deadline; +4. records `plan.observed_generation` after the pass has evaluated that shape; +5. schedules `plan.requeue_at`, if any; and +6. immediately runs again if marked dirty while processing. + +Applying `Insert`, `BeginDelete`, `MarkDestroyed`, or `Restart` marks the cluster +dirty so its next lifecycle step cannot depend on an external event. `Reap` +needs no follow-up unless another trigger is already pending. `Dispatch` waits +for its correlated result or stored deadline. + +Observations are folded by serialized driver transitions before they can affect +a later snapshot. A periodic tick is a safety net, not the primary progress +mechanism. + +The cluster is converged for a generation when every desired node has the same +spec snapshot, is ready, has active intent, and has no pending operation; no +undesired or deleting nodes remain. `observed_generation == generation` alone +does not mean converged. + +## 10. Effect identity and idempotency + +A deterministic plan is not by itself an idempotent side effect. Safety comes +from observed facts, pending-operation state, and executor behavior. + +- The driver records an operation before dispatch and never emits a second + operation for that node while one is pending. +- The executor deduplicates repeated submissions of the same `OperationId` + within the driver process and returns the recorded outcome when known. +- The executor permits at most one live lease and one live bootstrap session + for `(run_id, logical_node_id, attempt)`, including across retries with newer + operation IDs. +- Provider creates use a deterministic external label or request token derived + from `(run_id, logical_node_id, attempt)`. After an ambiguous create outcome, + the executor looks up and adopts that attempt before creating again. +- Destroy and cancel treat "already absent" as success. +- Completion observations are correlated to both attempt and operation. Late + observations cannot mutate a replacement attempt. +- An executor timeout may clear pending state only after the executor has + stopped the operation or classified its outcome as ambiguous. Retrying an + ambiguous create or bootstrap start performs lookup/adoption first; it never + runs a concurrent blind duplicate. + +Because persistence is out of scope, v1 does not claim recovery from a process +crash between an external side effect and its in-memory observation. The stable +identities and adoption rule are the required shape for adding that guarantee +later. + +## 11. Failure and backoff + +Failure is observed state, but not every failed external call destroys the +whole attempt. The driver classifies by operation: + +| failure | retained state and retry behavior | +|---|---| +| lease creation | retain no lease; retry `CreateLease` after backoff | +| endpoint lookup / endpoint not ready | retain lease; retry lookup after backoff / probe interval | +| bootstrap start, bootstrap runtime, or join | mark attempt failed and begin cleanup immediately; delay only the later restart | +| bootstrap cancel | remain deleting; retry cancel after backoff | +| lease destroy | remain deleting with lease facts; retry destroy after backoff | + +Attempt-ending failure records `failed_reason` and `failed_at` and stores +`restart_at`; the next decision returns `BeginDelete`. It never resets a record +with live resources directly to `New`. Cleanup begins on the following pass and +is not delayed by the restart backoff. Once cleanup reaches `Destroyed`, the +node waits until `restart_at`, then starts a fresh attempt if it is still +desired. + +`RetryPolicy` is driver configuration: initial delay, cap, jitter, executor +operation timeout, and endpoint probe interval. When an observation is folded, +the driver computes the next deadline once and stores it in `RetryState`. +Random jitter is therefore not sampled by `reconcile`, preserving determinism. +Success clears `next_effect_at`; reaching ready steady state resets consecutive +failure count. + +Backoff is per node. A node waiting for a deadline contributes `requeue_at` but +does not prevent actions for other nodes. Provider work is dispatched outside +the pass, so slow or failed I/O for one node cannot block reconciliation of +another. + +## 12. Relationship to existing code + +| existing abstraction | role under this spec | +|---|---| +| `NodeRecord` / `NodeStage` | provider-neutral lifecycle facts retained inside `ManagedNode` | +| `NodeManagerCommand` | effect payload retained inside identity-bearing `PlannedEffect` | +| `RunNodeGroupSpec` / `expand_node_group` / `LogicalNodeSpec` | desired templates validated and wrapped by `ClusterShape` | +| provider and provisioning plugins | executor implementations behind effect dispatch and result correlation | +| `apps/myelin` orchestration | driver integration, executor wiring, and observation routing | + +`NodeManager::handle(msg) -> Vec` currently couples +observation mutation and effect selection. That shape cannot serve as the +observation-only half of this contract: some handlers advance state past the +point where a later level decision would emit the returned command. + +Implementation therefore performs a clean split: + +1. move the authoritative record into `ManagedNode`; +2. extract observation-only mutation into `observe`; +3. extract pure level decisions into `reconcile_node`; and +4. retire `handle` after callers migrate. + +This is one lifecycle expressed as a reducer plus a decider, not parallel edge +and level state machines. + +## 13. Kubernetes references (non-normative) + +The mechanics in §4 are adapted from: + +- [Kubernetes API conventions](https://github.com/kubernetes/community/blob/master/contributors/devel/sig-architecture/api-conventions.md), especially spec/status, generation, level-based behavior, and operation sequencing; +- [controller-runtime's reconcile contract](https://github.com/kubernetes-sigs/controller-runtime/blob/main/pkg/reconcile/reconcile.go), especially key-only requests and requeue semantics; +- [client-go workqueue](https://github.com/kubernetes/client-go/blob/master/util/workqueue/queue.go), especially dirty-key coalescing and per-key serialization; +- [Kubernetes finalizers](https://kubernetes.io/docs/concepts/overview/working-with-objects/finalizers/), especially durable cleanup-before-delete; and +- [Cluster API's InfraMachine contract](https://cluster-api.sigs.k8s.io/developer/providers/contracts/infra-machine), the closest analogue for provider-backed machine lifecycle. + +The cited systems persist controller objects in an API server. This spec adopts +their state-machine mechanics over an in-process Swactor/Myelin substrate; it +does not import their storage or network architecture. + +## 14. Invariants (normative) + +1. **Purity.** `reconcile` and `reconcile_node` perform no I/O, read no global + state, and mutate no input. +2. **Level-based decisions.** Effect selection depends on the latest desired and + observed state, never on the triggering event. +3. **Latest desired wins.** Intermediate desired generations need not be + visited. +4. **Single observed-state owner.** Only the driver mutates `ClusterState`. +5. **Observation/decision separation.** Observation folding emits no effects. +6. **One action per node per pass.** Convergence occurs across passes. +7. **Record before dispatch.** Every external effect has a pending operation in + observed state before execution begins. +8. **Attempt correlation.** An observation from an old attempt cannot mutate a + newer attempt. +9. **Non-reentrant, coalesced driving.** One cluster pass runs at a time; a + trigger during a pass guarantees a later pass without concurrent mutation. +10. **Cleanup before reuse.** A logical slot is not restarted or reaped until + bootstrap and lease cleanup for its old attempt is observed complete. +11. **Distinct identities.** Logical-node, node-attempt, and operation identity + are not interchangeable. +12. **Per-node failure isolation.** Retry or I/O for one node does not prevent + progress for another. +13. **No waiting in reconciliation.** Provider I/O and timers live in executor + or driver/engine work, never inside a reconcile pass. +14. **Observed generation is acknowledgment, not readiness.** Convergence is + determined from node state and pending topology work. +15. **Determinism.** Identical `(observed, desired, now)` inputs produce the same + plan; sampled retry deadlines are stored before reconciliation reads them. diff --git a/docs/specs/drafts/RECONCILER_SPEC.md b/docs/specs/drafts/RECONCILER_SPEC.md deleted file mode 100644 index 6abd601..0000000 --- a/docs/specs/drafts/RECONCILER_SPEC.md +++ /dev/null @@ -1,304 +0,0 @@ -# cluster reconciler — specification - -Id: 3 -Last modified: b887e941cbe6f1e209339abd0375507aca9bfe52 -Last reviewed: -> Any edit to this spec must update `Last modified` above to the current `git HEAD` commit. - -**Scope:** a level-triggered reconciler that drives a declared cluster shape toward -convergence over the existing node lifecycle, living in `crates/provisioning` -alongside `NodeManager`. - -## 1. Purpose - -Today, node lifecycle is **edge-triggered and imperative**: `NodeManager` reacts to -discrete events (`LeaseCreated`, `EndpointKnown`, `BootstrapObserved`…) and emits -commands, and `apps/myelin`'s orchestration module drives those managers by hand, -deciding *when* to start, retry, and tear down each node. There is no object that -owns "the cluster should look like *this*." Scaling, replacement, and recovery are -woven into imperative workflow code. - -This spec introduces a **reconciler**: a pure, level-triggered function that, given -a desired cluster shape and the currently observed cluster state, emits the effects -that move observed → desired. A **driver** calls it repeatedly (periodically and on -events), an **executor** applies the effects against providers, and observations -fold back into state for the next pass. The system converges; it does not execute a -script. - -The mindset shift: **edge-triggered imperative workflow → level-triggered -declarative convergence.** The reconciler never asks "what event just happened?" It -asks "given where this node is and where it must be, what is the next step?" - -## 2. Scope - -**In scope** - -- The reconciler contract: desired state, observed state, the pure `reconcile` - function, and the effect vocabulary. -- Driving semantics: periodic + event triggers, one-step-per-pass convergence, - idempotency, failure/backoff. -- Cluster-topology reconciliation: scale up/down across node groups, replacement on - shape change. -- The responsibility split between the reconciler (pure decider), the driver - (stateful owner of the node fleet), and the executor (applies effects to - providers). -- How the reconciler maps onto the existing `NodeManager` / `NodeStage` state - machine without inventing a parallel lifecycle. - -**Out of scope** - -- Actor workload placement on reconciled nodes (future scope; topology only for v1). -- Data-plane / connectivity reconciliation (iroh mesh, datastream links) as part of - the shape (future scope). -- Specific provider adapters (Vast.ai, Docker). Those implement the existing - `ProviderPlugin` / `ProvisionPlugin` executor seams. -- Backpressure and admission policy across the whole cluster. -- Persistence and leader election (single driver instance assumed for v1). - -## 3. Model - -Three roles, one invariant. - -- **Reconciler** — a pure, deterministic function - `reconcile(observed, desired) → plan`. No I/O, no clocks beyond an injected - `now()`, no mutation of inputs. Given identical inputs it yields an identical - plan. This is the load-bearing seam: everything testable and provider-neutral - lives here. -- **Driver** — the stateful loop. It owns the fleet of per-node state machines - (today: a `NodeManager` per logical node), calls the reconciler each pass, - dispatches the plan to the executor, folds provider/bootstrap observations back - into observed state, and decides *when* to run (periodic tick + event triggers). - The driver is the only writer of observed state. -- **Executor** — applies effects against reality. Maps to the existing seams: - `ProviderPlugin` (lease lifecycle: create / lookup / destroy) and - `ProvisionPlugin` (node process) plus bootstrap sessions. `apps/myelin`'s - orchestration module becomes this layer. - -**Invariant — the reconciler is pure; the driver owns all state and waiting.** -This mirrors the engine split (ENGINE_SPEC §3): core never waits, the engine owns -all waiting. Here, the reconciler never waits or mutates; the driver owns the -node fleet, the clocks, and the retry timers. - -## 4. State - -### Desired state - -Authoritative, supplied by the caller, held immutably between shape edits: - -```rust -// already defined in node.rs — unchanged -pub struct RunNodeGroupSpec { /* run_id, group_id, role, count, provider, shape, boot, swarm_join */ } -pub fn expand_node_group(group: &RunNodeGroupSpec) -> Vec; - -// new -pub struct ClusterShape { - pub run_id: RunId, - pub groups: Vec, -} -impl ClusterShape { - pub fn expand(&self) -> Vec { /* flatMap expand_node_group */ } -} -``` - -`ClusterShape` is a thin bag over the existing group spec; `expand` reuses -`expand_node_group`. The expanded `Vec` is the set of logical nodes -the cluster **should** contain. - -### Observed state - -The set of logical nodes the cluster **does** contain, each with its lifecycle -facts. `NodeRecord` already is per-node observed state. The cluster wraps it: - -```rust -// NodeRecord already holds: desired, stage, ready, lease, connection, bootstrap, -// swactor, failed_reason, destroyed_at. - -pub struct ClusterState { - pub nodes: BTreeMap, -} -``` - -The driver is the sole writer of `ClusterState`. Observations (lease results, -endpoints, bootstrap progress, swactor joins, failures) are folded into `NodeRecord` -between passes — exactly the work `NodeManager::handle` already does internally; -under the reconciler that folding is the driver's job (see §7). - -## 5. The reconciler contract - -```rust -pub struct ReconcilePlan { - /// Desired nodes with no observed record: begin their lifecycle. - pub to_start: Vec, - /// Observed nodes with no desired entry: tear them down. - pub to_destroy: Vec, - /// Live nodes: the next one or more commands to advance each toward desired. - pub per_node: Vec<(LogicalNodeId, Vec)>, -} - -/// Pure. Deterministic. No I/O. -pub fn reconcile( - observed: &ClusterState, - desired: &ClusterShape, - now: SystemTime, -) -> ReconcilePlan; -``` - -`reconcile` is **level-triggered**: it reads only `observed` + `desired` (+ `now` -for backoff; see §8). It does not know which event triggered the pass. It is -**idempotent**: re-running with unchanged inputs yields the same plan, and an effect -whose result is already reflected in observed state is never re-emitted (e.g. a node -whose `record.lease` is `Some` never yields `CreateLease` again). - -### Per-node reconcile - -For each live node, `reconcile` computes the next step from `(NodeRecord, -LogicalNodeSpec)` — a pure function over the existing `NodeStage` machine: - -```rust -/// One pass = at most one lifecycle step per node. Convergence happens across -/// passes, not within one. -fn reconcile_node(record: &NodeRecord, desired: &LogicalNodeSpec, now: SystemTime) - -> Vec; -``` - -The mapping reuses the existing `NodeManagerCommand` vocabulary and the existing -`NodeStage` transitions — it is the **level** reading of the same state machine that -`NodeManager::handle` expresses in **edge** form: - -| observed (`record`) | next effect(s) | -|---|---| -| `New`, no lease | `CreateLease` | -| `LeaseCreated`, lease carries endpoint | `StartBootstrap` | -| `LeaseCreated`, endpoint unknown | `LookupEndpoint` | -| `EndpointKnown` / `BootstrapRunning` | (none — awaiting bootstrap observation) | -| `BootstrapRunning`, swactor joined | `BootstrapConvergenceObserved` | -| `HandedOff` / `Dormant`, `ready` | (none — steady state) | -| `Failed`, within backoff window | (none — waiting; see §8) | -| `Failed`, backoff elapsed | reset to `New` → `CreateLease` (retry) | -| destroy requested | `CancelBootstrap` (if active) + `DestroyLease` | - -**One step per node per pass.** This is the heart of the level-triggered model: -the reconciler never waits within a pass. It emits the step the current observed -state permits; the executor applies it; observation updates the record; the next -pass emits the next step. Ordering across the lease → bootstrap → join chain falls -out of convergence, not from an explicit workflow. - -## 6. Topology reconciliation - -`reconcile` first diffs the expanded desired set against the observed set by -`LogicalNodeId`: - -- **desired, not observed** → `to_start`. The driver instantiates a `NodeManager` - for the spec; the first pass emits `CreateLease`. -- **observed, not desired** → `to_destroy`. The driver runs the destroy path - (`CancelBootstrap` + `DestroyLease`); once `stage == Destroyed` the record is - reaped. -- **both** → `per_node` via `reconcile_node`. - -**Scale policy (v1, deliberately simple):** logical node identity is -`{group_id}-{index}`. Scaling a group up adds higher indices; scaling down removes -the **highest** indices first. A group's `shape` is treated as immutable per node: -changing a field that is not achievable in place (image, gpu, disk) is a -**replacement** — the affected logical nodes move to `to_destroy` and fresh specs to -`to_start` — not an in-place mutation. This is k8s-style immutable-spec rolling -replacement, kept coarse for v1. - -## 7. Driving semantics - -The driver runs a pass when **either** (a) a periodic tick fires, or (b) an event -arrives — a desired-shape edit, or a provider/bootstrap observation that changed -observed state. Each pass: - -1. Snapshot current `ClusterState` and `ClusterShape`. -2. Call `reconcile(&observed, &desired, now)` → `ReconcilePlan`. -3. Apply the plan: create `NodeManager`s for `to_start`, drive destroy for - `to_destroy`, dispatch each `per_node` command to the executor. -4. Fold executor results + pending observations into `NodeRecord`s (the driver's - only write). -5. Repeat. Terminal when every desired node is `ready` and no orphans remain. - -**Observation folding is the bridge from edge to level.** Provider/bootstrap events -arrive as the existing observation types (`CreateLeaseResult`, `SshEndpoint`, -`BootstrapObservation`, swactor-join, failures). The driver folds each into the -node's `NodeRecord` between passes — the same field updates `NodeManager::handle` -performs today (`record.lease = …`, `record.bootstrap.last_stage = …`, etc.). The -reconciler then reads the updated record and emits the next step. The edge-triggered -`NodeManager::handle` and the level-triggered `reconcile_node` are two readings of -one state machine; see §10 for the migration choice. - -**Non-reentrancy.** A pass is synchronous and exclusive: the driver never runs two -passes concurrently. This matches the worker non-reentrancy invariant (ENGINE_SPEC -§5). - -## 8. Failure and backoff - -Failure is **observed state**, not a control-flow signal. A node reaching -`NodeStage::Failed` records `failed_reason` and `failed_at`. The reconciler emits no -command for that node **until its per-node backoff window elapses** (hence `now` in -the signature); after the window it resets the node to `New` and re-emits -`CreateLease`. Backoff is **per-node and isolated** — one failed node never blocks -another (this is the pay-off of the hybrid granularity chosen in §9). Destroyed -nodes that were desired are simply re-started by the topology diff. - -Backoff parameters (initial delay, cap, jitter) are driver configuration, not -reconciler logic — the reconciler only reads `failed_at` + the configured window -and decides "retry now" vs "wait." v1 uses a fixed window; exponential backoff is a -driver-side refinement. - -## 9. Granularity and growth path - -**Hybrid, by design.** v1 ships one top-level `reconcile` whose body is: topology -diff + fan-out to `reconcile_node`. The contract — pure function of (observed, -desired) → effects — is **identical at every level**, so the growth ladder is -internal refactor, never a contract change: - -1. **Now** — one loop, topology diff + per-node reconcile. Simple. -2. **When shapes diversify** — fan out to sub-reconcilers per concern (node-groups, - roles, future: data-plane links), each with the same signature; the top level - merges their effect streams. -3. **If independent backoff/isolation/work-queues are ever needed** — promote a - sub-reconciler to its own loop + driver. Same contract; the migration is - mechanical. - -`NodeManager` is already a sub-reconciler in waiting. The contract is what must not -ossify; the loop count is cheap to grow. - -## 10. Relationship to existing code - -| exists today | role under the reconciler | -|---|---| -| `NodeManager` + `NodeStage` | per-node observed state + lifecycle transitions. Kept. | -| `NodeRecord` | per-node observed state record. Kept; the driver writes it. | -| `NodeManagerCommand` | the effect vocabulary. Reused verbatim by `reconcile_node`. | -| `RunNodeGroupSpec` / `expand_node_group` / `LogicalNodeSpec` | desired state. Wrapped by `ClusterShape`; reused. | -| `ProviderPlugin` / `ProvisionPlugin` | executor seams. Unchanged; the driver calls them. | -| `apps/myelin` orchestration | becomes the driver + executor. Imperative workflow code is replaced by `reconcile` calls. | - -**Open design choice — edge handle vs. level reconcile for `NodeManager`:** -`NodeManager::handle(msg) → Vec` is edge-triggered; the -reconciler needs the level form. Two options, to be settled at implementation: - -- **(A) Add `reconcile_node` alongside `handle`.** `handle` keeps folding streaming - observations into the record (the parts that genuinely need event semantics, e.g. - bootstrap seq numbers); `reconcile_node` reads the record and emits the next step. - Minimal churn; two readings of one machine coexist. *Recommended for v1.* -- **(B) Split into `observe(&mut record, obs)` + `reconcile_node(&record, desired)`** - and retire `handle`. One level-triggered path; more churn, cleaner end state. - -Either way the *contract* in §5 is unchanged; only `NodeManager`'s internal shape -differs. - -## 11. Invariants (normative) - -1. **Purity.** `reconcile` performs no I/O, reads no global state, mutates no input. - `now` is its only non-input dependency. -2. **Level-triggered.** `reconcile` is a function of `(observed, desired, now)`, - never of "which event fired." Safe to call at any time. -3. **Idempotent.** Re-running with unchanged inputs yields the same plan; effects - already reflected in observed state are not re-emitted. -4. **One step per node per pass.** Convergence is across passes, not within one. -5. **Driver is the sole writer of observed state.** The reconciler and executor - never mutate `ClusterState`. -6. **Non-reentrant passes.** The driver never runs two passes concurrently. -7. **Failure isolation.** A failed node's backoff never blocks another node's - progress. diff --git a/tools/vastai/src/client.rs b/tools/vastai/src/client.rs index 24717b1..a4c551c 100644 --- a/tools/vastai/src/client.rs +++ b/tools/vastai/src/client.rs @@ -1,8 +1,8 @@ use std::time::Duration; use crate::types::{ - LabeledInstance, LifecyclePolicy, Offer, ProviderInstanceStatus, ProvisionRequest, - ProvisionedFleet, RunningInstance, + CreateInstanceRequest, InstanceInfo, LabeledInstance, LifecyclePolicy, Offer, + ProviderInstanceStatus, ProvisionRequest, ProvisionedFleet, RunningInstance, }; /// Small convenience wrapper around a reqwest client + vast.ai endpoint. @@ -30,16 +30,26 @@ impl VastClient { } } - pub fn http(&self) -> &reqwest::Client { - &self.http + pub async fn create_instance( + &self, + req: &CreateInstanceRequest, + ) -> Result { + crate::provision::create_instance(&self.http, &self.base_url, &self.api_key, req).await } - pub fn base_url(&self) -> &str { - &self.base_url + pub async fn destroy_instance(&self, contract_id: u64) -> Result<(), String> { + crate::teardown::destroy_instance(&self.http, &self.base_url, &self.api_key, contract_id) + .await } - pub fn api_key(&self) -> &str { - &self.api_key + pub async fn destroy_instance_with_retry(&self, contract_id: u64) -> Result<(), String> { + crate::teardown::destroy_instance_with_retry( + &self.http, + &self.base_url, + &self.api_key, + contract_id, + ) + .await } pub async fn search_offers(