diff --git a/crates/mvp-system/src/chat/node_image.rs b/crates/mvp-system/src/chat/node_image.rs index 904b917..84e592f 100644 --- a/crates/mvp-system/src/chat/node_image.rs +++ b/crates/mvp-system/src/chat/node_image.rs @@ -103,8 +103,7 @@ fn prepare_node_image_inner( )); } - run_status( - progress, + run_status_command( &root, "cargo", &[ @@ -117,6 +116,7 @@ fn prepare_node_image_inner( ], "build mvp-worker-node", None, + progress, )?; let base_hash = content_hash_for_inputs(&root, BASE_IMAGE_SOURCE_INPUTS)?; @@ -124,12 +124,13 @@ fn prepare_node_image_inner( let tag = image_version_tag(&root, &image_content_hash)?; let image_ref = image.ref_for_tag(&tag); emit_image_reference(progress, "resolved", &image_ref); - let worker_hash = hash_relative_files( + let worker_hash = hash_relative_files_with_salts( &root, vec![relative_path( &root, &root.join("apps/mvp-node/tinygrad_worker.py"), )?], + &[], )?; let expected_node_labels = vec![ (NODE_IMAGE_TAG_LABEL, tag.as_str()), @@ -168,18 +169,19 @@ fn prepare_node_image_inner( let base_image_matches = docker_image_labels_match(&root, &request.base_image, &expected_base_labels)?; if !base_image_matches { + let base_source_hash_label = format!("{BASE_IMAGE_SOURCE_HASH_LABEL}={base_hash}"); run_status_command( &root, "docker", - &vec![ - "build".to_owned(), - "-f".to_owned(), - "apps/mvp-node/Dockerfile.base".to_owned(), - "--label".to_owned(), - format!("{BASE_IMAGE_SOURCE_HASH_LABEL}={base_hash}"), - "-t".to_owned(), - request.base_image.clone(), - ".".to_owned(), + &[ + "build", + "-f", + "apps/mvp-node/Dockerfile.base", + "--label", + base_source_hash_label.as_str(), + "-t", + request.base_image.as_str(), + ".", ], "build mvp node base image", Some(&request.base_image), @@ -187,7 +189,14 @@ fn prepare_node_image_inner( )?; } - let node_bin = docker_build_context_path(&root, &request.node_bin)?; + let node_bin = { + let full = if request.node_bin.is_absolute() { + request.node_bin.to_path_buf() + } else { + root.join(&request.node_bin) + }; + relative_path(&root, &full).map(|relative| relative.to_string_lossy().to_string()) + }?; let mut build_args = vec![ "build".to_owned(), "-f".to_owned(), @@ -205,7 +214,7 @@ fn prepare_node_image_inner( run_status_command( &root, "docker", - &build_args, + &build_args.iter().map(String::as_str).collect::>(), "build mvp node image", Some(&image_ref), progress, @@ -302,10 +311,6 @@ fn content_hash_for_inputs(root: &Path, inputs: &[&str]) -> Result) -> Result { hash_relative_files_with_salts(root, files, &[]) } @@ -360,7 +365,7 @@ fn collect_hash_inputs(root: &Path, path: &Path, out: &mut Vec) -> Resu } return Ok(()); } - if !metadata.is_dir() || skip_dir(path) { + if !metadata.is_dir() || matches!(path.file_name().and_then(|name| name.to_str()), Some(".git" | "target" | "__pycache__")) { return Ok(()); } let entries = fs::read_dir(path).map_err(|e| format!("read dir {display}: {e}"))?; @@ -381,15 +386,6 @@ fn relative_path(root: &Path, path: &Path) -> Result { }) } -fn docker_build_context_path(root: &Path, path: &Path) -> Result { - let full = if path.is_absolute() { - path.to_path_buf() - } else { - root.join(path) - }; - relative_path(root, &full).map(|relative| relative.to_string_lossy().to_string()) -} - fn display_workspace_path(root: &Path, path: &Path) -> String { match path.strip_prefix(root) { Ok(relative) if relative.as_os_str().is_empty() => ".".to_owned(), @@ -398,13 +394,6 @@ fn display_workspace_path(root: &Path, path: &Path) -> String { } } -fn skip_dir(path: &Path) -> bool { - matches!( - path.file_name().and_then(|name| name.to_str()), - Some(".git" | "target" | "__pycache__") - ) -} - fn alias_tags( image: &ImageName, extra_tag: Option<&str>, @@ -444,13 +433,13 @@ fn ensure_aliases_local( ) -> Result<(), String> { for alias in alias_refs(image, alias_tags) { if alias != source_ref { - run_status( - progress, + run_status_command( root, "docker", &["tag", source_ref, &alias], "tag mvp node image", Some(&alias), + progress, )?; } } @@ -468,13 +457,13 @@ fn ensure_aliases_for_remote( return Ok(false); } if !docker_image_exists(root, source_ref) { - run_status( - progress, + run_status_command( root, "docker", &["pull", source_ref], "pull mvp node image", Some(source_ref), + progress, )?; } ensure_aliases_local(progress, root, source_ref, image, alias_tags)?; @@ -496,13 +485,13 @@ fn push_image( root: &Path, image_ref: &str, ) -> Result<(), String> { - run_status( - progress, + run_status_command( root, "docker", &["push", image_ref], "push mvp node image", Some(image_ref), + progress, ) } @@ -581,18 +570,6 @@ fn prune_old_dirty_images(root: &Path, image: &ImageName, keep_tag: &str) { } } -fn run_status( - progress: &mut Option<&mut dyn NodeImageProgressSink>, - root: &Path, - program: &str, - args: &[&str], - label: &str, - image_ref: Option<&str>, -) -> Result<(), String> { - let args = args.iter().map(|arg| (*arg).to_owned()).collect::>(); - run_status_command(root, program, &args, label, image_ref, progress) -} - fn emit_image_reference( progress: &mut Option<&mut dyn NodeImageProgressSink>, role: &str, @@ -686,16 +663,17 @@ fn drain_command_lines( fn run_status_command( root: &Path, program: &str, - args: &[String], + args: &[&str], label: &str, image_ref: Option<&str>, progress: &mut Option<&mut dyn NodeImageProgressSink>, ) -> Result<(), String> { + let args: Vec = args.iter().map(|arg| (*arg).to_owned()).collect(); eprintln!("mvp-node-image: {label}"); if progress.is_none() { let status = Command::new(program) .current_dir(root) - .args(args) + .args(&args) .stdin(Stdio::null()) .stdout(Stdio::inherit()) .stderr(Stdio::inherit()) @@ -721,7 +699,7 @@ fn run_status_command( ); let mut child = match Command::new(program) .current_dir(root) - .args(args) + .args(&args) .stdin(Stdio::null()) .stdout(Stdio::piped()) .stderr(Stdio::piped()) diff --git a/crates/mvp-system/src/chat/runtime.rs b/crates/mvp-system/src/chat/runtime.rs index c3b904e..c6940c3 100644 --- a/crates/mvp-system/src/chat/runtime.rs +++ b/crates/mvp-system/src/chat/runtime.rs @@ -1351,6 +1351,33 @@ struct InProcessOrch { cleaned: bool, } +fn wait_for_rpc_ready( + rpc_addr: &str, + mut check_dead: impl FnMut() -> Result<(), 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.to_owned()); + } + 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}")), + } + check_dead()?; + thread::sleep(Duration::from_millis(100)); + } +} + impl InProcessOrch { fn spawn(config: &Config, image_ref: &str) -> Result { let args = config.orchestrator_cli_args(image_ref); @@ -1366,24 +1393,7 @@ impl InProcessOrch { } 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}")), - } + wait_for_rpc_ready(&rpc_addr, || { if let Some(result) = self.take_finished_result() { let reason = match result { Ok(()) => "completed successfully".to_owned(), @@ -1393,8 +1403,8 @@ impl InProcessOrch { "in-process orchestrator exited before prompt RPC ready: {reason}" )); } - thread::sleep(Duration::from_millis(100)); - } + Ok(()) + }) } fn shutdown(&mut self) { @@ -1478,24 +1488,7 @@ impl OrchChild { } 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}")), - } + wait_for_rpc_ready(&rpc_addr, || { if let Some(status) = self .child .try_wait() @@ -1505,8 +1498,8 @@ impl OrchChild { "orchestrator exited before prompt RPC ready: {status}" )); } - thread::sleep(Duration::from_millis(100)); - } + Ok(()) + }) } // The orchestrator shutdown spec is still pending. Replace this with the approved @@ -1566,6 +1559,60 @@ fn signal_orch_process_group(child: &Child, signal: libc::c_int) -> io::Result<( } } +fn ensure_binary_with_progress( + progress: &mut Option<&mut ChatDatastream>, + phase: &str, + mode: &str, + verbose: bool, + skip_rebuild: bool, + bin: &Path, + label: &str, + cargo_args: &[&str], +) -> Result<(), String> { + let started = Instant::now(); + emit_chat_progress( + progress, + CHAT_RUNTIME_CHANNEL, + phase, + "started", + if verbose { + json!({"mode": mode, "command_label": phase}) + } else { + json!({"mode": mode}) + }, + ); + match ensure_runtime_binary(skip_rebuild, bin, label, cargo_args) { + Ok(()) => { + emit_chat_progress( + progress, + CHAT_RUNTIME_CHANNEL, + phase, + "ready", + if verbose { + json!({"mode": mode, "command_label": phase, "elapsed_ms": started.elapsed().as_millis()}) + } else { + json!({"mode": mode}) + }, + ); + Ok(()) + } + Err(error) => { + emit_chat_progress( + progress, + CHAT_RUNTIME_CHANNEL, + phase, + "failed", + if verbose { + json!({"mode": mode, "command_label": phase, "elapsed_ms": started.elapsed().as_millis(), "error": error.as_str()}) + } else { + json!({"mode": mode, "error": error.as_str()}) + }, + ); + Err(error) + } + } +} + fn prepare_runtime_with_progress( config: &Config, mut prepare_node_image_fn: F, @@ -1596,15 +1643,11 @@ where json!({"mode": config.orchestrator_launch_mode()}), ); } else { - let ensure_orch_started = Instant::now(); - emit_chat_progress( + ensure_binary_with_progress( &mut progress, - CHAT_RUNTIME_CHANNEL, "ensure_orch_binary", - "started", - json!({"mode": binary_mode, "command_label": "ensure_orch_binary"}), - ); - match ensure_runtime_binary( + binary_mode, + true, config.skip_rebuild, &config.orch_bin, "mvp-orchestrator", @@ -1618,36 +1661,15 @@ where "--bin", "mvp-orchestrator", ], - ) { - Ok(()) => emit_chat_progress( - &mut progress, - CHAT_RUNTIME_CHANNEL, - "ensure_orch_binary", - "ready", - json!({"mode": binary_mode, "command_label": "ensure_orch_binary", "elapsed_ms": ensure_orch_started.elapsed().as_millis()}), - ), - Err(error) => { - emit_chat_progress( - &mut progress, - CHAT_RUNTIME_CHANNEL, - "ensure_orch_binary", - "failed", - json!({"mode": binary_mode, "command_label": "ensure_orch_binary", "elapsed_ms": ensure_orch_started.elapsed().as_millis(), "error": error.as_str()}), - ); - return Err(error); - } - } + )?; } if config.provider == provider_kind::process() { - emit_chat_progress( + ensure_binary_with_progress( &mut progress, - CHAT_RUNTIME_CHANNEL, "ensure_worker_binary", - "started", - json!({"mode": binary_mode}), - ); - match ensure_runtime_binary( + binary_mode, + false, config.skip_rebuild, &config.worker_bin, "mvp-worker-node", @@ -1659,25 +1681,7 @@ where "--bin", "mvp-worker-node", ], - ) { - Ok(()) => emit_chat_progress( - &mut progress, - CHAT_RUNTIME_CHANNEL, - "ensure_worker_binary", - "ready", - json!({"mode": binary_mode}), - ), - Err(error) => { - emit_chat_progress( - &mut progress, - CHAT_RUNTIME_CHANNEL, - "ensure_worker_binary", - "failed", - json!({"mode": binary_mode, "error": error.as_str()}), - ); - return Err(error); - } - } + )?; emit_chat_progress( &mut progress, CHAT_RUNTIME_CHANNEL, @@ -1706,15 +1710,11 @@ where ); return Ok(config.node_image.clone()); } - let ensure_worker_started = Instant::now(); - emit_chat_progress( + ensure_binary_with_progress( &mut progress, - CHAT_RUNTIME_CHANNEL, "ensure_worker_binary", - "started", - json!({"mode": binary_mode, "command_label": "ensure_worker_binary"}), - ); - match ensure_runtime_binary( + binary_mode, + true, config.skip_rebuild, &config.worker_bin, "mvp-worker-node", @@ -1726,25 +1726,7 @@ where "--bin", "mvp-worker-node", ], - ) { - Ok(()) => emit_chat_progress( - &mut progress, - CHAT_RUNTIME_CHANNEL, - "ensure_worker_binary", - "ready", - json!({"mode": binary_mode, "command_label": "ensure_worker_binary", "elapsed_ms": ensure_worker_started.elapsed().as_millis()}), - ), - Err(error) => { - emit_chat_progress( - &mut progress, - CHAT_RUNTIME_CHANNEL, - "ensure_worker_binary", - "failed", - json!({"mode": binary_mode, "command_label": "ensure_worker_binary", "elapsed_ms": ensure_worker_started.elapsed().as_millis(), "error": error.as_str()}), - ); - return Err(error); - } - } + )?; emit_chat_progress( &mut progress, CHAT_RUNTIME_CHANNEL, @@ -1764,25 +1746,7 @@ where json!({"provider": config.provider.as_str(), "command_label": "prepare_node_image", "image_tag": config.image_tag.as_deref()}), ); let node_bin = default_worker_bin(); - let requires_registry_image = if config.provider == provider_kind::docker() { - false - } else if config.provider == provider_kind::vastai() { - true - } else { - let error = if config.provider == provider_kind::process() { - "process provider does not use node images" - } else { - "mvp-chat does not support mock provider" - }; - emit_chat_progress( - &mut progress, - CHAT_RUNTIME_CHANNEL, - "prepare_node_image", - "failed", - json!({"provider": config.provider.as_str(), "command_label": "prepare_node_image", "elapsed_ms": prepare_node_image_started.elapsed().as_millis(), "error": error}), - ); - return Err(error.to_owned()); - }; + let requires_registry_image = config.provider == provider_kind::vastai(); let prepared = { let command_progress = progress .as_deref_mut() @@ -1905,10 +1869,6 @@ fn emit_chat_progress( } } -fn prompt_hash_hex(prompt: &str) -> String { - blake3::hash(prompt.as_bytes()).to_hex().to_string() -} - fn run_chat_session_with_output_and_progress( writer: &mut impl Write, mut reader: impl BufRead, @@ -1920,16 +1880,11 @@ fn run_chat_session_with_output_and_progress( let mut progress = progress; let mut next_request_id = 1_u64; let mut next_prompt_index = 1_u64; + let prompt_exited = |progress: &mut Option<&mut ChatDatastream>, reason: &str| emit_chat_progress(progress, CHAT_PROMPT_CHANNEL, "prompt_loop", "exited", json!({"reason": reason})); loop { if STOP_REQUESTED.load(Ordering::SeqCst) { - emit_chat_progress( - &mut progress, - CHAT_PROMPT_CHANNEL, - "prompt_loop", - "exited", - json!({"reason": "stop_requested"}), - ); + prompt_exited(&mut progress, "stop_requested"); return Ok(()); } emit_chat_progress( @@ -1944,23 +1899,11 @@ fn run_chat_session_with_output_and_progress( let prompt = match input_rx.recv() { Ok(PromptInput::Line(line)) => line.trim_end().to_owned(), Ok(PromptInput::Closed) | Err(_) => { - emit_chat_progress( - &mut progress, - CHAT_PROMPT_CHANNEL, - "prompt_loop", - "exited", - json!({"reason": "input_closed"}), - ); + prompt_exited(&mut progress, "input_closed"); return Ok(()); } Ok(PromptInput::StopRequested) => { - emit_chat_progress( - &mut progress, - CHAT_PROMPT_CHANNEL, - "prompt_loop", - "exited", - json!({"reason": "stop_requested"}), - ); + prompt_exited(&mut progress, "stop_requested"); return Ok(()); } }; @@ -1972,7 +1915,7 @@ fn run_chat_session_with_output_and_progress( next_request_id = next_request_id.wrapping_add(1).max(1); let prompt_index = next_prompt_index; next_prompt_index = next_prompt_index.wrapping_add(1).max(1); - let prompt_hash = prompt_hash_hex(&prompt); + let prompt_hash = blake3::hash(prompt.as_bytes()).to_hex().to_string(); emit_chat_progress( &mut progress, CHAT_PROMPT_CHANNEL, @@ -2000,13 +1943,7 @@ fn run_chat_session_with_output_and_progress( loop { if STOP_REQUESTED.load(Ordering::SeqCst) { - emit_chat_progress( - &mut progress, - CHAT_PROMPT_CHANNEL, - "prompt_loop", - "exited", - json!({"reason": "stop_requested"}), - ); + prompt_exited(&mut progress, "stop_requested"); return Ok(()); } let mut line = String::new(); @@ -2269,16 +2206,6 @@ fn provider_from_sources( Ok(provider_kind::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() @@ -2302,6 +2229,16 @@ where .map_err(|e| format!("invalid {name}={value:?}: {e}")) } +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 parse_pipeline_stages_value( args: &mut impl Iterator, name: &str, diff --git a/crates/mvp-system/src/node/worker_node_runtime.rs b/crates/mvp-system/src/node/worker_node_runtime.rs index 9e510ff..8eef579 100644 --- a/crates/mvp-system/src/node/worker_node_runtime.rs +++ b/crates/mvp-system/src/node/worker_node_runtime.rs @@ -557,26 +557,75 @@ fn submit_sampler_sample_health( seq: u64, error: Option<&str>, ) { - match error { - Some(error) => submit_sampler_health( - producer, - health_channel, - context, - sampler, - sample_channel, + let (status, detail) = match error { + Some(error) => ( "failed", json!({"state":"error","sample_seq":seq,"error":error}), ), - None => submit_sampler_health( - producer, - health_channel, - context, - sampler, - sample_channel, + None => ( "ready", json!({"state":"sample_observed","sample_seq":seq}), ), - } + }; + submit_sampler_health( + producer, + health_channel, + context, + sampler, + sample_channel, + status, + detail, + ); +} + +fn spawn_blocking_sampler( + handle: tokio::runtime::Handle, + producer: DatastreamProducer, + channel: ChannelId, + health_channel: ChannelId, + health_context: SamplerHealthContext, + sampler: &'static str, + sample_channel: &'static str, + interval: Duration, + error_label: &'static str, + sample_fn: fn(u64) -> S, + error_fn: fn(u64, String) -> S, + error_of: fn(&S) -> Option<&str>, +) { + handle.spawn(async move { + submit_sampler_started( + &producer, + health_channel, + health_context, + sampler, + sample_channel, + interval, + ); + let mut seq = 0_u64; + let mut interval = tokio::time::interval(interval); + + loop { + interval.tick().await; + + let sample_seq = seq; + let sample = match tokio::task::spawn_blocking(move || sample_fn(sample_seq)).await { + Ok(sample) => sample, + Err(error) => error_fn(sample_seq, format!("{error_label}: {error}")), + }; + + submit_sampler_sample_health( + &producer, + health_channel, + health_context, + sampler, + sample_channel, + sample_seq, + error_of(&sample), + ); + seq = seq.saturating_add(1); + producer.submit_record(channel, &sample); + } + }); } fn spawn_host_gpu_sampler( @@ -586,48 +635,20 @@ fn spawn_host_gpu_sampler( health_channel: ChannelId, health_context: SamplerHealthContext, ) { - handle.spawn(async move { - let sample_channel = datastream::hardware::gpu::HOST_GPU_CHANNEL; - submit_sampler_started( - &producer, - health_channel, - health_context, - "gpu", - sample_channel, - datastream::hardware::gpu::GPU_SAMPLE_INTERVAL, - ); - let mut seq = 0_u64; - let mut interval = tokio::time::interval(datastream::hardware::gpu::GPU_SAMPLE_INTERVAL); - - loop { - interval.tick().await; - - let sample_seq = seq; - let sample = match tokio::task::spawn_blocking(move || { - datastream::hardware::gpu::sample(sample_seq) - }) - .await - { - Ok(sample) => sample, - Err(error) => datastream::hardware::gpu::HostGpuSample::error( - sample_seq, - format!("gpu sampler task failed: {error}"), - ), - }; - - submit_sampler_sample_health( - &producer, - health_channel, - health_context, - "gpu", - sample_channel, - sample_seq, - sample.error.as_deref(), - ); - seq = seq.saturating_add(1); - producer.submit_record(channel, &sample); - } - }); + spawn_blocking_sampler( + handle, + producer, + channel, + health_channel, + health_context, + "gpu", + datastream::hardware::gpu::HOST_GPU_CHANNEL, + datastream::hardware::gpu::GPU_SAMPLE_INTERVAL, + "gpu sampler task failed", + datastream::hardware::gpu::sample, + datastream::hardware::gpu::HostGpuSample::error, + |s| s.error.as_deref(), + ); } fn spawn_host_cpu_sampler( @@ -677,49 +698,20 @@ fn spawn_host_net_sampler( health_channel: ChannelId, health_context: SamplerHealthContext, ) { - handle.spawn(async move { - let sample_channel = datastream::hardware::net::HOST_NET_CHANNEL; - submit_sampler_started( - &producer, - health_channel, - health_context, - "net", - sample_channel, - datastream::hardware::net::HOST_NET_SAMPLE_INTERVAL, - ); - let mut seq = 0_u64; - let mut interval = - tokio::time::interval(datastream::hardware::net::HOST_NET_SAMPLE_INTERVAL); - - loop { - interval.tick().await; - - let sample_seq = seq; - let sample = match tokio::task::spawn_blocking(move || { - datastream::hardware::net::sample(sample_seq) - }) - .await - { - Ok(sample) => sample, - Err(error) => datastream::hardware::net::HostNetSample::error( - sample_seq, - format!("network sampler task failed: {error}"), - ), - }; - - submit_sampler_sample_health( - &producer, - health_channel, - health_context, - "net", - sample_channel, - sample_seq, - sample.error.as_deref(), - ); - seq = seq.saturating_add(1); - producer.submit_record(channel, &sample); - } - }); + spawn_blocking_sampler( + handle, + producer, + channel, + health_channel, + health_context, + "net", + datastream::hardware::net::HOST_NET_CHANNEL, + datastream::hardware::net::HOST_NET_SAMPLE_INTERVAL, + "network sampler task failed", + datastream::hardware::net::sample, + datastream::hardware::net::HostNetSample::error, + |s| s.error.as_deref(), + ); } fn spawn_arena_sampler( @@ -857,7 +849,6 @@ impl WorkerEdgeRuntime { driver, )?; } - EdgeTransportEvent::StreamEnded { .. } => {} EdgeTransportEvent::StreamFault { edge_id: Some(edge_id), .. @@ -873,7 +864,8 @@ impl WorkerEdgeRuntime { driver, )?; } - EdgeTransportEvent::StreamFault { edge_id: None, .. } => {} + EdgeTransportEvent::StreamEnded { .. } + | EdgeTransportEvent::StreamFault { edge_id: None, .. } => {} } } Ok(()) @@ -1952,7 +1944,7 @@ fn run() -> Result<(), String> { datastream.channels.host_cpu, sampler_health_channel, sampler_health_context, - vec![std::process::id(), worker.pid()], + vec![std::process::id(), worker.child.id()], ); worker_evt( "worker_initialize", @@ -2032,7 +2024,7 @@ fn run() -> Result<(), String> { emit_swim_telemetry(&mut datastream, &stack, "main_loop"); drain_debug_join_commands(&mut debug_join_rx, &mut driver, &config, &mut datastream); datastream.tick(); - worker.drain_stderr(&config, &mut datastream); + drain_worker_stderr(&worker.stderr_rx, &config, &mut datastream); edge_runtime.poll_iroh( &mut driver, &stack, @@ -3171,15 +3163,6 @@ fn stop_stage_shard_child(child: &mut Option) { let _ = child.wait(); } -fn stage_shard_cache_path(plan: &StageShardPlan) -> PathBuf { - let root = std::env::var("MVP_MODEL_CACHE_DIR") - .ok() - .filter(|value| !value.trim().is_empty()) - .map(PathBuf::from) - .unwrap_or_else(|| PathBuf::from("/var/cache/mvp-models")); - root.join(plan.cache_file_name()) -} - fn spawn_stage_shard_reader( stream: StageShardProcessStream, reader: R, @@ -3251,7 +3234,12 @@ fn materialize_stage_shard_with_process( driver: &mut IrohDriver, stack: &DistributionRuntimeStack, ) -> Result { - let output_path = stage_shard_cache_path(plan); + let output_path = std::env::var("MVP_MODEL_CACHE_DIR") + .ok() + .filter(|value| !value.trim().is_empty()) + .map(PathBuf::from) + .unwrap_or_else(|| PathBuf::from("/var/cache/mvp-models")) + .join(plan.cache_file_name()); if output_path.is_file() { match validate_stage_shard_cache(&output_path, plan) { Ok(()) => { @@ -3678,6 +3666,13 @@ fn run_self_test( Ok(()) } +fn env_optional(name: &str) -> Option { + std::env::var(name) + .ok() + .map(|value| value.trim().to_owned()) + .filter(|value| !value.is_empty()) +} + #[derive(Clone)] struct DeploymentConfig { run_id: u64, @@ -4071,10 +4066,6 @@ impl TinygradWorker { }) } - fn pid(&self) -> u32 { - self.child.id() - } - fn initialize( &mut self, device: &str, @@ -4402,10 +4393,6 @@ impl TinygradWorker { .map_err(|e| format!("poll tinygrad helper: {e}")) } - fn drain_stderr(&mut self, config: &DeploymentConfig, datastream: &mut NodeDatastream) { - drain_worker_stderr(&self.stderr_rx, config, datastream); - } - fn command( &mut self, command: Value, @@ -4512,10 +4499,3 @@ fn spawn_stdin_shutdown_listener() -> Receiver<()> { }); rx } - -fn env_optional(name: &str) -> Option { - std::env::var(name) - .ok() - .map(|value| value.trim().to_owned()) - .filter(|value| !value.is_empty()) -} diff --git a/crates/mvp-system/src/observability/frame_archive.rs b/crates/mvp-system/src/observability/frame_archive.rs index 8d49995..8e5c5f4 100644 --- a/crates/mvp-system/src/observability/frame_archive.rs +++ b/crates/mvp-system/src/observability/frame_archive.rs @@ -19,10 +19,6 @@ pub(crate) struct FrameArchive { } impl FrameArchive { - pub(crate) fn open(path: &Path) -> Result { - Self::open_with_label(path, "datastream frame log") - } - pub(crate) fn open_with_label(path: &Path, label: &'static str) -> Result { if let Some(parent) = path.parent() && !parent.as_os_str().is_empty() diff --git a/crates/mvp-system/src/observability/lifecycle.rs b/crates/mvp-system/src/observability/lifecycle.rs index e837756..3c9f477 100644 --- a/crates/mvp-system/src/observability/lifecycle.rs +++ b/crates/mvp-system/src/observability/lifecycle.rs @@ -1,4 +1,3 @@ -#![allow(dead_code)] use serde::{Deserialize, Serialize}; @@ -179,11 +178,13 @@ impl Event { } } +#[cfg(test)] pub(crate) struct TraceBuilder { run_id: RunId, events: Vec, } +#[cfg(test)] impl TraceBuilder { pub(crate) fn new(run_id: RunId) -> Self { Self { @@ -473,6 +474,7 @@ impl TraceBuilder { } } +#[cfg(test)] pub(crate) fn requires_log_scraping(_events: &[Event]) -> bool { false } diff --git a/crates/mvp-system/src/observability/provisioning_logs.rs b/crates/mvp-system/src/observability/provisioning_logs.rs index aa33a57..d16e089 100644 --- a/crates/mvp-system/src/observability/provisioning_logs.rs +++ b/crates/mvp-system/src/observability/provisioning_logs.rs @@ -1,4 +1,3 @@ -#![allow(dead_code)] use std::io::{BufRead, BufReader, Read}; use std::thread::{self, JoinHandle}; @@ -14,12 +13,8 @@ use crate::provisioning::{ NodeProvisionSpec, PluginObservation, PluginSink, ProvisionLogLine, ProvisionLogStream, }; -pub(crate) fn node_datastream_id(node_id: u64) -> String { - node_id.to_string() -} - pub(crate) fn node_stream_id(run_id: u64, node_id: u64) -> StreamId { - StreamId::new(NodeId::new(&node_datastream_id(node_id)), Lifetime(run_id)) + StreamId::new(NodeId::new(&node_id.to_string()), Lifetime(run_id)) } #[derive(Clone)] @@ -46,10 +41,6 @@ impl BootstrapDatastreamBridge { &self.spec } - pub(crate) fn stream_id(&self) -> StreamId { - node_stream_id(self.spec.run_id, self.spec.node_id) - } - pub(crate) fn observe_stdout_line(&self, line: impl Into) { let line = line.into(); if let Some(frame) = parse_stdio_datastream_frame(&self.spec, &line) { @@ -185,6 +176,3 @@ pub(crate) fn parse_stdio_datastream_frame( }) } -pub(crate) fn bootstrap_log_channel(node_id: u64, stream: ProvisionLogStream) -> String { - mvp_provision_log_channel(node_id, stream) -} diff --git a/crates/mvp-system/src/observability/telemetry.rs b/crates/mvp-system/src/observability/telemetry.rs index ce10814..1d9888c 100644 --- a/crates/mvp-system/src/observability/telemetry.rs +++ b/crates/mvp-system/src/observability/telemetry.rs @@ -1,14 +1,11 @@ -#![allow(dead_code)] //! MVP-system-owned datastream channel records. -use datastream::hardware::net::HostNetSample; -use datastream::{ChannelRegistry, Record}; +use datastream::Record; use serde::{Deserialize, Serialize}; use crate::observability::lifecycle as obs; use crate::provisioning::{self, ProvisionLogStream}; -use data_plane::arena::ArenaSample; /// Structured MVP lifecycle facts: run, node, stage, edge, ring, object, step, and worker events. pub(crate) const MVP_LIFECYCLE: &str = "mvp.lifecycle"; @@ -24,22 +21,6 @@ pub(crate) struct MvpLifecycleRecord { pub event: obs::Event, } -impl MvpLifecycleRecord { - pub(crate) fn new(event: obs::Event) -> Self { - Self { event } - } - - pub(crate) fn kind(&self) -> obs::EventKind { - self.event.kind() - } -} - -impl From for MvpLifecycleRecord { - fn from(event: obs::Event) -> Self { - Self::new(event) - } -} - impl Record for MvpLifecycleRecord { const CHANNEL: &'static str = MVP_LIFECYCLE; } @@ -84,13 +65,3 @@ impl Record for MvpProvisionLogRecord { const CHANNEL: &'static str = MVP_PROVISIONING_LOGS; } -/// Registry fragment for consumers that want typed MVP datastream decoding. -pub(crate) fn channel_registry() -> ChannelRegistry { - let registry = ChannelRegistry::new() - .with_record::() - .with_record::() - .with_record::() - .with_record::() - .with_record::(); - registry -} diff --git a/crates/mvp-system/src/orchestration/actor.rs b/crates/mvp-system/src/orchestration/actor.rs index c8be606..e16b716 100644 --- a/crates/mvp-system/src/orchestration/actor.rs +++ b/crates/mvp-system/src/orchestration/actor.rs @@ -228,9 +228,7 @@ impl OrchestratorActor { run_id: core::RunId(run_id), stage_index, }), - OrchestratorMsg::ObserveNodeRuntimeReady { .. } => {} - OrchestratorMsg::ObserveNodeRuntimeReadyAck { .. } => {} - OrchestratorMsg::ObserveWeightsReady { .. } => {} + OrchestratorMsg::ObserveNodeRuntimeReady { .. } | OrchestratorMsg::ObserveNodeRuntimeReadyAck { .. } | OrchestratorMsg::ObserveWeightsReady { .. } | OrchestratorMsg::Snapshot { .. } => {} OrchestratorMsg::ObserveTokenInEndpointReady => { self.core.observe(core::RunEvent::TokenInEndpointReady) } @@ -275,7 +273,6 @@ impl OrchestratorActor { self.core.observe(core::RunEvent::TokenEndpointsStopped) } OrchestratorMsg::AdvanceTimeMs(delta) => self.core.advance_time_ms(delta), - OrchestratorMsg::Snapshot { .. } => {} } } diff --git a/crates/mvp-system/src/orchestration/app.rs b/crates/mvp-system/src/orchestration/app.rs index b868356..9911ff4 100644 --- a/crates/mvp-system/src/orchestration/app.rs +++ b/crates/mvp-system/src/orchestration/app.rs @@ -1792,8 +1792,8 @@ impl Config { } } - fn node_spec_env_keys(&self) -> Vec<&'static str> { - let mut keys = vec![ + fn node_spec_env_keys(&self) -> Vec { + let mut keys: Vec = vec![ "MVP_RUN_ID", "MVP_LOGICAL_NODE_ID", "MVP_NODE_PROVIDER", @@ -1804,90 +1804,18 @@ impl Config { "MVP_IROH_RELAY_MODE", MVP_IROH_ENDPOINT_ADDR_MASK_ENV, "MVP_PIPELINE_STAGES", - ]; - if self.relay.url.is_some() { - keys.push(MVP_IROH_RELAY_URL_ENV); - } - if self.provider.as_str() == "docker" { - keys.push("MVP_DOCKER_GPUS"); - } - if std::env::var_os("DEV").is_some() { - keys.push("DEV"); - } - if local_tinygrad_worker_env(self.provider.as_str()).is_some() { - keys.push("MVP_TINYGRAD_WORKER"); - } - for key in [ - "MVP_CPU_LINE_PROFILE", - "MVP_CPU_LINE_PROFILE_INTERVAL_MS", - "MVP_TOKEN_PROGRESS_EVERY", - "CUDA_DEVICE_SCHEDULE", - "MVP_MODEL_CACHE_DIR", - "HF_TOKEN", - ] { - if std::env::var_os(key).is_some() { - keys.push(key); - } - } - 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"); - } + ] + .into_iter() + .map(str::to_owned) + .collect(); + keys.extend(self.extra_worker_env().into_iter().map(|(k, _)| k)); keys } - fn node_spec_for_stage( - &self, - coordinator: EndpointAddr, - orchestrator_actor: ActorAddress, - logical_node_id: u64, - stage_index: u32, - ) -> Result { + /// Conditional worker env pairs shared by `node_spec_env_keys` and `node_spec_for_stage`. + fn extra_worker_env(&self) -> Vec<(String, String)> { let provider_name = self.provider.as_str(); - 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_IROH_ENDPOINT_ADDR_MASK_ENV.to_owned(), - self.endpoint_addr_mask.as_str().to_owned(), - ), - ("MVP_NODE_PROVIDER".to_owned(), provider_name.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(), - ), - ]; + let mut env = Vec::new(); if let Some(url) = &self.relay.url { env.push((MVP_IROH_RELAY_URL_ENV.to_owned(), url.clone())); } @@ -1932,6 +1860,50 @@ impl Config { if let Some(max_context) = self.max_context { env.push(("MVP_MAX_CONTEXT".to_owned(), max_context.to_string())); } + env + } + + fn node_spec_for_stage( + &self, + coordinator: EndpointAddr, + orchestrator_actor: ActorAddress, + logical_node_id: u64, + stage_index: u32, + ) -> Result { + let provider_name = self.provider.as_str(); + 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_IROH_ENDPOINT_ADDR_MASK_ENV.to_owned(), + self.endpoint_addr_mask.as_str().to_owned(), + ), + ("MVP_NODE_PROVIDER".to_owned(), provider_name.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(), + ), + ]; + env.extend(self.extra_worker_env()); let args = match provider_name { "vastai" => self .vastai @@ -2956,8 +2928,8 @@ fn wait_for_runtime_readies( while let Ok(observation) = obs_rx.try_recv() { emit_plugin_observation(orch_datastream, dashboard, provider, &observation); match observation { - PluginObservation::DatastreamFrame { .. } => {} - PluginObservation::ProviderLine { .. } + PluginObservation::DatastreamFrame { .. } + | PluginObservation::ProviderLine { .. } | PluginObservation::StdoutLine { .. } | PluginObservation::StderrLine { .. } => {} PluginObservation::Failed { reason, .. } => return Err(reason), @@ -3052,7 +3024,12 @@ fn wait_for_weights_loaded_count( } if last_resend.elapsed() >= Duration::from_secs(15) { resend_attempt += 1; - let pending = pending_pipeline_weight_load_stages(pipeline_plan, &loaded_stages); + let mut pending: Vec<&run_plan::StagePlan> = pipeline_plan + .stages + .iter() + .filter(|stage| !loaded_stages.contains(&stage.stage_index)) + .collect(); + pending.sort_by_key(|stage| stage.stage_index); if pending.is_empty() { return Err(format!( "missing unloaded pipeline weight stage; loaded {} of {expected_count}", @@ -3127,8 +3104,8 @@ fn wait_for_weights_loaded_count( ); return Err(reason); } - PluginObservation::DatastreamFrame { .. } => {} - PluginObservation::ProviderLine { .. } + PluginObservation::DatastreamFrame { .. } + | PluginObservation::ProviderLine { .. } | PluginObservation::StdoutLine { .. } | PluginObservation::StderrLine { .. } => {} } @@ -3166,19 +3143,6 @@ fn wait_for_weights_loaded_count( } } -fn pending_pipeline_weight_load_stages<'a>( - pipeline_plan: &'a run_plan::RunPlan, - loaded_stages: &BTreeSet, -) -> Vec<&'a run_plan::StagePlan> { - let mut pending = pipeline_plan - .stages - .iter() - .filter(|stage| !loaded_stages.contains(&stage.stage_index)) - .collect::>(); - pending.sort_by_key(|stage| stage.stage_index); - pending -} - struct PipelineStageProvision<'a> { driver: &'a mut IrohDriver, stack: &'a DistributionRuntimeStack, @@ -3732,7 +3696,7 @@ impl OrchDatastream { producer, channels: BTreeMap::new(), channel_names: BTreeMap::new(), - archive: frame_log.map(FrameArchive::open).transpose()?, + archive: frame_log.map(|p| FrameArchive::open_with_label(p, "datastream frame log")).transpose()? }; for name in [ MVP_PROVISIONING_EVENTS, @@ -4435,10 +4399,6 @@ impl PipelinePromptRuntime { }) } - fn is_active(&self) -> bool { - self.active.is_some() - } - fn note_progress(&mut self) { let now = Instant::now(); self.last_progress_at = Some(now); @@ -5026,7 +4986,7 @@ fn serve_prompts( if active.is_none() && pipeline_runtime .as_ref() - .is_none_or(|pipeline| !pipeline.is_active()) + .is_none_or(|pipeline| pipeline.active.is_none()) && let Ok(work) = work_rx.try_recv() { let request = work.request; @@ -5225,8 +5185,8 @@ fn drain_observations_with_exit( PluginObservation::Exited { node_id, status, .. } => return Err(exit_message(node_id, status)), - PluginObservation::DatastreamFrame { .. } => {} - PluginObservation::ProviderLine { .. } + PluginObservation::DatastreamFrame { .. } + | PluginObservation::ProviderLine { .. } | PluginObservation::StdoutLine { .. } | PluginObservation::StderrLine { .. } => {} } diff --git a/crates/mvp-system/src/orchestration/distribution_stack.rs b/crates/mvp-system/src/orchestration/distribution_stack.rs index b4747db..55898a4 100644 --- a/crates/mvp-system/src/orchestration/distribution_stack.rs +++ b/crates/mvp-system/src/orchestration/distribution_stack.rs @@ -1,4 +1,3 @@ -#![allow(dead_code)] //! MVP-system swactor distribution runtime wiring. //! @@ -58,10 +57,6 @@ pub(crate) struct DistributionRuntimeStack { } impl DistributionRuntimeStack { - pub(crate) fn new(node_id: NodeId, config: DistributedNodeConfig) -> Self { - Self::new_with_codecs(node_id, config, |_| {}) - } - pub(crate) fn new_with_codecs( node_id: NodeId, config: DistributedNodeConfig, @@ -210,16 +205,6 @@ impl DistributionRuntimeStack { .send_to(self.actors.directory, DirectoryIn::Register(entry)); } - pub(crate) fn alive_count(&self) -> usize { - self.membership_mirror - .lock() - .expect("membership mirror poisoned") - .all_members() - .iter() - .filter(|entry| entry.state == MemberState::Alive) - .count() - } - pub(crate) fn member_state(&self, node_id: NodeId) -> Option { self.membership_mirror .lock() diff --git a/crates/mvp-system/src/orchestration/engine_builder/mod.rs b/crates/mvp-system/src/orchestration/engine_builder/mod.rs index bb290a0..f85d250 100644 --- a/crates/mvp-system/src/orchestration/engine_builder/mod.rs +++ b/crates/mvp-system/src/orchestration/engine_builder/mod.rs @@ -1,4 +1,3 @@ -#![allow(dead_code)] //! Pool-based engine/node builder primitives. //! diff --git a/crates/mvp-system/src/orchestration/provider_adapters/relay.rs b/crates/mvp-system/src/orchestration/provider_adapters/relay.rs index 217cade..22a23f0 100644 --- a/crates/mvp-system/src/orchestration/provider_adapters/relay.rs +++ b/crates/mvp-system/src/orchestration/provider_adapters/relay.rs @@ -1,4 +1,3 @@ -#![allow(dead_code)] //! Relay provisioning shims for MVP runtimes. //! @@ -56,27 +55,8 @@ pub(crate) trait RelayProvider: Send { fn provision_relay(&mut self, request: RelayProvisionRequest) -> Result; fn relay_mode(&self, lease: &RelayLease) -> Result; - - fn release_relay(&mut self, _lease: RelayLease) -> Result<(), String> { - Ok(()) - } } -#[derive(Clone, Copy, Debug, Default)] -pub(crate) struct LocalShimRelayProvider; - -impl RelayProvider for LocalShimRelayProvider { - fn provision_relay(&mut self, request: RelayProvisionRequest) -> Result { - Ok(RelayLease { - id: RelayLeaseId(format!("local-shim:{}", request.run_id)), - endpoints: Vec::new(), - }) - } - - fn relay_mode(&self, _lease: &RelayLease) -> Result { - Ok(RelayMode::Disabled) - } -} #[derive(Clone, Debug)] pub(crate) struct StaticRelayProvider { @@ -128,7 +108,7 @@ impl RelayProvider for StaticRelayProvider { } pub(crate) fn relay_runtime_config_from_env(run_id: u64) -> Result { - let mode = relay_mode_setting_from_env(); + let mode = env_optional(MVP_IROH_RELAY_MODE_ENV).map(|value| value.to_ascii_lowercase()); let url = selected_relay_url_from_env(); relay_runtime_config_from_settings(run_id, mode.as_deref(), url.as_deref()) } @@ -183,10 +163,6 @@ fn relay_runtime_config_from_optional_static_provider( }) } -fn relay_mode_setting_from_env() -> Option { - env_optional(MVP_IROH_RELAY_MODE_ENV).map(|value| value.to_ascii_lowercase()) -} - fn env_optional(name: &str) -> Option { std::env::var(name) .ok() diff --git a/crates/mvp-system/src/orchestration/provider_adapters/vastai/mod.rs b/crates/mvp-system/src/orchestration/provider_adapters/vastai/mod.rs index 34c2d22..9abdf53 100644 --- a/crates/mvp-system/src/orchestration/provider_adapters/vastai/mod.rs +++ b/crates/mvp-system/src/orchestration/provider_adapters/vastai/mod.rs @@ -8,7 +8,6 @@ use std::thread::{self, JoinHandle}; use std::time::{Duration, Instant}; use datastream::DatastreamProducer; -use serde::{Deserialize, Serialize}; use swactor::actor::{ActorAddress, ActorInterface}; use swactor::runtime::{Ctx, ExternalSender, Runtime, RuntimeConfig, RuntimeHandle}; use swactor_vastai::{ @@ -517,7 +516,19 @@ impl VastAiLeaseClient for ToolsVastAiLeaseClient { label, lifecycle, ))?; - endpoint_from_parts(contract_id, endpoint.ip, endpoint.port, ssh_user) + let host = endpoint.ip; + let port = endpoint.port; + if host.is_empty() || host == "unknown" { + return Err(format!("vastai contract {contract_id} has no SSH host")); + } + if port == 0 { + return Err(format!("vastai contract {contract_id} has no SSH port")); + } + Ok(VastAiSshEndpoint { + host, + port, + user: ssh_user.to_owned(), + }) } fn spawn_provider_monitor( @@ -556,25 +567,6 @@ impl VastAiLeaseClient for ToolsVastAiLeaseClient { } } -fn endpoint_from_parts( - contract_id: u64, - host: String, - port: u16, - ssh_user: &str, -) -> Result { - if host.is_empty() || host == "unknown" { - return Err(format!("vastai contract {contract_id} has no SSH host")); - } - if port == 0 { - return Err(format!("vastai contract {contract_id} has no SSH port")); - } - Ok(VastAiSshEndpoint { - host, - port, - user: ssh_user.to_owned(), - }) -} - fn provider_terminal_start_error( contract_id: u64, actual: &str, @@ -611,12 +603,6 @@ fn provider_status_message_has_terminal_failure(message: &str) -> bool { }) } -#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] -pub(crate) enum BootstrapStopReason { - RuntimeReady, - NodeStop, -} - pub(crate) trait VastAiBootstrapLauncher: Send { type Handle: Send; @@ -629,7 +615,7 @@ pub(crate) trait VastAiBootstrapLauncher: Send { lifecycle: LifecyclePolicy, ) -> Result; - fn stop_bootstrap(&mut self, handle: &mut Self::Handle, reason: BootstrapStopReason); + fn stop_bootstrap(&mut self, handle: &mut Self::Handle); } #[derive(Clone, Copy, Debug, PartialEq, Eq)] @@ -672,7 +658,6 @@ struct SshBootstrapActor { backoff: Duration, observation_class: Option<&'static str>, stopped: bool, - start_on_boot: bool, } impl SshBootstrapActor { @@ -698,7 +683,6 @@ impl SshBootstrapActor { backoff: Duration::from_secs(1), observation_class: None, stopped: false, - start_on_boot: true, } } @@ -873,7 +857,7 @@ impl SshBootstrapActor { delay.as_secs(), self.attempt )); - self.backoff = next_ssh_backoff(self.backoff); + self.backoff = std::cmp::min(self.backoff.saturating_mul(2), Duration::from_secs(30)); self.attempt = self.attempt.saturating_add(1); schedule_ssh_message( self.sender.clone(), @@ -905,9 +889,7 @@ impl ActorInterface for SshBootstrapActor { type Response = (); fn on_start(&mut self, ctx: &Ctx) { - if self.start_on_boot { - let _ = ctx.send(ctx.self_addr(), SshBootstrapMsg::StartAttempt); - } + let _ = ctx.send(ctx.self_addr(), SshBootstrapMsg::StartAttempt); } fn handle(&mut self, ctx: &Ctx, msg: Self::Incoming) { @@ -996,7 +978,7 @@ impl VastAiBootstrapLauncher for SshCommandBootstrapLauncher { }) } - fn stop_bootstrap(&mut self, handle: &mut Self::Handle, _reason: BootstrapStopReason) { + fn stop_bootstrap(&mut self, handle: &mut Self::Handle) { let _ = handle.runtime.send_to(handle.actor, SshBootstrapMsg::Stop); handle.runtime.tick(); } @@ -1123,8 +1105,12 @@ fn ssh_bootstrap_args( args } -fn next_ssh_backoff(current: Duration) -> Duration { - std::cmp::min(current.saturating_mul(2), Duration::from_secs(30)) +fn emit_node_line(sink: &PluginSink, run_id: u64, node_id: u64, line: impl Into) { + sink.observe(PluginObservation::ProviderLine { + run_id, + node_id, + line: line.into(), + }); } pub(crate) struct VastAiProvisioningPlugin @@ -1253,54 +1239,33 @@ where } let stream_id = node_stream_id(spec.run_id, spec.node_id); let label = self.label_for(&spec); - sink.observe(PluginObservation::ProviderLine { - run_id: spec.run_id, - node_id: spec.node_id, - line: format!("vastai provisioning label={label} stream={stream_id}"), - }); + emit_node_line(&sink, spec.run_id, spec.node_id, format!("vastai provisioning label={label} stream={stream_id}")); let request = self.build_request(&spec, label.clone()); let instance = self.client.provision_one(request).map_err(|e| { classified_start_error(format!("vastai provision node {}: {e}", spec.node_id)) })?; - sink.observe(PluginObservation::ProviderLine { - run_id: spec.run_id, - node_id: spec.node_id, - line: format!( - "vastai contract {} ready for SSH lookup", - instance.contract_id - ), - }); - sink.observe(PluginObservation::ProviderLine { - run_id: spec.run_id, - node_id: spec.node_id, - line: serde_json::json!({ - "type": "VastAiLeaseReady", - "run_id": spec.run_id, - "node_id": spec.node_id, - "label": &label, - "image": &spec.image, - "contract_id": instance.contract_id, - "offer_id": instance.offer_id, - "host_id": instance.host_id, - "gpu_name": &instance.gpu_name, - "gpu_ram": instance.gpu_ram, - "dph_total": instance.dph_total, - }) - .to_string(), - }); - sink.observe(PluginObservation::ProviderLine { - run_id: spec.run_id, - node_id: spec.node_id, - line: serde_json::json!({ - "type": "VastAiSshEndpointDiscoveryStarted", - "run_id": spec.run_id, - "node_id": spec.node_id, - "contract_id": instance.contract_id, - "label": &label, - }) - .to_string(), - }); + emit_node_line(&sink, spec.run_id, spec.node_id, format!("vastai contract {} ready for SSH lookup", instance.contract_id)); + emit_node_line(&sink, spec.run_id, spec.node_id, serde_json::json!({ + "type": "VastAiLeaseReady", + "run_id": spec.run_id, + "node_id": spec.node_id, + "label": &label, + "image": &spec.image, + "contract_id": instance.contract_id, + "offer_id": instance.offer_id, + "host_id": instance.host_id, + "gpu_name": &instance.gpu_name, + "gpu_ram": instance.gpu_ram, + "dph_total": instance.dph_total, + }).to_string()); + emit_node_line(&sink, spec.run_id, spec.node_id, serde_json::json!({ + "type": "VastAiSshEndpointDiscoveryStarted", + "run_id": spec.run_id, + "node_id": spec.node_id, + "contract_id": instance.contract_id, + "label": &label, + }).to_string()); let endpoint = match self.client.ssh_endpoint( instance.contract_id, @@ -1322,35 +1287,25 @@ where )); } }; - sink.observe(PluginObservation::ProviderLine { - run_id: spec.run_id, - node_id: spec.node_id, - line: serde_json::json!({ - "type": "VastAiSshEndpointReady", - "run_id": spec.run_id, - "node_id": spec.node_id, - "contract_id": instance.contract_id, - "host": &endpoint.host, - "port": endpoint.port, - "user": &endpoint.user, - }) - .to_string(), - }); + emit_node_line(&sink, spec.run_id, spec.node_id, serde_json::json!({ + "type": "VastAiSshEndpointReady", + "run_id": spec.run_id, + "node_id": spec.node_id, + "contract_id": instance.contract_id, + "host": &endpoint.host, + "port": endpoint.port, + "user": &endpoint.user, + }).to_string()); - sink.observe(PluginObservation::ProviderLine { - run_id: spec.run_id, - node_id: spec.node_id, - line: serde_json::json!({ - "type": "VastAiBootstrapObservationStarted", - "run_id": spec.run_id, - "node_id": spec.node_id, - "contract_id": instance.contract_id, - "host": &endpoint.host, - "port": endpoint.port, - "user": &endpoint.user, - }) - .to_string(), - }); + emit_node_line(&sink, spec.run_id, spec.node_id, serde_json::json!({ + "type": "VastAiBootstrapObservationStarted", + "run_id": spec.run_id, + "node_id": spec.node_id, + "contract_id": instance.contract_id, + "host": &endpoint.host, + "port": endpoint.port, + "user": &endpoint.user, + }).to_string()); let bootstrap = match self.bootstrap.start_bootstrap( spec.clone(), @@ -1433,11 +1388,7 @@ where } let stream_id = node_stream_id(spec.run_id, spec.node_id); let label = self.label_for(&spec); - sink.observe(PluginObservation::ProviderLine { - run_id: spec.run_id, - node_id: spec.node_id, - line: format!("vastai provisioning label={label} stream={stream_id}"), - }); + emit_node_line(&sink, spec.run_id, spec.node_id, format!("vastai provisioning label={label} stream={stream_id}")); let request = self.build_request(&spec, label.clone()); start_inputs.push((index, spec, label, request)); } @@ -1450,13 +1401,7 @@ where Ok(plan) => plan, Err(error) => { for (_, spec, _, _) in &start_inputs { - sink.observe(PluginObservation::ProviderLine { - run_id: spec.run_id, - node_id: spec.node_id, - line: format!( - "vastai first-wave offer planning failed; falling back to per-node selection: {error}" - ), - }); + emit_node_line(&sink, spec.run_id, spec.node_id, format!("vastai first-wave offer planning failed; falling back to per-node selection: {error}")); } vec![None; start_inputs.len()] } @@ -1467,30 +1412,20 @@ where { request.preferred_offer_id = offer_plan.get(plan_index).copied().flatten(); if let Some(offer_id) = request.preferred_offer_id { - sink.observe(PluginObservation::ProviderLine { - run_id: spec.run_id, - node_id: spec.node_id, - line: serde_json::json!({ - "type": "VastAiFirstWaveOfferPlanned", - "run_id": spec.run_id, - "node_id": spec.node_id, - "label": &label, - "offer_id": offer_id, - }) - .to_string(), - }); + emit_node_line(&sink, spec.run_id, spec.node_id, serde_json::json!({ + "type": "VastAiFirstWaveOfferPlanned", + "run_id": spec.run_id, + "node_id": spec.node_id, + "label": &label, + "offer_id": offer_id, + }).to_string()); } else { - sink.observe(PluginObservation::ProviderLine { - run_id: spec.run_id, - node_id: spec.node_id, - line: serde_json::json!({ - "type": "VastAiFirstWaveOfferPlanUnavailable", - "run_id": spec.run_id, - "node_id": spec.node_id, - "label": &label, - }) - .to_string(), - }); + emit_node_line(&sink, spec.run_id, spec.node_id, serde_json::json!({ + "type": "VastAiFirstWaveOfferPlanUnavailable", + "run_id": spec.run_id, + "node_id": spec.node_id, + "label": &label, + }).to_string()); } let mut client = self.client.clone(); let config = self.config.clone(); @@ -1499,36 +1434,26 @@ where std::thread::spawn(move || { let started = match client.provision_one(request) { Ok(instance) => { - worker_sink.observe(PluginObservation::ProviderLine { - run_id: spec.run_id, - node_id: spec.node_id, - line: serde_json::json!({ - "type": "VastAiLeaseReady", - "run_id": spec.run_id, - "node_id": spec.node_id, - "label": &label, - "image": &spec.image, - "contract_id": instance.contract_id, - "offer_id": instance.offer_id, - "host_id": instance.host_id, - "gpu_name": &instance.gpu_name, - "gpu_ram": instance.gpu_ram, - "dph_total": instance.dph_total, - }) - .to_string(), - }); - worker_sink.observe(PluginObservation::ProviderLine { - run_id: spec.run_id, - node_id: spec.node_id, - line: serde_json::json!({ - "type": "VastAiSshEndpointDiscoveryStarted", - "run_id": spec.run_id, - "node_id": spec.node_id, - "contract_id": instance.contract_id, - "label": &label, - }) - .to_string(), - }); + emit_node_line(&worker_sink, spec.run_id, spec.node_id, serde_json::json!({ + "type": "VastAiLeaseReady", + "run_id": spec.run_id, + "node_id": spec.node_id, + "label": &label, + "image": &spec.image, + "contract_id": instance.contract_id, + "offer_id": instance.offer_id, + "host_id": instance.host_id, + "gpu_name": &instance.gpu_name, + "gpu_ram": instance.gpu_ram, + "dph_total": instance.dph_total, + }).to_string()); + emit_node_line(&worker_sink, spec.run_id, spec.node_id, serde_json::json!({ + "type": "VastAiSshEndpointDiscoveryStarted", + "run_id": spec.run_id, + "node_id": spec.node_id, + "contract_id": instance.contract_id, + "label": &label, + }).to_string()); match client.ssh_endpoint( instance.contract_id, &label, @@ -1575,34 +1500,24 @@ where for (index, spec, started) in completion_rx { match started { Ok(started) => { - sink.observe(PluginObservation::ProviderLine { - run_id: spec.run_id, - node_id: spec.node_id, - line: serde_json::json!({ - "type": "VastAiSshEndpointReady", - "run_id": spec.run_id, - "node_id": spec.node_id, - "contract_id": started.instance.contract_id, - "host": &started.endpoint.host, - "port": started.endpoint.port, - "user": &started.endpoint.user, - }) - .to_string(), - }); - sink.observe(PluginObservation::ProviderLine { - run_id: spec.run_id, - node_id: spec.node_id, - line: serde_json::json!({ - "type": "VastAiBootstrapObservationStarted", - "run_id": spec.run_id, - "node_id": spec.node_id, - "contract_id": started.instance.contract_id, - "host": &started.endpoint.host, - "port": started.endpoint.port, - "user": &started.endpoint.user, - }) - .to_string(), - }); + emit_node_line(&sink, spec.run_id, spec.node_id, serde_json::json!({ + "type": "VastAiSshEndpointReady", + "run_id": spec.run_id, + "node_id": spec.node_id, + "contract_id": started.instance.contract_id, + "host": &started.endpoint.host, + "port": started.endpoint.port, + "user": &started.endpoint.user, + }).to_string()); + emit_node_line(&sink, spec.run_id, spec.node_id, serde_json::json!({ + "type": "VastAiBootstrapObservationStarted", + "run_id": spec.run_id, + "node_id": spec.node_id, + "contract_id": started.instance.contract_id, + "host": &started.endpoint.host, + "port": started.endpoint.port, + "user": &started.endpoint.user, + }).to_string()); let bootstrap = match self.bootstrap.start_bootstrap( spec.clone(), @@ -1695,19 +1610,14 @@ where let Some(node) = self.nodes.get_mut(&handle.id) else { return Ok(()); }; - node.sink.observe(PluginObservation::ProviderLine { - run_id: node.run_id, - node_id: node.node_id, - line: serde_json::json!({ - "type": "VastAiRuntimeReadyAccepted", - "run_id": node.run_id, - "node_id": node.node_id, - "label": &node.label, - "contract_id": node.contract_id, - "classification": "runtime_ready_over_provider_staleness", - }) - .to_string(), - }); + emit_node_line(&node.sink, node.run_id, node.node_id, serde_json::json!({ + "type": "VastAiRuntimeReadyAccepted", + "run_id": node.run_id, + "node_id": node.node_id, + "label": &node.label, + "contract_id": node.contract_id, + "classification": "runtime_ready_over_provider_staleness", + }).to_string()); Ok(()) } @@ -1722,24 +1632,18 @@ where self.leased_host_ids.remove(&host_id); } if let Some(mut bootstrap) = node.bootstrap.take() { - self.bootstrap - .stop_bootstrap(&mut bootstrap, BootstrapStopReason::NodeStop); + self.bootstrap.stop_bootstrap(&mut bootstrap); } let result = self.client.destroy_contract(node.contract_id); - node.sink.observe(PluginObservation::ProviderLine { - run_id: node.run_id, - node_id: node.node_id, - line: serde_json::json!({ - "type": "VastAiContractCleanup", - "run_id": node.run_id, - "node_id": node.node_id, - "label": &node.label, - "contract_id": node.contract_id, - "result": if result.is_ok() { "ok" } else { "failed" }, - "error": result.as_ref().err(), - }) - .to_string(), - }); + emit_node_line(&node.sink, node.run_id, node.node_id, serde_json::json!({ + "type": "VastAiContractCleanup", + "run_id": node.run_id, + "node_id": node.node_id, + "label": &node.label, + "contract_id": node.contract_id, + "result": if result.is_ok() { "ok" } else { "failed" }, + "error": result.as_ref().err(), + }).to_string()); result } } diff --git a/crates/mvp-system/src/orchestration/run_fsm.rs b/crates/mvp-system/src/orchestration/run_fsm.rs index 2999650..35e133c 100644 --- a/crates/mvp-system/src/orchestration/run_fsm.rs +++ b/crates/mvp-system/src/orchestration/run_fsm.rs @@ -1,4 +1,3 @@ -#![allow(dead_code)] #[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] pub(crate) struct RunId(pub(crate) u64); @@ -279,7 +278,13 @@ impl OrchestratorRun { } fn stage_ready(&mut self, run_id: RunId, stage_index: u32) { - if run_id != self.config.run_id || !self.plan_has_stage(stage_index) { + if run_id != self.config.run_id + || !self.plan.as_ref().is_some_and(|plan| { + plan.stages + .iter() + .any(|stage| stage.stage_index == stage_index) + }) + { self.fault(RunFaultReason::UnknownStageReady { stage_index }); return; } @@ -295,7 +300,13 @@ impl OrchestratorRun { if eos { self.complete(); } else if (self.injected_sequences.len() as u64) < self.config.max_tokens { - self.inject_decode(sequence + 1, token_id, sequence); + self.inject(TokenObjectInjection { + sequence: sequence + 1, + payload: TokenObjectPayload::Decode { + token_id, + sampling: SamplingData { source_sequence: sequence }, + }, + }); } else { self.complete(); } @@ -344,30 +355,23 @@ impl OrchestratorRun { if self.terminal || !self.provisioned || !self.injected_sequences.is_empty() { return; } - if self.token_in_ready && self.token_out_ready && self.all_stages_ready() { - self.inject_prompt(0); + if self.token_in_ready + && self.token_out_ready + && self.plan.as_ref().is_some_and(|plan| { + plan.stages + .iter() + .all(|stage| self.ready_stages.contains(&stage.stage_index)) + }) + { + self.inject(TokenObjectInjection { + sequence: 0, + payload: TokenObjectPayload::Prompt { + tokens: self.config.prompt.clone(), + }, + }); } } - fn inject_prompt(&mut self, sequence: u64) { - self.inject(TokenObjectInjection { - sequence, - payload: TokenObjectPayload::Prompt { - tokens: self.config.prompt.clone(), - }, - }); - } - - fn inject_decode(&mut self, sequence: u64, token_id: u32, source_sequence: u64) { - self.inject(TokenObjectInjection { - sequence, - payload: TokenObjectPayload::Decode { - token_id, - sampling: SamplingData { source_sequence }, - }, - }); - } - fn inject(&mut self, object: TokenObjectInjection) { if self.terminal { return; @@ -434,7 +438,11 @@ impl OrchestratorRun { if !self.teardown_started || !self.token_endpoints_stopped { return; } - if !self.all_stages_stopped() { + if !self.plan.as_ref().is_some_and(|plan| { + plan.stages + .iter() + .all(|stage| self.stopped_stages.contains(&stage.stage_index)) + }) { return; } self.mark_torn_down(); @@ -451,28 +459,4 @@ impl OrchestratorRun { }); } } - - fn plan_has_stage(&self, stage_index: u32) -> bool { - self.plan.as_ref().is_some_and(|plan| { - plan.stages - .iter() - .any(|stage| stage.stage_index == stage_index) - }) - } - - fn all_stages_ready(&self) -> bool { - self.plan.as_ref().is_some_and(|plan| { - plan.stages - .iter() - .all(|stage| self.ready_stages.contains(&stage.stage_index)) - }) - } - - fn all_stages_stopped(&self) -> bool { - self.plan.as_ref().is_some_and(|plan| { - plan.stages - .iter() - .all(|stage| self.stopped_stages.contains(&stage.stage_index)) - }) - } } diff --git a/crates/mvp-system/src/orchestration/run_plan.rs b/crates/mvp-system/src/orchestration/run_plan.rs index 672f642..f3c4c0b 100644 --- a/crates/mvp-system/src/orchestration/run_plan.rs +++ b/crates/mvp-system/src/orchestration/run_plan.rs @@ -1,4 +1,3 @@ -#![allow(dead_code)] pub(crate) const MO01_HEADER_BYTES: u64 = 40; const TOKEN_ID_WIDTH_BYTES: u32 = 4; @@ -126,13 +125,11 @@ pub(crate) enum RingDirection { #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub(crate) enum HostPinning { Pageable, - PinnedRequired, } #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub(crate) enum WakeCoalescing { PendingBit, - ReadySet, } #[derive(Clone, Copy, Debug, PartialEq, Eq)] @@ -169,15 +166,12 @@ pub(crate) enum ObjectKind { Token, Activation, Weight, - ModelShard, } #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub(crate) enum ShapeRule { TokenIds, ActivationRows { max_seq_len: u32, hidden_dim: u32 }, - WeightTensor, - ModelShardBytes, } #[derive(Clone, Copy, Debug, PartialEq, Eq)] @@ -298,7 +292,6 @@ pub(crate) enum PlanRejectionKind { DuplicateStageAssignment, MissingStage, InvalidStageCount, - EdgeEndpointMismatch, ModelStageLayoutMismatch, InvalidObjectSpec, UnsupportedShapeOrLayout, @@ -439,7 +432,8 @@ pub(crate) fn plan_run(input: PlannerInput) -> Result { let mut stages = Vec::with_capacity(input.stage_count as usize); for placement in &placements { let stage_index = placement.stage_index; - let (start, end) = layer_range(input.model.num_layers, input.stage_count, stage_index); + let start = (u64::from(input.model.num_layers) * u64::from(stage_index) / u64::from(input.stage_count)) as u32; + let end = (u64::from(input.model.num_layers) * u64::from(stage_index + 1) / u64::from(input.stage_count)) as u32; let inbound_edge = if stage_index == 0 { token_in_edge } else { @@ -510,7 +504,11 @@ pub(crate) fn derive_stage_provision( outbound: OutboundEdgeProvision { edge_id: outbound.edge_id, kind: outbound.kind, - consumer_node_id: endpoint_node_id(&outbound.consumer), + consumer_node_id: match &outbound.consumer { + EdgeEndpoint::Orchestrator { node_id } | EdgeEndpoint::Stage { node_id, .. } => { + *node_id + } + }, object_spec: outbound.object_spec, ring_spec: ring_spec_for_direction(outbound.ring_spec, RingDirection::Egress), }, @@ -607,18 +605,6 @@ fn ring_spec_for_direction(mut spec: RingSpec, direction: RingDirection) -> Ring spec } -fn layer_range(num_layers: u32, stage_count: u32, stage_index: u32) -> (u32, u32) { - 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) -} - -fn endpoint_node_id(endpoint: &EdgeEndpoint) -> NodeId { - match endpoint { - EdgeEndpoint::Orchestrator { node_id } | EdgeEndpoint::Stage { node_id, .. } => *node_id, - } -} - fn valid_ring(spec: RingSpec) -> bool { spec.data_capacity > 0 && spec.alignment > 0 && spec.alignment.is_power_of_two() } diff --git a/crates/mvp-system/src/staging/actor.rs b/crates/mvp-system/src/staging/actor.rs deleted file mode 100644 index d0b1d1f..0000000 --- a/crates/mvp-system/src/staging/actor.rs +++ /dev/null @@ -1,79 +0,0 @@ -use swactor::actor::{ActorAddress, ActorInterface}; -use swactor::runtime::Ctx; - -use crate::staging::control as core; - -#[derive(Clone, Debug, PartialEq, Eq)] -pub(crate) enum StageControllerMsg { - Observe(core::StageEvent), - Snapshot { reply_to: ActorAddress }, -} - -#[derive(Clone, Debug, PartialEq, Eq)] -pub(crate) enum StageControllerReport { - Command(core::StageCommand), - Lifecycle(core::StageLifecycleEvent), - Snapshot { - commands: Vec, - events: Vec, - }, -} - -pub(crate) struct StageControllerActor { - core: core::StageController, - report_to: Option, - command_cursor: usize, - event_cursor: usize, -} - -impl StageControllerActor { - pub(crate) fn new(local_node_id: core::NodeId, report_to: Option) -> Self { - Self { - core: core::StageController::new(local_node_id), - report_to, - command_cursor: 0, - event_cursor: 0, - } - } - - fn drain_outputs(&mut self, ctx: &Ctx) { - let Some(report_to) = self.report_to else { - self.command_cursor = self.core.commands().len(); - self.event_cursor = self.core.events().len(); - return; - }; - - for command in &self.core.commands()[self.command_cursor..] { - let _ = ctx.send(report_to, StageControllerReport::Command(command.clone())); - } - self.command_cursor = self.core.commands().len(); - - for event in &self.core.events()[self.event_cursor..] { - let _ = ctx.send(report_to, StageControllerReport::Lifecycle(event.clone())); - } - self.event_cursor = self.core.events().len(); - } -} - -impl ActorInterface for StageControllerActor { - type Incoming = StageControllerMsg; - type Response = (); - - fn handle(&mut self, ctx: &Ctx, msg: Self::Incoming) { - match msg { - StageControllerMsg::Observe(event) => { - self.core.observe(event); - self.drain_outputs(ctx); - } - StageControllerMsg::Snapshot { reply_to } => { - let _ = ctx.send( - reply_to, - StageControllerReport::Snapshot { - commands: self.core.commands().to_vec(), - events: self.core.events().to_vec(), - }, - ); - } - } - } -} diff --git a/crates/mvp-system/src/staging/control.rs b/crates/mvp-system/src/staging/control.rs index 7aa700f..eb140eb 100644 --- a/crates/mvp-system/src/staging/control.rs +++ b/crates/mvp-system/src/staging/control.rs @@ -1,4 +1,3 @@ -#![allow(dead_code)] use crate::gguf_shard::StageShardPlan; use crate::run_plan::{GgufSource, TokenizerSource}; diff --git a/crates/mvp-system/src/staging/gguf_metadata.rs b/crates/mvp-system/src/staging/gguf_metadata.rs index 142f66c..ea36964 100644 --- a/crates/mvp-system/src/staging/gguf_metadata.rs +++ b/crates/mvp-system/src/staging/gguf_metadata.rs @@ -6,8 +6,8 @@ use std::path::Path; use crate::gguf_common::{GgufValueType, read_integer_value, read_u32, read_u64}; use crate::run_plan::{self, DTypeFamily, GgufSource, TokenizerSource}; -const GGUF_MAGIC: &[u8; 4] = b"GGUF"; -const SUPPORTED_GGUF_VERSION: u32 = 3; +pub(crate) const GGUF_MAGIC: &[u8; 4] = b"GGUF"; +pub(crate) const SUPPORTED_GGUF_VERSION: u32 = 3; const DEFAULT_EFFECTIVE_CONTEXT: u64 = 512; const MAX_METADATA_STRING_BYTES: u64 = 16 * 1024 * 1024; const MAX_METADATA_KEY_BYTES: u64 = 1024 * 1024; @@ -161,7 +161,7 @@ fn required_u32(map: &BTreeMap, key: &str, label: &str) -> Result(reader: &mut R, value_type: GgufValueType) -> Result<(), String> { +pub(crate) fn skip_scalar(reader: &mut R, value_type: GgufValueType) -> Result<(), String> { match value_type { GgufValueType::String => skip_gguf_string(reader), GgufValueType::Array => skip_array(reader), @@ -169,7 +169,7 @@ fn skip_scalar(reader: &mut R, value_type: GgufValueType) -> Res } } -fn skip_array(reader: &mut R) -> Result<(), String> { +pub(crate) fn skip_array(reader: &mut R) -> Result<(), String> { let element_type = GgufValueType::read(reader, "GGUF metadata value type")?; let len = read_u64(reader)?; match element_type { @@ -195,7 +195,7 @@ fn skip_array(reader: &mut R) -> Result<(), String> { } } -fn read_gguf_string(reader: &mut R, max_len: u64) -> Result { +pub(crate) fn read_gguf_string(reader: &mut R, max_len: u64) -> Result { let len = read_u64(reader)?; if len > max_len { return Err(format!( @@ -211,12 +211,12 @@ fn read_gguf_string(reader: &mut R, max_len: u64) -> Result(reader: &mut R) -> Result<(), String> { +pub(crate) fn skip_gguf_string(reader: &mut R) -> Result<(), String> { let len = read_u64(reader)?; skip_bytes(reader, len) } -fn skip_bytes(reader: &mut R, mut bytes: u64) -> Result<(), String> { +pub(crate) fn skip_bytes(reader: &mut R, mut bytes: u64) -> Result<(), String> { while bytes > 0 { let chunk = bytes.min(i64::MAX as u64); reader diff --git a/crates/mvp-system/src/staging/gguf_shard.rs b/crates/mvp-system/src/staging/gguf_shard.rs index 6e8d62f..8d76843 100644 --- a/crates/mvp-system/src/staging/gguf_shard.rs +++ b/crates/mvp-system/src/staging/gguf_shard.rs @@ -1,14 +1,13 @@ use std::fs::File; -use std::io::{Read, Seek, SeekFrom, Write}; +use std::io::{Read, Seek, Write}; use std::path::{Path, PathBuf}; use crate::gguf_common::{GgufValueType, read_integer_value, read_u32, read_u64}; use serde::{Deserialize, Serialize}; use crate::run_plan::GgufSource; +use crate::staging::gguf_metadata::{skip_scalar as skip_value, read_gguf_string, GGUF_MAGIC, SUPPORTED_GGUF_VERSION}; -const GGUF_MAGIC: &[u8; 4] = b"GGUF"; -const SUPPORTED_GGUF_VERSION: u32 = 3; const DEFAULT_ALIGNMENT: u64 = 32; const MAX_STRING_BYTES: u64 = 64 * 1024 * 1024; const STAGE_SHARD_CACHE_FORMAT_VERSION: &str = "stage-shard-cache-v2"; @@ -65,20 +64,9 @@ impl StageShardPlan { format!("{}.stage-{:05}.gguf", self.cache_key, self.stage_index) } - pub(crate) fn source_url(&self) -> Result { - source_url(&self.source) - } - - pub(crate) fn planned_tensor_fetch_bytes(&self) -> u64 { - self.merged_tensor_ranges - .iter() - .map(|range| range.len) - .sum() - } - pub(crate) fn planned_fetch_bytes(&self) -> u64 { self.metadata_end - .saturating_add(self.planned_tensor_fetch_bytes()) + .saturating_add(self.merged_tensor_ranges.iter().map(|r| r.len).sum()) } pub(crate) fn planned_range_count(&self) -> usize { @@ -216,7 +204,7 @@ pub(crate) fn source_url(source: &GgufSource) -> Result { } => Ok(format!( "https://huggingface.co/{repo}/resolve/{}/{}", revision.as_deref().unwrap_or("main"), - encode_hf_path(file) + file.split('/').map(percent_encode_path_segment).collect::>().join("/") )), GgufSource::LocalPath(path) => Err(format!( "stage shard range fetching requires a remote Hugging Face source; got local path {path:?}" @@ -224,13 +212,6 @@ pub(crate) fn source_url(source: &GgufSource) -> Result { } } -fn encode_hf_path(path: &str) -> String { - path.split('/') - .map(percent_encode_path_segment) - .collect::>() - .join("/") -} - fn percent_encode_path_segment(segment: &str) -> String { let mut out = String::new(); for byte in segment.bytes() { @@ -477,7 +458,7 @@ pub(crate) fn materialize_stage_shard_http( where F: FnMut(serde_json::Value), { - let url = plan.source_url()?; + let url = source_url(&plan.source)?; materialize_stage_shard_from_url(plan, &url, output_path, emit) } @@ -799,66 +780,6 @@ fn pad_writer_to_alignment(writer: &mut W, alignment: u64) -> R Ok(()) } -fn skip_value(reader: &mut R, value_type: GgufValueType) -> Result<(), String> { - match value_type { - GgufValueType::String => skip_gguf_string(reader), - GgufValueType::Array => skip_array(reader), - scalar => skip_bytes(reader, scalar.fixed_width().expect("scalar width")), - } -} - -fn skip_array(reader: &mut R) -> Result<(), String> { - let element_type = GgufValueType::read(reader, "GGUF value type")?; - let len = read_u64(reader)?; - match element_type { - GgufValueType::String => { - for _ in 0..len { - skip_gguf_string(reader)?; - } - Ok(()) - } - GgufValueType::Array => { - for _ in 0..len { - skip_array(reader)?; - } - Ok(()) - } - scalar => { - let width = scalar.fixed_width().expect("scalar array width"); - let bytes = width - .checked_mul(len) - .ok_or_else(|| "GGUF array byte count overflow".to_owned())?; - skip_bytes(reader, bytes) - } - } -} - -fn read_gguf_string(reader: &mut R, max_len: u64) -> Result { - let len = read_u64(reader)?; - if len > max_len { - return Err(format!("GGUF string length {len} exceeds {max_len}")); - } - let len = usize::try_from(len).map_err(|_| "GGUF string length exceeds usize".to_owned())?; - let mut bytes = vec![0_u8; len]; - reader - .read_exact(&mut bytes) - .map_err(|e| format!("read GGUF string: {e}"))?; - String::from_utf8(bytes).map_err(|e| format!("GGUF string is not UTF-8: {e}")) -} - -fn skip_gguf_string(reader: &mut R) -> Result<(), String> { - let len = read_u64(reader)?; - skip_bytes(reader, len) -} - -fn skip_bytes(reader: &mut R, bytes: u64) -> Result<(), String> { - let offset = i64::try_from(bytes).map_err(|_| format!("cannot seek over {bytes} bytes"))?; - reader - .seek(SeekFrom::Current(offset)) - .map_err(|e| format!("skip bytes: {e}"))?; - Ok(()) -} - fn align_to(value: u64, alignment: u64) -> Result { if alignment == 0 { return Err("GGUF alignment must be non-zero".to_owned()); diff --git a/crates/mvp-system/src/staging/mod.rs b/crates/mvp-system/src/staging/mod.rs index d7ba27c..d7d2dd4 100644 --- a/crates/mvp-system/src/staging/mod.rs +++ b/crates/mvp-system/src/staging/mod.rs @@ -1,9 +1,6 @@ -#![allow(dead_code)] //! MVP stage control, shard planning, and weight lifecycle public surface. -#[cfg(test)] -pub(crate) mod actor; pub(crate) mod control; pub(crate) mod gguf_metadata; #[cfg(test)]