zacheryasc/swactor
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

refactor ssh bootstrap logic

Browse source 
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 <zacheryasc@gmail.com>
This commit is contained in:
Zachery Aaron Shores-Chmielewski 2026-07-08 18:06:00 +04:00
parent 1ab6264ca0
commit f54f62b491
13 changed files with 613 additions and 207 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

28
apps/mvp-node/tinygrad_worker.py
View file

@ -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)

74
crates/mvp-system/src/actors/node_agent.rs
View file

@ -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::<NodeAgentMsg, _>(JsonCodec::<NodeAgentMsg>::default());
registry.register::<NodeAgentReport, _>(JsonCodec::<NodeAgentReport>::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::<OrchestratorMsg>()
.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,
})
);
}
}

93
crates/mvp-system/src/actors/orchestrator.rs
View file

@ -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<RunCommandWire>,
events: Vec<LifecycleEventWire>,
@ -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::<OrchestratorReport>()
.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);
}
}

95
crates/mvp-system/src/actors/provisioner.rs
View file

@ -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<P: ProvisionPlugin> ProvisionerActor<P> {
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<P: ProvisionPlugin> ProvisionerActor<P> {
}
}
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<P: ProvisionPlugin + 'static> ActorInterface for ProvisionerActor<P> {
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),
}
}

20
crates/mvp-system/src/bin/mvp_one_node_chat.rs
View file

@ -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()

172
crates/mvp-system/src/bin/orchestrator.rs
View file

@ -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::<OrchestratorReport>() {
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::<PromptEvent>() {
Ok(inbox) => inbox,
Err(error) => {
@ -265,7 +322,7 @@ fn run() -> Result<(), String> {
let (work_tx, work_rx) = mpsc::channel::<PromptWork>();
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<Box<dyn ProvisionPlugin>, String> {
fn build_provisioner(
&self,
bootstrap_runtime: Arc<swactor::runtime::Runtime>,
) -> Result<Box<dyn ProvisionPlugin>, 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<NodeProvisionSpec, String> {
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<PluginObservation>,
frame_rx: &mpsc::Receiver<(StreamId, Frame)>,
orchestrator_reports: &swactor::runtime::Inbox<OrchestratorReport>,
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<PathBuf>) -> Result<PathBuf, String> {
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,

45
crates/mvp-system/src/bin/worker_node.rs
View file

@ -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 || {

29
crates/mvp-system/src/bootstrap_datastream.rs
View file

@ -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<String>) {
@ -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<PluginObservation> {
let ready = serde_json::from_str::<RuntimeReadyLine>(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)
}

3
crates/mvp-system/src/config.rs
View file

@ -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()),

15
crates/mvp-system/src/provisioning.rs
View file

@ -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<u32>,
endpoint: EndpointAddr,
node_actor: ActorAddress,
},
Exited {
run_id: u64,
node_id: u64,
@ -144,6 +135,8 @@ pub trait ProvisionPlugin: Send {
sink: PluginSink,
) -> Result<PluginNodeHandle, String>;
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(());

25
crates/mvp-system/src/tests/bootstrap_datastream_guarantees.rs
View file

@ -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,
}]
);
}

43
crates/mvp-system/src/tests/vastai_provisioning_guarantees.rs
View file

@ -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<usize>,
stops: Vec<(usize, BootstrapStopReason)>,
fail: Option<String>,
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::<u64>::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);
}

178
crates/mvp-system/src/vastai_provisioning.rs
View file

@ -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<DatastreamProducer>,
) -> Result<Self::Handle, String>;
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<PathBuf>,
#[derive(Clone)]
enum SshBootstrapMsg {
Stop { reason: BootstrapStopReason },
}
pub struct SshCommandBootstrapHandle {
struct SshBootstrapActor {
child: Arc<Mutex<Option<Child>>>,
stopping: Arc<AtomicBool>,
}
impl SshBootstrapActor {
fn new(child: Arc<Mutex<Option<Child>>>, stopping: Arc<AtomicBool>) -> 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<Mutex<Option<Child>>>) {
let Some(mut child) = child_slot.lock().take() else {
return;
};
let _ = child.kill();
let _ = child.wait();
}
#[derive(Clone)]
pub struct SshCommandBootstrapLauncher {
ssh_identity: Option<PathBuf>,
runtime: Arc<Runtime>,
}
pub struct SshCommandBootstrapHandle {
actor: ActorAddress,
runtime: Arc<Runtime>,
}
impl SshCommandBootstrapLauncher {
pub fn new(ssh_identity: Option<PathBuf>) -> Self {
Self { ssh_identity }
pub fn new(ssh_identity: Option<PathBuf>, runtime: Arc<Runtime>) -> 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<AtomicBool>,
}
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"
);
}
}