refactor: actorization polish and stability

This commit is contained in:
Zachery Aaron Shores-Chmielewski 2026-07-29 12:14:52 +04:00
parent ce4845ab8d
commit d973340043
11 changed files with 1326 additions and 530 deletions

View file

@ -30,6 +30,10 @@ use crate::chat::node_image::{
use crate::observability::{benchmark, frame_archive::FrameArchive};
use crate::orchestration::node_provisioning::{ProviderKind, provider_kind};
use crate::orchestration::provider_adapters::vastai::config::ResolvedVastAiConfig;
use crate::orchestration::{
DEFAULT_PIPELINE_CACHED_MODEL_FILE, DEFAULT_PIPELINE_CACHED_MODEL_ID,
DEFAULT_PIPELINE_CACHED_MODEL_MAX_CONTEXT, DEFAULT_PIPELINE_CACHED_MODEL_REPO,
};
use crate::prompt::rpc::{PromptEvent, SubmitPrompt, write_json_line};
use crate::transport::endpoint_advertisement::EndpointAddrMask;
@ -814,6 +818,16 @@ impl Config {
let cached_model = cached_model_source
.map(CachedModelConfig::from_source)
.transpose()?;
let model = if provider == provider_kind::vastai() {
match &cached_model {
Some(cached_model) => {
vastai_model_config_for_cached_model(toml.model.clone(), cached_model)?
}
None => toml.model.clone(),
}
} else {
toml.model.clone()
};
let datastream_frame_log = if args.dump_logs {
Some(
args.dump_log_path
@ -847,7 +861,7 @@ impl Config {
vastai_yes: args.vastai_yes,
pipeline_stages,
max_tokens,
model: toml.model,
model,
vastai,
skip_rebuild: args.skip_rebuild,
gpu_run,
@ -1097,6 +1111,40 @@ impl ParsedArgs {
}
}
fn vastai_model_config_for_cached_model(
mut model: ChatModelConfig,
cached_model: &CachedModelConfig,
) -> Result<ChatModelConfig, String> {
let file_name = cached_model
.host_path
.file_name()
.and_then(|name| name.to_str())
.ok_or_else(|| {
format!(
"cached model path {} does not have a UTF-8 file name",
cached_model.host_path.display()
)
})?;
if file_name == DEFAULT_PIPELINE_CACHED_MODEL_FILE {
model.id = Some(DEFAULT_PIPELINE_CACHED_MODEL_ID.to_owned());
model.gguf_local_path = None;
model.gguf_repo = Some(DEFAULT_PIPELINE_CACHED_MODEL_REPO.to_owned());
model.gguf_file = Some(DEFAULT_PIPELINE_CACHED_MODEL_FILE.to_owned());
model.gguf_revision = None;
model.max_context = Some(DEFAULT_PIPELINE_CACHED_MODEL_MAX_CONTEXT);
return Ok(model);
}
if model.gguf_file.as_deref() == Some(file_name) {
model.gguf_local_path = None;
return Ok(model);
}
Err(format!(
"VastAI cached model {} does not match configured remote GGUF {}; use --cached-model=<matching .gguf> or configure [model].gguf_repo and [model].gguf_file for that cache",
cached_model.host_path.display(),
model.gguf_file.as_deref().unwrap_or("<unset>")
))
}
fn resolve_vastai_config(
file: &ChatVastAiConfig,
node_image: &str,
@ -3222,6 +3270,72 @@ relay_url = "https://relay.example"
});
}
#[test]
fn vastai_cached_pipeline_model_selects_matching_remote_gguf() {
let temp = TempDir::new("vastai-cached-pipeline-model");
let cached_path = temp.path().join(DEFAULT_PIPELINE_CACHED_MODEL_FILE);
fs::write(&cached_path, b"cached model").expect("write cached model");
let config_path = write_config(
&temp,
"chat.toml",
r#"
[provider]
kind = "vastai"
[image]
node = "docker.io/acme/node:latest"
[model]
id = "qwen2.5-7b-instruct-q4-k-m"
gguf_repo = "bartowski/Qwen2.5-7B-Instruct-GGUF"
gguf_file = "Qwen2.5-7B-Instruct-Q4_K_M.gguf"
max_context = 512
[vastai]
relay_url = "https://relay.example"
bootstrap_command = "boot"
"#,
);
let cached_arg = format!("--cached-model={}", cached_path.display());
with_process_state(&[("VAST_API_KEY", Some("secret"))], None, || {
let config_arg = config_path.to_string_lossy().into_owned();
let config = Config::from_args(strings(&[
"--config",
config_arg.as_str(),
cached_arg.as_str(),
"--yes",
]))
.expect("VastAI cached pipeline model resolves");
let args = config.orchestrator_cli_args("docker.io/acme/node:prepared");
assert_eq!(
config
.cached_model
.as_ref()
.and_then(|model| model.host_path.file_name()),
Some(OsStr::new(DEFAULT_PIPELINE_CACHED_MODEL_FILE))
);
assert!(
args.windows(2)
.any(|pair| pair == ["--model-id", DEFAULT_PIPELINE_CACHED_MODEL_ID])
);
assert!(
args.windows(2)
.any(|pair| pair == ["--gguf-repo", DEFAULT_PIPELINE_CACHED_MODEL_REPO])
);
assert!(
args.windows(2)
.any(|pair| pair == ["--gguf-file", DEFAULT_PIPELINE_CACHED_MODEL_FILE])
);
let expected_context = DEFAULT_PIPELINE_CACHED_MODEL_MAX_CONTEXT.to_string();
assert!(
args.windows(2)
.any(|pair| pair == ["--max-context", expected_context.as_str()])
);
});
}
fn panic_prepare_node_image(_: NodeImageRequest) -> Result<PreparedNodeImage, String> {
panic!("image preparer must not be called when --skip-rebuild is set")
}

View file

@ -166,7 +166,9 @@ pub enum NodeAgentMsg {
StepCompleted {
step_id: u64,
},
WorkerCrashed,
WorkerCrashed {
reason: Option<String>,
},
StopRun {
run_id: u64,
},
@ -323,6 +325,7 @@ pub struct NodeAgentActor {
outbound_edge: Option<StageOutboundEdgeWire>,
command_cursor: usize,
event_cursor: usize,
last_worker_crash: Option<String>,
}
impl NodeAgentActor {
@ -339,6 +342,7 @@ impl NodeAgentActor {
outbound_edge: None,
command_cursor: 0,
event_cursor: 0,
last_worker_crash: None,
}
}
@ -448,7 +452,10 @@ impl NodeAgentActor {
step_id: stage::StepId(step_id),
})
}
NodeAgentMsg::WorkerCrashed => self.core.observe(stage::StageEvent::WorkerCrashed),
NodeAgentMsg::WorkerCrashed { reason } => {
self.last_worker_crash = reason;
self.core.observe(stage::StageEvent::WorkerCrashed)
}
NodeAgentMsg::StopRun { run_id } => self.core.observe(stage::StageEvent::StopRun {
run_id: stage::RunId(run_id),
}),
@ -584,6 +591,7 @@ impl NodeAgentActor {
OrchestratorMsg::ObserveStageFault {
run_id: run_id.0,
stage_index: *stage_index,
reason: self.last_worker_crash.clone(),
},
);
}

