From f54f62b4917a51ee73b5cc53e9ebb28829ddeb4a Mon Sep 17 00:00:00 2001 From: Zachery Aaron Shores-Chmielewski Date: Wed, 8 Jul 2026 18:06:00 +0400 Subject: [PATCH] refactor ssh bootstrap logic Replace stdout-parsed runtime-ready detection with an explicit, plugin-driven bootstrap-completion step and actorize SSH bootstrap teardown. - provisioning: drop the `PluginObservation::RuntimeReady` variant and add `ProvisionPlugin::complete_bootstrap`, an explicit per-node completion hook (no-op for `LocalDockerPlugin`) - bootstrap_datastream: remove `parse_runtime_ready`/`RuntimeReadyLine` so bootstrap no longer infers readiness from a parsed stdout JSON line - vastai_provisioning: drop the `ReadyTrackingSink` ready-flag wrapper; the SSH retry loop now runs purely `while !stopping`, and `complete_bootstrap` stops the node's bootstrap with `BootstrapStopReason::RuntimeReady` - vastai_provisioning: actorize teardown as `SshBootstrapActor` on the swactor `Runtime` (handle holds an `ActorAddress`), with `stop_bootstrap(handle, reason)` delivering a `Stop` message; add `BootstrapStopReason::{RuntimeReady,NodeStop}` - actors/provisioner: replace the `RuntimeReady` observation arm with a `ProvisionerMsg::RuntimeReady` handler that calls `complete_bootstrap` then `mark_live`/emits NodeLive (or NodeFailed on error) - callers/tests: wire the new explicit ready flow through node_agent, the orchestrator/worker_node binaries, and `mvp_one_node_chat`; add the `ssh_bootstrap_actor_stop_kills_child` test Signed-off-by: Zachery Aaron Shores-Chmielewski --- apps/mvp-node/tinygrad_worker.py | 28 ++- crates/mvp-system/src/actors/node_agent.rs | 74 ++++++++ crates/mvp-system/src/actors/orchestrator.rs | 93 +++++++++ crates/mvp-system/src/actors/provisioner.rs | 95 +++++++--- .../mvp-system/src/bin/mvp_one_node_chat.rs | 20 +- crates/mvp-system/src/bin/orchestrator.rs | 172 +++++++++++++---- crates/mvp-system/src/bin/worker_node.rs | 45 +++-- crates/mvp-system/src/bootstrap_datastream.rs | 29 --- crates/mvp-system/src/config.rs | 3 +- crates/mvp-system/src/provisioning.rs | 15 +- .../tests/bootstrap_datastream_guarantees.rs | 25 +-- .../tests/vastai_provisioning_guarantees.rs | 43 ++++- crates/mvp-system/src/vastai_provisioning.rs | 178 +++++++++++++----- 13 files changed, 613 insertions(+), 207 deletions(-) diff --git a/apps/mvp-node/tinygrad_worker.py b/apps/mvp-node/tinygrad_worker.py index afc0852..024fc78 100755 --- a/apps/mvp-node/tinygrad_worker.py +++ b/apps/mvp-node/tinygrad_worker.py @@ -368,6 +368,9 @@ def load_weights(cmd: dict[str, Any]) -> None: ) tok_src = cmd.get("tokenizer", {"EmbeddedGguf": None}) if "EmbeddedGguf" in tok_src: + if kv.get("tokenizer.ggml.pre") == "smollm": + kv = dict(kv) + kv["tokenizer.ggml.pre"] = "qwen2" control(type="TokenizerBuildStarted", model_id=model_id, source="EmbeddedGguf") tokenizer = SimpleTokenizer.from_gguf_kv(kv) control(type="TokenizerBuildReady", model_id=model_id, source="EmbeddedGguf") @@ -396,7 +399,13 @@ def load_weights(cmd: dict[str, Any]) -> None: def prompt_template_name() -> str: - return os.environ.get("MVP_PROMPT_TEMPLATE", "llama3-chat").strip().lower() + explicit = os.environ.get("MVP_PROMPT_TEMPLATE") + if explicit is not None: + return explicit.strip().lower() + model_id = str(loaded.get("model_id", "")).lower() + if "smollm" in model_id: + return "smollm-chat" + return "llama3-chat" def model_prompt_text(prompt: str) -> tuple[str, str]: @@ -412,12 +421,27 @@ def model_prompt_text(prompt: str) -> tuple[str, str]: "<|start_header_id|>assistant<|end_header_id|>\n\n", "llama3-chat", ) + if template in {"smollm", "smollm-chat", "smollm2", "smollm2-chat"}: + return ( + "<|im_start|>user\n" + f"{prompt}" + "<|im_end|>\n" + "<|im_start|>assistant\n", + "smollm-chat", + ) return prompt, "raw" def strip_chat_stop_markers(text: str) -> str: cut = len(text) - for marker in ("<|eot_id|>", "<|end_of_text|>", "<|start_header_id|>"): + for marker in ( + "<|eot_id|>", + "<|end_of_text|>", + "<|start_header_id|>", + "<|im_end|>", + "<|endoftext|>", + "<|im_start|>", + ): index = text.find(marker) if index >= 0: cut = min(cut, index) diff --git a/crates/mvp-system/src/actors/node_agent.rs b/crates/mvp-system/src/actors/node_agent.rs index 6ce510b..f5c212f 100644 --- a/crates/mvp-system/src/actors/node_agent.rs +++ b/crates/mvp-system/src/actors/node_agent.rs @@ -1,3 +1,4 @@ +use iroh::EndpointAddr; use serde::{Deserialize, Serialize}; use swactor::actor::{ActorAddress, ActorInterface}; use swactor::runtime::Ctx; @@ -51,6 +52,13 @@ impl StageProvisionWire { pub enum NodeAgentMsg { ProvisionStage(StageProvisionWire), MarkWorkerReady, + RuntimeLoaded { + run_id: u64, + node_id: u64, + stage_index: u32, + endpoint: EndpointAddr, + node_actor: ActorAddress, + }, MarkWeightsReady, MarkInboundEdgeReady { edge_id: u64, @@ -220,6 +228,25 @@ impl NodeAgentActor { }) } NodeAgentMsg::MarkWorkerReady => self.core.observe(stage::StageEvent::WorkerReady), + NodeAgentMsg::RuntimeLoaded { + run_id, + node_id, + stage_index, + endpoint, + node_actor, + } => { + self.core.observe(stage::StageEvent::WorkerReady); + let _ = ctx.send( + self.orchestrator, + OrchestratorMsg::ObserveNodeRuntimeReady { + run_id, + node_id, + stage_index, + endpoint, + node_actor, + }, + ); + } NodeAgentMsg::MarkWeightsReady => self.core.observe(stage::StageEvent::WeightsReady), NodeAgentMsg::MarkInboundEdgeReady { edge_id } => { self.core.observe(stage::StageEvent::InboundEdgeReady { @@ -477,3 +504,50 @@ pub fn register_codecs(registry: &mut CodecRegistry) { registry.register::(JsonCodec::::default()); registry.register::(JsonCodec::::default()); } + +#[cfg(test)] +mod tests { + use super::*; + use iroh::SecretKey; + use swactor::config::RuntimeConfig; + use swactor::runtime::Runtime; + + #[test] + fn node_agent_runtime_loaded_reports_orchestrator() { + let runtime = Runtime::new(RuntimeConfig::default()); + let orchestrator_inbox = runtime + .new_inbox::() + .expect("orchestrator inbox"); + let orchestrator = *orchestrator_inbox.addr(); + let node_actor = ActorAddress::new_random(); + let endpoint = EndpointAddr::new(SecretKey::from_bytes(&[9; 32]).public()); + let actor = runtime + .spawn(NodeAgentActor::new(stage::NodeId(11), orchestrator, None)) + .expect("spawn node agent"); + + runtime + .send_to( + actor, + NodeAgentMsg::RuntimeLoaded { + run_id: 7, + node_id: 11, + stage_index: 3, + endpoint: endpoint.clone(), + node_actor, + }, + ) + .expect("send runtime loaded"); + runtime.tick(); + + assert_eq!( + orchestrator_inbox.try_recv(), + Some(OrchestratorMsg::ObserveNodeRuntimeReady { + run_id: 7, + node_id: 11, + stage_index: 3, + endpoint, + node_actor, + }) + ); + } +} diff --git a/crates/mvp-system/src/actors/orchestrator.rs b/crates/mvp-system/src/actors/orchestrator.rs index 375a255..d6fb52a 100644 --- a/crates/mvp-system/src/actors/orchestrator.rs +++ b/crates/mvp-system/src/actors/orchestrator.rs @@ -1,3 +1,4 @@ +use iroh::EndpointAddr; use serde::{Deserialize, Serialize}; use swactor::actor::{ActorAddress, ActorInterface}; use swactor::runtime::Ctx; @@ -26,6 +27,13 @@ pub enum OrchestratorMsg { run_id: u64, stage_index: u32, }, + ObserveNodeRuntimeReady { + run_id: u64, + node_id: u64, + stage_index: u32, + endpoint: EndpointAddr, + node_actor: ActorAddress, + }, ObserveTokenInEndpointReady, ObserveTokenOutEndpointReady, ObserveTokenReceived { @@ -120,6 +128,13 @@ pub enum LifecycleEventWire { pub enum OrchestratorReport { Command(RunCommandWire), Lifecycle(LifecycleEventWire), + NodeRuntimeReady { + run_id: u64, + node_id: u64, + stage_index: u32, + endpoint: EndpointAddr, + node_actor: ActorAddress, + }, Snapshot { commands: Vec, events: Vec, @@ -177,6 +192,7 @@ impl OrchestratorActor { run_id: core::RunId(run_id), stage_index, }), + OrchestratorMsg::ObserveNodeRuntimeReady { .. } => {} OrchestratorMsg::ObserveTokenInEndpointReady => { self.core.observe(core::RunEvent::TokenInEndpointReady) } @@ -247,6 +263,28 @@ impl ActorInterface for OrchestratorActor { type Response = (); fn handle(&mut self, ctx: &Ctx, msg: Self::Incoming) { + if let OrchestratorMsg::ObserveNodeRuntimeReady { + run_id, + node_id, + stage_index, + endpoint, + node_actor, + } = msg.clone() + { + if let Some(report_to) = self.report_to { + let _ = ctx.send( + report_to, + OrchestratorReport::NodeRuntimeReady { + run_id, + node_id, + stage_index, + endpoint, + node_actor, + }, + ); + } + return; + } if let OrchestratorMsg::Snapshot { reply_to } = msg.clone() { let _ = ctx.send( reply_to, @@ -346,3 +384,58 @@ impl From<&core::TokenObjectPayload> for TokenObjectPayloadWire { } } } + +#[cfg(test)] +mod tests { + use super::*; + use iroh::SecretKey; + use swactor::config::RuntimeConfig; + use swactor::runtime::Runtime; + + #[test] + fn orchestrator_actor_reports_node_runtime_ready() { + let runtime = Runtime::new(RuntimeConfig::default()); + let reports = runtime + .new_inbox::() + .expect("orchestrator report inbox"); + let report_to = *reports.addr(); + let actor = runtime + .spawn(OrchestratorActor::new( + core::RunConfig { + run_id: core::RunId(7), + max_tokens: 1, + prompt: Vec::new(), + }, + Some(report_to), + )) + .expect("spawn orchestrator actor"); + let endpoint = EndpointAddr::new(SecretKey::from_bytes(&[8; 32]).public()); + let node_actor = ActorAddress::new_random(); + + runtime + .send_to( + actor, + OrchestratorMsg::ObserveNodeRuntimeReady { + run_id: 7, + node_id: 11, + stage_index: 3, + endpoint: endpoint.clone(), + node_actor, + }, + ) + .expect("send runtime ready"); + runtime.tick(); + + assert_eq!( + reports.try_recv(), + Some(OrchestratorReport::NodeRuntimeReady { + run_id: 7, + node_id: 11, + stage_index: 3, + endpoint, + node_actor, + }) + ); + assert_eq!(reports.try_recv(), None); + } +} diff --git a/crates/mvp-system/src/actors/provisioner.rs b/crates/mvp-system/src/actors/provisioner.rs index 1ea94ea..aaa56bc 100644 --- a/crates/mvp-system/src/actors/provisioner.rs +++ b/crates/mvp-system/src/actors/provisioner.rs @@ -1,3 +1,4 @@ +use iroh::EndpointAddr; use std::collections::BTreeMap; use std::sync::Arc; @@ -25,6 +26,13 @@ pub enum ProvisionerMsg { run_id: u64, reply_to: ActorAddress, }, + RuntimeReady { + run_id: u64, + node_id: u64, + stage_index: u32, + endpoint: EndpointAddr, + node_actor: ActorAddress, + }, PluginObservation(PluginObservation), } @@ -201,35 +209,6 @@ impl ProvisionerActor

