refactor: actorization polish and stability
Convert the remaining thread/mpsc-based provider monitor and stage-shard fetch loops into swactor actors, propagate worker-crash reasons through the actor message chain, and validate cached stage shards before reuse.
- orchestration/provider_adapters/vastai: replace the thread+AtomicBool VastAiProviderMonitor with VastAiProviderMonitorActor driven by Poll/Stop messages on its own RuntimeHandle (schedule_provider_monitor_poll), preserving the status/terminal-failure observation logic
- node/worker_node_runtime: convert the blocking stage-shard-fetch mpsc loop into StageShardFetchActor (Start/PollChild/ProcessLine/ReaderError/ReaderClosed via ExternalSender, on_stop kills and joins the child) reporting Progress/Done/Failed
- staging/gguf_shard: add validate_stage_shard_cache (tensor count/alignment/name/dims/type) and mix STAGE_SHARD_CACHE_FORMAT_VERSION into the shard_cache_key; materialize_stage_shard_with_process now validates a cached shard, emitting StageShardCacheInvalid and refetching when stale
- node/actor + orchestration/actor: NodeAgentMsg::WorkerCrashed now carries Option<reason>, surfaced as ObserveStageFault{reason}/StageFault{reason} up to the orchestrator; NodeAgentActor records last_worker_crash
- orchestration/app: complete_bootstrap now runs after runtime-ready without dropping stop handles (ProvisionedClusterGuard) and folds worker-crash reasons into stage-fault errors; chat/runtime adds cached-model config selection
- xtask: relax the data-path requirement to accept activation-step-executed OR downstream-activation-loaded
Signed-off-by: Zachery Aaron Shores-Chmielewski <zacheryasc@gmail.com>
This commit is contained in:
parent
e323cd5e86
commit
ddec08ebf4
11 changed files with 1326 additions and 530 deletions
|
|
@ -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")
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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(),
|
||||
},
|
||||
);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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 } => {
|
||||
|
|
|
|||
|
|
@ -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 {
|
||||
stream: StageShardProcessStream,
|
||||
line: String,
|
||||
#[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 {
|
||||
stream,
|
||||
line: line.trim_end_matches(['\r', '\n']).to_owned(),
|
||||
});
|
||||
let _ = sender.send_to(
|
||||
actor,
|
||||
StageShardFetchMsg::ProcessLine {
|
||||
stream,
|
||||
line: line.trim_end_matches(['\r', '\n']).to_owned(),
|
||||
},
|
||||
);
|
||||
}
|
||||
Err(error) => {
|
||||
let _ = tx.send(StageShardProcessLine {
|
||||
stream,
|
||||
line: format!("reader error: {error}"),
|
||||
});
|
||||
let _ = sender.send_to(
|
||||
actor,
|
||||
StageShardFetchMsg::ReaderError {
|
||||
stream,
|
||||
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,14 +3480,34 @@ fn materialize_stage_shard_with_process(
|
|||
) -> Result<PathBuf, String> {
|
||||
let output_path = stage_shard_cache_path(plan);
|
||||
if output_path.is_file() {
|
||||
let event = json!({
|
||||
"type":"StageShardCacheReady",
|
||||
"stage_index":plan.stage_index,
|
||||
"path":output_path,
|
||||
"cache_hit":true,
|
||||
});
|
||||
publish_stage_shard_fetch_event(datastream, config, &event)?;
|
||||
return Ok(output_path);
|
||||
match validate_stage_shard_cache(&output_path, plan) {
|
||||
Ok(()) => {
|
||||
let event = json!({
|
||||
"type":"StageShardCacheReady",
|
||||
"stage_index":plan.stage_index,
|
||||
"path":output_path,
|
||||
"cache_hit":true,
|
||||
});
|
||||
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 {
|
||||
|
|
@ -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,
|
||||
)?;
|
||||
(
|
||||
GgufSource::LocalPath(local_path.to_string_lossy().into_owned()),
|
||||
true,
|
||||
)
|
||||
) {
|
||||
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);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
};
|
||||
|
|
|
|||
|
|
@ -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();
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load diff
|
|
@ -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> {
|
||||
|
|
|
|||
|
|
@ -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 })
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
Loading…
Reference in a new issue