View file

@ -55,7 +55,12 @@ impl ActorInterface for MvpDataPlaneReportSinkActor {
dp::DataPlaneReportMsg::ObjectProduced { .. } => {}
dp::DataPlaneReportMsg::EdgeFaulted { .. }
| dp::DataPlaneReportMsg::WorkerDataPlaneFaulted { .. } => {
let _ = ctx.send(self.node_agent, NodeAgentMsg::WorkerCrashed);
let _ = ctx.send(
self.node_agent,
NodeAgentMsg::WorkerCrashed {
reason: Some("data plane faulted".to_owned()),
},
);
}
dp::DataPlaneReportMsg::EdgeStopped { .. } => {}
dp::DataPlaneReportMsg::LocalEdgesStopped { run_id } => {

View file

@ -32,7 +32,9 @@ use crate::orchestration::provider_adapters::relay::relay_runtime_config_from_en
use crate::orchestration::run_plan::{GgufSource, TokenizerSource};
use crate::prompt::rpc::{PromptEvent, TokenizerEvent};
use crate::staging::control as stage;
use crate::staging::gguf_shard::{StageShardPlan, materialize_stage_shard_http};
use crate::staging::gguf_shard::{
StageShardPlan, materialize_stage_shard_http, validate_stage_shard_cache,
};
use crate::transport::codec_registry::register_mvp_actor_codecs;
use crate::transport::driver_pumps as driver_model;
use crate::transport::endpoint_advertisement::{
@ -52,7 +54,8 @@ use iroh_driver::{
};
use parking_lot::Mutex;
use serde_json::{Value, json};
use swactor::actor::ActorAddress;
use swactor::actor::{ActorAddress, ActorInterface};
use swactor::runtime::{Ctx, ExternalSender};
use tokio::io::{AsyncBufReadExt, AsyncWriteExt};
const DEFAULT_WORKER_SCRIPT: &str = "/usr/local/share/mvp/tinygrad_worker.py";
@ -1530,7 +1533,12 @@ impl WorkerEdgeRuntime {
edge::EdgeLifecycleEvent::EdgeFaulted { edge_id, reason } => {
stack
.runtime
.send_to(node_actor, NodeAgentMsg::WorkerCrashed)
.send_to(
node_actor,
NodeAgentMsg::WorkerCrashed {
reason: Some(format!("edge {} faulted: {reason:?}", edge_id.0)),
},
)
.map_err(|e| format!("mark worker crashed after edge fault: {e}"))?;
return Err(format!("edge {} faulted: {reason:?}", edge_id.0));
}
@ -2279,9 +2287,12 @@ fn run() -> Result<(), String> {
"failed",
json!({"exit_status":status.to_string()}),
);
let _ = stack
.runtime
.send_to(node_actor, NodeAgentMsg::WorkerCrashed);
let _ = stack.runtime.send_to(
node_actor,
NodeAgentMsg::WorkerCrashed {
reason: Some(format!("tinygrad helper exited with {status}")),
},
);
pump_network(&mut driver, &stack);
return Err(format!("tinygrad helper exited with {status}"));
}
@ -3101,15 +3112,290 @@ fn handle_decode_tokens_request(
.map_err(|e| format!("send tokenizer decode response: {e}"))
}
#[derive(Clone, Copy)]
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum StageShardProcessStream {
Stdout,
Stderr,
}
struct StageShardProcessLine {
#[derive(Clone)]
enum StageShardFetchMsg {
Start,
PollChild,
ProcessLine {
stream: StageShardProcessStream,
line: String,
},
ReaderError {
stream: StageShardProcessStream,
error: String,
},
ReaderClosed {
stream: StageShardProcessStream,
},
}
#[derive(Clone, Debug)]
enum StageShardFetchReport {
Progress(Value),
Done(PathBuf),
Failed(String),
}
struct StageShardFetchActor {
request_json: Vec<u8>,
output_path: PathBuf,
report_to: ActorAddress,
sender: ExternalSender,
child: Option<Child>,
stdout_reader: Option<thread::JoinHandle<()>>,
stderr_reader: Option<thread::JoinHandle<()>>,
stdout_closed: bool,
stderr_closed: bool,
ready_path: Option<PathBuf>,
exit_status: Option<std::process::ExitStatus>,
finished: bool,
}
impl StageShardFetchActor {
fn new(
request_json: Vec<u8>,
output_path: PathBuf,
report_to: ActorAddress,
sender: ExternalSender,
) -> Self {
Self {
request_json,
output_path,
report_to,
sender,
child: None,
stdout_reader: None,
stderr_reader: None,
stdout_closed: true,
stderr_closed: true,
ready_path: None,
exit_status: None,
finished: false,
}
}
fn start_fetch(&mut self, ctx: &Ctx) {
if self.finished {
return;
}
let exe = match std::env::current_exe() {
Ok(exe) => exe,
Err(error) => {
self.fail(ctx, format!("locate worker node executable: {error}"));
return;
}
};
let mut child = match Command::new(exe)
.arg("stage-shard-fetcher")
.stdin(Stdio::piped())
.stdout(Stdio::piped())
.stderr(Stdio::piped())
.spawn()
{
Ok(child) => child,
Err(error) => {
self.fail(ctx, format!("spawn stage shard fetcher: {error}"));
return;
}
};
match child.stdin.take() {
Some(mut stdin) => {
if let Err(error) = stdin.write_all(&self.request_json) {
let mut child = Some(child);
stop_stage_shard_child(&mut child);
self.fail(ctx, format!("write stage shard fetch request: {error}"));
return;
}
}
None => {
let mut child = Some(child);
stop_stage_shard_child(&mut child);
self.fail(ctx, "stage shard fetcher stdin missing".to_owned());
return;
}
}
self.stdout_closed = false;
self.stderr_closed = false;
self.ready_path = None;
self.exit_status = None;
if let Some(stdout) = child.stdout.take() {
self.stdout_reader = Some(spawn_stage_shard_reader(
StageShardProcessStream::Stdout,
stdout,
self.sender.clone(),
ctx.self_addr(),
));
} else {
self.stdout_closed = true;
}
if let Some(stderr) = child.stderr.take() {
self.stderr_reader = Some(spawn_stage_shard_reader(
StageShardProcessStream::Stderr,
stderr,
self.sender.clone(),
ctx.self_addr(),
));
} else {
self.stderr_closed = true;
}
self.child = Some(child);
schedule_stage_shard_message(
self.sender.clone(),
ctx.self_addr(),
StageShardFetchMsg::PollChild,
PUMP_INTERVAL,
);
}
fn poll_child(&mut self, ctx: &Ctx) {
if self.finished {
return;
}
let Some(child) = self.child.as_mut() else {
return;
};
match child.try_wait() {
Ok(Some(status)) => {
self.child = None;
self.exit_status = Some(status);
self.maybe_finish(ctx);
}
Ok(None) => schedule_stage_shard_message(
self.sender.clone(),
ctx.self_addr(),
StageShardFetchMsg::PollChild,
PUMP_INTERVAL,
),
Err(error) => {
self.child = None;
self.fail(ctx, format!("poll stage shard fetcher: {error}"));
}
}
}
fn handle_line(&mut self, ctx: &Ctx, stream: StageShardProcessStream, line: String) {
if line.is_empty() || self.finished {
return;
}
let event = match stream {
StageShardProcessStream::Stdout => match serde_json::from_str::<Value>(&line) {
Ok(value) => value,
Err(error) => {
json!({"type":"StageShardFetchOutputParseFailed","line":line,"error":error.to_string()})
}
},
StageShardProcessStream::Stderr => json!({"type":"StageShardFetchStderr","line":line}),
};
if event.get("type").and_then(Value::as_str) == Some("StageShardReady") {
self.ready_path = event
.get("path")
.and_then(Value::as_str)
.map(PathBuf::from)
.or_else(|| Some(self.output_path.clone()));
}
let _ = ctx.send(self.report_to, StageShardFetchReport::Progress(event));
}
fn handle_reader_error(&mut self, ctx: &Ctx, stream: StageShardProcessStream, error: String) {
self.handle_line(ctx, stream, format!("reader error: {error}"));
}
fn handle_reader_closed(&mut self, ctx: &Ctx, stream: StageShardProcessStream) {
match stream {
StageShardProcessStream::Stdout => self.stdout_closed = true,
StageShardProcessStream::Stderr => self.stderr_closed = true,
}
self.maybe_finish(ctx);
}
fn maybe_finish(&mut self, ctx: &Ctx) {
if self.finished || self.exit_status.is_none() || !self.stdout_closed || !self.stderr_closed
{
return;
}
self.join_readers();
let status = self.exit_status.take().expect("exit status checked");
if status.success() {
let path = self
.ready_path
.clone()
.unwrap_or_else(|| self.output_path.clone());
if path.is_file() {
self.finished = true;
let _ = ctx.send(self.report_to, StageShardFetchReport::Done(path));
ctx.stop_self();
return;
}
self.fail(
ctx,
format!(
"stage shard fetcher exited successfully but {} is missing",
path.display()
),
);
return;
}
self.fail(ctx, format!("stage shard fetcher exited with {status}"));
}
fn fail(&mut self, ctx: &Ctx, error: String) {
if self.finished {
return;
}
self.finished = true;
stop_stage_shard_child(&mut self.child);
self.join_readers();
let _ = ctx.send(self.report_to, StageShardFetchReport::Failed(error));
ctx.stop_self();
}
fn join_readers(&mut self) {
if let Some(reader) = self.stdout_reader.take() {
let _ = reader.join();
}
if let Some(reader) = self.stderr_reader.take() {
let _ = reader.join();
}
self.stdout_closed = true;
self.stderr_closed = true;
}
}
impl ActorInterface for StageShardFetchActor {
type Incoming = StageShardFetchMsg;
type Response = ();
fn handle(&mut self, ctx: &Ctx, msg: Self::Incoming) {
match msg {
StageShardFetchMsg::Start => self.start_fetch(ctx),
StageShardFetchMsg::PollChild => self.poll_child(ctx),
StageShardFetchMsg::ProcessLine { stream, line } => self.handle_line(ctx, stream, line),
StageShardFetchMsg::ReaderError { stream, error } => {
self.handle_reader_error(ctx, stream, error)
}
StageShardFetchMsg::ReaderClosed { stream } => self.handle_reader_closed(ctx, stream),
}
}
fn on_stop(&mut self, _ctx: &Ctx) {
stop_stage_shard_child(&mut self.child);
self.join_readers();
}
}
fn stop_stage_shard_child(child: &mut Option<Child>) {
let Some(mut child) = child.take() else {
return;
};
let _ = child.kill();
let _ = child.wait();
}
fn stage_shard_cache_path(plan: &StageShardPlan) -> PathBuf {
@ -3124,8 +3410,9 @@ fn stage_shard_cache_path(plan: &StageShardPlan) -> PathBuf {
fn spawn_stage_shard_reader<R: Read + Send + 'static>(
stream: StageShardProcessStream,
reader: R,
tx: mpsc::Sender<StageShardProcessLine>,
) {
sender: ExternalSender,
actor: ActorAddress,
) -> thread::JoinHandle<()> {
thread::spawn(move || {
let mut reader = BufReader::new(reader);
let mut line = String::new();
@ -3134,20 +3421,39 @@ fn spawn_stage_shard_reader<R: Read + Send + 'static>(
match reader.read_line(&mut line) {
Ok(0) => break,
Ok(_) => {
let _ = tx.send(StageShardProcessLine {
let _ = sender.send_to(
actor,
StageShardFetchMsg::ProcessLine {
stream,
line: line.trim_end_matches(['\r', '\n']).to_owned(),
});
},
);
}
Err(error) => {
let _ = tx.send(StageShardProcessLine {
let _ = sender.send_to(
actor,
StageShardFetchMsg::ReaderError {
stream,
line: format!("reader error: {error}"),
});
error: error.to_string(),
},
);
break;
}
}
}
let _ = sender.send_to(actor, StageShardFetchMsg::ReaderClosed { stream });
})
}
fn schedule_stage_shard_message(
sender: ExternalSender,
actor: ActorAddress,
msg: StageShardFetchMsg,
delay: Duration,
) {
thread::spawn(move || {
thread::sleep(delay);
let _ = sender.send_to(actor, msg);
});
}
@ -3174,6 +3480,8 @@ fn materialize_stage_shard_with_process(
) -> Result<PathBuf, String> {
let output_path = stage_shard_cache_path(plan);
if output_path.is_file() {
match validate_stage_shard_cache(&output_path, plan) {
Ok(()) => {
let event = json!({
"type":"StageShardCacheReady",
"stage_index":plan.stage_index,
@ -3183,6 +3491,24 @@ fn materialize_stage_shard_with_process(
publish_stage_shard_fetch_event(datastream, config, &event)?;
return Ok(output_path);
}
Err(error) => {
let event = json!({
"type":"StageShardCacheInvalid",
"stage_index":plan.stage_index,
"path":output_path,
"cache_hit":false,
"error":error,
});
publish_stage_shard_fetch_event(datastream, config, &event)?;
std::fs::remove_file(&output_path).map_err(|remove_error| {
format!(
"remove invalid stage shard cache {}: {remove_error}",
output_path.display()
)
})?;
}
}
}
let request = StageShardFetchRequest {
plan: plan.clone(),
@ -3190,95 +3516,39 @@ fn materialize_stage_shard_with_process(
};
let request_json = serde_json::to_vec(&request)
.map_err(|e| format!("serialize stage shard fetch request: {e}"))?;
let exe = std::env::current_exe().map_err(|e| format!("locate worker node executable: {e}"))?;
let mut child = Command::new(exe)
.arg("stage-shard-fetcher")
.stdin(Stdio::piped())
.stdout(Stdio::piped())
.stderr(Stdio::piped())
.spawn()
.map_err(|e| format!("spawn stage shard fetcher: {e}"))?;
if let Some(mut stdin) = child.stdin.take() {
stdin
.write_all(&request_json)
.map_err(|e| format!("write stage shard fetch request: {e}"))?;
}
let (tx, rx) = mpsc::channel::<StageShardProcessLine>();
if let Some(stdout) = child.stdout.take() {
spawn_stage_shard_reader(StageShardProcessStream::Stdout, stdout, tx.clone());
}
if let Some(stderr) = child.stderr.take() {
spawn_stage_shard_reader(StageShardProcessStream::Stderr, stderr, tx);
}
let mut ready_path = None;
let reports = stack
.runtime
.new_inbox::<StageShardFetchReport>()
.map_err(|e| format!("stage shard fetch report inbox: {e}"))?;
let actor = stack
.runtime
.spawn(StageShardFetchActor::new(
request_json,
output_path,
*reports.addr(),
stack.runtime.create_sender(),
))
.map_err(|e| format!("spawn stage shard fetch actor: {e}"))?;
stack
.runtime
.send_to(actor, StageShardFetchMsg::Start)
.map_err(|e| format!("start stage shard fetch actor: {e}"))?;
loop {
while let Ok(line) = rx.try_recv() {
if line.line.is_empty() {
continue;
}
let event = match line.stream {
StageShardProcessStream::Stdout => {
match serde_json::from_str::<Value>(&line.line) {
Ok(value) => value,
Err(error) => {
json!({"type":"StageShardFetchOutputParseFailed","line":line.line,"error":error.to_string()})
}
}
}
StageShardProcessStream::Stderr => {
json!({"type":"StageShardFetchStderr","line":line.line})
}
};
if event.get("type").and_then(Value::as_str) == Some("StageShardReady") {
ready_path = event
.get("path")
.and_then(Value::as_str)
.map(PathBuf::from)
.or_else(|| Some(output_path.clone()));
}
publish_stage_shard_fetch_event(datastream, config, &event)?;
}
if let Some(status) = child
.try_wait()
.map_err(|e| format!("poll stage shard fetcher: {e}"))?
{
while let Ok(line) = rx.try_recv() {
if line.line.is_empty() {
continue;
}
let event = match line.stream {
StageShardProcessStream::Stdout => serde_json::from_str::<Value>(&line.line)
.unwrap_or_else(|error| {
json!({"type":"StageShardFetchOutputParseFailed","line":line.line,"error":error.to_string()})
}),
StageShardProcessStream::Stderr => {
json!({"type":"StageShardFetchStderr","line":line.line})
}
};
if event.get("type").and_then(Value::as_str) == Some("StageShardReady") {
ready_path = event
.get("path")
.and_then(Value::as_str)
.map(PathBuf::from)
.or_else(|| Some(output_path.clone()));
}
publish_stage_shard_fetch_event(datastream, config, &event)?;
}
if status.success() {
let path = ready_path.unwrap_or_else(|| output_path.clone());
if path.is_file() {
return Ok(path);
}
return Err(format!(
"stage shard fetcher exited successfully but {} is missing",
path.display()
));
}
return Err(format!("stage shard fetcher exited with {status}"));
}
pump_network(driver, stack);
while let Some(report) = reports.try_recv() {
match report {
StageShardFetchReport::Progress(event) => {
if let Err(error) = publish_stage_shard_fetch_event(datastream, config, &event)
{
let _ = stack.runtime.stop_actor(actor);
return Err(error);
}
}
StageShardFetchReport::Done(path) => return Ok(path),
StageShardFetchReport::Failed(error) => return Err(error),
}
}
datastream.tick();
thread::sleep(PUMP_INTERVAL);
}
@ -3337,9 +3607,12 @@ fn handle_stage_command(
"failed",
json!({"error":error}),
);
let _ = stack
.runtime
.send_to(node_actor, NodeAgentMsg::WorkerCrashed);
let _ = stack.runtime.send_to(
node_actor,
NodeAgentMsg::WorkerCrashed {
reason: Some(error.clone()),
},
);
pump();
return Err(error);
}
@ -3367,17 +3640,36 @@ fn handle_stage_command(
};
let (resolved_gguf_source, using_stage_shard) =
if let Some(stage_plan) = stage_shard_plan {
let local_path = materialize_stage_shard_with_process(
match materialize_stage_shard_with_process(
&stage_plan,
config,
datastream,
driver,
stack,
)?;
(
) {
Ok(local_path) => (
GgufSource::LocalPath(local_path.to_string_lossy().into_owned()),
true,
)
),
Err(error) => {
emit_node_event(
datastream,
config,
NODE_STAGE_CHANNEL,
"load_weights",
"failed",
json!({"error":error,"stage_shard":true}),
);
let _ = stack.runtime.send_to(
node_actor,
NodeAgentMsg::WorkerCrashed {
reason: Some(error.clone()),
},
);
pump_network(driver, stack);
return Err(error);
}
}
} else {
(gguf_source, false)
};
@ -3417,9 +3709,12 @@ fn handle_stage_command(
"failed",
json!({"error":error}),
);
let _ = stack
.runtime
.send_to(node_actor, NodeAgentMsg::WorkerCrashed);
let _ = stack.runtime.send_to(
node_actor,
NodeAgentMsg::WorkerCrashed {
reason: Some(error.clone()),
},
);
pump();
return Err(error);
}

View file

@ -57,6 +57,7 @@ pub enum OrchestratorMsg {
ObserveStageFault {
run_id: u64,
stage_index: u32,
reason: Option<String>,
},
ObserveEndpointFault {
run_id: u64,
@ -168,6 +169,7 @@ pub enum OrchestratorReport {
StageFault {
run_id: u64,
stage_index: u32,
reason: Option<String>,
},
Snapshot {
commands: Vec<RunCommandWire>,
@ -247,6 +249,7 @@ impl OrchestratorActor {
OrchestratorMsg::ObserveStageFault {
run_id,
stage_index,
reason: _,
} => self.core.observe(core::RunEvent::StageFault {
run_id: core::RunId(run_id),
stage_index,
@ -396,9 +399,11 @@ impl ActorInterface for OrchestratorActor {
OrchestratorMsg::ObserveStageFault {
run_id,
stage_index,
reason,
} => Some(OrchestratorReport::StageFault {
run_id,
stage_index,
reason,
}),
_ => None,
};

View file

@ -77,7 +77,11 @@ const MVP_RUNTIME_CONFIG_ENV: &str = "MVP_RUNTIME_CONFIG";
const CACHED_MODEL_HOST_ENV: &str = "MVP_CACHED_MODEL_HOST_PATH";
const MVP_WORKER_BIN_ENV: &str = "MVP_WORKER_BIN";
const CACHED_MODEL_CONTAINER_DIR: &str = "/models/cached";
const DEFAULT_PIPELINE_CACHED_MODEL_FILE: &str = "SmolLM2-135M-Instruct.Q4_0.gguf";
pub(crate) const DEFAULT_PIPELINE_CACHED_MODEL_FILE: &str = "SmolLM2-135M-Instruct.Q4_0.gguf";
pub(crate) const DEFAULT_PIPELINE_CACHED_MODEL_REPO: &str =
"QuantFactory/SmolLM2-135M-Instruct-GGUF";
pub(crate) const DEFAULT_PIPELINE_CACHED_MODEL_ID: &str = "smollm2-135m-instruct-q4";
pub(crate) const DEFAULT_PIPELINE_CACHED_MODEL_MAX_CONTEXT: u32 = 256;
const DEFAULT_PIPELINE_MODEL_CACHE_DIR: &str = ".model-cache";
const DEFAULT_RPC_BIND: &str = "127.0.0.1:19777";
const DEFAULT_HF_REPO: &str = "bartowski/Llama-3.2-1B-Instruct-GGUF";
@ -2365,6 +2369,21 @@ impl ProvisionedClusterGuard {
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 {
@ -2591,7 +2610,7 @@ fn start_and_provision_workers(
}
pending_specs = start_outcome.failed_specs;
}
let provisioned_nodes = ProvisionedClusterGuard::new(provisioner, handles);
let mut provisioned_nodes = ProvisionedClusterGuard::new(provisioner, handles);
drain_orch_stdio_capture(
orch_stdio_rx,
orch_datastream,
@ -2702,6 +2721,9 @@ fn start_and_provision_workers(
&ack_targets,
&pipeline_coordinator,
)?;
provisioned_nodes
.complete_bootstrap()
.map_err(|e| format!("complete provider bootstrap after runtime-ready: {e}"))?;
orch_datastream.emit_bootstrap(
dashboard,
@ -3347,10 +3369,15 @@ fn wait_for_weights_loaded_count(
OrchestratorReport::StageFault {
run_id: report_run_id,
stage_index,
reason,
} if report_run_id == run_id && expected_stages.contains(&stage_index) => {
return Err(format!(
"stage {stage_index} faulted while loading pipeline weights"
));
let mut error =
format!("stage {stage_index} faulted while loading pipeline weights");
if let Some(reason) = reason {
error.push_str(": ");
error.push_str(&reason);
}
return Err(error);
}
_ => {}
}
@ -4586,10 +4613,15 @@ fn wait_for_weights_loaded(
OrchestratorReport::StageFault {
run_id: report_run_id,
stage_index: report_stage_index,
reason,
} if report_run_id == run_id && report_stage_index == stage_index => {
return Err(format!(
"stage {report_stage_index} faulted while loading weights"
));
let mut error =
format!("stage {report_stage_index} faulted while loading weights");
if let Some(reason) = reason {
error.push_str(": ");
error.push_str(&reason);
}
return Err(error);
}
_ => {}
}
@ -8958,9 +8990,16 @@ bootstrap_command = "/run"
}
}
#[derive(Default)]
struct FakeProvisionState {
stopped: Vec<u64>,
completed: Vec<u64>,
}
#[derive(Default)]
struct FakeProvisionPlugin {
stopped: Vec<u64>,
shared: Option<Arc<std::sync::Mutex<FakeProvisionState>>>,
}
impl ProvisionPlugin for FakeProvisionPlugin {
@ -8974,8 +9013,15 @@ bootstrap_command = "/run"
fn complete_bootstrap(
&mut self,
_handle: &crate::orchestration::provisioning::PluginNodeHandle,
handle: &crate::orchestration::provisioning::PluginNodeHandle,
) -> Result<(), String> {
if let Some(shared) = &self.shared {
shared
.lock()
.expect("fake provision state")
.completed
.push(handle.id);
}
Ok(())
}
@ -8984,6 +9030,13 @@ bootstrap_command = "/run"
handle: &crate::orchestration::provisioning::PluginNodeHandle,
) -> Result<(), String> {
self.stopped.push(handle.id);
if let Some(shared) = &self.shared {
shared
.lock()
.expect("fake provision state")
.stopped
.push(handle.id);
}
Ok(())
}
}
@ -9032,6 +9085,55 @@ bootstrap_command = "/run"
assert_eq!(outcome.results.len(), 2);
}
#[test]
fn provisioned_cluster_guard_completes_bootstrap_without_dropping_stop_handles() {
let state = Arc::new(std::sync::Mutex::new(FakeProvisionState::default()));
let plugin = FakeProvisionPlugin {
shared: Some(Arc::clone(&state)),
..FakeProvisionPlugin::default()
};
let mut guard = ProvisionedClusterGuard::new(
Box::new(plugin),
vec![
crate::orchestration::provisioning::PluginNodeHandle {
id: 7,
provider_process_id: None,
},
crate::orchestration::provisioning::PluginNodeHandle {
id: 8,
provider_process_id: None,
},
],
);
guard
.complete_bootstrap()
.expect("runtime-ready bootstrap completion succeeds");
assert_eq!(
state
.lock()
.expect("fake provision state")
.completed
.clone(),
vec![7, 8]
);
assert!(
state
.lock()
.expect("fake provision state")
.stopped
.is_empty()
);
guard
.stop()
.expect("node stop still succeeds after completion");
assert_eq!(
state.lock().expect("fake provision state").stopped.clone(),
vec![8, 7]
);
}
#[test]
fn provisioned_node_guard_stops_node_on_drop() {
let mut plugin = FakeProvisionPlugin::default();

View file

@ -10,6 +10,7 @@ const GGUF_MAGIC: &[u8; 4] = b"GGUF";
const SUPPORTED_GGUF_VERSION: u32 = 3;
const DEFAULT_ALIGNMENT: u64 = 32;
const MAX_STRING_BYTES: u64 = 64 * 1024 * 1024;
const STAGE_SHARD_CACHE_FORMAT_VERSION: &str = "stage-shard-cache-v2";
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct ByteRange {
@ -166,6 +167,45 @@ pub fn plan_stage_shard(
})
}
pub fn validate_stage_shard_cache(path: &Path, plan: &StageShardPlan) -> Result<(), String> {
let directory = read_gguf_directory(path)
.map_err(|error| format!("invalid cached stage shard {}: {error}", path.display()))?;
if directory.tensors.len() != plan.tensors.len() {
return Err(format!(
"cached stage shard {} has {} tensors; expected {}",
path.display(),
directory.tensors.len(),
plan.tensors.len()
));
}
if directory.alignment != plan.alignment {
return Err(format!(
"cached stage shard {} has alignment {}; expected {}",
path.display(),
directory.alignment,
plan.alignment
));
}
for (index, (actual, expected)) in directory.tensors.iter().zip(&plan.tensors).enumerate() {
if actual.name != expected.name
|| actual.dims != expected.dims
|| actual.ggml_type != expected.ggml_type
{
return Err(format!(
"cached stage shard {} tensor {index} is {} {:?} type {}; expected {} {:?} type {}",
path.display(),
actual.name,
actual.dims,
actual.ggml_type,
expected.name,
expected.dims,
expected.ggml_type
));
}
}
Ok(())
}
pub fn source_url(source: &GgufSource) -> Result<String, String> {
match source {
GgufSource::HuggingFaceGguf {
@ -405,6 +445,7 @@ fn shard_cache_key(
tensors: &[GgufTensorEntry],
) -> String {
let mut hasher = blake3::Hasher::new();
hasher.update(format!("format:{STAGE_SHARD_CACHE_FORMAT_VERSION}\n").as_bytes());
hasher.update(format!("source:{source:?}\n").as_bytes());
hasher.update(
format!("stage:{stage_index}/{stage_count}:{layer_start}-{layer_end_exclusive}\n")
@ -1128,6 +1169,24 @@ mod tests {
.sum::<u64>();
assert_eq!(total_requested, plan.planned_fetch_bytes());
assert!(total_requested < fixture.bytes_len);
validate_stage_shard_cache(&output_path, &plan).unwrap();
}
#[test]
fn corrupted_stage_shard_cache_is_rejected_before_reuse() {
let fixture = SyntheticGguf::new(4);
let source = GgufSource::HuggingFaceGguf {
repo: "org/repo".to_owned(),
file: "model.gguf".to_owned(),
revision: None,
};
let plan = plan_stage_shard(&fixture.path, source, 1, 2, 2, 4).unwrap();
let output_path = fixture.path.with_file_name(plan.cache_file_name());
std::fs::write(&output_path, b"not a gguf").unwrap();
let error = validate_stage_shard_cache(&output_path, &plan).unwrap_err();
assert!(error.contains("invalid cached stage shard"));
}
fn names(plan: &StageShardPlan) -> Vec<&str> {

View file

@ -88,7 +88,12 @@ fn data_plane_fault_and_stop_reports_map_to_mvp_lifecycle_messages() {
.expect("send stopped");
runtime.tick();
assert_eq!(node_messages.try_recv(), Some(NodeAgentMsg::WorkerCrashed));
assert_eq!(
node_messages.try_recv(),
Some(NodeAgentMsg::WorkerCrashed {
reason: Some("data plane faulted".to_owned())
})
);
assert_eq!(
node_messages.try_recv(),
Some(NodeAgentMsg::LocalEdgesStopped { run_id: 55 })

View file

@ -673,7 +673,7 @@ fn stop_destroys_known_vastai_contract_exactly_once() {
}
#[test]
fn vastai_complete_bootstrap_stops_optional_log_tail_before_node_stop() {
fn vastai_complete_bootstrap_keeps_optional_log_tail_until_node_stop() {
let mut plugin = VastAiProvisioningPlugin::new(
FakeLeaseClient::default().with_contract(100),
FakeBootstrap::default(),
@ -683,10 +683,7 @@ fn vastai_complete_bootstrap_stops_optional_log_tail_before_node_stop() {
plugin.complete_bootstrap(&handle).unwrap();
assert_eq!(
plugin.bootstrap().stops,
vec![(1, BootstrapStopReason::RuntimeReady)]
);
assert!(plugin.bootstrap().stops.is_empty());
assert_eq!(plugin.client().destroyed, Vec::<u64>::new());
assert_eq!(plugin.active_contract_count(), 1);
@ -695,7 +692,7 @@ fn vastai_complete_bootstrap_stops_optional_log_tail_before_node_stop() {
assert_eq!(plugin.client().destroyed, vec![100]);
assert_eq!(
plugin.bootstrap().stops,
vec![(1, BootstrapStopReason::RuntimeReady)]
vec![(1, BootstrapStopReason::NodeStop)]
);
assert_eq!(plugin.active_contract_count(), 0);
}

View file

@ -4689,20 +4689,16 @@ fn require_vastai_data_path_facts(facts: &DumpLogFacts) -> Result<(), String> {
"worker ingress ring installed",
)?;
require_dump_log_fact(facts.ring_installed_egress, "worker egress ring installed")?;
let downstream_activation = facts.activation_downstream_object_loaded
&& facts.max_activation_record_bytes >= DATA_PATH_MIN_PAYLOAD_BYTES;
require_dump_log_fact(
facts.activation_object_loaded || facts.worker_ingress_object_loaded,
"worker object loaded from ingress ring",
)?;
require_dump_log_fact(
facts.activation_step_executed,
"activation-producing worker step executed",
facts.activation_step_executed || downstream_activation,
"activation-producing worker step executed or downstream activation loaded",
)?;
let explicit_transport = facts.activation_egress_ring_read
&& facts.activation_iroh_edge_sent
&& facts.activation_iroh_edge_read
&& facts.activation_ingress_ring_write;
let downstream_activation = facts.activation_downstream_object_loaded
&& facts.max_activation_record_bytes >= DATA_PATH_MIN_PAYLOAD_BYTES;
let downstream_prompt_output = facts.activation_egress_record_written
&& facts.pipeline_token_out_requests.len() >= 2
&& facts.vastai_node_runtime_ready_nodes.len() >= 2;
@ -6100,7 +6096,6 @@ mod tests {
ring_installed_ingress: true,
ring_installed_egress: true,
activation_object_loaded: true,
activation_step_executed: true,
activation_downstream_object_loaded: true,
max_activation_record_bytes: DATA_PATH_MIN_PAYLOAD_BYTES,
..DumpLogFacts::default()