{ PluginObservation::DatastreamFrame { channel, payload, .. } => self.emit_datastream_frame(channel, payload), - PluginObservation::RuntimeReady { - run_id, - node_id, - stage_index, - endpoint, - node_actor, - } => { - let report = self.mark_live(run_id, node_id, stage_index); - if let Some((reply_to, provider_process_id, resolved_stage_index)) = report { - self.emit_event(ProvisionEvent { - run_id, - node_id, - kind: ProvisionEventKind::NodeLive, - provider: None, - message: None, - }); - let _ = ctx.send( - reply_to, - ProvisionerReport::NodeLive { - run_id, - node_id, - stage_index: resolved_stage_index, - endpoint, - node_actor, - provider_process_id, - }, - ); - } - } PluginObservation::Exited { run_id, node_id, @@ -273,6 +252,57 @@ impl ProvisionerActor

{ } } + fn runtime_ready( + &mut self, + ctx: &Ctx, + run_id: u64, + node_id: u64, + stage_index: u32, + endpoint: iroh::EndpointAddr, + node_actor: ActorAddress, + ) { + let Some((handle, reply_to)) = self.runs.get(&run_id).and_then(|run| { + run.nodes + .get(&node_id) + .map(|slot| (slot.handle.clone(), slot.reply_to)) + }) else { + return; + }; + if let Err(reason) = self.plugin.complete_bootstrap(&handle) { + self.emit_failed(run_id, node_id, &reason); + let _ = ctx.send( + reply_to, + ProvisionerReport::NodeFailed { + run_id, + node_id, + reason, + }, + ); + return; + } + let report = self.mark_live(run_id, node_id, Some(stage_index)); + if let Some((reply_to, provider_process_id, resolved_stage_index)) = report { + self.emit_event(ProvisionEvent { + run_id, + node_id, + kind: ProvisionEventKind::NodeLive, + provider: None, + message: None, + }); + let _ = ctx.send( + reply_to, + ProvisionerReport::NodeLive { + run_id, + node_id, + stage_index: resolved_stage_index, + endpoint, + node_actor, + provider_process_id, + }, + ); + } + } + fn mark_live( &mut self, run_id: u64, @@ -379,6 +409,13 @@ impl ActorInterface for ProvisionerActor

{ ProvisionerMsg::StopNodes { run_id, reply_to } => { self.stop_nodes(ctx, run_id, reply_to) } + ProvisionerMsg::RuntimeReady { + run_id, + node_id, + stage_index, + endpoint, + node_actor, + } => self.runtime_ready(ctx, run_id, node_id, stage_index, endpoint, node_actor), ProvisionerMsg::PluginObservation(observation) => self.observe_plugin(ctx, observation), } } diff --git a/crates/mvp-system/src/bin/mvp_one_node_chat.rs b/crates/mvp-system/src/bin/mvp_one_node_chat.rs index 5b442e5..45ed60d 100644 --- a/crates/mvp-system/src/bin/mvp_one_node_chat.rs +++ b/crates/mvp-system/src/bin/mvp_one_node_chat.rs @@ -24,7 +24,7 @@ const DEFAULT_RPC_ADDR: &str = "127.0.0.1:19777"; const DEFAULT_NODE_IMAGE: &str = "swactor-mvp-node:latest"; const BASE_NODE_IMAGE: &str = "swactor-mvp-node-base:cuda12.6"; const MVP_RUNTIME_CONFIG_ENV: &str = "MVP_RUNTIME_CONFIG"; -const DEFAULT_CACHED_MODEL_FILE: &str = "Llama-3.2-1B-Instruct-Q4_K_M.gguf"; +const DEFAULT_CACHED_MODEL_FILE: &str = "SmolLM2-135M-Instruct.Q4_0.gguf"; const REPO_MODEL_CACHE_DIR: &str = ".model-cache"; const DEFAULT_MAX_TOKENS: u32 = 64; const ORCH_REBUILD_INPUTS: &[&str] = &[ @@ -212,7 +212,12 @@ impl Config { (true, None) => Some(PathBuf::from("mvp-chat.log")), (false, explicit) => explicit .or_else(|| env_optional("MVP_DATASTREAM_FRAME_LOG").map(PathBuf::from)) - .or_else(|| toml.observability.datastream_frame_log.clone().map(PathBuf::from)), + .or_else(|| { + toml.observability + .datastream_frame_log + .clone() + .map(PathBuf::from) + }), }; let model_id = first_non_empty([env_optional("MVP_MODEL_ID"), toml.model.id.clone()]); let gguf_repo = @@ -224,7 +229,10 @@ impl Config { toml.model.gguf_revision.clone(), ]); let max_context = env_u32_optional("MVP_MAX_CONTEXT")?.or(toml.model.max_context); - let cached_model = match (args.cached_model, toml.docker.cached_model_host_path.clone()) { + let cached_model = match ( + args.cached_model, + toml.docker.cached_model_host_path.clone(), + ) { (Some(cached_model), _) => Some(cached_model), (None, Some(path)) => Some(CachedModelConfig::from_arg(Some(path))?), (None, None) => None, @@ -339,7 +347,10 @@ impl Config { ]); } if let Some(min_down_mbps) = vastai.min_down_mbps { - args.extend(["--vastai-min-down-mbps".to_owned(), min_down_mbps.to_string()]); + args.extend([ + "--vastai-min-down-mbps".to_owned(), + min_down_mbps.to_string(), + ]); } if let Some(min_up_mbps) = vastai.min_up_mbps { args.extend(["--vastai-min-up-mbps".to_owned(), min_up_mbps.to_string()]); @@ -1520,7 +1531,6 @@ mod tests { } } - fn assert_arg_value(args: &[String], flag: &str, expected: &str) { let flag_index = args .iter() diff --git a/crates/mvp-system/src/bin/orchestrator.rs b/crates/mvp-system/src/bin/orchestrator.rs index 1828f17..4dfc2e4 100644 --- a/crates/mvp-system/src/bin/orchestrator.rs +++ b/crates/mvp-system/src/bin/orchestrator.rs @@ -14,12 +14,14 @@ use distribution::node::DistributedNodeConfig; use iroh::EndpointAddr; use iroh_driver::{IrohDriver, IrohDriverConfig}; use mvp_system::actors::node_agent::{NodeAgentMsg, StageProvisionWire}; +use mvp_system::actors::orchestrator::{OrchestratorActor, OrchestratorReport}; use mvp_system::actors::register_mvp_actor_codecs; +use mvp_system::config::{DEFAULT_CONFIG_PATH, TomlConfigOverlay}; #[cfg(feature = "local-e2e")] use mvp_system::dashboard_view::MvpClusterDashboardView; -use mvp_system::config::{DEFAULT_CONFIG_PATH, TomlConfigOverlay}; use mvp_system::distribution_stack::DistributionRuntimeStack; use mvp_system::node_provisioning::ProviderKind; +use mvp_system::orchestrator_run_fsm::{RunConfig, RunId}; use mvp_system::prompt_rpc::{PromptEvent, SubmitPrompt, read_submit_prompt, write_json_line}; use mvp_system::provisioning::{ LocalDockerPlugin, NodeProvisionSpec, PluginObservation, PluginObservationSink, PluginSink, @@ -237,6 +239,61 @@ fn run() -> Result<(), String> { json!({"enabled":dashboard.is_some()}), ); + let orchestrator_reports = match stack.runtime.new_inbox::() { + Ok(inbox) => inbox, + Err(error) => { + orch_datastream.emit_bootstrap( + dashboard.as_ref(), + config.run_id, + config.node_id, + "orchestrator_report_actor", + "failed", + json!({"error":error.to_string()}), + ); + return Err(format!("orchestrator report inbox: {error}")); + } + }; + let orchestrator_report_actor = *orchestrator_reports.addr(); + stack.register_local_actor(driver.register_actor(orchestrator_report_actor, 1)); + orch_datastream.emit_bootstrap( + dashboard.as_ref(), + config.run_id, + config.node_id, + "orchestrator_report_actor", + "ready", + json!({"actor":orchestrator_report_actor}), + ); + let orchestrator_actor = match stack.runtime.spawn(OrchestratorActor::new( + RunConfig { + run_id: RunId(config.run_id), + max_tokens: u64::from(config.default_max_tokens), + prompt: Vec::new(), + }, + Some(orchestrator_report_actor), + )) { + Ok(actor) => actor, + Err(error) => { + orch_datastream.emit_bootstrap( + dashboard.as_ref(), + config.run_id, + config.node_id, + "orchestrator_actor", + "failed", + json!({"error":error.to_string()}), + ); + return Err(format!("spawn orchestrator actor: {error}")); + } + }; + stack.register_local_actor(driver.register_actor(orchestrator_actor, 1)); + orch_datastream.emit_bootstrap( + dashboard.as_ref(), + config.run_id, + config.node_id, + "orchestrator_actor", + "ready", + json!({"actor":orchestrator_actor}), + ); + let prompt_events = match stack.runtime.new_inbox::() { Ok(inbox) => inbox, Err(error) => { @@ -265,7 +322,7 @@ fn run() -> Result<(), String> { let (work_tx, work_rx) = mpsc::channel::(); let stop_rx = spawn_stop_listener(); - let mut provisioner = config.build_provisioner()?; + let mut provisioner = config.build_provisioner(Arc::clone(&stack.runtime))?; orch_datastream.emit_bootstrap( dashboard.as_ref(), config.run_id, @@ -282,7 +339,8 @@ fn run() -> Result<(), String> { let sink = PluginSink::new(Arc::new(ChannelObservationSink { tx: Mutex::new(obs_tx), })); - let node_spec = config.node_spec(driver.endpoint_addr(), datastream_sink)?; + let node_spec = + config.node_spec(driver.endpoint_addr(), datastream_sink, orchestrator_actor)?; orch_datastream.emit_event( dashboard.as_ref(), ProvisionEvent { @@ -379,6 +437,7 @@ fn run() -> Result<(), String> { &stack, &obs_rx, &frame_rx, + &orchestrator_reports, &stop_rx, dashboard.as_ref(), &mut orch_datastream, @@ -394,7 +453,7 @@ fn run() -> Result<(), String> { config.node_id, "node_runtime_ready", "ready", - json!({"endpoint":&ready.endpoint,"node_actor":ready.node_actor}), + json!({"endpoint":&ready.endpoint,"node_actor":ready.node_actor,"stage_index":ready.stage_index}), ); drain_orch_stdio_capture( orch_stdio_rx.as_ref(), @@ -424,6 +483,7 @@ fn run() -> Result<(), String> { return Err(error); } }; + provisioned_node.complete_bootstrap()?; driver.join(std::slice::from_ref(&ready.endpoint)); orch_datastream.emit_bootstrap( dashboard.as_ref(), @@ -786,7 +846,6 @@ impl RuntimeConfigProfile { } } - fn as_str(self) -> &'static str { match self { Self::Local => "local", @@ -1115,7 +1174,9 @@ impl ConfigBuilder { self.layer_end_exclusive = Self::parse_value("MVP_LAYER_END_EXCLUSIVE", &layer_end_exclusive)?; } - if let Some(provider) = env_optional("MVP_NODE_PROVIDER").or_else(|| env_optional("MVP_PROVIDER")) { + if let Some(provider) = + env_optional("MVP_NODE_PROVIDER").or_else(|| env_optional("MVP_PROVIDER")) + { self.provider = Some(ProviderKind::parse_deploy(&provider)?); } if let Some(image) = env_optional("MVP_NODE_IMAGE") { @@ -1169,7 +1230,9 @@ impl ConfigBuilder { { self.relay_url = Some(url); } - if let Some(api_key) = env_optional("MVP_VASTAI_API_KEY").or_else(|| env_optional("VASTAI_API_KEY")) { + if let Some(api_key) = + env_optional("MVP_VASTAI_API_KEY").or_else(|| env_optional("VASTAI_API_KEY")) + { self.vastai_api_key = Some(api_key); } if let Some(command) = env_optional("MVP_VASTAI_BOOTSTRAP_COMMAND") { @@ -1280,7 +1343,8 @@ impl ConfigBuilder { Some(next_arg(&mut args, "--vastai-bootstrap-command")?) } "--vastai-ssh-identity" => { - self.vastai_ssh_identity_raw = Some(next_arg(&mut args, "--vastai-ssh-identity")?); + self.vastai_ssh_identity_raw = + Some(next_arg(&mut args, "--vastai-ssh-identity")?); } "--vastai-disk-gb" => { self.vastai_disk_gb = Some(parse_next(&mut args, "--vastai-disk-gb")?); @@ -1314,8 +1378,7 @@ impl ConfigBuilder { self.vastai_min_down_mbps_raw = None; } "--vastai-min-up-mbps" => { - self.vastai_min_up_mbps = - Some(parse_next(&mut args, "--vastai-min-up-mbps")?); + self.vastai_min_up_mbps = Some(parse_next(&mut args, "--vastai-min-up-mbps")?); self.vastai_min_up_mbps_raw = None; } "--vastai-min-reliability" => { @@ -1541,7 +1604,10 @@ impl Config { Ok(()) } - fn build_provisioner(&self) -> Result, String> { + fn build_provisioner( + &self, + bootstrap_runtime: Arc, + ) -> Result, String> { match self.provider { ProviderKind::Docker => Ok(Box::new(LocalDockerPlugin::new("mvp-orchestrator"))), ProviderKind::VastAi => { @@ -1565,7 +1631,7 @@ impl Config { let client = ToolsVastAiLeaseClient::from_api_key(api_key)?; Ok(Box::new(VastAiProvisioningPlugin::new( client, - SshCommandBootstrapLauncher::new(Some(ssh_identity)), + SshCommandBootstrapLauncher::new(Some(ssh_identity), bootstrap_runtime), vastai.provisioning.clone(), ))) } @@ -1581,6 +1647,7 @@ impl Config { "MVP_STAGE_INDEX", "MVP_COORDINATOR_ENDPOINT", "MVP_DATASTREAM_SINK_ACTOR", + "MVP_ORCHESTRATOR_ACTOR", "MVP_MODEL_ID", "MVP_IROH_RELAY_MODE", ]; @@ -1634,6 +1701,7 @@ impl Config { &self, coordinator: EndpointAddr, datastream_sink: ActorAddress, + orchestrator_actor: ActorAddress, ) -> Result { let mut env = vec![ ("MVP_RUN_ID".to_owned(), self.run_id.to_string()), @@ -1653,6 +1721,11 @@ impl Config { serde_json::to_string(&datastream_sink) .map_err(|e| format!("serialize datastream sink actor: {e}"))?, ), + ( + "MVP_ORCHESTRATOR_ACTOR".to_owned(), + serde_json::to_string(&orchestrator_actor) + .map_err(|e| format!("serialize orchestrator actor: {e}"))?, + ), ("MVP_MODEL_ID".to_owned(), self.model_id.clone()), ( "MVP_IROH_RELAY_MODE".to_owned(), @@ -1731,6 +1804,7 @@ impl Config { struct RuntimeReady { endpoint: EndpointAddr, node_actor: ActorAddress, + stage_index: u32, } struct ProvisionedNodeGuard<'a> { @@ -1749,6 +1823,13 @@ impl<'a> ProvisionedNodeGuard<'a> { } } + fn complete_bootstrap(&mut self) -> Result<(), String> { + let Some(handle) = self.handle.as_ref() else { + return Ok(()); + }; + self.provisioner.complete_bootstrap(handle) + } + fn stop(&mut self) -> Result<(), String> { let Some(handle) = self.handle.take() else { return Ok(()); @@ -1816,6 +1897,13 @@ impl ProvisionPlugin for FailedProvisionPlugin { Err("provider start worker disconnected".to_owned()) } + fn complete_bootstrap( + &mut self, + _handle: &mvp_system::provisioning::PluginNodeHandle, + ) -> Result<(), String> { + Ok(()) + } + fn stop_node( &mut self, _handle: &mvp_system::provisioning::PluginNodeHandle, @@ -2197,6 +2285,7 @@ fn wait_for_runtime_ready( stack: &DistributionRuntimeStack, obs_rx: &mpsc::Receiver, frame_rx: &mpsc::Receiver<(StreamId, Frame)>, + orchestrator_reports: &swactor::runtime::Inbox, stop_rx: &mpsc::Receiver<()>, dashboard: Option<&DashboardSupport>, orch_datastream: &mut OrchDatastream, @@ -2215,16 +2304,6 @@ fn wait_for_runtime_ready( while let Ok(observation) = obs_rx.try_recv() { emit_plugin_observation(orch_datastream, dashboard, provider, &observation); match observation { - PluginObservation::RuntimeReady { - endpoint, - node_actor, - .. - } => { - return Ok(RuntimeReady { - endpoint, - node_actor, - }); - } PluginObservation::DatastreamFrame { .. } => {} PluginObservation::ProviderLine { .. } | PluginObservation::StdoutLine { .. } @@ -2235,6 +2314,24 @@ fn wait_for_runtime_ready( } } } + 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, + } = report + { + if report_run_id == run_id && report_node_id == node_id { + return Ok(RuntimeReady { + endpoint, + node_actor, + stage_index, + }); + } + } + } thread::sleep(PUMP_INTERVAL); } } @@ -2319,7 +2416,6 @@ fn wait_for_weights_loaded( PluginObservation::ProviderLine { .. } | PluginObservation::StdoutLine { .. } | PluginObservation::StderrLine { .. } => {} - PluginObservation::RuntimeReady { .. } => {} } } while let Ok((stream, frame)) = frame_rx.try_recv() { @@ -2587,7 +2683,6 @@ fn drain_observations( PluginObservation::ProviderLine { .. } | PluginObservation::StdoutLine { .. } | PluginObservation::StderrLine { .. } => {} - PluginObservation::RuntimeReady { .. } => {} } } Ok(()) @@ -2647,18 +2742,6 @@ fn emit_plugin_observation( payload.as_bytes().to_vec(), "node_bootstrap_stdio", ), - PluginObservation::RuntimeReady { - run_id, node_id, .. - } => orch_datastream.emit_event( - dashboard, - ProvisionEvent { - run_id: *run_id, - node_id: *node_id, - kind: ProvisionEventKind::NodeLive, - provider: Some(provider.as_str().to_owned()), - message: None, - }, - ), PluginObservation::Exited { run_id, node_id, @@ -2733,7 +2816,6 @@ fn optional_env(name: &str) -> Option<(String, String)> { env_optional(name).map(|value| (name.to_owned(), value)) } - fn resolve_vastai_ssh_identity(explicit: Option) -> Result { match explicit { Some(path) => Ok(path), @@ -3030,8 +3112,9 @@ mod tests { .expect("config parses"); let coordinator = EndpointAddr::new(iroh::SecretKey::from_bytes(&[9; 32]).public()); let datastream_sink = ActorAddress([11; 32]); + let orchestrator_actor = ActorAddress([12; 32]); config - .node_spec(coordinator, datastream_sink) + .node_spec(coordinator, datastream_sink, orchestrator_actor) .expect("node spec builds") .env }) @@ -3430,8 +3513,9 @@ bootstrap_command = "/run" }); let coordinator = EndpointAddr::new(iroh::SecretKey::from_bytes(&[3; 32]).public()); let datastream_sink = ActorAddress([17; 32]); + let orchestrator_actor = ActorAddress([18; 32]); let spec = config - .node_spec(coordinator, datastream_sink) + .node_spec(coordinator, datastream_sink, orchestrator_actor) .expect("node spec builds"); assert_eq!(env_value(&spec.env, "MVP_MAX_CONTEXT"), Some("256")); @@ -3527,8 +3611,9 @@ bootstrap_command = "/run" ); let coordinator = EndpointAddr::new(iroh::SecretKey::from_bytes(&[7; 32]).public()); let datastream_sink = ActorAddress([13; 32]); + let orchestrator_actor = ActorAddress([14; 32]); let spec = config - .node_spec(coordinator, datastream_sink) + .node_spec(coordinator, datastream_sink, orchestrator_actor) .expect("cached model node spec builds"); assert_eq!( @@ -3592,6 +3677,13 @@ bootstrap_command = "/run" unreachable!("guard tests construct handles directly") } + fn complete_bootstrap( + &mut self, + _handle: &mvp_system::provisioning::PluginNodeHandle, + ) -> Result<(), String> { + Ok(()) + } + fn stop_node( &mut self, handle: &mvp_system::provisioning::PluginNodeHandle, diff --git a/crates/mvp-system/src/bin/worker_node.rs b/crates/mvp-system/src/bin/worker_node.rs index bc235b9..6190981 100644 --- a/crates/mvp-system/src/bin/worker_node.rs +++ b/crates/mvp-system/src/bin/worker_node.rs @@ -406,10 +406,10 @@ fn run() -> Result<(), String> { return Err(format!("node report inbox: {error}")); } }; - let (orchestrator, orchestrator_source) = match config.orchestrator_actor { - Some(actor) => (actor, "env"), - None => (ActorAddress::new_random(), "generated_fallback"), - }; + let orchestrator = config.orchestrator_actor.ok_or_else(|| { + "MVP_ORCHESTRATOR_ACTOR is required for runtime readiness signaling".to_owned() + })?; + let orchestrator_source = "env"; emit_stdio_node_event( &config, NODE_BOOTSTRAP_CHANNEL, @@ -512,26 +512,32 @@ fn run() -> Result<(), String> { return Err(error); } } - match stack - .runtime - .send_to(node_actor, NodeAgentMsg::MarkWorkerReady) - { + match stack.runtime.send_to( + node_actor, + NodeAgentMsg::RuntimeLoaded { + run_id: config.run_id, + node_id: config.logical_node_id, + stage_index: config.stage_index, + endpoint: driver.endpoint_addr(), + node_actor, + }, + ) { Ok(()) => emit_stdio_node_event( &config, NODE_RUNTIME_CHANNEL, - "mark_worker_ready", + "runtime_loaded", "ready", - json!({"sent":"NodeAgentMsg::MarkWorkerReady","node_actor":node_actor}), + json!({"sent":"NodeAgentMsg::RuntimeLoaded","node_actor":node_actor}), )?, Err(error) => { emit_stdio_node_event( &config, NODE_RUNTIME_CHANNEL, - "mark_worker_ready", + "runtime_loaded", "failed", json!({"error":error.to_string()}), )?; - return Err(format!("mark initialized worker ready: {error}")); + return Err(format!("signal runtime loaded: {error}")); } } @@ -553,10 +559,8 @@ fn run() -> Result<(), String> { "node_actor":node_actor, "logical_node_id":config.logical_node_id, "stage_index":config.stage_index, - "stdio_ready_line_emitted":true, }), )?; - println!("{ready}"); datastream.submit_text(ChannelId::new("mvp.node.ready"), ready.to_string()); emit_stdio_node_event( &config, @@ -565,12 +569,16 @@ fn run() -> Result<(), String> { "ready", json!({"from":"stdio_envelope","to":"cluster_datastream","channel":NODE_BOOTSTRAP_CHANNEL}), )?; - std::io::stdout() - .flush() - .map_err(|e| format!("flush ready line: {e}"))?; if let Some(prompt) = &config.self_test_prompt { - run_self_test(&mut worker, &config, prompt, &mut datastream, &mut driver, &stack)?; + run_self_test( + &mut worker, + &config, + prompt, + &mut datastream, + &mut driver, + &stack, + )?; } let shutdown_rx = spawn_stdin_shutdown_listener(); @@ -1667,7 +1675,6 @@ impl Drop for TinygradWorker { } } - fn spawn_stdin_shutdown_listener() -> Receiver<()> { let (tx, rx) = mpsc::channel(); thread::spawn(move || { diff --git a/crates/mvp-system/src/bootstrap_datastream.rs b/crates/mvp-system/src/bootstrap_datastream.rs index eb0357c..7e7f5e1 100644 --- a/crates/mvp-system/src/bootstrap_datastream.rs +++ b/crates/mvp-system/src/bootstrap_datastream.rs @@ -2,10 +2,8 @@ use std::io::{BufRead, BufReader, Read}; use std::thread::{self, JoinHandle}; use datastream::{ChannelId, DatastreamProducer, Lifetime, NodeId, StreamId}; -use iroh::EndpointAddr; use serde::{Deserialize, Serialize}; use serde_json::Value; -use swactor::actor::ActorAddress; use crate::provisioning::{ NodeProvisionSpec, PluginObservation, PluginSink, ProvisionLogLine, ProvisionLogStream, @@ -60,9 +58,6 @@ impl BootstrapDatastreamBridge { node_id: self.spec.node_id, line: line.clone(), }); - if let Some(ready) = parse_runtime_ready(&self.spec, &line) { - self.sink.observe(ready); - } } pub fn observe_stderr_line(&self, line: impl Into) { @@ -183,30 +178,6 @@ pub fn parse_stdio_datastream_frame( }) } -#[derive(Deserialize)] -struct RuntimeReadyLine { - #[serde(rename = "type")] - kind: String, - endpoint: EndpointAddr, - node_actor: ActorAddress, - logical_node_id: u64, - stage_index: u32, -} - -pub fn parse_runtime_ready(spec: &NodeProvisionSpec, line: &str) -> Option { - let ready = serde_json::from_str::(line).ok()?; - if ready.kind != "ready" || ready.logical_node_id != spec.node_id { - return None; - } - Some(PluginObservation::RuntimeReady { - run_id: spec.run_id, - node_id: spec.node_id, - stage_index: Some(ready.stage_index), - endpoint: ready.endpoint, - node_actor: ready.node_actor, - }) -} - pub fn bootstrap_log_channel(node_id: u64, stream: ProvisionLogStream) -> ChannelId { mvp_provision_log_channel(node_id, stream) } diff --git a/crates/mvp-system/src/config.rs b/crates/mvp-system/src/config.rs index 0fbd995..8ec6ef9 100644 --- a/crates/mvp-system/src/config.rs +++ b/crates/mvp-system/src/config.rs @@ -414,7 +414,8 @@ poll_interval_secs = 30 )); let _ = fs::remove_file(&path); - let error = TomlConfigOverlay::load(Some(&path)).expect_err("missing explicit config path errors"); + let error = + TomlConfigOverlay::load(Some(&path)).expect_err("missing explicit config path errors"); assert!( error.contains(&path.display().to_string()), diff --git a/crates/mvp-system/src/provisioning.rs b/crates/mvp-system/src/provisioning.rs index 0b497d0..c4df3ca 100644 --- a/crates/mvp-system/src/provisioning.rs +++ b/crates/mvp-system/src/provisioning.rs @@ -12,9 +12,7 @@ use std::sync::Arc; use std::thread; use std::time::UNIX_EPOCH; -use iroh::EndpointAddr; use serde::{Deserialize, Serialize}; -use swactor::actor::ActorAddress; use crate::bootstrap_datastream::BootstrapDatastreamBridge; @@ -93,13 +91,6 @@ pub enum PluginObservation { node_id: u64, line: String, }, - RuntimeReady { - run_id: u64, - node_id: u64, - stage_index: Option, - endpoint: EndpointAddr, - node_actor: ActorAddress, - }, Exited { run_id: u64, node_id: u64, @@ -144,6 +135,8 @@ pub trait ProvisionPlugin: Send { sink: PluginSink, ) -> Result; + fn complete_bootstrap(&mut self, handle: &PluginNodeHandle) -> Result<(), String>; + fn stop_node(&mut self, handle: &PluginNodeHandle) -> Result<(), String>; } @@ -431,6 +424,10 @@ impl ProvisionPlugin for LocalDockerPlugin { Ok(handle) } + fn complete_bootstrap(&mut self, _handle: &PluginNodeHandle) -> Result<(), String> { + Ok(()) + } + fn stop_node(&mut self, handle: &PluginNodeHandle) -> Result<(), String> { let Some(mut node) = self.nodes.remove(&handle.id) else { return Ok(()); diff --git a/crates/mvp-system/src/tests/bootstrap_datastream_guarantees.rs b/crates/mvp-system/src/tests/bootstrap_datastream_guarantees.rs index feb8c0e..90b9b2c 100644 --- a/crates/mvp-system/src/tests/bootstrap_datastream_guarantees.rs +++ b/crates/mvp-system/src/tests/bootstrap_datastream_guarantees.rs @@ -9,7 +9,6 @@ use mvp_system::provisioning::{ use mvp_system::telemetry::MvpProvisionLogRecord; use parking_lot::Mutex; use serde_json::json; -use swactor::actor::ActorAddress; #[derive(Default)] struct RecordingSink { @@ -90,30 +89,26 @@ fn bootstrap_bridge_writes_node_stream_and_forwards_plugin_observations() { } #[test] -fn ready_json_on_stdout_emits_runtime_ready_through_plugin_sink() { +fn stdout_ready_json_is_log_only() { let (recording, sink) = recording_sink(); let bridge = BootstrapDatastreamBridge::new(spec(), sink, None); - let endpoint = EndpointAddr::new(SecretKey::from_bytes(&[7; 32]).public()); - let node_actor = ActorAddress::new_random(); let line = serde_json::to_string(&json!({ "type": "ready", - "endpoint": endpoint, - "node_actor": node_actor, + "endpoint": EndpointAddr::new(SecretKey::from_bytes(&[7; 32]).public()), + "node_actor": "ignored-by-stdout-bridge", "logical_node_id": 42, "stage_index": 3, })) .unwrap(); - bridge.observe_stdout_line(line); + bridge.observe_stdout_line(line.clone()); - assert!(recording.observations().iter().any(|observation| matches!( - observation, - PluginObservation::RuntimeReady { + assert_eq!( + recording.observations(), + vec![PluginObservation::StdoutLine { run_id: 7, node_id: 42, - stage_index: Some(3), - endpoint: observed_endpoint, - node_actor: observed_actor, - } if observed_endpoint == &endpoint && observed_actor == &node_actor - ))); + line, + }] + ); } diff --git a/crates/mvp-system/src/tests/vastai_provisioning_guarantees.rs b/crates/mvp-system/src/tests/vastai_provisioning_guarantees.rs index f209f3a..af3a56d 100644 --- a/crates/mvp-system/src/tests/vastai_provisioning_guarantees.rs +++ b/crates/mvp-system/src/tests/vastai_provisioning_guarantees.rs @@ -7,8 +7,8 @@ use mvp_system::provisioning::{ NodeProvisionSpec, PluginObservation, PluginObservationSink, PluginSink, ProvisionPlugin, }; use mvp_system::vastai_provisioning::{ - VastAiBootstrapLauncher, VastAiLeaseClient, VastAiProviderPlugin, VastAiProvisioningConfig, - VastAiProvisioningPlugin, VastAiSshEndpoint, + BootstrapStopReason, VastAiBootstrapLauncher, VastAiLeaseClient, VastAiProviderPlugin, + VastAiProvisioningConfig, VastAiProvisioningPlugin, VastAiSshEndpoint, }; use parking_lot::Mutex; use swactor_vastai::{LifecyclePolicy, ProvisionRequest, ProvisionedInstance, SelectionPolicy}; @@ -88,7 +88,7 @@ impl VastAiLeaseClient for FakeLeaseClient { #[derive(Default)] struct FakeBootstrap { starts: Vec<(NodeProvisionSpec, VastAiSshEndpoint)>, - stops: Vec, + stops: Vec<(usize, BootstrapStopReason)>, fail: Option, next_handle: usize, } @@ -111,8 +111,8 @@ impl VastAiBootstrapLauncher for FakeBootstrap { Ok(self.next_handle) } - fn stop_bootstrap(&mut self, handle: &mut Self::Handle) { - self.stops.push(*handle); + fn stop_bootstrap(&mut self, handle: &mut Self::Handle, reason: BootstrapStopReason) { + self.stops.push((*handle, reason)); } } @@ -251,7 +251,38 @@ fn stop_destroys_known_vastai_contract_exactly_once() { plugin.stop_node(&handle).unwrap(); assert_eq!(plugin.client().destroyed, vec![100]); - assert_eq!(plugin.bootstrap().stops, vec![1]); + assert_eq!( + plugin.bootstrap().stops, + vec![(1, BootstrapStopReason::NodeStop)] + ); + assert_eq!(plugin.active_contract_count(), 0); +} + +#[test] +fn vastai_complete_bootstrap_stops_bootstrap_without_destroying_contract() { + let mut plugin = VastAiProvisioningPlugin::new( + FakeLeaseClient::default().with_contract(100), + FakeBootstrap::default(), + config(), + ); + let handle = plugin.start_node(spec(), sink()).unwrap(); + + plugin.complete_bootstrap(&handle).unwrap(); + + assert_eq!( + plugin.bootstrap().stops, + vec![(1, BootstrapStopReason::RuntimeReady)] + ); + assert_eq!(plugin.client().destroyed, Vec::::new()); + assert_eq!(plugin.active_contract_count(), 1); + + plugin.stop_node(&handle).unwrap(); + + assert_eq!(plugin.client().destroyed, vec![100]); + assert_eq!( + plugin.bootstrap().stops, + vec![(1, BootstrapStopReason::RuntimeReady)] + ); assert_eq!(plugin.active_contract_count(), 0); } diff --git a/crates/mvp-system/src/vastai_provisioning.rs b/crates/mvp-system/src/vastai_provisioning.rs index 1c2843e..d831a97 100644 --- a/crates/mvp-system/src/vastai_provisioning.rs +++ b/crates/mvp-system/src/vastai_provisioning.rs @@ -9,6 +9,9 @@ use std::sync::{ use std::time::Duration; use datastream::DatastreamProducer; +use serde::{Deserialize, Serialize}; +use swactor::actor::{ActorAddress, ActorInterface}; +use swactor::runtime::{Ctx, Runtime}; use swactor_vastai::{LifecyclePolicy, ProvisionRequest, ProvisionedInstance, SelectionPolicy}; use crate::bootstrap_datastream::{BootstrapDatastreamBridge, node_stream_id}; @@ -17,8 +20,7 @@ use crate::node_provisioning::{ ProviderError, ProviderKind, ProviderLeaseId, ProviderPlugin, SshEndpoint, }; use crate::provisioning::{ - NodeProvisionSpec, PluginNodeHandle, PluginObservation, PluginObservationSink, PluginSink, - ProvisionPlugin, + NodeProvisionSpec, PluginNodeHandle, PluginObservation, PluginSink, ProvisionPlugin, }; #[derive(Clone, Debug)] @@ -363,6 +365,12 @@ where } } +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub enum BootstrapStopReason { + RuntimeReady, + NodeStop, +} + pub trait VastAiBootstrapLauncher: Send { type Handle: Send; @@ -374,21 +382,63 @@ pub trait VastAiBootstrapLauncher: Send { producer: Option, ) -> Result; - fn stop_bootstrap(&mut self, handle: &mut Self::Handle); + fn stop_bootstrap(&mut self, handle: &mut Self::Handle, reason: BootstrapStopReason); } -#[derive(Default, Debug, Clone)] -pub struct SshCommandBootstrapLauncher { - ssh_identity: Option, +#[derive(Clone)] +enum SshBootstrapMsg { + Stop { reason: BootstrapStopReason }, } -pub struct SshCommandBootstrapHandle { + +struct SshBootstrapActor { child: Arc>>, stopping: Arc, } +impl SshBootstrapActor { + fn new(child: Arc>>, stopping: Arc) -> Self { + Self { child, stopping } + } +} + +impl ActorInterface for SshBootstrapActor { + type Incoming = SshBootstrapMsg; + type Response = (); + + fn handle(&mut self, _ctx: &Ctx, msg: Self::Incoming) { + match msg { + SshBootstrapMsg::Stop { reason: _ } => { + self.stopping.store(true, Ordering::SeqCst); + stop_ssh_child(&self.child); + } + } + } +} + +fn stop_ssh_child(child_slot: &Arc>>) { + let Some(mut child) = child_slot.lock().take() else { + return; + }; + let _ = child.kill(); + let _ = child.wait(); +} + +#[derive(Clone)] +pub struct SshCommandBootstrapLauncher { + ssh_identity: Option, + runtime: Arc, +} +pub struct SshCommandBootstrapHandle { + actor: ActorAddress, + runtime: Arc, +} + impl SshCommandBootstrapLauncher { - pub fn new(ssh_identity: Option) -> Self { - Self { ssh_identity } + pub fn new(ssh_identity: Option, runtime: Arc) -> Self { + Self { + ssh_identity, + runtime, + } } } @@ -411,39 +461,31 @@ impl VastAiBootstrapLauncher for SshCommandBootstrapLauncher { let child = Arc::new(Mutex::new(None)); let stopping = Arc::new(AtomicBool::new(false)); + let actor = self + .runtime + .spawn(SshBootstrapActor::new(child.clone(), stopping.clone())) + .map_err(|e| format!("spawn VastAI SSH bootstrap actor: {e}"))?; spawn_retrying_ssh_bootstrap( spec, endpoint, sink, producer, self.ssh_identity.clone(), - child.clone(), - stopping.clone(), + child, + stopping, ); - Ok(SshCommandBootstrapHandle { child, stopping }) + Ok(SshCommandBootstrapHandle { + actor, + runtime: Arc::clone(&self.runtime), + }) } - fn stop_bootstrap(&mut self, handle: &mut Self::Handle) { - handle.stopping.store(true, Ordering::SeqCst); - if let Some(mut child) = handle.child.lock().take() { - let _ = child.kill(); - let _ = child.wait(); - } - } -} - -struct ReadyTrackingSink { - inner: PluginSink, - ready: Arc, -} - -impl PluginObservationSink for ReadyTrackingSink { - fn observe(&self, observation: PluginObservation) { - if matches!(observation, PluginObservation::RuntimeReady { .. }) { - self.ready.store(true, Ordering::SeqCst); - } - self.inner.observe(observation); + fn stop_bootstrap(&mut self, handle: &mut Self::Handle, reason: BootstrapStopReason) { + let _ = handle + .runtime + .send_to(handle.actor, SshBootstrapMsg::Stop { reason }); + handle.runtime.tick(); } } @@ -459,15 +501,10 @@ fn spawn_retrying_ssh_bootstrap( std::thread::spawn(move || { let run_id = spec.run_id; let node_id = spec.node_id; - let ready = Arc::new(AtomicBool::new(false)); - let tracking_sink = PluginSink::new(Arc::new(ReadyTrackingSink { - inner: sink.clone(), - ready: ready.clone(), - })); let mut attempt = 1u64; let mut backoff = Duration::from_secs(1); - while !stopping.load(Ordering::SeqCst) && !ready.load(Ordering::SeqCst) { + while !stopping.load(Ordering::SeqCst) { sink.observe(PluginObservation::ProviderLine { run_id, node_id, @@ -482,7 +519,7 @@ fn spawn_retrying_ssh_bootstrap( *child_slot.lock() = Some(child); let bridge = BootstrapDatastreamBridge::new( spec.clone(), - tracking_sink.clone(), + sink.clone(), producer.clone(), ); bridge.spawn_stdout_reader(stdout); @@ -492,16 +529,6 @@ fn spawn_retrying_ssh_bootstrap( if stopping.load(Ordering::SeqCst) { return; } - if ready.load(Ordering::SeqCst) { - sink.observe(PluginObservation::ProviderLine { - run_id, - node_id, - line: - "VastAI SSH bootstrap observed runtime ready; handoff complete" - .to_owned(), - }); - return; - } let wait_result = { let mut guard = child_slot.lock(); @@ -563,7 +590,7 @@ fn spawn_retrying_ssh_bootstrap( } } - if stopping.load(Ordering::SeqCst) || ready.load(Ordering::SeqCst) { + if stopping.load(Ordering::SeqCst) { return; } sink.observe(PluginObservation::ProviderLine { @@ -818,12 +845,24 @@ where Ok(handle) } + fn complete_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, BootstrapStopReason::RuntimeReady); + } + Ok(()) + } + fn stop_node(&mut self, handle: &PluginNodeHandle) -> Result<(), String> { let Some(mut node) = self.nodes.remove(&handle.id) else { return Ok(()); }; if let Some(mut bootstrap) = node.bootstrap.take() { - self.bootstrap.stop_bootstrap(&mut bootstrap); + self.bootstrap + .stop_bootstrap(&mut bootstrap, BootstrapStopReason::NodeStop); } self.client.destroy_contract(node.contract_id) } @@ -873,7 +912,7 @@ mod tests { panic!("build_request tests must not start SSH bootstrap") } - fn stop_bootstrap(&mut self, _handle: &mut Self::Handle) { + fn stop_bootstrap(&mut self, _handle: &mut Self::Handle, _reason: BootstrapStopReason) { panic!("build_request tests must not stop SSH bootstrap") } } @@ -994,4 +1033,39 @@ mod tests { assert!(!args.iter().any(|arg| arg == "-i")); assert!(!args.iter().any(|arg| arg == "IdentitiesOnly=yes")); } + + #[test] + fn ssh_bootstrap_actor_stop_kills_child() { + let runtime = Arc::new(swactor::runtime::Runtime::new( + swactor::config::RuntimeConfig::default(), + )); + let child = std::process::Command::new("sleep") + .arg("30") + .spawn() + .expect("spawn sleep child"); + let pid = child.id(); + let child_slot = Arc::new(Mutex::new(Some(child))); + let stopping = Arc::new(AtomicBool::new(false)); + let actor = runtime + .spawn(SshBootstrapActor::new(child_slot.clone(), stopping.clone())) + .expect("spawn ssh bootstrap actor"); + + runtime + .send_to( + actor, + SshBootstrapMsg::Stop { + reason: BootstrapStopReason::RuntimeReady, + }, + ) + .expect("send stop"); + runtime.tick(); + + assert!(stopping.load(Ordering::SeqCst)); + assert!(child_slot.lock().is_none()); + #[cfg(target_os = "linux")] + assert!( + !std::path::Path::new(&format!("/proc/{pid}")).exists(), + "child process should be reaped" + ); + } }