diff --git a/apps/mvp-node/tinygrad_worker.py b/apps/mvp-node/tinygrad_worker.py index 921f116..a587595 100755 --- a/apps/mvp-node/tinygrad_worker.py +++ b/apps/mvp-node/tinygrad_worker.py @@ -28,6 +28,7 @@ rings: dict[int, dict[str, Any]] = {} device_objects: dict[int, dict[str, Any]] = {} next_handle = 42 HEADER_LEN = 40 +FLAG_BEGIN_SEQUENCE = 1 << 1 WORKER_GENERATION = 1 BENCHMARK_SCHEMA = 1 _benchmark_start = time.monotonic() @@ -575,9 +576,19 @@ def load_weights(cmd: dict[str, Any]) -> None: layer_end_exclusive = int(cmd.get("layer_end_exclusive", 0)) try: control(type="TinygradLlmImportStarted", model_id=model_id) - from tinygrad.llm.cli import SimpleTokenizer + try: + from tinygrad.llm.cli import SimpleTokenizer - control(type="TinygradLlmImportReady", model_id=model_id) + Transformer = None + llm_backend = "tinygrad.llm" + except ModuleNotFoundError as exc: + if exc.name != "tinygrad.llm": + raise + from tinygrad.apps.llm import SimpleTokenizer, Transformer + + llm_backend = "tinygrad.apps.llm" + + control(type="TinygradLlmImportReady", model_id=model_id, backend=llm_backend) max_context_raw = os.environ.get("MVP_MAX_CONTEXT", "512") max_context = int(max_context_raw) if max_context_raw else 512 control( @@ -589,13 +600,30 @@ def load_weights(cmd: dict[str, Any]) -> None: layer_start=layer_start, layer_end_exclusive=layer_end_exclusive, requested_device=os.environ.get("DEV"), + llm_backend=llm_backend, ) - model, kv = load_pipeline_stage_model( - path, - max_context=max_context, - layer_start=layer_start, - layer_end_exclusive=layer_end_exclusive, - ) + if Transformer is None: + model, kv = load_pipeline_stage_model( + path, + max_context=max_context, + layer_start=layer_start, + layer_end_exclusive=layer_end_exclusive, + ) + else: + TensorCls = require_tinygrad() + model, kv = Transformer.from_gguf(TensorCls(path), max_context=max_context, realize=True) + total_layers = int(kv[f"{kv['general.architecture']}.block_count"]) - int( + kv.get(f"{kv['general.architecture']}.nextn_predict_layers", 0) + ) + if layer_start != 0 or layer_end_exclusive < total_layers: + fatal( + "TinygradAppsLlmPartialStageUnsupported", + layer_start=layer_start, + layer_end_exclusive=layer_end_exclusive, + total_layers=total_layers, + ) + model.first_stage = True + model.final_stage = True control( type="PipelineStageFromGgufReady", model_id=model_id, @@ -607,10 +635,12 @@ def load_weights(cmd: dict[str, Any]) -> None: first_stage=model.first_stage, final_stage=model.final_stage, requested_device=os.environ.get("DEV"), + llm_backend=llm_backend, ) tok_src = cmd.get("tokenizer", {"EmbeddedGguf": None}) if "EmbeddedGguf" in tok_src: - if kv.get("tokenizer.ggml.pre") == "smollm": + tokenizer_pre = str(kv.get("tokenizer.ggml.pre", "")).lower() + if tokenizer_pre == "smollm": kv = dict(kv) kv["tokenizer.ggml.pre"] = "qwen2" control(type="TokenizerBuildStarted", model_id=model_id, source="EmbeddedGguf") @@ -633,6 +663,7 @@ def load_weights(cmd: dict[str, Any]) -> None: hidden_dim=int(getattr(model, "hidden_dim", 0)), max_context=int(getattr(model, "max_context", 0)), eos_token_id=int(kv.get("tokenizer.ggml.eos_token_id", 0)), + tokenizer_pre=tokenizer_pre, ) control( type="WeightsLoaded", @@ -649,6 +680,11 @@ def prompt_template_name() -> str: explicit = os.environ.get("MVP_PROMPT_TEMPLATE") if explicit is not None: return explicit.strip().lower() + tokenizer_pre = str(loaded.get("tokenizer_pre", "")).lower() + if "smollm" in tokenizer_pre or tokenizer_pre == "qwen2": + return "smollm-chat" + if "llama" in tokenizer_pre: + return "llama3-chat" model_id = str(loaded.get("model_id", "")).lower() if "smollm" in model_id: return "smollm-chat" @@ -934,20 +970,26 @@ def payload_words(payload: bytes) -> list[int]: return [] return list(struct.unpack(f"<{len(payload) // 4}I", payload)) -def object_start_pos(sequence: int, token_count: int) -> int: - if sequence == 0: +def object_start_pos(sequence: int, token_count: int, flags: int) -> int: + if flags & FLAG_BEGIN_SEQUENCE or sequence == 0: role["prompt_tokens"] = token_count + role["prompt_decode_index"] = 0 return 0 - return int(role.get("prompt_tokens", 1)) + sequence - 1 + decode_index = int(role.get("prompt_decode_index", max(0, sequence - 1))) + role["prompt_decode_index"] = decode_index + 1 + return int(role.get("prompt_tokens", token_count)) + decode_index -def materialize_object(payload: bytes, sequence: int) -> dict[str, Any]: +def materialize_object(payload: bytes, sequence: int, flags: int) -> dict[str, Any]: if not isinstance(model, PipelineStageTinygradModel): + TensorCls = require_tinygrad() + tokens = payload_words(payload) + token_count = len(tokens) return { - "kind": "words", - "words": payload_words(payload), - "payload": payload, - "start_pos": object_start_pos(sequence, max(1, len(payload) // 4)), + "kind": "tokens", + "tokens": tokens, + "tensor": TensorCls([tokens], dtype="int32").realize(), + "start_pos": object_start_pos(sequence, token_count, flags), } TensorCls = require_tinygrad() if bool(getattr(model, "first_stage", False)) and int(role.get("layer_start", 0)) == 0: @@ -957,7 +999,7 @@ def materialize_object(payload: bytes, sequence: int) -> dict[str, Any]: "kind": "tokens", "tokens": tokens, "tensor": TensorCls([tokens], dtype="int32").realize(), - "start_pos": object_start_pos(sequence, token_count), + "start_pos": object_start_pos(sequence, token_count, flags), } import numpy as np @@ -972,7 +1014,7 @@ def materialize_object(payload: bytes, sequence: int) -> dict[str, Any]: return { "kind": "activation", "tensor": TensorCls(array).realize(), - "start_pos": object_start_pos(sequence, token_count), + "start_pos": object_start_pos(sequence, token_count, flags), } @@ -986,7 +1028,7 @@ def ring_readable(cmd: dict[str, Any]) -> None: object_id, sequence, extent, flags, payload = parse_record(ring) handle = next_handle next_handle += 1 - materialized = materialize_object(payload, sequence) + materialized = materialize_object(payload, sequence, flags) materialized.update( object_id=object_id, sequence=sequence, @@ -1048,19 +1090,19 @@ def execute_step(cmd: dict[str, Any]) -> None: if ring["direction"] != "egress": fatal("WrongRingDirection", ring_id=output_ring_id, direction=ring["direction"]) final_stage = bool(cmd.get("final_stage")) + execution_backend = "pipeline_stage" if not isinstance(model, PipelineStageTinygradModel): - if final_stage: - base = sum(int(word) for word in obj["words"]) + int(role.get("stage_index", 0)) - token = 6 if base % 2 else 8 - payload = struct.pack(" None: object_id=int(cmd["output_object_id"]), sequence=int(cmd["output_sequence"]), committed_bytes=committed, + execution_backend=execution_backend, ) diff --git a/crates/mvp-system/src/bin/mvp_chat.rs b/crates/mvp-system/src/bin/mvp_chat.rs index 59abb28..0eaefde 100644 --- a/crates/mvp-system/src/bin/mvp_chat.rs +++ b/crates/mvp-system/src/bin/mvp_chat.rs @@ -36,6 +36,22 @@ const BASE_NODE_IMAGE: &str = "swactor-mvp-node-base:cuda12.6"; const REPO_MODEL_CACHE_DIR: &str = ".model-cache"; const DEFAULT_MAX_TOKENS: u32 = 64; const ORCH_SHUTDOWN_GRACE_MS: u64 = 5_000; +const MVP_CHAT_GPU_RUN_ENV: &str = "MVP_CHAT_GPU_RUN"; +const MVP_CHAT_USAGE: &str = "\ +USAGE: cargo mvp-chat [OPTIONS] + +OPTIONS: + --gpu Run the local GPU path: in-process orchestrator plus DEV=CUDA worker selection + --process | --docker | --vastai + Select the runtime provider + --config Load config overlay + --pipeline-stages Number of pipeline stages + --cached-model[=] Use discovered or explicit cached GGUF model + --dump-logs[=] Write datastream frame log + --run-id Override run id + --skip-rebuild Reuse existing Cargo artifacts + --yes, -y Approve Vast.ai lease prompts + --help, -h Print this help"; const ORCH_SHUTDOWN_POLL_MS: u64 = 50; const CHAT_LIFECYCLE_CHANNEL: &str = "mvp.chat.lifecycle"; const CHAT_RUNTIME_CHANNEL: &str = "mvp.chat.runtime"; @@ -69,11 +85,54 @@ where } } +fn print_usage() { + println!("{MVP_CHAT_USAGE}"); +} + +fn is_help_request(args: &[String]) -> bool { + args.iter() + .any(|arg| matches!(arg.as_str(), "--help" | "-h" | "help")) +} + +struct RuntimeEnvGuard { + name: &'static str, + original: Option, +} + +impl RuntimeEnvGuard { + fn apply_gpu_defaults(gpu_run: bool) -> Option { + if !gpu_run || std::env::var_os("DEV").is_some() { + return None; + } + let guard = Self { + name: "DEV", + original: None, + }; + unsafe { std::env::set_var(guard.name, "CUDA") }; + Some(guard) + } +} + +impl Drop for RuntimeEnvGuard { + fn drop(&mut self) { + match &self.original { + Some(value) => unsafe { std::env::set_var(self.name, value) }, + None => unsafe { std::env::remove_var(self.name) }, + } + } +} + fn run(args: I) -> Result<(), String> where I: IntoIterator, { - let config = Config::from_args(args)?; + let provided_args = args.into_iter().collect::>(); + if is_help_request(&provided_args) { + print_usage(); + return Ok(()); + } + let config = Config::from_args(provided_args)?; + let _gpu_env = RuntimeEnvGuard::apply_gpu_defaults(config.gpu_run); let mut progress = ChatDatastream::new(config.run_id, config.datastream_frame_log.clone())?; progress.emit( CHAT_LIFECYCLE_CHANNEL, @@ -85,6 +144,7 @@ where "max_tokens": config.max_tokens, "cached_model": config.cached_model.as_ref().map(|model| model.host_path.to_string_lossy().to_string()), "dump_logs": config.datastream_frame_log.as_ref().map(|path| path.to_string_lossy().to_string()), + "gpu_run": config.gpu_run, }), ); confirm_vastai_if_needed(&config)?; @@ -120,21 +180,30 @@ where CHAT_COMPONENT_CHANNEL, "orchestrator_process_spawn", "started", - json!({"binary": config.orch_bin.to_string_lossy()}), + json!({ + "mode": config.orchestrator_launch_mode(), + "binary": config.orch_bin.to_string_lossy(), + }), ); - let mut orch = match OrchChild::spawn(&config, &image_ref) { + let mut orch = match OrchHandle::spawn(&config, &image_ref) { Ok(orch) => { progress.emit( CHAT_COMPONENT_CHANNEL, "orchestrator_process_spawn", "ready", - json!({"binary": config.orch_bin.to_string_lossy(), "pid": orch.child.id()}), + json!({ + "mode": config.orchestrator_launch_mode(), + "binary": config.orch_bin.to_string_lossy(), + }), ); progress.emit( CHAT_COMPONENT_CHANNEL, "orchestrator_process", "started", - json!({"binary": config.orch_bin.to_string_lossy()}), + json!({ + "mode": config.orchestrator_launch_mode(), + "binary": config.orch_bin.to_string_lossy(), + }), ); orch } @@ -143,13 +212,17 @@ where CHAT_COMPONENT_CHANNEL, "orchestrator_process_spawn", "failed", - json!({"binary": config.orch_bin.to_string_lossy(), "error": error}), + json!({ + "mode": config.orchestrator_launch_mode(), + "binary": config.orch_bin.to_string_lossy(), + "error": error, + }), ); progress.emit( CHAT_COMPONENT_CHANNEL, "orchestrator_process", "failed", - json!({"error": error}), + json!({"mode": config.orchestrator_launch_mode(), "error": error}), ); progress.archive_pending()?; return Err(error); @@ -257,6 +330,7 @@ struct Config { pipeline_stages: u32, max_tokens: u32, skip_rebuild: bool, + gpu_run: bool, } struct ChatDatastream { @@ -537,8 +611,15 @@ impl Config { if max_tokens == 0 { return Err("[runtime].max_tokens must be greater than 0".to_owned()); } - let cached_model = args - .cached_model + let gpu_run = args.gpu || env_flag(MVP_CHAT_GPU_RUN_ENV, false); + let cached_model_source = match args.cached_model { + Some(source) => Some(source), + None if gpu_run && provider == ProviderKind::Process => { + Some(CachedModelSource::Discover) + } + None => None, + }; + let cached_model = cached_model_source .map(CachedModelConfig::from_source) .transpose()?; let datastream_frame_log = if args.dump_logs { @@ -576,6 +657,7 @@ impl Config { max_tokens, vastai, skip_rebuild: args.skip_rebuild, + gpu_run, }) } @@ -666,6 +748,14 @@ impl Config { } args } + + fn orchestrator_launch_mode(&self) -> &'static str { + if self.gpu_run { + "in_process_actor" + } else { + "process_binary" + } + } } #[derive(Default, Debug)] @@ -679,6 +769,8 @@ struct ParsedArgs { run_id: Option, skip_rebuild: bool, cached_model: Option, + help: bool, + gpu: bool, } #[derive(Clone, Debug, PartialEq, Eq)] @@ -707,6 +799,8 @@ impl ParsedArgs { let mut args = provided_args.into_iter().peekable(); while let Some(arg) = args.next() { match arg.as_str() { + "--help" | "-h" | "help" => parsed.help = true, + "--gpu" => parsed.gpu = true, "--vastai" => parsed.set_provider_selector(ProviderKind::VastAi)?, "--process" => parsed.set_provider_selector(ProviderKind::Process)?, "--docker" => parsed.set_provider_selector(ProviderKind::Docker)?, @@ -860,6 +954,130 @@ fn parse_approval(input: &str) -> bool { matches!(input.trim().to_ascii_lowercase().as_str(), "y" | "yes") } +enum OrchHandle { + Process(OrchChild), + InProcess(InProcessOrch), +} + +impl OrchHandle { + fn spawn(config: &Config, image_ref: &str) -> Result { + if config.gpu_run { + InProcessOrch::spawn(config, image_ref).map(Self::InProcess) + } else { + OrchChild::spawn(config, image_ref).map(Self::Process) + } + } + + fn wait_ready(&mut self, rpc_addr: String) -> Result { + match self { + Self::Process(orch) => orch.wait_ready(rpc_addr), + Self::InProcess(orch) => orch.wait_ready(rpc_addr), + } + } + + fn shutdown(&mut self) { + match self { + Self::Process(orch) => orch.shutdown(), + Self::InProcess(orch) => orch.shutdown(), + } + } +} + +struct InProcessOrch { + stop_tx: Option>, + thread: Option>>, + cleaned: bool, +} + +impl InProcessOrch { + fn spawn(config: &Config, image_ref: &str) -> Result { + let args = config.orchestrator_cli_args(image_ref); + let (stop_tx, stop_rx) = mpsc::channel(); + let thread = thread::spawn(move || { + mvp_system::orchestrator_app::run_in_process_from_args(args, stop_rx) + }); + Ok(Self { + stop_tx: Some(stop_tx), + thread: Some(thread), + cleaned: false, + }) + } + + fn wait_ready(&mut self, rpc_addr: String) -> Result { + loop { + if STOP_REQUESTED.load(Ordering::SeqCst) { + return Err("interrupted before orchestrator became ready".to_owned()); + } + match TcpStream::connect(&rpc_addr) { + Ok(stream) => { + let _ = stream.shutdown(Shutdown::Both); + return Ok(rpc_addr); + } + Err(error) + if matches!( + error.kind(), + std::io::ErrorKind::ConnectionRefused + | std::io::ErrorKind::TimedOut + | std::io::ErrorKind::AddrNotAvailable + ) => {} + Err(error) => return Err(format!("connect prompt RPC {rpc_addr}: {error}")), + } + if let Some(result) = self.take_finished_result() { + return Err(format!( + "in-process orchestrator exited before prompt RPC ready: {}", + render_orch_thread_result(result) + )); + } + thread::sleep(Duration::from_millis(100)); + } + } + + fn shutdown(&mut self) { + if self.cleaned { + return; + } + self.cleaned = true; + let _ = self.stop_tx.take().map(|tx| tx.send(())); + let grace = Duration::from_millis(ORCH_SHUTDOWN_GRACE_MS); + let poll = Duration::from_millis(ORCH_SHUTDOWN_POLL_MS); + let started = Instant::now(); + while started.elapsed() < grace { + if self.take_finished_result().is_some() { + return; + } + thread::sleep(poll); + } + } + + fn take_finished_result(&mut self) -> Option> { + if !self + .thread + .as_ref() + .is_some_and(|thread| thread.is_finished()) + { + return None; + } + let thread = self.thread.take()?; + Some(match thread.join() { + Ok(result) => result, + Err(_) => Err("in-process orchestrator thread panicked".to_owned()), + }) + } +} + +impl Drop for InProcessOrch { + fn drop(&mut self) { + self.shutdown(); + } +} + +fn render_orch_thread_result(result: Result<(), String>) -> String { + match result { + Ok(()) => "completed successfully".to_owned(), + Err(error) => error, + } +} + struct OrchChild { child: Child, cleaned: bool, @@ -1007,30 +1225,47 @@ fn prepare_runtime_with_progress( } else { "cargo_build" }; - emit_chat_progress( - &mut progress, - CHAT_RUNTIME_CHANNEL, - "ensure_orch_binary", - "started", - json!({"mode": binary_mode}), - ); - match ensure_orch_binary(config) { - Ok(()) => emit_chat_progress( + if config.gpu_run { + emit_chat_progress( + &mut progress, + CHAT_RUNTIME_CHANNEL, + "ensure_orchestrator_actor", + "started", + json!({"mode": config.orchestrator_launch_mode()}), + ); + emit_chat_progress( + &mut progress, + CHAT_RUNTIME_CHANNEL, + "ensure_orchestrator_actor", + "ready", + json!({"mode": config.orchestrator_launch_mode()}), + ); + } else { + emit_chat_progress( &mut progress, CHAT_RUNTIME_CHANNEL, "ensure_orch_binary", - "ready", + "started", json!({"mode": binary_mode}), - ), - Err(error) => { - emit_chat_progress( + ); + match ensure_orch_binary(config) { + Ok(()) => emit_chat_progress( &mut progress, CHAT_RUNTIME_CHANNEL, "ensure_orch_binary", - "failed", - json!({"mode": binary_mode, "error": error.as_str()}), - ); - return Err(error); + "ready", + json!({"mode": binary_mode}), + ), + Err(error) => { + emit_chat_progress( + &mut progress, + CHAT_RUNTIME_CHANNEL, + "ensure_orch_binary", + "failed", + json!({"mode": binary_mode, "error": error.as_str()}), + ); + return Err(error); + } } } @@ -1709,6 +1944,16 @@ fn provider_from_sources( Ok(ProviderKind::Process) } +fn env_flag(name: &str, default: bool) -> bool { + match env_optional(name) { + Some(value) => !matches!( + value.to_ascii_lowercase().as_str(), + "0" | "false" | "no" | "off" + ), + None => default, + } +} + fn env_optional(name: &str) -> Option { std::env::var(name) .ok() @@ -1775,6 +2020,8 @@ mod tests { "VASTAI_API_KEY", "MVP_PIPELINE_STAGES", "MVP_RUNTIME_CONFIG", + "MVP_CHAT_GPU_RUN", + "DEV", ]; struct TempDir { @@ -1886,6 +2133,7 @@ mod tests { pipeline_stages: 1, max_tokens: DEFAULT_MAX_TOKENS, skip_rebuild: true, + gpu_run: false, } } @@ -1988,6 +2236,7 @@ mod tests { #[test] fn parsed_args_accepts_public_flags() { let parsed = ParsedArgs::parse(strings(&[ + "--gpu", "--docker", "--yes", "--config", @@ -2008,6 +2257,34 @@ mod tests { assert_eq!(parsed.dump_log_path, Some(PathBuf::from("logs.ndjson"))); assert_eq!(parsed.cached_model, Some(CachedModelSource::Discover)); assert!(parsed.skip_rebuild); + assert!(parsed.gpu); + + let help = ParsedArgs::parse(strings(&["--help"])).expect("help parses"); + assert!(help.help); + let short_help = ParsedArgs::parse(strings(&["-h"])).expect("short help parses"); + assert!(short_help.help); + } + + #[test] + fn config_gpu_flag_selects_in_process_gpu_run() { + let temp = TempDir::new("gpu-flag-config"); + let cache = temp.path().join(REPO_MODEL_CACHE_DIR); + fs::create_dir_all(&cache).expect("create model cache"); + let cached_path = cache.join("default.gguf"); + fs::write(&cached_path, b"cached model").expect("write cached model"); + with_process_state(&[], Some(temp.path()), || { + let config = Config::from_args(strings(&["--gpu", "--skip-rebuild"])) + .expect("gpu config resolves"); + assert!(config.gpu_run); + assert_eq!(config.orchestrator_launch_mode(), "in_process_actor"); + assert_eq!( + config + .cached_model + .as_ref() + .map(|model| model.host_path.clone()), + Some(cached_path.canonicalize().expect("canonical cached model")) + ); + }); } #[test] diff --git a/crates/mvp-system/src/bin/orchestrator.rs b/crates/mvp-system/src/bin/orchestrator.rs index 9d61eef..3d6c3aa 100644 --- a/crates/mvp-system/src/bin/orchestrator.rs +++ b/crates/mvp-system/src/bin/orchestrator.rs @@ -1,94 +1,7 @@ -use std::collections::{BTreeMap, BTreeSet}; -use std::fs::{File, OpenOptions}; -use std::io::{BufRead, BufReader, Write}; -use std::net::{SocketAddr, TcpListener, TcpStream}; -#[cfg(target_os = "linux")] -use std::os::fd::FromRawFd; -use std::path::{Path, PathBuf}; -use std::process::{Command, ExitCode}; -use std::sync::{Arc, mpsc}; -use std::thread; -use std::time::{Duration, Instant}; - -use datastream::{ - ChannelContent, ChannelId, ChannelRef, DatastreamEndpoint, DatastreamEvent, DatastreamProducer, - DatastreamPublisherMsg, DatastreamSubscribe, Frame, Lifetime, NodeId, StreamDescriptor, - StreamId, StreamOrigin, SubscriptionRequest, -}; -use distribution::node::DistributedNodeConfig; -use distribution::types::{MemberState, NodeId as DistNodeId}; -use iroh::EndpointAddr; -use iroh_driver::{ - DATASTREAM_ALPN, EDGE_ALPN, EdgeSendHandle, EdgeTransportEvent, IrohDriver, IrohDriverConfig, -}; -use mvp_system::actors::node_agent::{ - NodeAgentMsg, StageEdgeKindWire, StageInboundEdgeWire, StageObjectSpecWire, - StageOutboundEdgeWire, StageProvisionWire, StageRingSpecWire, -}; -use mvp_system::actors::orchestrator::{OrchestratorActor, OrchestratorReport}; -use mvp_system::actors::register_mvp_actor_codecs; -use mvp_system::benchmark_observability; -use mvp_system::config::{DEFAULT_CONFIG_PATH, TomlConfigOverlay}; -#[cfg(feature = "dashboard")] -use mvp_system::dashboard_view::MvpClusterDashboardView; -use mvp_system::distribution_stack::DistributionRuntimeStack; -use mvp_system::gpu_worker_ingress_parser as ingress; -use mvp_system::node_provisioning::ProviderKind; -use mvp_system::orchestrator_run_fsm::{RunConfig, RunId}; -use mvp_system::prompt_rpc::{ - PromptEvent, SubmitPrompt, TokenizerEvent, read_submit_prompt, write_json_line, -}; -use mvp_system::provisioning::{ - LocalDockerPlugin, LocalProcessPlugin, NodeProvisionSpec, PluginObservation, - PluginObservationSink, PluginSink, ProviderMount, ProvisionEvent, ProvisionEventKind, - ProvisionLogLine, ProvisionLogStream, ProvisionPlugin, -}; -#[cfg(test)] -use mvp_system::relay_provisioning::relay_runtime_config_from_env; -use mvp_system::relay_provisioning::{ - MVP_IROH_RELAY_URL_ENV, RelayRuntimeConfig, SWACTOR_IROH_RELAY_URL_ENV, relay_mode_env_value, - relay_runtime_config_from_settings, -}; -use mvp_system::run_plan::{self, GgufSource, TokenizerSource}; -use mvp_system::telemetry::{ - MVP_PROVISIONING_EVENTS, MvpProvisionEventRecord, MvpProvisionLogRecord, - mvp_provision_log_channel, -}; -use mvp_system::vastai_provisioning::{ - SshCommandBootstrapLauncher, ToolsVastAiLeaseClient, VastAiProvisioningConfig, - VastAiProvisioningPlugin, -}; -use parking_lot::Mutex; -use serde_json::{Value, json}; -use swactor::actor::ActorAddress; -#[cfg(test)] -use tokio::sync::mpsc as tokio_mpsc; - -const DEFAULT_IMAGE: &str = "swactor-mvp-node:latest"; -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"; -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"; -const DEFAULT_HF_FILE: &str = "Llama-3.2-1B-Instruct-Q4_K_M.gguf"; -const DEFAULT_MODEL_ID: &str = "llama-3.2-1b-instruct-q4"; -const DEFAULT_MAX_TOKENS: u32 = 64; -const PUMP_INTERVAL: Duration = Duration::from_millis(10); -const RUNTIME_READY_ACK_RETRY_INTERVAL: Duration = Duration::from_millis(250); -const RUNTIME_READY_ACK_TIMEOUT: Duration = Duration::from_secs(60); -const MVP_ORCH_BOOTSTRAP: &str = "mvp.orch.bootstrap"; -const MVP_ORCH_PROMPT: &str = "mvp.orch.prompt"; -const MVP_SWIM_MEMBERSHIP: &str = "mvp.swim.membership"; -const MVP_STAGE_ROUTE: &str = "mvp.orch.stage_route"; -const DATASTREAM_FRAME_LOG_ENV: &str = "MVP_DATASTREAM_FRAME_LOG"; -const DEFAULT_DOCKER_CONTAINER_PREFIX: &str = "mvp-orchestrator"; -const MVP_DOCKER_CONTAINER_PREFIX_ENV: &str = "MVP_DOCKER_CONTAINER_PREFIX"; +use std::process::ExitCode; fn main() -> ExitCode { - match run() { + match mvp_system::orchestrator_app::run_from_args(std::env::args().skip(1)) { Ok(()) => ExitCode::SUCCESS, Err(error) => { eprintln!("mvp-orchestrator: {error}"); @@ -96,7226 +9,3 @@ fn main() -> ExitCode { } } } - -fn run() -> Result<(), String> { - let mut config = Config::from_defaults_toml_env_args(std::env::args().skip(1))?; - config.prepare_vastai_ssh_key()?; - let orch_stdio_rx = install_orch_stdio_capture()?; - let mut orch_datastream = - OrchDatastream::new(config.run_id, config.datastream_frame_log.as_deref())?; - orch_datastream.emit_bootstrap( - None, - config.run_id, - config.node_id, - "config", - "ready", - json!({ - "config_profile":config.config_profile.as_str(), - "image":&config.image, - "provider":config.provider.as_str(), - "rpc_bind":config.rpc_bind.to_string(), - "model_id":&config.model_id, - "stage_index":config.stage_index, - "legacy_layer_end_exclusive":config.layer_end_exclusive, - "relay_mode":format!("{:?}", config.relay.mode), - "pipeline_stages":config.pipeline_stages, - "provider_config":config.provider_datastream_detail(), - }), - ); - drain_orch_stdio_capture( - orch_stdio_rx.as_ref(), - &mut orch_datastream, - None, - config.run_id, - config.node_id, - ); - let pipeline_plan = if config.uses_planned_execution() { - let plan = config.build_run_plan()?; - orch_datastream.emit_bootstrap( - None, - config.run_id, - config.node_id, - "run_plan", - "ready", - json!({ - "stage_count":plan.stages.len(), - "edge_count":plan.edges.len(), - "model_layers":plan.model.num_layers, - "hidden_dim":plan.model.hidden_dim, - "max_seq_len":plan.model.max_seq_len, - "eos_token_id":plan.model.eos_token_id, - }), - ); - Some(plan) - } else { - None - }; - let _ = &pipeline_plan; - - let tokio = match tokio::runtime::Runtime::new() { - Ok(runtime) => { - orch_datastream.emit_bootstrap( - None, - config.run_id, - config.node_id, - "tokio_runtime", - "ready", - json!({"runtime":"tokio"}), - ); - runtime - } - Err(error) => { - orch_datastream.emit_bootstrap( - None, - config.run_id, - config.node_id, - "tokio_runtime", - "failed", - json!({"error":error.to_string()}), - ); - return Err(format!("tokio runtime: {error}")); - } - }; - let mut driver = match IrohDriver::with_handle( - tokio.handle().clone(), - IrohDriverConfig { - secret_key: None, - relay_mode: config.relay.mode.clone(), - node: DistributedNodeConfig::default(), - peer_auth: None, - additional_alpns: vec![EDGE_ALPN.to_vec(), DATASTREAM_ALPN.to_vec()], - }, - ) { - Ok(driver) => { - orch_datastream.emit_bootstrap( - None, - config.run_id, - config.node_id, - "iroh_driver", - "ready", - json!({"relay_mode":format!("{:?}", config.relay.mode)}), - ); - driver - } - Err(error) => { - orch_datastream.emit_bootstrap( - None, - config.run_id, - config.node_id, - "iroh_driver", - "failed", - json!({"error":error.to_string()}), - ); - return Err(format!("create iroh driver: {error}")); - } - }; - let stack = DistributionRuntimeStack::new_with_codecs( - driver.node_id(), - DistributedNodeConfig::default(), - |registry| { - register_mvp_actor_codecs(registry); - datastream::wire::register_datastream_codec(registry); - }, - ); - orch_datastream.emit_bootstrap( - None, - config.run_id, - config.node_id, - "distribution_stack", - "ready", - json!({"actors":"initialized","route_view":"initialized","swim":"initialized"}), - ); - orch_datastream.emit_bootstrap( - None, - config.run_id, - config.node_id, - "codecs", - "ready", - json!({"registered":["node_agent","orchestrator","provisioner","prompt_rpc","datastream"]}), - ); - driver.enable_actor_bridge( - stack.runtime.clone(), - stack.codec.clone(), - stack.actor_bridge_routes(), - stack.actors.swim, - stack.relay_mirror.clone(), - stack.route_view.clone(), - ); - orch_datastream.emit_bootstrap( - None, - config.run_id, - config.node_id, - "actor_bridge", - "ready", - json!({"transport":"iroh","routes":"attached"}), - ); - - let (frame_tx, frame_rx) = mpsc::channel::(); - orch_datastream.emit_bootstrap( - None, - config.run_id, - config.node_id, - "datastream_collector", - "ready", - json!({"alpn":String::from_utf8_lossy(DATASTREAM_ALPN)}), - ); - let dashboard = DashboardSupport::start(config.dashboard)?; - orch_datastream.emit_bootstrap( - dashboard.as_ref(), - config.run_id, - config.node_id, - "dashboard", - "ready", - json!({"enabled":dashboard.is_some()}), - ); - - let orchestrator_reports = match stack.runtime.new_inbox::() { - Ok(inbox) => inbox, - Err(error) => { - orch_datastream.emit_bootstrap( - dashboard.as_ref(), - config.run_id, - config.node_id, - "orchestrator_report_actor", - "failed", - json!({"error":error.to_string()}), - ); - return Err(format!("orchestrator report inbox: {error}")); - } - }; - let orchestrator_report_actor = *orchestrator_reports.addr(); - stack.register_local_actor(driver.register_actor(orchestrator_report_actor, 1)); - orch_datastream.emit_bootstrap( - dashboard.as_ref(), - config.run_id, - config.node_id, - "orchestrator_report_actor", - "ready", - json!({"actor":orchestrator_report_actor}), - ); - let orchestrator_actor = match stack.runtime.spawn(OrchestratorActor::new( - RunConfig { - run_id: RunId(config.run_id), - max_tokens: u64::from(config.default_max_tokens), - prompt: Vec::new(), - }, - Some(orchestrator_report_actor), - )) { - Ok(actor) => actor, - Err(error) => { - orch_datastream.emit_bootstrap( - dashboard.as_ref(), - config.run_id, - config.node_id, - "orchestrator_actor", - "failed", - json!({"error":error.to_string()}), - ); - return Err(format!("spawn orchestrator actor: {error}")); - } - }; - stack.register_local_actor(driver.register_actor(orchestrator_actor, 1)); - orch_datastream.emit_bootstrap( - dashboard.as_ref(), - config.run_id, - config.node_id, - "orchestrator_actor", - "ready", - json!({"actor":orchestrator_actor}), - ); - - let prompt_events = match stack.runtime.new_inbox::() { - Ok(inbox) => inbox, - Err(error) => { - orch_datastream.emit_bootstrap( - dashboard.as_ref(), - config.run_id, - config.node_id, - "prompt_reply_actor", - "failed", - json!({"error":error.to_string()}), - ); - return Err(format!("prompt event inbox: {error}")); - } - }; - let prompt_reply_actor = *prompt_events.addr(); - stack.register_local_actor(driver.register_actor(prompt_reply_actor, 1)); - orch_datastream.emit_bootstrap( - dashboard.as_ref(), - config.run_id, - config.node_id, - "prompt_reply_actor", - "ready", - json!({"actor":prompt_reply_actor}), - ); - - let tokenizer_events = match stack.runtime.new_inbox::() { - Ok(inbox) => inbox, - Err(error) => { - orch_datastream.emit_bootstrap( - dashboard.as_ref(), - config.run_id, - config.node_id, - "tokenizer_reply_actor", - "failed", - json!({"error":error.to_string()}), - ); - return Err(format!("tokenizer event inbox: {error}")); - } - }; - let tokenizer_reply_actor = *tokenizer_events.addr(); - stack.register_local_actor(driver.register_actor(tokenizer_reply_actor, 1)); - orch_datastream.emit_bootstrap( - dashboard.as_ref(), - config.run_id, - config.node_id, - "tokenizer_reply_actor", - "ready", - json!({"actor":tokenizer_reply_actor}), - ); - - let (work_tx, work_rx) = mpsc::channel::(); - let stop_rx = spawn_stop_listener(); - - let provisioner = config.build_provisioner(Arc::clone(&stack.runtime))?; - orch_datastream.emit_bootstrap( - dashboard.as_ref(), - config.run_id, - config.node_id, - "node_provisioner", - "ready", - json!({ - "provider":config.provider.as_str(), - "owner":"mvp-orchestrator", - "config":config.provider_datastream_detail(), - }), - ); - let (obs_tx, obs_rx) = mpsc::channel::(); - let sink = PluginSink::new(Arc::new(ChannelObservationSink { - tx: Mutex::new(obs_tx), - })); - let coordinator_endpoint = driver.endpoint_addr(); - let pipeline_coordinator_endpoint = coordinator_endpoint.clone(); - let (mut provisioned_nodes, ready) = start_and_provision_workers( - provisioner, - &config, - pipeline_plan.as_ref(), - &mut driver, - &stack, - &obs_rx, - &frame_rx, - &frame_tx, - &orchestrator_reports, - &stop_rx, - dashboard.as_ref(), - &mut orch_datastream, - orch_stdio_rx.as_ref(), - sink, - coordinator_endpoint, - pipeline_coordinator_endpoint, - orchestrator_actor, - )?; - - orch_datastream.emit_bootstrap( - dashboard.as_ref(), - config.run_id, - config.node_id, - "prompt_rpc", - "started", - json!({ - "bind":config.rpc_bind.to_string(), - "default_max_tokens":config.default_max_tokens, - }), - ); - - let rpc_addr = match spawn_prompt_rpc(config.rpc_bind, work_tx, config.default_max_tokens) { - Ok(addr) => { - orch_datastream.emit_bootstrap( - dashboard.as_ref(), - config.run_id, - config.node_id, - "prompt_rpc", - "ready", - json!({ - "addr":addr.to_string(), - "default_max_tokens":config.default_max_tokens, - }), - ); - addr - } - Err(error) => { - orch_datastream.emit_bootstrap( - dashboard.as_ref(), - config.run_id, - config.node_id, - "prompt_rpc", - "failed", - json!({"error":error}), - ); - return Err(error); - } - }; - - orch_datastream.emit_bootstrap( - dashboard.as_ref(), - config.run_id, - config.node_id, - "prompt_loop", - "ready", - json!({"addr":rpc_addr.to_string(),"node_actor":ready.first_stage.node_actor}), - ); - - orch_datastream.emit_bootstrap( - dashboard.as_ref(), - config.run_id, - config.node_id, - "serve_prompts", - "started", - json!({"mode":"single_active_prompt","poll_interval_ms":PUMP_INTERVAL.as_millis()}), - ); - let result = serve_prompts( - &mut driver, - &stack, - &obs_rx, - &frame_rx, - &frame_tx, - &work_rx, - &prompt_events, - &stop_rx, - dashboard.as_ref(), - &mut orch_datastream, - orch_stdio_rx.as_ref(), - config.run_id, - config.node_id, - ready.first_stage.node_actor, - prompt_reply_actor, - &tokenizer_events, - ready.first_stage.node_actor, - ready.final_stage.node_actor, - tokenizer_reply_actor, - config.provider, - pipeline_plan.as_ref(), - ready.first_stage.endpoint.clone(), - ); - if let Err(error) = &result { - orch_datastream.emit_bootstrap( - dashboard.as_ref(), - config.run_id, - config.node_id, - "serve_prompts", - "failed", - json!({"error":error}), - ); - } - orch_datastream.emit_bootstrap( - dashboard.as_ref(), - config.run_id, - config.node_id, - "provider_stop", - "started", - json!({"provider":config.provider.as_str(),"node_id":config.node_id}), - ); - let stop_result = provisioned_nodes.stop(); - match &stop_result { - Ok(()) => orch_datastream.emit_bootstrap( - dashboard.as_ref(), - config.run_id, - config.node_id, - "provider_stop", - "ready", - json!({"provider":config.provider.as_str(),"node_id":config.node_id}), - ), - Err(error) => orch_datastream.emit_bootstrap( - dashboard.as_ref(), - config.run_id, - config.node_id, - "provider_stop", - "failed", - json!({"provider":config.provider.as_str(),"node_id":config.node_id,"error":error}), - ), - } - if result.is_ok() && stop_result.is_ok() { - orch_datastream.emit_bootstrap( - dashboard.as_ref(), - config.run_id, - config.node_id, - "orch_exit", - "ready", - json!({"result":"ok"}), - ); - } - result.and(stop_result) -} - -#[derive(Clone)] -struct VastAiRuntimeConfig { - api_key: Option, - provisioning: VastAiProvisioningConfig, - bootstrap_command: Option, - ssh_identity: Option, - ssh_public_key: Option, - ssh_public_fingerprint: Option, -} - -impl VastAiRuntimeConfig { - fn from_builder(builder: &ConfigBuilder) -> Result { - let mut provisioning = VastAiProvisioningConfig::default(); - let disk_gb = builder - .vastai_disk_gb_raw - .as_ref() - .map(|value| ConfigBuilder::parse_value("MVP_VASTAI_DISK_GB", value)) - .transpose()? - .or(builder.vastai_disk_gb); - if let Some(disk_gb) = disk_gb { - provisioning.disk_gb = disk_gb; - } - if let Some(ssh_user) = &builder.vastai_ssh_user { - provisioning.ssh_user = ssh_user.clone(); - } - let confirm_lease = builder - .vastai_confirm_lease_raw - .as_ref() - .map(|value| ConfigBuilder::parse_bool("MVP_VASTAI_CONFIRM_LEASE", value)) - .transpose()? - .or(builder.vastai_confirm_lease); - if let Some(confirm_lease) = confirm_lease { - provisioning.confirm_lease = confirm_lease; - } - provisioning.onstart = builder.vastai_onstart.clone(); - provisioning.selection.gpu_name = builder.vastai_gpu_name.clone(); - let min_gpu_ram_mb = builder - .vastai_min_gpu_ram_mb_raw - .as_ref() - .map(|value| ConfigBuilder::parse_value("MVP_VASTAI_MIN_GPU_RAM_MB", value)) - .transpose()? - .or(builder.vastai_min_gpu_ram_mb); - if let Some(min_gpu_ram_mb) = min_gpu_ram_mb { - provisioning.selection.min_gpu_ram_mb = Some(min_gpu_ram_mb); - } - let min_down_mbps = builder - .vastai_min_down_mbps_raw - .as_ref() - .map(|value| ConfigBuilder::parse_value("MVP_VASTAI_MIN_DOWN_MBPS", value)) - .transpose()? - .or(builder.vastai_min_down_mbps); - if let Some(min_down_mbps) = min_down_mbps { - provisioning.selection.min_down_mbps = min_down_mbps; - } - let min_up_mbps = builder - .vastai_min_up_mbps_raw - .as_ref() - .map(|value| ConfigBuilder::parse_value("MVP_VASTAI_MIN_UP_MBPS", value)) - .transpose()? - .or(builder.vastai_min_up_mbps); - if let Some(min_up_mbps) = min_up_mbps { - provisioning.selection.min_up_mbps = Some(min_up_mbps); - } - let min_reliability = builder - .vastai_min_reliability_raw - .as_ref() - .map(|value| ConfigBuilder::parse_value("MVP_VASTAI_MIN_RELIABILITY", value)) - .transpose()? - .or(builder.vastai_min_reliability); - if let Some(min_reliability) = min_reliability { - provisioning.selection.min_reliability = min_reliability; - } - let require_verified = builder - .vastai_require_verified_raw - .as_ref() - .map(|value| ConfigBuilder::parse_bool("MVP_VASTAI_REQUIRE_VERIFIED", value)) - .transpose()? - .or(builder.vastai_require_verified); - if let Some(require_verified) = require_verified { - provisioning.selection.require_verified = require_verified; - } - let poll_interval_secs = builder - .vastai_poll_interval_secs_raw - .as_ref() - .map(|value| ConfigBuilder::parse_value("MVP_VASTAI_POLL_INTERVAL_SECS", value)) - .transpose()? - .or(builder.vastai_poll_interval_secs); - if let Some(poll_interval_secs) = poll_interval_secs { - provisioning.lifecycle.poll_interval = Duration::from_secs(poll_interval_secs); - } - let ssh_identity = builder - .vastai_ssh_identity_raw - .as_ref() - .map(|value| expand_home_path(value)) - .transpose()?; - Ok(Self { - api_key: builder.vastai_api_key.clone(), - provisioning, - bootstrap_command: builder.vastai_bootstrap_command.clone(), - ssh_identity, - ssh_public_key: None, - ssh_public_fingerprint: None, - }) - } - - fn datastream_detail(&self) -> Value { - json!({ - "disk_gb": self.provisioning.disk_gb, - "ssh_user": &self.provisioning.ssh_user, - "gpu_name": &self.provisioning.selection.gpu_name, - "min_gpu_ram_mb": self.provisioning.selection.min_gpu_ram_mb, - "min_down_mbps": self.provisioning.selection.min_down_mbps, - "min_up_mbps": self.provisioning.selection.min_up_mbps, - "min_reliability": self.provisioning.selection.min_reliability, - "require_verified": self.provisioning.selection.require_verified, - "confirm_lease": self.provisioning.confirm_lease, - "has_api_key": self.api_key.is_some(), - "has_onstart": self.provisioning.onstart.is_some(), - "has_bootstrap_command": self.bootstrap_command.is_some(), - "has_ssh_identity": self.ssh_identity.is_some(), - "ssh_public_fingerprint": self.ssh_public_fingerprint.as_deref(), - }) - } -} - -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -enum RuntimeConfigProfile { - Local, - Deploy, -} - -impl RuntimeConfigProfile { - fn parse(value: &str) -> Result { - match value.trim().to_ascii_lowercase().as_str() { - "local" => Ok(Self::Local), - "deploy" => Ok(Self::Deploy), - other => Err(format!( - "unsupported {MVP_RUNTIME_CONFIG_ENV}={other:?}; use local or deploy" - )), - } - } - - fn as_str(self) -> &'static str { - match self { - Self::Local => "local", - Self::Deploy => "deploy", - } - } - - fn default_provider(self) -> ProviderKind { - match self { - Self::Local => ProviderKind::Process, - Self::Deploy => ProviderKind::VastAi, - } - } -} - -#[derive(Clone)] -struct CachedModelConfig { - host_path: PathBuf, - container_path: String, -} - -impl CachedModelConfig { - fn from_host_path(provider: ProviderKind, requested: PathBuf) -> Result { - if !matches!(provider, ProviderKind::Process | ProviderKind::Docker) { - return Err(format!( - "{CACHED_MODEL_HOST_ENV} is a host-local cache path and is only supported by provider=process or provider=docker" - )); - } - let host_path = requested.canonicalize().map_err(|e| { - format!( - "resolve {CACHED_MODEL_HOST_ENV} path {}: {e}", - requested.display() - ) - })?; - if !host_path.is_file() { - return Err(format!( - "{CACHED_MODEL_HOST_ENV} must point at a file: {}", - host_path.display() - )); - } - let container_path = cached_model_container_path(&host_path)?; - Ok(Self { - host_path, - container_path, - }) - } - - fn worker_path(&self, provider: ProviderKind) -> String { - match provider { - ProviderKind::Process => self.host_path.to_string_lossy().to_string(), - ProviderKind::Docker | ProviderKind::VastAi | ProviderKind::Mock => { - self.container_path.clone() - } - } - } - - fn datastream_detail(&self) -> Value { - json!({ - "host_path_present": true, - "file": self.host_path.file_name().and_then(|name| name.to_str()), - "container_path": &self.container_path, - }) - } -} - -fn cached_model_container_path(host_path: &Path) -> Result { - let file_name = host_path - .file_name() - .and_then(|name| name.to_str()) - .filter(|name| !name.is_empty()) - .ok_or_else(|| { - format!( - "cached model path has no file name: {}", - host_path.display() - ) - })?; - Ok(format!("{CACHED_MODEL_CONTAINER_DIR}/{file_name}")) -} - -fn default_worker_bin() -> Result { - let mut path = std::env::current_exe().map_err(|e| format!("current exe: {e}"))?; - path.set_file_name("mvp-worker-node"); - Ok(path) -} - -fn default_pipeline_cached_model_path() -> PathBuf { - let relative = PathBuf::from(".") - .join(DEFAULT_PIPELINE_MODEL_CACHE_DIR) - .join(DEFAULT_PIPELINE_CACHED_MODEL_FILE); - if relative.is_file() { - return relative; - } - PathBuf::from(env!("CARGO_MANIFEST_DIR")) - .join("..") - .join("..") - .join(DEFAULT_PIPELINE_MODEL_CACHE_DIR) - .join(DEFAULT_PIPELINE_CACHED_MODEL_FILE) -} - -fn gguf_source_is_default_hf(source: &GgufSource) -> bool { - matches!( - source, - GgufSource::HuggingFaceGguf { - repo, - file, - revision: None, - } if repo == DEFAULT_HF_REPO && file == DEFAULT_HF_FILE - ) -} - -fn gguf_source_matches_default_pipeline_cache(source: &GgufSource) -> bool { - matches!( - source, - GgufSource::HuggingFaceGguf { - file, - revision: None, - .. - } if file == DEFAULT_PIPELINE_CACHED_MODEL_FILE - ) -} - -#[derive(Clone)] -struct Config { - config_profile: RuntimeConfigProfile, - image: String, - docker_gpus: String, - provider: ProviderKind, - rpc_bind: SocketAddr, - run_id: u64, - node_id: u64, - stage_index: u32, - layer_end_exclusive: Option, - pipeline_stages: u32, - model_id: String, - gguf_source: GgufSource, - tokenizer: TokenizerSource, - default_max_tokens: u32, - dashboard: bool, - max_context: Option, - relay: RelayRuntimeConfig, - vastai: Option, - cached_model: Option, - datastream_frame_log: Option, - worker_bin: Option, -} - -#[derive(Clone)] -struct ConfigBuilder { - config_profile: RuntimeConfigProfile, - provider: Option, - image: String, - toml_vastai_image: Option, - image_overridden_after_toml: bool, - docker_gpus: String, - rpc_bind: String, - rpc_bind_label: &'static str, - run_id: u64, - node_id: u64, - stage_index: u32, - layer_end_exclusive: Option, - pipeline_stages: u32, - model_id: String, - gguf_source: GgufSource, - tokenizer: TokenizerSource, - default_max_tokens: u32, - dashboard: bool, - max_context: Option, - relay_mode: Option, - relay_url: Option, - vastai_api_key: Option, - vastai_bootstrap_command: Option, - vastai_disk_gb: Option, - vastai_disk_gb_raw: Option, - vastai_ssh_user: Option, - vastai_confirm_lease: Option, - vastai_confirm_lease_raw: Option, - vastai_onstart: Option, - vastai_ssh_identity_raw: Option, - vastai_gpu_name: Option, - vastai_min_gpu_ram_mb: Option, - vastai_min_gpu_ram_mb_raw: Option, - vastai_min_down_mbps: Option, - vastai_min_down_mbps_raw: Option, - vastai_min_up_mbps: Option, - vastai_min_up_mbps_raw: Option, - vastai_min_reliability: Option, - vastai_min_reliability_raw: Option, - vastai_require_verified: Option, - vastai_require_verified_raw: Option, - vastai_poll_interval_secs: Option, - vastai_poll_interval_secs_raw: Option, - cached_model_host_path: Option, - datastream_frame_log: Option, - worker_bin: Option, -} - -impl ConfigBuilder { - fn hardcoded_defaults() -> Self { - Self { - config_profile: RuntimeConfigProfile::Local, - provider: None, - image: DEFAULT_IMAGE.to_owned(), - toml_vastai_image: None, - image_overridden_after_toml: false, - docker_gpus: "all".to_owned(), - rpc_bind: DEFAULT_RPC_BIND.to_owned(), - rpc_bind_label: "MVP_PROMPT_RPC_BIND", - run_id: 1, - node_id: 1, - stage_index: 0, - layer_end_exclusive: None, - pipeline_stages: 1, - model_id: DEFAULT_MODEL_ID.to_owned(), - gguf_source: GgufSource::HuggingFaceGguf { - repo: DEFAULT_HF_REPO.to_owned(), - file: DEFAULT_HF_FILE.to_owned(), - revision: None, - }, - tokenizer: TokenizerSource::EmbeddedGguf, - default_max_tokens: DEFAULT_MAX_TOKENS, - dashboard: false, - max_context: None, - relay_mode: None, - relay_url: None, - vastai_api_key: None, - vastai_bootstrap_command: None, - vastai_disk_gb: None, - vastai_disk_gb_raw: None, - vastai_ssh_user: None, - vastai_confirm_lease: None, - vastai_confirm_lease_raw: None, - vastai_onstart: None, - vastai_ssh_identity_raw: None, - vastai_gpu_name: None, - vastai_min_gpu_ram_mb: None, - vastai_min_gpu_ram_mb_raw: None, - vastai_min_down_mbps: None, - vastai_min_down_mbps_raw: None, - vastai_min_up_mbps: None, - vastai_min_up_mbps_raw: None, - vastai_min_reliability: None, - vastai_min_reliability_raw: None, - vastai_require_verified: None, - vastai_require_verified_raw: None, - vastai_poll_interval_secs: None, - vastai_poll_interval_secs_raw: None, - cached_model_host_path: None, - datastream_frame_log: None, - worker_bin: None, - } - } - - fn overlay_toml(mut self, overlay: TomlConfigOverlay) -> Result { - if let Some(profile) = overlay.runtime.profile { - self.config_profile = RuntimeConfigProfile::parse(&profile)?; - } - if let Some(run_id) = overlay.runtime.run_id { - self.run_id = run_id; - } - if let Some(node_id) = overlay.runtime.node_id { - self.node_id = node_id; - } - if let Some(stage_index) = overlay.runtime.stage_index { - self.stage_index = stage_index; - } - if let Some(layer_end_exclusive) = overlay.runtime.layer_end_exclusive { - self.layer_end_exclusive = Some(layer_end_exclusive); - } - if let Some(pipeline_stages) = overlay.runtime.pipeline_stages { - self.pipeline_stages = pipeline_stages; - } - if let Some(provider) = overlay.provider.kind { - self.provider = Some(ProviderKind::parse_deploy(&provider)?); - } - if let Some(image) = overlay.image.node { - self.image = image; - } - if let Some(mode) = overlay.relay.mode { - self.relay_mode = Some(mode); - } - if let Some(url) = overlay.relay.url { - self.relay_url = Some(url); - } - if let Some(rpc_bind) = overlay.prompt.rpc_addr { - self.rpc_bind = rpc_bind; - self.rpc_bind_label = "[prompt].rpc_addr"; - } - if let Some(max_tokens) = overlay.prompt.max_tokens { - self.default_max_tokens = max_tokens; - } - if let Some(dashboard) = overlay.prompt.dashboard { - self.dashboard = dashboard; - } - if let Some(model_id) = overlay.model.id { - self.model_id = model_id; - } - if let Some(path) = overlay.model.gguf_local_path { - self.gguf_source = GgufSource::LocalPath(path); - } - if let Some(repo) = overlay.model.gguf_repo { - self.set_gguf_repo(repo); - } - if let Some(file) = overlay.model.gguf_file { - self.set_gguf_file(file); - } - if let Some(revision) = overlay.model.gguf_revision { - self.set_gguf_revision(Some(revision)); - } - if let Some(path) = overlay.model.tokenizer_local_path { - self.tokenizer = TokenizerSource::LocalPath(path); - } - if let Some(max_context) = overlay.model.max_context { - self.max_context = Some(max_context); - } - if let Some(gpus) = overlay.docker.gpus { - self.docker_gpus = gpus; - } - if let Some(path) = overlay.docker.cached_model_host_path { - self.cached_model_host_path = Some(PathBuf::from(path)); - } - if let Some(path) = overlay.observability.datastream_frame_log { - self.datastream_frame_log = Some(PathBuf::from(path)); - } - if let Some(image) = overlay.vastai.image { - self.toml_vastai_image = Some(image); - } - if let Some(api_key) = overlay.vastai.api_key { - self.vastai_api_key = Some(api_key); - } - if let Some(command) = overlay.vastai.bootstrap_command { - self.vastai_bootstrap_command = Some(command); - } - if let Some(disk_gb) = overlay.vastai.disk_gb { - self.vastai_disk_gb = Some(disk_gb); - } - if let Some(ssh_user) = overlay.vastai.ssh_user { - self.vastai_ssh_user = Some(ssh_user); - } - if let Some(confirm_lease) = overlay.vastai.confirm_lease { - self.vastai_confirm_lease = Some(confirm_lease); - } - if let Some(onstart) = overlay.vastai.onstart { - self.vastai_onstart = Some(onstart); - } - if let Some(identity) = overlay.vastai.ssh_identity { - self.vastai_ssh_identity_raw = Some(identity); - } - if let Some(gpu_name) = overlay.vastai.gpu_name { - self.vastai_gpu_name = Some(gpu_name); - } - if let Some(min_gpu_ram_mb) = overlay.vastai.min_gpu_ram_mb { - self.vastai_min_gpu_ram_mb = Some(min_gpu_ram_mb); - } - if let Some(min_down_mbps) = overlay.vastai.min_down_mbps { - self.vastai_min_down_mbps = Some(min_down_mbps); - } - if let Some(min_up_mbps) = overlay.vastai.min_up_mbps { - self.vastai_min_up_mbps = Some(min_up_mbps); - } - if let Some(min_reliability) = overlay.vastai.min_reliability { - self.vastai_min_reliability = Some(min_reliability); - } - if let Some(require_verified) = overlay.vastai.require_verified { - self.vastai_require_verified = Some(require_verified); - } - if let Some(poll_interval_secs) = overlay.vastai.poll_interval_secs { - self.vastai_poll_interval_secs = Some(poll_interval_secs); - } - Ok(self) - } - - fn overlay_env(mut self) -> Result { - if let Some(profile) = env_optional(MVP_RUNTIME_CONFIG_ENV) { - self.config_profile = RuntimeConfigProfile::parse(&profile)?; - } - if let Some(run_id) = env_optional("MVP_RUN_ID") { - self.run_id = Self::parse_value("MVP_RUN_ID", &run_id)?; - } - if let Some(node_id) = env_optional("MVP_LOGICAL_NODE_ID") { - self.node_id = Self::parse_value("MVP_LOGICAL_NODE_ID", &node_id)?; - } - if let Some(stage_index) = env_optional("MVP_STAGE_INDEX") { - self.stage_index = Self::parse_value("MVP_STAGE_INDEX", &stage_index)?; - } - if let Some(layer_end_exclusive) = env_optional("MVP_LAYER_END_EXCLUSIVE") { - self.layer_end_exclusive = Some(Self::parse_value( - "MVP_LAYER_END_EXCLUSIVE", - &layer_end_exclusive, - )?); - } - if let Some(pipeline_stages) = env_optional("MVP_PIPELINE_STAGES") { - self.pipeline_stages = Self::parse_value("MVP_PIPELINE_STAGES", &pipeline_stages)?; - } - 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") { - self.set_process_image(image); - } - if let Some(gpus) = env_optional("MVP_DOCKER_GPUS") { - self.docker_gpus = gpus; - } - if let Some(path) = env_optional(CACHED_MODEL_HOST_ENV) { - self.cached_model_host_path = Some(PathBuf::from(path)); - } - if let Some(path) = env_optional(MVP_WORKER_BIN_ENV) { - self.worker_bin = Some(PathBuf::from(path)); - } - if let Some(rpc_bind) = env_optional("MVP_PROMPT_RPC_BIND") { - self.rpc_bind = rpc_bind; - self.rpc_bind_label = "MVP_PROMPT_RPC_BIND"; - } - if let Some(max_tokens) = env_optional("MVP_PROMPT_MAX_TOKENS") { - self.default_max_tokens = Self::parse_value("MVP_PROMPT_MAX_TOKENS", &max_tokens)?; - } - if let Some(dashboard) = env_optional("MVP_DASHBOARD") { - self.dashboard = Self::parse_bool("MVP_DASHBOARD", &dashboard)?; - } - if let Some(path) = env_optional(DATASTREAM_FRAME_LOG_ENV) { - self.datastream_frame_log = Some(PathBuf::from(path)); - } - if let Some(model_id) = env_optional("MVP_MODEL_ID") { - self.model_id = model_id; - } - if let Some(path) = env_optional("MVP_GGUF_LOCAL_PATH") { - self.gguf_source = GgufSource::LocalPath(path); - } - if let Some(repo) = env_optional("MVP_GGUF_REPO") { - self.set_gguf_repo(repo); - } - if let Some(file) = env_optional("MVP_GGUF_FILE") { - self.set_gguf_file(file); - } - if let Some(revision) = env_optional("MVP_GGUF_REVISION") { - self.set_gguf_revision(Some(revision)); - } - if let Some(path) = env_optional("MVP_TOKENIZER_LOCAL_PATH") { - self.tokenizer = TokenizerSource::LocalPath(path); - } - if let Some(max_context) = env_optional("MVP_MAX_CONTEXT") { - self.max_context = Some(Self::parse_value("MVP_MAX_CONTEXT", &max_context)?); - } - if let Some(mode) = env_optional("MVP_IROH_RELAY_MODE") { - self.relay_mode = Some(mode.to_ascii_lowercase()); - } - if let Some(url) = env_optional(MVP_IROH_RELAY_URL_ENV) - .or_else(|| env_optional(SWACTOR_IROH_RELAY_URL_ENV)) - { - self.relay_url = Some(url); - } - 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") { - self.vastai_bootstrap_command = Some(command); - } - if let Some(identity) = env_optional("MVP_VASTAI_SSH_IDENTITY") { - self.vastai_ssh_identity_raw = Some(identity); - } - if let Some(disk_gb) = env_optional("MVP_VASTAI_DISK_GB") { - self.vastai_disk_gb_raw = Some(disk_gb); - } - if let Some(ssh_user) = env_optional("MVP_VASTAI_SSH_USER") { - self.vastai_ssh_user = Some(ssh_user); - } - if let Some(confirm_lease) = env_optional("MVP_VASTAI_CONFIRM_LEASE") { - self.vastai_confirm_lease_raw = Some(confirm_lease); - } - if let Some(onstart) = env_optional("MVP_VASTAI_ONSTART") { - self.vastai_onstart = Some(onstart); - } - if let Some(gpu_name) = env_optional("MVP_VASTAI_GPU_NAME") { - self.vastai_gpu_name = Some(gpu_name); - } - if let Some(min_gpu_ram_mb) = env_optional("MVP_VASTAI_MIN_GPU_RAM_MB") { - self.vastai_min_gpu_ram_mb_raw = Some(min_gpu_ram_mb); - } - if let Some(min_down_mbps) = env_optional("MVP_VASTAI_MIN_DOWN_MBPS") { - self.vastai_min_down_mbps_raw = Some(min_down_mbps); - } - if let Some(min_up_mbps) = env_optional("MVP_VASTAI_MIN_UP_MBPS") { - self.vastai_min_up_mbps_raw = Some(min_up_mbps); - } - if let Some(min_reliability) = env_optional("MVP_VASTAI_MIN_RELIABILITY") { - self.vastai_min_reliability_raw = Some(min_reliability); - } - if let Some(require_verified) = env_optional("MVP_VASTAI_REQUIRE_VERIFIED") { - self.vastai_require_verified_raw = Some(require_verified); - } - if let Some(poll_interval_secs) = env_optional("MVP_VASTAI_POLL_INTERVAL_SECS") { - self.vastai_poll_interval_secs_raw = Some(poll_interval_secs); - } - Ok(self) - } - - fn overlay_cli(mut self, args: impl IntoIterator) -> Result { - let mut args = args.into_iter(); - while let Some(arg) = args.next() { - match arg.as_str() { - "--runtime-config" => { - self.config_profile = - RuntimeConfigProfile::parse(&next_arg(&mut args, "--runtime-config")?)? - } - "--provider" => { - self.provider = Some(ProviderKind::parse_deploy(&next_arg( - &mut args, - "--provider", - )?)?) - } - "--worker-bin" => { - self.worker_bin = Some(PathBuf::from(next_arg(&mut args, "--worker-bin")?)) - } - "--image" => self.set_process_image(next_arg(&mut args, "--image")?), - "--gpus" => self.docker_gpus = next_arg(&mut args, "--gpus")?, - "--rpc-bind" => { - self.rpc_bind = next_arg(&mut args, "--rpc-bind")?; - self.rpc_bind_label = "--rpc-bind"; - } - "--run-id" => self.run_id = parse_next(&mut args, "--run-id")?, - "--node-id" => self.node_id = parse_next(&mut args, "--node-id")?, - "--stage-index" => self.stage_index = parse_next(&mut args, "--stage-index")?, - "--layer-end-exclusive" => { - self.layer_end_exclusive = Some(parse_next(&mut args, "--layer-end-exclusive")?) - } - "-N" | "--pipeline-stages" => { - self.pipeline_stages = parse_next(&mut args, arg.as_str())? - } - "--max-tokens" => self.default_max_tokens = parse_next(&mut args, "--max-tokens")?, - "--dashboard" => self.dashboard = true, - "--no-dashboard" => self.dashboard = false, - "--datastream-frame-log" => { - self.datastream_frame_log = Some(PathBuf::from(next_arg( - &mut args, - "--datastream-frame-log", - )?)); - } - "--model-id" => self.model_id = next_arg(&mut args, "--model-id")?, - "--gguf-local-path" => { - self.gguf_source = - GgufSource::LocalPath(next_arg(&mut args, "--gguf-local-path")?) - } - "--gguf-repo" => self.set_gguf_repo(next_arg(&mut args, "--gguf-repo")?), - "--gguf-file" => self.set_gguf_file(next_arg(&mut args, "--gguf-file")?), - "--gguf-revision" => { - self.set_gguf_revision(Some(next_arg(&mut args, "--gguf-revision")?)) - } - "--tokenizer-local-path" => { - self.tokenizer = - TokenizerSource::LocalPath(next_arg(&mut args, "--tokenizer-local-path")?) - } - "--max-context" => self.max_context = Some(parse_next(&mut args, "--max-context")?), - "--cached-model-host-path" => { - self.cached_model_host_path = Some(PathBuf::from(next_arg( - &mut args, - "--cached-model-host-path", - )?)); - } - "--relay-mode" => self.relay_mode = Some(next_arg(&mut args, "--relay-mode")?), - "--relay-url" => self.relay_url = Some(next_arg(&mut args, "--relay-url")?), - "--vastai-api-key" => { - self.vastai_api_key = Some(next_arg(&mut args, "--vastai-api-key")?) - } - "--vastai-bootstrap-command" => { - self.vastai_bootstrap_command = - Some(next_arg(&mut args, "--vastai-bootstrap-command")?) - } - "--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")?); - self.vastai_disk_gb_raw = None; - } - "--vastai-ssh-user" => { - self.vastai_ssh_user = Some(next_arg(&mut args, "--vastai-ssh-user")?) - } - "--vastai-confirm-lease" => { - self.vastai_confirm_lease = Some(true); - self.vastai_confirm_lease_raw = None; - } - "--no-vastai-confirm-lease" => { - self.vastai_confirm_lease = Some(false); - self.vastai_confirm_lease_raw = None; - } - "--vastai-onstart" => { - self.vastai_onstart = Some(next_arg(&mut args, "--vastai-onstart")?) - } - "--vastai-gpu-name" => { - self.vastai_gpu_name = Some(next_arg(&mut args, "--vastai-gpu-name")?) - } - "--vastai-min-gpu-ram-mb" => { - self.vastai_min_gpu_ram_mb = - Some(parse_next(&mut args, "--vastai-min-gpu-ram-mb")?); - self.vastai_min_gpu_ram_mb_raw = None; - } - "--vastai-min-down-mbps" => { - self.vastai_min_down_mbps = - Some(parse_next(&mut args, "--vastai-min-down-mbps")?); - 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_raw = None; - } - "--vastai-min-reliability" => { - self.vastai_min_reliability = - Some(parse_next(&mut args, "--vastai-min-reliability")?); - self.vastai_min_reliability_raw = None; - } - "--vastai-require-verified" => { - self.vastai_require_verified = Some(true); - self.vastai_require_verified_raw = None; - } - "--no-vastai-require-verified" => { - self.vastai_require_verified = Some(false); - self.vastai_require_verified_raw = None; - } - "--vastai-poll-interval-secs" => { - self.vastai_poll_interval_secs = - Some(parse_next(&mut args, "--vastai-poll-interval-secs")?); - self.vastai_poll_interval_secs_raw = None; - } - other => return Err(format!("unknown argument {other:?}")), - } - } - Ok(self) - } - - fn finalize(self) -> Result { - let provider = self - .provider - .unwrap_or_else(|| self.config_profile.default_provider()); - let mut image = self.image.clone(); - if provider == ProviderKind::VastAi && !self.image_overridden_after_toml { - if let Some(vastai_image) = &self.toml_vastai_image { - image = vastai_image.clone(); - } - } - if self.pipeline_stages == 0 { - return Err("--pipeline-stages must be greater than 0".to_owned()); - } - if provider == ProviderKind::VastAi && self.pipeline_stages > 1 { - return Err( - "pipeline stages greater than 1 are only supported with provider=docker; provider=vastai does not support pipelined provisioning yet".to_owned(), - ); - } - let mut cached_model_host_path = self.cached_model_host_path.clone(); - if matches!(provider, ProviderKind::Process | ProviderKind::Docker) - && self.pipeline_stages > 1 - && cached_model_host_path.is_none() - && (gguf_source_is_default_hf(&self.gguf_source) - || gguf_source_matches_default_pipeline_cache(&self.gguf_source)) - { - cached_model_host_path = Some(default_pipeline_cached_model_path()); - } - let cached_model = cached_model_host_path - .map(|path| CachedModelConfig::from_host_path(provider, path)) - .transpose()?; - let mut gguf_source = self.gguf_source.clone(); - if let Some(cached_model) = &cached_model { - gguf_source = GgufSource::LocalPath(cached_model.worker_path(provider)); - } - let relay = relay_runtime_config_from_settings( - self.run_id, - self.relay_mode.as_deref(), - self.relay_url.as_deref(), - )?; - let vastai = if provider == ProviderKind::VastAi { - Some(VastAiRuntimeConfig::from_builder(&self)?) - } else { - None - }; - Ok(Config { - config_profile: self.config_profile, - image, - docker_gpus: self.docker_gpus, - provider, - rpc_bind: self - .rpc_bind - .parse() - .map_err(|e| format!("invalid {}: {e}", self.rpc_bind_label))?, - run_id: self.run_id, - node_id: self.node_id, - stage_index: self.stage_index, - layer_end_exclusive: self.layer_end_exclusive, - pipeline_stages: self.pipeline_stages, - model_id: self.model_id, - gguf_source, - tokenizer: self.tokenizer, - default_max_tokens: self.default_max_tokens, - dashboard: self.dashboard, - max_context: self.max_context, - relay, - vastai, - cached_model, - worker_bin: self.worker_bin, - datastream_frame_log: self.datastream_frame_log, - }) - } - - fn set_process_image(&mut self, image: String) { - self.image = image; - self.image_overridden_after_toml = true; - } - - fn set_gguf_repo(&mut self, repo: String) { - let (file, revision) = match &self.gguf_source { - GgufSource::HuggingFaceGguf { file, revision, .. } => (file.clone(), revision.clone()), - GgufSource::LocalPath(_) => (DEFAULT_HF_FILE.to_owned(), None), - }; - self.gguf_source = GgufSource::HuggingFaceGguf { - repo, - file, - revision, - }; - } - - fn set_gguf_file(&mut self, file: String) { - let (repo, revision) = match &self.gguf_source { - GgufSource::HuggingFaceGguf { repo, revision, .. } => (repo.clone(), revision.clone()), - GgufSource::LocalPath(_) => (DEFAULT_HF_REPO.to_owned(), None), - }; - self.gguf_source = GgufSource::HuggingFaceGguf { - repo, - file, - revision, - }; - } - - fn set_gguf_revision(&mut self, revision: Option) { - let (repo, file) = match &self.gguf_source { - GgufSource::HuggingFaceGguf { repo, file, .. } => (repo.clone(), file.clone()), - GgufSource::LocalPath(_) => (DEFAULT_HF_REPO.to_owned(), DEFAULT_HF_FILE.to_owned()), - }; - self.gguf_source = GgufSource::HuggingFaceGguf { - repo, - file, - revision, - }; - } - - fn parse_value(name: &str, value: &str) -> Result - where - T: std::str::FromStr, - T::Err: std::fmt::Display, - { - value - .parse::() - .map_err(|e| format!("invalid {name}={value:?}: {e}")) - } - - fn parse_bool(name: &str, value: &str) -> Result { - match value.to_ascii_lowercase().as_str() { - "1" | "true" | "yes" | "on" => Ok(true), - "0" | "false" | "no" | "off" => Ok(false), - _ => Err(format!( - "invalid {name}={value:?}; use 1/0, true/false, yes/no, or on/off" - )), - } - } -} - -impl Config { - fn from_defaults_toml_env_args(args: impl IntoIterator) -> Result { - Self::from_layers_with_path_and_args(Some(Path::new(DEFAULT_CONFIG_PATH)), args) - } - - fn from_layers_with_path_and_args( - path: Option<&Path>, - args: impl IntoIterator, - ) -> Result { - let mut builder = Self::hardcoded_defaults(); - if let Some(path) = path { - if let Some(overlay) = TomlConfigOverlay::load_optional(path)? { - builder = builder.overlay_toml(overlay)?; - } - } - builder.overlay_env()?.overlay_cli(args)?.finalize() - } - - fn hardcoded_defaults() -> ConfigBuilder { - ConfigBuilder::hardcoded_defaults() - } - - fn uses_planned_execution(&self) -> bool { - self.cached_model.is_some() || self.pipeline_stages > 1 - } - - fn provider_datastream_detail(&self) -> Value { - match self.provider { - ProviderKind::Process => json!({ - "worker_bin": self.worker_bin.as_ref().map(|path| path.to_string_lossy().to_string()), - "cached_model": self.cached_model.as_ref().map(CachedModelConfig::datastream_detail), - }), - ProviderKind::Docker => json!({ - "docker_gpus": &self.docker_gpus, - "cached_model": self.cached_model.as_ref().map(CachedModelConfig::datastream_detail), - }), - ProviderKind::VastAi => self - .vastai - .as_ref() - .map_or_else(|| json!({}), VastAiRuntimeConfig::datastream_detail), - ProviderKind::Mock => json!({}), - } - } - - fn build_run_plan(&self) -> Result { - let host_path = self.local_planning_gguf_path()?; - let metadata = mvp_system::gguf_metadata::read_gguf_planning_metadata(&host_path)?; - let model = metadata.to_model_facts( - self.model_id.clone(), - self.gguf_source.clone(), - self.tokenizer.clone(), - self.max_context, - )?; - if self.pipeline_stages > model.num_layers { - return Err(format!( - "--pipeline-stages={} exceeds GGUF layer count {}; choose N <= {}", - self.pipeline_stages, model.num_layers, model.num_layers - )); - } - - let activation_extent = model - .max_seq_len - .checked_mul(model.hidden_dim) - .and_then(|value| value.checked_mul(model.dtype_width_bytes)) - .ok_or_else(|| "activation ring size overflow while planning pipeline".to_owned())?; - let activation_data_capacity = run_plan::MO01_HEADER_BYTES - .checked_add(activation_extent) - .ok_or_else(|| { - "activation ring data capacity overflow while planning pipeline".to_owned() - })?; - let token_extent = model - .max_seq_len - .checked_mul(4) - .ok_or_else(|| "token ring size overflow while planning pipeline".to_owned())?; - let token_data_capacity = run_plan::MO01_HEADER_BYTES - .checked_add(token_extent) - .ok_or_else(|| { - "token ring data capacity overflow while planning pipeline".to_owned() - })?; - let candidate_pool = (0..self.pipeline_stages) - .map(|stage_index| run_plan::NodeId(self.node_id + 1 + u64::from(stage_index))) - .collect::>(); - let placement = run_plan::PlacementInput::FixedLinear( - candidate_pool - .iter() - .enumerate() - .map(|(stage_index, node_id)| run_plan::StagePlacement { - stage_index: stage_index as u32, - node_id: *node_id, - }) - .collect(), - ); - - run_plan::plan_run(run_plan::PlannerInput { - run_id: run_plan::RunId(self.run_id), - orchestrator_node_id: run_plan::NodeId(self.node_id), - model, - runtime: run_plan::RuntimeConfig { - max_tokens: self.default_max_tokens, - prompt: run_plan::PromptSource::Inline(String::new()), - sampling: run_plan::SamplingPolicy { - temperature_millis: 0, - top_k: 1, - }, - token_output_policy: run_plan::TokenOutputPolicy::EmitAll, - }, - candidate_pool, - stage_count: self.pipeline_stages, - placement, - activation_ring: run_plan::RingSpec { - data_capacity: activation_data_capacity, - alignment: 64, - direction: run_plan::RingDirection::Egress, - host_pinning: run_plan::HostPinning::Pageable, - wake_coalescing: run_plan::WakeCoalescing::PendingBit, - }, - token_ring: run_plan::RingSpec { - data_capacity: token_data_capacity, - alignment: 8, - direction: run_plan::RingDirection::Egress, - host_pinning: run_plan::HostPinning::Pageable, - wake_coalescing: run_plan::WakeCoalescing::PendingBit, - }, - }) - .map_err(|e| format!("plan pipeline run: {:?}", e.kind())) - } - - fn local_planning_gguf_path(&self) -> Result { - if let Some(cached_model) = &self.cached_model { - return Ok(cached_model.host_path.clone()); - } - match &self.gguf_source { - GgufSource::LocalPath(path) => { - let host_path = PathBuf::from(path); - if host_path.is_file() { - Ok(host_path) - } else { - Err(format!( - "local pipeline requires a locally inspectable GGUF before provisioning; {path:?} is not a host file, so use --cached-model-host-path" - )) - } - } - GgufSource::HuggingFaceGguf { repo, file, .. } => Err(format!( - "local pipeline requires a locally inspectable GGUF before provisioning; selected source {repo}/{file} is remote, so use --cached-model-host-path" - )), - } - } - - fn prepare_vastai_ssh_key(&mut self) -> Result<(), String> { - if self.provider != ProviderKind::VastAi { - return Ok(()); - } - - let api_key = self - .vastai - .as_ref() - .and_then(|vastai| vastai.api_key.as_deref()) - .ok_or_else(|| { - "MVP_VASTAI_API_KEY or VASTAI_API_KEY is required when MVP_NODE_PROVIDER=vastai" - .to_owned() - })? - .to_owned(); - let identity = resolve_vastai_ssh_identity( - self.vastai - .as_ref() - .and_then(|vastai| vastai.ssh_identity.clone()), - )?; - if !identity.is_file() { - return Err(format!( - "missing VastAI SSH identity {}; create/register one with vastai create ssh-key or set MVP_VASTAI_SSH_IDENTITY", - identity.display() - )); - } - let public_key = derive_ssh_public_key(&identity)?; - let fingerprint = ssh_public_key_fingerprint(&public_key); - ensure_vastai_account_ssh_key(&api_key, &public_key)?; - - eprintln!( - "VastAI SSH identity {} fingerprint {} registered for account", - identity.display(), - fingerprint - ); - - let vastai = self - .vastai - .as_mut() - .expect("VastAI config exists when provider is vastai"); - vastai.ssh_identity = Some(identity); - vastai.provisioning.ssh_public_key = Some(public_key.clone()); - vastai.ssh_public_key = Some(public_key); - vastai.ssh_public_fingerprint = Some(fingerprint); - Ok(()) - } - - fn build_provisioner( - &self, - bootstrap_runtime: Arc, - ) -> Result, String> { - match self.provider { - ProviderKind::Process => { - let worker_bin = self.worker_bin.clone().unwrap_or(default_worker_bin()?); - if !worker_bin.is_file() { - return Err(format!( - "local process worker binary does not exist: {}", - worker_bin.display() - )); - } - Ok(Box::new(LocalProcessPlugin::new(worker_bin))) - } - ProviderKind::Docker => Ok(Box::new(LocalDockerPlugin::new(docker_container_prefix()))), - ProviderKind::VastAi => { - let vastai = self.vastai.as_ref().ok_or_else(|| { - "VastAI config was not resolved for provider vastai".to_owned() - })?; - if vastai.bootstrap_command.is_none() { - return Err( - "MVP_VASTAI_BOOTSTRAP_COMMAND is required when MVP_NODE_PROVIDER=vastai" - .to_owned(), - ); - } - let api_key = vastai.api_key.clone().ok_or_else(|| { - "MVP_VASTAI_API_KEY or VASTAI_API_KEY is required when MVP_NODE_PROVIDER=vastai" - .to_owned() - })?; - let ssh_identity = vastai - .ssh_identity - .clone() - .ok_or_else(|| "VastAI SSH identity was not prepared".to_owned())?; - let client = ToolsVastAiLeaseClient::from_api_key(api_key)?; - Ok(Box::new(VastAiProvisioningPlugin::new( - client, - SshCommandBootstrapLauncher::new(Some(ssh_identity), bootstrap_runtime), - vastai.provisioning.clone(), - ))) - } - ProviderKind::Mock => Err("mvp-orchestrator does not support mock provider".to_owned()), - } - } - - fn node_spec_env_keys(&self) -> Vec<&'static str> { - let mut keys = vec![ - "MVP_RUN_ID", - "MVP_LOGICAL_NODE_ID", - "MVP_NODE_PROVIDER", - "MVP_STAGE_INDEX", - "MVP_COORDINATOR_ENDPOINT", - "MVP_ORCHESTRATOR_ACTOR", - "MVP_MODEL_ID", - "MVP_IROH_RELAY_MODE", - "MVP_PIPELINE_STAGES", - ]; - if self.relay.url.is_some() { - keys.push(MVP_IROH_RELAY_URL_ENV); - } - if self.provider == ProviderKind::Docker { - keys.push("MVP_DOCKER_GPUS"); - } - if local_tinygrad_worker_env(self.provider).is_some() { - keys.push("MVP_TINYGRAD_WORKER"); - } - if std::env::var_os("MVP_CPU_LINE_PROFILE").is_some() { - keys.push("MVP_CPU_LINE_PROFILE"); - } - if std::env::var_os("MVP_CPU_LINE_PROFILE_INTERVAL_MS").is_some() { - keys.push("MVP_CPU_LINE_PROFILE_INTERVAL_MS"); - } - if std::env::var_os("MVP_TOKEN_PROGRESS_EVERY").is_some() { - keys.push("MVP_TOKEN_PROGRESS_EVERY"); - } - if std::env::var_os("CUDA_DEVICE_SCHEDULE").is_some() { - keys.push("CUDA_DEVICE_SCHEDULE"); - } - if std::env::var_os("MVP_MODEL_CACHE_DIR").is_some() { - keys.push("MVP_MODEL_CACHE_DIR"); - } - if std::env::var_os("HF_TOKEN").is_some() { - keys.push("HF_TOKEN"); - } - match &self.gguf_source { - GgufSource::LocalPath(_) => keys.push("MVP_GGUF_LOCAL_PATH"), - GgufSource::HuggingFaceGguf { revision, .. } => { - keys.push("MVP_GGUF_REPO"); - keys.push("MVP_GGUF_FILE"); - if revision.is_some() { - keys.push("MVP_GGUF_REVISION"); - } - } - } - if matches!(self.tokenizer, TokenizerSource::LocalPath(_)) { - keys.push("MVP_TOKENIZER_LOCAL_PATH"); - } - if self.max_context.is_some() { - keys.push("MVP_MAX_CONTEXT"); - } - keys - } - - fn node_spec( - &self, - coordinator: EndpointAddr, - orchestrator_actor: ActorAddress, - ) -> Result { - self.node_spec_for_stage( - coordinator, - orchestrator_actor, - self.node_id, - self.stage_index, - ) - } - - fn node_spec_for_stage( - &self, - coordinator: EndpointAddr, - orchestrator_actor: ActorAddress, - logical_node_id: u64, - stage_index: u32, - ) -> Result { - let mut env = vec![ - ("MVP_RUN_ID".to_owned(), self.run_id.to_string()), - ( - "MVP_LOGICAL_NODE_ID".to_owned(), - logical_node_id.to_string(), - ), - ("MVP_STAGE_INDEX".to_owned(), stage_index.to_string()), - ( - "MVP_PIPELINE_STAGES".to_owned(), - self.pipeline_stages.to_string(), - ), - ( - "MVP_NODE_PROVIDER".to_owned(), - self.provider.as_str().to_owned(), - ), - ( - "MVP_COORDINATOR_ENDPOINT".to_owned(), - serde_json::to_string(&coordinator) - .map_err(|e| format!("serialize coordinator endpoint: {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(), - relay_mode_env_value(&self.relay.mode).to_owned(), - ), - ]; - if let Some(url) = &self.relay.url { - env.push((MVP_IROH_RELAY_URL_ENV.to_owned(), url.clone())); - } - if self.provider == ProviderKind::Docker { - env.push(("MVP_DOCKER_GPUS".to_owned(), self.docker_gpus.clone())); - } - env.extend(local_tinygrad_worker_env(self.provider)); - env.extend(optional_env("MVP_CPU_LINE_PROFILE")); - env.extend(optional_env("MVP_CPU_LINE_PROFILE_INTERVAL_MS")); - env.extend(optional_env("MVP_TOKEN_PROGRESS_EVERY")); - env.extend(optional_env("CUDA_DEVICE_SCHEDULE")); - env.extend(optional_env("MVP_MODEL_CACHE_DIR")); - env.extend(optional_env("HF_TOKEN")); - match &self.gguf_source { - GgufSource::LocalPath(path) => { - env.push(("MVP_GGUF_LOCAL_PATH".to_owned(), path.clone())) - } - GgufSource::HuggingFaceGguf { - repo, - file, - revision, - } => { - env.push(("MVP_GGUF_REPO".to_owned(), repo.clone())); - env.push(("MVP_GGUF_FILE".to_owned(), file.clone())); - if let Some(revision) = revision { - env.push(("MVP_GGUF_REVISION".to_owned(), revision.clone())); - } - } - } - if let TokenizerSource::LocalPath(path) = &self.tokenizer { - env.push(("MVP_TOKENIZER_LOCAL_PATH".to_owned(), path.clone())); - } - if let Some(max_context) = self.max_context { - env.push(("MVP_MAX_CONTEXT".to_owned(), max_context.to_string())); - } - let args = match self.provider { - ProviderKind::VastAi => self - .vastai - .as_ref() - .and_then(|vastai| vastai.bootstrap_command.clone()) - .into_iter() - .collect(), - ProviderKind::Process | ProviderKind::Docker => Vec::new(), - ProviderKind::Mock => { - return Err("mvp-orchestrator does not support mock provider".to_owned()); - } - }; - let mounts = if self.provider == ProviderKind::Docker { - self.cached_model - .as_ref() - .map(|cached_model| { - vec![ProviderMount { - host_path: cached_model.host_path.to_string_lossy().to_string(), - container_path: cached_model.container_path.clone(), - readonly: self.uses_planned_execution(), - }] - }) - .unwrap_or_default() - } else { - Vec::new() - }; - Ok(NodeProvisionSpec { - run_id: self.run_id, - node_id: logical_node_id, - stage_index: Some(stage_index), - image: self.image.clone(), - env, - args, - mounts, - }) - } -} - -#[derive(Clone)] -struct RuntimeReady { - endpoint: EndpointAddr, - node_actor: ActorAddress, - datastream_publisher: ActorAddress, - stage_index: u32, - readiness_id: u64, - swim_node_id: DistNodeId, -} - -#[derive(Clone)] -struct PromptRuntimeReady { - first_stage: RuntimeReady, - final_stage: RuntimeReady, -} -#[derive(Clone)] -struct RuntimeReadyAckTarget { - node_id: u64, - ready: RuntimeReady, -} - -fn runtime_ready_barrier_met(stack: &DistributionRuntimeStack, ready: &RuntimeReady) -> bool { - stack.member_state(ready.swim_node_id) == Some(MemberState::Alive) - && stack.route_owner(ready.node_actor) == Some(ready.swim_node_id) -} - -fn enqueue_runtime_ready_ack( - stack: &DistributionRuntimeStack, - ready: &RuntimeReady, - run_id: u64, - node_id: u64, -) -> Result<(), String> { - stack - .runtime - .send_to( - ready.node_actor, - NodeAgentMsg::RuntimeReadyAck { - run_id, - node_id, - stage_index: ready.stage_index, - readiness_id: ready.readiness_id, - }, - ) - .map_err(|e| format!("send runtime ready ack: {e}")) -} - -fn enqueue_datastream_subscribe( - driver: &mut IrohDriver, - stack: &DistributionRuntimeStack, - ready: &RuntimeReady, - run_id: u64, - node_id: u64, -) -> Result<(), String> { - let mut flow_id = [0_u8; 16]; - flow_id[..8].copy_from_slice(&run_id.to_le_bytes()); - flow_id[8..].copy_from_slice(&node_id.to_le_bytes()); - stack - .runtime - .send_to( - ready.datastream_publisher, - DatastreamPublisherMsg::Subscribe(DatastreamSubscribe { - collector: driver.endpoint_addr(), - request: SubscriptionRequest::all(), - flow_id, - token: Vec::new(), - }), - ) - .map_err(|e| format!("send datastream subscribe: {e}")) -} - -#[allow(clippy::too_many_arguments)] -fn wait_for_runtime_ready_acks( - driver: &mut IrohDriver, - stack: &DistributionRuntimeStack, - obs_rx: &mpsc::Receiver, - frame_rx: &mpsc::Receiver, - frame_tx: &mpsc::Sender, - orchestrator_reports: &swactor::runtime::Inbox, - stop_rx: &mpsc::Receiver<()>, - dashboard: Option<&DashboardSupport>, - orch_datastream: &mut OrchDatastream, - orch_stdio_rx: Option<&mpsc::Receiver>, - run_id: u64, - orchestrator_node_id: u64, - provider: ProviderKind, - targets: &[RuntimeReadyAckTarget], -) -> Result<(), String> { - let mut pending = targets - .iter() - .cloned() - .map(|target| { - ( - ( - target.node_id, - target.ready.stage_index, - target.ready.readiness_id, - ), - target, - ) - }) - .collect::>(); - let mut attempts = BTreeMap::<(u64, u32, u64), u64>::new(); - let started = Instant::now(); - let mut last_send = None::; - - while !pending.is_empty() { - pump(driver, stack, frame_tx); - drain_orch_stdio_capture( - orch_stdio_rx, - orch_datastream, - dashboard, - run_id, - orchestrator_node_id, - ); - if stop_requested(stop_rx) { - return Err( - "shutdown requested while waiting for runtime-ready acknowledgements".to_owned(), - ); - } - while let Ok(observation) = obs_rx.try_recv() { - emit_plugin_observation(orch_datastream, dashboard, provider, &observation); - match observation { - PluginObservation::DatastreamFrame { .. } => {} - PluginObservation::ProviderLine { .. } - | PluginObservation::StdoutLine { .. } - | PluginObservation::StderrLine { .. } => {} - PluginObservation::Failed { reason, .. } => return Err(reason), - PluginObservation::Exited { - status, node_id, .. - } => { - return Err(format!("node {node_id} exited before ready: {status:?}")); - } - } - } - drain_frames(frame_rx, dashboard, orch_datastream); - while let Some(report) = orchestrator_reports.try_recv() { - let OrchestratorReport::NodeRuntimeReadyAck { - run_id: ack_run_id, - node_id, - stage_index, - readiness_id, - } = report - else { - continue; - }; - if ack_run_id != run_id { - continue; - } - let key = (node_id, stage_index, readiness_id); - let Some(target) = pending.remove(&key) else { - continue; - }; - orch_datastream.emit_bootstrap( - dashboard, - run_id, - orchestrator_node_id, - "runtime_ready_ack", - "ready", - json!({ - "node_id":node_id, - "node_actor":target.ready.node_actor, - "readiness_id":readiness_id, - "attempts":attempts.get(&key).copied().unwrap_or(0), - }), - ); - } - if pending.is_empty() { - return Ok(()); - } - if started.elapsed() >= RUNTIME_READY_ACK_TIMEOUT { - let pending_list = pending - .keys() - .map(|(node_id, stage_index, readiness_id)| { - format!("{node_id}/{stage_index}/{readiness_id}") - }) - .collect::>() - .join(", "); - return Err(format!( - "runtime_ready_ack timed out for node(s): {pending_list}" - )); - } - if last_send.is_none_or(|sent_at| sent_at.elapsed() >= RUNTIME_READY_ACK_RETRY_INTERVAL) { - for (key, target) in &pending { - if stack.route_owner(target.ready.datastream_publisher) - == Some(target.ready.swim_node_id) - && let Err(error) = enqueue_datastream_subscribe( - driver, - stack, - &target.ready, - run_id, - target.node_id, - ) - { - orch_datastream.emit_bootstrap( - dashboard, - run_id, - orchestrator_node_id, - "datastream_subscribe", - "failed", - json!({"node_id":target.node_id,"error":error}), - ); - } - enqueue_runtime_ready_ack(stack, &target.ready, run_id, target.node_id)?; - let attempt = attempts.entry(*key).or_default(); - *attempt += 1; - let attempt = *attempt; - orch_datastream.emit_bootstrap( - dashboard, - run_id, - orchestrator_node_id, - "runtime_ready_ack", - "sent", - json!({ - "node_id":target.node_id, - "node_actor":target.ready.node_actor, - "readiness_id":target.ready.readiness_id, - "attempt":attempt, - }), - ); - } - driver.drain_outbox(&stack.outbox); - last_send = Some(Instant::now()); - } - thread::sleep(PUMP_INTERVAL); - } - Ok(()) -} - -#[cfg(test)] -struct ProvisionedNodeGuard<'a> { - provisioner: &'a mut dyn ProvisionPlugin, - handle: Option, -} - -#[cfg(test)] -impl<'a> ProvisionedNodeGuard<'a> { - fn new( - provisioner: &'a mut dyn ProvisionPlugin, - handle: mvp_system::provisioning::PluginNodeHandle, - ) -> Self { - Self { - provisioner, - handle: Some(handle), - } - } - - fn stop(&mut self) -> Result<(), String> { - let Some(handle) = self.handle.take() else { - return Ok(()); - }; - self.provisioner.stop_node(&handle) - } -} - -#[cfg(test)] -impl Drop for ProvisionedNodeGuard<'_> { - fn drop(&mut self) { - let _ = self.stop(); - } -} - -struct ProvisionedClusterGuard { - provisioner: Box, - handles: Vec, -} - -impl ProvisionedClusterGuard { - fn new( - provisioner: Box, - handles: Vec, - ) -> Self { - Self { - provisioner, - handles, - } - } - - fn complete_bootstrap_all(&mut self) -> Result<(), String> { - for handle in &self.handles { - self.provisioner.complete_bootstrap(handle)?; - } - Ok(()) - } - - fn stop(&mut self) -> Result<(), String> { - let mut first_error = None; - while let Some(handle) = self.handles.pop() { - if let Err(error) = self.provisioner.stop_node(&handle) - && first_error.is_none() - { - first_error = Some(error); - } - } - match first_error { - Some(error) => Err(error), - None => Ok(()), - } - } -} - -impl Drop for ProvisionedClusterGuard { - fn drop(&mut self) { - let _ = self.stop(); - } -} - -fn start_node_with_stdio_capture( - provisioner: Box, - node_spec: NodeProvisionSpec, - sink: PluginSink, - orch_stdio_rx: Option<&mpsc::Receiver>, - dashboard: Option<&DashboardSupport>, - orch_datastream: &mut OrchDatastream, - run_id: u64, - node_id: u64, -) -> ( - Box, - Result, -) { - let (tx, rx) = mpsc::channel(); - thread::spawn(move || { - let mut provisioner = provisioner; - let result = provisioner.start_node(node_spec, sink); - let _ = tx.send((provisioner, result)); - }); - - loop { - match rx.recv_timeout(Duration::from_millis(100)) { - Ok(result) => return result, - Err(mpsc::RecvTimeoutError::Timeout) => { - drain_orch_stdio_capture( - orch_stdio_rx, - orch_datastream, - dashboard, - run_id, - node_id, - ); - } - Err(mpsc::RecvTimeoutError::Disconnected) => { - return ( - Box::new(FailedProvisionPlugin), - Err("provider start worker disconnected".to_owned()), - ); - } - } - } -} - -#[allow(clippy::too_many_arguments)] -fn start_and_provision_workers( - mut provisioner: Box, - config: &Config, - pipeline_plan: Option<&run_plan::RunPlan>, - driver: &mut IrohDriver, - stack: &DistributionRuntimeStack, - obs_rx: &mpsc::Receiver, - frame_rx: &mpsc::Receiver, - frame_tx: &mpsc::Sender, - orchestrator_reports: &swactor::runtime::Inbox, - stop_rx: &mpsc::Receiver<()>, - dashboard: Option<&DashboardSupport>, - orch_datastream: &mut OrchDatastream, - orch_stdio_rx: Option<&mpsc::Receiver>, - sink: PluginSink, - coordinator: EndpointAddr, - pipeline_coordinator: EndpointAddr, - orchestrator_actor: ActorAddress, -) -> Result<(ProvisionedClusterGuard, PromptRuntimeReady), String> { - let stage_specs = stage_node_specs( - config, - pipeline_plan, - coordinator.clone(), - orchestrator_actor, - )?; - let expected_node_ids = stage_specs - .iter() - .map(|spec| spec.node_id) - .collect::>(); - orch_datastream.emit_bootstrap( - dashboard, - config.run_id, - config.node_id, - "node_spec", - "ready", - json!({ - "provider":config.provider.as_str(), - "image":&config.image, - "relay_mode":relay_mode_env_value(&config.relay.mode), - "docker_gpus":if config.provider == ProviderKind::Docker { Some(config.docker_gpus.as_str()) } else { None }, - "provider_config":config.provider_datastream_detail(), - "env_keys":config.node_spec_env_keys(), - "worker_count":stage_specs.len(), - "worker_node_ids":expected_node_ids, - }), - ); - - let mut handles = Vec::with_capacity(stage_specs.len()); - for node_spec in stage_specs { - orch_datastream.emit_event( - dashboard, - ProvisionEvent { - run_id: config.run_id, - node_id: node_spec.node_id, - kind: ProvisionEventKind::ProvisionStart, - provider: Some(config.provider.as_str().to_owned()), - message: Some(format!( - "starting {} image {}", - config.provider.as_str(), - config.image - )), - }, - ); - orch_datastream.emit_bootstrap( - dashboard, - config.run_id, - config.node_id, - "provider_start", - "started", - json!({ - "provider":config.provider.as_str(), - "image":&config.image, - "node_id":node_spec.node_id, - "stage_index":node_spec.stage_index, - }), - ); - let (returned_provisioner, handle_result) = start_node_with_stdio_capture( - provisioner, - node_spec.clone(), - sink.clone(), - orch_stdio_rx, - dashboard, - orch_datastream, - config.run_id, - node_spec.node_id, - ); - provisioner = returned_provisioner; - match handle_result { - Ok(handle) => handles.push(handle), - Err(error) => { - orch_datastream.emit_bootstrap( - dashboard, - config.run_id, - config.node_id, - "provider_start", - "failed", - json!({"provider":config.provider.as_str(),"node_id":node_spec.node_id,"error":error}), - ); - stop_started_nodes(&mut *provisioner, &mut handles); - drain_orch_stdio_capture( - orch_stdio_rx, - orch_datastream, - dashboard, - config.run_id, - config.node_id, - ); - return Err(error); - } - } - } - let mut provisioned_nodes = ProvisionedClusterGuard::new(provisioner, handles); - drain_orch_stdio_capture( - orch_stdio_rx, - orch_datastream, - dashboard, - config.run_id, - config.node_id, - ); - - orch_datastream.emit_bootstrap( - dashboard, - config.run_id, - config.node_id, - "node_runtime_ready", - "started", - json!({"worker_count":expected_node_ids.len(),"node_ids":expected_node_ids}), - ); - let readies = if pipeline_plan.is_some() { - match wait_for_runtime_readies( - driver, - stack, - obs_rx, - frame_rx, - frame_tx, - orchestrator_reports, - stop_rx, - dashboard, - orch_datastream, - orch_stdio_rx, - config.run_id, - &expected_node_ids, - config.provider, - ) { - Ok(readies) => readies, - Err(error) => { - orch_datastream.emit_bootstrap( - dashboard, - config.run_id, - config.node_id, - "node_runtime_ready", - "failed", - json!({"error":error}), - ); - return Err(error); - } - } - } else { - let ready = match wait_for_runtime_ready( - driver, - stack, - obs_rx, - frame_rx, - frame_tx, - orchestrator_reports, - stop_rx, - dashboard, - orch_datastream, - orch_stdio_rx, - config.run_id, - config.node_id, - config.provider, - ) { - Ok(ready) => ready, - Err(error) => { - orch_datastream.emit_bootstrap( - dashboard, - config.run_id, - config.node_id, - "node_runtime_ready", - "failed", - json!({"error":error}), - ); - return Err(error); - } - }; - BTreeMap::from([(config.node_id, ready)]) - }; - for (node_id, ready) in &readies { - orch_datastream.emit_bootstrap( - dashboard, - config.run_id, - config.node_id, - "node_runtime_ready", - "ready", - json!({"endpoint":&ready.endpoint,"node_actor":ready.node_actor,"node_id":node_id,"stage_index":ready.stage_index}), - ); - } - provisioned_nodes.complete_bootstrap_all()?; - let ack_targets = readies - .iter() - .map(|(node_id, ready)| RuntimeReadyAckTarget { - node_id: *node_id, - ready: ready.clone(), - }) - .collect::>(); - wait_for_runtime_ready_acks( - driver, - stack, - obs_rx, - frame_rx, - frame_tx, - orchestrator_reports, - stop_rx, - dashboard, - orch_datastream, - orch_stdio_rx, - config.run_id, - config.node_id, - config.provider, - &ack_targets, - )?; - - orch_datastream.emit_bootstrap( - dashboard, - config.run_id, - config.node_id, - "stage_provision", - "started", - stage_provision_detail(config, pipeline_plan), - ); - if pipeline_plan.is_none() { - let ready = readies - .get(&config.node_id) - .ok_or_else(|| "missing runtime-ready node for single-stage run".to_owned())?; - provision_stage(stack, ready.node_actor, config)?; - } - - orch_datastream.emit_bootstrap( - dashboard, - config.run_id, - config.node_id, - "weights_loaded", - "started", - json!({"model_id":&config.model_id,"expected":expected_node_ids.len()}), - ); - let weights_result = if pipeline_plan.is_some() { - wait_for_weights_loaded_count( - driver, - stack, - obs_rx, - frame_rx, - frame_tx, - orchestrator_reports, - stop_rx, - dashboard, - orch_datastream, - orch_stdio_rx, - config.run_id, - config.node_id, - config.provider, - expected_node_ids.len(), - pipeline_plan.expect("pipeline mode requires plan"), - &readies, - &pipeline_coordinator, - ) - } else { - wait_for_weights_loaded( - driver, - stack, - obs_rx, - frame_rx, - frame_tx, - orchestrator_reports, - stop_rx, - dashboard, - orch_datastream, - orch_stdio_rx, - config.run_id, - config.node_id, - config.provider, - config.stage_index, - ) - }; - match weights_result { - Ok(()) => orch_datastream.emit_bootstrap( - dashboard, - config.run_id, - config.node_id, - "weights_loaded", - "ready", - json!({"source":"actor_stage_ready","model_id":&config.model_id,"expected":expected_node_ids.len()}), - ), - Err(error) => { - orch_datastream.emit_bootstrap( - dashboard, - config.run_id, - config.node_id, - "weights_loaded", - "failed", - json!({"error":error}), - ); - drain_orch_stdio_capture( - orch_stdio_rx, - orch_datastream, - dashboard, - config.run_id, - config.node_id, - ); - return Err(error); - } - } - - let prompt_ready = if let Some(plan) = pipeline_plan { - let first_node_id = plan - .stages - .iter() - .find(|stage| stage.stage_index == 0) - .map(|stage| stage.node_id.0) - .ok_or_else(|| "pipeline plan missing stage 0".to_owned())?; - let final_node_id = plan - .stages - .iter() - .find(|stage| stage.stage_index + 1 == stage.stage_count) - .map(|stage| stage.node_id.0) - .ok_or_else(|| "pipeline plan missing final stage".to_owned())?; - PromptRuntimeReady { - first_stage: readies - .get(&first_node_id) - .cloned() - .ok_or_else(|| "missing stage 0 runtime-ready node".to_owned())?, - final_stage: readies - .get(&final_node_id) - .cloned() - .ok_or_else(|| "missing final-stage runtime-ready node".to_owned())?, - } - } else { - let ready = readies - .get(&config.node_id) - .cloned() - .ok_or_else(|| "missing single-stage runtime-ready node".to_owned())?; - PromptRuntimeReady { - first_stage: ready.clone(), - final_stage: ready, - } - }; - Ok((provisioned_nodes, prompt_ready)) -} - -fn stage_node_specs( - config: &Config, - pipeline_plan: Option<&run_plan::RunPlan>, - coordinator: EndpointAddr, - orchestrator_actor: ActorAddress, -) -> Result, String> { - if let Some(plan) = pipeline_plan { - let mut stages = plan.stages.clone(); - stages.sort_by_key(|stage| stage.stage_index); - stages - .into_iter() - .map(|stage| { - config.node_spec_for_stage( - coordinator.clone(), - orchestrator_actor, - stage.node_id.0, - stage.stage_index, - ) - }) - .collect() - } else { - Ok(vec![config.node_spec(coordinator, orchestrator_actor)?]) - } -} - -fn stop_started_nodes( - provisioner: &mut dyn ProvisionPlugin, - handles: &mut Vec, -) { - while let Some(handle) = handles.pop() { - let _ = provisioner.stop_node(&handle); - } -} - -fn stage_provision_detail(config: &Config, pipeline_plan: Option<&run_plan::RunPlan>) -> Value { - if let Some(plan) = pipeline_plan { - json!({ - "run_id":config.run_id, - "stage_count":plan.stages.len(), - "stages":plan.stages.iter().map(|stage| { - json!({ - "node_id":stage.node_id.0, - "stage_index":stage.stage_index, - "layer_range":{"start":stage.layer_start,"end_exclusive":stage.layer_end_exclusive}, - "inbound_edge_id":stage.inbound_edge.0, - "outbound_edge_id":stage.outbound_edge.0, - }) - }).collect::>(), - "model_id":&config.model_id, - }) - } else { - json!({ - "run_id":config.run_id, - "node_id":config.node_id, - "stage_index":config.stage_index, - "stage_count":1, - "layer_range":{"start":0,"end_exclusive":config.layer_end_exclusive}, - "model_id":&config.model_id, - }) - } -} - -fn stage_provision_wire_from_plan( - plan: &run_plan::RunPlan, - stage_index: u32, - readies: &BTreeMap, - coordinator: &EndpointAddr, -) -> Result { - let provision = run_plan::derive_stage_provision(plan, stage_index) - .map_err(|e| format!("derive stage {stage_index} provision: {e:?}"))?; - Ok(StageProvisionWire { - run_id: provision.run_id.0, - authorized_orchestrator: 0, - node_id: provision.node_id.0, - stage_index: provision.stage_index, - stage_count: provision.stage_count, - layer_start: provision.layer_start, - layer_end_exclusive: provision.layer_end_exclusive, - inbound_edge_id: provision.inbound.edge_id.0, - outbound_edge_id: provision.outbound.edge_id.0, - inbound_edge: Some(StageInboundEdgeWire { - edge_id: provision.inbound.edge_id.0, - kind: stage_edge_kind_wire(provision.inbound.kind), - object_spec: stage_object_spec_wire(provision.inbound.object_spec), - ring_spec: stage_ring_spec_wire(provision.inbound.ring_spec), - }), - outbound_edge: Some(StageOutboundEdgeWire { - edge_id: provision.outbound.edge_id.0, - kind: stage_edge_kind_wire(provision.outbound.kind), - consumer_node_id: provision.outbound.consumer_node_id.0, - consumer_endpoint: stage_consumer_endpoint( - provision.outbound.consumer_node_id.0, - plan, - readies, - coordinator, - )?, - object_spec: stage_object_spec_wire(provision.outbound.object_spec), - ring_spec: stage_ring_spec_wire(provision.outbound.ring_spec), - }), - model_id: provision.model.model_id, - gguf_source: provision.gguf_source, - tokenizer: provision.tokenizer, - }) -} - -fn provision_stage_from_plan( - stack: &DistributionRuntimeStack, - node_actor: ActorAddress, - plan: &run_plan::RunPlan, - stage_index: u32, - readies: &BTreeMap, - coordinator: &EndpointAddr, -) -> Result<(), String> { - let provision = stage_provision_wire_from_plan(plan, stage_index, readies, coordinator)?; - stack - .runtime - .send_to(node_actor, NodeAgentMsg::ProvisionStage(provision)) - .map_err(|e| format!("send stage {stage_index} provision: {e}")) -} - -fn stage_consumer_endpoint( - consumer_node_id: u64, - plan: &run_plan::RunPlan, - readies: &BTreeMap, - coordinator: &EndpointAddr, -) -> Result, String> { - if consumer_node_id - == plan.stages.first().map_or(0, |stage| { - plan.edges - .iter() - .find(|edge| edge.kind == run_plan::EdgeKind::TokenIn) - .and_then(|edge| match edge.producer { - run_plan::EdgeEndpoint::Orchestrator { node_id } => Some(node_id.0), - run_plan::EdgeEndpoint::Stage { .. } => None, - }) - .unwrap_or(stage.node_id.0) - }) - { - return Ok(Some(coordinator.clone())); - } - readies - .get(&consumer_node_id) - .map(|ready| Some(ready.endpoint.clone())) - .ok_or_else(|| { - format!("missing runtime-ready endpoint for consumer node {consumer_node_id}") - }) -} - -fn stage_edge_kind_wire(kind: run_plan::EdgeKind) -> StageEdgeKindWire { - match kind { - run_plan::EdgeKind::TokenIn => StageEdgeKindWire::TokenIn, - run_plan::EdgeKind::Activation => StageEdgeKindWire::Activation, - run_plan::EdgeKind::TokenOut => StageEdgeKindWire::TokenOut, - } -} - -fn stage_object_spec_wire(spec: run_plan::ObjectSpec) -> StageObjectSpecWire { - StageObjectSpecWire { - max_extent: spec.max_extent, - alignment: spec.alignment, - } -} - -fn stage_ring_spec_wire(spec: run_plan::RingSpec) -> StageRingSpecWire { - StageRingSpecWire { - data_capacity: spec.data_capacity, - alignment: spec.alignment, - } -} - -#[allow(clippy::too_many_arguments)] -fn wait_for_runtime_readies( - driver: &mut IrohDriver, - stack: &DistributionRuntimeStack, - obs_rx: &mpsc::Receiver, - frame_rx: &mpsc::Receiver, - frame_tx: &mpsc::Sender, - orchestrator_reports: &swactor::runtime::Inbox, - stop_rx: &mpsc::Receiver<()>, - dashboard: Option<&DashboardSupport>, - orch_datastream: &mut OrchDatastream, - orch_stdio_rx: Option<&mpsc::Receiver>, - run_id: u64, - expected_node_ids: &[u64], - provider: ProviderKind, -) -> Result, String> { - let expected = expected_node_ids.iter().copied().collect::>(); - let mut pending = BTreeMap::::new(); - loop { - pump(driver, stack, frame_tx); - emit_swim_transitions( - orch_datastream, - dashboard, - run_id, - expected_node_ids.first().copied().unwrap_or(0), - stack, - ); - drain_frames(frame_rx, dashboard, orch_datastream); - drain_orch_stdio_capture( - orch_stdio_rx, - orch_datastream, - dashboard, - run_id, - expected_node_ids.first().copied().unwrap_or(0), - ); - if stop_requested(stop_rx) { - return Err("shutdown requested while waiting for pipeline nodes ready".to_owned()); - } - while let Ok(observation) = obs_rx.try_recv() { - emit_plugin_observation(orch_datastream, dashboard, provider, &observation); - match observation { - PluginObservation::DatastreamFrame { .. } => {} - PluginObservation::ProviderLine { .. } - | PluginObservation::StdoutLine { .. } - | PluginObservation::StderrLine { .. } => {} - PluginObservation::Failed { reason, .. } => return Err(reason), - PluginObservation::Exited { - status, node_id, .. - } => { - return Err(format!("node {node_id} exited before ready: {status:?}")); - } - } - } - while let Some(report) = orchestrator_reports.try_recv() { - if let OrchestratorReport::NodeRuntimeReady { - run_id: report_run_id, - node_id, - stage_index, - endpoint, - node_actor, - datastream_publisher, - readiness_id, - } = report - && report_run_id == run_id - && expected.contains(&node_id) - { - pending.insert( - node_id, - RuntimeReady { - endpoint: endpoint.clone(), - node_actor, - datastream_publisher, - stage_index, - readiness_id, - swim_node_id: DistNodeId(*endpoint.id.as_bytes()), - }, - ); - } - } - if expected.iter().all(|node_id| { - pending - .get(node_id) - .is_some_and(|ready| runtime_ready_barrier_met(stack, ready)) - }) { - return Ok(pending); - } - thread::sleep(PUMP_INTERVAL); - } -} - -#[allow(clippy::too_many_arguments)] -fn wait_for_weights_loaded_count( - driver: &mut IrohDriver, - stack: &DistributionRuntimeStack, - obs_rx: &mpsc::Receiver, - frame_rx: &mpsc::Receiver, - frame_tx: &mpsc::Sender, - orchestrator_reports: &swactor::runtime::Inbox, - stop_rx: &mpsc::Receiver<()>, - dashboard: Option<&DashboardSupport>, - orch_datastream: &mut OrchDatastream, - orch_stdio_rx: Option<&mpsc::Receiver>, - run_id: u64, - node_id: u64, - provider: ProviderKind, - expected_count: usize, - pipeline_plan: &run_plan::RunPlan, - readies: &BTreeMap, - pipeline_coordinator: &EndpointAddr, -) -> Result<(), String> { - let expected_stages = pipeline_plan - .stages - .iter() - .map(|stage| stage.stage_index) - .collect::>(); - let mut loaded_stages = BTreeSet::::new(); - let mut last_resend = Instant::now(); - let mut active_stage = None::; - let mut resend_attempt = 0_u64; - loop { - pump(driver, stack, frame_tx); - emit_swim_transitions(orch_datastream, dashboard, run_id, node_id, stack); - drain_orch_stdio_capture(orch_stdio_rx, orch_datastream, dashboard, run_id, node_id); - if stop_requested(stop_rx) { - return Err("shutdown requested while waiting for pipeline weights loaded".to_owned()); - } - if loaded_stages.len() >= expected_count { - return Ok(()); - } - if active_stage.map_or(true, |stage_index| loaded_stages.contains(&stage_index)) { - active_stage = - next_pipeline_weight_load_stage(pipeline_plan, &loaded_stages, active_stage) - .map(|stage| stage.stage_index); - last_resend = Instant::now() - Duration::from_secs(1); - } - if last_resend.elapsed() >= Duration::from_secs(1) { - resend_attempt += 1; - let Some(stage) = - next_pipeline_weight_load_stage(pipeline_plan, &loaded_stages, active_stage) - else { - return Err(format!( - "missing unloaded pipeline weight stage; loaded {} of {expected_count}", - loaded_stages.len() - )); - }; - let stage_node_id = stage.node_id.0; - let ready = readies.get(&stage_node_id).ok_or_else(|| { - format!("missing runtime-ready node for stage {}", stage.stage_index) - })?; - orch_datastream.emit_bootstrap( - dashboard, - run_id, - node_id, - "stage_provision_send", - "sent", - json!({ - "attempt":resend_attempt, - "stage_count":pipeline_plan.stages.len(), - "stage_index":stage.stage_index, - "loaded_stage_count":loaded_stages.len() - }), - ); - let route_owner = stack.route_owner(ready.node_actor); - let member_state = stack.member_state(ready.swim_node_id); - orch_datastream.emit_bootstrap_to_channel( - dashboard, - MVP_STAGE_ROUTE, - run_id, - node_id, - "stage_route_check", - "observed", - json!({ - "attempt":resend_attempt, - "stage_index":stage.stage_index, - "stage_node_id":stage_node_id, - "node_actor":ready.node_actor, - "datastream_publisher":ready.datastream_publisher, - "swim_node_id":format!("{:?}", ready.swim_node_id), - "member_state":member_state.map(|state| format!("{:?}", state)), - "route_owner":route_owner.map(|owner| format!("{:?}", owner)), - "datastream_route_owner":stack.route_owner(ready.datastream_publisher).map(|owner| format!("{:?}", owner)), - "route_matches_ready":route_owner == Some(ready.swim_node_id), - }), - ); - provision_stage_from_plan( - stack, - ready.node_actor, - pipeline_plan, - stage.stage_index, - readies, - pipeline_coordinator, - )?; - pump(driver, stack, frame_tx); - last_resend = Instant::now(); - } - while let Ok(observation) = obs_rx.try_recv() { - emit_plugin_observation(orch_datastream, dashboard, provider, &observation); - match observation { - PluginObservation::Failed { reason, .. } => return Err(reason), - PluginObservation::Exited { - node_id, status, .. - } => { - return Err(format!( - "node {node_id} exited while loading weights: {status:?}" - )); - } - PluginObservation::DatastreamFrame { .. } => {} - PluginObservation::ProviderLine { .. } - | PluginObservation::StdoutLine { .. } - | PluginObservation::StderrLine { .. } => {} - } - } - drain_frames(frame_rx, dashboard, orch_datastream); - while let Some(report) = orchestrator_reports.try_recv() { - match report { - OrchestratorReport::WeightsReady { - run_id: report_run_id, - node_id: _, - stage_index, - } if report_run_id == run_id && expected_stages.contains(&stage_index) => { - loaded_stages.insert(stage_index); - if active_stage == Some(stage_index) { - active_stage = None; - last_resend = Instant::now() - Duration::from_secs(1); - } - if loaded_stages.len() >= expected_count { - return Ok(()); - } - } - OrchestratorReport::StageFault { - run_id: report_run_id, - stage_index, - } if report_run_id == run_id && expected_stages.contains(&stage_index) => { - return Err(format!( - "stage {stage_index} faulted while loading pipeline weights" - )); - } - _ => {} - } - } - thread::sleep(PUMP_INTERVAL); - } -} - -fn next_pipeline_weight_load_stage<'a>( - pipeline_plan: &'a run_plan::RunPlan, - loaded_stages: &BTreeSet, - active_stage: Option, -) -> Option<&'a run_plan::StagePlan> { - if let Some(stage_index) = active_stage { - if !loaded_stages.contains(&stage_index) { - if let Some(stage) = pipeline_plan - .stages - .iter() - .find(|stage| stage.stage_index == stage_index) - { - return Some(stage); - } - } - } - pipeline_plan - .stages - .iter() - .filter(|stage| !loaded_stages.contains(&stage.stage_index)) - .min_by_key(|stage| stage.stage_index) -} - -struct FailedProvisionPlugin; - -impl ProvisionPlugin for FailedProvisionPlugin { - fn start_node( - &mut self, - _spec: NodeProvisionSpec, - _sink: PluginSink, - ) -> Result { - 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, - ) -> Result<(), String> { - Ok(()) - } -} - -struct PromptWork { - request: SubmitPrompt, - events: mpsc::Sender, -} - -struct ActivePrompt { - request: SubmitPrompt, - events: mpsc::Sender, -} - -#[derive(Clone, Debug)] -struct CollectedDatastreamFrame { - stream: StreamId, - channel_name: String, - frame: Frame, -} - -fn drain_datastream_connections( - driver: &mut IrohDriver, - frame_tx: &mpsc::Sender, -) { - driver.pump_datastream_ingress(); - for read in driver.drain_datastream_reads() { - let mut channels = read - .header - .channels - .iter() - .map(|descriptor| { - ( - ChannelRef { - stream: descriptor.stream.clone(), - channel: descriptor.id, - }, - descriptor.name.clone(), - ) - }) - .collect::>(); - for event in read.events { - match event { - DatastreamEvent::ChannelDeclared(descriptor) => { - channels.insert( - ChannelRef { - stream: descriptor.stream.clone(), - channel: descriptor.id, - }, - descriptor.name, - ); - } - DatastreamEvent::Frame(delivery) => { - let channel_name = channels - .get(&delivery.channel) - .cloned() - .unwrap_or_else(|| format!("channel#{}", delivery.channel.channel.0)); - let frame = Frame::new( - delivery.channel.channel, - delivery.position, - delivery.payload, - ); - if frame_tx - .send(CollectedDatastreamFrame { - stream: delivery.channel.stream, - channel_name, - frame, - }) - .is_err() - { - return; - } - } - DatastreamEvent::StreamDeclared(_) | DatastreamEvent::StreamEnded(_) => {} - } - } - } -} - -struct FrameArchive { - file: File, - next_seq: u64, -} - -impl FrameArchive { - fn open(path: &Path) -> Result { - if let Some(parent) = path.parent() - && !parent.as_os_str().is_empty() - { - std::fs::create_dir_all(parent).map_err(|e| { - format!("create datastream frame log dir {}: {e}", parent.display()) - })?; - } - let file = OpenOptions::new() - .create(true) - .append(true) - .open(path) - .map_err(|e| format!("open datastream frame log {}: {e}", path.display()))?; - Ok(Self { file, next_seq: 0 }) - } - - fn record(&mut self, source: &str, stream: &StreamId, channel: &str, frame: &Frame) { - let payload = match std::str::from_utf8(&frame.payload) { - Ok(text) => json!({"encoding":"utf8","value":text}), - Err(_) => json!({"encoding":"bytes","value":frame.payload}), - }; - let record = json!({ - "arrival_seq":self.next_seq, - "arrival_unix_ms":benchmark_observability::unix_ms_now(), - "source":source, - "stream":stream.to_string(), - "channel":channel, - "channel_id":frame.channel.0, - "position":frame.position.0, - "payload":payload, - }); - self.next_seq += 1; - let _ = serde_json::to_writer(&mut self.file, &record); - let _ = writeln!(self.file); - let _ = self.file.flush(); - } -} - -struct OrchDatastream { - stream: StreamId, - endpoint: DatastreamEndpoint, - producer: DatastreamProducer, - channels: BTreeMap, - channel_names: BTreeMap, - archive: Option, -} - -impl OrchDatastream { - fn new(run_id: u64, frame_log: Option<&Path>) -> Result { - let stream = StreamId::new(NodeId::new("mvp-orchestrator"), Lifetime(run_id)); - let endpoint = DatastreamEndpoint::with_descriptor( - StreamDescriptor { - stream: stream.clone(), - label: Some("mvp orchestrator".to_owned()), - origin: StreamOrigin::Orchestrator, - }, - 4096, - 1024, - ); - let producer = endpoint.producer(); - let mut out = Self { - stream, - endpoint, - producer, - channels: BTreeMap::new(), - channel_names: BTreeMap::new(), - archive: frame_log.map(FrameArchive::open).transpose()?, - }; - for name in [ - MVP_PROVISIONING_EVENTS, - MVP_ORCH_BOOTSTRAP, - MVP_ORCH_PROMPT, - MVP_SWIM_MEMBERSHIP, - MVP_STAGE_ROUTE, - ] { - out.channel_by_name(name); - } - Ok(out) - } - - fn channel_by_name(&mut self, name: &str) -> ChannelId { - if let Some(id) = self.channels.get(name).copied() { - return id; - } - let id = self.producer.register_channel( - name, - ChannelContent::JsonRecord { - schema: Some(name.to_owned()), - }, - ); - self.channels.insert(name.to_owned(), id); - self.channel_names.insert(id, name.to_owned()); - id - } - - fn emit_event(&mut self, dashboard: Option<&DashboardSupport>, event: ProvisionEvent) { - let payload = serde_json::to_vec(&MvpProvisionEventRecord::new(event)) - .expect("serialize provisioning event"); - self.emit_bytes(dashboard, MVP_PROVISIONING_EVENTS, payload); - } - - fn emit_log(&mut self, dashboard: Option<&DashboardSupport>, line: ProvisionLogLine) { - let channel = mvp_provision_log_channel(line.node_id, line.stream); - let payload = - serde_json::to_vec(&MvpProvisionLogRecord::new(line)).expect("serialize provision log"); - self.emit_bytes(dashboard, &channel, payload); - } - - fn emit_bootstrap( - &mut self, - dashboard: Option<&DashboardSupport>, - run_id: u64, - node_id: u64, - phase: &str, - status: &str, - detail: Value, - ) { - self.emit_bootstrap_to_channel( - dashboard, - MVP_ORCH_BOOTSTRAP, - run_id, - node_id, - phase, - status, - detail, - ); - } - - fn emit_bootstrap_to_channel( - &mut self, - dashboard: Option<&DashboardSupport>, - channel: &str, - run_id: u64, - node_id: u64, - phase: &str, - status: &str, - detail: Value, - ) { - let payload = serde_json::to_vec(&json!({ - "type":"OrchBootstrap", - "phase":phase, - "status":status, - "run_id":run_id, - "node_id":node_id, - "benchmark":benchmark_observability::stamp("mvp-orchestrator"), - "detail":detail, - })) - .expect("serialize orch bootstrap event"); - self.emit_bytes(dashboard, channel, payload); - } - - fn emit_prompt( - &mut self, - dashboard: Option<&DashboardSupport>, - run_id: u64, - node_id: u64, - request_id: u64, - phase: &str, - status: &str, - detail: Value, - ) { - let payload = serde_json::to_vec(&json!({ - "type":"OrchPromptEvent", - "phase":phase, - "status":status, - "run_id":run_id, - "node_id":node_id, - "request_id":request_id, - "benchmark":benchmark_observability::stamp("mvp-orchestrator"), - "detail":detail, - })) - .expect("serialize orch prompt event"); - self.emit_bytes(dashboard, MVP_ORCH_PROMPT, payload); - } - - fn emit_bytes( - &mut self, - dashboard: Option<&DashboardSupport>, - channel: &str, - payload: Vec, - ) { - self.emit_bytes_from(dashboard, channel, payload, "orchestrator"); - } - - fn emit_bytes_from( - &mut self, - dashboard: Option<&DashboardSupport>, - channel: &str, - payload: Vec, - source: &str, - ) { - let id = self.channel_by_name(channel); - self.producer.submit_bytes(id, payload); - self.flush(dashboard, source); - } - - fn flush(&mut self, dashboard: Option<&DashboardSupport>, source: &str) { - let stream = self.stream.clone(); - for frame in self.endpoint.mux().drain() { - let channel = self - .channel_names - .get(&frame.channel) - .cloned() - .unwrap_or_else(|| format!("channel#{}", frame.channel.0)); - ingest_dashboard_frame(dashboard, &stream, &channel, &frame); - self.archive_frame(source, &stream, &channel, &frame); - } - } - - fn archive_frame(&mut self, source: &str, stream: &StreamId, channel: &str, frame: &Frame) { - if let Some(archive) = &mut self.archive { - archive.record(source, stream, channel, frame); - } - } -} - -struct OrchStdioCapture; - -struct OrchStdioLine { - stream: ProvisionLogStream, - line: String, -} - -#[cfg(target_os = "linux")] -impl OrchStdioCapture { - fn install() -> Result>, String> { - let stdout_read = Self::redirect_stream(libc::STDOUT_FILENO, "stdout")?; - let stderr_read = Self::redirect_stream(libc::STDERR_FILENO, "stderr")?; - let (tx, rx) = mpsc::channel(); - Self::spawn_reader(stdout_read, ProvisionLogStream::Stdout, tx.clone()); - Self::spawn_reader(stderr_read, ProvisionLogStream::Stderr, tx); - Ok(Some(rx)) - } - - fn redirect_stream(fd: libc::c_int, name: &str) -> Result { - let mut pipe_fds = [0; 2]; - let pipe_result = unsafe { libc::pipe(pipe_fds.as_mut_ptr()) }; - if pipe_result != 0 { - return Err(format!( - "create orchestrator {name} capture pipe: {}", - std::io::Error::last_os_error() - )); - } - - let dup_result = unsafe { libc::dup2(pipe_fds[1], fd) }; - let close_write_result = unsafe { libc::close(pipe_fds[1]) }; - if dup_result < 0 { - let error = std::io::Error::last_os_error(); - let _ = unsafe { libc::close(pipe_fds[0]) }; - return Err(format!("redirect orchestrator {name}: {error}")); - } - if close_write_result != 0 { - let error = std::io::Error::last_os_error(); - let _ = unsafe { libc::close(pipe_fds[0]) }; - return Err(format!("close orchestrator {name} duplicate fd: {error}")); - } - - Ok(unsafe { File::from_raw_fd(pipe_fds[0]) }) - } - - fn spawn_reader(file: File, stream: ProvisionLogStream, tx: mpsc::Sender) { - thread::spawn(move || { - let reader = BufReader::new(file); - for line in reader.lines() { - let Ok(line) = line else { - break; - }; - if tx.send(OrchStdioLine { stream, line }).is_err() { - break; - } - } - }); - } -} - -#[cfg(not(target_os = "linux"))] -impl OrchStdioCapture { - fn install() -> Result>, String> { - Ok(None) - } -} - -fn install_orch_stdio_capture() -> Result>, String> { - OrchStdioCapture::install() -} - -fn drain_orch_stdio_capture( - rx: Option<&mpsc::Receiver>, - datastream: &mut OrchDatastream, - dashboard: Option<&DashboardSupport>, - run_id: u64, - node_id: u64, -) { - let Some(rx) = rx else { - return; - }; - while let Ok(line) = rx.try_recv() { - datastream.emit_log( - dashboard, - ProvisionLogLine { - run_id, - node_id, - stream: line.stream, - line: line.line, - }, - ); - } -} - -#[cfg(feature = "dashboard")] -struct DashboardSupport { - handle: dashboard::DashboardHandle, - _runtime: tokio::runtime::Runtime, -} - -#[cfg(feature = "dashboard")] -impl DashboardSupport { - fn start(enabled: bool) -> Result, String> { - if !enabled { - return Ok(None); - } - let mut config = dashboard::DashboardConfig::default(); - if let Some(port) = env_optional("MVP_DASHBOARD_PORT") { - config.port = port - .parse::() - .map_err(|e| format!("invalid MVP_DASHBOARD_PORT={port:?}: {e}"))?; - } - let runtime = tokio::runtime::Builder::new_multi_thread() - .enable_all() - .build() - .map_err(|e| format!("dashboard runtime: {e}"))?; - let handle = dashboard::DashboardHandle::new(config); - handle.register_view(Arc::new(MvpClusterDashboardView::new())); - handle.spawn_http(runtime.handle()); - Ok(Some(Self { - handle, - _runtime: runtime, - })) - } - - fn publish_frame(&self, stream: &StreamId, channel: &str, frame: &Frame) { - self.handle.publish(dashboard::FrameEvent { - stream: dashboard::StreamEvent { - node: stream.node.as_str().to_string(), - life: stream.life.0, - }, - channel: channel.to_owned(), - position: frame.position.0, - payload: frame.payload.clone(), - }); - } -} - -#[cfg(not(feature = "dashboard"))] -struct DashboardSupport; - -#[cfg(not(feature = "dashboard"))] -impl DashboardSupport { - fn start(enabled: bool) -> Result, String> { - if enabled { - return Err( - "MVP_DASHBOARD requires building mvp-system with feature dashboard".to_owned(), - ); - } - Ok(None) - } - - fn publish_frame(&self, _stream: &StreamId, _channel: &str, _frame: &Frame) {} -} - -struct ChannelObservationSink { - tx: Mutex>, -} - -impl PluginObservationSink for ChannelObservationSink { - fn observe(&self, observation: PluginObservation) { - let _ = self.tx.lock().send(observation); - } -} - -fn spawn_prompt_rpc( - bind: SocketAddr, - work_tx: mpsc::Sender, - default_max_tokens: u32, -) -> Result { - let listener = TcpListener::bind(bind).map_err(|e| format!("bind prompt RPC {bind}: {e}"))?; - let addr = listener - .local_addr() - .map_err(|e| format!("read prompt RPC addr: {e}"))?; - thread::spawn(move || { - for accepted in listener.incoming() { - match accepted { - Ok(stream) => { - let tx = work_tx.clone(); - thread::spawn(move || { - let _ = handle_prompt_connection(stream, tx, default_max_tokens); - }); - } - Err(_) => break, - } - } - }); - Ok(addr) -} -fn handle_prompt_connection( - stream: TcpStream, - work_tx: mpsc::Sender, - default_max_tokens: u32, -) -> Result<(), String> { - let mut reader = BufReader::new( - stream - .try_clone() - .map_err(|e| format!("clone prompt stream: {e}"))?, - ); - let mut writer = stream; - loop { - let request = match read_submit_prompt(&mut reader) { - Ok(Some(request)) => request, - Ok(None) => break, - Err(error) if error.contains("expected value at line 1 column 1") => break, - Err(error) => return Err(error), - }; - let request = request.with_defaults(default_max_tokens); - let (event_tx, event_rx) = mpsc::channel(); - work_tx - .send(PromptWork { - request, - events: event_tx, - }) - .map_err(|_| "prompt loop stopped".to_owned())?; - for event in event_rx { - let terminal = event.is_terminal(); - write_json_line(&mut writer, &event)?; - if terminal { - break; - } - } - } - Ok(()) -} - -fn wait_for_runtime_ready( - driver: &mut IrohDriver, - stack: &DistributionRuntimeStack, - obs_rx: &mpsc::Receiver, - frame_rx: &mpsc::Receiver, - frame_tx: &mpsc::Sender, - orchestrator_reports: &swactor::runtime::Inbox, - stop_rx: &mpsc::Receiver<()>, - dashboard: Option<&DashboardSupport>, - orch_datastream: &mut OrchDatastream, - orch_stdio_rx: Option<&mpsc::Receiver>, - run_id: u64, - node_id: u64, - provider: ProviderKind, -) -> Result { - let mut pending_ready: Option = None; - let mut node_swim_started = false; - let mut node_swim_ready = false; - let mut node_route_started = false; - loop { - pump(driver, stack, frame_tx); - drain_frames(frame_rx, dashboard, orch_datastream); - drain_orch_stdio_capture(orch_stdio_rx, orch_datastream, dashboard, run_id, node_id); - if stop_requested(stop_rx) { - return Err("shutdown requested while waiting for node ready".to_owned()); - } - while let Ok(observation) = obs_rx.try_recv() { - emit_plugin_observation(orch_datastream, dashboard, provider, &observation); - match observation { - PluginObservation::DatastreamFrame { .. } => {} - PluginObservation::ProviderLine { .. } - | PluginObservation::StdoutLine { .. } - | PluginObservation::StderrLine { .. } => {} - PluginObservation::Failed { reason, .. } => return Err(reason), - PluginObservation::Exited { status, .. } => { - return Err(format!("node exited before ready: {status:?}")); - } - } - } - while let Some(report) = orchestrator_reports.try_recv() { - if let OrchestratorReport::NodeRuntimeReady { - run_id: report_run_id, - node_id: report_node_id, - stage_index, - endpoint, - node_actor, - datastream_publisher, - readiness_id, - } = report - { - if report_run_id == run_id && report_node_id == node_id { - let reset_progress = pending_ready - .as_ref() - .map(|ready| ready.readiness_id != readiness_id) - .unwrap_or(true); - if reset_progress { - node_swim_started = false; - node_swim_ready = false; - node_route_started = false; - } - let swim_node_id = DistNodeId(*endpoint.id.as_bytes()); - pending_ready = Some(RuntimeReady { - endpoint, - node_actor, - datastream_publisher, - stage_index, - readiness_id, - swim_node_id, - }); - } - } - } - if let Some(ready) = pending_ready.as_ref() { - if runtime_ready_barrier_met(stack, ready) { - return Ok(ready.clone()); - } - let swim_ready = stack.member_state(ready.swim_node_id) == Some(MemberState::Alive); - let route_ready = stack.route_owner(ready.node_actor) == Some(ready.swim_node_id); - if !swim_ready { - if !node_swim_started { - orch_datastream.emit_bootstrap( - dashboard, - run_id, - node_id, - "node_swim", - "started", - json!({"node":format!("{:?}", ready.swim_node_id),"readiness_id":ready.readiness_id}), - ); - node_swim_started = true; - } - } else if !route_ready { - if !node_swim_ready { - orch_datastream.emit_bootstrap( - dashboard, - run_id, - node_id, - "node_swim", - "ready", - json!({"node":format!("{:?}", ready.swim_node_id),"readiness_id":ready.readiness_id}), - ); - node_swim_ready = true; - } - if !node_route_started { - orch_datastream.emit_bootstrap( - dashboard, - run_id, - node_id, - "node_route", - "started", - json!({"node_actor":ready.node_actor,"node":format!("{:?}", ready.swim_node_id),"readiness_id":ready.readiness_id}), - ); - node_route_started = true; - } - } - } - thread::sleep(PUMP_INTERVAL); - } -} - -fn provision_stage( - stack: &DistributionRuntimeStack, - node_actor: ActorAddress, - config: &Config, -) -> Result<(), String> { - stack - .runtime - .send_to( - node_actor, - NodeAgentMsg::ProvisionStage({ - let layer_end_exclusive = config.layer_end_exclusive.ok_or_else(|| { - "unplanned single-stage execution requires --layer-end-exclusive or MVP_LAYER_END_EXCLUSIVE; cached-model runs use metadata-derived planning".to_owned() - })?; - StageProvisionWire { - run_id: config.run_id, - authorized_orchestrator: 0, - node_id: config.node_id, - stage_index: config.stage_index, - stage_count: 1, - layer_start: 0, - layer_end_exclusive, - inbound_edge_id: 1, - outbound_edge_id: 2, - inbound_edge: None, - outbound_edge: None, - model_id: config.model_id.clone(), - gguf_source: config.gguf_source.clone(), - tokenizer: config.tokenizer.clone(), - } - }), - ) - .map_err(|e| format!("send stage provision: {e}")) -} - -fn wait_for_weights_loaded( - driver: &mut IrohDriver, - stack: &DistributionRuntimeStack, - obs_rx: &mpsc::Receiver, - frame_rx: &mpsc::Receiver, - frame_tx: &mpsc::Sender, - orchestrator_reports: &swactor::runtime::Inbox, - stop_rx: &mpsc::Receiver<()>, - dashboard: Option<&DashboardSupport>, - orch_datastream: &mut OrchDatastream, - orch_stdio_rx: Option<&mpsc::Receiver>, - run_id: u64, - node_id: u64, - provider: ProviderKind, - stage_index: u32, -) -> Result<(), String> { - loop { - pump(driver, stack, frame_tx); - drain_orch_stdio_capture(orch_stdio_rx, orch_datastream, dashboard, run_id, node_id); - if stop_requested(stop_rx) { - return Err("shutdown requested while waiting for weights loaded".to_owned()); - } - while let Ok(observation) = obs_rx.try_recv() { - emit_plugin_observation(orch_datastream, dashboard, provider, &observation); - match observation { - PluginObservation::Failed { reason, .. } => return Err(reason), - PluginObservation::Exited { status, .. } => { - return Err(format!("node exited while loading weights: {status:?}")); - } - PluginObservation::DatastreamFrame { .. } => {} - PluginObservation::ProviderLine { .. } - | PluginObservation::StdoutLine { .. } - | PluginObservation::StderrLine { .. } => {} - } - } - drain_frames(frame_rx, dashboard, orch_datastream); - while let Some(report) = orchestrator_reports.try_recv() { - match report { - OrchestratorReport::WeightsReady { - run_id: report_run_id, - node_id: _, - stage_index: report_stage_index, - } if report_run_id == run_id && report_stage_index == stage_index => { - return Ok(()); - } - OrchestratorReport::StageFault { - run_id: report_run_id, - stage_index: report_stage_index, - } if report_run_id == run_id && report_stage_index == stage_index => { - return Err(format!( - "stage {report_stage_index} faulted while loading weights" - )); - } - _ => {} - } - } - thread::sleep(PUMP_INTERVAL); - } -} - -#[derive(Debug)] -struct PipelineTokenRecord { - object_id: u64, - sequence: u64, - token_id: u32, - eos: bool, -} - -enum PipelineSendHandle { - Driver(EdgeSendHandle), - #[cfg(test)] - Channel(tokio_mpsc::UnboundedSender>), -} - -impl PipelineSendHandle { - fn send(&self, bytes: Vec) -> Result<(), String> { - match self { - Self::Driver(handle) => handle.send(bytes), - #[cfg(test)] - Self::Channel(tx) => tx - .send(bytes) - .map_err(|_| "pipeline token-in sender stopped".to_owned()), - } - } -} -struct PendingEncode { - request_id: u64, -} - -struct PendingDecode { - request_id: u64, - token_id: u32, - eos: bool, - reached_limit: bool, -} - -struct PipelinePromptRuntime { - token_in_edge_id: u64, - token_out_edge_id: u64, - token_spec: run_plan::ObjectSpec, - token_out_spec: run_plan::ObjectSpec, - token_in_sender: PipelineSendHandle, - recv_rx: mpsc::Receiver>, - recv_tx: mpsc::Sender>, - recv_buffer: Vec, - tokenizer_encode_actor: ActorAddress, - tokenizer_decode_actor: ActorAddress, - tokenizer_reply_to: ActorAddress, - pending_encode: Option, - pending_decode: Option, - next_sequence: u64, - generated_tokens: Vec, - final_text: String, - active: Option, - started_at: Option, -} - -impl PipelinePromptRuntime { - fn new( - driver: &IrohDriver, - plan: &run_plan::RunPlan, - first_stage_endpoint: EndpointAddr, - tokenizer_encode_actor: ActorAddress, - tokenizer_decode_actor: ActorAddress, - tokenizer_reply_to: ActorAddress, - ) -> Result { - let token_in_edge = plan - .edges - .iter() - .find(|edge| edge.kind == run_plan::EdgeKind::TokenIn) - .ok_or_else(|| "pipeline plan missing token-in edge".to_owned())?; - let token_out_edge = plan - .edges - .iter() - .find(|edge| edge.kind == run_plan::EdgeKind::TokenOut) - .ok_or_else(|| "pipeline plan missing token-out edge".to_owned())?; - let (recv_tx, recv_rx) = mpsc::channel(); - Ok(Self { - token_in_edge_id: token_in_edge.edge_id.0, - token_out_edge_id: token_out_edge.edge_id.0, - token_spec: token_in_edge.object_spec, - token_out_spec: token_out_edge.object_spec, - token_in_sender: PipelineSendHandle::Driver( - driver.spawn_edge_send_pump(first_stage_endpoint, token_in_edge.edge_id.0)?, - ), - recv_rx, - recv_tx, - tokenizer_encode_actor, - tokenizer_decode_actor, - tokenizer_reply_to, - pending_encode: None, - pending_decode: None, - recv_buffer: Vec::new(), - next_sequence: 0, - generated_tokens: Vec::new(), - final_text: String::new(), - active: None, - started_at: None, - }) - } - - fn is_active(&self) -> bool { - self.active.is_some() - } - - fn start_prompt( - &mut self, - request: SubmitPrompt, - events: mpsc::Sender, - runtime: &Arc, - dashboard: Option<&DashboardSupport>, - orch_datastream: &mut OrchDatastream, - run_id: u64, - node_id: u64, - ) -> Result<(), String> { - let request_id = request.request_id; - self.generated_tokens.clear(); - self.final_text.clear(); - self.recv_buffer.clear(); - self.pending_decode = None; - self.pending_encode = Some(PendingEncode { request_id }); - self.started_at = Some(Instant::now()); - orch_datastream.emit_prompt( - dashboard, - run_id, - node_id, - request_id, - "pipeline_tokenizer_encode", - "started", - json!({"node_actor":self.tokenizer_encode_actor,"reply_to":self.tokenizer_reply_to,"prompt_bytes":request.prompt_text.len()}), - ); - runtime - .send_to( - self.tokenizer_encode_actor, - NodeAgentMsg::EncodePrompt { - request_id, - prompt: request.prompt_text.clone(), - reply_to: self.tokenizer_reply_to, - }, - ) - .map_err(|e| format!("send tokenizer encode request: {e}"))?; - self.active = Some(ActivePrompt { request, events }); - Ok(()) - } - - fn drain_tokenizer_events( - &mut self, - runtime: &Arc, - tokenizer_events: &swactor::runtime::Inbox, - dashboard: Option<&DashboardSupport>, - orch_datastream: &mut OrchDatastream, - run_id: u64, - node_id: u64, - ) -> Result<(), String> { - while let Some(event) = tokenizer_events.try_recv() { - match event { - TokenizerEvent::PromptEncoded { request_id, tokens } => self - .handle_encoded_prompt( - request_id, - tokens, - dashboard, - orch_datastream, - run_id, - node_id, - )?, - TokenizerEvent::TokensDecoded { request_id, text } => self.handle_decoded_tokens( - runtime, - request_id, - text, - dashboard, - orch_datastream, - run_id, - node_id, - )?, - TokenizerEvent::Fault { request_id, error } => { - self.fault_active(request_id, format!("tokenizer request failed: {error}")); - } - } - } - Ok(()) - } - - fn handle_encoded_prompt( - &mut self, - request_id: u64, - tokens: Vec, - dashboard: Option<&DashboardSupport>, - orch_datastream: &mut OrchDatastream, - run_id: u64, - node_id: u64, - ) -> Result<(), String> { - let Some(pending) = self.pending_encode.take() else { - return Ok(()); - }; - if pending.request_id != request_id { - self.pending_encode = Some(pending); - return Ok(()); - } - let Some(active) = self.active.as_ref() else { - return Ok(()); - }; - if active.request.request_id != request_id { - return Ok(()); - } - orch_datastream.emit_prompt( - dashboard, - run_id, - node_id, - request_id, - "pipeline_tokenizer_encode", - "ready", - json!({"node_actor":self.tokenizer_encode_actor,"reply_to":self.tokenizer_reply_to,"tokens":tokens.len()}), - ); - let sequence = self.next_sequence; - orch_datastream.emit_prompt( - dashboard, - run_id, - node_id, - request_id, - "pipeline_token_in", - "started", - json!({"edge_id":self.token_in_edge_id,"sequence":sequence,"tokens":tokens.len()}), - ); - self.send_token_in(sequence, &tokens)?; - orch_datastream.emit_prompt( - dashboard, - run_id, - node_id, - request_id, - "pipeline_token_in", - "ready", - json!({"edge_id":self.token_in_edge_id,"sequence":sequence}), - ); - Ok(()) - } - - #[allow(clippy::too_many_arguments)] - fn handle_decoded_tokens( - &mut self, - _runtime: &Arc, - request_id: u64, - text: String, - dashboard: Option<&DashboardSupport>, - orch_datastream: &mut OrchDatastream, - run_id: u64, - node_id: u64, - ) -> Result<(), String> { - let Some(pending) = self.pending_decode.take() else { - return Ok(()); - }; - if pending.request_id != request_id { - self.pending_decode = Some(pending); - return Ok(()); - } - let Some(active) = self.active.as_ref() else { - return Ok(()); - }; - if active.request.request_id != request_id { - return Ok(()); - } - orch_datastream.emit_prompt( - dashboard, - run_id, - node_id, - request_id, - "pipeline_tokenizer_decode", - "ready", - json!({"node_actor":self.tokenizer_decode_actor,"reply_to":self.tokenizer_reply_to,"text_bytes":text.len()}), - ); - self.final_text.push_str(&text); - if !text.is_empty() { - let _ = active - .events - .send(PromptEvent::TextDelta { request_id, text }); - } - if pending.eos || pending.reached_limit { - let elapsed_ms = self - .started_at - .map(|started| started.elapsed().as_millis() as u64) - .unwrap_or(0); - let final_text = self.final_text.clone(); - let tokens_generated = self.generated_tokens.len() as u32; - orch_datastream.emit_prompt( - dashboard, - run_id, - node_id, - request_id, - "prompt_complete", - "ready", - json!({ - "event":"Done", - "terminal":true, - "tokens_generated":tokens_generated, - "elapsed_ms":elapsed_ms, - "final_text_bytes":final_text.len(), - }), - ); - let _ = active.events.send(PromptEvent::Done { - request_id, - final_text, - tokens_generated, - elapsed_ms, - }); - self.active = None; - return Ok(()); - } - self.send_token_in(self.next_sequence, &[pending.token_id])?; - Ok(()) - } - - fn send_token_in(&self, sequence: u64, tokens: &[u32]) -> Result<(), String> { - self.token_in_sender.send(encode_token_record( - self.token_spec, - self.token_in_edge_id, - sequence, - tokens, - false, - )?) - } - - fn request_decode( - &mut self, - runtime: &Arc, - request_id: u64, - token_id: u32, - eos: bool, - reached_limit: bool, - ) -> Result<(), String> { - self.pending_decode = Some(PendingDecode { - request_id, - token_id, - eos, - reached_limit, - }); - runtime - .send_to( - self.tokenizer_decode_actor, - NodeAgentMsg::DecodeTokens { - request_id, - tokens: vec![token_id], - reply_to: self.tokenizer_reply_to, - }, - ) - .map_err(|e| format!("send tokenizer decode request: {e}")) - } - - fn fault_active(&mut self, request_id: u64, error: String) { - if let Some(active) = self.active.take() - && active.request.request_id == request_id - { - let _ = active.events.send(PromptEvent::Fault { request_id, error }); - } - self.pending_encode = None; - self.pending_decode = None; - } - - fn poll_driver(&mut self, driver: &mut IrohDriver) { - driver.pump_edge_ingress(); - for event in driver.drain_edge_events() { - match event { - EdgeTransportEvent::BytesRead { edge_id, bytes, .. } - if edge_id == self.token_out_edge_id => - { - let _ = self.recv_tx.send(bytes); - } - EdgeTransportEvent::StreamFault { - edge_id: Some(edge_id), - reason, - .. - } if edge_id == self.token_out_edge_id => { - if let Some(request_id) = - self.active.as_ref().map(|active| active.request.request_id) - { - self.fault_active( - request_id, - format!("pipeline token-out stream fault: {reason:?}"), - ); - } - } - _ => {} - } - } - } - - fn drain_tokens( - &mut self, - runtime: &Arc, - dashboard: Option<&DashboardSupport>, - orch_datastream: &mut OrchDatastream, - run_id: u64, - node_id: u64, - ) -> Result<(), String> { - if self.pending_decode.is_some() { - return Ok(()); - } - while let Ok(bytes) = self.recv_rx.try_recv() { - self.recv_buffer.extend_from_slice(&bytes); - while self.pending_decode.is_none() - && let Some(record) = - take_pipeline_token_record(&mut self.recv_buffer, self.token_out_spec)? - { - if record.sequence != self.next_sequence { - return Err(format!( - "pipeline token sequence violation: expected {}, got {}", - self.next_sequence, record.sequence - )); - } - self.next_sequence = self.next_sequence.saturating_add(1); - let Some(active) = self.active.as_ref() else { - continue; - }; - let request_id = active.request.request_id; - orch_datastream.emit_prompt( - dashboard, - run_id, - node_id, - request_id, - "pipeline_token_out", - "observed", - json!({"edge_id":self.token_out_edge_id,"object_id":record.object_id,"sequence":record.sequence,"token_id":record.token_id,"eos":record.eos}), - ); - self.generated_tokens.push(record.token_id); - let reached_limit = self.generated_tokens.len() as u32 >= active.request.max_tokens; - orch_datastream.emit_prompt( - dashboard, - run_id, - node_id, - request_id, - "pipeline_tokenizer_decode", - "started", - json!({"node_actor":self.tokenizer_decode_actor,"reply_to":self.tokenizer_reply_to,"token_id":record.token_id}), - ); - self.request_decode( - runtime, - request_id, - record.token_id, - record.eos, - reached_limit, - )?; - } - } - Ok(()) - } -} - -fn encode_token_record( - spec: run_plan::ObjectSpec, - _edge_id: u64, - sequence: u64, - tokens: &[u32], - eos: bool, -) -> Result, String> { - let payload = tokens - .iter() - .flat_map(|token| token.to_le_bytes()) - .collect::>(); - let mut flags = ingress::ObjectFlags::default(); - flags.end_of_sequence = eos; - Ok( - ingress::ObjectRecordBuilder::new(ingress_object_spec_from_plan(spec)) - .object_id(ingress::ObjectId(9000_u64.saturating_add(sequence))) - .sequence(sequence) - .payload(payload) - .flags(flags) - .encode(), - ) -} - -fn ingress_object_spec_from_plan(spec: run_plan::ObjectSpec) -> ingress::ObjectSpec { - let extent_alignment = match spec.kind { - run_plan::ObjectKind::Token => u64::from(spec.dtype_width_bytes), - _ => u64::from(spec.alignment), - }; - ingress::ObjectSpec { - max_extent: spec.max_extent, - alignment: extent_alignment, - layout: ingress::ObjectLayout::Token, - } -} - -fn take_pipeline_token_record( - buffer: &mut Vec, - spec: run_plan::ObjectSpec, -) -> Result, String> { - let record = - match ingress::read_object_record(buffer, ingress_object_spec_from_plan(spec), false) - .map_err(|reason| format!("invalid token-out record: {reason:?}"))? - { - ingress::ObjectRecordRead::Incomplete => return Ok(None), - ingress::ObjectRecordRead::Complete(record) => record, - }; - let payload = record - .payload(buffer) - .ok_or_else(|| "token-out record payload missing".to_owned())?; - if payload.len() != 4 { - return Err(format!( - "token-out payload must be exactly one u32, got {}", - payload.len() - )); - } - let token_id = u32::from_le_bytes(payload.try_into().unwrap()); - let out = PipelineTokenRecord { - object_id: record.object_id.0, - sequence: record.sequence, - token_id, - eos: record.flags.end_of_sequence, - }; - buffer.drain(..record.total_len); - Ok(Some(out)) -} - -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -enum PromptRuntimeMode { - DirectInferPrompt, - PipelineTokenEdges, -} - -fn prompt_runtime_mode(pipeline_plan: Option<&run_plan::RunPlan>) -> PromptRuntimeMode { - if pipeline_plan.is_some() { - PromptRuntimeMode::PipelineTokenEdges - } else { - PromptRuntimeMode::DirectInferPrompt - } -} - -fn serve_prompts( - driver: &mut IrohDriver, - stack: &DistributionRuntimeStack, - obs_rx: &mpsc::Receiver, - frame_rx: &mpsc::Receiver, - frame_tx: &mpsc::Sender, - work_rx: &mpsc::Receiver, - prompt_events: &swactor::runtime::Inbox, - stop_rx: &mpsc::Receiver<()>, - dashboard: Option<&DashboardSupport>, - orch_datastream: &mut OrchDatastream, - orch_stdio_rx: Option<&mpsc::Receiver>, - run_id: u64, - node_id: u64, - node_actor: ActorAddress, - reply_to: ActorAddress, - tokenizer_events: &swactor::runtime::Inbox, - tokenizer_encode_actor: ActorAddress, - tokenizer_decode_actor: ActorAddress, - tokenizer_reply_to: ActorAddress, - provider: ProviderKind, - pipeline_plan: Option<&run_plan::RunPlan>, - prompt_endpoint: EndpointAddr, -) -> Result<(), String> { - let mut pipeline_runtime = match prompt_runtime_mode(pipeline_plan) { - PromptRuntimeMode::PipelineTokenEdges => Some(PipelinePromptRuntime::new( - driver, - pipeline_plan.expect("pipeline mode requires plan"), - prompt_endpoint, - tokenizer_encode_actor, - tokenizer_decode_actor, - tokenizer_reply_to, - )?), - PromptRuntimeMode::DirectInferPrompt => None, - }; - let mut active: Option = None; - loop { - pump(driver, stack, frame_tx); - if let Some(pipeline) = pipeline_runtime.as_mut() { - pipeline.poll_driver(driver); - pipeline.drain_tokenizer_events( - &stack.runtime, - tokenizer_events, - dashboard, - orch_datastream, - run_id, - node_id, - )?; - pipeline.drain_tokens(&stack.runtime, dashboard, orch_datastream, run_id, node_id)?; - } - drain_observations(obs_rx, dashboard, orch_datastream, provider)?; - drain_frames(frame_rx, dashboard, orch_datastream); - drain_orch_stdio_capture(orch_stdio_rx, orch_datastream, dashboard, run_id, node_id); - if stop_rx.try_recv().is_ok() { - orch_datastream.emit_bootstrap( - dashboard, - run_id, - node_id, - "shutdown", - "started", - json!({"source":"stdin"}), - ); - return Ok(()); - } - - if active.is_none() - && pipeline_runtime - .as_ref() - .is_none_or(|pipeline| !pipeline.is_active()) - && let Ok(work) = work_rx.try_recv() - { - let request = work.request; - let request_id = request.request_id; - orch_datastream.emit_prompt( - dashboard, - run_id, - node_id, - request_id, - "prompt_work", - "observed", - json!({ - "prompt_bytes":request.prompt_text.len(), - "max_tokens":request.max_tokens, - }), - ); - if let Some(pipeline) = pipeline_runtime.as_mut() { - pipeline.start_prompt( - request, - work.events, - &stack.runtime, - dashboard, - orch_datastream, - run_id, - node_id, - )?; - continue; - } - orch_datastream.emit_prompt( - dashboard, - run_id, - node_id, - request_id, - "node_prompt_send", - "started", - json!({"node_actor":node_actor,"reply_to":reply_to}), - ); - match stack.runtime.send_to( - node_actor, - NodeAgentMsg::InferPrompt { - request_id, - prompt: request.prompt_text.clone(), - max_tokens: request.max_tokens, - reply_to, - }, - ) { - Ok(()) => { - orch_datastream.emit_prompt( - dashboard, - run_id, - node_id, - request_id, - "node_prompt_send", - "ready", - json!({"node_actor":node_actor,"reply_to":reply_to}), - ); - active = Some(ActivePrompt { - request, - events: work.events, - }); - } - Err(error) => { - orch_datastream.emit_prompt( - dashboard, - run_id, - node_id, - request_id, - "node_prompt_send", - "failed", - json!({"node_actor":node_actor,"reply_to":reply_to,"error":error.to_string()}), - ); - return Err(format!("send prompt request: {error}")); - } - } - } - - while let Some(event) = prompt_events.try_recv() { - let request_id = event.request_id(); - let Some(current) = active.as_ref() else { - orch_datastream.emit_prompt( - dashboard, - run_id, - node_id, - request_id, - "node_prompt_event", - "dropped", - json!({"reason":"no_active_prompt","event":prompt_event_name(&event)}), - ); - continue; - }; - if request_id != current.request.request_id { - orch_datastream.emit_prompt( - dashboard, - run_id, - node_id, - request_id, - "node_prompt_event", - "dropped", - json!({ - "reason":"request_mismatch", - "event":prompt_event_name(&event), - "active_request_id":current.request.request_id, - }), - ); - continue; - } - - let terminal = event.is_terminal(); - let completion_status = prompt_completion_status(&event); - let completion_detail = prompt_completion_detail(&event); - orch_datastream.emit_prompt( - dashboard, - run_id, - node_id, - request_id, - "node_prompt_event", - "observed", - prompt_event_detail(&event), - ); - let _ = current.events.send(event); - if terminal { - orch_datastream.emit_prompt( - dashboard, - run_id, - node_id, - request_id, - "prompt_complete", - completion_status, - completion_detail, - ); - active = None; - } - } - - thread::sleep(PUMP_INTERVAL); - } -} - -fn prompt_event_name(event: &PromptEvent) -> &'static str { - match event { - PromptEvent::TextDelta { .. } => "TextDelta", - PromptEvent::Done { .. } => "Done", - PromptEvent::Fault { .. } => "Fault", - } -} - -fn prompt_event_detail(event: &PromptEvent) -> Value { - match event { - PromptEvent::TextDelta { text, .. } => { - json!({"event":"TextDelta","terminal":false,"text_bytes":text.len()}) - } - PromptEvent::Done { - final_text, - tokens_generated, - elapsed_ms, - .. - } => json!({ - "event":"Done", - "terminal":true, - "tokens_generated":tokens_generated, - "elapsed_ms":elapsed_ms, - "final_text_bytes":final_text.len(), - }), - PromptEvent::Fault { error, .. } => { - json!({"event":"Fault","terminal":true,"error":error}) - } - } -} - -fn prompt_completion_status(event: &PromptEvent) -> &'static str { - match event { - PromptEvent::Done { .. } => "ready", - PromptEvent::Fault { .. } => "failed", - PromptEvent::TextDelta { .. } => "observed", - } -} - -fn prompt_completion_detail(event: &PromptEvent) -> Value { - match event { - PromptEvent::Done { .. } => json!({"event":"Done"}), - PromptEvent::Fault { error, .. } => json!({"event":"Fault","error":error}), - PromptEvent::TextDelta { .. } => json!({"event":"TextDelta"}), - } -} - -fn stop_requested(stop_rx: &mpsc::Receiver<()>) -> bool { - stop_rx.try_recv().is_ok() -} - -fn spawn_stop_listener() -> mpsc::Receiver<()> { - let (tx, rx) = mpsc::channel(); - thread::spawn(move || { - let stdin = std::io::stdin(); - for line in stdin.lock().lines().map_while(Result::ok) { - let trimmed = line.trim(); - if trimmed.eq_ignore_ascii_case("stop") - || trimmed.eq_ignore_ascii_case("shutdown") - || trimmed.eq_ignore_ascii_case("quit") - { - let _ = tx.send(()); - break; - } - } - }); - rx -} - -fn drain_observations( - obs_rx: &mpsc::Receiver, - dashboard: Option<&DashboardSupport>, - orch_datastream: &mut OrchDatastream, - provider: ProviderKind, -) -> Result<(), String> { - while let Ok(observation) = obs_rx.try_recv() { - emit_plugin_observation(orch_datastream, dashboard, provider, &observation); - match observation { - PluginObservation::Failed { reason, .. } => return Err(reason), - PluginObservation::Exited { status, .. } => { - return Err(format!("node exited: {status:?}")); - } - PluginObservation::DatastreamFrame { .. } => {} - PluginObservation::ProviderLine { .. } - | PluginObservation::StdoutLine { .. } - | PluginObservation::StderrLine { .. } => {} - } - } - Ok(()) -} - -fn emit_plugin_observation( - orch_datastream: &mut OrchDatastream, - dashboard: Option<&DashboardSupport>, - provider: ProviderKind, - observation: &PluginObservation, -) { - match observation { - PluginObservation::StdoutLine { - run_id, - node_id, - line, - } => orch_datastream.emit_log( - dashboard, - ProvisionLogLine { - run_id: *run_id, - node_id: *node_id, - stream: ProvisionLogStream::Stdout, - line: line.clone(), - }, - ), - PluginObservation::StderrLine { - run_id, - node_id, - line, - } => orch_datastream.emit_log( - dashboard, - ProvisionLogLine { - run_id: *run_id, - node_id: *node_id, - stream: ProvisionLogStream::Stderr, - line: line.clone(), - }, - ), - PluginObservation::ProviderLine { - run_id, - node_id, - line, - } => orch_datastream.emit_log( - dashboard, - ProvisionLogLine { - run_id: *run_id, - node_id: *node_id, - stream: ProvisionLogStream::Provider, - line: line.clone(), - }, - ), - PluginObservation::DatastreamFrame { - channel, payload, .. - } => orch_datastream.emit_bytes_from( - dashboard, - channel, - payload.as_bytes().to_vec(), - "node_bootstrap_stdio", - ), - PluginObservation::Exited { - run_id, - node_id, - status, - } => orch_datastream.emit_event( - dashboard, - ProvisionEvent { - run_id: *run_id, - node_id: *node_id, - kind: ProvisionEventKind::NodeStopped, - provider: Some(provider.as_str().to_owned()), - message: Some(format!("node process exited with {status:?}")), - }, - ), - PluginObservation::Failed { - run_id, - node_id, - reason, - } => orch_datastream.emit_event( - dashboard, - ProvisionEvent { - run_id: *run_id, - node_id: *node_id, - kind: ProvisionEventKind::ProvisionFailed, - provider: Some(provider.as_str().to_owned()), - message: Some(reason.clone()), - }, - ), - } -} - -fn drain_frames( - frame_rx: &mpsc::Receiver, - dashboard: Option<&DashboardSupport>, - orch_datastream: &mut OrchDatastream, -) { - while let Ok(collected) = frame_rx.try_recv() { - ingest_dashboard_frame( - dashboard, - &collected.stream, - &collected.channel_name, - &collected.frame, - ); - orch_datastream.archive_frame( - "node_cluster", - &collected.stream, - &collected.channel_name, - &collected.frame, - ); - } -} - -fn ingest_dashboard_frame( - dashboard: Option<&DashboardSupport>, - stream: &StreamId, - channel: &str, - frame: &Frame, -) { - if let Some(dashboard) = dashboard { - dashboard.publish_frame(stream, channel, frame); - } -} - -fn emit_swim_transitions( - orch_datastream: &mut OrchDatastream, - dashboard: Option<&DashboardSupport>, - run_id: u64, - node_id: u64, - stack: &DistributionRuntimeStack, -) { - for transition in stack.drain_swim_transitions() { - orch_datastream.emit_bootstrap_to_channel( - dashboard, - MVP_SWIM_MEMBERSHIP, - run_id, - node_id, - "membership_transition", - "observed", - json!({ - "peer":format!("{:?}", transition.peer), - "from":transition.from.map(|state| format!("{:?}", state)), - "to":format!("{:?}", transition.to), - "reason":transition.reason, - }), - ); - } -} - -fn pump( - driver: &mut IrohDriver, - stack: &DistributionRuntimeStack, - frame_tx: &mpsc::Sender, -) { - stack.tick_protocol_actors(Instant::now()); - driver.pump_inbound_to_actors(); - stack.pump_runtime_once(); - driver.drain_outbox(&stack.outbox); - drain_datastream_connections(driver, frame_tx); -} - -fn env_optional(name: &str) -> Option { - std::env::var(name) - .ok() - .map(|value| value.trim().to_owned()) - .filter(|value| !value.is_empty()) -} - -fn docker_container_prefix() -> String { - env_optional(MVP_DOCKER_CONTAINER_PREFIX_ENV) - .unwrap_or_else(|| DEFAULT_DOCKER_CONTAINER_PREFIX.to_owned()) -} - -fn optional_env(name: &str) -> Option<(String, String)> { - env_optional(name).map(|value| (name.to_owned(), value)) -} - -fn local_tinygrad_worker_env(provider: ProviderKind) -> Option<(String, String)> { - optional_env("MVP_TINYGRAD_WORKER").or_else(|| { - if provider != ProviderKind::Process { - return None; - } - default_local_tinygrad_worker_path().map(|path| { - ( - "MVP_TINYGRAD_WORKER".to_owned(), - path.to_string_lossy().to_string(), - ) - }) - }) -} - -fn default_local_tinygrad_worker_path() -> Option { - let mut candidates = Vec::new(); - if let Ok(cwd) = std::env::current_dir() { - candidates.push(cwd.join("apps").join("mvp-node").join("tinygrad_worker.py")); - } - candidates.push( - PathBuf::from(env!("CARGO_MANIFEST_DIR")) - .join("..") - .join("..") - .join("apps") - .join("mvp-node") - .join("tinygrad_worker.py"), - ); - for candidate in candidates { - if candidate.is_file() { - return Some(candidate.canonicalize().unwrap_or(candidate)); - } - } - None -} - -fn resolve_vastai_ssh_identity(explicit: Option) -> Result { - match explicit { - Some(path) => Ok(path), - None => { - let home = std::env::var_os("HOME") - .filter(|value| !value.is_empty()) - .ok_or_else(|| { - "MVP_VASTAI_SSH_IDENTITY is required because HOME is unset".to_owned() - })?; - Ok(PathBuf::from(home).join(".ssh").join("id_ed25519")) - } - } -} - -fn expand_home_path(value: &str) -> Result { - let trimmed = value.trim(); - if let Some(rest) = trimmed.strip_prefix("~/") { - let home = std::env::var_os("HOME") - .filter(|value| !value.is_empty()) - .ok_or_else(|| "MVP_VASTAI_SSH_IDENTITY uses ~/ but HOME is unset".to_owned())?; - return Ok(PathBuf::from(home).join(rest)); - } - Ok(PathBuf::from(trimmed)) -} - -fn derive_ssh_public_key(identity: &Path) -> Result { - let output = Command::new("ssh-keygen") - .arg("-y") - .arg("-f") - .arg(identity) - .output() - .map_err(|e| { - format!( - "derive VastAI SSH public key from {}: {e}", - identity.display() - ) - })?; - let public_key = String::from_utf8_lossy(&output.stdout) - .trim_end_matches(['\r', '\n']) - .to_owned(); - if !output.status.success() || public_key.trim().is_empty() { - return Err(format!( - "derive VastAI SSH public key from {}: {}", - identity.display(), - command_output_failure_detail(&output, None) - )); - } - Ok(public_key) -} - -fn ssh_public_key_fingerprint(public_key: &str) -> String { - let path = std::env::temp_dir().join(format!("mvp-vastai-ssh-key-{}.pub", std::process::id())); - if std::fs::write(&path, format!("{public_key}\n")).is_err() { - return "unavailable".to_owned(); - } - let output = Command::new("ssh-keygen") - .arg("-l") - .arg("-f") - .arg(&path) - .output(); - let _ = std::fs::remove_file(&path); - let Ok(output) = output else { - return "unavailable".to_owned(); - }; - if !output.status.success() { - return "unavailable".to_owned(); - } - let stdout = String::from_utf8_lossy(&output.stdout); - let mut fields = stdout.split_whitespace(); - match (fields.next(), fields.next()) { - (Some(bits), Some(fingerprint)) => format!("{bits} {fingerprint}"), - _ => "unavailable".to_owned(), - } -} - -fn vastai_account_has_ssh_key(api_key: &str, public_key: &str) -> Result { - let output = Command::new("vastai") - .args(["show", "ssh-keys", "--raw", "--api-key", api_key]) - .output() - .map_err(vastai_cli_error)?; - if !output.status.success() { - return Err(format!( - "vastai show ssh-keys failed: {}", - command_output_failure_detail(&output, Some(api_key)) - )); - } - let stdout = String::from_utf8_lossy(&output.stdout); - Ok(account_ssh_keys_output_contains_public_key( - &stdout, public_key, - )) -} - -fn ensure_vastai_account_ssh_key(api_key: &str, public_key: &str) -> Result<(), String> { - if vastai_account_has_ssh_key(api_key, public_key)? { - return Ok(()); - } - - let output = Command::new("vastai") - .args(["create", "ssh-key"]) - .arg(public_key) - .args(["-y", "--api-key", api_key]) - .output() - .map_err(vastai_cli_error)?; - if !output.status.success() { - return Err(format!( - "vastai create ssh-key failed: {}", - command_output_failure_detail(&output, Some(api_key)) - )); - } - - if vastai_account_has_ssh_key(api_key, public_key)? { - Ok(()) - } else { - Err( - "VastAI SSH key registration did not make the selected key visible in vastai show ssh-keys" - .to_owned(), - ) - } -} - -fn account_ssh_keys_output_contains_public_key(output: &str, public_key: &str) -> bool { - let public_key = public_key.trim(); - if public_key.is_empty() { - return false; - } - if output.contains(public_key) { - return true; - } - public_key - .split_whitespace() - .nth(1) - .is_some_and(|body| !body.is_empty() && output.contains(body)) -} - -fn vastai_cli_error(error: std::io::Error) -> String { - if error.kind() == std::io::ErrorKind::NotFound { - "vastai CLI is required to verify/register MVP_VASTAI_SSH_IDENTITY; install with pip install vastai" - .to_owned() - } else { - format!("run vastai CLI: {error}") - } -} - -fn command_output_failure_detail(output: &std::process::Output, secret: Option<&str>) -> String { - let mut detail = String::from_utf8_lossy(&output.stderr).trim().to_owned(); - if detail.is_empty() { - detail = output.status.to_string(); - } - if let Some(secret) = secret.filter(|secret| !secret.is_empty()) { - detail = detail.replace(secret, ""); - } - detail -} - -#[cfg(test)] -fn relay_mode_from_env() -> Result { - relay_runtime_config_from_env(1).map(|relay| relay.mode) -} - -fn next_arg(args: &mut impl Iterator, name: &str) -> Result { - args.next() - .ok_or_else(|| format!("missing value after {name}")) -} - -fn parse_next(args: &mut impl Iterator, name: &str) -> Result -where - T: std::str::FromStr, - T::Err: std::fmt::Display, -{ - let value = next_arg(args, name)?; - value - .parse::() - .map_err(|e| format!("invalid {name}={value:?}: {e}")) -} - -#[cfg(test)] -mod tests { - use super::*; - use std::{ffi::OsString, path::PathBuf}; - - static ENV_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(()); - - const ENV_KEYS: &[&str] = &[ - CACHED_MODEL_HOST_ENV, - "HOME", - "HF_TOKEN", - "MVP_CPU_LINE_PROFILE", - "MVP_CPU_LINE_PROFILE_INTERVAL_MS", - "CUDA_DEVICE_SCHEDULE", - "MVP_DASHBOARD", - "MVP_DOCKER_GPUS", - "MVP_GGUF_FILE", - "MVP_GGUF_LOCAL_PATH", - "MVP_GGUF_REPO", - "MVP_GGUF_REVISION", - "MVP_IROH_RELAY_MODE", - MVP_IROH_RELAY_URL_ENV, - "MVP_LAYER_END_EXCLUSIVE", - "MVP_LOGICAL_NODE_ID", - "MVP_MODEL_CACHE_DIR", - "MVP_MAX_CONTEXT", - "MVP_MODEL_ID", - "MVP_NODE_IMAGE", - "MVP_NODE_PROVIDER", - "MVP_PROVIDER", - "MVP_PROMPT_MAX_TOKENS", - "MVP_PROMPT_RPC_BIND", - "MVP_RUN_ID", - "MVP_RUNTIME_CONFIG", - "MVP_PIPELINE_STAGES", - "MVP_STAGE_INDEX", - "MVP_TOKEN_PROGRESS_EVERY", - "MVP_TINYGRAD_WORKER", - "MVP_TOKENIZER_LOCAL_PATH", - "MVP_VASTAI_API_KEY", - "MVP_VASTAI_BOOTSTRAP_COMMAND", - "MVP_VASTAI_CONFIRM_LEASE", - "MVP_VASTAI_DISK_GB", - "MVP_VASTAI_GPU_NAME", - "MVP_VASTAI_MIN_DOWN_MBPS", - "MVP_VASTAI_MIN_GPU_RAM_MB", - "MVP_VASTAI_MIN_RELIABILITY", - "MVP_VASTAI_MIN_UP_MBPS", - "MVP_VASTAI_ONSTART", - "MVP_VASTAI_POLL_INTERVAL_SECS", - "MVP_VASTAI_REQUIRE_VERIFIED", - "MVP_VASTAI_SSH_USER", - "MVP_VASTAI_SSH_IDENTITY", - "VASTAI_API_KEY", - SWACTOR_IROH_RELAY_URL_ENV, - MVP_DOCKER_CONTAINER_PREFIX_ENV, - MVP_WORKER_BIN_ENV, - ]; - - struct RestoreEnv { - saved: Vec<(&'static str, Option)>, - } - - impl Drop for RestoreEnv { - fn drop(&mut self) { - for (key, value) in &self.saved { - match value { - Some(value) => unsafe { std::env::set_var(key, value) }, - None => unsafe { std::env::remove_var(key) }, - } - } - } - } - - fn with_clean_env(settings: &[(&'static str, &'static str)], test: impl FnOnce() -> T) -> T { - let settings = settings - .iter() - .map(|(key, value)| (*key, OsString::from(value))) - .collect::>(); - with_clean_env_os(&settings, test) - } - - fn with_clean_env_os(settings: &[(&'static str, OsString)], test: impl FnOnce() -> T) -> T { - let _lock = ENV_LOCK - .lock() - .unwrap_or_else(|poisoned| poisoned.into_inner()); - let saved = ENV_KEYS - .iter() - .map(|&key| (key, std::env::var_os(key))) - .collect::>(); - for key in ENV_KEYS { - unsafe { std::env::remove_var(key) }; - } - for (key, value) in settings { - assert!( - ENV_KEYS.contains(key), - "test env key {key} must be restored" - ); - unsafe { std::env::set_var(key, value) }; - } - let _restore = RestoreEnv { saved }; - test() - } - - fn selected_provider(settings: &[(&'static str, &'static str)]) -> ProviderKind { - with_clean_env(settings, || { - Config::from_layers_with_path_and_args(None, std::iter::empty::()) - .expect("config parses") - .provider - }) - } - - fn canonical_relay_url(raw: &str) -> String { - raw.parse::() - .expect("fixture relay URL parses") - .to_string() - } - - fn node_spec_env(settings: &[(&'static str, &'static str)]) -> Vec<(String, String)> { - with_clean_env(settings, || { - let config = Config::from_layers_with_path_and_args(None, std::iter::empty::()) - .expect("config parses"); - let coordinator = EndpointAddr::new(iroh::SecretKey::from_bytes(&[9; 32]).public()); - let orchestrator_actor = ActorAddress([12; 32]); - config - .node_spec(coordinator, orchestrator_actor) - .expect("node spec builds") - .env - }) - } - - fn env_value<'a>(env: &'a [(String, String)], key: &str) -> Option<&'a str> { - env.iter() - .find(|(env_key, _)| env_key == key) - .map(|(_, value)| value.as_str()) - } - - fn cached_model_config_with_args(model: &TempModelFile, args: &[&str]) -> Config { - with_clean_env_os( - &[ - ("MVP_RUNTIME_CONFIG", OsString::from("local")), - ("MVP_NODE_PROVIDER", OsString::from("docker")), - (CACHED_MODEL_HOST_ENV, model.raw_path.as_os_str().to_owned()), - ], - || { - Config::from_layers_with_path_and_args( - None, - args.iter().copied().map(str::to_owned), - ) - .expect("cached-model config parses") - }, - ) - } - - fn pipeline_config_with_cached_model(model: &TempModelFile, pipeline_stages: u32) -> Config { - let stages_arg = pipeline_stages.to_string(); - cached_model_config_with_args( - model, - &[ - "--pipeline-stages", - stages_arg.as_str(), - "--max-context", - "512", - ], - ) - } - - fn expected_layer_ranges(num_layers: u32, stage_count: u32) -> Vec<(u32, u32)> { - (0..stage_count) - .map(|stage_index| { - let start = (u64::from(num_layers) * u64::from(stage_index) - / u64::from(stage_count)) as u32; - let end = (u64::from(num_layers) * u64::from(stage_index + 1) - / u64::from(stage_count)) as u32; - (start, end) - }) - .collect() - } - #[test] - fn pipeline_weight_load_scheduler_keeps_one_active_stage() { - let model = TempModelFile::with_metadata( - "seven-stage-scheduler.gguf", - TestGgufMetadata { - num_layers: 30, - ..TestGgufMetadata::default() - }, - ); - let config = pipeline_config_with_cached_model(&model, 7); - let plan = config.build_run_plan().expect("seven-stage plan builds"); - let mut loaded = BTreeSet::new(); - - let first = - next_pipeline_weight_load_stage(&plan, &loaded, None).expect("first stage selected"); - assert_eq!(first.stage_index, 0); - - let resent = next_pipeline_weight_load_stage(&plan, &loaded, Some(first.stage_index)) - .expect("active stage is resent before it loads"); - assert_eq!(resent.stage_index, 0); - - for expected_stage in 0..7 { - let active = next_pipeline_weight_load_stage(&plan, &loaded, None) - .expect("next unloaded stage selected"); - assert_eq!(active.stage_index, expected_stage); - loaded.insert(expected_stage); - } - - assert!( - next_pipeline_weight_load_stage(&plan, &loaded, None).is_none(), - "all stages loaded should leave no active load" - ); - } - - fn assert_plan_matches_metadata( - config: &Config, - plan: &run_plan::RunPlan, - metadata: TestGgufMetadata, - stage_count: u32, - requested_context: u64, - ) { - assert_eq!(plan.run_id, run_plan::RunId(config.run_id)); - assert_eq!(plan.model.model_id, config.model_id); - assert_eq!(plan.model.gguf_source, config.gguf_source); - assert_eq!(plan.model.num_layers, metadata.num_layers); - assert_eq!(plan.model.hidden_dim, metadata.hidden_dim as u32); - assert_eq!( - plan.model.max_seq_len, - requested_context.min(metadata.context_length) as u32 - ); - assert_eq!(plan.model.eos_token_id, metadata.eos_token_id); - assert_eq!(plan.stages.len(), stage_count as usize); - assert_eq!(plan.edges.len(), stage_count as usize + 1); - assert_eq!(plan.max_tokens, config.default_max_tokens); - - let mut stages = plan.stages.clone(); - stages.sort_by_key(|stage| stage.stage_index); - let expected_ranges = expected_layer_ranges(metadata.num_layers, stage_count); - for (stage, expected_range) in stages.iter().zip(expected_ranges) { - assert_eq!(stage.stage_count, stage_count); - assert_eq!( - stage.node_id, - run_plan::NodeId(config.node_id + 1 + u64::from(stage.stage_index)) - ); - assert_eq!( - (stage.layer_start, stage.layer_end_exclusive), - expected_range - ); - assert!( - stage.layer_start < stage.layer_end_exclusive, - "stage {} must have a non-empty layer range", - stage.stage_index - ); - } - } - - fn token_object_spec(max_extent: u64, alignment: u32) -> run_plan::ObjectSpec { - run_plan::ObjectSpec { - kind: run_plan::ObjectKind::Token, - max_extent, - dtype_family: run_plan::DTypeFamily::BFloat, - dtype_width_bytes: 4, - shape: run_plan::ShapeRule::TokenIds, - layout: run_plan::LayoutRule::Contiguous, - alignment, - sequence_policy: run_plan::SequencePolicy::Ordered, - } - } - - fn decode_token_record_payload( - bytes: &[u8], - spec: run_plan::ObjectSpec, - ) -> (u64, Vec, bool) { - let read = ingress::read_object_record(bytes, ingress_object_spec_from_plan(spec), false) - .expect("token record should be valid"); - let ingress::ObjectRecordRead::Complete(record) = read else { - panic!("token record should be complete"); - }; - let payload = record.payload(bytes).expect("token payload is present"); - let tokens = payload - .chunks_exact(4) - .map(|chunk| u32::from_le_bytes(chunk.try_into().unwrap())) - .collect::>(); - assert!(payload.chunks_exact(4).remainder().is_empty()); - (record.sequence, tokens, record.flags.end_of_sequence) - } - - struct PipelineRuntimeTestFixture { - actor_runtime: Arc, - tokenizer_events: swactor::runtime::Inbox, - encode_requests: swactor::runtime::Inbox, - decode_requests: swactor::runtime::Inbox, - runtime: PipelinePromptRuntime, - token_in_rx: tokio_mpsc::UnboundedReceiver>, - } - - fn pipeline_runtime_fixture() -> PipelineRuntimeTestFixture { - let spec = token_object_spec(64, 64); - let (token_in_tx, token_in_rx) = tokio_mpsc::unbounded_channel(); - let (recv_tx, recv_rx) = mpsc::channel(); - pipeline_runtime_fixture_from_specs( - 11, - 12, - spec, - spec, - token_in_tx, - token_in_rx, - recv_tx, - recv_rx, - ) - } - - fn pipeline_runtime_fixture_from_plan(stage_count: u32) -> PipelineRuntimeTestFixture { - let model = TempModelFile::new(&format!("prompt-plan-{stage_count}.gguf")); - let config = pipeline_config_with_cached_model(&model, stage_count); - let plan = config.build_run_plan().expect("prompt runtime plan builds"); - let token_in = plan - .edges - .iter() - .find(|edge| edge.kind == run_plan::EdgeKind::TokenIn) - .expect("token-in edge"); - let token_out = plan - .edges - .iter() - .find(|edge| edge.kind == run_plan::EdgeKind::TokenOut) - .expect("token-out edge"); - let (token_in_tx, token_in_rx) = tokio_mpsc::unbounded_channel(); - let (recv_tx, recv_rx) = mpsc::channel(); - pipeline_runtime_fixture_from_specs( - token_in.edge_id.0, - token_out.edge_id.0, - token_in.object_spec, - token_out.object_spec, - token_in_tx, - token_in_rx, - recv_tx, - recv_rx, - ) - } - - #[allow(clippy::too_many_arguments)] - fn pipeline_runtime_fixture_from_specs( - token_in_edge_id: u64, - token_out_edge_id: u64, - token_spec: run_plan::ObjectSpec, - token_out_spec: run_plan::ObjectSpec, - token_in_tx: tokio_mpsc::UnboundedSender>, - token_in_rx: tokio_mpsc::UnboundedReceiver>, - recv_tx: mpsc::Sender>, - recv_rx: mpsc::Receiver>, - ) -> PipelineRuntimeTestFixture { - let actor_runtime = Arc::new(swactor::runtime::Runtime::new( - swactor::config::RuntimeConfig::default(), - )); - let encode_requests = actor_runtime - .new_inbox::() - .expect("encode request inbox"); - let decode_requests = actor_runtime - .new_inbox::() - .expect("decode request inbox"); - let tokenizer_events = actor_runtime - .new_inbox::() - .expect("tokenizer event inbox"); - PipelineRuntimeTestFixture { - runtime: PipelinePromptRuntime { - token_in_edge_id, - token_out_edge_id, - token_spec, - token_out_spec, - token_in_sender: PipelineSendHandle::Channel(token_in_tx), - recv_rx, - recv_tx, - recv_buffer: Vec::new(), - tokenizer_encode_actor: *encode_requests.addr(), - tokenizer_decode_actor: *decode_requests.addr(), - tokenizer_reply_to: *tokenizer_events.addr(), - pending_encode: None, - pending_decode: None, - next_sequence: 0, - generated_tokens: Vec::new(), - final_text: String::new(), - active: None, - started_at: None, - }, - actor_runtime, - tokenizer_events, - encode_requests, - decode_requests, - token_in_rx, - } - } - - fn start_fixture_prompt( - fixture: &mut PipelineRuntimeTestFixture, - request: SubmitPrompt, - events: mpsc::Sender, - datastream: &mut OrchDatastream, - run_id: u64, - node_id: u64, - ) { - fixture - .runtime - .start_prompt( - request, - events, - &fixture.actor_runtime, - None, - datastream, - run_id, - node_id, - ) - .expect("pipeline prompt starts"); - } - - fn drain_fixture_tokenizer( - fixture: &mut PipelineRuntimeTestFixture, - datastream: &mut OrchDatastream, - run_id: u64, - node_id: u64, - ) { - fixture - .runtime - .drain_tokenizer_events( - &fixture.actor_runtime, - &fixture.tokenizer_events, - None, - datastream, - run_id, - node_id, - ) - .expect("tokenizer events drain"); - } - - fn send_encoded_tokens( - fixture: &PipelineRuntimeTestFixture, - request_id: u64, - tokens: Vec, - ) { - fixture - .actor_runtime - .send_to( - *fixture.tokenizer_events.addr(), - TokenizerEvent::PromptEncoded { request_id, tokens }, - ) - .expect("send encoded tokens"); - } - - fn send_decoded_text(fixture: &PipelineRuntimeTestFixture, request_id: u64, text: &str) { - fixture - .actor_runtime - .send_to( - *fixture.tokenizer_events.addr(), - TokenizerEvent::TokensDecoded { - request_id, - text: text.to_owned(), - }, - ) - .expect("send decoded text"); - } - - fn assert_encode_request(fixture: &PipelineRuntimeTestFixture, request_id: u64, prompt: &str) { - match fixture.encode_requests.try_recv() { - Some(NodeAgentMsg::EncodePrompt { - request_id: actual_request_id, - prompt: actual_prompt, - reply_to, - }) => { - assert_eq!(actual_request_id, request_id); - assert_eq!(actual_prompt, prompt); - assert_eq!(reply_to, *fixture.tokenizer_events.addr()); - } - other => panic!("expected EncodePrompt request, got {other:?}"), - } - } - - fn assert_decode_request(fixture: &PipelineRuntimeTestFixture, request_id: u64, token: u32) { - match fixture.decode_requests.try_recv() { - Some(NodeAgentMsg::DecodeTokens { - request_id: actual_request_id, - tokens, - reply_to, - }) => { - assert_eq!(actual_request_id, request_id); - assert_eq!(tokens, vec![token]); - assert_eq!(reply_to, *fixture.tokenizer_events.addr()); - } - other => panic!("expected DecodeTokens request, got {other:?}"), - } - } - - #[test] - fn pipeline_prompt_token_record_round_trips_one_token_with_eos_and_plan_ring_alignment() { - let spec = token_object_spec(64, 64); - let bytes = encode_token_record(spec, 12, 7, &[513], true).expect("token record encodes"); - assert_eq!(&bytes[0..4], b"MO01"); - - let mut partial = bytes[..run_plan::MO01_HEADER_BYTES as usize + 2].to_vec(); - assert!( - take_pipeline_token_record(&mut partial, spec) - .expect("partial token record is not malformed") - .is_none() - ); - assert_eq!(partial.len(), run_plan::MO01_HEADER_BYTES as usize + 2); - - let mut buffer = bytes; - let record = take_pipeline_token_record(&mut buffer, spec) - .expect("one-u32 token-out payload should parse despite ring alignment") - .expect("complete token-out record is available"); - - assert_eq!(record.object_id, 9007); - assert_eq!(record.sequence, 7); - assert_eq!(record.token_id, 513); - assert!(record.eos); - assert!(buffer.is_empty()); - } - - #[test] - fn plan_derived_mo01_specs_accept_token_and_activation_record_extents() { - let metadata = TestGgufMetadata { - num_layers: 7, - hidden_dim: 13, - context_length: 64, - eos_token_id: 11, - }; - let model = TempModelFile::with_metadata("plan-mo01.gguf", metadata); - let config = pipeline_config_with_cached_model(&model, 3); - let plan = config - .build_run_plan() - .expect("plan-derived MO01 plan builds"); - let token_in = plan - .edges - .iter() - .find(|edge| edge.kind == run_plan::EdgeKind::TokenIn) - .expect("token-in edge"); - let activation = plan - .edges - .iter() - .find(|edge| edge.kind == run_plan::EdgeKind::Activation) - .expect("activation edge"); - - let token_record = encode_token_record( - token_in.object_spec, - token_in.edge_id.0, - 0, - &[65, 195, 169], - false, - ) - .expect("plan token record encodes"); - assert_eq!( - decode_token_record_payload(&token_record, token_in.object_spec), - (0, vec![65, 195, 169], false) - ); - - let activation_payload = vec![ - 0_u8; - usize::try_from(metadata.hidden_dim * 2) - .expect("activation payload size fits usize") - ]; - let activation_record = ingress::ObjectRecordBuilder::new(ingress_object_spec_from_plan( - activation.object_spec, - )) - .object_id(ingress::ObjectId(9100)) - .sequence(0) - .payload(activation_payload) - .encode(); - let read = ingress::read_object_record( - &activation_record, - ingress_object_spec_from_plan(activation.object_spec), - false, - ) - .expect("plan activation record parses"); - let ingress::ObjectRecordRead::Complete(record) = read else { - panic!("activation record should be complete"); - }; - assert_eq!(record.payload(&activation_record).unwrap().len(), 26); - } - - #[test] - fn pipeline_prompt_token_out_parser_rejects_payloads_that_are_not_one_u32() { - let spec = token_object_spec(64, 64); - let mut buffer = - encode_token_record(spec, 12, 0, &[65, 66], false).expect("multi-token record encodes"); - - let error = take_pipeline_token_record(&mut buffer, spec) - .expect_err("token-out records must carry exactly one generated token"); - - assert!( - error.contains("token-out payload must be exactly one u32, got 8"), - "unexpected error: {error}" - ); - } - - #[test] - fn pipeline_prompt_runtime_uses_tokenizer_events_for_encode_decode_and_continuations() { - let mut fixture = pipeline_runtime_fixture(); - let (event_tx, event_rx) = mpsc::channel(); - let mut datastream = OrchDatastream::new(91, None).expect("datastream opens"); - start_fixture_prompt( - &mut fixture, - SubmitPrompt { - request_id: 42, - prompt_text: "Hi".to_owned(), - max_tokens: 2, - }, - event_tx, - &mut datastream, - 91, - 3, - ); - assert_encode_request(&fixture, 42, "Hi"); - send_encoded_tokens(&fixture, 42, vec![1001, 1002]); - drain_fixture_tokenizer(&mut fixture, &mut datastream, 91, 3); - - let initial = fixture - .token_in_rx - .try_recv() - .expect("tokenizer tokens are sent to token-in"); - assert_eq!( - decode_token_record_payload(&initial, fixture.runtime.token_spec), - (0, vec![1001, 1002], false) - ); - assert!(fixture.token_in_rx.try_recv().is_err()); - - fixture - .runtime - .recv_tx - .send( - encode_token_record( - fixture.runtime.token_out_spec, - fixture.runtime.token_out_edge_id, - 0, - &[79], - false, - ) - .expect("first token-out record encodes"), - ) - .expect("token-out bytes enqueue"); - fixture - .runtime - .drain_tokens(&fixture.actor_runtime, None, &mut datastream, 91, 3) - .expect("first token drains"); - assert_decode_request(&fixture, 42, 79); - assert!(event_rx.try_recv().is_err()); - - send_decoded_text(&fixture, 42, "O"); - drain_fixture_tokenizer(&mut fixture, &mut datastream, 91, 3); - assert_eq!( - event_rx.try_recv().expect("first delta event"), - PromptEvent::TextDelta { - request_id: 42, - text: "O".to_owned(), - } - ); - let continuation = fixture - .token_in_rx - .try_recv() - .expect("non-terminal token is fed back to token-in"); - assert_eq!( - decode_token_record_payload(&continuation, fixture.runtime.token_spec), - (1, vec![79], false) - ); - - fixture - .runtime - .recv_tx - .send( - encode_token_record( - fixture.runtime.token_out_spec, - fixture.runtime.token_out_edge_id, - 1, - &[75], - false, - ) - .expect("second token-out record encodes"), - ) - .expect("token-out bytes enqueue"); - fixture - .runtime - .drain_tokens(&fixture.actor_runtime, None, &mut datastream, 91, 3) - .expect("second token drains"); - assert_decode_request(&fixture, 42, 75); - - send_decoded_text(&fixture, 42, "K"); - drain_fixture_tokenizer(&mut fixture, &mut datastream, 91, 3); - assert_eq!( - event_rx.try_recv().expect("second delta event"), - PromptEvent::TextDelta { - request_id: 42, - text: "K".to_owned(), - } - ); - match event_rx.try_recv().expect("done event at max tokens") { - PromptEvent::Done { - request_id, - final_text, - tokens_generated, - .. - } => { - assert_eq!(request_id, 42); - assert_eq!(final_text, "OK"); - assert_eq!(tokens_generated, 2); - } - event => panic!("expected Done at max tokens, got {event:?}"), - } - assert!(event_rx.try_recv().is_err()); - assert!(fixture.token_in_rx.try_recv().is_err()); - assert!(!fixture.runtime.is_active()); - } - - #[test] - fn prompt_runtime_uses_plan_derived_token_edges_for_cached_n_one() { - let mut fixture = pipeline_runtime_fixture_from_plan(1); - let (event_tx, event_rx) = mpsc::channel(); - let mut datastream = OrchDatastream::new(93, None).expect("datastream opens"); - start_fixture_prompt( - &mut fixture, - SubmitPrompt { - request_id: 88, - prompt_text: "Aé".to_owned(), - max_tokens: 1, - }, - event_tx, - &mut datastream, - 93, - 3, - ); - assert_encode_request(&fixture, 88, "Aé"); - send_encoded_tokens(&fixture, 88, vec![321, 654]); - drain_fixture_tokenizer(&mut fixture, &mut datastream, 93, 3); - - let initial = fixture - .token_in_rx - .try_recv() - .expect("plan-derived token-in record is emitted"); - assert_eq!( - decode_token_record_payload(&initial, fixture.runtime.token_spec), - (0, vec![321, 654], false) - ); - assert_eq!(fixture.runtime.token_in_edge_id, 1); - assert_eq!(fixture.runtime.token_out_edge_id, 2); - - fixture - .runtime - .recv_tx - .send( - encode_token_record( - fixture.runtime.token_out_spec, - fixture.runtime.token_out_edge_id, - 0, - &[33], - false, - ) - .expect("plan-derived token-out record encodes"), - ) - .expect("token-out bytes enqueue"); - fixture - .runtime - .drain_tokens(&fixture.actor_runtime, None, &mut datastream, 93, 3) - .expect("plan-derived token drains"); - assert_decode_request(&fixture, 88, 33); - send_decoded_text(&fixture, 88, "!"); - drain_fixture_tokenizer(&mut fixture, &mut datastream, 93, 3); - - assert_eq!( - event_rx.try_recv().expect("delta event"), - PromptEvent::TextDelta { - request_id: 88, - text: "!".to_owned(), - } - ); - match event_rx.try_recv().expect("done event at max tokens") { - PromptEvent::Done { - request_id, - final_text, - tokens_generated, - .. - } => { - assert_eq!(request_id, 88); - assert_eq!(final_text, "!"); - assert_eq!(tokens_generated, 1); - } - event => panic!("expected Done at max tokens, got {event:?}"), - } - assert!(fixture.token_in_rx.try_recv().is_err()); - } - - #[test] - fn pipeline_prompt_runtime_finishes_on_eos_without_feedback_token() { - let mut fixture = pipeline_runtime_fixture(); - let (event_tx, event_rx) = mpsc::channel(); - let mut datastream = OrchDatastream::new(92, None).expect("datastream opens"); - start_fixture_prompt( - &mut fixture, - SubmitPrompt { - request_id: 77, - prompt_text: "go".to_owned(), - max_tokens: 8, - }, - event_tx, - &mut datastream, - 92, - 4, - ); - assert_encode_request(&fixture, 77, "go"); - send_encoded_tokens(&fixture, 77, vec![700]); - drain_fixture_tokenizer(&mut fixture, &mut datastream, 92, 4); - fixture - .token_in_rx - .try_recv() - .expect("initial prompt token-in record is sent"); - - fixture - .runtime - .recv_tx - .send( - encode_token_record( - fixture.runtime.token_out_spec, - fixture.runtime.token_out_edge_id, - 0, - &[33], - true, - ) - .expect("eos token-out record encodes"), - ) - .expect("token-out bytes enqueue"); - fixture - .runtime - .drain_tokens(&fixture.actor_runtime, None, &mut datastream, 92, 4) - .expect("eos token drains"); - assert_decode_request(&fixture, 77, 33); - send_decoded_text(&fixture, 77, "!"); - drain_fixture_tokenizer(&mut fixture, &mut datastream, 92, 4); - - assert_eq!( - event_rx.try_recv().expect("delta event before eos done"), - PromptEvent::TextDelta { - request_id: 77, - text: "!".to_owned(), - } - ); - match event_rx.try_recv().expect("done event on eos") { - PromptEvent::Done { - request_id, - final_text, - tokens_generated, - .. - } => { - assert_eq!(request_id, 77); - assert_eq!(final_text, "!"); - assert_eq!(tokens_generated, 1); - } - event => panic!("expected Done on eos, got {event:?}"), - } - assert!(event_rx.try_recv().is_err()); - assert!(fixture.token_in_rx.try_recv().is_err()); - assert!(!fixture.runtime.is_active()); - } - - #[test] - fn planned_prompt_runtime_uses_pipeline_edges_for_single_or_multi_stage_cached_models() { - assert_eq!( - prompt_runtime_mode(None), - PromptRuntimeMode::DirectInferPrompt - ); - - for stage_count in [1_u32, 3] { - let model = TempModelFile::new(&format!("prompt-runtime-{stage_count}.gguf")); - let config = pipeline_config_with_cached_model(&model, stage_count); - let plan = config - .build_run_plan() - .expect("cached-model planned prompt runtime plan builds"); - - assert_eq!( - prompt_runtime_mode(Some(&plan)), - PromptRuntimeMode::PipelineTokenEdges, - "cached-model N={stage_count} must use the planned token-edge runtime" - ); - } - } - - #[test] - fn docker_container_prefix_defaults_and_trims_env_override() { - with_clean_env(&[], || { - assert_eq!(docker_container_prefix(), DEFAULT_DOCKER_CONTAINER_PREFIX); - }); - with_clean_env( - &[(MVP_DOCKER_CONTAINER_PREFIX_ENV, " custom-prefix ")], - || { - assert_eq!(docker_container_prefix(), "custom-prefix"); - }, - ); - } - - #[test] - fn expand_home_path_expands_leading_home_segment() { - with_clean_env(&[("HOME", "/tmp/mvp-vastai-home")], || { - assert_eq!( - expand_home_path("~/keys/deploy").expect("home path expands"), - PathBuf::from("/tmp/mvp-vastai-home/keys/deploy") - ); - assert_eq!( - expand_home_path("/tmp/not-~/expanded").expect("literal path stays literal"), - PathBuf::from("/tmp/not-~/expanded") - ); - }); - } - - #[test] - fn account_ssh_keys_output_contains_public_key_matches_exact_key_material() { - let public_key = "ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAITestKeyBody vastai"; - - assert!(account_ssh_keys_output_contains_public_key( - public_key, public_key - )); - assert!(account_ssh_keys_output_contains_public_key( - r#"{"keys":[{"public_key":"AAAAC3NzaC1lZDI1NTE5AAAAITestKeyBody"}]}"#, - public_key - )); - assert!(!account_ssh_keys_output_contains_public_key( - r#"{"keys":[{"public_key":"AAAAC3NzaC1lZDI1NTE5AAAADifferent"}]}"#, - public_key - )); - } - - #[test] - fn vastai_config_reads_ssh_identity_without_runtime_preparation() { - let config = with_clean_env( - &[ - ("MVP_RUNTIME_CONFIG", "deploy"), - ("MVP_NODE_PROVIDER", "vastai"), - ("MVP_VASTAI_API_KEY", "vast-key"), - ("MVP_VASTAI_BOOTSTRAP_COMMAND", "/usr/local/bin/mvp-node"), - ("MVP_VASTAI_SSH_IDENTITY", "/tmp/mvp-vastai-key"), - ], - || { - Config::from_layers_with_path_and_args(None, std::iter::empty::()) - .expect("vastai config parses without ssh-keygen or vastai CLI") - }, - ); - let vastai = config.vastai.expect("vastai config is present"); - assert_eq!( - vastai.ssh_identity.as_deref(), - Some(Path::new("/tmp/mvp-vastai-key")) - ); - assert_eq!(vastai.ssh_public_key, None); - assert_eq!(vastai.ssh_public_fingerprint, None); - } - - #[derive(Clone, Copy)] - struct TestGgufMetadata { - num_layers: u32, - hidden_dim: u64, - context_length: u64, - eos_token_id: u32, - } - - impl Default for TestGgufMetadata { - fn default() -> Self { - Self { - num_layers: 7, - hidden_dim: 13, - context_length: 64, - eos_token_id: 11, - } - } - } - - struct TempModelFile { - root: PathBuf, - raw_path: PathBuf, - canonical_path: PathBuf, - metadata: TestGgufMetadata, - } - - impl TempModelFile { - fn new(file_name: &str) -> Self { - Self::with_metadata(file_name, TestGgufMetadata::default()) - } - - fn with_metadata(file_name: &str, metadata: TestGgufMetadata) -> Self { - let root = std::env::temp_dir().join(format!( - "mvp-cached-model-test-{}-{}", - std::process::id(), - std::thread::current().name().unwrap_or("unnamed") - )); - let _ = std::fs::remove_dir_all(&root); - std::fs::create_dir_all(root.join("nested")).expect("create temp model dir"); - let canonical_path = root.join(file_name); - std::fs::write(&canonical_path, minimal_gguf(metadata)) - .expect("write temp GGUF metadata file"); - let raw_path = root.join("nested").join("..").join(file_name); - Self { - root, - raw_path, - canonical_path: canonical_path - .canonicalize() - .expect("canonicalize temp model file"), - metadata, - } - } - } - - fn minimal_gguf(metadata: TestGgufMetadata) -> Vec { - let mut bytes = Vec::new(); - bytes.extend_from_slice(b"GGUF"); - bytes.extend_from_slice(&3_u32.to_le_bytes()); - bytes.extend_from_slice(&0_u64.to_le_bytes()); - bytes.extend_from_slice(&6_u64.to_le_bytes()); - push_string_kv(&mut bytes, "general.architecture", "llama"); - push_string_kv(&mut bytes, "general.name", "fixture"); - push_u32_kv(&mut bytes, "llama.block_count", metadata.num_layers); - push_u32_kv( - &mut bytes, - "llama.embedding_length", - metadata - .hidden_dim - .try_into() - .expect("test hidden dimension fits u32"), - ); - push_u32_kv( - &mut bytes, - "llama.context_length", - metadata - .context_length - .try_into() - .expect("test context length fits u32"), - ); - push_u32_kv( - &mut bytes, - "tokenizer.ggml.eos_token_id", - metadata.eos_token_id, - ); - bytes - } - - fn push_string_kv(bytes: &mut Vec, key: &str, value: &str) { - push_gguf_string(bytes, key); - bytes.extend_from_slice(&8_u32.to_le_bytes()); - push_gguf_string(bytes, value); - } - - fn push_u32_kv(bytes: &mut Vec, key: &str, value: u32) { - push_gguf_string(bytes, key); - bytes.extend_from_slice(&4_u32.to_le_bytes()); - bytes.extend_from_slice(&value.to_le_bytes()); - } - - fn push_gguf_string(bytes: &mut Vec, value: &str) { - bytes.extend_from_slice(&(value.len() as u64).to_le_bytes()); - bytes.extend_from_slice(value.as_bytes()); - } - - impl Drop for TempModelFile { - fn drop(&mut self) { - let _ = std::fs::remove_dir_all(&self.root); - } - } - - struct TempTomlFile { - path: PathBuf, - } - - impl TempTomlFile { - fn new(file_name: &str, contents: &str) -> Self { - let path = std::env::temp_dir().join(format!( - "mvp-orchestrator-config-test-{}-{}-{file_name}", - std::process::id(), - std::thread::current().name().unwrap_or("unnamed") - )); - let _ = std::fs::remove_file(&path); - std::fs::write(&path, contents).expect("write temp TOML config"); - Self { path } - } - } - - impl Drop for TempTomlFile { - fn drop(&mut self) { - let _ = std::fs::remove_file(&self.path); - } - } - - #[test] - fn frame_archive_writes_jsonl_records_for_text_and_binary_payloads() { - static NEXT_TEMP_FILE: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0); - - let suffix = NEXT_TEMP_FILE.fetch_add(1, std::sync::atomic::Ordering::Relaxed); - let path = std::env::temp_dir().join(format!( - "mvp-frame-archive-test-{}-{suffix}.jsonl", - std::process::id() - )); - let _ = std::fs::remove_file(&path); - - let stream = StreamId::new("test-node", Lifetime(42)); - let mut archive = FrameArchive::open(&path).expect("frame archive opens"); - archive.record( - "orchestrator", - &stream, - "stdout", - &Frame::new( - ChannelId(1), - datastream::Position(7), - b"hello \xce\xbb".to_vec(), - ), - ); - archive.record( - "orchestrator", - &stream, - "stderr", - &Frame::new( - ChannelId(2), - datastream::Position(8), - vec![0xff, 0x00, b'A'], - ), - ); - drop(archive); - - let contents = std::fs::read_to_string(&path).expect("read frame archive jsonl"); - let records = contents - .lines() - .map(|line| serde_json::from_str::(line).expect("archive line is json")) - .collect::>(); - let _ = std::fs::remove_file(&path); - - assert_eq!(records.len(), 2); - assert_eq!(records[0]["arrival_seq"], json!(0)); - assert!( - records[0]["arrival_unix_ms"] - .as_u64() - .is_some_and(|value| value > 0) - ); - assert_eq!(records[0]["source"], json!("orchestrator")); - assert_eq!(records[0]["stream"], json!("test-node#42")); - assert_eq!(records[0]["channel"], json!("stdout")); - assert_eq!(records[0]["channel_id"], json!(1)); - assert_eq!(records[0]["position"], json!(7)); - assert_eq!( - records[0]["payload"], - json!({"encoding":"utf8","value":"hello λ"}) - ); - - assert_eq!(records[1]["arrival_seq"], json!(1)); - assert!( - records[1]["arrival_unix_ms"] - .as_u64() - .is_some_and(|value| value > 0) - ); - assert_eq!(records[1]["source"], json!("orchestrator")); - assert_eq!(records[1]["stream"], json!("test-node#42")); - assert_eq!(records[1]["channel"], json!("stderr")); - assert_eq!(records[1]["channel_id"], json!(2)); - assert_eq!(records[1]["position"], json!(8)); - assert_eq!( - records[1]["payload"], - json!({"encoding":"bytes","value":[255,0,65]}) - ); - } - - #[test] - fn orchestrator_stdio_drain_archives_stdout_and_stderr_as_provision_logs() { - static NEXT_TEMP_FILE: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0); - - let suffix = NEXT_TEMP_FILE.fetch_add(1, std::sync::atomic::Ordering::Relaxed); - let path = std::env::temp_dir().join(format!( - "mvp-orch-stdio-test-{}-{suffix}.jsonl", - std::process::id() - )); - let _ = std::fs::remove_file(&path); - let (tx, rx) = mpsc::channel(); - tx.send(OrchStdioLine { - stream: ProvisionLogStream::Stdout, - line: "offer pool selected".to_owned(), - }) - .expect("send stdout line"); - tx.send(OrchStdioLine { - stream: ProvisionLogStream::Stderr, - line: "lease chain detail".to_owned(), - }) - .expect("send stderr line"); - - let mut datastream = OrchDatastream::new(77, Some(&path)).expect("datastream opens"); - drain_orch_stdio_capture(Some(&rx), &mut datastream, None, 77, 9); - drop(datastream); - - let contents = std::fs::read_to_string(&path).expect("read frame archive jsonl"); - let records = contents - .lines() - .map(|line| serde_json::from_str::(line).expect("archive line is json")) - .collect::>(); - let _ = std::fs::remove_file(&path); - - let stdout_record = records - .iter() - .rev() - .find(|record| record["channel"] == "mvp.provisioning.logs.node.9.stdout") - .expect("stdout provisioning log frame archived"); - let stderr_record = records - .iter() - .rev() - .find(|record| record["channel"] == "mvp.provisioning.logs.node.9.stderr") - .expect("stderr provisioning log frame archived"); - assert_eq!(stdout_record["source"], "orchestrator"); - assert_eq!(stderr_record["source"], "orchestrator"); - - let stdout_payload = serde_json::from_str::( - stdout_record["payload"]["value"] - .as_str() - .expect("stdout payload is archived as text"), - ) - .expect("stdout payload is log record json"); - let stderr_payload = serde_json::from_str::( - stderr_record["payload"]["value"] - .as_str() - .expect("stderr payload is archived as text"), - ) - .expect("stderr payload is log record json"); - assert_eq!(stdout_payload["line"]["run_id"], 77); - assert_eq!(stdout_payload["line"]["node_id"], 9); - assert_eq!(stdout_payload["line"]["stream"], "Stdout"); - assert_eq!(stdout_payload["line"]["line"], "offer pool selected"); - assert_eq!(stderr_payload["line"]["stream"], "Stderr"); - assert_eq!(stderr_payload["line"]["line"], "lease chain detail"); - } - - #[test] - fn config_layers_defaults_toml_env_then_cli() { - let toml = TempTomlFile::new( - "layering.toml", - r#" -[runtime] -profile = "deploy" -run_id = 41 -node_id = 9 -stage_index = 3 -layer_end_exclusive = 24 - -[provider] -kind = "vastai" - -[image] -node = "docker.io/example/from-image-node:toml" - -[vastai] -image = "docker.io/example/from-vastai-image:toml" -api_key = "toml-key" -bootstrap_command = "/toml/bootstrap" -disk_gb = 60 -gpu_name = "RTX 4090" - -[prompt] -rpc_addr = "127.0.0.1:19999" -max_tokens = 17 -dashboard = true - -[model] -id = "toml-model" -gguf_repo = "toml/repo" -gguf_file = "toml.gguf" -gguf_revision = "toml-rev" -max_context = 384 - -[relay] -mode = "disabled" -"#, - ); - - let config = with_clean_env( - &[ - ("MVP_NODE_PROVIDER", "docker"), - ("MVP_NODE_IMAGE", "docker.io/example/from-env:latest"), - ("MVP_PROMPT_MAX_TOKENS", "23"), - ("MVP_MODEL_ID", "env-model"), - ("MVP_IROH_RELAY_MODE", "default"), - ], - || { - Config::from_layers_with_path_and_args( - Some(&toml.path), - [ - "--image", - "docker.io/example/from-cli:latest", - "--max-tokens", - "31", - "--model-id", - "cli-model", - "--max-context", - "768", - ] - .into_iter() - .map(str::to_owned), - ) - .expect("layered config parses") - }, - ); - - assert_eq!(config.provider, ProviderKind::Docker); - assert_eq!(config.image, "docker.io/example/from-cli:latest"); - assert_eq!(config.default_max_tokens, 31); - assert_eq!(config.model_id, "cli-model"); - assert_eq!(config.run_id, 41); - assert_eq!(config.node_id, 9); - assert_eq!(config.stage_index, 3); - assert_eq!(config.layer_end_exclusive, Some(24)); - assert!(config.dashboard); - assert_eq!(config.max_context, Some(768)); - assert!(matches!(config.relay.mode, iroh::RelayMode::Default)); - assert!(config.vastai.is_none()); - } - - #[test] - fn pipeline_stages_layers_toml_env_then_cli_aliases() { - let toml = TempTomlFile::new( - "pipeline-stages.toml", - r#" -[runtime] -pipeline_stages = 2 - -[provider] -kind = "docker" -"#, - ); - - for (case, settings, args, expected) in [ - ("toml", vec![], vec![], 2), - ( - "env-over-toml", - vec![("MVP_PIPELINE_STAGES", "3")], - vec![], - 3, - ), - ( - "long-cli-over-env", - vec![("MVP_PIPELINE_STAGES", "3")], - vec!["--pipeline-stages", "4"], - 4, - ), - ( - "short-cli-over-env", - vec![("MVP_PIPELINE_STAGES", "3")], - vec!["-N", "5"], - 5, - ), - ] { - let config = with_clean_env(&settings, || { - Config::from_layers_with_path_and_args( - Some(&toml.path), - args.iter().copied().map(str::to_owned), - ) - .unwrap_or_else(|error| { - panic!("{case} pipeline stage config should parse: {error}") - }) - }); - - assert_eq!(config.pipeline_stages, expected, "{case}"); - } - } - - #[test] - fn pipeline_stages_rejects_zero_missing_and_non_numeric_values() { - for (case, settings, args, expected) in [ - ( - "env-zero", - vec![("MVP_PIPELINE_STAGES", "0")], - vec![], - "--pipeline-stages must be greater than 0", - ), - ( - "env-non-numeric", - vec![("MVP_PIPELINE_STAGES", "many")], - vec![], - "invalid MVP_PIPELINE_STAGES=\"many\"", - ), - ( - "long-cli-zero", - vec![], - vec!["--pipeline-stages", "0"], - "--pipeline-stages must be greater than 0", - ), - ( - "long-cli-non-numeric", - vec![], - vec!["--pipeline-stages", "many"], - "invalid --pipeline-stages=\"many\"", - ), - ( - "short-cli-missing", - vec![], - vec!["-N"], - "missing value after -N", - ), - ] { - let error = - with_clean_env(&settings, || { - match Config::from_layers_with_path_and_args( - None, - args.iter().copied().map(str::to_owned), - ) { - Ok(_) => panic!("invalid pipeline stages setting must fail"), - Err(error) => error, - } - }); - - assert!( - error.contains(expected), - "{case} error {error:?} should contain {expected:?}" - ); - } - } - - #[test] - fn vastai_rejects_pipeline_stages_count_above_one() { - let error = with_clean_env(&[], || { - match Config::from_layers_with_path_and_args( - None, - ["--provider", "vastai", "--pipeline-stages", "2"] - .into_iter() - .map(str::to_owned), - ) { - Ok(_) => panic!("VastAI cannot provision more than one pipeline stage yet"), - Err(error) => error, - } - }); - - assert!( - error.contains("provider=vastai does not support pipelined provisioning yet"), - "unexpected error: {error}" - ); - } - - #[test] - fn pipeline_cached_model_resolution_uses_toml_smollm2_gguf_file_with_cli_stage_count() { - let toml = TempTomlFile::new( - "pipeline-smollm2-model.toml", - r#" -[model] -gguf_repo = "QuantFactory/SmolLM2-135M-Instruct-GGUF" -gguf_file = "SmolLM2-135M-Instruct.Q4_0.gguf" -"#, - ); - - let config = with_clean_env(&[], || { - Config::from_layers_with_path_and_args( - Some(&toml.path), - ["--pipeline-stages", "3"].into_iter().map(str::to_owned), - ) - .expect("TOML SmolLM2 pipeline config parses") - }); - let expected_cached_host_path = default_pipeline_cached_model_path() - .canonicalize() - .expect("default cached SmolLM2 GGUF is present for cache resolution tests"); - let cached_model = config - .cached_model - .as_ref() - .expect("matching TOML SmolLM2 pipeline source should resolve the default cache"); - - assert_eq!(config.pipeline_stages, 3); - assert_eq!(config.provider, ProviderKind::Process); - assert_eq!(cached_model.host_path, expected_cached_host_path); - assert_eq!( - cached_model.container_path, - "/models/cached/SmolLM2-135M-Instruct.Q4_0.gguf" - ); - assert_eq!( - config.gguf_source, - GgufSource::LocalPath(expected_cached_host_path.to_string_lossy().to_string()) - ); - } - - #[test] - fn pipeline_cached_model_resolution_does_not_use_default_cache_for_other_toml_gguf_file() { - let toml = TempTomlFile::new( - "pipeline-other-model.toml", - r#" -[model] -gguf_repo = "QuantFactory/SmolLM2-135M-Instruct-GGUF" -gguf_file = "SmolLM2-135M-Instruct.Q8_0.gguf" -"#, - ); - - let config = with_clean_env(&[], || { - Config::from_layers_with_path_and_args( - Some(&toml.path), - ["--pipeline-stages", "3"].into_iter().map(str::to_owned), - ) - .expect("non-default TOML GGUF pipeline config parses") - }); - - assert_eq!(config.pipeline_stages, 3); - assert!(config.cached_model.is_none()); - assert_eq!( - config.gguf_source, - GgufSource::HuggingFaceGguf { - repo: "QuantFactory/SmolLM2-135M-Instruct-GGUF".to_owned(), - file: "SmolLM2-135M-Instruct.Q8_0.gguf".to_owned(), - revision: None, - } - ); - - let error = config - .build_run_plan() - .expect_err("non-default remote TOML GGUF should still require an explicit cache"); - - assert!( - error.contains( - "QuantFactory/SmolLM2-135M-Instruct-GGUF/SmolLM2-135M-Instruct.Q8_0.gguf" - ), - "unexpected error: {error}" - ); - assert!( - error.contains("--cached-model-host-path"), - "unexpected error: {error}" - ); - } - - #[test] - fn cached_model_plan_phase_is_metadata_driven_for_each_pipeline_width() { - let metadata = TestGgufMetadata { - num_layers: 7, - hidden_dim: 13, - context_length: 64, - eos_token_id: 11, - }; - let model = TempModelFile::with_metadata("metadata-driven.gguf", metadata); - - for stage_count in [1_u32, 3, metadata.num_layers] { - let config = pipeline_config_with_cached_model(&model, stage_count); - assert!(config.uses_planned_execution()); - let plan = config - .build_run_plan() - .expect("metadata-derived cached-model plan builds"); - assert_plan_matches_metadata(&config, &plan, metadata, stage_count, 512); - } - } - - #[test] - fn cached_model_implicit_single_stage_matches_explicit_n_one_plan() { - let model = TempModelFile::new("implicit-n-one.gguf"); - let implicit = cached_model_config_with_args(&model, &["--max-context", "32"]); - let explicit = cached_model_config_with_args( - &model, - &["--pipeline-stages", "1", "--max-context", "32"], - ); - - let implicit_plan = implicit - .build_run_plan() - .expect("implicit cached-model N=1 plan builds"); - let explicit_plan = explicit - .build_run_plan() - .expect("explicit cached-model N=1 plan builds"); - - assert!(implicit.uses_planned_execution()); - assert!(explicit.uses_planned_execution()); - assert_plan_matches_metadata(&implicit, &implicit_plan, model.metadata, 1, 32); - assert_plan_matches_metadata(&explicit, &explicit_plan, model.metadata, 1, 32); - assert_eq!(implicit_plan.model, explicit_plan.model); - assert_eq!(implicit_plan.edges, explicit_plan.edges); - assert_eq!(implicit_plan.stages, explicit_plan.stages); - } - - #[test] - fn cached_model_plan_rejects_pipeline_width_above_model_layers() { - let metadata = TestGgufMetadata { - num_layers: 2, - ..TestGgufMetadata::default() - }; - let model = TempModelFile::with_metadata("too-many-stages.gguf", metadata); - let config = pipeline_config_with_cached_model(&model, metadata.num_layers + 1); - - let error = config - .build_run_plan() - .expect_err("pipeline width above layer count must reject before provisioning"); - - assert!( - error.contains("--pipeline-stages=3 exceeds GGUF layer count 2"), - "unexpected error: {error}" - ); - } - - #[test] - fn pipeline_planning_rejects_remote_gguf_source_without_local_cache() { - let config = with_clean_env(&[], || { - Config::from_layers_with_path_and_args( - None, - [ - "--pipeline-stages", - "3", - "--gguf-repo", - "example/remote", - "--gguf-file", - "remote.gguf", - ] - .into_iter() - .map(str::to_owned), - ) - .expect("non-default remote pipeline config parses") - }); - - let error = config - .build_run_plan() - .expect_err("remote GGUF source cannot be inspected before provisioning"); - - assert!( - error.contains("locally inspectable GGUF before provisioning"), - "unexpected error: {error}" - ); - assert!( - error.contains("example/remote/remote.gguf"), - "unexpected error: {error}" - ); - assert!( - error.contains("--cached-model-host-path"), - "unexpected error: {error}" - ); - } - - #[test] - fn toml_vastai_image_overrides_image_node_for_vastai_provider() { - let toml = TempTomlFile::new( - "vastai-image.toml", - r#" -[runtime] -profile = "deploy" - -[provider] -kind = "vastai" - -[image] -node = "docker.io/example/generic:toml" - -[vastai] -image = "docker.io/example/vastai:toml" -api_key = "k" -bootstrap_command = "/run" -"#, - ); - - let config = with_clean_env(&[], || { - Config::from_layers_with_path_and_args(Some(&toml.path), std::iter::empty::()) - .expect("VastAI TOML config parses") - }); - - assert_eq!(config.provider, ProviderKind::VastAi); - assert_eq!(config.image, "docker.io/example/vastai:toml"); - } - - #[test] - fn missing_toml_uses_hardcoded_defaults() { - let missing_path = std::env::temp_dir().join(format!( - "mvp-orchestrator-missing-config-{}-{}.toml", - std::process::id(), - std::thread::current().name().unwrap_or("unnamed") - )); - let _ = std::fs::remove_file(&missing_path); - - let config = with_clean_env(&[], || { - Config::from_layers_with_path_and_args( - Some(&missing_path), - std::iter::empty::(), - ) - .expect("missing optional TOML config uses defaults") - }); - - assert_eq!(config.provider, ProviderKind::Process); - assert_eq!(config.image, DEFAULT_IMAGE); - assert_eq!(config.model_id, DEFAULT_MODEL_ID); - assert_eq!(config.default_max_tokens, DEFAULT_MAX_TOKENS); - assert!(!config.dashboard); - assert_eq!(config.max_context, None); - } - - #[test] - fn node_spec_propagates_max_context_when_configured() { - let config = with_clean_env(&[], || { - Config::from_layers_with_path_and_args( - None, - ["--max-context", "256"].into_iter().map(str::to_owned), - ) - .expect("CLI max context config parses") - }); - let coordinator = EndpointAddr::new(iroh::SecretKey::from_bytes(&[3; 32]).public()); - let orchestrator_actor = ActorAddress([18; 32]); - let spec = config - .node_spec(coordinator, orchestrator_actor) - .expect("node spec builds"); - - assert_eq!(env_value(&spec.env, "MVP_MAX_CONTEXT"), Some("256")); - } - - #[test] - fn runtime_profile_selects_provider_and_node_provider_takes_precedence() { - assert_eq!( - selected_provider(&[("MVP_RUNTIME_CONFIG", "local")]), - ProviderKind::Process - ); - assert_eq!( - selected_provider(&[("MVP_RUNTIME_CONFIG", "deploy")]), - ProviderKind::VastAi - ); - assert_eq!( - selected_provider(&[ - ("MVP_RUNTIME_CONFIG", "deploy"), - ("MVP_NODE_PROVIDER", "docker"), - ]), - ProviderKind::Docker - ); - } - - #[test] - fn relay_mode_env_uses_default_relay_and_accepts_disabled() { - with_clean_env(&[], || { - assert!(matches!( - relay_mode_from_env().expect("unset relay mode parses"), - iroh::RelayMode::Default - )); - }); - with_clean_env(&[("MVP_IROH_RELAY_MODE", "disabled")], || { - assert!(matches!( - relay_mode_from_env().expect("disabled relay mode parses"), - iroh::RelayMode::Disabled - )); - }); - } - - #[test] - fn node_spec_env_propagates_relay_url_only_for_custom_relay_config() { - const RELAY_URL: &str = "https://relay-node-spec.example.com"; - - let custom_env = node_spec_env(&[(MVP_IROH_RELAY_URL_ENV, RELAY_URL)]); - let expected_url = canonical_relay_url(RELAY_URL); - assert_eq!( - env_value(&custom_env, MVP_IROH_RELAY_URL_ENV), - Some(expected_url.as_str()) - ); - - let disabled_env = node_spec_env(&[ - ("MVP_IROH_RELAY_MODE", "disabled"), - (MVP_IROH_RELAY_URL_ENV, RELAY_URL), - ]); - assert_eq!(env_value(&disabled_env, MVP_IROH_RELAY_URL_ENV), None); - } - - #[test] - fn node_spec_preserves_relay_transport_in_coordinator_endpoint() { - with_clean_env(&[], || { - let config = Config::from_layers_with_path_and_args(None, std::iter::empty::()) - .expect("config parses"); - let secret = iroh::SecretKey::from_bytes(&[10; 32]); - let coordinator = EndpointAddr::new(secret.public()).with_relay_url( - "http://relay.example.com" - .parse::() - .unwrap(), - ); - let orchestrator_actor = ActorAddress([20; 32]); - - let spec = config - .node_spec(coordinator, orchestrator_actor) - .expect("node spec builds"); - let coordinator_endpoint_json = env_value(&spec.env, "MVP_COORDINATOR_ENDPOINT") - .expect("coordinator endpoint env is present"); - let coordinator_endpoint = - serde_json::from_str::(coordinator_endpoint_json) - .expect("coordinator endpoint env deserializes"); - - assert_eq!( - coordinator_endpoint - .relay_urls() - .next() - .map(|url| url.to_string()), - Some("http://relay.example.com/".to_owned()) - ); - }); - } - - #[test] - fn docker_config_construction_ignores_malformed_vastai_environment() { - let config = with_clean_env( - &[ - ("MVP_RUNTIME_CONFIG", "local"), - ("MVP_NODE_PROVIDER", "docker"), - ("MVP_VASTAI_CONFIRM_LEASE", "definitely-not-a-bool"), - ("MVP_VASTAI_DISK_GB", "not-a-u32"), - ("MVP_VASTAI_MIN_DOWN_MBPS", "not-a-float"), - ], - || { - Config::from_layers_with_path_and_args(None, std::iter::empty::()) - .expect("docker config ignores VastAI-only env") - }, - ); - - assert_eq!(config.provider, ProviderKind::Docker); - assert!(config.vastai.is_none()); - } - - #[test] - fn docker_cached_model_builds_local_gguf_env_and_planned_file_mount_from_canonical_host_path() { - let model = TempModelFile::new("weights-q4.gguf"); - let config = with_clean_env_os( - &[ - ("MVP_RUNTIME_CONFIG", OsString::from("local")), - ("MVP_NODE_PROVIDER", OsString::from("docker")), - (CACHED_MODEL_HOST_ENV, model.raw_path.as_os_str().to_owned()), - ], - || { - Config::from_layers_with_path_and_args(None, std::iter::empty::()) - .expect("docker cached model config parses") - }, - ); - let coordinator = EndpointAddr::new(iroh::SecretKey::from_bytes(&[7; 32]).public()); - let orchestrator_actor = ActorAddress([14; 32]); - let spec = config - .node_spec(coordinator, orchestrator_actor) - .expect("cached model node spec builds"); - - assert_eq!( - env_value(&spec.env, "MVP_GGUF_LOCAL_PATH"), - Some("/models/cached/weights-q4.gguf") - ); - assert_eq!(env_value(&spec.env, "MVP_GGUF_REPO"), None); - assert_eq!(env_value(&spec.env, "MVP_GGUF_FILE"), None); - assert_eq!( - spec.mounts, - vec![ProviderMount { - host_path: model.canonical_path.to_string_lossy().to_string(), - container_path: "/models/cached/weights-q4.gguf".to_owned(), - readonly: true, - }] - ); - } - - #[test] - fn process_cached_model_builds_host_gguf_env_without_mounts() { - let model = TempModelFile::new("process-weights-q4.gguf"); - let config = with_clean_env_os( - &[ - ("MVP_RUNTIME_CONFIG", OsString::from("local")), - (CACHED_MODEL_HOST_ENV, model.raw_path.as_os_str().to_owned()), - ], - || { - Config::from_layers_with_path_and_args(None, std::iter::empty::()) - .expect("process cached model config parses") - }, - ); - let coordinator = EndpointAddr::new(iroh::SecretKey::from_bytes(&[8; 32]).public()); - let orchestrator_actor = ActorAddress([15; 32]); - let spec = config - .node_spec(coordinator, orchestrator_actor) - .expect("process cached model node spec builds"); - let host_path = model.canonical_path.to_string_lossy().to_string(); - - assert_eq!(config.provider, ProviderKind::Process); - assert_eq!(env_value(&spec.env, "MVP_NODE_PROVIDER"), Some("process")); - assert_eq!( - env_value(&spec.env, "MVP_GGUF_LOCAL_PATH"), - Some(host_path.as_str()) - ); - assert_eq!(env_value(&spec.env, "MVP_DOCKER_GPUS"), None); - assert!(spec.mounts.is_empty(), "process workers use host paths"); - } - - #[test] - fn vectorized_local_docker_stage_node_specs_follow_three_stage_plan() { - let model = TempModelFile::new("pipeline-node-spec.gguf"); - let config = pipeline_config_with_cached_model(&model, 3); - let plan = config - .build_run_plan() - .expect("stage node spec plan builds"); - let coordinator = EndpointAddr::new(iroh::SecretKey::from_bytes(&[21; 32]).public()); - let orchestrator_actor = ActorAddress([23; 32]); - let coordinator_env = - serde_json::to_string(&coordinator).expect("coordinator endpoint serializes"); - let orchestrator_actor_env = - serde_json::to_string(&orchestrator_actor).expect("orchestrator actor serializes"); - let expected_run_id = config.run_id.to_string(); - let expected_mount = ProviderMount { - host_path: model.canonical_path.to_string_lossy().to_string(), - container_path: "/models/cached/pipeline-node-spec.gguf".to_owned(), - readonly: true, - }; - - let specs = stage_node_specs(&config, Some(&plan), coordinator, orchestrator_actor) - .expect("pipeline stage node specs build"); - - assert_eq!( - specs.len(), - 3, - "pipeline_stages=3 should provision three Docker workers, not one" - ); - assert!( - specs.iter().all(|spec| spec.node_id != config.node_id), - "coordinator node {} must not be counted as a worker: {specs:?}", - config.node_id - ); - - for (expected_stage_index, spec) in specs.iter().enumerate() { - let expected_stage_index = - u32::try_from(expected_stage_index).expect("fixture stage index fits u32"); - let expected_node_id = config.node_id + 1 + u64::from(expected_stage_index); - let expected_node_id_env = expected_node_id.to_string(); - let expected_stage_index_env = expected_stage_index.to_string(); - - assert_eq!(spec.run_id, config.run_id); - assert_eq!(spec.node_id, expected_node_id); - assert_eq!(spec.stage_index, Some(expected_stage_index)); - assert_eq!( - env_value(&spec.env, "MVP_RUN_ID"), - Some(expected_run_id.as_str()) - ); - assert_eq!( - env_value(&spec.env, "MVP_LOGICAL_NODE_ID"), - Some(expected_node_id_env.as_str()) - ); - assert_eq!( - env_value(&spec.env, "MVP_STAGE_INDEX"), - Some(expected_stage_index_env.as_str()) - ); - assert_eq!(env_value(&spec.env, "MVP_PIPELINE_STAGES"), Some("3")); - assert_eq!(env_value(&spec.env, "MVP_NODE_PROVIDER"), Some("docker")); - assert_eq!( - env_value(&spec.env, "MVP_MODEL_ID"), - Some(config.model_id.as_str()) - ); - assert_eq!( - env_value(&spec.env, "MVP_GGUF_LOCAL_PATH"), - Some("/models/cached/pipeline-node-spec.gguf") - ); - assert_eq!(env_value(&spec.env, "MVP_MAX_CONTEXT"), Some("512")); - assert_eq!( - env_value(&spec.env, "MVP_COORDINATOR_ENDPOINT"), - Some(coordinator_env.as_str()) - ); - assert_eq!( - env_value(&spec.env, "MVP_ORCHESTRATOR_ACTOR"), - Some(orchestrator_actor_env.as_str()) - ); - assert_eq!( - spec.mounts, - vec![expected_mount.clone()], - "pipeline cached GGUF mount should be read-only for stage {expected_stage_index}" - ); - } - } - - #[test] - fn vectorized_local_process_stage_node_specs_follow_three_stage_plan() { - let model = TempModelFile::new("process-pipeline-node-spec.gguf"); - let config = with_clean_env_os( - &[ - ("MVP_RUNTIME_CONFIG", OsString::from("local")), - (CACHED_MODEL_HOST_ENV, model.raw_path.as_os_str().to_owned()), - ], - || { - Config::from_layers_with_path_and_args( - None, - ["--pipeline-stages", "3", "--max-context", "512"] - .into_iter() - .map(str::to_owned), - ) - .expect("process pipeline config parses") - }, - ); - let plan = config - .build_run_plan() - .expect("process stage node spec plan builds"); - let coordinator = EndpointAddr::new(iroh::SecretKey::from_bytes(&[22; 32]).public()); - let orchestrator_actor = ActorAddress([24; 32]); - let host_path = model.canonical_path.to_string_lossy().to_string(); - - let specs = stage_node_specs(&config, Some(&plan), coordinator, orchestrator_actor) - .expect("process pipeline stage node specs build"); - - assert_eq!(config.provider, ProviderKind::Process); - assert_eq!(specs.len(), 3); - for (expected_stage_index, spec) in specs.iter().enumerate() { - let expected_stage_index = - u32::try_from(expected_stage_index).expect("fixture stage index fits u32"); - assert_eq!(spec.stage_index, Some(expected_stage_index)); - assert_eq!(env_value(&spec.env, "MVP_NODE_PROVIDER"), Some("process")); - assert_eq!(env_value(&spec.env, "MVP_PIPELINE_STAGES"), Some("3")); - assert_eq!( - env_value(&spec.env, "MVP_GGUF_LOCAL_PATH"), - Some(host_path.as_str()) - ); - assert_eq!(env_value(&spec.env, "MVP_MAX_CONTEXT"), Some("512")); - assert!(spec.mounts.is_empty(), "process stage specs must not mount"); - } - } - - #[test] - fn planned_single_stage_cached_model_node_spec_follows_generated_plan() { - let model = TempModelFile::new("single-stage-node-spec.gguf"); - let config = pipeline_config_with_cached_model(&model, 1); - let plan = config - .build_run_plan() - .expect("single-stage cached plan builds"); - let coordinator = EndpointAddr::new(iroh::SecretKey::from_bytes(&[31; 32]).public()); - let orchestrator_actor = ActorAddress([33; 32]); - - let specs = stage_node_specs(&config, Some(&plan), coordinator, orchestrator_actor) - .expect("single planned stage node spec builds"); - - assert_eq!(specs.len(), 1); - let spec = &specs[0]; - let stage = &plan.stages[0]; - assert_eq!(stage.stage_index, 0); - assert_eq!(stage.layer_start, 0); - assert_eq!(stage.layer_end_exclusive, model.metadata.num_layers); - assert_eq!(spec.node_id, stage.node_id.0); - assert_ne!( - spec.node_id, config.node_id, - "planned cached N=1 still provisions a worker separate from the coordinator" - ); - assert_eq!(spec.stage_index, Some(0)); - assert_eq!(env_value(&spec.env, "MVP_PIPELINE_STAGES"), Some("1")); - assert_eq!( - env_value(&spec.env, "MVP_GGUF_LOCAL_PATH"), - Some("/models/cached/single-stage-node-spec.gguf") - ); - assert_eq!( - spec.mounts, - vec![ProviderMount { - host_path: model.canonical_path.to_string_lossy().to_string(), - container_path: "/models/cached/single-stage-node-spec.gguf".to_owned(), - readonly: true, - }] - ); - } - - #[test] - fn vectorized_local_docker_stage_provision_detail_preserves_plan_edges_and_layers() { - let model = TempModelFile::new("pipeline-stage-detail.gguf"); - let config = pipeline_config_with_cached_model(&model, 3); - let plan = config.build_run_plan().expect("stage detail plan builds"); - - let detail = stage_provision_detail(&config, Some(&plan)); - - assert_eq!( - detail, - json!({ - "run_id": config.run_id, - "stage_count": 3, - "stages": plan.stages.iter().map(|stage| { - json!({ - "node_id": stage.node_id.0, - "stage_index": stage.stage_index, - "layer_range": { - "start": stage.layer_start, - "end_exclusive": stage.layer_end_exclusive - }, - "inbound_edge_id": stage.inbound_edge.0, - "outbound_edge_id": stage.outbound_edge.0, - }) - }).collect::>(), - "model_id": config.model_id, - }) - ); - } - - #[test] - fn stage_provision_wire_preserves_plan_edges_for_single_and_multi_stage_cached_models() { - for stage_count in [1_u32, 3] { - let model = TempModelFile::new(&format!("wire-{stage_count}.gguf")); - let config = pipeline_config_with_cached_model(&model, stage_count); - let plan = config.build_run_plan().expect("wire test plan builds"); - let coordinator = EndpointAddr::new(iroh::SecretKey::from_bytes(&[41; 32]).public()); - let readies = plan - .stages - .iter() - .map(|stage| { - let byte = u8::try_from(stage.stage_index + 42).expect("fixture byte fits"); - let endpoint = - EndpointAddr::new(iroh::SecretKey::from_bytes(&[byte; 32]).public()); - ( - stage.node_id.0, - RuntimeReady { - endpoint: endpoint.clone(), - node_actor: ActorAddress([byte; 32]), - datastream_publisher: ActorAddress([byte.wrapping_add(80); 32]), - stage_index: stage.stage_index, - readiness_id: u64::from(stage.stage_index) + 100, - swim_node_id: DistNodeId(*endpoint.id.as_bytes()), - }, - ) - }) - .collect::>(); - - for stage in &plan.stages { - let wire = stage_provision_wire_from_plan( - &plan, - stage.stage_index, - &readies, - &coordinator, - ) - .expect("stage provision wire builds from plan"); - let inbound = wire.inbound_edge.as_ref().expect("planned inbound edge"); - let outbound = wire.outbound_edge.as_ref().expect("planned outbound edge"); - - assert_eq!(wire.stage_count, stage_count); - assert_eq!(wire.node_id, stage.node_id.0); - assert_eq!(wire.layer_start, stage.layer_start); - assert_eq!(wire.layer_end_exclusive, stage.layer_end_exclusive); - assert_eq!(inbound.edge_id, stage.inbound_edge.0); - assert_eq!(outbound.edge_id, stage.outbound_edge.0); - assert_eq!( - inbound.kind, - if stage.stage_index == 0 { - StageEdgeKindWire::TokenIn - } else { - StageEdgeKindWire::Activation - } - ); - assert_eq!( - outbound.kind, - if stage.stage_index + 1 == stage_count { - StageEdgeKindWire::TokenOut - } else { - StageEdgeKindWire::Activation - } - ); - assert_eq!(wire.model_id, config.model_id); - assert_eq!(wire.gguf_source, config.gguf_source); - assert_eq!(wire.tokenizer, config.tokenizer); - } - } - } - - #[test] - fn cached_model_with_deploy_provider_is_rejected_before_vastai_env_is_parsed() { - let model = TempModelFile::new("deploy-rejected.gguf"); - let error = with_clean_env_os( - &[ - ("MVP_RUNTIME_CONFIG", OsString::from("deploy")), - (CACHED_MODEL_HOST_ENV, model.raw_path.as_os_str().to_owned()), - ( - "MVP_VASTAI_CONFIRM_LEASE", - OsString::from("definitely-not-a-bool"), - ), - ("MVP_VASTAI_DISK_GB", OsString::from("not-a-u32")), - ], - || match Config::from_layers_with_path_and_args(None, std::iter::empty::()) { - Ok(_) => panic!("deploy cached model must be rejected"), - Err(error) => error, - }, - ); - - assert!( - error.contains( - "MVP_CACHED_MODEL_HOST_PATH is a host-local cache path and is only supported by provider=process or provider=docker" - ), - "unexpected error: {error}" - ); - assert!( - !error.contains("MVP_VASTAI_CONFIRM_LEASE") && !error.contains("MVP_VASTAI_DISK_GB"), - "cached-model rejection should not require valid VastAI env, got: {error}" - ); - } - - #[derive(Default)] - struct FakeProvisionPlugin { - stopped: Vec, - } - - impl ProvisionPlugin for FakeProvisionPlugin { - fn start_node( - &mut self, - _spec: NodeProvisionSpec, - _sink: PluginSink, - ) -> Result { - 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, - ) -> Result<(), String> { - self.stopped.push(handle.id); - Ok(()) - } - } - - #[test] - fn provisioned_node_guard_stops_node_on_drop() { - let mut plugin = FakeProvisionPlugin::default(); - { - let _guard = ProvisionedNodeGuard::new( - &mut plugin, - mvp_system::provisioning::PluginNodeHandle { - id: 7, - provider_process_id: Some(99), - }, - ); - } - assert_eq!(plugin.stopped, vec![7]); - } - - #[test] - fn provisioned_node_guard_explicit_stop_runs_once() { - let mut plugin = FakeProvisionPlugin::default(); - { - let mut guard = ProvisionedNodeGuard::new( - &mut plugin, - mvp_system::provisioning::PluginNodeHandle { - id: 8, - provider_process_id: None, - }, - ); - guard.stop().expect("first stop succeeds"); - guard.stop().expect("second stop is a no-op"); - } - assert_eq!(plugin.stopped, vec![8]); - } - - #[test] - fn stop_requested_observes_shutdown_signal_only() { - let (_tx, rx) = mpsc::channel(); - assert!(!stop_requested(&rx)); - - let (tx, rx) = mpsc::channel(); - tx.send(()).expect("send shutdown"); - assert!(stop_requested(&rx)); - assert!(!stop_requested(&rx)); - } - - fn endpoint(seed: u8) -> EndpointAddr { - EndpointAddr::new(iroh::SecretKey::from_bytes(&[seed; 32]).public()) - } - - fn insert_route(stack: &DistributionRuntimeStack, actor: ActorAddress, owner: DistNodeId) { - let mut route_view = match stack.route_view.write() { - Ok(route_view) => route_view, - Err(poisoned) => poisoned.into_inner(), - }; - route_view.insert(actor, owner); - } - - fn mark_alive(stack: &DistributionRuntimeStack, node_id: DistNodeId) { - stack - .runtime - .send_to( - stack.actors.membership_fanout, - distribution::swim::actor::MembershipChanged { - node_id, - state: MemberState::Alive, - incarnation: 1, - }, - ) - .expect("send membership change"); - stack.pump_runtime_once(); - } - - #[test] - fn runtime_ready_barrier_waits_for_specific_swim_and_route() { - let remote = DistNodeId([2; 32]); - let node_actor = ActorAddress::new_random(); - let datastream_publisher = ActorAddress::new_random(); - let ready = RuntimeReady { - endpoint: endpoint(2), - node_actor, - datastream_publisher, - stage_index: 3, - readiness_id: 99, - swim_node_id: remote, - }; - - let stack = - DistributionRuntimeStack::new(DistNodeId([1; 32]), DistributedNodeConfig::default()); - assert!(!runtime_ready_barrier_met(&stack, &ready)); - - let stack = - DistributionRuntimeStack::new(DistNodeId([1; 32]), DistributedNodeConfig::default()); - mark_alive(&stack, remote); - assert!(!runtime_ready_barrier_met(&stack, &ready)); - - let stack = - DistributionRuntimeStack::new(DistNodeId([1; 32]), DistributedNodeConfig::default()); - insert_route(&stack, node_actor, remote); - assert!(!runtime_ready_barrier_met(&stack, &ready)); - - let stack = - DistributionRuntimeStack::new(DistNodeId([1; 32]), DistributedNodeConfig::default()); - mark_alive(&stack, remote); - insert_route(&stack, node_actor, remote); - assert!(runtime_ready_barrier_met(&stack, &ready)); - } - - #[test] - fn enqueue_runtime_ready_ack_reports_to_node_agent() { - use mvp_system::actors::node_agent::{NodeAgentActor, NodeAgentReport}; - use mvp_system::actors::orchestrator::OrchestratorMsg; - use mvp_system::stage_controller as stage; - - let stack = - DistributionRuntimeStack::new(DistNodeId([1; 32]), DistributedNodeConfig::default()); - let orchestrator_inbox = stack - .runtime - .new_inbox::() - .expect("orchestrator inbox"); - let reports = stack - .runtime - .new_inbox::() - .expect("node report inbox"); - let node_actor = stack - .runtime - .spawn(NodeAgentActor::new( - stage::NodeId(11), - *orchestrator_inbox.addr(), - Some(*reports.addr()), - )) - .expect("spawn node agent"); - let ready = RuntimeReady { - endpoint: endpoint(9), - node_actor, - datastream_publisher: ActorAddress::new_random(), - stage_index: 3, - readiness_id: 99, - swim_node_id: DistNodeId([2; 32]), - }; - - enqueue_runtime_ready_ack(&stack, &ready, 7, 11).expect("enqueue runtime ready ack"); - stack.pump_runtime_once(); - - assert_eq!( - reports.try_recv(), - Some(NodeAgentReport::RuntimeReadyAck { - run_id: 7, - node_id: 11, - stage_index: ready.stage_index, - readiness_id: 99, - }) - ); - - assert_eq!( - orchestrator_inbox.try_recv(), - Some(OrchestratorMsg::ObserveNodeRuntimeReadyAck { - run_id: 7, - node_id: 11, - stage_index: ready.stage_index, - readiness_id: 99, - }) - ); - } -} diff --git a/crates/mvp-system/src/gpu_worker_ingress_parser.rs b/crates/mvp-system/src/gpu_worker_ingress_parser.rs index 4ae0998..794b408 100644 --- a/crates/mvp-system/src/gpu_worker_ingress_parser.rs +++ b/crates/mvp-system/src/gpu_worker_ingress_parser.rs @@ -56,20 +56,25 @@ pub const OBJECT_MAGIC_BYTES: [u8; 4] = *b"MO01"; pub const OBJECT_MAGIC: u32 = u32::from_le_bytes(OBJECT_MAGIC_BYTES); pub const OBJECT_VERSION: u16 = 1; pub const FLAG_END_OF_SEQUENCE: u32 = 1; -pub const KNOWN_FLAGS_MASK: u32 = FLAG_END_OF_SEQUENCE; +pub const FLAG_BEGIN_SEQUENCE: u32 = 1 << 1; +pub const KNOWN_FLAGS_MASK: u32 = FLAG_END_OF_SEQUENCE | FLAG_BEGIN_SEQUENCE; #[derive(Clone, Copy, Debug, PartialEq, Eq, Default)] pub struct ObjectFlags { pub end_of_sequence: bool, + pub begin_sequence: bool, } impl ObjectFlags { pub fn bits(self) -> u32 { + let mut bits = 0; if self.end_of_sequence { - FLAG_END_OF_SEQUENCE - } else { - 0 + bits |= FLAG_END_OF_SEQUENCE; } + if self.begin_sequence { + bits |= FLAG_BEGIN_SEQUENCE; + } + bits } pub fn from_bits(bits: u32) -> Option { @@ -78,6 +83,7 @@ impl ObjectFlags { } Some(Self { end_of_sequence: bits & FLAG_END_OF_SEQUENCE != 0, + begin_sequence: bits & FLAG_BEGIN_SEQUENCE != 0, }) } } diff --git a/crates/mvp-system/src/lib.rs b/crates/mvp-system/src/lib.rs index 19f7f9c..37995ea 100644 --- a/crates/mvp-system/src/lib.rs +++ b/crates/mvp-system/src/lib.rs @@ -24,6 +24,7 @@ pub mod node_boot_lifecycle; pub mod node_image; pub mod node_provisioning; pub mod observability_surface; +pub mod orchestrator_app; pub mod orchestrator_run_fsm; pub mod orchestrator_token_endpoint; pub mod prompt_rpc; diff --git a/crates/mvp-system/src/orchestrator_app.rs b/crates/mvp-system/src/orchestrator_app.rs new file mode 100644 index 0000000..cbe6b0d --- /dev/null +++ b/crates/mvp-system/src/orchestrator_app.rs @@ -0,0 +1,7483 @@ +use std::collections::{BTreeMap, BTreeSet}; +use std::fs::{File, OpenOptions}; +use std::io::{BufRead, BufReader, Write}; +use std::net::{SocketAddr, TcpListener, TcpStream}; +#[cfg(target_os = "linux")] +use std::os::fd::FromRawFd; +use std::path::{Path, PathBuf}; +use std::process::Command; +use std::sync::{Arc, mpsc}; +use std::thread; +use std::time::{Duration, Instant}; + +use crate::actors::node_agent::{ + NodeAgentMsg, StageEdgeKindWire, StageInboundEdgeWire, StageObjectSpecWire, + StageOutboundEdgeWire, StageProvisionWire, StageRingSpecWire, +}; +use crate::actors::orchestrator::{OrchestratorActor, OrchestratorReport}; +use crate::actors::register_mvp_actor_codecs; +use crate::benchmark_observability; +use crate::config::{DEFAULT_CONFIG_PATH, TomlConfigOverlay}; +#[cfg(feature = "dashboard")] +use crate::dashboard_view::MvpClusterDashboardView; +use crate::distribution_stack::DistributionRuntimeStack; +use crate::gpu_worker_ingress_parser as ingress; +use crate::node_provisioning::ProviderKind; +use crate::orchestrator_run_fsm::{RunConfig, RunId}; +use crate::prompt_rpc::{ + PromptEvent, SubmitPrompt, TokenizerEvent, read_submit_prompt, write_json_line, +}; +use crate::provisioning::{ + LocalDockerPlugin, LocalProcessPlugin, NodeProvisionSpec, PluginObservation, + PluginObservationSink, PluginSink, ProviderMount, ProvisionEvent, ProvisionEventKind, + ProvisionLogLine, ProvisionLogStream, ProvisionPlugin, +}; +#[cfg(test)] +use crate::relay_provisioning::relay_runtime_config_from_env; +use crate::relay_provisioning::{ + MVP_IROH_RELAY_URL_ENV, RelayRuntimeConfig, SWACTOR_IROH_RELAY_URL_ENV, relay_mode_env_value, + relay_runtime_config_from_settings, +}; +use crate::run_plan::{self, GgufSource, TokenizerSource}; +use crate::telemetry::{ + MVP_PROVISIONING_EVENTS, MvpProvisionEventRecord, MvpProvisionLogRecord, + mvp_provision_log_channel, +}; +use crate::vastai_provisioning::{ + SshCommandBootstrapLauncher, ToolsVastAiLeaseClient, VastAiProvisioningConfig, + VastAiProvisioningPlugin, +}; +use datastream::{ + ChannelContent, ChannelId, ChannelRef, DatastreamEndpoint, DatastreamEvent, DatastreamProducer, + DatastreamPublisherMsg, DatastreamSubscribe, Frame, Lifetime, NodeId, StreamDescriptor, + StreamId, StreamOrigin, SubscriptionRequest, +}; +use distribution::node::DistributedNodeConfig; +use distribution::types::{MemberState, NodeId as DistNodeId}; +use iroh::EndpointAddr; +use iroh_driver::{ + DATASTREAM_ALPN, EDGE_ALPN, EdgeSendHandle, EdgeTransportEvent, IrohDriver, IrohDriverConfig, +}; +use parking_lot::Mutex; +use serde_json::{Value, json}; +use swactor::actor::ActorAddress; +#[cfg(test)] +use tokio::sync::mpsc as tokio_mpsc; + +const DEFAULT_IMAGE: &str = "swactor-mvp-node:latest"; +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"; +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"; +const DEFAULT_HF_FILE: &str = "Llama-3.2-1B-Instruct-Q4_K_M.gguf"; +const DEFAULT_MODEL_ID: &str = "llama-3.2-1b-instruct-q4"; +const DEFAULT_MAX_TOKENS: u32 = 64; +const PUMP_INTERVAL: Duration = Duration::from_millis(10); +const RUNTIME_READY_ACK_RETRY_INTERVAL: Duration = Duration::from_millis(250); +const RUNTIME_READY_ACK_TIMEOUT: Duration = Duration::from_secs(60); +const MVP_ORCH_BOOTSTRAP: &str = "mvp.orch.bootstrap"; +const MVP_ORCH_PROMPT: &str = "mvp.orch.prompt"; +const MVP_SWIM_MEMBERSHIP: &str = "mvp.swim.membership"; +const MVP_STAGE_ROUTE: &str = "mvp.orch.stage_route"; +const DATASTREAM_FRAME_LOG_ENV: &str = "MVP_DATASTREAM_FRAME_LOG"; +const DEFAULT_DOCKER_CONTAINER_PREFIX: &str = "mvp-orchestrator"; +const MVP_DOCKER_CONTAINER_PREFIX_ENV: &str = "MVP_DOCKER_CONTAINER_PREFIX"; + +pub struct OrchestratorRunOptions { + pub capture_stdio: bool, + pub stop_rx: Option>, +} + +impl Default for OrchestratorRunOptions { + fn default() -> Self { + Self { + capture_stdio: true, + stop_rx: None, + } + } +} + +pub fn run_from_args(args: I) -> Result<(), String> +where + I: IntoIterator, +{ + run_with_options(args, OrchestratorRunOptions::default()) +} + +pub fn run_in_process_from_args(args: I, stop_rx: mpsc::Receiver<()>) -> Result<(), String> +where + I: IntoIterator, +{ + run_with_options( + args, + OrchestratorRunOptions { + capture_stdio: false, + stop_rx: Some(stop_rx), + }, + ) +} + +pub fn run_with_options(args: I, options: OrchestratorRunOptions) -> Result<(), String> +where + I: IntoIterator, +{ + let mut config = Config::from_defaults_toml_env_args(args)?; + config.prepare_vastai_ssh_key()?; + let orch_stdio_rx = if options.capture_stdio { + install_orch_stdio_capture()? + } else { + None + }; + let mut orch_datastream = + OrchDatastream::new(config.run_id, config.datastream_frame_log.as_deref())?; + orch_datastream.emit_bootstrap( + None, + config.run_id, + config.node_id, + "config", + "ready", + json!({ + "config_profile":config.config_profile.as_str(), + "image":&config.image, + "provider":config.provider.as_str(), + "rpc_bind":config.rpc_bind.to_string(), + "model_id":&config.model_id, + "stage_index":config.stage_index, + "legacy_layer_end_exclusive":config.layer_end_exclusive, + "relay_mode":format!("{:?}", config.relay.mode), + "pipeline_stages":config.pipeline_stages, + "provider_config":config.provider_datastream_detail(), + }), + ); + drain_orch_stdio_capture( + orch_stdio_rx.as_ref(), + &mut orch_datastream, + None, + config.run_id, + config.node_id, + ); + let pipeline_plan = if config.uses_planned_execution() { + let plan = config.build_run_plan()?; + orch_datastream.emit_bootstrap( + None, + config.run_id, + config.node_id, + "run_plan", + "ready", + json!({ + "stage_count":plan.stages.len(), + "edge_count":plan.edges.len(), + "model_layers":plan.model.num_layers, + "hidden_dim":plan.model.hidden_dim, + "max_seq_len":plan.model.max_seq_len, + "eos_token_id":plan.model.eos_token_id, + }), + ); + Some(plan) + } else { + None + }; + let _ = &pipeline_plan; + + let tokio = match tokio::runtime::Runtime::new() { + Ok(runtime) => { + orch_datastream.emit_bootstrap( + None, + config.run_id, + config.node_id, + "tokio_runtime", + "ready", + json!({"runtime":"tokio"}), + ); + runtime + } + Err(error) => { + orch_datastream.emit_bootstrap( + None, + config.run_id, + config.node_id, + "tokio_runtime", + "failed", + json!({"error":error.to_string()}), + ); + return Err(format!("tokio runtime: {error}")); + } + }; + let mut driver = match IrohDriver::with_handle( + tokio.handle().clone(), + IrohDriverConfig { + secret_key: None, + relay_mode: config.relay.mode.clone(), + node: DistributedNodeConfig::default(), + peer_auth: None, + additional_alpns: vec![EDGE_ALPN.to_vec(), DATASTREAM_ALPN.to_vec()], + }, + ) { + Ok(driver) => { + orch_datastream.emit_bootstrap( + None, + config.run_id, + config.node_id, + "iroh_driver", + "ready", + json!({"relay_mode":format!("{:?}", config.relay.mode)}), + ); + driver + } + Err(error) => { + orch_datastream.emit_bootstrap( + None, + config.run_id, + config.node_id, + "iroh_driver", + "failed", + json!({"error":error.to_string()}), + ); + return Err(format!("create iroh driver: {error}")); + } + }; + let stack = DistributionRuntimeStack::new_with_codecs( + driver.node_id(), + DistributedNodeConfig::default(), + |registry| { + register_mvp_actor_codecs(registry); + datastream::wire::register_datastream_codec(registry); + }, + ); + orch_datastream.emit_bootstrap( + None, + config.run_id, + config.node_id, + "distribution_stack", + "ready", + json!({"actors":"initialized","route_view":"initialized","swim":"initialized"}), + ); + orch_datastream.emit_bootstrap( + None, + config.run_id, + config.node_id, + "codecs", + "ready", + json!({"registered":["node_agent","orchestrator","provisioner","prompt_rpc","datastream"]}), + ); + driver.enable_actor_bridge( + stack.runtime.clone(), + stack.codec.clone(), + stack.actor_bridge_routes(), + stack.actors.swim, + stack.relay_mirror.clone(), + stack.route_view.clone(), + ); + orch_datastream.emit_bootstrap( + None, + config.run_id, + config.node_id, + "actor_bridge", + "ready", + json!({"transport":"iroh","routes":"attached"}), + ); + + let (frame_tx, frame_rx) = mpsc::channel::(); + orch_datastream.emit_bootstrap( + None, + config.run_id, + config.node_id, + "datastream_collector", + "ready", + json!({"alpn":String::from_utf8_lossy(DATASTREAM_ALPN)}), + ); + let dashboard = DashboardSupport::start(config.dashboard)?; + orch_datastream.emit_bootstrap( + dashboard.as_ref(), + config.run_id, + config.node_id, + "dashboard", + "ready", + json!({"enabled":dashboard.is_some()}), + ); + + let orchestrator_reports = match stack.runtime.new_inbox::() { + Ok(inbox) => inbox, + Err(error) => { + orch_datastream.emit_bootstrap( + dashboard.as_ref(), + config.run_id, + config.node_id, + "orchestrator_report_actor", + "failed", + json!({"error":error.to_string()}), + ); + return Err(format!("orchestrator report inbox: {error}")); + } + }; + let orchestrator_report_actor = *orchestrator_reports.addr(); + stack.register_local_actor(driver.register_actor(orchestrator_report_actor, 1)); + orch_datastream.emit_bootstrap( + dashboard.as_ref(), + config.run_id, + config.node_id, + "orchestrator_report_actor", + "ready", + json!({"actor":orchestrator_report_actor}), + ); + let orchestrator_actor = match stack.runtime.spawn(OrchestratorActor::new( + RunConfig { + run_id: RunId(config.run_id), + max_tokens: u64::from(config.default_max_tokens), + prompt: Vec::new(), + }, + Some(orchestrator_report_actor), + )) { + Ok(actor) => actor, + Err(error) => { + orch_datastream.emit_bootstrap( + dashboard.as_ref(), + config.run_id, + config.node_id, + "orchestrator_actor", + "failed", + json!({"error":error.to_string()}), + ); + return Err(format!("spawn orchestrator actor: {error}")); + } + }; + stack.register_local_actor(driver.register_actor(orchestrator_actor, 1)); + orch_datastream.emit_bootstrap( + dashboard.as_ref(), + config.run_id, + config.node_id, + "orchestrator_actor", + "ready", + json!({"actor":orchestrator_actor}), + ); + + let prompt_events = match stack.runtime.new_inbox::() { + Ok(inbox) => inbox, + Err(error) => { + orch_datastream.emit_bootstrap( + dashboard.as_ref(), + config.run_id, + config.node_id, + "prompt_reply_actor", + "failed", + json!({"error":error.to_string()}), + ); + return Err(format!("prompt event inbox: {error}")); + } + }; + let prompt_reply_actor = *prompt_events.addr(); + stack.register_local_actor(driver.register_actor(prompt_reply_actor, 1)); + orch_datastream.emit_bootstrap( + dashboard.as_ref(), + config.run_id, + config.node_id, + "prompt_reply_actor", + "ready", + json!({"actor":prompt_reply_actor}), + ); + + let tokenizer_events = match stack.runtime.new_inbox::() { + Ok(inbox) => inbox, + Err(error) => { + orch_datastream.emit_bootstrap( + dashboard.as_ref(), + config.run_id, + config.node_id, + "tokenizer_reply_actor", + "failed", + json!({"error":error.to_string()}), + ); + return Err(format!("tokenizer event inbox: {error}")); + } + }; + let tokenizer_reply_actor = *tokenizer_events.addr(); + stack.register_local_actor(driver.register_actor(tokenizer_reply_actor, 1)); + orch_datastream.emit_bootstrap( + dashboard.as_ref(), + config.run_id, + config.node_id, + "tokenizer_reply_actor", + "ready", + json!({"actor":tokenizer_reply_actor}), + ); + + let (work_tx, work_rx) = mpsc::channel::(); + let stop_rx = options.stop_rx.unwrap_or_else(spawn_stop_listener); + + let provisioner = config.build_provisioner(Arc::clone(&stack.runtime))?; + orch_datastream.emit_bootstrap( + dashboard.as_ref(), + config.run_id, + config.node_id, + "node_provisioner", + "ready", + json!({ + "provider":config.provider.as_str(), + "owner":"mvp-orchestrator", + "config":config.provider_datastream_detail(), + }), + ); + let (obs_tx, obs_rx) = mpsc::channel::(); + let sink = PluginSink::new(Arc::new(ChannelObservationSink { + tx: Mutex::new(obs_tx), + })); + let coordinator_endpoint = driver.endpoint_addr(); + let pipeline_coordinator_endpoint = coordinator_endpoint.clone(); + let (mut provisioned_nodes, ready) = start_and_provision_workers( + provisioner, + &config, + pipeline_plan.as_ref(), + &mut driver, + &stack, + &obs_rx, + &frame_rx, + &frame_tx, + &orchestrator_reports, + &stop_rx, + dashboard.as_ref(), + &mut orch_datastream, + orch_stdio_rx.as_ref(), + sink, + coordinator_endpoint, + pipeline_coordinator_endpoint, + orchestrator_actor, + )?; + + orch_datastream.emit_bootstrap( + dashboard.as_ref(), + config.run_id, + config.node_id, + "prompt_rpc", + "started", + json!({ + "bind":config.rpc_bind.to_string(), + "default_max_tokens":config.default_max_tokens, + }), + ); + + let rpc_addr = match spawn_prompt_rpc(config.rpc_bind, work_tx, config.default_max_tokens) { + Ok(addr) => { + orch_datastream.emit_bootstrap( + dashboard.as_ref(), + config.run_id, + config.node_id, + "prompt_rpc", + "ready", + json!({ + "addr":addr.to_string(), + "default_max_tokens":config.default_max_tokens, + }), + ); + addr + } + Err(error) => { + orch_datastream.emit_bootstrap( + dashboard.as_ref(), + config.run_id, + config.node_id, + "prompt_rpc", + "failed", + json!({"error":error}), + ); + return Err(error); + } + }; + + orch_datastream.emit_bootstrap( + dashboard.as_ref(), + config.run_id, + config.node_id, + "prompt_loop", + "ready", + json!({"addr":rpc_addr.to_string(),"node_actor":ready.first_stage.node_actor}), + ); + + orch_datastream.emit_bootstrap( + dashboard.as_ref(), + config.run_id, + config.node_id, + "serve_prompts", + "started", + json!({"mode":"single_active_prompt","poll_interval_ms":PUMP_INTERVAL.as_millis()}), + ); + let result = serve_prompts( + &mut driver, + &stack, + &obs_rx, + &frame_rx, + &frame_tx, + &work_rx, + &prompt_events, + &stop_rx, + dashboard.as_ref(), + &mut orch_datastream, + orch_stdio_rx.as_ref(), + config.run_id, + config.node_id, + ready.first_stage.node_actor, + prompt_reply_actor, + &tokenizer_events, + ready.first_stage.node_actor, + ready.final_stage.node_actor, + tokenizer_reply_actor, + config.provider, + pipeline_plan.as_ref(), + ready.first_stage.endpoint.clone(), + ); + if let Err(error) = &result { + orch_datastream.emit_bootstrap( + dashboard.as_ref(), + config.run_id, + config.node_id, + "serve_prompts", + "failed", + json!({"error":error}), + ); + } + orch_datastream.emit_bootstrap( + dashboard.as_ref(), + config.run_id, + config.node_id, + "provider_stop", + "started", + json!({"provider":config.provider.as_str(),"node_id":config.node_id}), + ); + let stop_result = provisioned_nodes.stop(); + match &stop_result { + Ok(()) => orch_datastream.emit_bootstrap( + dashboard.as_ref(), + config.run_id, + config.node_id, + "provider_stop", + "ready", + json!({"provider":config.provider.as_str(),"node_id":config.node_id}), + ), + Err(error) => orch_datastream.emit_bootstrap( + dashboard.as_ref(), + config.run_id, + config.node_id, + "provider_stop", + "failed", + json!({"provider":config.provider.as_str(),"node_id":config.node_id,"error":error}), + ), + } + if result.is_ok() && stop_result.is_ok() { + orch_datastream.emit_bootstrap( + dashboard.as_ref(), + config.run_id, + config.node_id, + "orch_exit", + "ready", + json!({"result":"ok"}), + ); + } + result.and(stop_result) +} + +#[derive(Clone)] +struct VastAiRuntimeConfig { + api_key: Option, + provisioning: VastAiProvisioningConfig, + bootstrap_command: Option, + ssh_identity: Option, + ssh_public_key: Option, + ssh_public_fingerprint: Option, +} + +impl VastAiRuntimeConfig { + fn from_builder(builder: &ConfigBuilder) -> Result { + let mut provisioning = VastAiProvisioningConfig::default(); + let disk_gb = builder + .vastai_disk_gb_raw + .as_ref() + .map(|value| ConfigBuilder::parse_value("MVP_VASTAI_DISK_GB", value)) + .transpose()? + .or(builder.vastai_disk_gb); + if let Some(disk_gb) = disk_gb { + provisioning.disk_gb = disk_gb; + } + if let Some(ssh_user) = &builder.vastai_ssh_user { + provisioning.ssh_user = ssh_user.clone(); + } + let confirm_lease = builder + .vastai_confirm_lease_raw + .as_ref() + .map(|value| ConfigBuilder::parse_bool("MVP_VASTAI_CONFIRM_LEASE", value)) + .transpose()? + .or(builder.vastai_confirm_lease); + if let Some(confirm_lease) = confirm_lease { + provisioning.confirm_lease = confirm_lease; + } + provisioning.onstart = builder.vastai_onstart.clone(); + provisioning.selection.gpu_name = builder.vastai_gpu_name.clone(); + let min_gpu_ram_mb = builder + .vastai_min_gpu_ram_mb_raw + .as_ref() + .map(|value| ConfigBuilder::parse_value("MVP_VASTAI_MIN_GPU_RAM_MB", value)) + .transpose()? + .or(builder.vastai_min_gpu_ram_mb); + if let Some(min_gpu_ram_mb) = min_gpu_ram_mb { + provisioning.selection.min_gpu_ram_mb = Some(min_gpu_ram_mb); + } + let min_down_mbps = builder + .vastai_min_down_mbps_raw + .as_ref() + .map(|value| ConfigBuilder::parse_value("MVP_VASTAI_MIN_DOWN_MBPS", value)) + .transpose()? + .or(builder.vastai_min_down_mbps); + if let Some(min_down_mbps) = min_down_mbps { + provisioning.selection.min_down_mbps = min_down_mbps; + } + let min_up_mbps = builder + .vastai_min_up_mbps_raw + .as_ref() + .map(|value| ConfigBuilder::parse_value("MVP_VASTAI_MIN_UP_MBPS", value)) + .transpose()? + .or(builder.vastai_min_up_mbps); + if let Some(min_up_mbps) = min_up_mbps { + provisioning.selection.min_up_mbps = Some(min_up_mbps); + } + let min_reliability = builder + .vastai_min_reliability_raw + .as_ref() + .map(|value| ConfigBuilder::parse_value("MVP_VASTAI_MIN_RELIABILITY", value)) + .transpose()? + .or(builder.vastai_min_reliability); + if let Some(min_reliability) = min_reliability { + provisioning.selection.min_reliability = min_reliability; + } + let require_verified = builder + .vastai_require_verified_raw + .as_ref() + .map(|value| ConfigBuilder::parse_bool("MVP_VASTAI_REQUIRE_VERIFIED", value)) + .transpose()? + .or(builder.vastai_require_verified); + if let Some(require_verified) = require_verified { + provisioning.selection.require_verified = require_verified; + } + let poll_interval_secs = builder + .vastai_poll_interval_secs_raw + .as_ref() + .map(|value| ConfigBuilder::parse_value("MVP_VASTAI_POLL_INTERVAL_SECS", value)) + .transpose()? + .or(builder.vastai_poll_interval_secs); + if let Some(poll_interval_secs) = poll_interval_secs { + provisioning.lifecycle.poll_interval = Duration::from_secs(poll_interval_secs); + } + let ssh_identity = builder + .vastai_ssh_identity_raw + .as_ref() + .map(|value| expand_home_path(value)) + .transpose()?; + Ok(Self { + api_key: builder.vastai_api_key.clone(), + provisioning, + bootstrap_command: builder.vastai_bootstrap_command.clone(), + ssh_identity, + ssh_public_key: None, + ssh_public_fingerprint: None, + }) + } + + fn datastream_detail(&self) -> Value { + json!({ + "disk_gb": self.provisioning.disk_gb, + "ssh_user": &self.provisioning.ssh_user, + "gpu_name": &self.provisioning.selection.gpu_name, + "min_gpu_ram_mb": self.provisioning.selection.min_gpu_ram_mb, + "min_down_mbps": self.provisioning.selection.min_down_mbps, + "min_up_mbps": self.provisioning.selection.min_up_mbps, + "min_reliability": self.provisioning.selection.min_reliability, + "require_verified": self.provisioning.selection.require_verified, + "confirm_lease": self.provisioning.confirm_lease, + "has_api_key": self.api_key.is_some(), + "has_onstart": self.provisioning.onstart.is_some(), + "has_bootstrap_command": self.bootstrap_command.is_some(), + "has_ssh_identity": self.ssh_identity.is_some(), + "ssh_public_fingerprint": self.ssh_public_fingerprint.as_deref(), + }) + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum RuntimeConfigProfile { + Local, + Deploy, +} + +impl RuntimeConfigProfile { + fn parse(value: &str) -> Result { + match value.trim().to_ascii_lowercase().as_str() { + "local" => Ok(Self::Local), + "deploy" => Ok(Self::Deploy), + other => Err(format!( + "unsupported {MVP_RUNTIME_CONFIG_ENV}={other:?}; use local or deploy" + )), + } + } + + fn as_str(self) -> &'static str { + match self { + Self::Local => "local", + Self::Deploy => "deploy", + } + } + + fn default_provider(self) -> ProviderKind { + match self { + Self::Local => ProviderKind::Process, + Self::Deploy => ProviderKind::VastAi, + } + } +} + +#[derive(Clone)] +struct CachedModelConfig { + host_path: PathBuf, + container_path: String, +} + +impl CachedModelConfig { + fn from_host_path(provider: ProviderKind, requested: PathBuf) -> Result { + if !matches!(provider, ProviderKind::Process | ProviderKind::Docker) { + return Err(format!( + "{CACHED_MODEL_HOST_ENV} is a host-local cache path and is only supported by provider=process or provider=docker" + )); + } + let host_path = requested.canonicalize().map_err(|e| { + format!( + "resolve {CACHED_MODEL_HOST_ENV} path {}: {e}", + requested.display() + ) + })?; + if !host_path.is_file() { + return Err(format!( + "{CACHED_MODEL_HOST_ENV} must point at a file: {}", + host_path.display() + )); + } + let container_path = cached_model_container_path(&host_path)?; + Ok(Self { + host_path, + container_path, + }) + } + + fn worker_path(&self, provider: ProviderKind) -> String { + match provider { + ProviderKind::Process => self.host_path.to_string_lossy().to_string(), + ProviderKind::Docker | ProviderKind::VastAi | ProviderKind::Mock => { + self.container_path.clone() + } + } + } + + fn datastream_detail(&self) -> Value { + json!({ + "host_path_present": true, + "file": self.host_path.file_name().and_then(|name| name.to_str()), + "container_path": &self.container_path, + }) + } +} + +fn cached_model_container_path(host_path: &Path) -> Result { + let file_name = host_path + .file_name() + .and_then(|name| name.to_str()) + .filter(|name| !name.is_empty()) + .ok_or_else(|| { + format!( + "cached model path has no file name: {}", + host_path.display() + ) + })?; + Ok(format!("{CACHED_MODEL_CONTAINER_DIR}/{file_name}")) +} + +fn default_worker_bin() -> Result { + let mut path = std::env::current_exe().map_err(|e| format!("current exe: {e}"))?; + path.set_file_name("mvp-worker-node"); + Ok(path) +} + +fn default_pipeline_cached_model_path() -> PathBuf { + let relative = PathBuf::from(".") + .join(DEFAULT_PIPELINE_MODEL_CACHE_DIR) + .join(DEFAULT_PIPELINE_CACHED_MODEL_FILE); + if relative.is_file() { + return relative; + } + PathBuf::from(env!("CARGO_MANIFEST_DIR")) + .join("..") + .join("..") + .join(DEFAULT_PIPELINE_MODEL_CACHE_DIR) + .join(DEFAULT_PIPELINE_CACHED_MODEL_FILE) +} + +fn gguf_source_is_default_hf(source: &GgufSource) -> bool { + matches!( + source, + GgufSource::HuggingFaceGguf { + repo, + file, + revision: None, + } if repo == DEFAULT_HF_REPO && file == DEFAULT_HF_FILE + ) +} + +fn gguf_source_matches_default_pipeline_cache(source: &GgufSource) -> bool { + matches!( + source, + GgufSource::HuggingFaceGguf { + file, + revision: None, + .. + } if file == DEFAULT_PIPELINE_CACHED_MODEL_FILE + ) +} + +#[derive(Clone)] +struct Config { + config_profile: RuntimeConfigProfile, + image: String, + docker_gpus: String, + provider: ProviderKind, + rpc_bind: SocketAddr, + run_id: u64, + node_id: u64, + stage_index: u32, + layer_end_exclusive: Option, + pipeline_stages: u32, + model_id: String, + gguf_source: GgufSource, + tokenizer: TokenizerSource, + default_max_tokens: u32, + dashboard: bool, + max_context: Option, + relay: RelayRuntimeConfig, + vastai: Option, + cached_model: Option, + datastream_frame_log: Option, + worker_bin: Option, +} + +#[derive(Clone)] +struct ConfigBuilder { + config_profile: RuntimeConfigProfile, + provider: Option, + image: String, + toml_vastai_image: Option, + image_overridden_after_toml: bool, + docker_gpus: String, + rpc_bind: String, + rpc_bind_label: &'static str, + run_id: u64, + node_id: u64, + stage_index: u32, + layer_end_exclusive: Option, + pipeline_stages: u32, + model_id: String, + gguf_source: GgufSource, + tokenizer: TokenizerSource, + default_max_tokens: u32, + dashboard: bool, + max_context: Option, + relay_mode: Option, + relay_url: Option, + vastai_api_key: Option, + vastai_bootstrap_command: Option, + vastai_disk_gb: Option, + vastai_disk_gb_raw: Option, + vastai_ssh_user: Option, + vastai_confirm_lease: Option, + vastai_confirm_lease_raw: Option, + vastai_onstart: Option, + vastai_ssh_identity_raw: Option, + vastai_gpu_name: Option, + vastai_min_gpu_ram_mb: Option, + vastai_min_gpu_ram_mb_raw: Option, + vastai_min_down_mbps: Option, + vastai_min_down_mbps_raw: Option, + vastai_min_up_mbps: Option, + vastai_min_up_mbps_raw: Option, + vastai_min_reliability: Option, + vastai_min_reliability_raw: Option, + vastai_require_verified: Option, + vastai_require_verified_raw: Option, + vastai_poll_interval_secs: Option, + vastai_poll_interval_secs_raw: Option, + cached_model_host_path: Option, + datastream_frame_log: Option, + worker_bin: Option, +} + +impl ConfigBuilder { + fn hardcoded_defaults() -> Self { + Self { + config_profile: RuntimeConfigProfile::Local, + provider: None, + image: DEFAULT_IMAGE.to_owned(), + toml_vastai_image: None, + image_overridden_after_toml: false, + docker_gpus: "all".to_owned(), + rpc_bind: DEFAULT_RPC_BIND.to_owned(), + rpc_bind_label: "MVP_PROMPT_RPC_BIND", + run_id: 1, + node_id: 1, + stage_index: 0, + layer_end_exclusive: None, + pipeline_stages: 1, + model_id: DEFAULT_MODEL_ID.to_owned(), + gguf_source: GgufSource::HuggingFaceGguf { + repo: DEFAULT_HF_REPO.to_owned(), + file: DEFAULT_HF_FILE.to_owned(), + revision: None, + }, + tokenizer: TokenizerSource::EmbeddedGguf, + default_max_tokens: DEFAULT_MAX_TOKENS, + dashboard: false, + max_context: None, + relay_mode: None, + relay_url: None, + vastai_api_key: None, + vastai_bootstrap_command: None, + vastai_disk_gb: None, + vastai_disk_gb_raw: None, + vastai_ssh_user: None, + vastai_confirm_lease: None, + vastai_confirm_lease_raw: None, + vastai_onstart: None, + vastai_ssh_identity_raw: None, + vastai_gpu_name: None, + vastai_min_gpu_ram_mb: None, + vastai_min_gpu_ram_mb_raw: None, + vastai_min_down_mbps: None, + vastai_min_down_mbps_raw: None, + vastai_min_up_mbps: None, + vastai_min_up_mbps_raw: None, + vastai_min_reliability: None, + vastai_min_reliability_raw: None, + vastai_require_verified: None, + vastai_require_verified_raw: None, + vastai_poll_interval_secs: None, + vastai_poll_interval_secs_raw: None, + cached_model_host_path: None, + datastream_frame_log: None, + worker_bin: None, + } + } + + fn overlay_toml(mut self, overlay: TomlConfigOverlay) -> Result { + if let Some(profile) = overlay.runtime.profile { + self.config_profile = RuntimeConfigProfile::parse(&profile)?; + } + if let Some(run_id) = overlay.runtime.run_id { + self.run_id = run_id; + } + if let Some(node_id) = overlay.runtime.node_id { + self.node_id = node_id; + } + if let Some(stage_index) = overlay.runtime.stage_index { + self.stage_index = stage_index; + } + if let Some(layer_end_exclusive) = overlay.runtime.layer_end_exclusive { + self.layer_end_exclusive = Some(layer_end_exclusive); + } + if let Some(pipeline_stages) = overlay.runtime.pipeline_stages { + self.pipeline_stages = pipeline_stages; + } + if let Some(provider) = overlay.provider.kind { + self.provider = Some(ProviderKind::parse_deploy(&provider)?); + } + if let Some(image) = overlay.image.node { + self.image = image; + } + if let Some(mode) = overlay.relay.mode { + self.relay_mode = Some(mode); + } + if let Some(url) = overlay.relay.url { + self.relay_url = Some(url); + } + if let Some(rpc_bind) = overlay.prompt.rpc_addr { + self.rpc_bind = rpc_bind; + self.rpc_bind_label = "[prompt].rpc_addr"; + } + if let Some(max_tokens) = overlay.prompt.max_tokens { + self.default_max_tokens = max_tokens; + } + if let Some(dashboard) = overlay.prompt.dashboard { + self.dashboard = dashboard; + } + if let Some(model_id) = overlay.model.id { + self.model_id = model_id; + } + if let Some(path) = overlay.model.gguf_local_path { + self.gguf_source = GgufSource::LocalPath(path); + } + if let Some(repo) = overlay.model.gguf_repo { + self.set_gguf_repo(repo); + } + if let Some(file) = overlay.model.gguf_file { + self.set_gguf_file(file); + } + if let Some(revision) = overlay.model.gguf_revision { + self.set_gguf_revision(Some(revision)); + } + if let Some(path) = overlay.model.tokenizer_local_path { + self.tokenizer = TokenizerSource::LocalPath(path); + } + if let Some(max_context) = overlay.model.max_context { + self.max_context = Some(max_context); + } + if let Some(gpus) = overlay.docker.gpus { + self.docker_gpus = gpus; + } + if let Some(path) = overlay.docker.cached_model_host_path { + self.cached_model_host_path = Some(PathBuf::from(path)); + } + if let Some(path) = overlay.observability.datastream_frame_log { + self.datastream_frame_log = Some(PathBuf::from(path)); + } + if let Some(image) = overlay.vastai.image { + self.toml_vastai_image = Some(image); + } + if let Some(api_key) = overlay.vastai.api_key { + self.vastai_api_key = Some(api_key); + } + if let Some(command) = overlay.vastai.bootstrap_command { + self.vastai_bootstrap_command = Some(command); + } + if let Some(disk_gb) = overlay.vastai.disk_gb { + self.vastai_disk_gb = Some(disk_gb); + } + if let Some(ssh_user) = overlay.vastai.ssh_user { + self.vastai_ssh_user = Some(ssh_user); + } + if let Some(confirm_lease) = overlay.vastai.confirm_lease { + self.vastai_confirm_lease = Some(confirm_lease); + } + if let Some(onstart) = overlay.vastai.onstart { + self.vastai_onstart = Some(onstart); + } + if let Some(identity) = overlay.vastai.ssh_identity { + self.vastai_ssh_identity_raw = Some(identity); + } + if let Some(gpu_name) = overlay.vastai.gpu_name { + self.vastai_gpu_name = Some(gpu_name); + } + if let Some(min_gpu_ram_mb) = overlay.vastai.min_gpu_ram_mb { + self.vastai_min_gpu_ram_mb = Some(min_gpu_ram_mb); + } + if let Some(min_down_mbps) = overlay.vastai.min_down_mbps { + self.vastai_min_down_mbps = Some(min_down_mbps); + } + if let Some(min_up_mbps) = overlay.vastai.min_up_mbps { + self.vastai_min_up_mbps = Some(min_up_mbps); + } + if let Some(min_reliability) = overlay.vastai.min_reliability { + self.vastai_min_reliability = Some(min_reliability); + } + if let Some(require_verified) = overlay.vastai.require_verified { + self.vastai_require_verified = Some(require_verified); + } + if let Some(poll_interval_secs) = overlay.vastai.poll_interval_secs { + self.vastai_poll_interval_secs = Some(poll_interval_secs); + } + Ok(self) + } + + fn overlay_env(mut self) -> Result { + if let Some(profile) = env_optional(MVP_RUNTIME_CONFIG_ENV) { + self.config_profile = RuntimeConfigProfile::parse(&profile)?; + } + if let Some(run_id) = env_optional("MVP_RUN_ID") { + self.run_id = Self::parse_value("MVP_RUN_ID", &run_id)?; + } + if let Some(node_id) = env_optional("MVP_LOGICAL_NODE_ID") { + self.node_id = Self::parse_value("MVP_LOGICAL_NODE_ID", &node_id)?; + } + if let Some(stage_index) = env_optional("MVP_STAGE_INDEX") { + self.stage_index = Self::parse_value("MVP_STAGE_INDEX", &stage_index)?; + } + if let Some(layer_end_exclusive) = env_optional("MVP_LAYER_END_EXCLUSIVE") { + self.layer_end_exclusive = Some(Self::parse_value( + "MVP_LAYER_END_EXCLUSIVE", + &layer_end_exclusive, + )?); + } + if let Some(pipeline_stages) = env_optional("MVP_PIPELINE_STAGES") { + self.pipeline_stages = Self::parse_value("MVP_PIPELINE_STAGES", &pipeline_stages)?; + } + 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") { + self.set_process_image(image); + } + if let Some(gpus) = env_optional("MVP_DOCKER_GPUS") { + self.docker_gpus = gpus; + } + if let Some(path) = env_optional(CACHED_MODEL_HOST_ENV) { + self.cached_model_host_path = Some(PathBuf::from(path)); + } + if let Some(path) = env_optional(MVP_WORKER_BIN_ENV) { + self.worker_bin = Some(PathBuf::from(path)); + } + if let Some(rpc_bind) = env_optional("MVP_PROMPT_RPC_BIND") { + self.rpc_bind = rpc_bind; + self.rpc_bind_label = "MVP_PROMPT_RPC_BIND"; + } + if let Some(max_tokens) = env_optional("MVP_PROMPT_MAX_TOKENS") { + self.default_max_tokens = Self::parse_value("MVP_PROMPT_MAX_TOKENS", &max_tokens)?; + } + if let Some(dashboard) = env_optional("MVP_DASHBOARD") { + self.dashboard = Self::parse_bool("MVP_DASHBOARD", &dashboard)?; + } + if let Some(path) = env_optional(DATASTREAM_FRAME_LOG_ENV) { + self.datastream_frame_log = Some(PathBuf::from(path)); + } + if let Some(model_id) = env_optional("MVP_MODEL_ID") { + self.model_id = model_id; + } + if let Some(path) = env_optional("MVP_GGUF_LOCAL_PATH") { + self.gguf_source = GgufSource::LocalPath(path); + } + if let Some(repo) = env_optional("MVP_GGUF_REPO") { + self.set_gguf_repo(repo); + } + if let Some(file) = env_optional("MVP_GGUF_FILE") { + self.set_gguf_file(file); + } + if let Some(revision) = env_optional("MVP_GGUF_REVISION") { + self.set_gguf_revision(Some(revision)); + } + if let Some(path) = env_optional("MVP_TOKENIZER_LOCAL_PATH") { + self.tokenizer = TokenizerSource::LocalPath(path); + } + if let Some(max_context) = env_optional("MVP_MAX_CONTEXT") { + self.max_context = Some(Self::parse_value("MVP_MAX_CONTEXT", &max_context)?); + } + if let Some(mode) = env_optional("MVP_IROH_RELAY_MODE") { + self.relay_mode = Some(mode.to_ascii_lowercase()); + } + if let Some(url) = env_optional(MVP_IROH_RELAY_URL_ENV) + .or_else(|| env_optional(SWACTOR_IROH_RELAY_URL_ENV)) + { + self.relay_url = Some(url); + } + 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") { + self.vastai_bootstrap_command = Some(command); + } + if let Some(identity) = env_optional("MVP_VASTAI_SSH_IDENTITY") { + self.vastai_ssh_identity_raw = Some(identity); + } + if let Some(disk_gb) = env_optional("MVP_VASTAI_DISK_GB") { + self.vastai_disk_gb_raw = Some(disk_gb); + } + if let Some(ssh_user) = env_optional("MVP_VASTAI_SSH_USER") { + self.vastai_ssh_user = Some(ssh_user); + } + if let Some(confirm_lease) = env_optional("MVP_VASTAI_CONFIRM_LEASE") { + self.vastai_confirm_lease_raw = Some(confirm_lease); + } + if let Some(onstart) = env_optional("MVP_VASTAI_ONSTART") { + self.vastai_onstart = Some(onstart); + } + if let Some(gpu_name) = env_optional("MVP_VASTAI_GPU_NAME") { + self.vastai_gpu_name = Some(gpu_name); + } + if let Some(min_gpu_ram_mb) = env_optional("MVP_VASTAI_MIN_GPU_RAM_MB") { + self.vastai_min_gpu_ram_mb_raw = Some(min_gpu_ram_mb); + } + if let Some(min_down_mbps) = env_optional("MVP_VASTAI_MIN_DOWN_MBPS") { + self.vastai_min_down_mbps_raw = Some(min_down_mbps); + } + if let Some(min_up_mbps) = env_optional("MVP_VASTAI_MIN_UP_MBPS") { + self.vastai_min_up_mbps_raw = Some(min_up_mbps); + } + if let Some(min_reliability) = env_optional("MVP_VASTAI_MIN_RELIABILITY") { + self.vastai_min_reliability_raw = Some(min_reliability); + } + if let Some(require_verified) = env_optional("MVP_VASTAI_REQUIRE_VERIFIED") { + self.vastai_require_verified_raw = Some(require_verified); + } + if let Some(poll_interval_secs) = env_optional("MVP_VASTAI_POLL_INTERVAL_SECS") { + self.vastai_poll_interval_secs_raw = Some(poll_interval_secs); + } + Ok(self) + } + + fn overlay_cli(mut self, args: impl IntoIterator) -> Result { + let mut args = args.into_iter(); + while let Some(arg) = args.next() { + match arg.as_str() { + "--runtime-config" => { + self.config_profile = + RuntimeConfigProfile::parse(&next_arg(&mut args, "--runtime-config")?)? + } + "--provider" => { + self.provider = Some(ProviderKind::parse_deploy(&next_arg( + &mut args, + "--provider", + )?)?) + } + "--worker-bin" => { + self.worker_bin = Some(PathBuf::from(next_arg(&mut args, "--worker-bin")?)) + } + "--image" => self.set_process_image(next_arg(&mut args, "--image")?), + "--gpus" => self.docker_gpus = next_arg(&mut args, "--gpus")?, + "--rpc-bind" => { + self.rpc_bind = next_arg(&mut args, "--rpc-bind")?; + self.rpc_bind_label = "--rpc-bind"; + } + "--run-id" => self.run_id = parse_next(&mut args, "--run-id")?, + "--node-id" => self.node_id = parse_next(&mut args, "--node-id")?, + "--stage-index" => self.stage_index = parse_next(&mut args, "--stage-index")?, + "--layer-end-exclusive" => { + self.layer_end_exclusive = Some(parse_next(&mut args, "--layer-end-exclusive")?) + } + "-N" | "--pipeline-stages" => { + self.pipeline_stages = parse_next(&mut args, arg.as_str())? + } + "--max-tokens" => self.default_max_tokens = parse_next(&mut args, "--max-tokens")?, + "--dashboard" => self.dashboard = true, + "--no-dashboard" => self.dashboard = false, + "--datastream-frame-log" => { + self.datastream_frame_log = Some(PathBuf::from(next_arg( + &mut args, + "--datastream-frame-log", + )?)); + } + "--model-id" => self.model_id = next_arg(&mut args, "--model-id")?, + "--gguf-local-path" => { + self.gguf_source = + GgufSource::LocalPath(next_arg(&mut args, "--gguf-local-path")?) + } + "--gguf-repo" => self.set_gguf_repo(next_arg(&mut args, "--gguf-repo")?), + "--gguf-file" => self.set_gguf_file(next_arg(&mut args, "--gguf-file")?), + "--gguf-revision" => { + self.set_gguf_revision(Some(next_arg(&mut args, "--gguf-revision")?)) + } + "--tokenizer-local-path" => { + self.tokenizer = + TokenizerSource::LocalPath(next_arg(&mut args, "--tokenizer-local-path")?) + } + "--max-context" => self.max_context = Some(parse_next(&mut args, "--max-context")?), + "--cached-model-host-path" => { + self.cached_model_host_path = Some(PathBuf::from(next_arg( + &mut args, + "--cached-model-host-path", + )?)); + } + "--relay-mode" => self.relay_mode = Some(next_arg(&mut args, "--relay-mode")?), + "--relay-url" => self.relay_url = Some(next_arg(&mut args, "--relay-url")?), + "--vastai-api-key" => { + self.vastai_api_key = Some(next_arg(&mut args, "--vastai-api-key")?) + } + "--vastai-bootstrap-command" => { + self.vastai_bootstrap_command = + Some(next_arg(&mut args, "--vastai-bootstrap-command")?) + } + "--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")?); + self.vastai_disk_gb_raw = None; + } + "--vastai-ssh-user" => { + self.vastai_ssh_user = Some(next_arg(&mut args, "--vastai-ssh-user")?) + } + "--vastai-confirm-lease" => { + self.vastai_confirm_lease = Some(true); + self.vastai_confirm_lease_raw = None; + } + "--no-vastai-confirm-lease" => { + self.vastai_confirm_lease = Some(false); + self.vastai_confirm_lease_raw = None; + } + "--vastai-onstart" => { + self.vastai_onstart = Some(next_arg(&mut args, "--vastai-onstart")?) + } + "--vastai-gpu-name" => { + self.vastai_gpu_name = Some(next_arg(&mut args, "--vastai-gpu-name")?) + } + "--vastai-min-gpu-ram-mb" => { + self.vastai_min_gpu_ram_mb = + Some(parse_next(&mut args, "--vastai-min-gpu-ram-mb")?); + self.vastai_min_gpu_ram_mb_raw = None; + } + "--vastai-min-down-mbps" => { + self.vastai_min_down_mbps = + Some(parse_next(&mut args, "--vastai-min-down-mbps")?); + 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_raw = None; + } + "--vastai-min-reliability" => { + self.vastai_min_reliability = + Some(parse_next(&mut args, "--vastai-min-reliability")?); + self.vastai_min_reliability_raw = None; + } + "--vastai-require-verified" => { + self.vastai_require_verified = Some(true); + self.vastai_require_verified_raw = None; + } + "--no-vastai-require-verified" => { + self.vastai_require_verified = Some(false); + self.vastai_require_verified_raw = None; + } + "--vastai-poll-interval-secs" => { + self.vastai_poll_interval_secs = + Some(parse_next(&mut args, "--vastai-poll-interval-secs")?); + self.vastai_poll_interval_secs_raw = None; + } + other => return Err(format!("unknown argument {other:?}")), + } + } + Ok(self) + } + + fn finalize(self) -> Result { + let provider = self + .provider + .unwrap_or_else(|| self.config_profile.default_provider()); + let mut image = self.image.clone(); + if provider == ProviderKind::VastAi && !self.image_overridden_after_toml { + if let Some(vastai_image) = &self.toml_vastai_image { + image = vastai_image.clone(); + } + } + if self.pipeline_stages == 0 { + return Err("--pipeline-stages must be greater than 0".to_owned()); + } + if provider == ProviderKind::VastAi && self.pipeline_stages > 1 { + return Err( + "pipeline stages greater than 1 are only supported with provider=docker; provider=vastai does not support pipelined provisioning yet".to_owned(), + ); + } + let mut cached_model_host_path = self.cached_model_host_path.clone(); + if matches!(provider, ProviderKind::Process | ProviderKind::Docker) + && self.pipeline_stages > 1 + && cached_model_host_path.is_none() + && (gguf_source_is_default_hf(&self.gguf_source) + || gguf_source_matches_default_pipeline_cache(&self.gguf_source)) + { + cached_model_host_path = Some(default_pipeline_cached_model_path()); + } + let cached_model = cached_model_host_path + .map(|path| CachedModelConfig::from_host_path(provider, path)) + .transpose()?; + let mut gguf_source = self.gguf_source.clone(); + if let Some(cached_model) = &cached_model { + gguf_source = GgufSource::LocalPath(cached_model.worker_path(provider)); + } + let relay = relay_runtime_config_from_settings( + self.run_id, + self.relay_mode.as_deref(), + self.relay_url.as_deref(), + )?; + let vastai = if provider == ProviderKind::VastAi { + Some(VastAiRuntimeConfig::from_builder(&self)?) + } else { + None + }; + Ok(Config { + config_profile: self.config_profile, + image, + docker_gpus: self.docker_gpus, + provider, + rpc_bind: self + .rpc_bind + .parse() + .map_err(|e| format!("invalid {}: {e}", self.rpc_bind_label))?, + run_id: self.run_id, + node_id: self.node_id, + stage_index: self.stage_index, + layer_end_exclusive: self.layer_end_exclusive, + pipeline_stages: self.pipeline_stages, + model_id: self.model_id, + gguf_source, + tokenizer: self.tokenizer, + default_max_tokens: self.default_max_tokens, + dashboard: self.dashboard, + max_context: self.max_context, + relay, + vastai, + cached_model, + worker_bin: self.worker_bin, + datastream_frame_log: self.datastream_frame_log, + }) + } + + fn set_process_image(&mut self, image: String) { + self.image = image; + self.image_overridden_after_toml = true; + } + + fn set_gguf_repo(&mut self, repo: String) { + let (file, revision) = match &self.gguf_source { + GgufSource::HuggingFaceGguf { file, revision, .. } => (file.clone(), revision.clone()), + GgufSource::LocalPath(_) => (DEFAULT_HF_FILE.to_owned(), None), + }; + self.gguf_source = GgufSource::HuggingFaceGguf { + repo, + file, + revision, + }; + } + + fn set_gguf_file(&mut self, file: String) { + let (repo, revision) = match &self.gguf_source { + GgufSource::HuggingFaceGguf { repo, revision, .. } => (repo.clone(), revision.clone()), + GgufSource::LocalPath(_) => (DEFAULT_HF_REPO.to_owned(), None), + }; + self.gguf_source = GgufSource::HuggingFaceGguf { + repo, + file, + revision, + }; + } + + fn set_gguf_revision(&mut self, revision: Option) { + let (repo, file) = match &self.gguf_source { + GgufSource::HuggingFaceGguf { repo, file, .. } => (repo.clone(), file.clone()), + GgufSource::LocalPath(_) => (DEFAULT_HF_REPO.to_owned(), DEFAULT_HF_FILE.to_owned()), + }; + self.gguf_source = GgufSource::HuggingFaceGguf { + repo, + file, + revision, + }; + } + + fn parse_value(name: &str, value: &str) -> Result + where + T: std::str::FromStr, + T::Err: std::fmt::Display, + { + value + .parse::() + .map_err(|e| format!("invalid {name}={value:?}: {e}")) + } + + fn parse_bool(name: &str, value: &str) -> Result { + match value.to_ascii_lowercase().as_str() { + "1" | "true" | "yes" | "on" => Ok(true), + "0" | "false" | "no" | "off" => Ok(false), + _ => Err(format!( + "invalid {name}={value:?}; use 1/0, true/false, yes/no, or on/off" + )), + } + } +} + +impl Config { + fn from_defaults_toml_env_args(args: impl IntoIterator) -> Result { + Self::from_layers_with_path_and_args(Some(Path::new(DEFAULT_CONFIG_PATH)), args) + } + + fn from_layers_with_path_and_args( + path: Option<&Path>, + args: impl IntoIterator, + ) -> Result { + let mut builder = Self::hardcoded_defaults(); + if let Some(path) = path { + if let Some(overlay) = TomlConfigOverlay::load_optional(path)? { + builder = builder.overlay_toml(overlay)?; + } + } + builder.overlay_env()?.overlay_cli(args)?.finalize() + } + + fn hardcoded_defaults() -> ConfigBuilder { + ConfigBuilder::hardcoded_defaults() + } + + fn uses_planned_execution(&self) -> bool { + self.cached_model.is_some() || self.pipeline_stages > 1 + } + + fn provider_datastream_detail(&self) -> Value { + match self.provider { + ProviderKind::Process => json!({ + "worker_bin": self.worker_bin.as_ref().map(|path| path.to_string_lossy().to_string()), + "cached_model": self.cached_model.as_ref().map(CachedModelConfig::datastream_detail), + }), + ProviderKind::Docker => json!({ + "docker_gpus": &self.docker_gpus, + "cached_model": self.cached_model.as_ref().map(CachedModelConfig::datastream_detail), + }), + ProviderKind::VastAi => self + .vastai + .as_ref() + .map_or_else(|| json!({}), VastAiRuntimeConfig::datastream_detail), + ProviderKind::Mock => json!({}), + } + } + + fn build_run_plan(&self) -> Result { + let host_path = self.local_planning_gguf_path()?; + let metadata = crate::gguf_metadata::read_gguf_planning_metadata(&host_path)?; + let model = metadata.to_model_facts( + self.model_id.clone(), + self.gguf_source.clone(), + self.tokenizer.clone(), + self.max_context, + )?; + if self.pipeline_stages > model.num_layers { + return Err(format!( + "--pipeline-stages={} exceeds GGUF layer count {}; choose N <= {}", + self.pipeline_stages, model.num_layers, model.num_layers + )); + } + + let activation_extent = model + .max_seq_len + .checked_mul(model.hidden_dim) + .and_then(|value| value.checked_mul(model.dtype_width_bytes)) + .ok_or_else(|| "activation ring size overflow while planning pipeline".to_owned())?; + let activation_data_capacity = run_plan::MO01_HEADER_BYTES + .checked_add(activation_extent) + .ok_or_else(|| { + "activation ring data capacity overflow while planning pipeline".to_owned() + })?; + let token_extent = model + .max_seq_len + .checked_mul(4) + .ok_or_else(|| "token ring size overflow while planning pipeline".to_owned())?; + let token_data_capacity = run_plan::MO01_HEADER_BYTES + .checked_add(token_extent) + .ok_or_else(|| { + "token ring data capacity overflow while planning pipeline".to_owned() + })?; + let candidate_pool = (0..self.pipeline_stages) + .map(|stage_index| run_plan::NodeId(self.node_id + 1 + u64::from(stage_index))) + .collect::>(); + let placement = run_plan::PlacementInput::FixedLinear( + candidate_pool + .iter() + .enumerate() + .map(|(stage_index, node_id)| run_plan::StagePlacement { + stage_index: stage_index as u32, + node_id: *node_id, + }) + .collect(), + ); + + run_plan::plan_run(run_plan::PlannerInput { + run_id: run_plan::RunId(self.run_id), + orchestrator_node_id: run_plan::NodeId(self.node_id), + model, + runtime: run_plan::RuntimeConfig { + max_tokens: self.default_max_tokens, + prompt: run_plan::PromptSource::Inline(String::new()), + sampling: run_plan::SamplingPolicy { + temperature_millis: 0, + top_k: 1, + }, + token_output_policy: run_plan::TokenOutputPolicy::EmitAll, + }, + candidate_pool, + stage_count: self.pipeline_stages, + placement, + activation_ring: run_plan::RingSpec { + data_capacity: activation_data_capacity, + alignment: 64, + direction: run_plan::RingDirection::Egress, + host_pinning: run_plan::HostPinning::Pageable, + wake_coalescing: run_plan::WakeCoalescing::PendingBit, + }, + token_ring: run_plan::RingSpec { + data_capacity: token_data_capacity, + alignment: 8, + direction: run_plan::RingDirection::Egress, + host_pinning: run_plan::HostPinning::Pageable, + wake_coalescing: run_plan::WakeCoalescing::PendingBit, + }, + }) + .map_err(|e| format!("plan pipeline run: {:?}", e.kind())) + } + + fn local_planning_gguf_path(&self) -> Result { + if let Some(cached_model) = &self.cached_model { + return Ok(cached_model.host_path.clone()); + } + match &self.gguf_source { + GgufSource::LocalPath(path) => { + let host_path = PathBuf::from(path); + if host_path.is_file() { + Ok(host_path) + } else { + Err(format!( + "local pipeline requires a locally inspectable GGUF before provisioning; {path:?} is not a host file, so use --cached-model-host-path" + )) + } + } + GgufSource::HuggingFaceGguf { repo, file, .. } => Err(format!( + "local pipeline requires a locally inspectable GGUF before provisioning; selected source {repo}/{file} is remote, so use --cached-model-host-path" + )), + } + } + + fn prepare_vastai_ssh_key(&mut self) -> Result<(), String> { + if self.provider != ProviderKind::VastAi { + return Ok(()); + } + + let api_key = self + .vastai + .as_ref() + .and_then(|vastai| vastai.api_key.as_deref()) + .ok_or_else(|| { + "MVP_VASTAI_API_KEY or VASTAI_API_KEY is required when MVP_NODE_PROVIDER=vastai" + .to_owned() + })? + .to_owned(); + let identity = resolve_vastai_ssh_identity( + self.vastai + .as_ref() + .and_then(|vastai| vastai.ssh_identity.clone()), + )?; + if !identity.is_file() { + return Err(format!( + "missing VastAI SSH identity {}; create/register one with vastai create ssh-key or set MVP_VASTAI_SSH_IDENTITY", + identity.display() + )); + } + let public_key = derive_ssh_public_key(&identity)?; + let fingerprint = ssh_public_key_fingerprint(&public_key); + ensure_vastai_account_ssh_key(&api_key, &public_key)?; + + eprintln!( + "VastAI SSH identity {} fingerprint {} registered for account", + identity.display(), + fingerprint + ); + + let vastai = self + .vastai + .as_mut() + .expect("VastAI config exists when provider is vastai"); + vastai.ssh_identity = Some(identity); + vastai.provisioning.ssh_public_key = Some(public_key.clone()); + vastai.ssh_public_key = Some(public_key); + vastai.ssh_public_fingerprint = Some(fingerprint); + Ok(()) + } + + fn build_provisioner( + &self, + bootstrap_runtime: Arc, + ) -> Result, String> { + match self.provider { + ProviderKind::Process => { + let worker_bin = self.worker_bin.clone().unwrap_or(default_worker_bin()?); + if !worker_bin.is_file() { + return Err(format!( + "local process worker binary does not exist: {}", + worker_bin.display() + )); + } + Ok(Box::new(LocalProcessPlugin::new(worker_bin))) + } + ProviderKind::Docker => Ok(Box::new(LocalDockerPlugin::new(docker_container_prefix()))), + ProviderKind::VastAi => { + let vastai = self.vastai.as_ref().ok_or_else(|| { + "VastAI config was not resolved for provider vastai".to_owned() + })?; + if vastai.bootstrap_command.is_none() { + return Err( + "MVP_VASTAI_BOOTSTRAP_COMMAND is required when MVP_NODE_PROVIDER=vastai" + .to_owned(), + ); + } + let api_key = vastai.api_key.clone().ok_or_else(|| { + "MVP_VASTAI_API_KEY or VASTAI_API_KEY is required when MVP_NODE_PROVIDER=vastai" + .to_owned() + })?; + let ssh_identity = vastai + .ssh_identity + .clone() + .ok_or_else(|| "VastAI SSH identity was not prepared".to_owned())?; + let client = ToolsVastAiLeaseClient::from_api_key(api_key)?; + Ok(Box::new(VastAiProvisioningPlugin::new( + client, + SshCommandBootstrapLauncher::new(Some(ssh_identity), bootstrap_runtime), + vastai.provisioning.clone(), + ))) + } + ProviderKind::Mock => Err("mvp-orchestrator does not support mock provider".to_owned()), + } + } + + fn node_spec_env_keys(&self) -> Vec<&'static str> { + let mut keys = vec![ + "MVP_RUN_ID", + "MVP_LOGICAL_NODE_ID", + "MVP_NODE_PROVIDER", + "MVP_STAGE_INDEX", + "MVP_COORDINATOR_ENDPOINT", + "MVP_ORCHESTRATOR_ACTOR", + "MVP_MODEL_ID", + "MVP_IROH_RELAY_MODE", + "MVP_PIPELINE_STAGES", + ]; + if self.relay.url.is_some() { + keys.push(MVP_IROH_RELAY_URL_ENV); + } + if self.provider == ProviderKind::Docker { + keys.push("MVP_DOCKER_GPUS"); + } + if std::env::var_os("DEV").is_some() { + keys.push("DEV"); + } + if local_tinygrad_worker_env(self.provider).is_some() { + keys.push("MVP_TINYGRAD_WORKER"); + } + if std::env::var_os("MVP_CPU_LINE_PROFILE").is_some() { + keys.push("MVP_CPU_LINE_PROFILE"); + } + if std::env::var_os("MVP_CPU_LINE_PROFILE_INTERVAL_MS").is_some() { + keys.push("MVP_CPU_LINE_PROFILE_INTERVAL_MS"); + } + if std::env::var_os("MVP_TOKEN_PROGRESS_EVERY").is_some() { + keys.push("MVP_TOKEN_PROGRESS_EVERY"); + } + if std::env::var_os("CUDA_DEVICE_SCHEDULE").is_some() { + keys.push("CUDA_DEVICE_SCHEDULE"); + } + if std::env::var_os("MVP_MODEL_CACHE_DIR").is_some() { + keys.push("MVP_MODEL_CACHE_DIR"); + } + if std::env::var_os("HF_TOKEN").is_some() { + keys.push("HF_TOKEN"); + } + match &self.gguf_source { + GgufSource::LocalPath(_) => keys.push("MVP_GGUF_LOCAL_PATH"), + GgufSource::HuggingFaceGguf { revision, .. } => { + keys.push("MVP_GGUF_REPO"); + keys.push("MVP_GGUF_FILE"); + if revision.is_some() { + keys.push("MVP_GGUF_REVISION"); + } + } + } + if matches!(self.tokenizer, TokenizerSource::LocalPath(_)) { + keys.push("MVP_TOKENIZER_LOCAL_PATH"); + } + if self.max_context.is_some() { + keys.push("MVP_MAX_CONTEXT"); + } + keys + } + + fn node_spec( + &self, + coordinator: EndpointAddr, + orchestrator_actor: ActorAddress, + ) -> Result { + self.node_spec_for_stage( + coordinator, + orchestrator_actor, + self.node_id, + self.stage_index, + ) + } + + fn node_spec_for_stage( + &self, + coordinator: EndpointAddr, + orchestrator_actor: ActorAddress, + logical_node_id: u64, + stage_index: u32, + ) -> Result { + let mut env = vec![ + ("MVP_RUN_ID".to_owned(), self.run_id.to_string()), + ( + "MVP_LOGICAL_NODE_ID".to_owned(), + logical_node_id.to_string(), + ), + ("MVP_STAGE_INDEX".to_owned(), stage_index.to_string()), + ( + "MVP_PIPELINE_STAGES".to_owned(), + self.pipeline_stages.to_string(), + ), + ( + "MVP_NODE_PROVIDER".to_owned(), + self.provider.as_str().to_owned(), + ), + ( + "MVP_COORDINATOR_ENDPOINT".to_owned(), + serde_json::to_string(&coordinator) + .map_err(|e| format!("serialize coordinator endpoint: {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(), + relay_mode_env_value(&self.relay.mode).to_owned(), + ), + ]; + if let Some(url) = &self.relay.url { + env.push((MVP_IROH_RELAY_URL_ENV.to_owned(), url.clone())); + } + if self.provider == ProviderKind::Docker { + env.push(("MVP_DOCKER_GPUS".to_owned(), self.docker_gpus.clone())); + } + env.extend(optional_env("DEV")); + env.extend(local_tinygrad_worker_env(self.provider)); + env.extend(optional_env("MVP_CPU_LINE_PROFILE")); + env.extend(optional_env("MVP_CPU_LINE_PROFILE_INTERVAL_MS")); + env.extend(optional_env("MVP_TOKEN_PROGRESS_EVERY")); + env.extend(optional_env("CUDA_DEVICE_SCHEDULE")); + env.extend(optional_env("MVP_MODEL_CACHE_DIR")); + env.extend(optional_env("HF_TOKEN")); + match &self.gguf_source { + GgufSource::LocalPath(path) => { + env.push(("MVP_GGUF_LOCAL_PATH".to_owned(), path.clone())) + } + GgufSource::HuggingFaceGguf { + repo, + file, + revision, + } => { + env.push(("MVP_GGUF_REPO".to_owned(), repo.clone())); + env.push(("MVP_GGUF_FILE".to_owned(), file.clone())); + if let Some(revision) = revision { + env.push(("MVP_GGUF_REVISION".to_owned(), revision.clone())); + } + } + } + if let TokenizerSource::LocalPath(path) = &self.tokenizer { + env.push(("MVP_TOKENIZER_LOCAL_PATH".to_owned(), path.clone())); + } + if let Some(max_context) = self.max_context { + env.push(("MVP_MAX_CONTEXT".to_owned(), max_context.to_string())); + } + let args = match self.provider { + ProviderKind::VastAi => self + .vastai + .as_ref() + .and_then(|vastai| vastai.bootstrap_command.clone()) + .into_iter() + .collect(), + ProviderKind::Process | ProviderKind::Docker => Vec::new(), + ProviderKind::Mock => { + return Err("mvp-orchestrator does not support mock provider".to_owned()); + } + }; + let mounts = if self.provider == ProviderKind::Docker { + self.cached_model + .as_ref() + .map(|cached_model| { + vec![ProviderMount { + host_path: cached_model.host_path.to_string_lossy().to_string(), + container_path: cached_model.container_path.clone(), + readonly: self.uses_planned_execution(), + }] + }) + .unwrap_or_default() + } else { + Vec::new() + }; + Ok(NodeProvisionSpec { + run_id: self.run_id, + node_id: logical_node_id, + stage_index: Some(stage_index), + image: self.image.clone(), + env, + args, + mounts, + }) + } +} + +#[derive(Clone)] +struct RuntimeReady { + endpoint: EndpointAddr, + node_actor: ActorAddress, + datastream_publisher: ActorAddress, + stage_index: u32, + readiness_id: u64, + swim_node_id: DistNodeId, +} + +#[derive(Clone)] +struct PromptRuntimeReady { + first_stage: RuntimeReady, + final_stage: RuntimeReady, +} +#[derive(Clone)] +struct RuntimeReadyAckTarget { + node_id: u64, + ready: RuntimeReady, +} + +fn runtime_ready_barrier_met(stack: &DistributionRuntimeStack, ready: &RuntimeReady) -> bool { + stack.member_state(ready.swim_node_id) == Some(MemberState::Alive) + && stack.route_owner(ready.node_actor) == Some(ready.swim_node_id) +} + +fn enqueue_runtime_ready_ack( + stack: &DistributionRuntimeStack, + ready: &RuntimeReady, + run_id: u64, + node_id: u64, +) -> Result<(), String> { + stack + .runtime + .send_to( + ready.node_actor, + NodeAgentMsg::RuntimeReadyAck { + run_id, + node_id, + stage_index: ready.stage_index, + readiness_id: ready.readiness_id, + }, + ) + .map_err(|e| format!("send runtime ready ack: {e}")) +} + +fn enqueue_datastream_subscribe( + driver: &mut IrohDriver, + stack: &DistributionRuntimeStack, + ready: &RuntimeReady, + run_id: u64, + node_id: u64, +) -> Result<(), String> { + let mut flow_id = [0_u8; 16]; + flow_id[..8].copy_from_slice(&run_id.to_le_bytes()); + flow_id[8..].copy_from_slice(&node_id.to_le_bytes()); + stack + .runtime + .send_to( + ready.datastream_publisher, + DatastreamPublisherMsg::Subscribe(DatastreamSubscribe { + collector: driver.endpoint_addr(), + request: SubscriptionRequest::all(), + flow_id, + token: Vec::new(), + }), + ) + .map_err(|e| format!("send datastream subscribe: {e}")) +} + +#[allow(clippy::too_many_arguments)] +fn wait_for_runtime_ready_acks( + driver: &mut IrohDriver, + stack: &DistributionRuntimeStack, + obs_rx: &mpsc::Receiver, + frame_rx: &mpsc::Receiver, + frame_tx: &mpsc::Sender, + orchestrator_reports: &swactor::runtime::Inbox, + stop_rx: &mpsc::Receiver<()>, + dashboard: Option<&DashboardSupport>, + orch_datastream: &mut OrchDatastream, + orch_stdio_rx: Option<&mpsc::Receiver>, + run_id: u64, + orchestrator_node_id: u64, + provider: ProviderKind, + targets: &[RuntimeReadyAckTarget], +) -> Result<(), String> { + let mut pending = targets + .iter() + .cloned() + .map(|target| { + ( + ( + target.node_id, + target.ready.stage_index, + target.ready.readiness_id, + ), + target, + ) + }) + .collect::>(); + let mut attempts = BTreeMap::<(u64, u32, u64), u64>::new(); + let started = Instant::now(); + let mut last_send = None::; + + while !pending.is_empty() { + pump(driver, stack, frame_tx); + drain_orch_stdio_capture( + orch_stdio_rx, + orch_datastream, + dashboard, + run_id, + orchestrator_node_id, + ); + if stop_requested(stop_rx) { + return Err( + "shutdown requested while waiting for runtime-ready acknowledgements".to_owned(), + ); + } + while let Ok(observation) = obs_rx.try_recv() { + emit_plugin_observation(orch_datastream, dashboard, provider, &observation); + match observation { + PluginObservation::DatastreamFrame { .. } => {} + PluginObservation::ProviderLine { .. } + | PluginObservation::StdoutLine { .. } + | PluginObservation::StderrLine { .. } => {} + PluginObservation::Failed { reason, .. } => return Err(reason), + PluginObservation::Exited { + status, node_id, .. + } => { + return Err(format!("node {node_id} exited before ready: {status:?}")); + } + } + } + drain_frames(frame_rx, dashboard, orch_datastream); + while let Some(report) = orchestrator_reports.try_recv() { + let OrchestratorReport::NodeRuntimeReadyAck { + run_id: ack_run_id, + node_id, + stage_index, + readiness_id, + } = report + else { + continue; + }; + if ack_run_id != run_id { + continue; + } + let key = (node_id, stage_index, readiness_id); + let Some(target) = pending.remove(&key) else { + continue; + }; + orch_datastream.emit_bootstrap( + dashboard, + run_id, + orchestrator_node_id, + "runtime_ready_ack", + "ready", + json!({ + "node_id":node_id, + "node_actor":target.ready.node_actor, + "readiness_id":readiness_id, + "attempts":attempts.get(&key).copied().unwrap_or(0), + }), + ); + } + if pending.is_empty() { + return Ok(()); + } + if started.elapsed() >= RUNTIME_READY_ACK_TIMEOUT { + let pending_list = pending + .keys() + .map(|(node_id, stage_index, readiness_id)| { + format!("{node_id}/{stage_index}/{readiness_id}") + }) + .collect::>() + .join(", "); + return Err(format!( + "runtime_ready_ack timed out for node(s): {pending_list}" + )); + } + if last_send.is_none_or(|sent_at| sent_at.elapsed() >= RUNTIME_READY_ACK_RETRY_INTERVAL) { + for (key, target) in &pending { + if stack.route_owner(target.ready.datastream_publisher) + == Some(target.ready.swim_node_id) + && let Err(error) = enqueue_datastream_subscribe( + driver, + stack, + &target.ready, + run_id, + target.node_id, + ) + { + orch_datastream.emit_bootstrap( + dashboard, + run_id, + orchestrator_node_id, + "datastream_subscribe", + "failed", + json!({"node_id":target.node_id,"error":error}), + ); + } + enqueue_runtime_ready_ack(stack, &target.ready, run_id, target.node_id)?; + let attempt = attempts.entry(*key).or_default(); + *attempt += 1; + let attempt = *attempt; + orch_datastream.emit_bootstrap( + dashboard, + run_id, + orchestrator_node_id, + "runtime_ready_ack", + "sent", + json!({ + "node_id":target.node_id, + "node_actor":target.ready.node_actor, + "readiness_id":target.ready.readiness_id, + "attempt":attempt, + }), + ); + } + driver.drain_outbox(&stack.outbox); + last_send = Some(Instant::now()); + } + thread::sleep(PUMP_INTERVAL); + } + Ok(()) +} + +#[cfg(test)] +struct ProvisionedNodeGuard<'a> { + provisioner: &'a mut dyn ProvisionPlugin, + handle: Option, +} + +#[cfg(test)] +impl<'a> ProvisionedNodeGuard<'a> { + fn new( + provisioner: &'a mut dyn ProvisionPlugin, + handle: crate::provisioning::PluginNodeHandle, + ) -> Self { + Self { + provisioner, + handle: Some(handle), + } + } + + fn stop(&mut self) -> Result<(), String> { + let Some(handle) = self.handle.take() else { + return Ok(()); + }; + self.provisioner.stop_node(&handle) + } +} + +#[cfg(test)] +impl Drop for ProvisionedNodeGuard<'_> { + fn drop(&mut self) { + let _ = self.stop(); + } +} + +struct ProvisionedClusterGuard { + provisioner: Box, + handles: Vec, +} + +impl ProvisionedClusterGuard { + fn new( + provisioner: Box, + handles: Vec, + ) -> Self { + Self { + provisioner, + handles, + } + } + + fn complete_bootstrap_all(&mut self) -> Result<(), String> { + for handle in &self.handles { + self.provisioner.complete_bootstrap(handle)?; + } + Ok(()) + } + + fn stop(&mut self) -> Result<(), String> { + let mut first_error = None; + while let Some(handle) = self.handles.pop() { + if let Err(error) = self.provisioner.stop_node(&handle) + && first_error.is_none() + { + first_error = Some(error); + } + } + match first_error { + Some(error) => Err(error), + None => Ok(()), + } + } +} + +impl Drop for ProvisionedClusterGuard { + fn drop(&mut self) { + let _ = self.stop(); + } +} + +fn start_node_with_stdio_capture( + provisioner: Box, + node_spec: NodeProvisionSpec, + sink: PluginSink, + orch_stdio_rx: Option<&mpsc::Receiver>, + dashboard: Option<&DashboardSupport>, + orch_datastream: &mut OrchDatastream, + run_id: u64, + node_id: u64, +) -> ( + Box, + Result, +) { + let (tx, rx) = mpsc::channel(); + thread::spawn(move || { + let mut provisioner = provisioner; + let result = provisioner.start_node(node_spec, sink); + let _ = tx.send((provisioner, result)); + }); + + loop { + match rx.recv_timeout(Duration::from_millis(100)) { + Ok(result) => return result, + Err(mpsc::RecvTimeoutError::Timeout) => { + drain_orch_stdio_capture( + orch_stdio_rx, + orch_datastream, + dashboard, + run_id, + node_id, + ); + } + Err(mpsc::RecvTimeoutError::Disconnected) => { + return ( + Box::new(FailedProvisionPlugin), + Err("provider start worker disconnected".to_owned()), + ); + } + } + } +} + +#[allow(clippy::too_many_arguments)] +fn start_and_provision_workers( + mut provisioner: Box, + config: &Config, + pipeline_plan: Option<&run_plan::RunPlan>, + driver: &mut IrohDriver, + stack: &DistributionRuntimeStack, + obs_rx: &mpsc::Receiver, + frame_rx: &mpsc::Receiver, + frame_tx: &mpsc::Sender, + orchestrator_reports: &swactor::runtime::Inbox, + stop_rx: &mpsc::Receiver<()>, + dashboard: Option<&DashboardSupport>, + orch_datastream: &mut OrchDatastream, + orch_stdio_rx: Option<&mpsc::Receiver>, + sink: PluginSink, + coordinator: EndpointAddr, + pipeline_coordinator: EndpointAddr, + orchestrator_actor: ActorAddress, +) -> Result<(ProvisionedClusterGuard, PromptRuntimeReady), String> { + let stage_specs = stage_node_specs( + config, + pipeline_plan, + coordinator.clone(), + orchestrator_actor, + )?; + let expected_node_ids = stage_specs + .iter() + .map(|spec| spec.node_id) + .collect::>(); + orch_datastream.emit_bootstrap( + dashboard, + config.run_id, + config.node_id, + "node_spec", + "ready", + json!({ + "provider":config.provider.as_str(), + "image":&config.image, + "relay_mode":relay_mode_env_value(&config.relay.mode), + "docker_gpus":if config.provider == ProviderKind::Docker { Some(config.docker_gpus.as_str()) } else { None }, + "provider_config":config.provider_datastream_detail(), + "env_keys":config.node_spec_env_keys(), + "worker_count":stage_specs.len(), + "worker_node_ids":expected_node_ids, + }), + ); + + let mut handles = Vec::with_capacity(stage_specs.len()); + for node_spec in stage_specs { + orch_datastream.emit_event( + dashboard, + ProvisionEvent { + run_id: config.run_id, + node_id: node_spec.node_id, + kind: ProvisionEventKind::ProvisionStart, + provider: Some(config.provider.as_str().to_owned()), + message: Some(format!( + "starting {} image {}", + config.provider.as_str(), + config.image + )), + }, + ); + orch_datastream.emit_bootstrap( + dashboard, + config.run_id, + config.node_id, + "provider_start", + "started", + json!({ + "provider":config.provider.as_str(), + "image":&config.image, + "node_id":node_spec.node_id, + "stage_index":node_spec.stage_index, + }), + ); + let (returned_provisioner, handle_result) = start_node_with_stdio_capture( + provisioner, + node_spec.clone(), + sink.clone(), + orch_stdio_rx, + dashboard, + orch_datastream, + config.run_id, + node_spec.node_id, + ); + provisioner = returned_provisioner; + match handle_result { + Ok(handle) => handles.push(handle), + Err(error) => { + orch_datastream.emit_bootstrap( + dashboard, + config.run_id, + config.node_id, + "provider_start", + "failed", + json!({"provider":config.provider.as_str(),"node_id":node_spec.node_id,"error":error}), + ); + stop_started_nodes(&mut *provisioner, &mut handles); + drain_orch_stdio_capture( + orch_stdio_rx, + orch_datastream, + dashboard, + config.run_id, + config.node_id, + ); + return Err(error); + } + } + } + let mut provisioned_nodes = ProvisionedClusterGuard::new(provisioner, handles); + drain_orch_stdio_capture( + orch_stdio_rx, + orch_datastream, + dashboard, + config.run_id, + config.node_id, + ); + + orch_datastream.emit_bootstrap( + dashboard, + config.run_id, + config.node_id, + "node_runtime_ready", + "started", + json!({"worker_count":expected_node_ids.len(),"node_ids":expected_node_ids}), + ); + let readies = if pipeline_plan.is_some() { + match wait_for_runtime_readies( + driver, + stack, + obs_rx, + frame_rx, + frame_tx, + orchestrator_reports, + stop_rx, + dashboard, + orch_datastream, + orch_stdio_rx, + config.run_id, + &expected_node_ids, + config.provider, + ) { + Ok(readies) => readies, + Err(error) => { + orch_datastream.emit_bootstrap( + dashboard, + config.run_id, + config.node_id, + "node_runtime_ready", + "failed", + json!({"error":error}), + ); + return Err(error); + } + } + } else { + let ready = match wait_for_runtime_ready( + driver, + stack, + obs_rx, + frame_rx, + frame_tx, + orchestrator_reports, + stop_rx, + dashboard, + orch_datastream, + orch_stdio_rx, + config.run_id, + config.node_id, + config.provider, + ) { + Ok(ready) => ready, + Err(error) => { + orch_datastream.emit_bootstrap( + dashboard, + config.run_id, + config.node_id, + "node_runtime_ready", + "failed", + json!({"error":error}), + ); + return Err(error); + } + }; + BTreeMap::from([(config.node_id, ready)]) + }; + for (node_id, ready) in &readies { + orch_datastream.emit_bootstrap( + dashboard, + config.run_id, + config.node_id, + "node_runtime_ready", + "ready", + json!({"endpoint":&ready.endpoint,"node_actor":ready.node_actor,"node_id":node_id,"stage_index":ready.stage_index}), + ); + } + provisioned_nodes.complete_bootstrap_all()?; + let ack_targets = readies + .iter() + .map(|(node_id, ready)| RuntimeReadyAckTarget { + node_id: *node_id, + ready: ready.clone(), + }) + .collect::>(); + wait_for_runtime_ready_acks( + driver, + stack, + obs_rx, + frame_rx, + frame_tx, + orchestrator_reports, + stop_rx, + dashboard, + orch_datastream, + orch_stdio_rx, + config.run_id, + config.node_id, + config.provider, + &ack_targets, + )?; + + orch_datastream.emit_bootstrap( + dashboard, + config.run_id, + config.node_id, + "stage_provision", + "started", + stage_provision_detail(config, pipeline_plan), + ); + if pipeline_plan.is_none() { + let ready = readies + .get(&config.node_id) + .ok_or_else(|| "missing runtime-ready node for single-stage run".to_owned())?; + provision_stage(stack, ready.node_actor, config)?; + } + + orch_datastream.emit_bootstrap( + dashboard, + config.run_id, + config.node_id, + "weights_loaded", + "started", + json!({"model_id":&config.model_id,"expected":expected_node_ids.len()}), + ); + let weights_result = if pipeline_plan.is_some() { + wait_for_weights_loaded_count( + driver, + stack, + obs_rx, + frame_rx, + frame_tx, + orchestrator_reports, + stop_rx, + dashboard, + orch_datastream, + orch_stdio_rx, + config.run_id, + config.node_id, + config.provider, + expected_node_ids.len(), + pipeline_plan.expect("pipeline mode requires plan"), + &readies, + &pipeline_coordinator, + ) + } else { + wait_for_weights_loaded( + driver, + stack, + obs_rx, + frame_rx, + frame_tx, + orchestrator_reports, + stop_rx, + dashboard, + orch_datastream, + orch_stdio_rx, + config.run_id, + config.node_id, + config.provider, + config.stage_index, + ) + }; + match weights_result { + Ok(()) => orch_datastream.emit_bootstrap( + dashboard, + config.run_id, + config.node_id, + "weights_loaded", + "ready", + json!({"source":"actor_stage_ready","model_id":&config.model_id,"expected":expected_node_ids.len()}), + ), + Err(error) => { + orch_datastream.emit_bootstrap( + dashboard, + config.run_id, + config.node_id, + "weights_loaded", + "failed", + json!({"error":error}), + ); + drain_orch_stdio_capture( + orch_stdio_rx, + orch_datastream, + dashboard, + config.run_id, + config.node_id, + ); + return Err(error); + } + } + + let prompt_ready = if let Some(plan) = pipeline_plan { + let first_node_id = plan + .stages + .iter() + .find(|stage| stage.stage_index == 0) + .map(|stage| stage.node_id.0) + .ok_or_else(|| "pipeline plan missing stage 0".to_owned())?; + let final_node_id = plan + .stages + .iter() + .find(|stage| stage.stage_index + 1 == stage.stage_count) + .map(|stage| stage.node_id.0) + .ok_or_else(|| "pipeline plan missing final stage".to_owned())?; + PromptRuntimeReady { + first_stage: readies + .get(&first_node_id) + .cloned() + .ok_or_else(|| "missing stage 0 runtime-ready node".to_owned())?, + final_stage: readies + .get(&final_node_id) + .cloned() + .ok_or_else(|| "missing final-stage runtime-ready node".to_owned())?, + } + } else { + let ready = readies + .get(&config.node_id) + .cloned() + .ok_or_else(|| "missing single-stage runtime-ready node".to_owned())?; + PromptRuntimeReady { + first_stage: ready.clone(), + final_stage: ready, + } + }; + Ok((provisioned_nodes, prompt_ready)) +} + +fn stage_node_specs( + config: &Config, + pipeline_plan: Option<&run_plan::RunPlan>, + coordinator: EndpointAddr, + orchestrator_actor: ActorAddress, +) -> Result, String> { + if let Some(plan) = pipeline_plan { + let mut stages = plan.stages.clone(); + stages.sort_by_key(|stage| stage.stage_index); + stages + .into_iter() + .map(|stage| { + config.node_spec_for_stage( + coordinator.clone(), + orchestrator_actor, + stage.node_id.0, + stage.stage_index, + ) + }) + .collect() + } else { + Ok(vec![config.node_spec(coordinator, orchestrator_actor)?]) + } +} + +fn stop_started_nodes( + provisioner: &mut dyn ProvisionPlugin, + handles: &mut Vec, +) { + while let Some(handle) = handles.pop() { + let _ = provisioner.stop_node(&handle); + } +} + +fn stage_provision_detail(config: &Config, pipeline_plan: Option<&run_plan::RunPlan>) -> Value { + if let Some(plan) = pipeline_plan { + json!({ + "run_id":config.run_id, + "stage_count":plan.stages.len(), + "stages":plan.stages.iter().map(|stage| { + json!({ + "node_id":stage.node_id.0, + "stage_index":stage.stage_index, + "layer_range":{"start":stage.layer_start,"end_exclusive":stage.layer_end_exclusive}, + "inbound_edge_id":stage.inbound_edge.0, + "outbound_edge_id":stage.outbound_edge.0, + }) + }).collect::>(), + "model_id":&config.model_id, + }) + } else { + json!({ + "run_id":config.run_id, + "node_id":config.node_id, + "stage_index":config.stage_index, + "stage_count":1, + "layer_range":{"start":0,"end_exclusive":config.layer_end_exclusive}, + "model_id":&config.model_id, + }) + } +} + +fn stage_provision_wire_from_plan( + plan: &run_plan::RunPlan, + stage_index: u32, + readies: &BTreeMap, + coordinator: &EndpointAddr, +) -> Result { + let provision = run_plan::derive_stage_provision(plan, stage_index) + .map_err(|e| format!("derive stage {stage_index} provision: {e:?}"))?; + Ok(StageProvisionWire { + run_id: provision.run_id.0, + authorized_orchestrator: 0, + node_id: provision.node_id.0, + stage_index: provision.stage_index, + stage_count: provision.stage_count, + layer_start: provision.layer_start, + layer_end_exclusive: provision.layer_end_exclusive, + inbound_edge_id: provision.inbound.edge_id.0, + outbound_edge_id: provision.outbound.edge_id.0, + inbound_edge: Some(StageInboundEdgeWire { + edge_id: provision.inbound.edge_id.0, + kind: stage_edge_kind_wire(provision.inbound.kind), + object_spec: stage_object_spec_wire(provision.inbound.object_spec), + ring_spec: stage_ring_spec_wire(provision.inbound.ring_spec), + }), + outbound_edge: Some(StageOutboundEdgeWire { + edge_id: provision.outbound.edge_id.0, + kind: stage_edge_kind_wire(provision.outbound.kind), + consumer_node_id: provision.outbound.consumer_node_id.0, + consumer_endpoint: stage_consumer_endpoint( + provision.outbound.consumer_node_id.0, + plan, + readies, + coordinator, + )?, + object_spec: stage_object_spec_wire(provision.outbound.object_spec), + ring_spec: stage_ring_spec_wire(provision.outbound.ring_spec), + }), + model_id: provision.model.model_id, + gguf_source: provision.gguf_source, + tokenizer: provision.tokenizer, + }) +} + +fn provision_stage_from_plan( + stack: &DistributionRuntimeStack, + node_actor: ActorAddress, + plan: &run_plan::RunPlan, + stage_index: u32, + readies: &BTreeMap, + coordinator: &EndpointAddr, +) -> Result<(), String> { + let provision = stage_provision_wire_from_plan(plan, stage_index, readies, coordinator)?; + stack + .runtime + .send_to(node_actor, NodeAgentMsg::ProvisionStage(provision)) + .map_err(|e| format!("send stage {stage_index} provision: {e}")) +} + +fn stage_consumer_endpoint( + consumer_node_id: u64, + plan: &run_plan::RunPlan, + readies: &BTreeMap, + coordinator: &EndpointAddr, +) -> Result, String> { + if consumer_node_id + == plan.stages.first().map_or(0, |stage| { + plan.edges + .iter() + .find(|edge| edge.kind == run_plan::EdgeKind::TokenIn) + .and_then(|edge| match edge.producer { + run_plan::EdgeEndpoint::Orchestrator { node_id } => Some(node_id.0), + run_plan::EdgeEndpoint::Stage { .. } => None, + }) + .unwrap_or(stage.node_id.0) + }) + { + return Ok(Some(coordinator.clone())); + } + readies + .get(&consumer_node_id) + .map(|ready| Some(ready.endpoint.clone())) + .ok_or_else(|| { + format!("missing runtime-ready endpoint for consumer node {consumer_node_id}") + }) +} + +fn stage_edge_kind_wire(kind: run_plan::EdgeKind) -> StageEdgeKindWire { + match kind { + run_plan::EdgeKind::TokenIn => StageEdgeKindWire::TokenIn, + run_plan::EdgeKind::Activation => StageEdgeKindWire::Activation, + run_plan::EdgeKind::TokenOut => StageEdgeKindWire::TokenOut, + } +} + +fn stage_object_spec_wire(spec: run_plan::ObjectSpec) -> StageObjectSpecWire { + StageObjectSpecWire { + max_extent: spec.max_extent, + alignment: spec.alignment, + } +} + +fn stage_ring_spec_wire(spec: run_plan::RingSpec) -> StageRingSpecWire { + StageRingSpecWire { + data_capacity: spec.data_capacity, + alignment: spec.alignment, + } +} + +#[allow(clippy::too_many_arguments)] +fn wait_for_runtime_readies( + driver: &mut IrohDriver, + stack: &DistributionRuntimeStack, + obs_rx: &mpsc::Receiver, + frame_rx: &mpsc::Receiver, + frame_tx: &mpsc::Sender, + orchestrator_reports: &swactor::runtime::Inbox, + stop_rx: &mpsc::Receiver<()>, + dashboard: Option<&DashboardSupport>, + orch_datastream: &mut OrchDatastream, + orch_stdio_rx: Option<&mpsc::Receiver>, + run_id: u64, + expected_node_ids: &[u64], + provider: ProviderKind, +) -> Result, String> { + let expected = expected_node_ids.iter().copied().collect::>(); + let mut pending = BTreeMap::::new(); + loop { + pump(driver, stack, frame_tx); + emit_swim_transitions( + orch_datastream, + dashboard, + run_id, + expected_node_ids.first().copied().unwrap_or(0), + stack, + ); + drain_frames(frame_rx, dashboard, orch_datastream); + drain_orch_stdio_capture( + orch_stdio_rx, + orch_datastream, + dashboard, + run_id, + expected_node_ids.first().copied().unwrap_or(0), + ); + if stop_requested(stop_rx) { + return Err("shutdown requested while waiting for pipeline nodes ready".to_owned()); + } + while let Ok(observation) = obs_rx.try_recv() { + emit_plugin_observation(orch_datastream, dashboard, provider, &observation); + match observation { + PluginObservation::DatastreamFrame { .. } => {} + PluginObservation::ProviderLine { .. } + | PluginObservation::StdoutLine { .. } + | PluginObservation::StderrLine { .. } => {} + PluginObservation::Failed { reason, .. } => return Err(reason), + PluginObservation::Exited { + status, node_id, .. + } => { + return Err(format!("node {node_id} exited before ready: {status:?}")); + } + } + } + while let Some(report) = orchestrator_reports.try_recv() { + if let OrchestratorReport::NodeRuntimeReady { + run_id: report_run_id, + node_id, + stage_index, + endpoint, + node_actor, + datastream_publisher, + readiness_id, + } = report + && report_run_id == run_id + && expected.contains(&node_id) + { + pending.insert( + node_id, + RuntimeReady { + endpoint: endpoint.clone(), + node_actor, + datastream_publisher, + stage_index, + readiness_id, + swim_node_id: DistNodeId(*endpoint.id.as_bytes()), + }, + ); + } + } + if expected.iter().all(|node_id| { + pending + .get(node_id) + .is_some_and(|ready| runtime_ready_barrier_met(stack, ready)) + }) { + return Ok(pending); + } + thread::sleep(PUMP_INTERVAL); + } +} + +#[allow(clippy::too_many_arguments)] +fn wait_for_weights_loaded_count( + driver: &mut IrohDriver, + stack: &DistributionRuntimeStack, + obs_rx: &mpsc::Receiver, + frame_rx: &mpsc::Receiver, + frame_tx: &mpsc::Sender, + orchestrator_reports: &swactor::runtime::Inbox, + stop_rx: &mpsc::Receiver<()>, + dashboard: Option<&DashboardSupport>, + orch_datastream: &mut OrchDatastream, + orch_stdio_rx: Option<&mpsc::Receiver>, + run_id: u64, + node_id: u64, + provider: ProviderKind, + expected_count: usize, + pipeline_plan: &run_plan::RunPlan, + readies: &BTreeMap, + pipeline_coordinator: &EndpointAddr, +) -> Result<(), String> { + let expected_stages = pipeline_plan + .stages + .iter() + .map(|stage| stage.stage_index) + .collect::>(); + let mut loaded_stages = BTreeSet::::new(); + let mut last_resend = Instant::now(); + let mut active_stage = None::; + let mut resend_attempt = 0_u64; + loop { + pump(driver, stack, frame_tx); + emit_swim_transitions(orch_datastream, dashboard, run_id, node_id, stack); + drain_orch_stdio_capture(orch_stdio_rx, orch_datastream, dashboard, run_id, node_id); + if stop_requested(stop_rx) { + return Err("shutdown requested while waiting for pipeline weights loaded".to_owned()); + } + if loaded_stages.len() >= expected_count { + return Ok(()); + } + if active_stage.map_or(true, |stage_index| loaded_stages.contains(&stage_index)) { + active_stage = + next_pipeline_weight_load_stage(pipeline_plan, &loaded_stages, active_stage) + .map(|stage| stage.stage_index); + last_resend = Instant::now() - Duration::from_secs(1); + } + if last_resend.elapsed() >= Duration::from_secs(1) { + resend_attempt += 1; + let Some(stage) = + next_pipeline_weight_load_stage(pipeline_plan, &loaded_stages, active_stage) + else { + return Err(format!( + "missing unloaded pipeline weight stage; loaded {} of {expected_count}", + loaded_stages.len() + )); + }; + let stage_node_id = stage.node_id.0; + let ready = readies.get(&stage_node_id).ok_or_else(|| { + format!("missing runtime-ready node for stage {}", stage.stage_index) + })?; + orch_datastream.emit_bootstrap( + dashboard, + run_id, + node_id, + "stage_provision_send", + "sent", + json!({ + "attempt":resend_attempt, + "stage_count":pipeline_plan.stages.len(), + "stage_index":stage.stage_index, + "loaded_stage_count":loaded_stages.len() + }), + ); + let route_owner = stack.route_owner(ready.node_actor); + let member_state = stack.member_state(ready.swim_node_id); + orch_datastream.emit_bootstrap_to_channel( + dashboard, + MVP_STAGE_ROUTE, + run_id, + node_id, + "stage_route_check", + "observed", + json!({ + "attempt":resend_attempt, + "stage_index":stage.stage_index, + "stage_node_id":stage_node_id, + "node_actor":ready.node_actor, + "datastream_publisher":ready.datastream_publisher, + "swim_node_id":format!("{:?}", ready.swim_node_id), + "member_state":member_state.map(|state| format!("{:?}", state)), + "route_owner":route_owner.map(|owner| format!("{:?}", owner)), + "datastream_route_owner":stack.route_owner(ready.datastream_publisher).map(|owner| format!("{:?}", owner)), + "route_matches_ready":route_owner == Some(ready.swim_node_id), + }), + ); + provision_stage_from_plan( + stack, + ready.node_actor, + pipeline_plan, + stage.stage_index, + readies, + pipeline_coordinator, + )?; + pump(driver, stack, frame_tx); + last_resend = Instant::now(); + } + while let Ok(observation) = obs_rx.try_recv() { + emit_plugin_observation(orch_datastream, dashboard, provider, &observation); + match observation { + PluginObservation::Failed { reason, .. } => return Err(reason), + PluginObservation::Exited { + node_id, status, .. + } => { + return Err(format!( + "node {node_id} exited while loading weights: {status:?}" + )); + } + PluginObservation::DatastreamFrame { .. } => {} + PluginObservation::ProviderLine { .. } + | PluginObservation::StdoutLine { .. } + | PluginObservation::StderrLine { .. } => {} + } + } + drain_frames(frame_rx, dashboard, orch_datastream); + while let Some(report) = orchestrator_reports.try_recv() { + match report { + OrchestratorReport::WeightsReady { + run_id: report_run_id, + node_id: _, + stage_index, + } if report_run_id == run_id && expected_stages.contains(&stage_index) => { + loaded_stages.insert(stage_index); + if active_stage == Some(stage_index) { + active_stage = None; + last_resend = Instant::now() - Duration::from_secs(1); + } + if loaded_stages.len() >= expected_count { + return Ok(()); + } + } + OrchestratorReport::StageFault { + run_id: report_run_id, + stage_index, + } if report_run_id == run_id && expected_stages.contains(&stage_index) => { + return Err(format!( + "stage {stage_index} faulted while loading pipeline weights" + )); + } + _ => {} + } + } + thread::sleep(PUMP_INTERVAL); + } +} + +fn next_pipeline_weight_load_stage<'a>( + pipeline_plan: &'a run_plan::RunPlan, + loaded_stages: &BTreeSet, + active_stage: Option, +) -> Option<&'a run_plan::StagePlan> { + if let Some(stage_index) = active_stage { + if !loaded_stages.contains(&stage_index) { + if let Some(stage) = pipeline_plan + .stages + .iter() + .find(|stage| stage.stage_index == stage_index) + { + return Some(stage); + } + } + } + pipeline_plan + .stages + .iter() + .filter(|stage| !loaded_stages.contains(&stage.stage_index)) + .min_by_key(|stage| stage.stage_index) +} + +struct FailedProvisionPlugin; + +impl ProvisionPlugin for FailedProvisionPlugin { + fn start_node( + &mut self, + _spec: NodeProvisionSpec, + _sink: PluginSink, + ) -> Result { + Err("provider start worker disconnected".to_owned()) + } + + fn complete_bootstrap( + &mut self, + _handle: &crate::provisioning::PluginNodeHandle, + ) -> Result<(), String> { + Ok(()) + } + + fn stop_node(&mut self, _handle: &crate::provisioning::PluginNodeHandle) -> Result<(), String> { + Ok(()) + } +} + +struct PromptWork { + request: SubmitPrompt, + events: mpsc::Sender, +} + +struct ActivePrompt { + request: SubmitPrompt, + events: mpsc::Sender, +} + +#[derive(Clone, Debug)] +struct CollectedDatastreamFrame { + stream: StreamId, + channel_name: String, + frame: Frame, +} + +fn drain_datastream_connections( + driver: &mut IrohDriver, + frame_tx: &mpsc::Sender, +) { + driver.pump_datastream_ingress(); + for read in driver.drain_datastream_reads() { + let mut channels = read + .header + .channels + .iter() + .map(|descriptor| { + ( + ChannelRef { + stream: descriptor.stream.clone(), + channel: descriptor.id, + }, + descriptor.name.clone(), + ) + }) + .collect::>(); + for event in read.events { + match event { + DatastreamEvent::ChannelDeclared(descriptor) => { + channels.insert( + ChannelRef { + stream: descriptor.stream.clone(), + channel: descriptor.id, + }, + descriptor.name, + ); + } + DatastreamEvent::Frame(delivery) => { + let channel_name = channels + .get(&delivery.channel) + .cloned() + .unwrap_or_else(|| format!("channel#{}", delivery.channel.channel.0)); + let frame = Frame::new( + delivery.channel.channel, + delivery.position, + delivery.payload, + ); + if frame_tx + .send(CollectedDatastreamFrame { + stream: delivery.channel.stream, + channel_name, + frame, + }) + .is_err() + { + return; + } + } + DatastreamEvent::StreamDeclared(_) | DatastreamEvent::StreamEnded(_) => {} + } + } + } +} + +struct FrameArchive { + file: File, + next_seq: u64, +} + +impl FrameArchive { + fn open(path: &Path) -> Result { + if let Some(parent) = path.parent() + && !parent.as_os_str().is_empty() + { + std::fs::create_dir_all(parent).map_err(|e| { + format!("create datastream frame log dir {}: {e}", parent.display()) + })?; + } + let file = OpenOptions::new() + .create(true) + .append(true) + .open(path) + .map_err(|e| format!("open datastream frame log {}: {e}", path.display()))?; + Ok(Self { file, next_seq: 0 }) + } + + fn record(&mut self, source: &str, stream: &StreamId, channel: &str, frame: &Frame) { + let payload = match std::str::from_utf8(&frame.payload) { + Ok(text) => json!({"encoding":"utf8","value":text}), + Err(_) => json!({"encoding":"bytes","value":frame.payload}), + }; + let record = json!({ + "arrival_seq":self.next_seq, + "arrival_unix_ms":benchmark_observability::unix_ms_now(), + "source":source, + "stream":stream.to_string(), + "channel":channel, + "channel_id":frame.channel.0, + "position":frame.position.0, + "payload":payload, + }); + self.next_seq += 1; + let _ = serde_json::to_writer(&mut self.file, &record); + let _ = writeln!(self.file); + let _ = self.file.flush(); + } +} + +struct OrchDatastream { + stream: StreamId, + endpoint: DatastreamEndpoint, + producer: DatastreamProducer, + channels: BTreeMap, + channel_names: BTreeMap, + archive: Option, +} + +impl OrchDatastream { + fn new(run_id: u64, frame_log: Option<&Path>) -> Result { + let stream = StreamId::new(NodeId::new("mvp-orchestrator"), Lifetime(run_id)); + let endpoint = DatastreamEndpoint::with_descriptor( + StreamDescriptor { + stream: stream.clone(), + label: Some("mvp orchestrator".to_owned()), + origin: StreamOrigin::Orchestrator, + }, + 4096, + 1024, + ); + let producer = endpoint.producer(); + let mut out = Self { + stream, + endpoint, + producer, + channels: BTreeMap::new(), + channel_names: BTreeMap::new(), + archive: frame_log.map(FrameArchive::open).transpose()?, + }; + for name in [ + MVP_PROVISIONING_EVENTS, + MVP_ORCH_BOOTSTRAP, + MVP_ORCH_PROMPT, + MVP_SWIM_MEMBERSHIP, + MVP_STAGE_ROUTE, + ] { + out.channel_by_name(name); + } + Ok(out) + } + + fn channel_by_name(&mut self, name: &str) -> ChannelId { + if let Some(id) = self.channels.get(name).copied() { + return id; + } + let id = self.producer.register_channel( + name, + ChannelContent::JsonRecord { + schema: Some(name.to_owned()), + }, + ); + self.channels.insert(name.to_owned(), id); + self.channel_names.insert(id, name.to_owned()); + id + } + + fn emit_event(&mut self, dashboard: Option<&DashboardSupport>, event: ProvisionEvent) { + let payload = serde_json::to_vec(&MvpProvisionEventRecord::new(event)) + .expect("serialize provisioning event"); + self.emit_bytes(dashboard, MVP_PROVISIONING_EVENTS, payload); + } + + fn emit_log(&mut self, dashboard: Option<&DashboardSupport>, line: ProvisionLogLine) { + let channel = mvp_provision_log_channel(line.node_id, line.stream); + let payload = + serde_json::to_vec(&MvpProvisionLogRecord::new(line)).expect("serialize provision log"); + self.emit_bytes(dashboard, &channel, payload); + } + + fn emit_bootstrap( + &mut self, + dashboard: Option<&DashboardSupport>, + run_id: u64, + node_id: u64, + phase: &str, + status: &str, + detail: Value, + ) { + self.emit_bootstrap_to_channel( + dashboard, + MVP_ORCH_BOOTSTRAP, + run_id, + node_id, + phase, + status, + detail, + ); + } + + fn emit_bootstrap_to_channel( + &mut self, + dashboard: Option<&DashboardSupport>, + channel: &str, + run_id: u64, + node_id: u64, + phase: &str, + status: &str, + detail: Value, + ) { + let payload = serde_json::to_vec(&json!({ + "type":"OrchBootstrap", + "phase":phase, + "status":status, + "run_id":run_id, + "node_id":node_id, + "benchmark":benchmark_observability::stamp("mvp-orchestrator"), + "detail":detail, + })) + .expect("serialize orch bootstrap event"); + self.emit_bytes(dashboard, channel, payload); + } + + fn emit_prompt( + &mut self, + dashboard: Option<&DashboardSupport>, + run_id: u64, + node_id: u64, + request_id: u64, + phase: &str, + status: &str, + detail: Value, + ) { + let payload = serde_json::to_vec(&json!({ + "type":"OrchPromptEvent", + "phase":phase, + "status":status, + "run_id":run_id, + "node_id":node_id, + "request_id":request_id, + "benchmark":benchmark_observability::stamp("mvp-orchestrator"), + "detail":detail, + })) + .expect("serialize orch prompt event"); + self.emit_bytes(dashboard, MVP_ORCH_PROMPT, payload); + } + + fn emit_bytes( + &mut self, + dashboard: Option<&DashboardSupport>, + channel: &str, + payload: Vec, + ) { + self.emit_bytes_from(dashboard, channel, payload, "orchestrator"); + } + + fn emit_bytes_from( + &mut self, + dashboard: Option<&DashboardSupport>, + channel: &str, + payload: Vec, + source: &str, + ) { + let id = self.channel_by_name(channel); + self.producer.submit_bytes(id, payload); + self.flush(dashboard, source); + } + + fn flush(&mut self, dashboard: Option<&DashboardSupport>, source: &str) { + let stream = self.stream.clone(); + for frame in self.endpoint.mux().drain() { + let channel = self + .channel_names + .get(&frame.channel) + .cloned() + .unwrap_or_else(|| format!("channel#{}", frame.channel.0)); + ingest_dashboard_frame(dashboard, &stream, &channel, &frame); + self.archive_frame(source, &stream, &channel, &frame); + } + } + + fn archive_frame(&mut self, source: &str, stream: &StreamId, channel: &str, frame: &Frame) { + if let Some(archive) = &mut self.archive { + archive.record(source, stream, channel, frame); + } + } +} + +struct OrchStdioCapture; + +struct OrchStdioLine { + stream: ProvisionLogStream, + line: String, +} + +#[cfg(target_os = "linux")] +impl OrchStdioCapture { + fn install() -> Result>, String> { + let stdout_read = Self::redirect_stream(libc::STDOUT_FILENO, "stdout")?; + let stderr_read = Self::redirect_stream(libc::STDERR_FILENO, "stderr")?; + let (tx, rx) = mpsc::channel(); + Self::spawn_reader(stdout_read, ProvisionLogStream::Stdout, tx.clone()); + Self::spawn_reader(stderr_read, ProvisionLogStream::Stderr, tx); + Ok(Some(rx)) + } + + fn redirect_stream(fd: libc::c_int, name: &str) -> Result { + let mut pipe_fds = [0; 2]; + let pipe_result = unsafe { libc::pipe(pipe_fds.as_mut_ptr()) }; + if pipe_result != 0 { + return Err(format!( + "create orchestrator {name} capture pipe: {}", + std::io::Error::last_os_error() + )); + } + + let dup_result = unsafe { libc::dup2(pipe_fds[1], fd) }; + let close_write_result = unsafe { libc::close(pipe_fds[1]) }; + if dup_result < 0 { + let error = std::io::Error::last_os_error(); + let _ = unsafe { libc::close(pipe_fds[0]) }; + return Err(format!("redirect orchestrator {name}: {error}")); + } + if close_write_result != 0 { + let error = std::io::Error::last_os_error(); + let _ = unsafe { libc::close(pipe_fds[0]) }; + return Err(format!("close orchestrator {name} duplicate fd: {error}")); + } + + Ok(unsafe { File::from_raw_fd(pipe_fds[0]) }) + } + + fn spawn_reader(file: File, stream: ProvisionLogStream, tx: mpsc::Sender) { + thread::spawn(move || { + let reader = BufReader::new(file); + for line in reader.lines() { + let Ok(line) = line else { + break; + }; + if tx.send(OrchStdioLine { stream, line }).is_err() { + break; + } + } + }); + } +} + +#[cfg(not(target_os = "linux"))] +impl OrchStdioCapture { + fn install() -> Result>, String> { + Ok(None) + } +} + +fn install_orch_stdio_capture() -> Result>, String> { + OrchStdioCapture::install() +} + +fn drain_orch_stdio_capture( + rx: Option<&mpsc::Receiver>, + datastream: &mut OrchDatastream, + dashboard: Option<&DashboardSupport>, + run_id: u64, + node_id: u64, +) { + let Some(rx) = rx else { + return; + }; + while let Ok(line) = rx.try_recv() { + datastream.emit_log( + dashboard, + ProvisionLogLine { + run_id, + node_id, + stream: line.stream, + line: line.line, + }, + ); + } +} + +#[cfg(feature = "dashboard")] +struct DashboardSupport { + handle: dashboard::DashboardHandle, + _runtime: tokio::runtime::Runtime, +} + +#[cfg(feature = "dashboard")] +impl DashboardSupport { + fn start(enabled: bool) -> Result, String> { + if !enabled { + return Ok(None); + } + let mut config = dashboard::DashboardConfig::default(); + if let Some(port) = env_optional("MVP_DASHBOARD_PORT") { + config.port = port + .parse::() + .map_err(|e| format!("invalid MVP_DASHBOARD_PORT={port:?}: {e}"))?; + } + let runtime = tokio::runtime::Builder::new_multi_thread() + .enable_all() + .build() + .map_err(|e| format!("dashboard runtime: {e}"))?; + let handle = dashboard::DashboardHandle::new(config); + handle.register_view(Arc::new(MvpClusterDashboardView::new())); + handle.spawn_http(runtime.handle()); + Ok(Some(Self { + handle, + _runtime: runtime, + })) + } + + fn publish_frame(&self, stream: &StreamId, channel: &str, frame: &Frame) { + self.handle.publish(dashboard::FrameEvent { + stream: dashboard::StreamEvent { + node: stream.node.as_str().to_string(), + life: stream.life.0, + }, + channel: channel.to_owned(), + position: frame.position.0, + payload: frame.payload.clone(), + }); + } +} + +#[cfg(not(feature = "dashboard"))] +struct DashboardSupport; + +#[cfg(not(feature = "dashboard"))] +impl DashboardSupport { + fn start(enabled: bool) -> Result, String> { + if enabled { + return Err( + "MVP_DASHBOARD requires building mvp-system with feature dashboard".to_owned(), + ); + } + Ok(None) + } + + fn publish_frame(&self, _stream: &StreamId, _channel: &str, _frame: &Frame) {} +} + +struct ChannelObservationSink { + tx: Mutex>, +} + +impl PluginObservationSink for ChannelObservationSink { + fn observe(&self, observation: PluginObservation) { + let _ = self.tx.lock().send(observation); + } +} + +fn spawn_prompt_rpc( + bind: SocketAddr, + work_tx: mpsc::Sender, + default_max_tokens: u32, +) -> Result { + let listener = TcpListener::bind(bind).map_err(|e| format!("bind prompt RPC {bind}: {e}"))?; + let addr = listener + .local_addr() + .map_err(|e| format!("read prompt RPC addr: {e}"))?; + thread::spawn(move || { + for accepted in listener.incoming() { + match accepted { + Ok(stream) => { + let tx = work_tx.clone(); + thread::spawn(move || { + let _ = handle_prompt_connection(stream, tx, default_max_tokens); + }); + } + Err(_) => break, + } + } + }); + Ok(addr) +} +fn handle_prompt_connection( + stream: TcpStream, + work_tx: mpsc::Sender, + default_max_tokens: u32, +) -> Result<(), String> { + let mut reader = BufReader::new( + stream + .try_clone() + .map_err(|e| format!("clone prompt stream: {e}"))?, + ); + let mut writer = stream; + loop { + let request = match read_submit_prompt(&mut reader) { + Ok(Some(request)) => request, + Ok(None) => break, + Err(error) if error.contains("expected value at line 1 column 1") => break, + Err(error) => return Err(error), + }; + let request = request.with_defaults(default_max_tokens); + let (event_tx, event_rx) = mpsc::channel(); + work_tx + .send(PromptWork { + request, + events: event_tx, + }) + .map_err(|_| "prompt loop stopped".to_owned())?; + for event in event_rx { + let terminal = event.is_terminal(); + write_json_line(&mut writer, &event)?; + if terminal { + break; + } + } + } + Ok(()) +} + +fn wait_for_runtime_ready( + driver: &mut IrohDriver, + stack: &DistributionRuntimeStack, + obs_rx: &mpsc::Receiver, + frame_rx: &mpsc::Receiver, + frame_tx: &mpsc::Sender, + orchestrator_reports: &swactor::runtime::Inbox, + stop_rx: &mpsc::Receiver<()>, + dashboard: Option<&DashboardSupport>, + orch_datastream: &mut OrchDatastream, + orch_stdio_rx: Option<&mpsc::Receiver>, + run_id: u64, + node_id: u64, + provider: ProviderKind, +) -> Result { + let mut pending_ready: Option = None; + let mut node_swim_started = false; + let mut node_swim_ready = false; + let mut node_route_started = false; + loop { + pump(driver, stack, frame_tx); + drain_frames(frame_rx, dashboard, orch_datastream); + drain_orch_stdio_capture(orch_stdio_rx, orch_datastream, dashboard, run_id, node_id); + if stop_requested(stop_rx) { + return Err("shutdown requested while waiting for node ready".to_owned()); + } + while let Ok(observation) = obs_rx.try_recv() { + emit_plugin_observation(orch_datastream, dashboard, provider, &observation); + match observation { + PluginObservation::DatastreamFrame { .. } => {} + PluginObservation::ProviderLine { .. } + | PluginObservation::StdoutLine { .. } + | PluginObservation::StderrLine { .. } => {} + PluginObservation::Failed { reason, .. } => return Err(reason), + PluginObservation::Exited { status, .. } => { + return Err(format!("node exited before ready: {status:?}")); + } + } + } + while let Some(report) = orchestrator_reports.try_recv() { + if let OrchestratorReport::NodeRuntimeReady { + run_id: report_run_id, + node_id: report_node_id, + stage_index, + endpoint, + node_actor, + datastream_publisher, + readiness_id, + } = report + { + if report_run_id == run_id && report_node_id == node_id { + let reset_progress = pending_ready + .as_ref() + .map(|ready| ready.readiness_id != readiness_id) + .unwrap_or(true); + if reset_progress { + node_swim_started = false; + node_swim_ready = false; + node_route_started = false; + } + let swim_node_id = DistNodeId(*endpoint.id.as_bytes()); + pending_ready = Some(RuntimeReady { + endpoint, + node_actor, + datastream_publisher, + stage_index, + readiness_id, + swim_node_id, + }); + } + } + } + if let Some(ready) = pending_ready.as_ref() { + if runtime_ready_barrier_met(stack, ready) { + return Ok(ready.clone()); + } + let swim_ready = stack.member_state(ready.swim_node_id) == Some(MemberState::Alive); + let route_ready = stack.route_owner(ready.node_actor) == Some(ready.swim_node_id); + if !swim_ready { + if !node_swim_started { + orch_datastream.emit_bootstrap( + dashboard, + run_id, + node_id, + "node_swim", + "started", + json!({"node":format!("{:?}", ready.swim_node_id),"readiness_id":ready.readiness_id}), + ); + node_swim_started = true; + } + } else if !route_ready { + if !node_swim_ready { + orch_datastream.emit_bootstrap( + dashboard, + run_id, + node_id, + "node_swim", + "ready", + json!({"node":format!("{:?}", ready.swim_node_id),"readiness_id":ready.readiness_id}), + ); + node_swim_ready = true; + } + if !node_route_started { + orch_datastream.emit_bootstrap( + dashboard, + run_id, + node_id, + "node_route", + "started", + json!({"node_actor":ready.node_actor,"node":format!("{:?}", ready.swim_node_id),"readiness_id":ready.readiness_id}), + ); + node_route_started = true; + } + } + } + thread::sleep(PUMP_INTERVAL); + } +} + +fn provision_stage( + stack: &DistributionRuntimeStack, + node_actor: ActorAddress, + config: &Config, +) -> Result<(), String> { + stack + .runtime + .send_to( + node_actor, + NodeAgentMsg::ProvisionStage({ + let layer_end_exclusive = config.layer_end_exclusive.ok_or_else(|| { + "unplanned single-stage execution requires --layer-end-exclusive or MVP_LAYER_END_EXCLUSIVE; cached-model runs use metadata-derived planning".to_owned() + })?; + StageProvisionWire { + run_id: config.run_id, + authorized_orchestrator: 0, + node_id: config.node_id, + stage_index: config.stage_index, + stage_count: 1, + layer_start: 0, + layer_end_exclusive, + inbound_edge_id: 1, + outbound_edge_id: 2, + inbound_edge: None, + outbound_edge: None, + model_id: config.model_id.clone(), + gguf_source: config.gguf_source.clone(), + tokenizer: config.tokenizer.clone(), + } + }), + ) + .map_err(|e| format!("send stage provision: {e}")) +} + +fn wait_for_weights_loaded( + driver: &mut IrohDriver, + stack: &DistributionRuntimeStack, + obs_rx: &mpsc::Receiver, + frame_rx: &mpsc::Receiver, + frame_tx: &mpsc::Sender, + orchestrator_reports: &swactor::runtime::Inbox, + stop_rx: &mpsc::Receiver<()>, + dashboard: Option<&DashboardSupport>, + orch_datastream: &mut OrchDatastream, + orch_stdio_rx: Option<&mpsc::Receiver>, + run_id: u64, + node_id: u64, + provider: ProviderKind, + stage_index: u32, +) -> Result<(), String> { + loop { + pump(driver, stack, frame_tx); + drain_orch_stdio_capture(orch_stdio_rx, orch_datastream, dashboard, run_id, node_id); + if stop_requested(stop_rx) { + return Err("shutdown requested while waiting for weights loaded".to_owned()); + } + while let Ok(observation) = obs_rx.try_recv() { + emit_plugin_observation(orch_datastream, dashboard, provider, &observation); + match observation { + PluginObservation::Failed { reason, .. } => return Err(reason), + PluginObservation::Exited { status, .. } => { + return Err(format!("node exited while loading weights: {status:?}")); + } + PluginObservation::DatastreamFrame { .. } => {} + PluginObservation::ProviderLine { .. } + | PluginObservation::StdoutLine { .. } + | PluginObservation::StderrLine { .. } => {} + } + } + drain_frames(frame_rx, dashboard, orch_datastream); + while let Some(report) = orchestrator_reports.try_recv() { + match report { + OrchestratorReport::WeightsReady { + run_id: report_run_id, + node_id: _, + stage_index: report_stage_index, + } if report_run_id == run_id && report_stage_index == stage_index => { + return Ok(()); + } + OrchestratorReport::StageFault { + run_id: report_run_id, + stage_index: report_stage_index, + } if report_run_id == run_id && report_stage_index == stage_index => { + return Err(format!( + "stage {report_stage_index} faulted while loading weights" + )); + } + _ => {} + } + } + thread::sleep(PUMP_INTERVAL); + } +} + +#[derive(Debug)] +struct PipelineTokenRecord { + object_id: u64, + sequence: u64, + token_id: u32, + eos: bool, +} + +enum PipelineSendHandle { + Driver(EdgeSendHandle), + #[cfg(test)] + Channel(tokio_mpsc::UnboundedSender>), +} + +impl PipelineSendHandle { + fn send(&self, bytes: Vec) -> Result<(), String> { + match self { + Self::Driver(handle) => handle.send(bytes), + #[cfg(test)] + Self::Channel(tx) => tx + .send(bytes) + .map_err(|_| "pipeline token-in sender stopped".to_owned()), + } + } +} +struct PendingEncode { + request_id: u64, +} + +struct PendingDecode { + request_id: u64, + token_id: u32, + eos: bool, + reached_limit: bool, +} + +struct PipelinePromptRuntime { + token_in_edge_id: u64, + token_out_edge_id: u64, + token_spec: run_plan::ObjectSpec, + token_out_spec: run_plan::ObjectSpec, + token_in_sender: PipelineSendHandle, + recv_rx: mpsc::Receiver>, + recv_tx: mpsc::Sender>, + recv_buffer: Vec, + tokenizer_encode_actor: ActorAddress, + tokenizer_decode_actor: ActorAddress, + tokenizer_reply_to: ActorAddress, + pending_encode: Option, + pending_decode: Option, + next_sequence: u64, + generated_tokens: Vec, + final_text: String, + active: Option, + started_at: Option, +} + +impl PipelinePromptRuntime { + fn new( + driver: &IrohDriver, + plan: &run_plan::RunPlan, + first_stage_endpoint: EndpointAddr, + tokenizer_encode_actor: ActorAddress, + tokenizer_decode_actor: ActorAddress, + tokenizer_reply_to: ActorAddress, + ) -> Result { + let token_in_edge = plan + .edges + .iter() + .find(|edge| edge.kind == run_plan::EdgeKind::TokenIn) + .ok_or_else(|| "pipeline plan missing token-in edge".to_owned())?; + let token_out_edge = plan + .edges + .iter() + .find(|edge| edge.kind == run_plan::EdgeKind::TokenOut) + .ok_or_else(|| "pipeline plan missing token-out edge".to_owned())?; + let (recv_tx, recv_rx) = mpsc::channel(); + Ok(Self { + token_in_edge_id: token_in_edge.edge_id.0, + token_out_edge_id: token_out_edge.edge_id.0, + token_spec: token_in_edge.object_spec, + token_out_spec: token_out_edge.object_spec, + token_in_sender: PipelineSendHandle::Driver( + driver.spawn_edge_send_pump(first_stage_endpoint, token_in_edge.edge_id.0)?, + ), + recv_rx, + recv_tx, + tokenizer_encode_actor, + tokenizer_decode_actor, + tokenizer_reply_to, + pending_encode: None, + pending_decode: None, + recv_buffer: Vec::new(), + next_sequence: 0, + generated_tokens: Vec::new(), + final_text: String::new(), + active: None, + started_at: None, + }) + } + + fn is_active(&self) -> bool { + self.active.is_some() + } + + fn start_prompt( + &mut self, + request: SubmitPrompt, + events: mpsc::Sender, + runtime: &Arc, + dashboard: Option<&DashboardSupport>, + orch_datastream: &mut OrchDatastream, + run_id: u64, + node_id: u64, + ) -> Result<(), String> { + let request_id = request.request_id; + self.generated_tokens.clear(); + self.final_text.clear(); + self.recv_buffer.clear(); + self.pending_decode = None; + self.pending_encode = Some(PendingEncode { request_id }); + self.started_at = Some(Instant::now()); + orch_datastream.emit_prompt( + dashboard, + run_id, + node_id, + request_id, + "pipeline_tokenizer_encode", + "started", + json!({"node_actor":self.tokenizer_encode_actor,"reply_to":self.tokenizer_reply_to,"prompt_bytes":request.prompt_text.len()}), + ); + runtime + .send_to( + self.tokenizer_encode_actor, + NodeAgentMsg::EncodePrompt { + request_id, + prompt: request.prompt_text.clone(), + reply_to: self.tokenizer_reply_to, + }, + ) + .map_err(|e| format!("send tokenizer encode request: {e}"))?; + self.active = Some(ActivePrompt { request, events }); + Ok(()) + } + + fn drain_tokenizer_events( + &mut self, + runtime: &Arc, + tokenizer_events: &swactor::runtime::Inbox, + dashboard: Option<&DashboardSupport>, + orch_datastream: &mut OrchDatastream, + run_id: u64, + node_id: u64, + ) -> Result<(), String> { + while let Some(event) = tokenizer_events.try_recv() { + match event { + TokenizerEvent::PromptEncoded { request_id, tokens } => self + .handle_encoded_prompt( + request_id, + tokens, + dashboard, + orch_datastream, + run_id, + node_id, + )?, + TokenizerEvent::TokensDecoded { request_id, text } => self.handle_decoded_tokens( + runtime, + request_id, + text, + dashboard, + orch_datastream, + run_id, + node_id, + )?, + TokenizerEvent::Fault { request_id, error } => { + self.fault_active(request_id, format!("tokenizer request failed: {error}")); + } + } + } + Ok(()) + } + + fn handle_encoded_prompt( + &mut self, + request_id: u64, + tokens: Vec, + dashboard: Option<&DashboardSupport>, + orch_datastream: &mut OrchDatastream, + run_id: u64, + node_id: u64, + ) -> Result<(), String> { + let Some(pending) = self.pending_encode.take() else { + return Ok(()); + }; + if pending.request_id != request_id { + self.pending_encode = Some(pending); + return Ok(()); + } + let Some(active) = self.active.as_ref() else { + return Ok(()); + }; + if active.request.request_id != request_id { + return Ok(()); + } + orch_datastream.emit_prompt( + dashboard, + run_id, + node_id, + request_id, + "pipeline_tokenizer_encode", + "ready", + json!({"node_actor":self.tokenizer_encode_actor,"reply_to":self.tokenizer_reply_to,"tokens":tokens.len()}), + ); + let sequence = self.next_sequence; + orch_datastream.emit_prompt( + dashboard, + run_id, + node_id, + request_id, + "pipeline_token_in", + "started", + json!({"edge_id":self.token_in_edge_id,"sequence":sequence,"tokens":tokens.len(),"begin_sequence":true}), + ); + self.send_token_in(sequence, &tokens, true)?; + orch_datastream.emit_prompt( + dashboard, + run_id, + node_id, + request_id, + "pipeline_token_in", + "ready", + json!({"edge_id":self.token_in_edge_id,"sequence":sequence,"begin_sequence":true}), + ); + Ok(()) + } + + #[allow(clippy::too_many_arguments)] + fn handle_decoded_tokens( + &mut self, + _runtime: &Arc, + request_id: u64, + text: String, + dashboard: Option<&DashboardSupport>, + orch_datastream: &mut OrchDatastream, + run_id: u64, + node_id: u64, + ) -> Result<(), String> { + let Some(pending) = self.pending_decode.take() else { + return Ok(()); + }; + if pending.request_id != request_id { + self.pending_decode = Some(pending); + return Ok(()); + } + let Some(active) = self.active.as_ref() else { + return Ok(()); + }; + if active.request.request_id != request_id { + return Ok(()); + } + orch_datastream.emit_prompt( + dashboard, + run_id, + node_id, + request_id, + "pipeline_tokenizer_decode", + "ready", + json!({"node_actor":self.tokenizer_decode_actor,"reply_to":self.tokenizer_reply_to,"text_bytes":text.len()}), + ); + self.final_text.push_str(&text); + if !text.is_empty() { + let _ = active + .events + .send(PromptEvent::TextDelta { request_id, text }); + } + if pending.eos || pending.reached_limit { + let elapsed_ms = self + .started_at + .map(|started| started.elapsed().as_millis() as u64) + .unwrap_or(0); + let final_text = self.final_text.clone(); + let tokens_generated = self.generated_tokens.len() as u32; + orch_datastream.emit_prompt( + dashboard, + run_id, + node_id, + request_id, + "prompt_complete", + "ready", + json!({ + "event":"Done", + "terminal":true, + "tokens_generated":tokens_generated, + "elapsed_ms":elapsed_ms, + "final_text_bytes":final_text.len(), + }), + ); + let _ = active.events.send(PromptEvent::Done { + request_id, + final_text, + tokens_generated, + elapsed_ms, + }); + self.active = None; + return Ok(()); + } + self.send_token_in(self.next_sequence, &[pending.token_id], false)?; + Ok(()) + } + fn send_token_in( + &self, + sequence: u64, + tokens: &[u32], + begin_sequence: bool, + ) -> Result<(), String> { + self.token_in_sender.send(encode_token_record_with_flags( + self.token_spec, + sequence, + tokens, + false, + begin_sequence, + )?) + } + + fn request_decode( + &mut self, + runtime: &Arc, + request_id: u64, + token_id: u32, + eos: bool, + reached_limit: bool, + ) -> Result<(), String> { + self.pending_decode = Some(PendingDecode { + request_id, + token_id, + eos, + reached_limit, + }); + runtime + .send_to( + self.tokenizer_decode_actor, + NodeAgentMsg::DecodeTokens { + request_id, + tokens: vec![token_id], + reply_to: self.tokenizer_reply_to, + }, + ) + .map_err(|e| format!("send tokenizer decode request: {e}")) + } + + fn fault_active(&mut self, request_id: u64, error: String) { + if let Some(active) = self.active.take() + && active.request.request_id == request_id + { + let _ = active.events.send(PromptEvent::Fault { request_id, error }); + } + self.pending_encode = None; + self.pending_decode = None; + } + + fn poll_driver(&mut self, driver: &mut IrohDriver) { + driver.pump_edge_ingress(); + for event in driver.drain_edge_events() { + match event { + EdgeTransportEvent::BytesRead { edge_id, bytes, .. } + if edge_id == self.token_out_edge_id => + { + let _ = self.recv_tx.send(bytes); + } + EdgeTransportEvent::StreamFault { + edge_id: Some(edge_id), + reason, + .. + } if edge_id == self.token_out_edge_id => { + if let Some(request_id) = + self.active.as_ref().map(|active| active.request.request_id) + { + self.fault_active( + request_id, + format!("pipeline token-out stream fault: {reason:?}"), + ); + } + } + _ => {} + } + } + } + + fn drain_tokens( + &mut self, + runtime: &Arc, + dashboard: Option<&DashboardSupport>, + orch_datastream: &mut OrchDatastream, + run_id: u64, + node_id: u64, + ) -> Result<(), String> { + if self.pending_decode.is_some() { + return Ok(()); + } + while let Ok(bytes) = self.recv_rx.try_recv() { + self.recv_buffer.extend_from_slice(&bytes); + while self.pending_decode.is_none() + && let Some(record) = + take_pipeline_token_record(&mut self.recv_buffer, self.token_out_spec)? + { + if record.sequence != self.next_sequence { + return Err(format!( + "pipeline token sequence violation: expected {}, got {}", + self.next_sequence, record.sequence + )); + } + self.next_sequence = self.next_sequence.saturating_add(1); + let Some(active) = self.active.as_ref() else { + continue; + }; + let request_id = active.request.request_id; + orch_datastream.emit_prompt( + dashboard, + run_id, + node_id, + request_id, + "pipeline_token_out", + "observed", + json!({"edge_id":self.token_out_edge_id,"object_id":record.object_id,"sequence":record.sequence,"token_id":record.token_id,"eos":record.eos}), + ); + self.generated_tokens.push(record.token_id); + let reached_limit = self.generated_tokens.len() as u32 >= active.request.max_tokens; + orch_datastream.emit_prompt( + dashboard, + run_id, + node_id, + request_id, + "pipeline_tokenizer_decode", + "started", + json!({"node_actor":self.tokenizer_decode_actor,"reply_to":self.tokenizer_reply_to,"token_id":record.token_id}), + ); + self.request_decode( + runtime, + request_id, + record.token_id, + record.eos, + reached_limit, + )?; + } + } + Ok(()) + } +} + +#[cfg(test)] +fn encode_token_record( + spec: run_plan::ObjectSpec, + _edge_id: u64, + sequence: u64, + tokens: &[u32], + eos: bool, +) -> Result, String> { + encode_token_record_with_flags(spec, sequence, tokens, eos, false) +} + +fn encode_token_record_with_flags( + spec: run_plan::ObjectSpec, + sequence: u64, + tokens: &[u32], + eos: bool, + begin_sequence: bool, +) -> Result, String> { + let payload = tokens + .iter() + .flat_map(|token| token.to_le_bytes()) + .collect::>(); + let mut flags = ingress::ObjectFlags::default(); + flags.end_of_sequence = eos; + flags.begin_sequence = begin_sequence; + Ok( + ingress::ObjectRecordBuilder::new(ingress_object_spec_from_plan(spec)) + .object_id(ingress::ObjectId(9000_u64.saturating_add(sequence))) + .sequence(sequence) + .payload(payload) + .flags(flags) + .encode(), + ) +} + +fn ingress_object_spec_from_plan(spec: run_plan::ObjectSpec) -> ingress::ObjectSpec { + let extent_alignment = match spec.kind { + run_plan::ObjectKind::Token => u64::from(spec.dtype_width_bytes), + _ => u64::from(spec.alignment), + }; + ingress::ObjectSpec { + max_extent: spec.max_extent, + alignment: extent_alignment, + layout: ingress::ObjectLayout::Token, + } +} + +fn take_pipeline_token_record( + buffer: &mut Vec, + spec: run_plan::ObjectSpec, +) -> Result, String> { + let record = + match ingress::read_object_record(buffer, ingress_object_spec_from_plan(spec), false) + .map_err(|reason| format!("invalid token-out record: {reason:?}"))? + { + ingress::ObjectRecordRead::Incomplete => return Ok(None), + ingress::ObjectRecordRead::Complete(record) => record, + }; + let payload = record + .payload(buffer) + .ok_or_else(|| "token-out record payload missing".to_owned())?; + if payload.len() != 4 { + return Err(format!( + "token-out payload must be exactly one u32, got {}", + payload.len() + )); + } + let token_id = u32::from_le_bytes(payload.try_into().unwrap()); + let out = PipelineTokenRecord { + object_id: record.object_id.0, + sequence: record.sequence, + token_id, + eos: record.flags.end_of_sequence, + }; + buffer.drain(..record.total_len); + Ok(Some(out)) +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum PromptRuntimeMode { + DirectInferPrompt, + PipelineTokenEdges, +} + +fn prompt_runtime_mode(pipeline_plan: Option<&run_plan::RunPlan>) -> PromptRuntimeMode { + if pipeline_plan.is_some() { + PromptRuntimeMode::PipelineTokenEdges + } else { + PromptRuntimeMode::DirectInferPrompt + } +} + +fn serve_prompts( + driver: &mut IrohDriver, + stack: &DistributionRuntimeStack, + obs_rx: &mpsc::Receiver, + frame_rx: &mpsc::Receiver, + frame_tx: &mpsc::Sender, + work_rx: &mpsc::Receiver, + prompt_events: &swactor::runtime::Inbox, + stop_rx: &mpsc::Receiver<()>, + dashboard: Option<&DashboardSupport>, + orch_datastream: &mut OrchDatastream, + orch_stdio_rx: Option<&mpsc::Receiver>, + run_id: u64, + node_id: u64, + node_actor: ActorAddress, + reply_to: ActorAddress, + tokenizer_events: &swactor::runtime::Inbox, + tokenizer_encode_actor: ActorAddress, + tokenizer_decode_actor: ActorAddress, + tokenizer_reply_to: ActorAddress, + provider: ProviderKind, + pipeline_plan: Option<&run_plan::RunPlan>, + prompt_endpoint: EndpointAddr, +) -> Result<(), String> { + let mut pipeline_runtime = match prompt_runtime_mode(pipeline_plan) { + PromptRuntimeMode::PipelineTokenEdges => Some(PipelinePromptRuntime::new( + driver, + pipeline_plan.expect("pipeline mode requires plan"), + prompt_endpoint, + tokenizer_encode_actor, + tokenizer_decode_actor, + tokenizer_reply_to, + )?), + PromptRuntimeMode::DirectInferPrompt => None, + }; + let mut active: Option = None; + loop { + pump(driver, stack, frame_tx); + if let Some(pipeline) = pipeline_runtime.as_mut() { + pipeline.poll_driver(driver); + pipeline.drain_tokenizer_events( + &stack.runtime, + tokenizer_events, + dashboard, + orch_datastream, + run_id, + node_id, + )?; + pipeline.drain_tokens(&stack.runtime, dashboard, orch_datastream, run_id, node_id)?; + } + drain_observations(obs_rx, dashboard, orch_datastream, provider)?; + drain_frames(frame_rx, dashboard, orch_datastream); + drain_orch_stdio_capture(orch_stdio_rx, orch_datastream, dashboard, run_id, node_id); + if stop_rx.try_recv().is_ok() { + orch_datastream.emit_bootstrap( + dashboard, + run_id, + node_id, + "shutdown", + "started", + json!({"source":"stdin"}), + ); + return Ok(()); + } + + if active.is_none() + && pipeline_runtime + .as_ref() + .is_none_or(|pipeline| !pipeline.is_active()) + && let Ok(work) = work_rx.try_recv() + { + let request = work.request; + let request_id = request.request_id; + orch_datastream.emit_prompt( + dashboard, + run_id, + node_id, + request_id, + "prompt_work", + "observed", + json!({ + "prompt_bytes":request.prompt_text.len(), + "max_tokens":request.max_tokens, + }), + ); + if let Some(pipeline) = pipeline_runtime.as_mut() { + pipeline.start_prompt( + request, + work.events, + &stack.runtime, + dashboard, + orch_datastream, + run_id, + node_id, + )?; + continue; + } + orch_datastream.emit_prompt( + dashboard, + run_id, + node_id, + request_id, + "node_prompt_send", + "started", + json!({"node_actor":node_actor,"reply_to":reply_to}), + ); + match stack.runtime.send_to( + node_actor, + NodeAgentMsg::InferPrompt { + request_id, + prompt: request.prompt_text.clone(), + max_tokens: request.max_tokens, + reply_to, + }, + ) { + Ok(()) => { + orch_datastream.emit_prompt( + dashboard, + run_id, + node_id, + request_id, + "node_prompt_send", + "ready", + json!({"node_actor":node_actor,"reply_to":reply_to}), + ); + active = Some(ActivePrompt { + request, + events: work.events, + }); + } + Err(error) => { + orch_datastream.emit_prompt( + dashboard, + run_id, + node_id, + request_id, + "node_prompt_send", + "failed", + json!({"node_actor":node_actor,"reply_to":reply_to,"error":error.to_string()}), + ); + return Err(format!("send prompt request: {error}")); + } + } + } + + while let Some(event) = prompt_events.try_recv() { + let request_id = event.request_id(); + let Some(current) = active.as_ref() else { + orch_datastream.emit_prompt( + dashboard, + run_id, + node_id, + request_id, + "node_prompt_event", + "dropped", + json!({"reason":"no_active_prompt","event":prompt_event_name(&event)}), + ); + continue; + }; + if request_id != current.request.request_id { + orch_datastream.emit_prompt( + dashboard, + run_id, + node_id, + request_id, + "node_prompt_event", + "dropped", + json!({ + "reason":"request_mismatch", + "event":prompt_event_name(&event), + "active_request_id":current.request.request_id, + }), + ); + continue; + } + + let terminal = event.is_terminal(); + let completion_status = prompt_completion_status(&event); + let completion_detail = prompt_completion_detail(&event); + orch_datastream.emit_prompt( + dashboard, + run_id, + node_id, + request_id, + "node_prompt_event", + "observed", + prompt_event_detail(&event), + ); + let _ = current.events.send(event); + if terminal { + orch_datastream.emit_prompt( + dashboard, + run_id, + node_id, + request_id, + "prompt_complete", + completion_status, + completion_detail, + ); + active = None; + } + } + + thread::sleep(PUMP_INTERVAL); + } +} + +fn prompt_event_name(event: &PromptEvent) -> &'static str { + match event { + PromptEvent::TextDelta { .. } => "TextDelta", + PromptEvent::Done { .. } => "Done", + PromptEvent::Fault { .. } => "Fault", + } +} + +fn prompt_event_detail(event: &PromptEvent) -> Value { + match event { + PromptEvent::TextDelta { text, .. } => { + json!({"event":"TextDelta","terminal":false,"text_bytes":text.len()}) + } + PromptEvent::Done { + final_text, + tokens_generated, + elapsed_ms, + .. + } => json!({ + "event":"Done", + "terminal":true, + "tokens_generated":tokens_generated, + "elapsed_ms":elapsed_ms, + "final_text_bytes":final_text.len(), + }), + PromptEvent::Fault { error, .. } => { + json!({"event":"Fault","terminal":true,"error":error}) + } + } +} + +fn prompt_completion_status(event: &PromptEvent) -> &'static str { + match event { + PromptEvent::Done { .. } => "ready", + PromptEvent::Fault { .. } => "failed", + PromptEvent::TextDelta { .. } => "observed", + } +} + +fn prompt_completion_detail(event: &PromptEvent) -> Value { + match event { + PromptEvent::Done { .. } => json!({"event":"Done"}), + PromptEvent::Fault { error, .. } => json!({"event":"Fault","error":error}), + PromptEvent::TextDelta { .. } => json!({"event":"TextDelta"}), + } +} + +fn stop_requested(stop_rx: &mpsc::Receiver<()>) -> bool { + stop_rx.try_recv().is_ok() +} + +fn spawn_stop_listener() -> mpsc::Receiver<()> { + let (tx, rx) = mpsc::channel(); + thread::spawn(move || { + let stdin = std::io::stdin(); + for line in stdin.lock().lines().map_while(Result::ok) { + let trimmed = line.trim(); + if trimmed.eq_ignore_ascii_case("stop") + || trimmed.eq_ignore_ascii_case("shutdown") + || trimmed.eq_ignore_ascii_case("quit") + { + let _ = tx.send(()); + break; + } + } + }); + rx +} + +fn drain_observations( + obs_rx: &mpsc::Receiver, + dashboard: Option<&DashboardSupport>, + orch_datastream: &mut OrchDatastream, + provider: ProviderKind, +) -> Result<(), String> { + while let Ok(observation) = obs_rx.try_recv() { + emit_plugin_observation(orch_datastream, dashboard, provider, &observation); + match observation { + PluginObservation::Failed { reason, .. } => return Err(reason), + PluginObservation::Exited { status, .. } => { + return Err(format!("node exited: {status:?}")); + } + PluginObservation::DatastreamFrame { .. } => {} + PluginObservation::ProviderLine { .. } + | PluginObservation::StdoutLine { .. } + | PluginObservation::StderrLine { .. } => {} + } + } + Ok(()) +} + +fn emit_plugin_observation( + orch_datastream: &mut OrchDatastream, + dashboard: Option<&DashboardSupport>, + provider: ProviderKind, + observation: &PluginObservation, +) { + match observation { + PluginObservation::StdoutLine { + run_id, + node_id, + line, + } => orch_datastream.emit_log( + dashboard, + ProvisionLogLine { + run_id: *run_id, + node_id: *node_id, + stream: ProvisionLogStream::Stdout, + line: line.clone(), + }, + ), + PluginObservation::StderrLine { + run_id, + node_id, + line, + } => orch_datastream.emit_log( + dashboard, + ProvisionLogLine { + run_id: *run_id, + node_id: *node_id, + stream: ProvisionLogStream::Stderr, + line: line.clone(), + }, + ), + PluginObservation::ProviderLine { + run_id, + node_id, + line, + } => orch_datastream.emit_log( + dashboard, + ProvisionLogLine { + run_id: *run_id, + node_id: *node_id, + stream: ProvisionLogStream::Provider, + line: line.clone(), + }, + ), + PluginObservation::DatastreamFrame { + channel, payload, .. + } => orch_datastream.emit_bytes_from( + dashboard, + channel, + payload.as_bytes().to_vec(), + "node_bootstrap_stdio", + ), + PluginObservation::Exited { + run_id, + node_id, + status, + } => orch_datastream.emit_event( + dashboard, + ProvisionEvent { + run_id: *run_id, + node_id: *node_id, + kind: ProvisionEventKind::NodeStopped, + provider: Some(provider.as_str().to_owned()), + message: Some(format!("node process exited with {status:?}")), + }, + ), + PluginObservation::Failed { + run_id, + node_id, + reason, + } => orch_datastream.emit_event( + dashboard, + ProvisionEvent { + run_id: *run_id, + node_id: *node_id, + kind: ProvisionEventKind::ProvisionFailed, + provider: Some(provider.as_str().to_owned()), + message: Some(reason.clone()), + }, + ), + } +} + +fn drain_frames( + frame_rx: &mpsc::Receiver, + dashboard: Option<&DashboardSupport>, + orch_datastream: &mut OrchDatastream, +) { + while let Ok(collected) = frame_rx.try_recv() { + ingest_dashboard_frame( + dashboard, + &collected.stream, + &collected.channel_name, + &collected.frame, + ); + orch_datastream.archive_frame( + "node_cluster", + &collected.stream, + &collected.channel_name, + &collected.frame, + ); + } +} + +fn ingest_dashboard_frame( + dashboard: Option<&DashboardSupport>, + stream: &StreamId, + channel: &str, + frame: &Frame, +) { + if let Some(dashboard) = dashboard { + dashboard.publish_frame(stream, channel, frame); + } +} + +fn emit_swim_transitions( + orch_datastream: &mut OrchDatastream, + dashboard: Option<&DashboardSupport>, + run_id: u64, + node_id: u64, + stack: &DistributionRuntimeStack, +) { + for transition in stack.drain_swim_transitions() { + orch_datastream.emit_bootstrap_to_channel( + dashboard, + MVP_SWIM_MEMBERSHIP, + run_id, + node_id, + "membership_transition", + "observed", + json!({ + "peer":format!("{:?}", transition.peer), + "from":transition.from.map(|state| format!("{:?}", state)), + "to":format!("{:?}", transition.to), + "reason":transition.reason, + }), + ); + } +} + +fn pump( + driver: &mut IrohDriver, + stack: &DistributionRuntimeStack, + frame_tx: &mpsc::Sender, +) { + stack.tick_protocol_actors(Instant::now()); + driver.pump_inbound_to_actors(); + stack.pump_runtime_once(); + driver.drain_outbox(&stack.outbox); + drain_datastream_connections(driver, frame_tx); +} + +fn env_optional(name: &str) -> Option { + std::env::var(name) + .ok() + .map(|value| value.trim().to_owned()) + .filter(|value| !value.is_empty()) +} + +fn docker_container_prefix() -> String { + env_optional(MVP_DOCKER_CONTAINER_PREFIX_ENV) + .unwrap_or_else(|| DEFAULT_DOCKER_CONTAINER_PREFIX.to_owned()) +} + +fn optional_env(name: &str) -> Option<(String, String)> { + env_optional(name).map(|value| (name.to_owned(), value)) +} + +fn local_tinygrad_worker_env(provider: ProviderKind) -> Option<(String, String)> { + optional_env("MVP_TINYGRAD_WORKER").or_else(|| { + if provider != ProviderKind::Process { + return None; + } + default_local_tinygrad_worker_path().map(|path| { + ( + "MVP_TINYGRAD_WORKER".to_owned(), + path.to_string_lossy().to_string(), + ) + }) + }) +} + +fn default_local_tinygrad_worker_path() -> Option { + let mut candidates = Vec::new(); + if let Ok(cwd) = std::env::current_dir() { + candidates.push(cwd.join("apps").join("mvp-node").join("tinygrad_worker.py")); + } + candidates.push( + PathBuf::from(env!("CARGO_MANIFEST_DIR")) + .join("..") + .join("..") + .join("apps") + .join("mvp-node") + .join("tinygrad_worker.py"), + ); + for candidate in candidates { + if candidate.is_file() { + return Some(candidate.canonicalize().unwrap_or(candidate)); + } + } + None +} + +fn resolve_vastai_ssh_identity(explicit: Option) -> Result { + match explicit { + Some(path) => Ok(path), + None => { + let home = std::env::var_os("HOME") + .filter(|value| !value.is_empty()) + .ok_or_else(|| { + "MVP_VASTAI_SSH_IDENTITY is required because HOME is unset".to_owned() + })?; + Ok(PathBuf::from(home).join(".ssh").join("id_ed25519")) + } + } +} + +fn expand_home_path(value: &str) -> Result { + let trimmed = value.trim(); + if let Some(rest) = trimmed.strip_prefix("~/") { + let home = std::env::var_os("HOME") + .filter(|value| !value.is_empty()) + .ok_or_else(|| "MVP_VASTAI_SSH_IDENTITY uses ~/ but HOME is unset".to_owned())?; + return Ok(PathBuf::from(home).join(rest)); + } + Ok(PathBuf::from(trimmed)) +} + +fn derive_ssh_public_key(identity: &Path) -> Result { + let output = Command::new("ssh-keygen") + .arg("-y") + .arg("-f") + .arg(identity) + .output() + .map_err(|e| { + format!( + "derive VastAI SSH public key from {}: {e}", + identity.display() + ) + })?; + let public_key = String::from_utf8_lossy(&output.stdout) + .trim_end_matches(['\r', '\n']) + .to_owned(); + if !output.status.success() || public_key.trim().is_empty() { + return Err(format!( + "derive VastAI SSH public key from {}: {}", + identity.display(), + command_output_failure_detail(&output, None) + )); + } + Ok(public_key) +} + +fn ssh_public_key_fingerprint(public_key: &str) -> String { + let path = std::env::temp_dir().join(format!("mvp-vastai-ssh-key-{}.pub", std::process::id())); + if std::fs::write(&path, format!("{public_key}\n")).is_err() { + return "unavailable".to_owned(); + } + let output = Command::new("ssh-keygen") + .arg("-l") + .arg("-f") + .arg(&path) + .output(); + let _ = std::fs::remove_file(&path); + let Ok(output) = output else { + return "unavailable".to_owned(); + }; + if !output.status.success() { + return "unavailable".to_owned(); + } + let stdout = String::from_utf8_lossy(&output.stdout); + let mut fields = stdout.split_whitespace(); + match (fields.next(), fields.next()) { + (Some(bits), Some(fingerprint)) => format!("{bits} {fingerprint}"), + _ => "unavailable".to_owned(), + } +} + +fn vastai_account_has_ssh_key(api_key: &str, public_key: &str) -> Result { + let output = Command::new("vastai") + .args(["show", "ssh-keys", "--raw", "--api-key", api_key]) + .output() + .map_err(vastai_cli_error)?; + if !output.status.success() { + return Err(format!( + "vastai show ssh-keys failed: {}", + command_output_failure_detail(&output, Some(api_key)) + )); + } + let stdout = String::from_utf8_lossy(&output.stdout); + Ok(account_ssh_keys_output_contains_public_key( + &stdout, public_key, + )) +} + +fn ensure_vastai_account_ssh_key(api_key: &str, public_key: &str) -> Result<(), String> { + if vastai_account_has_ssh_key(api_key, public_key)? { + return Ok(()); + } + + let output = Command::new("vastai") + .args(["create", "ssh-key"]) + .arg(public_key) + .args(["-y", "--api-key", api_key]) + .output() + .map_err(vastai_cli_error)?; + if !output.status.success() { + return Err(format!( + "vastai create ssh-key failed: {}", + command_output_failure_detail(&output, Some(api_key)) + )); + } + + if vastai_account_has_ssh_key(api_key, public_key)? { + Ok(()) + } else { + Err( + "VastAI SSH key registration did not make the selected key visible in vastai show ssh-keys" + .to_owned(), + ) + } +} + +fn account_ssh_keys_output_contains_public_key(output: &str, public_key: &str) -> bool { + let public_key = public_key.trim(); + if public_key.is_empty() { + return false; + } + if output.contains(public_key) { + return true; + } + public_key + .split_whitespace() + .nth(1) + .is_some_and(|body| !body.is_empty() && output.contains(body)) +} + +fn vastai_cli_error(error: std::io::Error) -> String { + if error.kind() == std::io::ErrorKind::NotFound { + "vastai CLI is required to verify/register MVP_VASTAI_SSH_IDENTITY; install with pip install vastai" + .to_owned() + } else { + format!("run vastai CLI: {error}") + } +} + +fn command_output_failure_detail(output: &std::process::Output, secret: Option<&str>) -> String { + let mut detail = String::from_utf8_lossy(&output.stderr).trim().to_owned(); + if detail.is_empty() { + detail = output.status.to_string(); + } + if let Some(secret) = secret.filter(|secret| !secret.is_empty()) { + detail = detail.replace(secret, ""); + } + detail +} + +#[cfg(test)] +fn relay_mode_from_env() -> Result { + relay_runtime_config_from_env(1).map(|relay| relay.mode) +} + +fn next_arg(args: &mut impl Iterator, name: &str) -> Result { + args.next() + .ok_or_else(|| format!("missing value after {name}")) +} + +fn parse_next(args: &mut impl Iterator, name: &str) -> Result +where + T: std::str::FromStr, + T::Err: std::fmt::Display, +{ + let value = next_arg(args, name)?; + value + .parse::() + .map_err(|e| format!("invalid {name}={value:?}: {e}")) +} + +#[cfg(test)] +mod tests { + use super::*; + use std::{ffi::OsString, path::PathBuf}; + + static ENV_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(()); + + const ENV_KEYS: &[&str] = &[ + CACHED_MODEL_HOST_ENV, + "HOME", + "HF_TOKEN", + "MVP_CPU_LINE_PROFILE", + "MVP_CPU_LINE_PROFILE_INTERVAL_MS", + "CUDA_DEVICE_SCHEDULE", + "MVP_DASHBOARD", + "MVP_DOCKER_GPUS", + "MVP_GGUF_FILE", + "MVP_GGUF_LOCAL_PATH", + "MVP_GGUF_REPO", + "MVP_GGUF_REVISION", + "MVP_IROH_RELAY_MODE", + MVP_IROH_RELAY_URL_ENV, + "MVP_LAYER_END_EXCLUSIVE", + "MVP_LOGICAL_NODE_ID", + "MVP_MODEL_CACHE_DIR", + "MVP_MAX_CONTEXT", + "MVP_MODEL_ID", + "MVP_NODE_IMAGE", + "MVP_NODE_PROVIDER", + "MVP_PROVIDER", + "MVP_PROMPT_MAX_TOKENS", + "MVP_PROMPT_RPC_BIND", + "MVP_RUN_ID", + "MVP_RUNTIME_CONFIG", + "MVP_PIPELINE_STAGES", + "MVP_STAGE_INDEX", + "MVP_TOKEN_PROGRESS_EVERY", + "MVP_TINYGRAD_WORKER", + "MVP_TOKENIZER_LOCAL_PATH", + "MVP_VASTAI_API_KEY", + "MVP_VASTAI_BOOTSTRAP_COMMAND", + "MVP_VASTAI_CONFIRM_LEASE", + "MVP_VASTAI_DISK_GB", + "MVP_VASTAI_GPU_NAME", + "MVP_VASTAI_MIN_DOWN_MBPS", + "MVP_VASTAI_MIN_GPU_RAM_MB", + "MVP_VASTAI_MIN_RELIABILITY", + "MVP_VASTAI_MIN_UP_MBPS", + "MVP_VASTAI_ONSTART", + "MVP_VASTAI_POLL_INTERVAL_SECS", + "MVP_VASTAI_REQUIRE_VERIFIED", + "MVP_VASTAI_SSH_USER", + "MVP_VASTAI_SSH_IDENTITY", + "VASTAI_API_KEY", + SWACTOR_IROH_RELAY_URL_ENV, + MVP_DOCKER_CONTAINER_PREFIX_ENV, + MVP_WORKER_BIN_ENV, + ]; + + struct RestoreEnv { + saved: Vec<(&'static str, Option)>, + } + + impl Drop for RestoreEnv { + fn drop(&mut self) { + for (key, value) in &self.saved { + match value { + Some(value) => unsafe { std::env::set_var(key, value) }, + None => unsafe { std::env::remove_var(key) }, + } + } + } + } + + fn with_clean_env(settings: &[(&'static str, &'static str)], test: impl FnOnce() -> T) -> T { + let settings = settings + .iter() + .map(|(key, value)| (*key, OsString::from(value))) + .collect::>(); + with_clean_env_os(&settings, test) + } + + fn with_clean_env_os(settings: &[(&'static str, OsString)], test: impl FnOnce() -> T) -> T { + let _lock = ENV_LOCK + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let saved = ENV_KEYS + .iter() + .map(|&key| (key, std::env::var_os(key))) + .collect::>(); + for key in ENV_KEYS { + unsafe { std::env::remove_var(key) }; + } + for (key, value) in settings { + assert!( + ENV_KEYS.contains(key), + "test env key {key} must be restored" + ); + unsafe { std::env::set_var(key, value) }; + } + let _restore = RestoreEnv { saved }; + test() + } + + fn selected_provider(settings: &[(&'static str, &'static str)]) -> ProviderKind { + with_clean_env(settings, || { + Config::from_layers_with_path_and_args(None, std::iter::empty::()) + .expect("config parses") + .provider + }) + } + + fn canonical_relay_url(raw: &str) -> String { + raw.parse::() + .expect("fixture relay URL parses") + .to_string() + } + + fn node_spec_env(settings: &[(&'static str, &'static str)]) -> Vec<(String, String)> { + with_clean_env(settings, || { + let config = Config::from_layers_with_path_and_args(None, std::iter::empty::()) + .expect("config parses"); + let coordinator = EndpointAddr::new(iroh::SecretKey::from_bytes(&[9; 32]).public()); + let orchestrator_actor = ActorAddress([12; 32]); + config + .node_spec(coordinator, orchestrator_actor) + .expect("node spec builds") + .env + }) + } + + fn env_value<'a>(env: &'a [(String, String)], key: &str) -> Option<&'a str> { + env.iter() + .find(|(env_key, _)| env_key == key) + .map(|(_, value)| value.as_str()) + } + + fn cached_model_config_with_args(model: &TempModelFile, args: &[&str]) -> Config { + with_clean_env_os( + &[ + ("MVP_RUNTIME_CONFIG", OsString::from("local")), + ("MVP_NODE_PROVIDER", OsString::from("docker")), + (CACHED_MODEL_HOST_ENV, model.raw_path.as_os_str().to_owned()), + ], + || { + Config::from_layers_with_path_and_args( + None, + args.iter().copied().map(str::to_owned), + ) + .expect("cached-model config parses") + }, + ) + } + + fn pipeline_config_with_cached_model(model: &TempModelFile, pipeline_stages: u32) -> Config { + let stages_arg = pipeline_stages.to_string(); + cached_model_config_with_args( + model, + &[ + "--pipeline-stages", + stages_arg.as_str(), + "--max-context", + "512", + ], + ) + } + + fn expected_layer_ranges(num_layers: u32, stage_count: u32) -> Vec<(u32, u32)> { + (0..stage_count) + .map(|stage_index| { + let start = (u64::from(num_layers) * u64::from(stage_index) + / u64::from(stage_count)) as u32; + let end = (u64::from(num_layers) * u64::from(stage_index + 1) + / u64::from(stage_count)) as u32; + (start, end) + }) + .collect() + } + #[test] + fn pipeline_weight_load_scheduler_keeps_one_active_stage() { + let model = TempModelFile::with_metadata( + "seven-stage-scheduler.gguf", + TestGgufMetadata { + num_layers: 30, + ..TestGgufMetadata::default() + }, + ); + let config = pipeline_config_with_cached_model(&model, 7); + let plan = config.build_run_plan().expect("seven-stage plan builds"); + let mut loaded = BTreeSet::new(); + + let first = + next_pipeline_weight_load_stage(&plan, &loaded, None).expect("first stage selected"); + assert_eq!(first.stage_index, 0); + + let resent = next_pipeline_weight_load_stage(&plan, &loaded, Some(first.stage_index)) + .expect("active stage is resent before it loads"); + assert_eq!(resent.stage_index, 0); + + for expected_stage in 0..7 { + let active = next_pipeline_weight_load_stage(&plan, &loaded, None) + .expect("next unloaded stage selected"); + assert_eq!(active.stage_index, expected_stage); + loaded.insert(expected_stage); + } + + assert!( + next_pipeline_weight_load_stage(&plan, &loaded, None).is_none(), + "all stages loaded should leave no active load" + ); + } + + fn assert_plan_matches_metadata( + config: &Config, + plan: &run_plan::RunPlan, + metadata: TestGgufMetadata, + stage_count: u32, + requested_context: u64, + ) { + assert_eq!(plan.run_id, run_plan::RunId(config.run_id)); + assert_eq!(plan.model.model_id, config.model_id); + assert_eq!(plan.model.gguf_source, config.gguf_source); + assert_eq!(plan.model.num_layers, metadata.num_layers); + assert_eq!(plan.model.hidden_dim, metadata.hidden_dim as u32); + assert_eq!( + plan.model.max_seq_len, + requested_context.min(metadata.context_length) as u32 + ); + assert_eq!(plan.model.eos_token_id, metadata.eos_token_id); + assert_eq!(plan.stages.len(), stage_count as usize); + assert_eq!(plan.edges.len(), stage_count as usize + 1); + assert_eq!(plan.max_tokens, config.default_max_tokens); + + let mut stages = plan.stages.clone(); + stages.sort_by_key(|stage| stage.stage_index); + let expected_ranges = expected_layer_ranges(metadata.num_layers, stage_count); + for (stage, expected_range) in stages.iter().zip(expected_ranges) { + assert_eq!(stage.stage_count, stage_count); + assert_eq!( + stage.node_id, + run_plan::NodeId(config.node_id + 1 + u64::from(stage.stage_index)) + ); + assert_eq!( + (stage.layer_start, stage.layer_end_exclusive), + expected_range + ); + assert!( + stage.layer_start < stage.layer_end_exclusive, + "stage {} must have a non-empty layer range", + stage.stage_index + ); + } + } + + fn token_object_spec(max_extent: u64, alignment: u32) -> run_plan::ObjectSpec { + run_plan::ObjectSpec { + kind: run_plan::ObjectKind::Token, + max_extent, + dtype_family: run_plan::DTypeFamily::BFloat, + dtype_width_bytes: 4, + shape: run_plan::ShapeRule::TokenIds, + layout: run_plan::LayoutRule::Contiguous, + alignment, + sequence_policy: run_plan::SequencePolicy::Ordered, + } + } + + fn decode_token_record_payload( + bytes: &[u8], + spec: run_plan::ObjectSpec, + ) -> (u64, Vec, bool) { + let read = ingress::read_object_record(bytes, ingress_object_spec_from_plan(spec), false) + .expect("token record should be valid"); + let ingress::ObjectRecordRead::Complete(record) = read else { + panic!("token record should be complete"); + }; + let payload = record.payload(bytes).expect("token payload is present"); + let tokens = payload + .chunks_exact(4) + .map(|chunk| u32::from_le_bytes(chunk.try_into().unwrap())) + .collect::>(); + assert!(payload.chunks_exact(4).remainder().is_empty()); + (record.sequence, tokens, record.flags.end_of_sequence) + } + + fn decode_token_record_with_flags( + bytes: &[u8], + spec: run_plan::ObjectSpec, + ) -> (u64, Vec, bool, bool) { + let read = ingress::read_object_record(bytes, ingress_object_spec_from_plan(spec), false) + .expect("token record should be valid"); + let ingress::ObjectRecordRead::Complete(record) = read else { + panic!("token record should be complete"); + }; + let payload = record.payload(bytes).expect("token payload is present"); + let tokens = payload + .chunks_exact(4) + .map(|chunk| u32::from_le_bytes(chunk.try_into().unwrap())) + .collect::>(); + assert!(payload.chunks_exact(4).remainder().is_empty()); + ( + record.sequence, + tokens, + record.flags.end_of_sequence, + record.flags.begin_sequence, + ) + } + + struct PipelineRuntimeTestFixture { + actor_runtime: Arc, + tokenizer_events: swactor::runtime::Inbox, + encode_requests: swactor::runtime::Inbox, + decode_requests: swactor::runtime::Inbox, + runtime: PipelinePromptRuntime, + token_in_rx: tokio_mpsc::UnboundedReceiver>, + } + + fn pipeline_runtime_fixture() -> PipelineRuntimeTestFixture { + let spec = token_object_spec(64, 64); + let (token_in_tx, token_in_rx) = tokio_mpsc::unbounded_channel(); + let (recv_tx, recv_rx) = mpsc::channel(); + pipeline_runtime_fixture_from_specs( + 11, + 12, + spec, + spec, + token_in_tx, + token_in_rx, + recv_tx, + recv_rx, + ) + } + + fn pipeline_runtime_fixture_from_plan(stage_count: u32) -> PipelineRuntimeTestFixture { + let model = TempModelFile::new(&format!("prompt-plan-{stage_count}.gguf")); + let config = pipeline_config_with_cached_model(&model, stage_count); + let plan = config.build_run_plan().expect("prompt runtime plan builds"); + let token_in = plan + .edges + .iter() + .find(|edge| edge.kind == run_plan::EdgeKind::TokenIn) + .expect("token-in edge"); + let token_out = plan + .edges + .iter() + .find(|edge| edge.kind == run_plan::EdgeKind::TokenOut) + .expect("token-out edge"); + let (token_in_tx, token_in_rx) = tokio_mpsc::unbounded_channel(); + let (recv_tx, recv_rx) = mpsc::channel(); + pipeline_runtime_fixture_from_specs( + token_in.edge_id.0, + token_out.edge_id.0, + token_in.object_spec, + token_out.object_spec, + token_in_tx, + token_in_rx, + recv_tx, + recv_rx, + ) + } + + #[allow(clippy::too_many_arguments)] + fn pipeline_runtime_fixture_from_specs( + token_in_edge_id: u64, + token_out_edge_id: u64, + token_spec: run_plan::ObjectSpec, + token_out_spec: run_plan::ObjectSpec, + token_in_tx: tokio_mpsc::UnboundedSender>, + token_in_rx: tokio_mpsc::UnboundedReceiver>, + recv_tx: mpsc::Sender>, + recv_rx: mpsc::Receiver>, + ) -> PipelineRuntimeTestFixture { + let actor_runtime = Arc::new(swactor::runtime::Runtime::new( + swactor::config::RuntimeConfig::default(), + )); + let encode_requests = actor_runtime + .new_inbox::() + .expect("encode request inbox"); + let decode_requests = actor_runtime + .new_inbox::() + .expect("decode request inbox"); + let tokenizer_events = actor_runtime + .new_inbox::() + .expect("tokenizer event inbox"); + PipelineRuntimeTestFixture { + runtime: PipelinePromptRuntime { + token_in_edge_id, + token_out_edge_id, + token_spec, + token_out_spec, + token_in_sender: PipelineSendHandle::Channel(token_in_tx), + recv_rx, + recv_tx, + recv_buffer: Vec::new(), + tokenizer_encode_actor: *encode_requests.addr(), + tokenizer_decode_actor: *decode_requests.addr(), + tokenizer_reply_to: *tokenizer_events.addr(), + pending_encode: None, + pending_decode: None, + next_sequence: 0, + generated_tokens: Vec::new(), + final_text: String::new(), + active: None, + started_at: None, + }, + actor_runtime, + tokenizer_events, + encode_requests, + decode_requests, + token_in_rx, + } + } + + fn start_fixture_prompt( + fixture: &mut PipelineRuntimeTestFixture, + request: SubmitPrompt, + events: mpsc::Sender, + datastream: &mut OrchDatastream, + run_id: u64, + node_id: u64, + ) { + fixture + .runtime + .start_prompt( + request, + events, + &fixture.actor_runtime, + None, + datastream, + run_id, + node_id, + ) + .expect("pipeline prompt starts"); + } + + fn drain_fixture_tokenizer( + fixture: &mut PipelineRuntimeTestFixture, + datastream: &mut OrchDatastream, + run_id: u64, + node_id: u64, + ) { + fixture + .runtime + .drain_tokenizer_events( + &fixture.actor_runtime, + &fixture.tokenizer_events, + None, + datastream, + run_id, + node_id, + ) + .expect("tokenizer events drain"); + } + + fn send_encoded_tokens( + fixture: &PipelineRuntimeTestFixture, + request_id: u64, + tokens: Vec, + ) { + fixture + .actor_runtime + .send_to( + *fixture.tokenizer_events.addr(), + TokenizerEvent::PromptEncoded { request_id, tokens }, + ) + .expect("send encoded tokens"); + } + + fn send_decoded_text(fixture: &PipelineRuntimeTestFixture, request_id: u64, text: &str) { + fixture + .actor_runtime + .send_to( + *fixture.tokenizer_events.addr(), + TokenizerEvent::TokensDecoded { + request_id, + text: text.to_owned(), + }, + ) + .expect("send decoded text"); + } + + fn assert_encode_request(fixture: &PipelineRuntimeTestFixture, request_id: u64, prompt: &str) { + match fixture.encode_requests.try_recv() { + Some(NodeAgentMsg::EncodePrompt { + request_id: actual_request_id, + prompt: actual_prompt, + reply_to, + }) => { + assert_eq!(actual_request_id, request_id); + assert_eq!(actual_prompt, prompt); + assert_eq!(reply_to, *fixture.tokenizer_events.addr()); + } + other => panic!("expected EncodePrompt request, got {other:?}"), + } + } + + fn assert_decode_request(fixture: &PipelineRuntimeTestFixture, request_id: u64, token: u32) { + match fixture.decode_requests.try_recv() { + Some(NodeAgentMsg::DecodeTokens { + request_id: actual_request_id, + tokens, + reply_to, + }) => { + assert_eq!(actual_request_id, request_id); + assert_eq!(tokens, vec![token]); + assert_eq!(reply_to, *fixture.tokenizer_events.addr()); + } + other => panic!("expected DecodeTokens request, got {other:?}"), + } + } + + #[test] + fn pipeline_prompt_token_record_round_trips_one_token_with_eos_and_plan_ring_alignment() { + let spec = token_object_spec(64, 64); + let bytes = encode_token_record(spec, 12, 7, &[513], true).expect("token record encodes"); + assert_eq!(&bytes[0..4], b"MO01"); + + let mut partial = bytes[..run_plan::MO01_HEADER_BYTES as usize + 2].to_vec(); + assert!( + take_pipeline_token_record(&mut partial, spec) + .expect("partial token record is not malformed") + .is_none() + ); + assert_eq!(partial.len(), run_plan::MO01_HEADER_BYTES as usize + 2); + + let mut buffer = bytes; + let record = take_pipeline_token_record(&mut buffer, spec) + .expect("one-u32 token-out payload should parse despite ring alignment") + .expect("complete token-out record is available"); + + assert_eq!(record.object_id, 9007); + assert_eq!(record.sequence, 7); + assert_eq!(record.token_id, 513); + assert!(record.eos); + assert!(buffer.is_empty()); + } + + #[test] + fn plan_derived_mo01_specs_accept_token_and_activation_record_extents() { + let metadata = TestGgufMetadata { + num_layers: 7, + hidden_dim: 13, + context_length: 64, + eos_token_id: 11, + }; + let model = TempModelFile::with_metadata("plan-mo01.gguf", metadata); + let config = pipeline_config_with_cached_model(&model, 3); + let plan = config + .build_run_plan() + .expect("plan-derived MO01 plan builds"); + let token_in = plan + .edges + .iter() + .find(|edge| edge.kind == run_plan::EdgeKind::TokenIn) + .expect("token-in edge"); + let activation = plan + .edges + .iter() + .find(|edge| edge.kind == run_plan::EdgeKind::Activation) + .expect("activation edge"); + + let token_record = encode_token_record( + token_in.object_spec, + token_in.edge_id.0, + 0, + &[65, 195, 169], + false, + ) + .expect("plan token record encodes"); + assert_eq!( + decode_token_record_payload(&token_record, token_in.object_spec), + (0, vec![65, 195, 169], false) + ); + + let activation_payload = vec![ + 0_u8; + usize::try_from(metadata.hidden_dim * 2) + .expect("activation payload size fits usize") + ]; + let activation_record = ingress::ObjectRecordBuilder::new(ingress_object_spec_from_plan( + activation.object_spec, + )) + .object_id(ingress::ObjectId(9100)) + .sequence(0) + .payload(activation_payload) + .encode(); + let read = ingress::read_object_record( + &activation_record, + ingress_object_spec_from_plan(activation.object_spec), + false, + ) + .expect("plan activation record parses"); + let ingress::ObjectRecordRead::Complete(record) = read else { + panic!("activation record should be complete"); + }; + assert_eq!(record.payload(&activation_record).unwrap().len(), 26); + } + + #[test] + fn pipeline_prompt_token_out_parser_rejects_payloads_that_are_not_one_u32() { + let spec = token_object_spec(64, 64); + let mut buffer = + encode_token_record(spec, 12, 0, &[65, 66], false).expect("multi-token record encodes"); + + let error = take_pipeline_token_record(&mut buffer, spec) + .expect_err("token-out records must carry exactly one generated token"); + + assert!( + error.contains("token-out payload must be exactly one u32, got 8"), + "unexpected error: {error}" + ); + } + + #[test] + fn pipeline_prompt_runtime_uses_tokenizer_events_for_encode_decode_and_continuations() { + let mut fixture = pipeline_runtime_fixture(); + let (event_tx, event_rx) = mpsc::channel(); + let mut datastream = OrchDatastream::new(91, None).expect("datastream opens"); + start_fixture_prompt( + &mut fixture, + SubmitPrompt { + request_id: 42, + prompt_text: "Hi".to_owned(), + max_tokens: 2, + }, + event_tx, + &mut datastream, + 91, + 3, + ); + assert_encode_request(&fixture, 42, "Hi"); + send_encoded_tokens(&fixture, 42, vec![1001, 1002]); + drain_fixture_tokenizer(&mut fixture, &mut datastream, 91, 3); + + let initial = fixture + .token_in_rx + .try_recv() + .expect("tokenizer tokens are sent to token-in"); + assert_eq!( + decode_token_record_payload(&initial, fixture.runtime.token_spec), + (0, vec![1001, 1002], false) + ); + assert!(fixture.token_in_rx.try_recv().is_err()); + + fixture + .runtime + .recv_tx + .send( + encode_token_record( + fixture.runtime.token_out_spec, + fixture.runtime.token_out_edge_id, + 0, + &[79], + false, + ) + .expect("first token-out record encodes"), + ) + .expect("token-out bytes enqueue"); + fixture + .runtime + .drain_tokens(&fixture.actor_runtime, None, &mut datastream, 91, 3) + .expect("first token drains"); + assert_decode_request(&fixture, 42, 79); + assert!(event_rx.try_recv().is_err()); + + send_decoded_text(&fixture, 42, "O"); + drain_fixture_tokenizer(&mut fixture, &mut datastream, 91, 3); + assert_eq!( + event_rx.try_recv().expect("first delta event"), + PromptEvent::TextDelta { + request_id: 42, + text: "O".to_owned(), + } + ); + let continuation = fixture + .token_in_rx + .try_recv() + .expect("non-terminal token is fed back to token-in"); + assert_eq!( + decode_token_record_payload(&continuation, fixture.runtime.token_spec), + (1, vec![79], false) + ); + + fixture + .runtime + .recv_tx + .send( + encode_token_record( + fixture.runtime.token_out_spec, + fixture.runtime.token_out_edge_id, + 1, + &[75], + false, + ) + .expect("second token-out record encodes"), + ) + .expect("token-out bytes enqueue"); + fixture + .runtime + .drain_tokens(&fixture.actor_runtime, None, &mut datastream, 91, 3) + .expect("second token drains"); + assert_decode_request(&fixture, 42, 75); + + send_decoded_text(&fixture, 42, "K"); + drain_fixture_tokenizer(&mut fixture, &mut datastream, 91, 3); + assert_eq!( + event_rx.try_recv().expect("second delta event"), + PromptEvent::TextDelta { + request_id: 42, + text: "K".to_owned(), + } + ); + match event_rx.try_recv().expect("done event at max tokens") { + PromptEvent::Done { + request_id, + final_text, + tokens_generated, + .. + } => { + assert_eq!(request_id, 42); + assert_eq!(final_text, "OK"); + assert_eq!(tokens_generated, 2); + } + event => panic!("expected Done at max tokens, got {event:?}"), + } + assert!(event_rx.try_recv().is_err()); + assert!(fixture.token_in_rx.try_recv().is_err()); + assert!(!fixture.runtime.is_active()); + } + + #[test] + fn pipeline_prompt_runtime_marks_each_prompt_start_without_resetting_stream_sequence() { + let mut fixture = pipeline_runtime_fixture(); + let (first_event_tx, first_event_rx) = mpsc::channel(); + let mut datastream = OrchDatastream::new(94, None).expect("datastream opens"); + start_fixture_prompt( + &mut fixture, + SubmitPrompt { + request_id: 101, + prompt_text: "first".to_owned(), + max_tokens: 1, + }, + first_event_tx, + &mut datastream, + 94, + 3, + ); + assert_encode_request(&fixture, 101, "first"); + send_encoded_tokens(&fixture, 101, vec![10, 11]); + drain_fixture_tokenizer(&mut fixture, &mut datastream, 94, 3); + let first_initial = fixture + .token_in_rx + .try_recv() + .expect("first prompt token-in record is emitted"); + assert_eq!( + decode_token_record_with_flags(&first_initial, fixture.runtime.token_spec), + (0, vec![10, 11], false, true) + ); + fixture + .runtime + .recv_tx + .send( + encode_token_record( + fixture.runtime.token_out_spec, + fixture.runtime.token_out_edge_id, + 0, + &[21], + false, + ) + .expect("first token-out record encodes"), + ) + .expect("first token-out bytes enqueue"); + fixture + .runtime + .drain_tokens(&fixture.actor_runtime, None, &mut datastream, 94, 3) + .expect("first token drains"); + assert_decode_request(&fixture, 101, 21); + send_decoded_text(&fixture, 101, "A"); + drain_fixture_tokenizer(&mut fixture, &mut datastream, 94, 3); + assert!(matches!( + first_event_rx.try_recv().expect("first delta"), + PromptEvent::TextDelta { + request_id: 101, + .. + } + )); + assert!(matches!( + first_event_rx.try_recv().expect("first done"), + PromptEvent::Done { + request_id: 101, + .. + } + )); + assert!(!fixture.runtime.is_active()); + + let (second_event_tx, _second_event_rx) = mpsc::channel(); + start_fixture_prompt( + &mut fixture, + SubmitPrompt { + request_id: 102, + prompt_text: "second".to_owned(), + max_tokens: 1, + }, + second_event_tx, + &mut datastream, + 94, + 3, + ); + assert_encode_request(&fixture, 102, "second"); + send_encoded_tokens(&fixture, 102, vec![20, 22, 24]); + drain_fixture_tokenizer(&mut fixture, &mut datastream, 94, 3); + let second_initial = fixture + .token_in_rx + .try_recv() + .expect("second prompt token-in record is emitted"); + assert_eq!( + decode_token_record_with_flags(&second_initial, fixture.runtime.token_spec), + (1, vec![20, 22, 24], false, true) + ); + } + + #[test] + fn prompt_runtime_uses_plan_derived_token_edges_for_cached_n_one() { + let mut fixture = pipeline_runtime_fixture_from_plan(1); + let (event_tx, event_rx) = mpsc::channel(); + let mut datastream = OrchDatastream::new(93, None).expect("datastream opens"); + start_fixture_prompt( + &mut fixture, + SubmitPrompt { + request_id: 88, + prompt_text: "Aé".to_owned(), + max_tokens: 1, + }, + event_tx, + &mut datastream, + 93, + 3, + ); + assert_encode_request(&fixture, 88, "Aé"); + send_encoded_tokens(&fixture, 88, vec![321, 654]); + drain_fixture_tokenizer(&mut fixture, &mut datastream, 93, 3); + + let initial = fixture + .token_in_rx + .try_recv() + .expect("plan-derived token-in record is emitted"); + assert_eq!( + decode_token_record_payload(&initial, fixture.runtime.token_spec), + (0, vec![321, 654], false) + ); + assert_eq!(fixture.runtime.token_in_edge_id, 1); + assert_eq!(fixture.runtime.token_out_edge_id, 2); + + fixture + .runtime + .recv_tx + .send( + encode_token_record( + fixture.runtime.token_out_spec, + fixture.runtime.token_out_edge_id, + 0, + &[33], + false, + ) + .expect("plan-derived token-out record encodes"), + ) + .expect("token-out bytes enqueue"); + fixture + .runtime + .drain_tokens(&fixture.actor_runtime, None, &mut datastream, 93, 3) + .expect("plan-derived token drains"); + assert_decode_request(&fixture, 88, 33); + send_decoded_text(&fixture, 88, "!"); + drain_fixture_tokenizer(&mut fixture, &mut datastream, 93, 3); + + assert_eq!( + event_rx.try_recv().expect("delta event"), + PromptEvent::TextDelta { + request_id: 88, + text: "!".to_owned(), + } + ); + match event_rx.try_recv().expect("done event at max tokens") { + PromptEvent::Done { + request_id, + final_text, + tokens_generated, + .. + } => { + assert_eq!(request_id, 88); + assert_eq!(final_text, "!"); + assert_eq!(tokens_generated, 1); + } + event => panic!("expected Done at max tokens, got {event:?}"), + } + assert!(fixture.token_in_rx.try_recv().is_err()); + } + + #[test] + fn pipeline_prompt_runtime_finishes_on_eos_without_feedback_token() { + let mut fixture = pipeline_runtime_fixture(); + let (event_tx, event_rx) = mpsc::channel(); + let mut datastream = OrchDatastream::new(92, None).expect("datastream opens"); + start_fixture_prompt( + &mut fixture, + SubmitPrompt { + request_id: 77, + prompt_text: "go".to_owned(), + max_tokens: 8, + }, + event_tx, + &mut datastream, + 92, + 4, + ); + assert_encode_request(&fixture, 77, "go"); + send_encoded_tokens(&fixture, 77, vec![700]); + drain_fixture_tokenizer(&mut fixture, &mut datastream, 92, 4); + fixture + .token_in_rx + .try_recv() + .expect("initial prompt token-in record is sent"); + + fixture + .runtime + .recv_tx + .send( + encode_token_record( + fixture.runtime.token_out_spec, + fixture.runtime.token_out_edge_id, + 0, + &[33], + true, + ) + .expect("eos token-out record encodes"), + ) + .expect("token-out bytes enqueue"); + fixture + .runtime + .drain_tokens(&fixture.actor_runtime, None, &mut datastream, 92, 4) + .expect("eos token drains"); + assert_decode_request(&fixture, 77, 33); + send_decoded_text(&fixture, 77, "!"); + drain_fixture_tokenizer(&mut fixture, &mut datastream, 92, 4); + + assert_eq!( + event_rx.try_recv().expect("delta event before eos done"), + PromptEvent::TextDelta { + request_id: 77, + text: "!".to_owned(), + } + ); + match event_rx.try_recv().expect("done event on eos") { + PromptEvent::Done { + request_id, + final_text, + tokens_generated, + .. + } => { + assert_eq!(request_id, 77); + assert_eq!(final_text, "!"); + assert_eq!(tokens_generated, 1); + } + event => panic!("expected Done on eos, got {event:?}"), + } + assert!(event_rx.try_recv().is_err()); + assert!(fixture.token_in_rx.try_recv().is_err()); + assert!(!fixture.runtime.is_active()); + } + + #[test] + fn planned_prompt_runtime_uses_pipeline_edges_for_single_or_multi_stage_cached_models() { + assert_eq!( + prompt_runtime_mode(None), + PromptRuntimeMode::DirectInferPrompt + ); + + for stage_count in [1_u32, 3] { + let model = TempModelFile::new(&format!("prompt-runtime-{stage_count}.gguf")); + let config = pipeline_config_with_cached_model(&model, stage_count); + let plan = config + .build_run_plan() + .expect("cached-model planned prompt runtime plan builds"); + + assert_eq!( + prompt_runtime_mode(Some(&plan)), + PromptRuntimeMode::PipelineTokenEdges, + "cached-model N={stage_count} must use the planned token-edge runtime" + ); + } + } + + #[test] + fn docker_container_prefix_defaults_and_trims_env_override() { + with_clean_env(&[], || { + assert_eq!(docker_container_prefix(), DEFAULT_DOCKER_CONTAINER_PREFIX); + }); + with_clean_env( + &[(MVP_DOCKER_CONTAINER_PREFIX_ENV, " custom-prefix ")], + || { + assert_eq!(docker_container_prefix(), "custom-prefix"); + }, + ); + } + + #[test] + fn expand_home_path_expands_leading_home_segment() { + with_clean_env(&[("HOME", "/tmp/mvp-vastai-home")], || { + assert_eq!( + expand_home_path("~/keys/deploy").expect("home path expands"), + PathBuf::from("/tmp/mvp-vastai-home/keys/deploy") + ); + assert_eq!( + expand_home_path("/tmp/not-~/expanded").expect("literal path stays literal"), + PathBuf::from("/tmp/not-~/expanded") + ); + }); + } + + #[test] + fn account_ssh_keys_output_contains_public_key_matches_exact_key_material() { + let public_key = "ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAITestKeyBody vastai"; + + assert!(account_ssh_keys_output_contains_public_key( + public_key, public_key + )); + assert!(account_ssh_keys_output_contains_public_key( + r#"{"keys":[{"public_key":"AAAAC3NzaC1lZDI1NTE5AAAAITestKeyBody"}]}"#, + public_key + )); + assert!(!account_ssh_keys_output_contains_public_key( + r#"{"keys":[{"public_key":"AAAAC3NzaC1lZDI1NTE5AAAADifferent"}]}"#, + public_key + )); + } + + #[test] + fn vastai_config_reads_ssh_identity_without_runtime_preparation() { + let config = with_clean_env( + &[ + ("MVP_RUNTIME_CONFIG", "deploy"), + ("MVP_NODE_PROVIDER", "vastai"), + ("MVP_VASTAI_API_KEY", "vast-key"), + ("MVP_VASTAI_BOOTSTRAP_COMMAND", "/usr/local/bin/mvp-node"), + ("MVP_VASTAI_SSH_IDENTITY", "/tmp/mvp-vastai-key"), + ], + || { + Config::from_layers_with_path_and_args(None, std::iter::empty::()) + .expect("vastai config parses without ssh-keygen or vastai CLI") + }, + ); + let vastai = config.vastai.expect("vastai config is present"); + assert_eq!( + vastai.ssh_identity.as_deref(), + Some(Path::new("/tmp/mvp-vastai-key")) + ); + assert_eq!(vastai.ssh_public_key, None); + assert_eq!(vastai.ssh_public_fingerprint, None); + } + + #[derive(Clone, Copy)] + struct TestGgufMetadata { + num_layers: u32, + hidden_dim: u64, + context_length: u64, + eos_token_id: u32, + } + + impl Default for TestGgufMetadata { + fn default() -> Self { + Self { + num_layers: 7, + hidden_dim: 13, + context_length: 64, + eos_token_id: 11, + } + } + } + + struct TempModelFile { + root: PathBuf, + raw_path: PathBuf, + canonical_path: PathBuf, + metadata: TestGgufMetadata, + } + + impl TempModelFile { + fn new(file_name: &str) -> Self { + Self::with_metadata(file_name, TestGgufMetadata::default()) + } + + fn with_metadata(file_name: &str, metadata: TestGgufMetadata) -> Self { + let root = std::env::temp_dir().join(format!( + "mvp-cached-model-test-{}-{}", + std::process::id(), + std::thread::current().name().unwrap_or("unnamed") + )); + let _ = std::fs::remove_dir_all(&root); + std::fs::create_dir_all(root.join("nested")).expect("create temp model dir"); + let canonical_path = root.join(file_name); + std::fs::write(&canonical_path, minimal_gguf(metadata)) + .expect("write temp GGUF metadata file"); + let raw_path = root.join("nested").join("..").join(file_name); + Self { + root, + raw_path, + canonical_path: canonical_path + .canonicalize() + .expect("canonicalize temp model file"), + metadata, + } + } + } + + fn minimal_gguf(metadata: TestGgufMetadata) -> Vec { + let mut bytes = Vec::new(); + bytes.extend_from_slice(b"GGUF"); + bytes.extend_from_slice(&3_u32.to_le_bytes()); + bytes.extend_from_slice(&0_u64.to_le_bytes()); + bytes.extend_from_slice(&6_u64.to_le_bytes()); + push_string_kv(&mut bytes, "general.architecture", "llama"); + push_string_kv(&mut bytes, "general.name", "fixture"); + push_u32_kv(&mut bytes, "llama.block_count", metadata.num_layers); + push_u32_kv( + &mut bytes, + "llama.embedding_length", + metadata + .hidden_dim + .try_into() + .expect("test hidden dimension fits u32"), + ); + push_u32_kv( + &mut bytes, + "llama.context_length", + metadata + .context_length + .try_into() + .expect("test context length fits u32"), + ); + push_u32_kv( + &mut bytes, + "tokenizer.ggml.eos_token_id", + metadata.eos_token_id, + ); + bytes + } + + fn push_string_kv(bytes: &mut Vec, key: &str, value: &str) { + push_gguf_string(bytes, key); + bytes.extend_from_slice(&8_u32.to_le_bytes()); + push_gguf_string(bytes, value); + } + + fn push_u32_kv(bytes: &mut Vec, key: &str, value: u32) { + push_gguf_string(bytes, key); + bytes.extend_from_slice(&4_u32.to_le_bytes()); + bytes.extend_from_slice(&value.to_le_bytes()); + } + + fn push_gguf_string(bytes: &mut Vec, value: &str) { + bytes.extend_from_slice(&(value.len() as u64).to_le_bytes()); + bytes.extend_from_slice(value.as_bytes()); + } + + impl Drop for TempModelFile { + fn drop(&mut self) { + let _ = std::fs::remove_dir_all(&self.root); + } + } + + struct TempTomlFile { + path: PathBuf, + } + + impl TempTomlFile { + fn new(file_name: &str, contents: &str) -> Self { + let path = std::env::temp_dir().join(format!( + "mvp-orchestrator-config-test-{}-{}-{file_name}", + std::process::id(), + std::thread::current().name().unwrap_or("unnamed") + )); + let _ = std::fs::remove_file(&path); + std::fs::write(&path, contents).expect("write temp TOML config"); + Self { path } + } + } + + impl Drop for TempTomlFile { + fn drop(&mut self) { + let _ = std::fs::remove_file(&self.path); + } + } + + #[test] + fn frame_archive_writes_jsonl_records_for_text_and_binary_payloads() { + static NEXT_TEMP_FILE: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0); + + let suffix = NEXT_TEMP_FILE.fetch_add(1, std::sync::atomic::Ordering::Relaxed); + let path = std::env::temp_dir().join(format!( + "mvp-frame-archive-test-{}-{suffix}.jsonl", + std::process::id() + )); + let _ = std::fs::remove_file(&path); + + let stream = StreamId::new("test-node", Lifetime(42)); + let mut archive = FrameArchive::open(&path).expect("frame archive opens"); + archive.record( + "orchestrator", + &stream, + "stdout", + &Frame::new( + ChannelId(1), + datastream::Position(7), + b"hello \xce\xbb".to_vec(), + ), + ); + archive.record( + "orchestrator", + &stream, + "stderr", + &Frame::new( + ChannelId(2), + datastream::Position(8), + vec![0xff, 0x00, b'A'], + ), + ); + drop(archive); + + let contents = std::fs::read_to_string(&path).expect("read frame archive jsonl"); + let records = contents + .lines() + .map(|line| serde_json::from_str::(line).expect("archive line is json")) + .collect::>(); + let _ = std::fs::remove_file(&path); + + assert_eq!(records.len(), 2); + assert_eq!(records[0]["arrival_seq"], json!(0)); + assert!( + records[0]["arrival_unix_ms"] + .as_u64() + .is_some_and(|value| value > 0) + ); + assert_eq!(records[0]["source"], json!("orchestrator")); + assert_eq!(records[0]["stream"], json!("test-node#42")); + assert_eq!(records[0]["channel"], json!("stdout")); + assert_eq!(records[0]["channel_id"], json!(1)); + assert_eq!(records[0]["position"], json!(7)); + assert_eq!( + records[0]["payload"], + json!({"encoding":"utf8","value":"hello λ"}) + ); + + assert_eq!(records[1]["arrival_seq"], json!(1)); + assert!( + records[1]["arrival_unix_ms"] + .as_u64() + .is_some_and(|value| value > 0) + ); + assert_eq!(records[1]["source"], json!("orchestrator")); + assert_eq!(records[1]["stream"], json!("test-node#42")); + assert_eq!(records[1]["channel"], json!("stderr")); + assert_eq!(records[1]["channel_id"], json!(2)); + assert_eq!(records[1]["position"], json!(8)); + assert_eq!( + records[1]["payload"], + json!({"encoding":"bytes","value":[255,0,65]}) + ); + } + + #[test] + fn orchestrator_stdio_drain_archives_stdout_and_stderr_as_provision_logs() { + static NEXT_TEMP_FILE: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0); + + let suffix = NEXT_TEMP_FILE.fetch_add(1, std::sync::atomic::Ordering::Relaxed); + let path = std::env::temp_dir().join(format!( + "mvp-orch-stdio-test-{}-{suffix}.jsonl", + std::process::id() + )); + let _ = std::fs::remove_file(&path); + let (tx, rx) = mpsc::channel(); + tx.send(OrchStdioLine { + stream: ProvisionLogStream::Stdout, + line: "offer pool selected".to_owned(), + }) + .expect("send stdout line"); + tx.send(OrchStdioLine { + stream: ProvisionLogStream::Stderr, + line: "lease chain detail".to_owned(), + }) + .expect("send stderr line"); + + let mut datastream = OrchDatastream::new(77, Some(&path)).expect("datastream opens"); + drain_orch_stdio_capture(Some(&rx), &mut datastream, None, 77, 9); + drop(datastream); + + let contents = std::fs::read_to_string(&path).expect("read frame archive jsonl"); + let records = contents + .lines() + .map(|line| serde_json::from_str::(line).expect("archive line is json")) + .collect::>(); + let _ = std::fs::remove_file(&path); + + let stdout_record = records + .iter() + .rev() + .find(|record| record["channel"] == "mvp.provisioning.logs.node.9.stdout") + .expect("stdout provisioning log frame archived"); + let stderr_record = records + .iter() + .rev() + .find(|record| record["channel"] == "mvp.provisioning.logs.node.9.stderr") + .expect("stderr provisioning log frame archived"); + assert_eq!(stdout_record["source"], "orchestrator"); + assert_eq!(stderr_record["source"], "orchestrator"); + + let stdout_payload = serde_json::from_str::( + stdout_record["payload"]["value"] + .as_str() + .expect("stdout payload is archived as text"), + ) + .expect("stdout payload is log record json"); + let stderr_payload = serde_json::from_str::( + stderr_record["payload"]["value"] + .as_str() + .expect("stderr payload is archived as text"), + ) + .expect("stderr payload is log record json"); + assert_eq!(stdout_payload["line"]["run_id"], 77); + assert_eq!(stdout_payload["line"]["node_id"], 9); + assert_eq!(stdout_payload["line"]["stream"], "Stdout"); + assert_eq!(stdout_payload["line"]["line"], "offer pool selected"); + assert_eq!(stderr_payload["line"]["stream"], "Stderr"); + assert_eq!(stderr_payload["line"]["line"], "lease chain detail"); + } + + #[test] + fn config_layers_defaults_toml_env_then_cli() { + let toml = TempTomlFile::new( + "layering.toml", + r#" +[runtime] +profile = "deploy" +run_id = 41 +node_id = 9 +stage_index = 3 +layer_end_exclusive = 24 + +[provider] +kind = "vastai" + +[image] +node = "docker.io/example/from-image-node:toml" + +[vastai] +image = "docker.io/example/from-vastai-image:toml" +api_key = "toml-key" +bootstrap_command = "/toml/bootstrap" +disk_gb = 60 +gpu_name = "RTX 4090" + +[prompt] +rpc_addr = "127.0.0.1:19999" +max_tokens = 17 +dashboard = true + +[model] +id = "toml-model" +gguf_repo = "toml/repo" +gguf_file = "toml.gguf" +gguf_revision = "toml-rev" +max_context = 384 + +[relay] +mode = "disabled" +"#, + ); + + let config = with_clean_env( + &[ + ("MVP_NODE_PROVIDER", "docker"), + ("MVP_NODE_IMAGE", "docker.io/example/from-env:latest"), + ("MVP_PROMPT_MAX_TOKENS", "23"), + ("MVP_MODEL_ID", "env-model"), + ("MVP_IROH_RELAY_MODE", "default"), + ], + || { + Config::from_layers_with_path_and_args( + Some(&toml.path), + [ + "--image", + "docker.io/example/from-cli:latest", + "--max-tokens", + "31", + "--model-id", + "cli-model", + "--max-context", + "768", + ] + .into_iter() + .map(str::to_owned), + ) + .expect("layered config parses") + }, + ); + + assert_eq!(config.provider, ProviderKind::Docker); + assert_eq!(config.image, "docker.io/example/from-cli:latest"); + assert_eq!(config.default_max_tokens, 31); + assert_eq!(config.model_id, "cli-model"); + assert_eq!(config.run_id, 41); + assert_eq!(config.node_id, 9); + assert_eq!(config.stage_index, 3); + assert_eq!(config.layer_end_exclusive, Some(24)); + assert!(config.dashboard); + assert_eq!(config.max_context, Some(768)); + assert!(matches!(config.relay.mode, iroh::RelayMode::Default)); + assert!(config.vastai.is_none()); + } + + #[test] + fn pipeline_stages_layers_toml_env_then_cli_aliases() { + let toml = TempTomlFile::new( + "pipeline-stages.toml", + r#" +[runtime] +pipeline_stages = 2 + +[provider] +kind = "docker" +"#, + ); + + for (case, settings, args, expected) in [ + ("toml", vec![], vec![], 2), + ( + "env-over-toml", + vec![("MVP_PIPELINE_STAGES", "3")], + vec![], + 3, + ), + ( + "long-cli-over-env", + vec![("MVP_PIPELINE_STAGES", "3")], + vec!["--pipeline-stages", "4"], + 4, + ), + ( + "short-cli-over-env", + vec![("MVP_PIPELINE_STAGES", "3")], + vec!["-N", "5"], + 5, + ), + ] { + let config = with_clean_env(&settings, || { + Config::from_layers_with_path_and_args( + Some(&toml.path), + args.iter().copied().map(str::to_owned), + ) + .unwrap_or_else(|error| { + panic!("{case} pipeline stage config should parse: {error}") + }) + }); + + assert_eq!(config.pipeline_stages, expected, "{case}"); + } + } + + #[test] + fn pipeline_stages_rejects_zero_missing_and_non_numeric_values() { + for (case, settings, args, expected) in [ + ( + "env-zero", + vec![("MVP_PIPELINE_STAGES", "0")], + vec![], + "--pipeline-stages must be greater than 0", + ), + ( + "env-non-numeric", + vec![("MVP_PIPELINE_STAGES", "many")], + vec![], + "invalid MVP_PIPELINE_STAGES=\"many\"", + ), + ( + "long-cli-zero", + vec![], + vec!["--pipeline-stages", "0"], + "--pipeline-stages must be greater than 0", + ), + ( + "long-cli-non-numeric", + vec![], + vec!["--pipeline-stages", "many"], + "invalid --pipeline-stages=\"many\"", + ), + ( + "short-cli-missing", + vec![], + vec!["-N"], + "missing value after -N", + ), + ] { + let error = + with_clean_env(&settings, || { + match Config::from_layers_with_path_and_args( + None, + args.iter().copied().map(str::to_owned), + ) { + Ok(_) => panic!("invalid pipeline stages setting must fail"), + Err(error) => error, + } + }); + + assert!( + error.contains(expected), + "{case} error {error:?} should contain {expected:?}" + ); + } + } + + #[test] + fn vastai_rejects_pipeline_stages_count_above_one() { + let error = with_clean_env(&[], || { + match Config::from_layers_with_path_and_args( + None, + ["--provider", "vastai", "--pipeline-stages", "2"] + .into_iter() + .map(str::to_owned), + ) { + Ok(_) => panic!("VastAI cannot provision more than one pipeline stage yet"), + Err(error) => error, + } + }); + + assert!( + error.contains("provider=vastai does not support pipelined provisioning yet"), + "unexpected error: {error}" + ); + } + + #[test] + fn pipeline_cached_model_resolution_uses_toml_smollm2_gguf_file_with_cli_stage_count() { + let toml = TempTomlFile::new( + "pipeline-smollm2-model.toml", + r#" +[model] +gguf_repo = "QuantFactory/SmolLM2-135M-Instruct-GGUF" +gguf_file = "SmolLM2-135M-Instruct.Q4_0.gguf" +"#, + ); + + let config = with_clean_env(&[], || { + Config::from_layers_with_path_and_args( + Some(&toml.path), + ["--pipeline-stages", "3"].into_iter().map(str::to_owned), + ) + .expect("TOML SmolLM2 pipeline config parses") + }); + let expected_cached_host_path = default_pipeline_cached_model_path() + .canonicalize() + .expect("default cached SmolLM2 GGUF is present for cache resolution tests"); + let cached_model = config + .cached_model + .as_ref() + .expect("matching TOML SmolLM2 pipeline source should resolve the default cache"); + + assert_eq!(config.pipeline_stages, 3); + assert_eq!(config.provider, ProviderKind::Process); + assert_eq!(cached_model.host_path, expected_cached_host_path); + assert_eq!( + cached_model.container_path, + "/models/cached/SmolLM2-135M-Instruct.Q4_0.gguf" + ); + assert_eq!( + config.gguf_source, + GgufSource::LocalPath(expected_cached_host_path.to_string_lossy().to_string()) + ); + } + + #[test] + fn pipeline_cached_model_resolution_does_not_use_default_cache_for_other_toml_gguf_file() { + let toml = TempTomlFile::new( + "pipeline-other-model.toml", + r#" +[model] +gguf_repo = "QuantFactory/SmolLM2-135M-Instruct-GGUF" +gguf_file = "SmolLM2-135M-Instruct.Q8_0.gguf" +"#, + ); + + let config = with_clean_env(&[], || { + Config::from_layers_with_path_and_args( + Some(&toml.path), + ["--pipeline-stages", "3"].into_iter().map(str::to_owned), + ) + .expect("non-default TOML GGUF pipeline config parses") + }); + + assert_eq!(config.pipeline_stages, 3); + assert!(config.cached_model.is_none()); + assert_eq!( + config.gguf_source, + GgufSource::HuggingFaceGguf { + repo: "QuantFactory/SmolLM2-135M-Instruct-GGUF".to_owned(), + file: "SmolLM2-135M-Instruct.Q8_0.gguf".to_owned(), + revision: None, + } + ); + + let error = config + .build_run_plan() + .expect_err("non-default remote TOML GGUF should still require an explicit cache"); + + assert!( + error.contains( + "QuantFactory/SmolLM2-135M-Instruct-GGUF/SmolLM2-135M-Instruct.Q8_0.gguf" + ), + "unexpected error: {error}" + ); + assert!( + error.contains("--cached-model-host-path"), + "unexpected error: {error}" + ); + } + + #[test] + fn cached_model_plan_phase_is_metadata_driven_for_each_pipeline_width() { + let metadata = TestGgufMetadata { + num_layers: 7, + hidden_dim: 13, + context_length: 64, + eos_token_id: 11, + }; + let model = TempModelFile::with_metadata("metadata-driven.gguf", metadata); + + for stage_count in [1_u32, 3, metadata.num_layers] { + let config = pipeline_config_with_cached_model(&model, stage_count); + assert!(config.uses_planned_execution()); + let plan = config + .build_run_plan() + .expect("metadata-derived cached-model plan builds"); + assert_plan_matches_metadata(&config, &plan, metadata, stage_count, 512); + } + } + + #[test] + fn cached_model_implicit_single_stage_matches_explicit_n_one_plan() { + let model = TempModelFile::new("implicit-n-one.gguf"); + let implicit = cached_model_config_with_args(&model, &["--max-context", "32"]); + let explicit = cached_model_config_with_args( + &model, + &["--pipeline-stages", "1", "--max-context", "32"], + ); + + let implicit_plan = implicit + .build_run_plan() + .expect("implicit cached-model N=1 plan builds"); + let explicit_plan = explicit + .build_run_plan() + .expect("explicit cached-model N=1 plan builds"); + + assert!(implicit.uses_planned_execution()); + assert!(explicit.uses_planned_execution()); + assert_plan_matches_metadata(&implicit, &implicit_plan, model.metadata, 1, 32); + assert_plan_matches_metadata(&explicit, &explicit_plan, model.metadata, 1, 32); + assert_eq!(implicit_plan.model, explicit_plan.model); + assert_eq!(implicit_plan.edges, explicit_plan.edges); + assert_eq!(implicit_plan.stages, explicit_plan.stages); + } + + #[test] + fn cached_model_plan_rejects_pipeline_width_above_model_layers() { + let metadata = TestGgufMetadata { + num_layers: 2, + ..TestGgufMetadata::default() + }; + let model = TempModelFile::with_metadata("too-many-stages.gguf", metadata); + let config = pipeline_config_with_cached_model(&model, metadata.num_layers + 1); + + let error = config + .build_run_plan() + .expect_err("pipeline width above layer count must reject before provisioning"); + + assert!( + error.contains("--pipeline-stages=3 exceeds GGUF layer count 2"), + "unexpected error: {error}" + ); + } + + #[test] + fn pipeline_planning_rejects_remote_gguf_source_without_local_cache() { + let config = with_clean_env(&[], || { + Config::from_layers_with_path_and_args( + None, + [ + "--pipeline-stages", + "3", + "--gguf-repo", + "example/remote", + "--gguf-file", + "remote.gguf", + ] + .into_iter() + .map(str::to_owned), + ) + .expect("non-default remote pipeline config parses") + }); + + let error = config + .build_run_plan() + .expect_err("remote GGUF source cannot be inspected before provisioning"); + + assert!( + error.contains("locally inspectable GGUF before provisioning"), + "unexpected error: {error}" + ); + assert!( + error.contains("example/remote/remote.gguf"), + "unexpected error: {error}" + ); + assert!( + error.contains("--cached-model-host-path"), + "unexpected error: {error}" + ); + } + + #[test] + fn toml_vastai_image_overrides_image_node_for_vastai_provider() { + let toml = TempTomlFile::new( + "vastai-image.toml", + r#" +[runtime] +profile = "deploy" + +[provider] +kind = "vastai" + +[image] +node = "docker.io/example/generic:toml" + +[vastai] +image = "docker.io/example/vastai:toml" +api_key = "k" +bootstrap_command = "/run" +"#, + ); + + let config = with_clean_env(&[], || { + Config::from_layers_with_path_and_args(Some(&toml.path), std::iter::empty::()) + .expect("VastAI TOML config parses") + }); + + assert_eq!(config.provider, ProviderKind::VastAi); + assert_eq!(config.image, "docker.io/example/vastai:toml"); + } + + #[test] + fn missing_toml_uses_hardcoded_defaults() { + let missing_path = std::env::temp_dir().join(format!( + "mvp-orchestrator-missing-config-{}-{}.toml", + std::process::id(), + std::thread::current().name().unwrap_or("unnamed") + )); + let _ = std::fs::remove_file(&missing_path); + + let config = with_clean_env(&[], || { + Config::from_layers_with_path_and_args( + Some(&missing_path), + std::iter::empty::(), + ) + .expect("missing optional TOML config uses defaults") + }); + + assert_eq!(config.provider, ProviderKind::Process); + assert_eq!(config.image, DEFAULT_IMAGE); + assert_eq!(config.model_id, DEFAULT_MODEL_ID); + assert_eq!(config.default_max_tokens, DEFAULT_MAX_TOKENS); + assert!(!config.dashboard); + assert_eq!(config.max_context, None); + } + + #[test] + fn node_spec_propagates_max_context_when_configured() { + let config = with_clean_env(&[], || { + Config::from_layers_with_path_and_args( + None, + ["--max-context", "256"].into_iter().map(str::to_owned), + ) + .expect("CLI max context config parses") + }); + let coordinator = EndpointAddr::new(iroh::SecretKey::from_bytes(&[3; 32]).public()); + let orchestrator_actor = ActorAddress([18; 32]); + let spec = config + .node_spec(coordinator, orchestrator_actor) + .expect("node spec builds"); + + assert_eq!(env_value(&spec.env, "MVP_MAX_CONTEXT"), Some("256")); + } + + #[test] + fn runtime_profile_selects_provider_and_node_provider_takes_precedence() { + assert_eq!( + selected_provider(&[("MVP_RUNTIME_CONFIG", "local")]), + ProviderKind::Process + ); + assert_eq!( + selected_provider(&[("MVP_RUNTIME_CONFIG", "deploy")]), + ProviderKind::VastAi + ); + assert_eq!( + selected_provider(&[ + ("MVP_RUNTIME_CONFIG", "deploy"), + ("MVP_NODE_PROVIDER", "docker"), + ]), + ProviderKind::Docker + ); + } + + #[test] + fn relay_mode_env_uses_default_relay_and_accepts_disabled() { + with_clean_env(&[], || { + assert!(matches!( + relay_mode_from_env().expect("unset relay mode parses"), + iroh::RelayMode::Default + )); + }); + with_clean_env(&[("MVP_IROH_RELAY_MODE", "disabled")], || { + assert!(matches!( + relay_mode_from_env().expect("disabled relay mode parses"), + iroh::RelayMode::Disabled + )); + }); + } + + #[test] + fn node_spec_env_propagates_relay_url_only_for_custom_relay_config() { + const RELAY_URL: &str = "https://relay-node-spec.example.com"; + + let custom_env = node_spec_env(&[(MVP_IROH_RELAY_URL_ENV, RELAY_URL)]); + let expected_url = canonical_relay_url(RELAY_URL); + assert_eq!( + env_value(&custom_env, MVP_IROH_RELAY_URL_ENV), + Some(expected_url.as_str()) + ); + + let disabled_env = node_spec_env(&[ + ("MVP_IROH_RELAY_MODE", "disabled"), + (MVP_IROH_RELAY_URL_ENV, RELAY_URL), + ]); + assert_eq!(env_value(&disabled_env, MVP_IROH_RELAY_URL_ENV), None); + } + + #[test] + fn node_spec_preserves_relay_transport_in_coordinator_endpoint() { + with_clean_env(&[], || { + let config = Config::from_layers_with_path_and_args(None, std::iter::empty::()) + .expect("config parses"); + let secret = iroh::SecretKey::from_bytes(&[10; 32]); + let coordinator = EndpointAddr::new(secret.public()).with_relay_url( + "http://relay.example.com" + .parse::() + .unwrap(), + ); + let orchestrator_actor = ActorAddress([20; 32]); + + let spec = config + .node_spec(coordinator, orchestrator_actor) + .expect("node spec builds"); + let coordinator_endpoint_json = env_value(&spec.env, "MVP_COORDINATOR_ENDPOINT") + .expect("coordinator endpoint env is present"); + let coordinator_endpoint = + serde_json::from_str::(coordinator_endpoint_json) + .expect("coordinator endpoint env deserializes"); + + assert_eq!( + coordinator_endpoint + .relay_urls() + .next() + .map(|url| url.to_string()), + Some("http://relay.example.com/".to_owned()) + ); + }); + } + + #[test] + fn docker_config_construction_ignores_malformed_vastai_environment() { + let config = with_clean_env( + &[ + ("MVP_RUNTIME_CONFIG", "local"), + ("MVP_NODE_PROVIDER", "docker"), + ("MVP_VASTAI_CONFIRM_LEASE", "definitely-not-a-bool"), + ("MVP_VASTAI_DISK_GB", "not-a-u32"), + ("MVP_VASTAI_MIN_DOWN_MBPS", "not-a-float"), + ], + || { + Config::from_layers_with_path_and_args(None, std::iter::empty::()) + .expect("docker config ignores VastAI-only env") + }, + ); + + assert_eq!(config.provider, ProviderKind::Docker); + assert!(config.vastai.is_none()); + } + + #[test] + fn docker_cached_model_builds_local_gguf_env_and_planned_file_mount_from_canonical_host_path() { + let model = TempModelFile::new("weights-q4.gguf"); + let config = with_clean_env_os( + &[ + ("MVP_RUNTIME_CONFIG", OsString::from("local")), + ("MVP_NODE_PROVIDER", OsString::from("docker")), + (CACHED_MODEL_HOST_ENV, model.raw_path.as_os_str().to_owned()), + ], + || { + Config::from_layers_with_path_and_args(None, std::iter::empty::()) + .expect("docker cached model config parses") + }, + ); + let coordinator = EndpointAddr::new(iroh::SecretKey::from_bytes(&[7; 32]).public()); + let orchestrator_actor = ActorAddress([14; 32]); + let spec = config + .node_spec(coordinator, orchestrator_actor) + .expect("cached model node spec builds"); + + assert_eq!( + env_value(&spec.env, "MVP_GGUF_LOCAL_PATH"), + Some("/models/cached/weights-q4.gguf") + ); + assert_eq!(env_value(&spec.env, "MVP_GGUF_REPO"), None); + assert_eq!(env_value(&spec.env, "MVP_GGUF_FILE"), None); + assert_eq!( + spec.mounts, + vec![ProviderMount { + host_path: model.canonical_path.to_string_lossy().to_string(), + container_path: "/models/cached/weights-q4.gguf".to_owned(), + readonly: true, + }] + ); + } + + #[test] + fn process_cached_model_builds_host_gguf_env_without_mounts() { + let model = TempModelFile::new("process-weights-q4.gguf"); + let config = with_clean_env_os( + &[ + ("MVP_RUNTIME_CONFIG", OsString::from("local")), + (CACHED_MODEL_HOST_ENV, model.raw_path.as_os_str().to_owned()), + ], + || { + Config::from_layers_with_path_and_args(None, std::iter::empty::()) + .expect("process cached model config parses") + }, + ); + let coordinator = EndpointAddr::new(iroh::SecretKey::from_bytes(&[8; 32]).public()); + let orchestrator_actor = ActorAddress([15; 32]); + let spec = config + .node_spec(coordinator, orchestrator_actor) + .expect("process cached model node spec builds"); + let host_path = model.canonical_path.to_string_lossy().to_string(); + + assert_eq!(config.provider, ProviderKind::Process); + assert_eq!(env_value(&spec.env, "MVP_NODE_PROVIDER"), Some("process")); + assert_eq!( + env_value(&spec.env, "MVP_GGUF_LOCAL_PATH"), + Some(host_path.as_str()) + ); + assert_eq!(env_value(&spec.env, "MVP_DOCKER_GPUS"), None); + assert!(spec.mounts.is_empty(), "process workers use host paths"); + } + + #[test] + fn vectorized_local_docker_stage_node_specs_follow_three_stage_plan() { + let model = TempModelFile::new("pipeline-node-spec.gguf"); + let config = pipeline_config_with_cached_model(&model, 3); + let plan = config + .build_run_plan() + .expect("stage node spec plan builds"); + let coordinator = EndpointAddr::new(iroh::SecretKey::from_bytes(&[21; 32]).public()); + let orchestrator_actor = ActorAddress([23; 32]); + let coordinator_env = + serde_json::to_string(&coordinator).expect("coordinator endpoint serializes"); + let orchestrator_actor_env = + serde_json::to_string(&orchestrator_actor).expect("orchestrator actor serializes"); + let expected_run_id = config.run_id.to_string(); + let expected_mount = ProviderMount { + host_path: model.canonical_path.to_string_lossy().to_string(), + container_path: "/models/cached/pipeline-node-spec.gguf".to_owned(), + readonly: true, + }; + + let specs = stage_node_specs(&config, Some(&plan), coordinator, orchestrator_actor) + .expect("pipeline stage node specs build"); + + assert_eq!( + specs.len(), + 3, + "pipeline_stages=3 should provision three Docker workers, not one" + ); + assert!( + specs.iter().all(|spec| spec.node_id != config.node_id), + "coordinator node {} must not be counted as a worker: {specs:?}", + config.node_id + ); + + for (expected_stage_index, spec) in specs.iter().enumerate() { + let expected_stage_index = + u32::try_from(expected_stage_index).expect("fixture stage index fits u32"); + let expected_node_id = config.node_id + 1 + u64::from(expected_stage_index); + let expected_node_id_env = expected_node_id.to_string(); + let expected_stage_index_env = expected_stage_index.to_string(); + + assert_eq!(spec.run_id, config.run_id); + assert_eq!(spec.node_id, expected_node_id); + assert_eq!(spec.stage_index, Some(expected_stage_index)); + assert_eq!( + env_value(&spec.env, "MVP_RUN_ID"), + Some(expected_run_id.as_str()) + ); + assert_eq!( + env_value(&spec.env, "MVP_LOGICAL_NODE_ID"), + Some(expected_node_id_env.as_str()) + ); + assert_eq!( + env_value(&spec.env, "MVP_STAGE_INDEX"), + Some(expected_stage_index_env.as_str()) + ); + assert_eq!(env_value(&spec.env, "MVP_PIPELINE_STAGES"), Some("3")); + assert_eq!(env_value(&spec.env, "MVP_NODE_PROVIDER"), Some("docker")); + assert_eq!( + env_value(&spec.env, "MVP_MODEL_ID"), + Some(config.model_id.as_str()) + ); + assert_eq!( + env_value(&spec.env, "MVP_GGUF_LOCAL_PATH"), + Some("/models/cached/pipeline-node-spec.gguf") + ); + assert_eq!(env_value(&spec.env, "MVP_MAX_CONTEXT"), Some("512")); + assert_eq!( + env_value(&spec.env, "MVP_COORDINATOR_ENDPOINT"), + Some(coordinator_env.as_str()) + ); + assert_eq!( + env_value(&spec.env, "MVP_ORCHESTRATOR_ACTOR"), + Some(orchestrator_actor_env.as_str()) + ); + assert_eq!( + spec.mounts, + vec![expected_mount.clone()], + "pipeline cached GGUF mount should be read-only for stage {expected_stage_index}" + ); + } + } + + #[test] + fn vectorized_local_process_stage_node_specs_follow_three_stage_plan() { + let model = TempModelFile::new("process-pipeline-node-spec.gguf"); + let config = with_clean_env_os( + &[ + ("MVP_RUNTIME_CONFIG", OsString::from("local")), + (CACHED_MODEL_HOST_ENV, model.raw_path.as_os_str().to_owned()), + ], + || { + Config::from_layers_with_path_and_args( + None, + ["--pipeline-stages", "3", "--max-context", "512"] + .into_iter() + .map(str::to_owned), + ) + .expect("process pipeline config parses") + }, + ); + let plan = config + .build_run_plan() + .expect("process stage node spec plan builds"); + let coordinator = EndpointAddr::new(iroh::SecretKey::from_bytes(&[22; 32]).public()); + let orchestrator_actor = ActorAddress([24; 32]); + let host_path = model.canonical_path.to_string_lossy().to_string(); + + let specs = stage_node_specs(&config, Some(&plan), coordinator, orchestrator_actor) + .expect("process pipeline stage node specs build"); + + assert_eq!(config.provider, ProviderKind::Process); + assert_eq!(specs.len(), 3); + for (expected_stage_index, spec) in specs.iter().enumerate() { + let expected_stage_index = + u32::try_from(expected_stage_index).expect("fixture stage index fits u32"); + assert_eq!(spec.stage_index, Some(expected_stage_index)); + assert_eq!(env_value(&spec.env, "MVP_NODE_PROVIDER"), Some("process")); + assert_eq!(env_value(&spec.env, "MVP_PIPELINE_STAGES"), Some("3")); + assert_eq!( + env_value(&spec.env, "MVP_GGUF_LOCAL_PATH"), + Some(host_path.as_str()) + ); + assert_eq!(env_value(&spec.env, "MVP_MAX_CONTEXT"), Some("512")); + assert!(spec.mounts.is_empty(), "process stage specs must not mount"); + } + } + + #[test] + fn planned_single_stage_cached_model_node_spec_follows_generated_plan() { + let model = TempModelFile::new("single-stage-node-spec.gguf"); + let config = pipeline_config_with_cached_model(&model, 1); + let plan = config + .build_run_plan() + .expect("single-stage cached plan builds"); + let coordinator = EndpointAddr::new(iroh::SecretKey::from_bytes(&[31; 32]).public()); + let orchestrator_actor = ActorAddress([33; 32]); + + let specs = stage_node_specs(&config, Some(&plan), coordinator, orchestrator_actor) + .expect("single planned stage node spec builds"); + + assert_eq!(specs.len(), 1); + let spec = &specs[0]; + let stage = &plan.stages[0]; + assert_eq!(stage.stage_index, 0); + assert_eq!(stage.layer_start, 0); + assert_eq!(stage.layer_end_exclusive, model.metadata.num_layers); + assert_eq!(spec.node_id, stage.node_id.0); + assert_ne!( + spec.node_id, config.node_id, + "planned cached N=1 still provisions a worker separate from the coordinator" + ); + assert_eq!(spec.stage_index, Some(0)); + assert_eq!(env_value(&spec.env, "MVP_PIPELINE_STAGES"), Some("1")); + assert_eq!( + env_value(&spec.env, "MVP_GGUF_LOCAL_PATH"), + Some("/models/cached/single-stage-node-spec.gguf") + ); + assert_eq!( + spec.mounts, + vec![ProviderMount { + host_path: model.canonical_path.to_string_lossy().to_string(), + container_path: "/models/cached/single-stage-node-spec.gguf".to_owned(), + readonly: true, + }] + ); + } + + #[test] + fn vectorized_local_docker_stage_provision_detail_preserves_plan_edges_and_layers() { + let model = TempModelFile::new("pipeline-stage-detail.gguf"); + let config = pipeline_config_with_cached_model(&model, 3); + let plan = config.build_run_plan().expect("stage detail plan builds"); + + let detail = stage_provision_detail(&config, Some(&plan)); + + assert_eq!( + detail, + json!({ + "run_id": config.run_id, + "stage_count": 3, + "stages": plan.stages.iter().map(|stage| { + json!({ + "node_id": stage.node_id.0, + "stage_index": stage.stage_index, + "layer_range": { + "start": stage.layer_start, + "end_exclusive": stage.layer_end_exclusive + }, + "inbound_edge_id": stage.inbound_edge.0, + "outbound_edge_id": stage.outbound_edge.0, + }) + }).collect::>(), + "model_id": config.model_id, + }) + ); + } + + #[test] + fn stage_provision_wire_preserves_plan_edges_for_single_and_multi_stage_cached_models() { + for stage_count in [1_u32, 3] { + let model = TempModelFile::new(&format!("wire-{stage_count}.gguf")); + let config = pipeline_config_with_cached_model(&model, stage_count); + let plan = config.build_run_plan().expect("wire test plan builds"); + let coordinator = EndpointAddr::new(iroh::SecretKey::from_bytes(&[41; 32]).public()); + let readies = plan + .stages + .iter() + .map(|stage| { + let byte = u8::try_from(stage.stage_index + 42).expect("fixture byte fits"); + let endpoint = + EndpointAddr::new(iroh::SecretKey::from_bytes(&[byte; 32]).public()); + ( + stage.node_id.0, + RuntimeReady { + endpoint: endpoint.clone(), + node_actor: ActorAddress([byte; 32]), + datastream_publisher: ActorAddress([byte.wrapping_add(80); 32]), + stage_index: stage.stage_index, + readiness_id: u64::from(stage.stage_index) + 100, + swim_node_id: DistNodeId(*endpoint.id.as_bytes()), + }, + ) + }) + .collect::>(); + + for stage in &plan.stages { + let wire = stage_provision_wire_from_plan( + &plan, + stage.stage_index, + &readies, + &coordinator, + ) + .expect("stage provision wire builds from plan"); + let inbound = wire.inbound_edge.as_ref().expect("planned inbound edge"); + let outbound = wire.outbound_edge.as_ref().expect("planned outbound edge"); + + assert_eq!(wire.stage_count, stage_count); + assert_eq!(wire.node_id, stage.node_id.0); + assert_eq!(wire.layer_start, stage.layer_start); + assert_eq!(wire.layer_end_exclusive, stage.layer_end_exclusive); + assert_eq!(inbound.edge_id, stage.inbound_edge.0); + assert_eq!(outbound.edge_id, stage.outbound_edge.0); + assert_eq!( + inbound.kind, + if stage.stage_index == 0 { + StageEdgeKindWire::TokenIn + } else { + StageEdgeKindWire::Activation + } + ); + assert_eq!( + outbound.kind, + if stage.stage_index + 1 == stage_count { + StageEdgeKindWire::TokenOut + } else { + StageEdgeKindWire::Activation + } + ); + assert_eq!(wire.model_id, config.model_id); + assert_eq!(wire.gguf_source, config.gguf_source); + assert_eq!(wire.tokenizer, config.tokenizer); + } + } + } + + #[test] + fn cached_model_with_deploy_provider_is_rejected_before_vastai_env_is_parsed() { + let model = TempModelFile::new("deploy-rejected.gguf"); + let error = with_clean_env_os( + &[ + ("MVP_RUNTIME_CONFIG", OsString::from("deploy")), + (CACHED_MODEL_HOST_ENV, model.raw_path.as_os_str().to_owned()), + ( + "MVP_VASTAI_CONFIRM_LEASE", + OsString::from("definitely-not-a-bool"), + ), + ("MVP_VASTAI_DISK_GB", OsString::from("not-a-u32")), + ], + || match Config::from_layers_with_path_and_args(None, std::iter::empty::()) { + Ok(_) => panic!("deploy cached model must be rejected"), + Err(error) => error, + }, + ); + + assert!( + error.contains( + "MVP_CACHED_MODEL_HOST_PATH is a host-local cache path and is only supported by provider=process or provider=docker" + ), + "unexpected error: {error}" + ); + assert!( + !error.contains("MVP_VASTAI_CONFIRM_LEASE") && !error.contains("MVP_VASTAI_DISK_GB"), + "cached-model rejection should not require valid VastAI env, got: {error}" + ); + } + + #[derive(Default)] + struct FakeProvisionPlugin { + stopped: Vec, + } + + impl ProvisionPlugin for FakeProvisionPlugin { + fn start_node( + &mut self, + _spec: NodeProvisionSpec, + _sink: PluginSink, + ) -> Result { + unreachable!("guard tests construct handles directly") + } + + fn complete_bootstrap( + &mut self, + _handle: &crate::provisioning::PluginNodeHandle, + ) -> Result<(), String> { + Ok(()) + } + + fn stop_node( + &mut self, + handle: &crate::provisioning::PluginNodeHandle, + ) -> Result<(), String> { + self.stopped.push(handle.id); + Ok(()) + } + } + + #[test] + fn provisioned_node_guard_stops_node_on_drop() { + let mut plugin = FakeProvisionPlugin::default(); + { + let _guard = ProvisionedNodeGuard::new( + &mut plugin, + crate::provisioning::PluginNodeHandle { + id: 7, + provider_process_id: Some(99), + }, + ); + } + assert_eq!(plugin.stopped, vec![7]); + } + + #[test] + fn provisioned_node_guard_explicit_stop_runs_once() { + let mut plugin = FakeProvisionPlugin::default(); + { + let mut guard = ProvisionedNodeGuard::new( + &mut plugin, + crate::provisioning::PluginNodeHandle { + id: 8, + provider_process_id: None, + }, + ); + guard.stop().expect("first stop succeeds"); + guard.stop().expect("second stop is a no-op"); + } + assert_eq!(plugin.stopped, vec![8]); + } + + #[test] + fn stop_requested_observes_shutdown_signal_only() { + let (_tx, rx) = mpsc::channel(); + assert!(!stop_requested(&rx)); + + let (tx, rx) = mpsc::channel(); + tx.send(()).expect("send shutdown"); + assert!(stop_requested(&rx)); + assert!(!stop_requested(&rx)); + } + + fn endpoint(seed: u8) -> EndpointAddr { + EndpointAddr::new(iroh::SecretKey::from_bytes(&[seed; 32]).public()) + } + + fn insert_route(stack: &DistributionRuntimeStack, actor: ActorAddress, owner: DistNodeId) { + let mut route_view = match stack.route_view.write() { + Ok(route_view) => route_view, + Err(poisoned) => poisoned.into_inner(), + }; + route_view.insert(actor, owner); + } + + fn mark_alive(stack: &DistributionRuntimeStack, node_id: DistNodeId) { + stack + .runtime + .send_to( + stack.actors.membership_fanout, + distribution::swim::actor::MembershipChanged { + node_id, + state: MemberState::Alive, + incarnation: 1, + }, + ) + .expect("send membership change"); + stack.pump_runtime_once(); + } + + #[test] + fn runtime_ready_barrier_waits_for_specific_swim_and_route() { + let remote = DistNodeId([2; 32]); + let node_actor = ActorAddress::new_random(); + let datastream_publisher = ActorAddress::new_random(); + let ready = RuntimeReady { + endpoint: endpoint(2), + node_actor, + datastream_publisher, + stage_index: 3, + readiness_id: 99, + swim_node_id: remote, + }; + + let stack = + DistributionRuntimeStack::new(DistNodeId([1; 32]), DistributedNodeConfig::default()); + assert!(!runtime_ready_barrier_met(&stack, &ready)); + + let stack = + DistributionRuntimeStack::new(DistNodeId([1; 32]), DistributedNodeConfig::default()); + mark_alive(&stack, remote); + assert!(!runtime_ready_barrier_met(&stack, &ready)); + + let stack = + DistributionRuntimeStack::new(DistNodeId([1; 32]), DistributedNodeConfig::default()); + insert_route(&stack, node_actor, remote); + assert!(!runtime_ready_barrier_met(&stack, &ready)); + + let stack = + DistributionRuntimeStack::new(DistNodeId([1; 32]), DistributedNodeConfig::default()); + mark_alive(&stack, remote); + insert_route(&stack, node_actor, remote); + assert!(runtime_ready_barrier_met(&stack, &ready)); + } + + #[test] + fn enqueue_runtime_ready_ack_reports_to_node_agent() { + use crate::actors::node_agent::{NodeAgentActor, NodeAgentReport}; + use crate::actors::orchestrator::OrchestratorMsg; + use crate::stage_controller as stage; + + let stack = + DistributionRuntimeStack::new(DistNodeId([1; 32]), DistributedNodeConfig::default()); + let orchestrator_inbox = stack + .runtime + .new_inbox::() + .expect("orchestrator inbox"); + let reports = stack + .runtime + .new_inbox::() + .expect("node report inbox"); + let node_actor = stack + .runtime + .spawn(NodeAgentActor::new( + stage::NodeId(11), + *orchestrator_inbox.addr(), + Some(*reports.addr()), + )) + .expect("spawn node agent"); + let ready = RuntimeReady { + endpoint: endpoint(9), + node_actor, + datastream_publisher: ActorAddress::new_random(), + stage_index: 3, + readiness_id: 99, + swim_node_id: DistNodeId([2; 32]), + }; + + enqueue_runtime_ready_ack(&stack, &ready, 7, 11).expect("enqueue runtime ready ack"); + stack.pump_runtime_once(); + + assert_eq!( + reports.try_recv(), + Some(NodeAgentReport::RuntimeReadyAck { + run_id: 7, + node_id: 11, + stage_index: ready.stage_index, + readiness_id: 99, + }) + ); + + assert_eq!( + orchestrator_inbox.try_recv(), + Some(OrchestratorMsg::ObserveNodeRuntimeReadyAck { + run_id: 7, + node_id: 11, + stage_index: ready.stage_index, + readiness_id: 99, + }) + ); + } +} diff --git a/crates/mvp-system/src/tests/gpu_worker_ingress_parser_guarantees.rs b/crates/mvp-system/src/tests/gpu_worker_ingress_parser_guarantees.rs index 98510a5..80acd78 100644 --- a/crates/mvp-system/src/tests/gpu_worker_ingress_parser_guarantees.rs +++ b/crates/mvp-system/src/tests/gpu_worker_ingress_parser_guarantees.rs @@ -52,6 +52,7 @@ fn flagged_record(sequence: u64, extent: u64) -> Vec { .extent(extent) .flags(ingress::ObjectFlags { end_of_sequence: true, + begin_sequence: true, }) .payload(vec![7; extent as usize]) .encode() @@ -167,7 +168,8 @@ fn parser_decodes_and_propagates_object_header_flags() { assert_eq!( parsed.flags, ingress::ObjectFlags { - end_of_sequence: true + end_of_sequence: true, + begin_sequence: true, } ); } diff --git a/xtask/src/main.rs b/xtask/src/main.rs index dc5469c..de9cf34 100644 --- a/xtask/src/main.rs +++ b/xtask/src/main.rs @@ -1,4 +1,4 @@ -use std::collections::BTreeMap; +use std::collections::{BTreeMap, BTreeSet}; use std::fs; use std::io::{Read, Write}; use std::path::{Path, PathBuf}; @@ -37,6 +37,78 @@ struct MvpChatCheckOutput { stdin_error: Option, } +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum MvpChatCheckScenario { + ProcessBaseline, + Gpu, + Multinode, + MultinodeDocker, +} + +impl MvpChatCheckScenario { + fn parse_args(args: Vec) -> Result { + let mut scenario = Self::ProcessBaseline; + for arg in args { + let selected = match arg.as_str() { + "--gpu" => Self::Gpu, + "--multinode" => Self::Multinode, + "--multinode-docker" => Self::MultinodeDocker, + other => return Err(format!("unsupported mvp-chat-check argument {other:?}")), + }; + if scenario != Self::ProcessBaseline { + return Err( + "mvp-chat-check accepts at most one scenario flag: --gpu, --multinode, or --multinode-docker" + .to_owned(), + ); + } + scenario = selected; + } + Ok(scenario) + } + + fn name(self) -> &'static str { + match self { + Self::ProcessBaseline => "process", + Self::Gpu => "gpu", + Self::Multinode => "multinode", + Self::MultinodeDocker => "multinode-docker", + } + } + + fn mvp_chat_args(self, run_id: u64, dump_log: &Path) -> Vec { + let mut args = Vec::new(); + match self { + Self::ProcessBaseline | Self::Multinode => { + args.push("--process".to_owned()); + } + Self::Gpu => { + args.extend(["--process".to_owned(), "--gpu".to_owned()]); + } + Self::MultinodeDocker => { + args.push("--docker".to_owned()); + } + } + if matches!(self, Self::Multinode | Self::MultinodeDocker) { + args.extend(["--pipeline-stages".to_owned(), "2".to_owned()]); + } + if !matches!(self, Self::Gpu) { + args.push("--cached-model".to_owned()); + } + args.extend([ + "--run-id".to_owned(), + run_id.to_string(), + format!("--dump-logs={}", dump_log.display()), + ]); + args + } + + fn env_overrides(self) -> &'static [(&'static str, &'static str)] { + match self { + Self::ProcessBaseline | Self::Gpu | Self::Multinode | Self::MultinodeDocker => &[], + } + } +} + const BASIC_TESTS: &[TestStep] = &[ TestStep { label: "root crate", @@ -92,13 +164,39 @@ fn print_usage() { USAGE: cargo xtask COMMANDS: - mvp-chat [--process|--docker|--vastai] [--pipeline-stages n] [--cached-model] [-- args...] Run the human chat wrapper against the real orchestrator/worker bins. - mvp-chat-check Run real cargo mvp-chat acceptance check. + mvp-chat [--gpu] [--process|--docker|--vastai] [--pipeline-stages n] [--cached-model] [-- args...] Run the human chat wrapper against the real orchestrator/worker bins. + mvp-chat-check [--gpu|--multinode|--multinode-docker] + Run real cargo mvp-chat acceptance check for one explicit scenario. test Run the basic non-binding test barrier: root crate plus each non-binding repository package with `cargo test -p`." ); } +const MVP_CHAT_USAGE: &str = "\ +USAGE: cargo mvp-chat [OPTIONS] + +OPTIONS: + --gpu Run the local GPU path: in-process orchestrator plus DEV=CUDA worker selection + --process | --docker | --vastai + Select the runtime provider + --config Load config overlay + --pipeline-stages Number of pipeline stages + --cached-model[=] Use discovered or explicit cached GGUF model + --dump-logs[=] Write datastream frame log + --run-id Override run id + --skip-rebuild Reuse existing Cargo artifacts + --yes, -y Approve Vast.ai lease prompts + --help, -h Print this help"; + +fn print_mvp_chat_usage() { + println!("{MVP_CHAT_USAGE}"); +} + +fn is_mvp_chat_help_request(args: &[String]) -> bool { + args.iter() + .any(|arg| matches!(arg.as_str(), "--help" | "-h" | "help")) +} + fn run_step(step: &TestStep) -> bool { println!("\n=== {} ===", step.label); println!(" cargo {}", step.args.join(" ")); @@ -115,7 +213,7 @@ fn run_step(step: &TestStep) -> bool { fn run_tests() -> ExitCode { let start = Instant::now(); - let check = run_mvp_chat_check(); + let check = run_mvp_chat_check(Vec::new()); if check != ExitCode::SUCCESS { return check; } @@ -139,13 +237,17 @@ fn run_tests() -> ExitCode { } fn run_mvp_chat(args: Vec) -> ExitCode { - let mut command = Command::new(cargo_bin()); - command.args(["run", "--package", "mvp-system", "--bin", "mvp-chat", "--"]); let forwarded = if args.first().is_some_and(|arg| arg == "--") { args[1..].to_vec() } else { args }; + if is_mvp_chat_help_request(&forwarded) { + print_mvp_chat_usage(); + return ExitCode::SUCCESS; + } + let mut command = Command::new(cargo_bin()); + command.args(["run", "--package", "mvp-system", "--bin", "mvp-chat", "--"]); let dump_log_path = explicit_dump_log_path_from_mvp_chat_args(&forwarded); let run_id = run_id_from_mvp_chat_args(&forwarded); let benchmark_target = dump_log_path.as_deref().zip(run_id); @@ -374,7 +476,15 @@ fn write_mvp_chat_check_paths(root: &Path) -> Result }) } -fn run_mvp_chat_check() -> ExitCode { +fn run_mvp_chat_check(args: Vec) -> ExitCode { + let scenario = match MvpChatCheckScenario::parse_args(args) { + Ok(scenario) => scenario, + Err(error) => { + eprintln!("mvp-chat-check: failed: {error}"); + print_usage(); + return ExitCode::from(1); + } + }; let workspace = workspace_root(); let temp_root = unique_temp_dir("mvp-chat-check"); let paths = match write_mvp_chat_check_paths(&temp_root) { @@ -389,8 +499,9 @@ fn run_mvp_chat_check() -> ExitCode { } }; let run_id = mvp_chat_check_run_id(); + println!("mvp-chat-check: scenario {}", scenario.name()); - let output = match run_mvp_chat_check_process(&workspace, &paths, run_id) { + let output = match run_mvp_chat_check_process(&workspace, &paths, run_id, scenario) { Ok(output) => output, Err(error) => return fail_mvp_chat_check(&error, &paths, "", "", None), }; @@ -436,7 +547,7 @@ fn run_mvp_chat_check() -> ExitCode { ); } }; - let events = match assert_dump_log_facts(&paths.dump_log) { + let events = match assert_dump_log_facts(&paths.dump_log, scenario) { Ok(events) => events, Err(error) => { return fail_mvp_chat_check( @@ -448,7 +559,7 @@ fn run_mvp_chat_check() -> ExitCode { ); } }; - let report = match build_benchmark_report(&events, output.child_elapsed_ms, run_id) { + let report = match build_benchmark_report(&events, output.child_elapsed_ms, run_id, scenario) { Ok(report) => report, Err(error) => { return fail_mvp_chat_check( @@ -483,13 +594,17 @@ fn run_mvp_chat_check_process( workspace: &Path, paths: &MvpChatCheckPaths, run_id: u64, + scenario: MvpChatCheckScenario, ) -> Result { let mut command = Command::new(cargo_bin()); + command.current_dir(workspace).arg("mvp-chat").arg("--"); + for arg in scenario.mvp_chat_args(run_id, &paths.dump_log) { + command.arg(arg); + } + for &(key, value) in scenario.env_overrides() { + command.env(key, value); + } command - .current_dir(workspace) - .args(["mvp-chat", "--", "--cached-model", "--run-id"]) - .arg(run_id.to_string()) - .arg(format!("--dump-logs={}", paths.dump_log.display())) .stdin(Stdio::piped()) .stdout(Stdio::piped()) .stderr(Stdio::piped()); @@ -1086,7 +1201,10 @@ fn dump_log_inner_payload_text<'a>( }) } -fn assert_dump_log_facts(path: &Path) -> Result, String> { +fn assert_dump_log_facts( + path: &Path, + scenario: MvpChatCheckScenario, +) -> Result, String> { let events = parse_dump_log_events(path)?; let mut facts = DumpLogFacts::default(); for record in &events { @@ -1116,6 +1234,9 @@ fn assert_dump_log_facts(path: &Path) -> Result, String> { require_dump_log_fact(facts.request_completed_2, "request_completed request_id=2")?; require_dump_log_fact(facts.shutdown_requested, "shutdown requested")?; require_dump_log_fact(facts.orchestrator_stopped, "orchestrator_process stopped")?; + if scenario == MvpChatCheckScenario::Gpu { + require_gpu_dump_log_facts(&facts)?; + } Ok(events) } @@ -1123,6 +1244,7 @@ fn build_benchmark_report( events: &[DumpLogEvent], child_elapsed_ms: u64, run_id: u64, + scenario: MvpChatCheckScenario, ) -> Result { let facts = BenchmarkFacts::from_events(events, run_id); facts.require_span( @@ -1149,12 +1271,21 @@ fn build_benchmark_report( "prepare_runtime", "ready", )?; - facts.require_span( - "mvp.chat.runtime", - "ChatProgress", - "ensure_orch_binary", - "ready", - )?; + if scenario == MvpChatCheckScenario::Gpu { + facts.require_span( + "mvp.chat.runtime", + "ChatProgress", + "ensure_orchestrator_actor", + "ready", + )?; + } else { + facts.require_span( + "mvp.chat.runtime", + "ChatProgress", + "ensure_orch_binary", + "ready", + )?; + } facts.require_span( "mvp.chat.runtime", "ChatProgress", @@ -1438,6 +1569,22 @@ fn missing_benchmark_event(event: String) -> String { #[derive(Default)] struct DumpLogFacts { + gpu_worker_device_requested: bool, + gpu_import_ready: bool, + gpu_probe_ready: bool, + gpu_worker_ready: bool, + gpu_cpu_fallback_seen: bool, + gpu_decode_started: BTreeSet, + gpu_first_token_ready: BTreeSet, + gpu_decode_ready: BTreeSet, + gpu_prompt_completed: BTreeSet, + gpu_pipeline_prompt_encoded: BTreeSet, + gpu_pipeline_prompt_begin: BTreeSet, + gpu_pipeline_token_in: BTreeSet, + gpu_pipeline_token_out: BTreeSet, + gpu_pipeline_tokenizer_decode_ready: BTreeSet, + gpu_pipeline_tokens_decoded: BTreeSet, + gpu_pipeline_real_worker_step_seen: bool, chat_config_ready: bool, prepare_runtime_ready: bool, prompt_rpc_ready: bool, @@ -1458,6 +1605,7 @@ fn record_dump_log_event( event: &Value, facts: &mut DumpLogFacts, ) -> Result<(), String> { + record_gpu_dump_log_event(channel, event, facts); let event_type = event.get("type").and_then(Value::as_str); let phase = event.get("phase").and_then(Value::as_str); let status = event.get("status").and_then(Value::as_str); @@ -1522,12 +1670,245 @@ fn record_dump_log_event( Ok(()) } +fn record_gpu_dump_log_event(channel: &str, event: &Value, facts: &mut DumpLogFacts) { + let event_type = event.get("type").and_then(Value::as_str); + let phase = event.get("phase").and_then(Value::as_str); + let status = event.get("status").and_then(Value::as_str); + if event_type == Some("TinygradCpuCompilerSelected") { + facts.gpu_cpu_fallback_seen = true; + } + + match (channel, event_type) { + ("mvp.worker.initialize", Some("TinygradImportStarted")) + if event_requested_device_is_cuda(event) => + { + facts.gpu_worker_device_requested = true; + } + ("mvp.worker.initialize", Some("TinygradImportReady")) + if event_requested_device_is_cuda(event) && event_env_dev_is_cuda(event) => + { + facts.gpu_import_ready = true; + } + ("mvp.worker.initialize", Some("TinygradDeviceProbeReady")) + if event_requested_device_is_cuda(event) && event_probe_result_is_one(event) => + { + facts.gpu_probe_ready = true; + } + ("mvp.worker.initialize", Some("WorkerReady")) if worker_ready_backend_is_cuda(event) => { + facts.gpu_worker_ready = true; + } + ("mvp.worker.prompt", Some("DecodeStarted")) + if event_positive_u64(event, "prompt_tokens") + && event_positive_u64(event, "max_tokens") + && event + .get("decode_impl") + .and_then(Value::as_str) + .is_some_and(|decode_impl| !decode_impl.is_empty()) => + { + if let Some(request_id) = event.get("request_id").and_then(Value::as_u64) { + facts.gpu_decode_started.insert(request_id); + } + } + ("mvp.worker.prompt", Some("FirstTokenReady")) + if event.get("token_index").and_then(Value::as_u64) == Some(1) => + { + if let Some(request_id) = event.get("request_id").and_then(Value::as_u64) { + facts.gpu_first_token_ready.insert(request_id); + } + } + ("mvp.worker.prompt", Some("DecodeReady")) + if event_positive_u64(event, "tokens_generated") => + { + if let Some(request_id) = event.get("request_id").and_then(Value::as_u64) { + facts.gpu_decode_ready.insert(request_id); + } + } + ("mvp.worker.prompt", Some("PromptCompleted")) + if event + .get("generated_tokens") + .and_then(Value::as_array) + .is_some_and(|tokens| !tokens.is_empty()) => + { + if let Some(request_id) = event.get("request_id").and_then(Value::as_u64) { + facts.gpu_prompt_completed.insert(request_id); + } + } + ("mvp.worker.tokenizer", Some("PromptEncoded")) + if event + .get("tokens") + .and_then(Value::as_array) + .is_some_and(|tokens| !tokens.is_empty()) => + { + if let Some(request_id) = event_request_id(event) { + facts.gpu_pipeline_prompt_encoded.insert(request_id); + } + } + ("mvp.worker.tokenizer", Some("TokensDecoded")) + if event + .get("text") + .and_then(Value::as_str) + .is_some_and(|text| !text.is_empty()) => + { + if let Some(request_id) = event_request_id(event) { + facts.gpu_pipeline_tokens_decoded.insert(request_id); + } + } + ("mvp.worker.step", Some("StepExecuted")) + if event_positive_u64(event, "committed_bytes") => + { + let backend = event.get("execution_backend").and_then(Value::as_str); + if matches!(backend, Some("pipeline_stage" | "full_transformer")) { + facts.gpu_pipeline_real_worker_step_seen = true; + } + } + ("mvp.orch.prompt", Some("OrchPromptEvent")) + if phase == Some("pipeline_token_in") + && status == Some("ready") + && event_request_id(event).is_some() => + { + let request_id = event_request_id(event).expect("guarded request_id"); + facts.gpu_pipeline_token_in.insert(request_id); + if event + .get("detail") + .and_then(|detail| detail.get("begin_sequence")) + .and_then(Value::as_bool) + == Some(true) + { + facts.gpu_pipeline_prompt_begin.insert(request_id); + } + } + ("mvp.orch.prompt", Some("OrchPromptEvent")) + if phase == Some("pipeline_token_out") + && status == Some("observed") + && event + .get("detail") + .and_then(|detail| detail.get("token_id")) + .and_then(Value::as_u64) + .is_some() + && event_request_id(event).is_some() => + { + facts + .gpu_pipeline_token_out + .insert(event_request_id(event).expect("guarded request_id")); + } + ("mvp.orch.prompt", Some("OrchPromptEvent")) + if phase == Some("pipeline_tokenizer_decode") + && status == Some("ready") + && event + .get("detail") + .and_then(|detail| detail.get("text_bytes")) + .and_then(Value::as_u64) + .is_some_and(|bytes| bytes > 0) + && event_request_id(event).is_some() => + { + facts + .gpu_pipeline_tokenizer_decode_ready + .insert(event_request_id(event).expect("guarded request_id")); + } + _ => {} + } +} + +fn require_gpu_dump_log_facts(facts: &DumpLogFacts) -> Result<(), String> { + if facts.gpu_cpu_fallback_seen { + return Err("mvp-chat-check: GPU run fell back to the tinygrad CPU compiler".to_owned()); + } + require_dump_log_fact( + facts.gpu_worker_device_requested, + "TinygradImportStarted requested_device CUDA", + )?; + require_dump_log_fact(facts.gpu_import_ready, "TinygradImportReady env_DEV CUDA")?; + require_dump_log_fact(facts.gpu_probe_ready, "TinygradDeviceProbeReady CUDA probe")?; + require_dump_log_fact(facts.gpu_worker_ready, "WorkerReady CUDA backend")?; + for request_id in 1..=2 { + let direct_decode = facts.gpu_decode_started.contains(&request_id) + && facts.gpu_first_token_ready.contains(&request_id) + && facts.gpu_decode_ready.contains(&request_id) + && facts.gpu_prompt_completed.contains(&request_id); + let pipeline_decode = facts.gpu_pipeline_real_worker_step_seen + && facts.gpu_pipeline_prompt_encoded.contains(&request_id) + && facts.gpu_pipeline_prompt_begin.contains(&request_id) + && facts.gpu_pipeline_token_in.contains(&request_id) + && facts.gpu_pipeline_token_out.contains(&request_id) + && facts + .gpu_pipeline_tokenizer_decode_ready + .contains(&request_id) + && facts.gpu_pipeline_tokens_decoded.contains(&request_id); + require_dump_log_fact( + direct_decode || pipeline_decode, + &format!("GPU decode/token evidence request_id={request_id}"), + )?; + } + Ok(()) +} + +fn event_requested_device_is_cuda(event: &Value) -> bool { + event + .get("requested_device") + .and_then(Value::as_str) + .is_some_and(is_cuda_device) +} + +fn event_env_dev_is_cuda(event: &Value) -> bool { + event + .get("env_DEV") + .and_then(Value::as_str) + .is_some_and(is_cuda_device) +} + +fn event_probe_result_is_one(event: &Value) -> bool { + event + .get("probe_result") + .and_then(Value::as_array) + .is_some_and(|values| values.iter().any(|value| value.as_i64() == Some(1))) +} + +fn worker_ready_backend_is_cuda(event: &Value) -> bool { + let Some(backend) = event.get("backend") else { + return false; + }; + backend + .get("requested_device") + .and_then(Value::as_str) + .is_some_and(is_cuda_device) + && backend + .get("env_DEV") + .and_then(Value::as_str) + .is_some_and(is_cuda_device) + && backend + .get("tinygrad_device") + .and_then(Value::as_str) + .is_some_and(is_cuda_device) + && event_probe_result_is_one_from_key(event, "cuda_probe") +} + +fn event_probe_result_is_one_from_key(event: &Value, key: &str) -> bool { + event + .get(key) + .and_then(Value::as_array) + .is_some_and(|values| values.iter().any(|value| value.as_i64() == Some(1))) +} + +fn event_positive_u64(event: &Value, key: &str) -> bool { + event + .get(key) + .and_then(Value::as_u64) + .is_some_and(|value| value > 0) +} + +fn is_cuda_device(value: &str) -> bool { + value.to_ascii_uppercase().contains("CUDA") +} + fn dump_log_request_id(event: &Value) -> Option { event .get("detail") .and_then(|detail| detail.get("request_id")) .and_then(Value::as_u64) } +fn event_request_id(event: &Value) -> Option { + event.get("request_id").and_then(Value::as_u64) +} fn require_dump_log_fact(found: bool, fact: &str) -> Result<(), String> { if found { @@ -1547,6 +1928,85 @@ mod tests { values.iter().map(|value| (*value).to_owned()).collect() } + #[test] + fn scenario_flags_select_expected_launch_contract() { + let dump_log = Path::new("/tmp/mvp-chat-check.ndjson"); + + let baseline = + MvpChatCheckScenario::parse_args(Vec::new()).expect("default scenario parses"); + assert_eq!(baseline, MvpChatCheckScenario::ProcessBaseline); + assert_eq!( + baseline.mvp_chat_args(42, dump_log), + strings(&[ + "--process", + "--cached-model", + "--run-id", + "42", + "--dump-logs=/tmp/mvp-chat-check.ndjson", + ]) + ); + + let gpu = MvpChatCheckScenario::parse_args(strings(&["--gpu"])).expect("gpu parses"); + assert_eq!(gpu, MvpChatCheckScenario::Gpu); + assert!(gpu.env_overrides().is_empty()); + assert_eq!( + gpu.mvp_chat_args(42, dump_log), + strings(&[ + "--process", + "--gpu", + "--run-id", + "42", + "--dump-logs=/tmp/mvp-chat-check.ndjson", + ]) + ); + + let multinode = + MvpChatCheckScenario::parse_args(strings(&["--multinode"])).expect("multinode parses"); + assert_eq!(multinode, MvpChatCheckScenario::Multinode); + assert_eq!( + multinode.mvp_chat_args(42, dump_log), + strings(&[ + "--process", + "--pipeline-stages", + "2", + "--cached-model", + "--run-id", + "42", + "--dump-logs=/tmp/mvp-chat-check.ndjson", + ]) + ); + + let multinode_docker = MvpChatCheckScenario::parse_args(strings(&["--multinode-docker"])) + .expect("multinode docker parses"); + assert_eq!(multinode_docker, MvpChatCheckScenario::MultinodeDocker); + assert_eq!( + multinode_docker.mvp_chat_args(42, dump_log), + strings(&[ + "--docker", + "--pipeline-stages", + "2", + "--cached-model", + "--run-id", + "42", + "--dump-logs=/tmp/mvp-chat-check.ndjson", + ]) + ); + } + + #[test] + fn scenario_flags_reject_unknown_or_ambiguous_invocations() { + assert!( + MvpChatCheckScenario::parse_args(strings(&["--docker"])) + .expect_err("unknown flag fails") + .contains("unsupported mvp-chat-check argument") + ); + assert!( + MvpChatCheckScenario::parse_args(strings(&["--gpu", "--multinode"])) + .expect_err("multiple scenarios fail") + .contains("at most one scenario flag") + ); + } + fn temp_path(label: &str) -> PathBuf { let id = NEXT_TEST_FILE.fetch_add(1, Ordering::Relaxed); std::env::temp_dir().join(format!( @@ -1583,6 +2043,13 @@ mod tests { label: &str, events: Vec<(&'static str, Value)>, ) -> Vec { + let path = write_synthetic_event_dump(label, events); + let parsed = parse_dump_log_events(&path).expect("parse synthetic dump log"); + let _ = fs::remove_file(path); + parsed + } + + fn write_synthetic_event_dump(label: &str, events: Vec<(&'static str, Value)>) -> PathBuf { let lines = events .into_iter() .enumerate() @@ -1594,10 +2061,7 @@ mod tests { ) }) .collect::>(); - let path = write_dump_log(label, lines); - let parsed = parse_dump_log_events(&path).expect("parse synthetic dump log"); - let _ = fs::remove_file(path); - parsed + write_dump_log(label, lines) } fn benchmark(component: &str, wall_unix_ms: u64, mono_ms: u64) -> Value { @@ -1680,11 +2144,23 @@ mod tests { wall_unix_ms, mono_ms, ); + let object = event.as_object_mut().expect("event object"); + match event_type { + "DecodeStarted" => { + object.insert("prompt_tokens".to_owned(), json!(4)); + object.insert("max_tokens".to_owned(), json!(8)); + object.insert("decode_impl".to_owned(), json!("device_resident_greedy")); + } + "FirstTokenReady" => { + object.insert("token_index".to_owned(), json!(1)); + } + "DecodeReady" => { + object.insert("tokens_generated".to_owned(), json!(3)); + } + _ => {} + } if event_type == "PromptCompleted" { - event - .as_object_mut() - .expect("event object") - .insert("generated_tokens".to_owned(), json!([1, 2, 3])); + object.insert("generated_tokens".to_owned(), json!([1, 2, 3])); } event } @@ -1953,6 +2429,7 @@ mod tests { "type": "OrchBootstrap", "phase": "weights_loaded", "status": "ready", + "run_id": 9, "node_id": 1, "detail": {}, @@ -1969,6 +2446,157 @@ mod tests { ] } + fn dump_log_fact_events( + include_gpu: bool, + include_cpu_fallback: bool, + ) -> Vec<(&'static str, Value)> { + let mut events = vec![ + ("mvp.chat.lifecycle", chat_span("config", "ready", 1_000, 0)), + ( + "mvp.chat.runtime", + chat_span("prepare_runtime", "ready", 1_010, 10), + ), + ( + "mvp.chat.runtime", + chat_span("prompt_rpc", "ready", 1_020, 20), + ), + ( + "mvp.orch.bootstrap", + stamped( + json!({"type":"OrchBootstrap","phase":"iroh_driver","status":"ready","run_id":9,"node_id":1,"detail":{}}), + "mvp-orchestrator", + 1_030, + 30, + ), + ), + ( + "mvp.node.bootstrap", + stamped( + json!({"type":"NodeEvent","phase":"iroh_driver","status":"ready","run_id":9,"node_id":3,"stage_index":1,"detail":{}}), + "mvp-worker-node", + 1_040, + 40, + ), + ), + ( + "mvp.node.worker", + stamped( + json!({"type":"NodeEvent","phase":"worker_initialize","status":"ready","run_id":9,"node_id":3,"stage_index":1,"detail":{"device":"CUDA"}}), + "mvp-worker-node", + 1_050, + 50, + ), + ), + ( + "mvp.orch.bootstrap", + stamped( + json!({"type":"OrchBootstrap","phase":"weights_loaded","status":"ready","run_id":9,"node_id":1,"detail":{}}), + "mvp-orchestrator", + 1_060, + 60, + ), + ), + ( + "mvp.chat.prompt", + prompt_chat_span("response_text", "observed", 1, 1_200, 200), + ), + ( + "mvp.chat.prompt", + prompt_chat_span("request_completed", "ready", 1, 1_210, 210), + ), + ( + "mvp.chat.prompt", + prompt_chat_span("response_text", "observed", 2, 1_300, 300), + ), + ( + "mvp.chat.prompt", + prompt_chat_span("request_completed", "ready", 2, 1_310, 310), + ), + ( + "mvp.chat.lifecycle", + chat_span("shutdown", "requested", 1_400, 400), + ), + ( + "mvp.chat.component", + chat_span("orchestrator_process", "stopped", 1_410, 410), + ), + ]; + + if include_gpu { + events.extend([ + ( + "mvp.worker.initialize", + stamped( + json!({"type":"TinygradImportStarted","run_id":9,"node_id":3,"stage_index":1,"requested_device":"CUDA","env_DEV":"CUDA"}), + "tinygrad-worker", + 1_070, + 70, + ), + ), + ( + "mvp.worker.initialize", + stamped( + json!({"type":"TinygradImportReady","run_id":9,"node_id":3,"stage_index":1,"requested_device":"CUDA","env_DEV":"CUDA"}), + "tinygrad-worker", + 1_080, + 80, + ), + ), + ( + "mvp.worker.initialize", + stamped( + json!({"type":"TinygradDeviceProbeReady","run_id":9,"node_id":3,"stage_index":1,"requested_device":"CUDA","probe_result":[1]}), + "tinygrad-worker", + 1_090, + 90, + ), + ), + ( + "mvp.worker.initialize", + stamped( + json!({"type":"WorkerReady","run_id":9,"node_id":3,"stage_index":1,"backend":{"requested_device":"CUDA","env_DEV":"CUDA","tinygrad_device":"CUDA"},"cuda_probe":[1]}), + "tinygrad-worker", + 1_100, + 100, + ), + ), + ]); + if include_cpu_fallback { + events.push(( + "mvp.worker.initialize", + stamped( + json!({"type":"TinygradCpuCompilerSelected","run_id":9,"node_id":3,"stage_index":1,"requested_device":"CUDA","selected_device":"CPU:X86"}), + "tinygrad-worker", + 1_105, + 105, + ), + )); + } + for request_id in 1..=2 { + let base = 1_200 + request_id * 100; + events.extend([ + ( + "mvp.worker.prompt", + worker_prompt_event("DecodeStarted", request_id, base + 20, base - 980), + ), + ( + "mvp.worker.prompt", + worker_prompt_event("FirstTokenReady", request_id, base + 25, base - 975), + ), + ( + "mvp.worker.prompt", + worker_prompt_event("DecodeReady", request_id, base + 40, base - 960), + ), + ( + "mvp.worker.prompt", + worker_prompt_event("PromptCompleted", request_id, base + 50, base - 950), + ), + ]); + } + } + events + } + #[test] fn benchmark_observability_dump_log_path_parser_finds_equals_and_separate_forms() { assert_eq!( @@ -2042,12 +2670,91 @@ mod tests { ); } + #[test] + fn benchmark_observability_gpu_dump_facts_require_cuda_worker_and_decode_cycles() { + let path = write_synthetic_event_dump("gpu-dump-facts", dump_log_fact_events(true, false)); + + let events = assert_dump_log_facts(&path, MvpChatCheckScenario::Gpu) + .expect("GPU dump log facts pass"); + let _ = fs::remove_file(path); + + assert!(events.iter().any(|event| { + event.channel == "mvp.worker.initialize" + && event.event.get("type").and_then(Value::as_str) == Some("WorkerReady") + })); + } + + fn gpu_pipeline_only_facts( + real_worker_backend: bool, + prompt_begin_markers: bool, + ) -> DumpLogFacts { + let mut facts = DumpLogFacts { + gpu_worker_device_requested: true, + gpu_import_ready: true, + gpu_probe_ready: true, + gpu_worker_ready: true, + ..DumpLogFacts::default() + }; + if real_worker_backend { + facts.gpu_pipeline_real_worker_step_seen = true; + } + for request_id in 1..=2 { + facts.gpu_pipeline_prompt_encoded.insert(request_id); + facts.gpu_pipeline_token_in.insert(request_id); + facts.gpu_pipeline_token_out.insert(request_id); + facts.gpu_pipeline_tokenizer_decode_ready.insert(request_id); + facts.gpu_pipeline_tokens_decoded.insert(request_id); + if prompt_begin_markers { + facts.gpu_pipeline_prompt_begin.insert(request_id); + } + } + facts + } + + #[test] + fn benchmark_observability_gpu_pipeline_facts_require_real_steps_and_prompt_begin_markers() { + let valid = gpu_pipeline_only_facts(true, true); + require_gpu_dump_log_facts(&valid).expect("real pipeline facts pass"); + + let missing_real_backend = gpu_pipeline_only_facts(false, true); + let error = require_gpu_dump_log_facts(&missing_real_backend) + .expect_err("missing real worker backend should fail"); + assert!( + error.contains("GPU decode/token evidence request_id=1"), + "unexpected error: {error}" + ); + + let missing_prompt_begin = gpu_pipeline_only_facts(true, false); + let error = require_gpu_dump_log_facts(&missing_prompt_begin) + .expect_err("missing prompt begin marker should fail"); + assert!( + error.contains("GPU decode/token evidence request_id=1"), + "unexpected error: {error}" + ); + } + + #[test] + fn benchmark_observability_gpu_dump_facts_reject_cpu_fallback() { + let path = write_synthetic_event_dump("gpu-cpu-fallback", dump_log_fact_events(true, true)); + + let error = match assert_dump_log_facts(&path, MvpChatCheckScenario::Gpu) { + Ok(_) => panic!("CPU fallback should fail GPU check"), + Err(error) => error, + }; + let _ = fs::remove_file(path); + + assert!( + error.contains("fell back to the tinygrad CPU compiler"), + "{error}" + ); + } + #[test] fn benchmark_observability_report_requires_granular_decode_events() { let events = parse_synthetic_events("missing-first-token", benchmark_report_events(false)); - let error = - build_benchmark_report(&events, 80, 9).expect_err("missing first token should fail"); + let error = build_benchmark_report(&events, 80, 9, MvpChatCheckScenario::ProcessBaseline) + .expect_err("missing first token should fail"); assert!( error.starts_with("mvp-chat-check: missing benchmark event "), @@ -2064,7 +2771,8 @@ mod tests { let events = parse_synthetic_events("pipeline-report", pipeline_benchmark_report_events(true)); - let report = build_benchmark_report(&events, 80, 9).expect("pipeline report builds"); + let report = build_benchmark_report(&events, 80, 9, MvpChatCheckScenario::ProcessBaseline) + .expect("pipeline report builds"); assert!( report @@ -2089,7 +2797,7 @@ fn main() -> ExitCode { let mut args = std::env::args().skip(1); match args.next().as_deref() { Some("test") if args.next().is_none() => run_tests(), - Some("mvp-chat-check") if args.next().is_none() => run_mvp_chat_check(), + Some("mvp-chat-check") => run_mvp_chat_check(args.collect()), Some("mvp-chat") => run_mvp_chat(args.collect()), Some("help" | "--help" | "-h") | None => { print_usage();