diff --git a/crates/mvp-system/src/actors/provisioner.rs b/crates/mvp-system/src/actors/provisioner.rs index fa372d2..dc9e8eb 100644 --- a/crates/mvp-system/src/actors/provisioner.rs +++ b/crates/mvp-system/src/actors/provisioner.rs @@ -11,7 +11,7 @@ use crate::provisioning::{ NodeProvisionSpec, PluginNodeHandle, PluginObservation, PluginObservationSink, PluginSink, ProvisionEvent, ProvisionEventKind, ProvisionLogLine, ProvisionLogStream, ProvisionPlugin, }; -use crate::telemetry::{MvpProvisionEventRecord, MvpProvisionLogRecord}; +use crate::telemetry::{MvpProvisionEventRecord, MvpProvisionLogRecord, mvp_provision_log_channel}; use super::codec::JsonCodec; @@ -346,7 +346,10 @@ impl ProvisionerActor

{ fn emit_log(&self, line: ProvisionLogLine) { if let Some(producer) = &self.telemetry { - producer.submit_record(&MvpProvisionLogRecord::new(line)); + let channel = mvp_provision_log_channel(line.node_id, line.stream); + let record = MvpProvisionLogRecord::new(line); + let payload = serde_json::to_vec(&record).expect("serialize provisioning log record"); + producer.submit_bytes(channel, payload); } } } diff --git a/crates/mvp-system/src/bin/local_e2e_cluster.rs b/crates/mvp-system/src/bin/local_e2e_cluster.rs index 6ba5f3b..23b194d 100644 --- a/crates/mvp-system/src/bin/local_e2e_cluster.rs +++ b/crates/mvp-system/src/bin/local_e2e_cluster.rs @@ -1,4 +1,4 @@ -use std::collections::{HashMap, VecDeque}; +use std::collections::{BTreeMap, HashMap, VecDeque}; use std::io::{BufRead, BufReader, Write}; use std::path::PathBuf; use std::process::{Child, ChildStdin, ChildStdout, Command, ExitCode, Stdio}; @@ -24,19 +24,20 @@ use mvp_system::actors::orchestrator::{ LifecycleEventWire, OrchestratorActor, OrchestratorMsg, OrchestratorReport, RunCommandWire, StageRefWire, }; -use mvp_system::actors::provisioner::{ProvisionerActor, ProvisionerMsg, ProvisionerReport}; use mvp_system::actors::register_mvp_actor_codecs; use mvp_system::arena_manager as arena; use mvp_system::dashboard_view::MvpClusterDashboardView; use mvp_system::distribution_stack::DistributionRuntimeStack; +use mvp_system::docker_cluster_provisioning as docker_provision; use mvp_system::driver_pumps as driver_model; use mvp_system::edge_establisher as edge; use mvp_system::engine_builder as engine; use mvp_system::gpu_worker_ctl as worker_ctl; use mvp_system::gpu_worker_ingress_parser as ingress; +use mvp_system::node_provisioning as node_provision; use mvp_system::observability_surface as obs; use mvp_system::orchestrator_run_fsm as fsm; -use mvp_system::provisioning::{LocalDockerPlugin, NodeProvisionSpec, ProvisionLogStream}; +use mvp_system::provisioning::{NodeProvisionSpec, ProvisionLogStream}; use mvp_system::run_plan as plan; use mvp_system::stage_controller as stage; use mvp_system::tx_rx_edge_actor as edge_actor; @@ -58,6 +59,7 @@ const OBJECT_ALIGNMENT: u64 = 4; const ARENA_BYTES: usize = 16 * 1024; const RING_BYTES: usize = 4096; const EDGE_READY_TIMEOUT: Duration = Duration::from_secs(20); +const DEFAULT_RUNTIME_SNAPSHOT_INTERVAL: Duration = Duration::from_millis(500); struct MvpDashboard { url: String, @@ -66,6 +68,8 @@ struct MvpDashboard { subscription: DatastreamSubscription, handle: dashboard::DashboardHandle, runtime_position: u64, + runtime_snapshot_interval: Duration, + last_runtime_snapshot: Option, } impl MvpDashboard { @@ -78,6 +82,7 @@ impl MvpDashboard { .map_err(|e| format!("invalid MVP_DASHBOARD_PORT: {e}"))?; config.port = port; } + let runtime_snapshot_interval = runtime_snapshot_interval_from_env()?; let url = format!("http://127.0.0.1:{}/view/datastream/live", config.port); let handle = dashboard::start_dashboard(config.clone()); handle.register_view(Arc::new(MvpClusterDashboardView::new())); @@ -96,15 +101,14 @@ impl MvpDashboard { subscription, handle, runtime_position: 0, + runtime_snapshot_interval, + last_runtime_snapshot: None, }) } fn url(&self) -> &str { &self.url } - fn producer(&self) -> DatastreamProducer { - self.producer.clone() - } fn drain(&mut self) { self.endpoint.tick(); @@ -119,6 +123,23 @@ impl MvpDashboard { self.drain(); } + fn record_provision_log(&mut self, line: mvp_system::provisioning::ProvisionLogLine) { + let channel = mvp_system::telemetry::mvp_provision_log_channel(line.node_id, line.stream); + let record = mvp_system::telemetry::MvpProvisionLogRecord::new(line); + let payload = serde_json::to_vec(&record).expect("serialize provisioning log record"); + self.producer.submit_bytes(channel, payload); + self.drain(); + } + + fn publish_runtime_snapshot_throttled(&mut self, stack: &DistributionRuntimeStack) { + let due = self.last_runtime_snapshot.map_or(true, |last| { + last.elapsed() >= self.runtime_snapshot_interval + }); + if self.runtime_snapshot_interval.is_zero() || due { + self.publish_runtime_snapshot(stack); + } + } + fn publish_runtime_snapshot(&mut self, stack: &DistributionRuntimeStack) { let stats = stack.runtime.stats(); let actors = stats @@ -149,6 +170,7 @@ impl MvpDashboard { if !stats.actor_details.is_empty() { self.ingest_runtime_json(RUNTIME_ACTORS, json!({ "actors": &stats.actor_details })); } + self.last_runtime_snapshot = Some(Instant::now()); } fn ingest_runtime_json(&mut self, channel: &str, value: Value) { @@ -163,6 +185,17 @@ impl MvpDashboard { } } +fn runtime_snapshot_interval_from_env() -> Result { + let Some(value) = std::env::var_os("MVP_RUNTIME_SNAPSHOT_MS") else { + return Ok(DEFAULT_RUNTIME_SNAPSHOT_INTERVAL); + }; + let millis = value + .to_string_lossy() + .parse::() + .map_err(|e| format!("invalid MVP_RUNTIME_SNAPSHOT_MS: {e}"))?; + Ok(Duration::from_millis(millis)) +} + fn main() -> ExitCode { let args = std::env::args().collect::>(); let result = if args.iter().any(|arg| arg == "--role=node") { @@ -191,14 +224,25 @@ struct NodeStdoutLine { kind: String, stage_index: Option, event: Option, + endpoint: Option, + node_actor: Option, + logical_node_id: Option, } -#[derive(Clone, Debug)] + +type LocalDockerNodeProvisioner = docker_provision::DockerNodeProvisioner< + LocalE2eDockerCli, + LocalE2eBootstrapFactory, + LocalE2eBootstrapDatastream, +>; + struct ProvisionedDockerNode { node_id: u64, stage_index: u32, endpoint: EndpointAddr, node_actor: ActorAddress, provider_process_id: Option, + provisioner: LocalDockerNodeProvisioner, + events: Receiver, } #[derive(Default)] @@ -208,6 +252,13 @@ struct ProvisionStats { stderr_line_count: usize, } +#[derive(Clone, Debug)] +enum LocalDockerNodeEvent { + Stdout(String), + Stderr(String), + Exited(Option), +} + #[derive(Clone, Debug)] enum DriverIngressEvent { StreamArrived { @@ -275,6 +326,23 @@ fn record_dashboard_event(dashboard: &mut Option<&mut MvpDashboard>, event: obs: dashboard.record_event(event); } } +fn record_dashboard_provision_log( + dashboard: &mut Option<&mut MvpDashboard>, + run_id: u64, + node_id: u64, + stream: ProvisionLogStream, + line: &str, +) { + if let Some(dashboard) = dashboard.as_deref_mut() { + dashboard.record_provision_log(mvp_system::provisioning::ProvisionLogLine { + run_id, + node_id, + stream, + line: line.to_owned(), + }); + } +} + fn drain_dashboard(dashboard: &mut Option<&mut MvpDashboard>) { if let Some(dashboard) = dashboard.as_deref_mut() { dashboard.drain(); @@ -290,6 +358,15 @@ fn publish_runtime_snapshot( } } +fn publish_runtime_snapshot_throttled( + dashboard: &mut Option<&mut MvpDashboard>, + stack: &DistributionRuntimeStack, +) { + if let Some(dashboard) = dashboard.as_deref_mut() { + dashboard.publish_runtime_snapshot_throttled(stack); + } +} + fn run_event(run_id: u64, kind: obs::EventKind) -> obs::Event { obs::Event::RunScoped { kind, @@ -398,18 +475,6 @@ fn run_supervisor_once( )) .map_err(|e| format!("spawn orchestrator actor: {e}"))?; stack.register_local_actor(driver.register_actor(orchestrator_addr, 1)); - let provisioner_report = stack - .runtime - .new_inbox::() - .map_err(|e| format!("provisioner report inbox: {e}"))?; - let provisioner_addr = stack - .runtime - .spawn(ProvisionerActor::new( - LocalDockerPlugin::new(format!("mvp-local-e2e-cluster-{}", std::process::id())), - stack.runtime.create_sender(), - dashboard.as_ref().map(|dashboard| dashboard.producer()), - )) - .map_err(|e| format!("spawn provisioner actor: {e}"))?; publish_runtime_snapshot(&mut dashboard, &stack); let mut provision_stats = ProvisionStats::default(); @@ -433,9 +498,7 @@ fn run_supervisor_once( let orchestrator_actor_json = serde_json::to_string(&orchestrator_addr) .map_err(|e| format!("serialize orchestrator actor: {e}"))?; - let node1 = provision_local_docker_node( - provisioner_addr, - &provisioner_report, + let mut node1 = provision_local_docker_node( &mut driver, &stack, &mut dashboard, @@ -451,9 +514,7 @@ fn run_supervisor_once( )?; let node1_endpoint_json = serde_json::to_string(&node1.endpoint) .map_err(|e| format!("serialize node1 endpoint: {e}"))?; - let node0 = provision_local_docker_node( - provisioner_addr, - &provisioner_report, + let mut node0 = match provision_local_docker_node( &mut driver, &stack, &mut dashboard, @@ -466,7 +527,13 @@ fn run_supervisor_once( &self_endpoint_json, &orchestrator_actor_json, ), - )?; + ) { + Ok(node) => node, + Err(error) => { + let _ = stop_provisioned_nodes(&mut [&mut node1], &mut driver, &stack, &mut dashboard); + return Err(error); + } + }; for stage in [&stage1, &stage0] { record_dashboard_event( @@ -619,14 +686,14 @@ fn run_supervisor_once( } } } - stage_ready_count += drain_provisioner_reports( - &provisioner_report, + stage_ready_count += drain_provisioned_node_events( + &mut [&mut node0, &mut node1], run_id, &mut dashboard, &mut provision_stats, )?; drain_dashboard(&mut dashboard); - publish_runtime_snapshot(&mut dashboard, &stack); + publish_runtime_snapshot_throttled(&mut dashboard, &stack); while let Some(report) = orchestrator_report.try_recv() { match report { @@ -739,13 +806,10 @@ fn run_supervisor_once( && sent_stop_to_node1 { stop_provisioned_nodes( - run_id, - provisioner_addr, - &provisioner_report, + &mut [&mut node0, &mut node1], &mut driver, &stack, &mut dashboard, - &mut provision_stats, )?; let builder_stage_assignments = topology .events @@ -816,13 +880,10 @@ fn run_supervisor_once( record_dashboard_event(&mut dashboard, run_fault_event(run_id)); let _ = stop_provisioned_nodes( - run_id, - provisioner_addr, - &provisioner_report, + &mut [&mut node0, &mut node1], &mut driver, &stack, &mut dashboard, - &mut provision_stats, ); Err(format!( "timed out: injected={injected} token_received={token_received} completed={completed} torn_down={torn_down} stop0={sent_stop_to_node0} stop1={sent_stop_to_node1}" @@ -2140,9 +2201,296 @@ fn local_docker_spec( } } +struct LocalE2eContainer { + container_name: String, + stdin: ChildStdin, +} + +struct LocalE2eDockerCli { + spec: NodeProvisionSpec, + events: Sender, + nodes: HashMap, + provider_process_id: Option, +} + +impl LocalE2eDockerCli { + fn new(spec: NodeProvisionSpec, events: Sender) -> Self { + Self { + spec, + events, + nodes: HashMap::new(), + provider_process_id: None, + } + } + + fn provider_process_id(&self) -> Option { + self.provider_process_id + } +} + +impl docker_provision::DockerCli for LocalE2eDockerCli { + fn run_container( + &mut self, + request: docker_provision::DockerRunRequest, + ) -> Result { + let mut env = request.env.clone(); + for (key, value) in &self.spec.env { + env.insert(key.clone(), value.clone()); + } + + let mut command = Command::new("docker"); + command + .arg("run") + .arg("--rm") + .arg("--add-host") + .arg("host.docker.internal:host-gateway") + .arg("--name") + .arg(&request.container_name) + .arg("-i"); + for (key, value) in &request.labels { + command.arg("--label").arg(format!("{key}={value}")); + } + for (key, value) in &env { + command.arg("-e").arg(format!("{key}={value}")); + } + command.arg(&self.spec.image); + for arg in &self.spec.args { + command.arg(arg); + } + + let mut child = command + .stdin(Stdio::piped()) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .spawn() + .map_err(|e| { + docker_provision::DockerCliError::new(format!( + "spawn Docker node {}: {e}", + self.spec.node_id + )) + })?; + + let provider_process_id = child.id(); + self.provider_process_id = Some(provider_process_id); + let stdin = child.stdin.take().ok_or_else(|| { + docker_provision::DockerCliError::new(format!( + "Docker node {} stdin missing", + self.spec.node_id + )) + })?; + let stdout = child.stdout.take().ok_or_else(|| { + docker_provision::DockerCliError::new(format!( + "Docker node {} stdout missing", + self.spec.node_id + )) + })?; + let stderr = child.stderr.take().ok_or_else(|| { + docker_provision::DockerCliError::new(format!( + "Docker node {} stderr missing", + self.spec.node_id + )) + })?; + + self.nodes.insert( + request.container_name.clone(), + LocalE2eContainer { + container_name: request.container_name.clone(), + stdin, + }, + ); + + spawn_local_docker_reader(stdout, self.events.clone(), true); + spawn_local_docker_reader(stderr, self.events.clone(), false); + let events = self.events.clone(); + thread::spawn(move || match child.wait() { + Ok(status) => { + let _ = events.send(LocalDockerNodeEvent::Exited(status.code())); + } + Err(error) => { + let _ = events.send(LocalDockerNodeEvent::Stderr(format!( + "wait Docker node: {error}" + ))); + let _ = events.send(LocalDockerNodeEvent::Exited(None)); + } + }); + + Ok(docker_provision::DockerRunResult { + container_id: request.container_name, + }) + } + + fn inspect_ssh_endpoint( + &mut self, + container_id: &str, + ) -> Result, docker_provision::DockerCliError> { + Ok(Some(node_provision::SshEndpoint { + host: "127.0.0.1".to_owned(), + port: 0, + user: "local-e2e".to_owned(), + auth_ref: format!("local-docker:{container_id}"), + })) + } + + fn remove_force(&mut self, container_id: &str) -> Result<(), docker_provision::DockerCliError> { + let Some(mut node) = self.nodes.remove(container_id) else { + return Ok(()); + }; + let _ = writeln!(node.stdin, "shutdown"); + let _ = node.stdin.flush(); + let status = Command::new("docker") + .arg("stop") + .arg("-t") + .arg("2") + .arg(&node.container_name) + .stdout(Stdio::null()) + .stderr(Stdio::null()) + .status() + .map_err(|e| { + docker_provision::DockerCliError::new(format!( + "docker stop {}: {e}", + node.container_name + )) + })?; + if status.success() { + Ok(()) + } else { + Err(docker_provision::DockerCliError::new(format!( + "docker stop {} exited with {status}", + node.container_name + ))) + } + } +} + +#[derive(Default)] +struct LocalE2eBootstrapFactory; + +impl docker_provision::SshBootstrapClientFactory for LocalE2eBootstrapFactory { + type Client = LocalE2eBootstrapClient; + + fn client_for(&mut self, _spec: &node_provision::BootstrapSessionSpec) -> Self::Client { + LocalE2eBootstrapClient + } +} + +struct LocalE2eBootstrapClient; + +impl docker_provision::BootstrapSshClient for LocalE2eBootstrapClient { + fn connect( + &mut self, + _endpoint: &node_provision::SshEndpoint, + ) -> Result<(), docker_provision::BootstrapSshError> { + Ok(()) + } + + fn probe_stdout(&mut self) -> Result<(), docker_provision::BootstrapSshError> { + Ok(()) + } + + fn read_bootstrap_logs( + &mut self, + _stdout_sources: &[String], + _stderr_sources: &[String], + ) -> Result< + Vec<(node_provision::BootstrapLogStream, String)>, + docker_provision::BootstrapSshError, + > { + Ok(Vec::new()) + } + + fn run_verify_commands( + &mut self, + _commands: &[String], + ) -> Result<(), docker_provision::BootstrapSshError> { + Ok(()) + } + + fn start_swactor( + &mut self, + _command: &str, + _join: &node_provision::SwarmJoinSpec, + ) -> Result<(), docker_provision::BootstrapSshError> { + Ok(()) + } + + fn close(&mut self) {} +} + +#[derive(Default)] +struct LocalE2eBootstrapDatastream; + +impl node_provision::BootstrapDatastreamSink for LocalE2eBootstrapDatastream { + fn record(&mut self, _record: node_provision::BootstrapLogRecord) {} + + fn flush(&mut self) {} +} + +fn spawn_local_docker_reader( + stream: impl std::io::Read + Send + 'static, + events: Sender, + stdout: bool, +) { + thread::spawn(move || { + let reader = BufReader::new(stream); + for next in reader.lines() { + let Ok(line) = next else { + break; + }; + let event = if stdout { + LocalDockerNodeEvent::Stdout(line) + } else { + LocalDockerNodeEvent::Stderr(line) + }; + if events.send(event).is_err() { + break; + } + } + }); +} + +fn logical_node_spec_from_local(spec: &NodeProvisionSpec) -> node_provision::LogicalNodeSpec { + let logical_node_id = node_provision::LogicalNodeId(spec.node_id.to_string()); + node_provision::LogicalNodeSpec { + run_id: node_provision::RunId(spec.run_id), + logical_node_id: logical_node_id.clone(), + group_id: node_provision::NodeGroupId("local-e2e".to_owned()), + role: node_provision::RoleId("stage-worker".to_owned()), + provider: node_provision::ProviderKind::Docker, + shape: node_provision::DesiredNodeShape { + image: spec.image.clone(), + disk_gb: 0, + gpu_name: None, + min_gpu_ram_mb: None, + min_down_mbps: None, + min_up_mbps: None, + min_reliability: None, + require_verified: false, + provider_labels: BTreeMap::from([ + ("mvp.local_e2e".to_owned(), "true".to_owned()), + ("mvp.node_id".to_owned(), spec.node_id.to_string()), + ]), + }, + boot: node_provision::BootSpec { + ssh_user: "local-e2e".to_owned(), + verify_commands: Vec::new(), + start_swactor_command: spec.args.join(" "), + stdout_sources: Vec::new(), + stderr_sources: Vec::new(), + timeout_policy: node_provision::BootstrapTimeoutPolicy { + ssh_connect_secs: 1, + boot_check_secs: 1, + swactor_join_secs: 30, + }, + }, + swarm_join: node_provision::SwarmJoinSpec { + orch_swactor_addr: "local-e2e".to_owned(), + join_token_ref: "local-e2e".to_owned(), + expected_logical_node_id: logical_node_id, + }, + } +} + fn provision_local_docker_node( - provisioner_addr: ActorAddress, - provisioner_report: &swactor::runtime::Inbox, driver: &mut IrohDriver, stack: &DistributionRuntimeStack, dashboard: &mut Option<&mut MvpDashboard>, @@ -2151,50 +2499,71 @@ fn provision_local_docker_node( ) -> Result { let expected_run_id = spec.run_id; let expected_node_id = spec.node_id; - stack - .runtime - .send_to( - provisioner_addr, - ProvisionerMsg::StartNodes { - nodes: vec![spec], - reply_to: *provisioner_report.addr(), - }, - ) - .map_err(|e| format!("start provisioning node {expected_node_id}: {e}"))?; + let expected_stage_index = spec.stage_index.unwrap_or_default(); + let logical_spec = logical_node_spec_from_local(&spec); + let (events_tx, events_rx) = mpsc::channel(); + let cli = LocalE2eDockerCli::new(spec, events_tx); + let provider = docker_provision::DockerProvider::new(cli); + let mut provisioner = docker_provision::DockerNodeProvisioner::new( + provider, + LocalE2eBootstrapFactory, + LocalE2eBootstrapDatastream, + ); + provisioner + .start(logical_spec) + .map_err(|e| format!("start Docker node {expected_node_id}: {e:?}"))?; + let provider_process_id = provisioner.provider().cli().provider_process_id(); let started = Instant::now(); while started.elapsed() < Duration::from_secs(30) { pump_network(driver, stack); drain_dashboard(dashboard); - publish_runtime_snapshot(dashboard, stack); - while let Some(report) = provisioner_report.try_recv() { - match report { - ProvisionerReport::NodeLive { - run_id, - node_id, - stage_index, - endpoint, - node_actor, - provider_process_id, - } if run_id == expected_run_id && node_id == expected_node_id => { - stats.node_live_count += 1; - return Ok(ProvisionedDockerNode { - node_id, - stage_index: stage_index.unwrap_or_default(), - endpoint, - node_actor, - provider_process_id, - }); + publish_runtime_snapshot_throttled(dashboard, stack); + while let Ok(event) = events_rx.try_recv() { + match event { + LocalDockerNodeEvent::Stdout(line) => { + handle_local_node_stdout( + expected_run_id, + expected_node_id, + &line, + dashboard, + stats, + ); + if let Some((endpoint, node_actor, stage_index)) = + ready_node_from_stdout(expected_node_id, expected_stage_index, &line)? + { + provisioner + .observe_swactor_join(node_provision::SwactorId(format!( + "local-e2e-node-{expected_node_id}" + ))) + .map_err(|e| { + format!("complete Docker node {expected_node_id} handoff: {e:?}") + })?; + stats.node_live_count += 1; + return Ok(ProvisionedDockerNode { + node_id: expected_node_id, + stage_index, + endpoint, + node_actor, + provider_process_id, + provisioner, + events: events_rx, + }); + } } - ProvisionerReport::NodeFailed { - run_id, - node_id, - reason, - } if run_id == expected_run_id && node_id == expected_node_id => { - return Err(format!("node {node_id} provisioning failed: {reason}")); + LocalDockerNodeEvent::Stderr(line) => { + handle_local_node_stderr( + expected_run_id, + expected_node_id, + &line, + dashboard, + stats, + ); } - other => { - let _ = handle_provisioner_report(other, expected_run_id, dashboard, stats)?; + LocalDockerNodeEvent::Exited(status) => { + return Err(format!( + "node {expected_node_id} process exited before ready: {status:?}" + )); } } } @@ -2203,49 +2572,83 @@ fn provision_local_docker_node( Err(format!("timed out provisioning node {expected_node_id}")) } -fn drain_provisioner_reports( - provisioner_report: &swactor::runtime::Inbox, +fn ready_node_from_stdout( + expected_node_id: u64, + default_stage_index: u32, + line: &str, +) -> Result, String> { + let Ok(line) = serde_json::from_str::(line) else { + return Ok(None); + }; + if line.kind != "ready" { + return Ok(None); + } + if line.logical_node_id != Some(expected_node_id) { + return Ok(None); + } + let endpoint = line + .endpoint + .ok_or_else(|| format!("ready line for node {expected_node_id} missing endpoint"))?; + let node_actor = line + .node_actor + .ok_or_else(|| format!("ready line for node {expected_node_id} missing node actor"))?; + Ok(Some(( + endpoint, + node_actor, + line.stage_index.unwrap_or(default_stage_index), + ))) +} + +fn drain_provisioned_node_events( + nodes: &mut [&mut ProvisionedDockerNode], run_id: u64, dashboard: &mut Option<&mut MvpDashboard>, stats: &mut ProvisionStats, ) -> Result { let mut stage_ready_count = 0usize; - while let Some(report) = provisioner_report.try_recv() { - stage_ready_count += handle_provisioner_report(report, run_id, dashboard, stats)?; + for node in nodes.iter_mut() { + while let Ok(event) = node.events.try_recv() { + match event { + LocalDockerNodeEvent::Stdout(line) => { + stage_ready_count += + handle_local_node_stdout(run_id, node.node_id, &line, dashboard, stats); + } + LocalDockerNodeEvent::Stderr(line) => { + handle_local_node_stderr(run_id, node.node_id, &line, dashboard, stats); + } + LocalDockerNodeEvent::Exited(status) => { + return Err(format!( + "node {} process exited before stop: {status:?}", + node.node_id + )); + } + } + } } Ok(stage_ready_count) } -fn handle_provisioner_report( - report: ProvisionerReport, +fn handle_local_node_stdout( run_id: u64, + node_id: u64, + line: &str, dashboard: &mut Option<&mut MvpDashboard>, stats: &mut ProvisionStats, -) -> Result { - match report { - ProvisionerReport::LogLine { - run_id: line_run_id, - stream, - line, - .. - } if line_run_id == run_id => { - match stream { - ProvisionLogStream::Stdout => stats.stdout_line_count += 1, - ProvisionLogStream::Stderr => stats.stderr_line_count += 1, - ProvisionLogStream::Provider => {} - } - Ok(record_stage_ready_from_stdout( - run_id, stream, &line, dashboard, - )) - } - ProvisionerReport::NodeFailed { - run_id: failed_run_id, - node_id, - reason, - } if failed_run_id == run_id => Err(format!("node {node_id} failed: {reason}")), - ProvisionerReport::NodeLive { .. } | ProvisionerReport::NodesStopped { .. } => Ok(0), - ProvisionerReport::LogLine { .. } | ProvisionerReport::NodeFailed { .. } => Ok(0), - } +) -> usize { + stats.stdout_line_count += 1; + record_dashboard_provision_log(dashboard, run_id, node_id, ProvisionLogStream::Stdout, line); + record_stage_ready_from_stdout(run_id, ProvisionLogStream::Stdout, line, dashboard) +} + +fn handle_local_node_stderr( + run_id: u64, + node_id: u64, + line: &str, + dashboard: &mut Option<&mut MvpDashboard>, + stats: &mut ProvisionStats, +) { + stats.stderr_line_count += 1; + record_dashboard_provision_log(dashboard, run_id, node_id, ProvisionLogStream::Stderr, line); } fn record_stage_ready_from_stdout( @@ -2282,42 +2685,20 @@ fn record_stage_ready_from_stdout( } fn stop_provisioned_nodes( - run_id: u64, - provisioner_addr: ActorAddress, - provisioner_report: &swactor::runtime::Inbox, + nodes: &mut [&mut ProvisionedDockerNode], driver: &mut IrohDriver, stack: &DistributionRuntimeStack, dashboard: &mut Option<&mut MvpDashboard>, - stats: &mut ProvisionStats, ) -> Result<(), String> { - stack - .runtime - .send_to( - provisioner_addr, - ProvisionerMsg::StopNodes { - run_id, - reply_to: *provisioner_report.addr(), - }, - ) - .map_err(|e| format!("stop provisioned nodes: {e}"))?; - let started = Instant::now(); - while started.elapsed() < Duration::from_secs(10) { - pump_network(driver, stack); - drain_dashboard(dashboard); - publish_runtime_snapshot(dashboard, stack); - while let Some(report) = provisioner_report.try_recv() { - match report { - ProvisionerReport::NodesStopped { - run_id: stopped_run_id, - } if stopped_run_id == run_id => return Ok(()), - other => { - let _ = handle_provisioner_report(other, run_id, dashboard, stats)?; - } - } - } - thread::sleep(Duration::from_millis(10)); + for node in nodes.iter_mut().rev() { + node.provisioner + .stop() + .map_err(|e| format!("stop node {}: {e:?}", node.node_id))?; } - Err("timed out stopping provisioned nodes".to_owned()) + pump_network(driver, stack); + drain_dashboard(dashboard); + publish_runtime_snapshot(dashboard, stack); + Ok(()) } fn spawn_send_pump( diff --git a/crates/mvp-system/src/dashboard_view.rs b/crates/mvp-system/src/dashboard_view.rs index 88270ad..e6a5ec2 100644 --- a/crates/mvp-system/src/dashboard_view.rs +++ b/crates/mvp-system/src/dashboard_view.rs @@ -17,11 +17,8 @@ use crate::telemetry::{ MvpProvisionEventRecord, MvpProvisionLogRecord, }; -const CHANNELS: &[&str] = &[ - MVP_LIFECYCLE, - MVP_PROVISIONING_EVENTS, - MVP_PROVISIONING_LOGS, -]; +const CHANNELS: &[&str] = &[]; +const PROVISIONING_LOG_PREFIX: &str = "mvp.provisioning.logs.node."; const EVENT_LOG_CAP: usize = 256; const LOG_TAIL_CAP: usize = 128; @@ -141,7 +138,10 @@ impl DashboardView for MvpClusterDashboardView { state.apply_provision_event(frame, record); } } - MVP_PROVISIONING_LOGS => { + channel + if channel == MVP_PROVISIONING_LOGS + || channel.starts_with(PROVISIONING_LOG_PREFIX) => + { if let Ok(record) = MvpProvisionLogRecord::decode(&frame.payload) { state.apply_provision_log(record); } diff --git a/crates/mvp-system/src/docker_cluster_provisioning.rs b/crates/mvp-system/src/docker_cluster_provisioning.rs index bac478e..b766cb1 100644 --- a/crates/mvp-system/src/docker_cluster_provisioning.rs +++ b/crates/mvp-system/src/docker_cluster_provisioning.rs @@ -1,11 +1,12 @@ -//! Docker-backed provider and bootstrap wiring for node provisioning tests. +//! Docker-backed provider and single-node bootstrap wiring. //! //! Docker is treated as a concrete provider adapter here: it creates and destroys -//! real Docker container leases through a `DockerCli` boundary. Unit tests use a -//! deterministic CLI implementation, but the provider behavior remains the same -//! provider contract as a remote adapter. +//! real Docker container leases through a `DockerCli` boundary. Higher-level +//! callers may start many nodes, but this module intentionally exposes only +//! per-node ownership primitives so deployment code does not grow a cluster +//! provisioning layer. -use std::collections::{BTreeMap, BTreeSet}; +use std::collections::BTreeMap; use std::time::SystemTime; use crate::node_provisioning::{ @@ -13,8 +14,7 @@ use crate::node_provisioning::{ BootstrapObservation, BootstrapSessionEvent, BootstrapSessionSpec, BootstrapStage, CreateLeaseRequest, CreateLeaseResult, DesiredNodeShape, DestroyHandle, LeaseFacts, LogicalNodeId, NodeManager, NodeManagerCommand, NodeManagerMsg, NodeRecord, ProviderError, - ProviderKind, ProviderLeaseId, ProviderPlugin, RunId, RunNodeGroupSpec, SshEndpoint, SwactorId, - SwarmJoinSpec, expand_node_group, + ProviderKind, ProviderLeaseId, ProviderPlugin, RunId, SshEndpoint, SwactorId, SwarmJoinSpec, }; #[derive(Clone, Debug, PartialEq, Eq)] @@ -391,65 +391,23 @@ impl SshBootstrapSession { } } -#[derive(Clone, Debug, PartialEq, Eq)] -pub struct SwactorJoinEvent { - pub logical_node_id: LogicalNodeId, - pub swactor_id: SwactorId, -} - -#[derive(Clone, Debug, Default)] -pub struct SwactorJoinRouter { - expected_nodes: BTreeSet, -} - -impl SwactorJoinRouter { - pub fn new() -> Self { - Self::default() - } - - pub fn register(&mut self, logical_node_id: LogicalNodeId) { - self.expected_nodes.insert(logical_node_id); - } - - pub fn route( - &self, - event: SwactorJoinEvent, - manager: &mut NodeManager, - ) -> Result, DockerClusterError> { - if !self.expected_nodes.contains(&event.logical_node_id) { - return Err(DockerClusterError::UnknownNode(event.logical_node_id)); - } - manager - .handle(NodeManagerMsg::SwactorJoined { - logical_node_id: event.logical_node_id, - swactor_id: event.swactor_id, - }) - .map_err(|error| DockerClusterError::Node(error.reason)) - } -} - pub trait SshBootstrapClientFactory { type Client: BootstrapSshClient; fn client_for(&mut self, spec: &BootstrapSessionSpec) -> Self::Client; } -#[derive(Clone, Debug)] -pub struct ManagedDockerNode { - pub manager: NodeManager, - pub bootstrap: Option>, -} - #[derive(Clone, Debug, PartialEq, Eq)] -pub enum DockerClusterError { +pub enum DockerNodeProvisionError { Provider(String), Node(String), - UnknownNode(LogicalNodeId), EndpointUnavailable(ProviderLeaseId), MissingBootstrap(LogicalNodeId), + Bootstrap(String), + UnexpectedCommand(String), } -pub struct DockerClusterHarness +pub struct DockerNodeProvisioner where D: DockerCli, F: SshBootstrapClientFactory, @@ -458,12 +416,12 @@ where provider: DockerProvider, client_factory: F, datastream: S, - router: SwactorJoinRouter, - nodes: BTreeMap>, + manager: NodeManager, + bootstrap: Option>, teardown_complete: bool, } -impl DockerClusterHarness +impl DockerNodeProvisioner where D: DockerCli, F: SshBootstrapClientFactory, @@ -474,9 +432,9 @@ where provider, client_factory, datastream, - router: SwactorJoinRouter::new(), - nodes: BTreeMap::new(), - teardown_complete: false, + manager: NodeManager::new(), + bootstrap: None, + teardown_complete: true, } } @@ -484,131 +442,96 @@ where &self.provider } + pub fn provider_mut(&mut self) -> &mut DockerProvider { + &mut self.provider + } + pub fn datastream(&self) -> &S { &self.datastream } - pub fn nodes(&self) -> &BTreeMap> { - &self.nodes + pub fn manager(&self) -> &NodeManager { + &self.manager } - pub fn records(&self) -> Vec { - self.nodes - .values() - .filter_map(|node| node.manager.record().cloned()) - .collect() + pub fn bootstrap(&self) -> Option<&SshBootstrapSession> { + self.bootstrap.as_ref() } - pub fn all_ready(&self) -> bool { - !self.nodes.is_empty() && self.nodes.values().all(|node| node.manager.is_ready()) + pub fn record(&self) -> Option<&NodeRecord> { + self.manager.record() } - pub fn start_group(&mut self, group: &RunNodeGroupSpec) -> Result<(), DockerClusterError> { - for spec in expand_node_group(group) { - let logical_node_id = spec.logical_node_id.clone(); - self.router.register(logical_node_id.clone()); - let mut manager = NodeManager::new(); - let commands = manager - .handle(NodeManagerMsg::Start(spec)) - .map_err(|error| DockerClusterError::Node(error.reason))?; - match self.process_start_commands(&mut manager, commands) { - Ok(bootstrap) => { - self.nodes - .insert(logical_node_id, ManagedDockerNode { manager, bootstrap }); - } - Err(error) => { - self.nodes.insert( - logical_node_id, - ManagedDockerNode { - manager, - bootstrap: None, - }, - ); - let _ = self.teardown(); - return Err(error); - } - } - } + pub fn is_ready(&self) -> bool { + self.manager.is_ready() + } + + pub fn start( + &mut self, + spec: crate::node_provisioning::LogicalNodeSpec, + ) -> Result<(), DockerNodeProvisionError> { + let commands = self + .manager + .handle(NodeManagerMsg::Start(spec)) + .map_err(|error| DockerNodeProvisionError::Node(error.reason))?; self.teardown_complete = false; - Ok(()) - } - - pub fn route_join(&mut self, event: SwactorJoinEvent) -> Result<(), DockerClusterError> { - let logical_node_id = event.logical_node_id.clone(); - let node = self - .nodes - .get_mut(&logical_node_id) - .ok_or_else(|| DockerClusterError::UnknownNode(logical_node_id.clone()))?; - let commands = self.router.route(event, &mut node.manager)?; - for command in commands { - match command { - NodeManagerCommand::BootstrapConvergenceObserved { swactor_id, .. } => { - let bootstrap = node.bootstrap.as_mut().ok_or_else(|| { - DockerClusterError::MissingBootstrap(logical_node_id.clone()) - })?; - let events = bootstrap.convergence_observed(swactor_id, &mut self.datastream); - Self::feed_bootstrap_events(&mut node.manager, events)?; - } - other => { - return Err(DockerClusterError::Node(format!( - "unexpected command {other:?}" - ))); - } - } + if let Err(error) = self.process_commands(commands) { + let _ = self.stop(); + return Err(error); } Ok(()) } - pub fn teardown(&mut self) -> Result<(), DockerClusterError> { - for node in self.nodes.values_mut() { - let commands = node - .manager - .handle(NodeManagerMsg::Destroy) - .map_err(|error| DockerClusterError::Node(error.reason))?; - for command in commands { - match command { - NodeManagerCommand::CancelBootstrap { .. } => { - if let Some(bootstrap) = node.bootstrap.as_mut() { - let _ = bootstrap.cancel(); - } - } - NodeManagerCommand::DestroyLease(handle) => { - self.provider - .destroy_lease(&handle) - .map_err(|error| DockerClusterError::Provider(error.reason))?; - node.manager - .handle(NodeManagerMsg::LeaseDestroyed) - .map_err(|error| DockerClusterError::Node(error.reason))?; - } - other => { - return Err(DockerClusterError::Node(format!( - "unexpected command {other:?}" - ))); - } - } - } + pub fn observe_swactor_join( + &mut self, + swactor_id: SwactorId, + ) -> Result<(), DockerNodeProvisionError> { + let logical_node_id = self + .manager + .record() + .ok_or_else(|| DockerNodeProvisionError::Node("node manager not started".into()))? + .logical_node_id + .clone(); + let commands = self + .manager + .handle(NodeManagerMsg::SwactorJoined { + logical_node_id, + swactor_id, + }) + .map_err(|error| DockerNodeProvisionError::Node(error.reason))?; + self.process_commands(commands) + } + + pub fn stop(&mut self) -> Result<(), DockerNodeProvisionError> { + if self.manager.record().is_none() { + self.teardown_complete = true; + return Ok(()); } + let commands = self + .manager + .handle(NodeManagerMsg::Destroy) + .map_err(|error| DockerNodeProvisionError::Node(error.reason))?; + self.process_stop_commands(commands)?; self.teardown_complete = true; Ok(()) } - fn process_start_commands( + fn process_commands( &mut self, - manager: &mut NodeManager, commands: Vec, - ) -> Result>, DockerClusterError> { + ) -> Result<(), DockerNodeProvisionError> { let mut pending = commands; - let mut bootstrap = None; while let Some(command) = pending.pop() { match command { NodeManagerCommand::CreateLease(request) => { let result = self .provider .create_lease(request) - .map_err(|error| DockerClusterError::Provider(error.reason))?; - let more = manager + .map_err(|error| DockerNodeProvisionError::Provider(error.reason))?; + let more = self + .manager .handle(NodeManagerMsg::LeaseCreated(result)) - .map_err(|error| DockerClusterError::Node(error.reason))?; + .map_err(|error| DockerNodeProvisionError::Node(error.reason))?; pending.extend(more); } NodeManagerCommand::LookupEndpoint(lease) => { @@ -616,58 +539,115 @@ where let endpoint = match self .provider .lookup_endpoint(&lease) - .map_err(|error| DockerClusterError::Provider(error.reason))? + .map_err(|error| DockerNodeProvisionError::Provider(error.reason))? { Some(endpoint) => endpoint, None => { - let _ = manager.handle(NodeManagerMsg::EndpointFailed( + let _ = self.manager.handle(NodeManagerMsg::EndpointFailed( "docker ssh endpoint unavailable".into(), )); - return Err(DockerClusterError::EndpointUnavailable(lease_id)); + return Err(DockerNodeProvisionError::EndpointUnavailable(lease_id)); } }; - let more = manager + let more = self + .manager .handle(NodeManagerMsg::EndpointKnown(endpoint)) - .map_err(|error| DockerClusterError::Node(error.reason))?; + .map_err(|error| DockerNodeProvisionError::Node(error.reason))?; pending.extend(more); } NodeManagerCommand::StartBootstrap(spec) => { + let logical_node_id = spec.logical_node_id.clone(); let client = self.client_factory.client_for(&spec); let mut session = SshBootstrapSession::new(spec, client); let events = session.start(&mut self.datastream); - Self::feed_bootstrap_events(manager, events)?; - bootstrap = Some(session); + let bootstrap_error = events.iter().find_map(|event| match event { + BootstrapSessionEvent::Failed(reason) => Some(reason.clone()), + BootstrapSessionEvent::Observed(_) | BootstrapSessionEvent::Closed => None, + }); + Self::feed_bootstrap_events(&mut self.manager, events)?; + self.bootstrap = Some(session); + if let Some(reason) = bootstrap_error { + return Err(DockerNodeProvisionError::Bootstrap(format!( + "{}: {reason}", + logical_node_id.0 + ))); + } } - other => { - return Err(DockerClusterError::Node(format!( - "unexpected command {other:?}" + NodeManagerCommand::BootstrapConvergenceObserved { swactor_id, .. } => { + let logical_node_id = self + .manager + .record() + .ok_or_else(|| { + DockerNodeProvisionError::Node("node manager not started".into()) + })? + .logical_node_id + .clone(); + let bootstrap = self.bootstrap.as_mut().ok_or_else(|| { + DockerNodeProvisionError::MissingBootstrap(logical_node_id.clone()) + })?; + let events = bootstrap.convergence_observed(swactor_id, &mut self.datastream); + Self::feed_bootstrap_events(&mut self.manager, events)?; + } + NodeManagerCommand::CancelBootstrap { .. } + | NodeManagerCommand::DestroyLease(_) => { + return Err(DockerNodeProvisionError::UnexpectedCommand(format!( + "lifecycle command {command:?} outside stop" ))); } } } - Ok(bootstrap) + Ok(()) + } + + fn process_stop_commands( + &mut self, + commands: Vec, + ) -> Result<(), DockerNodeProvisionError> { + for command in commands { + match command { + NodeManagerCommand::CancelBootstrap { .. } => { + if let Some(bootstrap) = self.bootstrap.as_mut() { + let _ = bootstrap.cancel(); + } + } + NodeManagerCommand::DestroyLease(handle) => { + self.provider + .destroy_lease(&handle) + .map_err(|error| DockerNodeProvisionError::Provider(error.reason))?; + self.manager + .handle(NodeManagerMsg::LeaseDestroyed) + .map_err(|error| DockerNodeProvisionError::Node(error.reason))?; + } + other => { + return Err(DockerNodeProvisionError::UnexpectedCommand(format!( + "unexpected stop command {other:?}" + ))); + } + } + } + Ok(()) } fn feed_bootstrap_events( manager: &mut NodeManager, events: Vec, - ) -> Result<(), DockerClusterError> { + ) -> Result<(), DockerNodeProvisionError> { for event in events { match event { BootstrapSessionEvent::Observed(observation) => { manager .handle(NodeManagerMsg::BootstrapObserved(observation)) - .map_err(|error| DockerClusterError::Node(error.reason))?; + .map_err(|error| DockerNodeProvisionError::Node(error.reason))?; } BootstrapSessionEvent::Failed(reason) => { manager .handle(NodeManagerMsg::BootstrapFailed(reason)) - .map_err(|error| DockerClusterError::Node(error.reason))?; + .map_err(|error| DockerNodeProvisionError::Node(error.reason))?; } BootstrapSessionEvent::Closed => { manager .handle(NodeManagerMsg::BootstrapClosed) - .map_err(|error| DockerClusterError::Node(error.reason))?; + .map_err(|error| DockerNodeProvisionError::Node(error.reason))?; } } } @@ -675,7 +655,7 @@ where } } -impl Drop for DockerClusterHarness +impl Drop for DockerNodeProvisioner where D: DockerCli, F: SshBootstrapClientFactory, @@ -683,7 +663,7 @@ where { fn drop(&mut self) { if !self.teardown_complete { - let _ = self.teardown(); + let _ = self.stop(); } } } diff --git a/crates/mvp-system/src/gpu_worker_ingress_parser.rs b/crates/mvp-system/src/gpu_worker_ingress_parser.rs index 750fca2..4ae0998 100644 --- a/crates/mvp-system/src/gpu_worker_ingress_parser.rs +++ b/crates/mvp-system/src/gpu_worker_ingress_parser.rs @@ -210,7 +210,6 @@ pub struct DeviceCopyLog { #[derive(Clone, Debug)] pub struct ObjectRecordBuilder { - spec: ObjectSpec, object_id: ObjectId, sequence: u64, extent: u64, @@ -222,9 +221,8 @@ pub struct ObjectRecordBuilder { } impl ObjectRecordBuilder { - pub fn new(spec: ObjectSpec) -> Self { + pub fn new(_spec: ObjectSpec) -> Self { Self { - spec, object_id: ObjectId(9000), sequence: 0, extent: 0, @@ -557,6 +555,7 @@ pub fn read_object_record( })) } +#[cfg(test)] fn object_failure_metadata( bytes: &[u8], reason: ObjectFailureReason, diff --git a/crates/mvp-system/src/lib.rs b/crates/mvp-system/src/lib.rs index 4971001..0a8f4d8 100644 --- a/crates/mvp-system/src/lib.rs +++ b/crates/mvp-system/src/lib.rs @@ -24,11 +24,14 @@ pub mod orchestrator_token_endpoint; pub mod provisioning; pub mod resource_inventory; pub mod run_plan; +pub mod shard_fetch; +pub mod shard_weight_lifecycle; pub mod shared_ring_helper_abi; pub mod stage_controller; pub mod telemetry; pub mod tx_rx_edge_actor; pub mod weight_lifecycle; +pub mod weight_shards; #[cfg(test)] mod tests; @@ -36,3 +39,13 @@ mod tests; #[cfg(test)] #[path = "tests/driver_pumps_guarantees.rs"] mod driver_pumps_guarantees; + +#[cfg(test)] +#[path = "tests/shard_fetch_guarantees.rs"] +mod shard_fetch_guarantees; +#[cfg(test)] +#[path = "tests/shard_weight_lifecycle_guarantees.rs"] +mod shard_weight_lifecycle_guarantees; +#[cfg(test)] +#[path = "tests/weight_shards_guarantees.rs"] +mod weight_shards_guarantees; diff --git a/crates/mvp-system/src/membership_pool_readiness.rs b/crates/mvp-system/src/membership_pool_readiness.rs index 2c9636c..f1292c4 100644 --- a/crates/mvp-system/src/membership_pool_readiness.rs +++ b/crates/mvp-system/src/membership_pool_readiness.rs @@ -86,6 +86,7 @@ pub enum ReadinessCommand { AssignObjectSpec { node_id: NodeId }, } +#[cfg(test)] #[derive(Clone, Default)] struct NodeFacts { known: bool, diff --git a/crates/mvp-system/src/shard_fetch.rs b/crates/mvp-system/src/shard_fetch.rs new file mode 100644 index 0000000..365a80f --- /dev/null +++ b/crates/mvp-system/src/shard_fetch.rs @@ -0,0 +1,108 @@ +use crate::weight_shards::{ShardAssignment, ShardManifest}; + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct ShardLocation { + pub uri: String, + pub cache_key: String, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct FetchShard { + pub location: ShardLocation, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct FetchedShard { + pub local_path: String, + pub manifest: ShardManifest, +} + +impl FetchedShard { + pub fn new(local_path: impl Into, manifest: ShardManifest) -> Self { + Self { + local_path: local_path.into(), + manifest, + } + } +} + +pub struct ShardLocator; + +impl ShardLocator { + pub fn locate(assignment: &ShardAssignment) -> ShardLocation { + let uri = assignment + .model_ref + .shard_uri(&assignment.split_id, assignment.stage_index); + let cache_key = format!( + "{}:{}:{:05}", + assignment.expected_model_digest().as_str(), + assignment.split_id.as_str(), + assignment.stage_index, + ); + ShardLocation { uri, cache_key } + } +} + +pub trait ShardCache { + fn get(&self, cache_key: &str) -> Option; + fn insert(&mut self, cache_key: String, shard: FetchedShard); +} + +pub trait ShardFetcher { + fn fetch(&mut self, request: &FetchShard) -> Result; +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum FetchError { + Unauthorized, + NotFound, + Unavailable, + IntegrityMismatch, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum ShardFetchStatus { + CacheHit, + Downloaded, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct ShardFetchOutcome { + pub shard: FetchedShard, + pub location: ShardLocation, + pub status: ShardFetchStatus, +} + +pub struct ShardFetchCoordinator; + +impl ShardFetchCoordinator { + pub fn get_or_fetch( + assignment: &ShardAssignment, + cache: &mut C, + fetcher: &mut F, + ) -> Result + where + C: ShardCache, + F: ShardFetcher, + { + let location = ShardLocator::locate(assignment); + if let Some(shard) = cache.get(&location.cache_key) { + return Ok(ShardFetchOutcome { + shard, + location, + status: ShardFetchStatus::CacheHit, + }); + } + + let request = FetchShard { + location: location.clone(), + }; + let shard = fetcher.fetch(&request)?; + cache.insert(location.cache_key.clone(), shard.clone()); + Ok(ShardFetchOutcome { + shard, + location, + status: ShardFetchStatus::Downloaded, + }) + } +} diff --git a/crates/mvp-system/src/shard_weight_lifecycle.rs b/crates/mvp-system/src/shard_weight_lifecycle.rs new file mode 100644 index 0000000..58d02a8 --- /dev/null +++ b/crates/mvp-system/src/shard_weight_lifecycle.rs @@ -0,0 +1,165 @@ +use crate::shard_fetch::{ + FetchError, ShardCache, ShardFetchCoordinator, ShardFetchStatus, ShardFetcher, ShardLocation, +}; +use crate::weight_shards::{ShardAssignment, ShardValidationError, ValidatedShard}; + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum ShardLifecycleState { + Idle, + Assigned, + Located, + Fetching, + Fetched, + Validating, + Binding, + Ready, + Faulted, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum ShardLifecycleEvent { + Assigned { assignment: ShardAssignment }, + Located { location: ShardLocation }, + Fetching { location: ShardLocation }, + CacheHit { cache_key: String }, + Fetched { uri: String, local_path: String }, + Validated { local_path: String }, + Binding { local_path: String }, + Ready { local_path: String }, + Faulted { reason: ShardLifecycleFault }, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum ShardLifecycleFault { + Fetch(FetchError), + Validation(ShardValidationError), + Bind(BindError), +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum BindError { + WorkerRejected, + DeviceAllocationFailed, +} + +pub trait WorkerShardBinder { + fn bind(&mut self, shard: &ValidatedShard) -> Result<(), BindError>; +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct ShardWeightLifecycle { + state: ShardLifecycleState, + events: Vec, +} + +impl ShardWeightLifecycle { + pub fn new() -> Self { + Self { + state: ShardLifecycleState::Idle, + events: Vec::new(), + } + } + + pub fn state(&self) -> ShardLifecycleState { + self.state + } + + pub fn events(&self) -> &[ShardLifecycleEvent] { + &self.events + } + + pub fn load( + &mut self, + assignment: ShardAssignment, + cache: &mut C, + fetcher: &mut F, + binder: &mut B, + ) where + C: ShardCache, + F: ShardFetcher, + B: WorkerShardBinder, + { + if matches!( + self.state, + ShardLifecycleState::Ready | ShardLifecycleState::Faulted + ) { + return; + } + + self.state = ShardLifecycleState::Assigned; + self.events.push(ShardLifecycleEvent::Assigned { + assignment: assignment.clone(), + }); + + let location = crate::shard_fetch::ShardLocator::locate(&assignment); + self.state = ShardLifecycleState::Located; + self.events.push(ShardLifecycleEvent::Located { + location: location.clone(), + }); + + self.state = ShardLifecycleState::Fetching; + self.events.push(ShardLifecycleEvent::Fetching { + location: location.clone(), + }); + + let outcome = match ShardFetchCoordinator::get_or_fetch(&assignment, cache, fetcher) { + Ok(outcome) => outcome, + Err(error) => { + self.fault(ShardLifecycleFault::Fetch(error)); + return; + } + }; + + match outcome.status { + ShardFetchStatus::CacheHit => self.events.push(ShardLifecycleEvent::CacheHit { + cache_key: outcome.location.cache_key, + }), + ShardFetchStatus::Downloaded => self.events.push(ShardLifecycleEvent::Fetched { + uri: outcome.location.uri, + local_path: outcome.shard.local_path.clone(), + }), + } + + self.state = ShardLifecycleState::Fetched; + self.state = ShardLifecycleState::Validating; + let validated = match ValidatedShard::new( + assignment, + outcome.shard.manifest, + outcome.shard.local_path.clone(), + ) { + Ok(validated) => validated, + Err(error) => { + self.fault(ShardLifecycleFault::Validation(error)); + return; + } + }; + self.events.push(ShardLifecycleEvent::Validated { + local_path: validated.local_path.clone(), + }); + + self.state = ShardLifecycleState::Binding; + self.events.push(ShardLifecycleEvent::Binding { + local_path: validated.local_path.clone(), + }); + if let Err(error) = binder.bind(&validated) { + self.fault(ShardLifecycleFault::Bind(error)); + return; + } + + self.state = ShardLifecycleState::Ready; + self.events.push(ShardLifecycleEvent::Ready { + local_path: validated.local_path, + }); + } + + fn fault(&mut self, reason: ShardLifecycleFault) { + self.state = ShardLifecycleState::Faulted; + self.events.push(ShardLifecycleEvent::Faulted { reason }); + } +} + +impl Default for ShardWeightLifecycle { + fn default() -> Self { + Self::new() + } +} diff --git a/crates/mvp-system/src/telemetry.rs b/crates/mvp-system/src/telemetry.rs index 66b5a35..e6849a8 100644 --- a/crates/mvp-system/src/telemetry.rs +++ b/crates/mvp-system/src/telemetry.rs @@ -1,10 +1,11 @@ //! MVP-system-owned datastream channel records. +use datastream::frame::ChannelId; use datastream::{ChannelRegistry, Record}; use serde::{Deserialize, Serialize}; use crate::observability_surface as obs; -use crate::provisioning; +use crate::provisioning::{self, ProvisionLogStream}; /// Structured MVP lifecycle facts: run, node, stage, edge, ring, object, step, and worker events. pub const MVP_LIFECYCLE: &str = "mvp.lifecycle"; @@ -67,6 +68,15 @@ impl MvpProvisionLogRecord { } } +pub fn mvp_provision_log_channel(node_id: u64, stream: ProvisionLogStream) -> ChannelId { + let stream = match stream { + ProvisionLogStream::Stdout => "stdout", + ProvisionLogStream::Stderr => "stderr", + ProvisionLogStream::Provider => "provider", + }; + ChannelId::new(format!("mvp.provisioning.logs.node.{node_id}.{stream}")) +} + impl Record for MvpProvisionLogRecord { const CHANNEL: &'static str = MVP_PROVISIONING_LOGS; } diff --git a/crates/mvp-system/src/tests/docker_cluster_provisioning_guarantees.rs b/crates/mvp-system/src/tests/docker_cluster_provisioning_guarantees.rs index 53943ad..1775567 100644 --- a/crates/mvp-system/src/tests/docker_cluster_provisioning_guarantees.rs +++ b/crates/mvp-system/src/tests/docker_cluster_provisioning_guarantees.rs @@ -205,30 +205,6 @@ impl docker::SshBootstrapClientFactory for FakeSshFactory { } } -fn started_manager() -> (provision::NodeManager, provision::LogicalNodeId) { - let spec = one_spec(); - let logical_node_id = spec.logical_node_id.clone(); - let mut manager = provision::NodeManager::new(); - let start_commands = manager - .handle(provision::NodeManagerMsg::Start(spec.clone())) - .expect("start succeeds"); - assert!(matches!( - start_commands.as_slice(), - [provision::NodeManagerCommand::CreateLease(_)] - )); - let lease = docker::DockerProvider::new(FakeDockerCli::new()) - .create_lease(provision::CreateLeaseRequest { spec }) - .expect("docker provider creates lease"); - let bootstrap_commands = manager - .handle(provision::NodeManagerMsg::LeaseCreated(lease)) - .expect("lease accepted"); - assert!(matches!( - bootstrap_commands.as_slice(), - [provision::NodeManagerCommand::StartBootstrap(_)] - )); - (manager, logical_node_id) -} - #[test] fn docker_provider_maps_node_spec_to_container_lease_and_destroy_handle() { let spec = one_spec(); @@ -330,137 +306,110 @@ fn ssh_bootstrap_session_runs_pre_handoff_steps_and_closes_on_convergence() { } #[test] -fn join_router_routes_known_logical_nodes_and_rejects_unknown_nodes() { - let (mut manager, logical_node_id) = started_manager(); - let mut router = docker::SwactorJoinRouter::new(); - router.register(logical_node_id.clone()); - - let unknown = router.route( - docker::SwactorJoinEvent { - logical_node_id: provision::LogicalNodeId("workers-99".into()), - swactor_id: provision::SwactorId("swactor-wrong".into()), - }, - &mut manager, +fn docker_node_provisioner_rejects_join_for_wrong_logical_node() { + let spec = one_spec(); + let mut node = docker::DockerNodeProvisioner::new( + docker::DockerProvider::new(FakeDockerCli::new()), + FakeSshFactory::default(), + provision::InMemoryBootstrapDatastream::default(), ); - assert!(matches!( - unknown, - Err(docker::DockerClusterError::UnknownNode(id)) if id == provision::LogicalNodeId("workers-99".into()) - )); + + node.start(spec).expect("node starts"); + let mut manager = node.manager().clone(); + let wrong_join = manager.handle(provision::NodeManagerMsg::SwactorJoined { + logical_node_id: provision::LogicalNodeId("workers-99".into()), + swactor_id: provision::SwactorId("swactor-wrong".into()), + }); + + assert!(wrong_join.is_err()); assert_eq!( - manager.record().expect("record exists").stage, + node.record().expect("record exists").stage, provision::NodeStage::BootstrapRunning ); - - let commands = router - .route( - docker::SwactorJoinEvent { - logical_node_id, - swactor_id: provision::SwactorId("swactor-workers-0".into()), - }, - &mut manager, - ) - .expect("known join routes"); - assert!(matches!( - commands.as_slice(), - [provision::NodeManagerCommand::BootstrapConvergenceObserved { .. }] - )); } #[test] -fn docker_cluster_harness_wires_provider_bootstrap_join_and_teardown() { +fn docker_node_provisioner_wires_provider_bootstrap_join_and_teardown() { let provider = docker::DockerProvider::new(FakeDockerCli::new()); let datastream = provision::InMemoryBootstrapDatastream::default(); - let mut harness = - docker::DockerClusterHarness::new(provider, FakeSshFactory::default(), datastream); + let mut node = + docker::DockerNodeProvisioner::new(provider, FakeSshFactory::default(), datastream); - harness - .start_group(&docker_group_spec(2)) - .expect("cluster starts"); - assert!(!harness.all_ready()); - assert_eq!(harness.nodes().len(), 2); - assert_eq!(harness.provider().cli().run_requests.len(), 2); - assert_eq!(harness.datastream().records().len(), 4); + node.start(one_spec()).expect("node starts"); - let ids: Vec<_> = harness.nodes().keys().cloned().collect(); - for id in ids { - harness - .route_join(docker::SwactorJoinEvent { - swactor_id: provision::SwactorId(format!("swactor-{}", id.0)), - logical_node_id: id, - }) - .expect("join routes"); - } - - assert!(harness.all_ready()); - assert_eq!(harness.datastream().flush_count(), 2); - for record in harness.records() { - assert_eq!(record.stage, provision::NodeStage::Dormant); - assert!(record.ready); - assert_eq!( - record.lease.as_ref().unwrap().provider, - provision::ProviderKind::Docker - ); - } - - harness.teardown().expect("teardown succeeds"); + assert!(!node.is_ready()); + assert_eq!(node.provider().cli().run_requests.len(), 1); + assert_eq!(node.datastream().records().len(), 2); assert_eq!( - harness.provider().cli().removed, - vec!["container-1", "container-2"] - ); - for record in harness.records() { - assert_eq!(record.stage, provision::NodeStage::Destroyed); - assert!(!record.ready); - } -} - -#[test] -fn docker_cluster_harness_teardown_cleans_known_leases_before_join() { - let provider = docker::DockerProvider::new(FakeDockerCli::new()); - let datastream = provision::InMemoryBootstrapDatastream::default(); - let mut harness = - docker::DockerClusterHarness::new(provider, FakeSshFactory::default(), datastream); - - harness - .start_group(&docker_group_spec(1)) - .expect("cluster starts"); - let id = provision::LogicalNodeId("workers-0".into()); - assert_eq!( - harness.nodes()[&id].manager.record().unwrap().stage, + node.record().expect("record exists").stage, provision::NodeStage::BootstrapRunning ); - harness.teardown().expect("teardown succeeds before join"); + node.observe_swactor_join(provision::SwactorId("swactor-workers-0".into())) + .expect("join completes handoff"); - let node = &harness.nodes()[&id]; + assert!(node.is_ready()); + assert_eq!(node.datastream().flush_count(), 1); + let record = node.record().expect("record exists"); + assert_eq!(record.stage, provision::NodeStage::Dormant); + assert!(record.ready); assert_eq!( - node.manager.record().unwrap().stage, + record.lease.as_ref().unwrap().provider, + provision::ProviderKind::Docker + ); + assert!(node.bootstrap().expect("bootstrap exists").is_closed()); + assert!(node.bootstrap().unwrap().client().closed); + + node.stop().expect("teardown succeeds"); + assert_eq!(node.provider().cli().removed, vec!["container-1"]); + let record = node.record().expect("record exists"); + assert_eq!(record.stage, provision::NodeStage::Destroyed); + assert!(!record.ready); +} + +#[test] +fn docker_node_provisioner_teardown_cleans_known_lease_before_join() { + let provider = docker::DockerProvider::new(FakeDockerCli::new()); + let datastream = provision::InMemoryBootstrapDatastream::default(); + let mut node = + docker::DockerNodeProvisioner::new(provider, FakeSshFactory::default(), datastream); + + node.start(one_spec()).expect("node starts"); + assert_eq!( + node.record().unwrap().stage, + provision::NodeStage::BootstrapRunning + ); + + node.stop().expect("teardown succeeds before join"); + + assert_eq!( + node.record().unwrap().stage, provision::NodeStage::Destroyed ); - assert!(node.bootstrap.as_ref().unwrap().is_closed()); - assert!(node.bootstrap.as_ref().unwrap().client().closed); - assert_eq!(harness.provider().cli().removed, vec!["container-1"]); + assert!(node.bootstrap().unwrap().is_closed()); + assert!(node.bootstrap().unwrap().client().closed); + assert_eq!(node.provider().cli().removed, vec!["container-1"]); } #[test] -fn docker_cluster_harness_cleans_known_lease_when_endpoint_never_appears() { +fn docker_node_provisioner_cleans_known_lease_when_endpoint_never_appears() { let mut cli = FakeDockerCli::new(); cli.endpoints.insert("container-1".into(), None); let provider = docker::DockerProvider::new(cli); let datastream = provision::InMemoryBootstrapDatastream::default(); - let mut harness = - docker::DockerClusterHarness::new(provider, FakeSshFactory::default(), datastream); + let mut node = + docker::DockerNodeProvisioner::new(provider, FakeSshFactory::default(), datastream); - let result = harness.start_group(&docker_group_spec(1)); + let result = node.start(one_spec()); assert!(matches!( result, - Err(docker::DockerClusterError::EndpointUnavailable(id)) + Err(docker::DockerNodeProvisionError::EndpointUnavailable(id)) if id == provision::ProviderLeaseId("docker:container-1".into()) )); - assert_eq!(harness.provider().cli().removed, vec!["container-1"]); - let id = provision::LogicalNodeId("workers-0".into()); + assert_eq!(node.provider().cli().removed, vec!["container-1"]); assert_eq!( - harness.nodes()[&id].manager.record().unwrap().stage, + node.record().unwrap().stage, provision::NodeStage::Destroyed ); } diff --git a/crates/mvp-system/src/tests/shard_fetch_guarantees.rs b/crates/mvp-system/src/tests/shard_fetch_guarantees.rs new file mode 100644 index 0000000..0292124 --- /dev/null +++ b/crates/mvp-system/src/tests/shard_fetch_guarantees.rs @@ -0,0 +1,168 @@ +use std::collections::BTreeMap; + +use mvp_system::shard_fetch as fetch; +use mvp_system::weight_shards as shards; + +fn assignment(stage_index: u32) -> shards::ShardAssignment { + let model_ref = + shards::ModelArtifactRef::parse("hf://org/repo@abcdef123456/model.gguf").unwrap(); + let split_scheme = shards::SplitScheme::GgufLayerContiguousV1; + let split_id = shards::SplitId::derive(&model_ref, split_scheme); + shards::ShardAssignment::new( + model_ref, + split_id, + split_scheme, + stage_index, + 8, + shards::LayerRange::new(stage_index * 4, stage_index * 4 + 4).unwrap(), + ) + .unwrap() +} + +fn fetched_for(assignment: &shards::ShardAssignment) -> fetch::FetchedShard { + fetch::FetchedShard::new( + format!("/cache/stage-{:05}.gguf", assignment.stage_index), + shards::ShardManifest::for_assignment( + assignment, + shards::ContentHash::literal(format!("sha256:stage-{}", assignment.stage_index)) + .unwrap(), + ), + ) +} + +#[derive(Default)] +struct MemoryCache { + entries: BTreeMap, + inserts: usize, +} + +impl fetch::ShardCache for MemoryCache { + fn get(&self, cache_key: &str) -> Option { + self.entries.get(cache_key).cloned() + } + + fn insert(&mut self, cache_key: String, shard: fetch::FetchedShard) { + self.inserts += 1; + self.entries.insert(cache_key, shard); + } +} + +struct RecordingFetcher { + calls: Vec, + result: Result, +} + +impl fetch::ShardFetcher for RecordingFetcher { + fn fetch( + &mut self, + request: &fetch::FetchShard, + ) -> Result { + self.calls.push(request.clone()); + self.result.clone() + } +} + +#[test] +fn shard_locator_derives_expected_uri() { + let assignment = assignment(3); + let location = fetch::ShardLocator::locate(&assignment); + + assert_eq!( + location.uri, + format!( + "hf://org/repo@abcdef123456/shards/{}/stage-00003.gguf", + assignment.split_id.as_str() + ) + ); + assert!(location.cache_key.contains(assignment.split_id.as_str())); + assert!(location.cache_key.ends_with(":00003")); +} + +#[test] +fn cache_key_depends_on_model_split_and_stage() { + let stage_three = assignment(3); + let stage_four = assignment(4); + + let different_model_ref = + shards::ModelArtifactRef::parse("hf://other/repo@abcdef123456/model.gguf").unwrap(); + let split_scheme = shards::SplitScheme::GgufLayerContiguousV1; + let different_model = shards::ShardAssignment::new( + different_model_ref.clone(), + shards::SplitId::derive(&different_model_ref, split_scheme), + split_scheme, + 3, + 8, + shards::LayerRange::new(12, 16).unwrap(), + ) + .unwrap(); + + let different_split = shards::ShardAssignment::new( + stage_three.model_ref.clone(), + shards::SplitId::literal("split-other").unwrap(), + stage_three.split_scheme, + 3, + 8, + stage_three.layer_range, + ) + .unwrap(); + + let keys = [ + fetch::ShardLocator::locate(&stage_three).cache_key, + fetch::ShardLocator::locate(&stage_four).cache_key, + fetch::ShardLocator::locate(&different_model).cache_key, + fetch::ShardLocator::locate(&different_split).cache_key, + ]; + + assert_eq!( + keys.iter().collect::>().len(), + keys.len() + ); +} + +#[test] +fn coordinator_downloads_and_caches_missing_shard() { + let assignment = assignment(2); + let expected = fetched_for(&assignment); + let mut cache = MemoryCache::default(); + let mut fetcher = RecordingFetcher { + calls: Vec::new(), + result: Ok(expected.clone()), + }; + + let outcome = + fetch::ShardFetchCoordinator::get_or_fetch(&assignment, &mut cache, &mut fetcher).unwrap(); + + assert_eq!(outcome.status, fetch::ShardFetchStatus::Downloaded); + assert_eq!(outcome.shard, expected); + assert_eq!(fetcher.calls.len(), 1); + assert_eq!(fetcher.calls[0].location, outcome.location); + assert_eq!(cache.inserts, 1); + assert_eq!( + cache.entries.get(&outcome.location.cache_key), + Some(&expected) + ); +} + +#[test] +fn coordinator_uses_cache_hit_without_fetching() { + let assignment = assignment(5); + let location = fetch::ShardLocator::locate(&assignment); + let cached = fetched_for(&assignment); + let mut cache = MemoryCache::default(); + cache + .entries + .insert(location.cache_key.clone(), cached.clone()); + let mut fetcher = RecordingFetcher { + calls: Vec::new(), + result: Err(fetch::FetchError::Unavailable), + }; + + let outcome = + fetch::ShardFetchCoordinator::get_or_fetch(&assignment, &mut cache, &mut fetcher).unwrap(); + + assert_eq!(outcome.status, fetch::ShardFetchStatus::CacheHit); + assert_eq!(outcome.location, location); + assert_eq!(outcome.shard, cached); + assert!(fetcher.calls.is_empty()); + assert_eq!(cache.inserts, 0); +} diff --git a/crates/mvp-system/src/tests/shard_weight_lifecycle_guarantees.rs b/crates/mvp-system/src/tests/shard_weight_lifecycle_guarantees.rs new file mode 100644 index 0000000..a4d1db1 --- /dev/null +++ b/crates/mvp-system/src/tests/shard_weight_lifecycle_guarantees.rs @@ -0,0 +1,308 @@ +use std::collections::BTreeMap; + +use mvp_system::shard_fetch as fetch; +use mvp_system::shard_weight_lifecycle as lifecycle; +use mvp_system::weight_shards as shards; + +fn assignment() -> shards::ShardAssignment { + let model_ref = + shards::ModelArtifactRef::parse("hf://org/repo@abcdef123456/model.gguf").unwrap(); + let split_scheme = shards::SplitScheme::GgufLayerContiguousV1; + let split_id = shards::SplitId::derive(&model_ref, split_scheme); + shards::ShardAssignment::new( + model_ref, + split_id, + split_scheme, + 3, + 8, + shards::LayerRange::new(12, 16).unwrap(), + ) + .unwrap() +} + +fn fetched_for(assignment: &shards::ShardAssignment) -> fetch::FetchedShard { + fetch::FetchedShard::new( + "/cache/stage-00003.gguf", + shards::ShardManifest::for_assignment( + assignment, + shards::ContentHash::literal("sha256:stage-3").unwrap(), + ), + ) +} + +fn event_kinds(events: &[lifecycle::ShardLifecycleEvent]) -> Vec<&'static str> { + events + .iter() + .map(|event| match event { + lifecycle::ShardLifecycleEvent::Assigned { .. } => "assigned", + lifecycle::ShardLifecycleEvent::Located { .. } => "located", + lifecycle::ShardLifecycleEvent::Fetching { .. } => "fetching", + lifecycle::ShardLifecycleEvent::CacheHit { .. } => "cache-hit", + lifecycle::ShardLifecycleEvent::Fetched { .. } => "fetched", + lifecycle::ShardLifecycleEvent::Validated { .. } => "validated", + lifecycle::ShardLifecycleEvent::Binding { .. } => "binding", + lifecycle::ShardLifecycleEvent::Ready { .. } => "ready", + lifecycle::ShardLifecycleEvent::Faulted { .. } => "faulted", + }) + .collect() +} + +#[derive(Default)] +struct MemoryCache { + entries: BTreeMap, + inserts: usize, +} + +impl fetch::ShardCache for MemoryCache { + fn get(&self, cache_key: &str) -> Option { + self.entries.get(cache_key).cloned() + } + + fn insert(&mut self, cache_key: String, shard: fetch::FetchedShard) { + self.inserts += 1; + self.entries.insert(cache_key, shard); + } +} + +struct RecordingFetcher { + calls: Vec, + result: Result, +} + +impl fetch::ShardFetcher for RecordingFetcher { + fn fetch( + &mut self, + request: &fetch::FetchShard, + ) -> Result { + self.calls.push(request.clone()); + self.result.clone() + } +} + +struct RecordingBinder { + calls: Vec, + result: Result<(), lifecycle::BindError>, +} + +impl lifecycle::WorkerShardBinder for RecordingBinder { + fn bind(&mut self, shard: &shards::ValidatedShard) -> Result<(), lifecycle::BindError> { + self.calls.push(shard.clone()); + self.result.clone() + } +} + +#[test] +fn lifecycle_happy_path_reaches_ready() { + let assignment = assignment(); + let fetched = fetched_for(&assignment); + let mut cache = MemoryCache::default(); + let mut fetcher = RecordingFetcher { + calls: Vec::new(), + result: Ok(fetched.clone()), + }; + let mut binder = RecordingBinder { + calls: Vec::new(), + result: Ok(()), + }; + let mut lifecycle = lifecycle::ShardWeightLifecycle::new(); + + lifecycle.load(assignment.clone(), &mut cache, &mut fetcher, &mut binder); + + assert_eq!(lifecycle.state(), lifecycle::ShardLifecycleState::Ready); + assert_eq!( + event_kinds(lifecycle.events()), + vec![ + "assigned", + "located", + "fetching", + "fetched", + "validated", + "binding", + "ready" + ] + ); + assert_eq!(fetcher.calls.len(), 1); + assert_eq!(cache.inserts, 1); + assert_eq!(binder.calls.len(), 1); + assert_eq!(binder.calls[0].assignment, assignment); + assert_eq!(binder.calls[0].local_path, fetched.local_path); +} + +#[test] +fn lifecycle_cache_hit_skips_fetch_but_still_validates_and_binds() { + let assignment = assignment(); + let fetched = fetched_for(&assignment); + let location = fetch::ShardLocator::locate(&assignment); + let mut cache = MemoryCache::default(); + cache.entries.insert(location.cache_key.clone(), fetched); + let mut fetcher = RecordingFetcher { + calls: Vec::new(), + result: Err(fetch::FetchError::Unavailable), + }; + let mut binder = RecordingBinder { + calls: Vec::new(), + result: Ok(()), + }; + let mut lifecycle = lifecycle::ShardWeightLifecycle::new(); + + lifecycle.load(assignment, &mut cache, &mut fetcher, &mut binder); + + assert_eq!(lifecycle.state(), lifecycle::ShardLifecycleState::Ready); + assert_eq!( + event_kinds(lifecycle.events()), + vec![ + "assigned", + "located", + "fetching", + "cache-hit", + "validated", + "binding", + "ready" + ] + ); + assert!(fetcher.calls.is_empty()); + assert_eq!(cache.inserts, 0); + assert_eq!(binder.calls.len(), 1); +} + +#[test] +fn lifecycle_fetch_failure_faults_without_validation_or_bind() { + let assignment = assignment(); + let mut cache = MemoryCache::default(); + let mut fetcher = RecordingFetcher { + calls: Vec::new(), + result: Err(fetch::FetchError::NotFound), + }; + let mut binder = RecordingBinder { + calls: Vec::new(), + result: Ok(()), + }; + let mut lifecycle = lifecycle::ShardWeightLifecycle::new(); + + lifecycle.load(assignment, &mut cache, &mut fetcher, &mut binder); + + assert_eq!(lifecycle.state(), lifecycle::ShardLifecycleState::Faulted); + assert_eq!(fetcher.calls.len(), 1); + assert!(binder.calls.is_empty()); + assert!(matches!( + lifecycle.events().last(), + Some(lifecycle::ShardLifecycleEvent::Faulted { + reason: lifecycle::ShardLifecycleFault::Fetch(fetch::FetchError::NotFound) + }) + )); +} + +#[test] +fn lifecycle_validation_failure_faults_without_bind() { + let assignment = assignment(); + let mut bad_manifest = shards::ShardManifest::for_assignment( + &assignment, + shards::ContentHash::literal("sha256:stage-3").unwrap(), + ); + bad_manifest.stage_index = 4; + let mut cache = MemoryCache::default(); + let mut fetcher = RecordingFetcher { + calls: Vec::new(), + result: Ok(fetch::FetchedShard::new( + "/cache/bad-stage.gguf", + bad_manifest, + )), + }; + let mut binder = RecordingBinder { + calls: Vec::new(), + result: Ok(()), + }; + let mut lifecycle = lifecycle::ShardWeightLifecycle::new(); + + lifecycle.load(assignment, &mut cache, &mut fetcher, &mut binder); + + assert_eq!(lifecycle.state(), lifecycle::ShardLifecycleState::Faulted); + assert!(binder.calls.is_empty()); + assert!(matches!( + lifecycle.events().last(), + Some(lifecycle::ShardLifecycleEvent::Faulted { + reason: lifecycle::ShardLifecycleFault::Validation( + shards::ShardValidationError::StageIndexMismatch + ) + }) + )); +} + +#[test] +fn lifecycle_bind_failure_faults_after_validation() { + let assignment = assignment(); + let fetched = fetched_for(&assignment); + let mut cache = MemoryCache::default(); + let mut fetcher = RecordingFetcher { + calls: Vec::new(), + result: Ok(fetched), + }; + let mut binder = RecordingBinder { + calls: Vec::new(), + result: Err(lifecycle::BindError::WorkerRejected), + }; + let mut lifecycle = lifecycle::ShardWeightLifecycle::new(); + + lifecycle.load(assignment, &mut cache, &mut fetcher, &mut binder); + + assert_eq!(lifecycle.state(), lifecycle::ShardLifecycleState::Faulted); + assert_eq!(binder.calls.len(), 1); + assert!(event_kinds(lifecycle.events()).contains(&"validated")); + assert!(matches!( + lifecycle.events().last(), + Some(lifecycle::ShardLifecycleEvent::Faulted { + reason: lifecycle::ShardLifecycleFault::Bind(lifecycle::BindError::WorkerRejected) + }) + )); +} + +#[test] +fn ready_terminal_state_is_idempotent() { + let assignment = assignment(); + let fetched = fetched_for(&assignment); + let mut cache = MemoryCache::default(); + let mut fetcher = RecordingFetcher { + calls: Vec::new(), + result: Ok(fetched.clone()), + }; + let mut binder = RecordingBinder { + calls: Vec::new(), + result: Ok(()), + }; + let mut lifecycle = lifecycle::ShardWeightLifecycle::new(); + + lifecycle.load(assignment.clone(), &mut cache, &mut fetcher, &mut binder); + let event_count = lifecycle.events().len(); + lifecycle.load(assignment, &mut cache, &mut fetcher, &mut binder); + + assert_eq!(lifecycle.state(), lifecycle::ShardLifecycleState::Ready); + assert_eq!(lifecycle.events().len(), event_count); + assert_eq!(fetcher.calls.len(), 1); + assert_eq!(binder.calls.len(), 1); +} + +#[test] +fn fault_terminal_state_is_idempotent() { + let assignment = assignment(); + let fetched = fetched_for(&assignment); + let mut cache = MemoryCache::default(); + let mut fetcher = RecordingFetcher { + calls: Vec::new(), + result: Err(fetch::FetchError::Unavailable), + }; + let mut binder = RecordingBinder { + calls: Vec::new(), + result: Ok(()), + }; + let mut lifecycle = lifecycle::ShardWeightLifecycle::new(); + + lifecycle.load(assignment.clone(), &mut cache, &mut fetcher, &mut binder); + let event_count = lifecycle.events().len(); + fetcher.result = Ok(fetched); + lifecycle.load(assignment, &mut cache, &mut fetcher, &mut binder); + + assert_eq!(lifecycle.state(), lifecycle::ShardLifecycleState::Faulted); + assert_eq!(lifecycle.events().len(), event_count); + assert_eq!(fetcher.calls.len(), 1); + assert!(binder.calls.is_empty()); +} diff --git a/crates/mvp-system/src/tests/telemetry_guarantees.rs b/crates/mvp-system/src/tests/telemetry_guarantees.rs index 6d526db..a610929 100644 --- a/crates/mvp-system/src/tests/telemetry_guarantees.rs +++ b/crates/mvp-system/src/tests/telemetry_guarantees.rs @@ -3,8 +3,9 @@ use mvp_system::observability_surface as obs; use mvp_system::provisioning::{ ProvisionEvent, ProvisionEventKind, ProvisionLogLine, ProvisionLogStream, }; -use mvp_system::telemetry::{self, MvpLifecycleRecord}; -use mvp_system::telemetry::{MvpProvisionEventRecord, MvpProvisionLogRecord}; +use mvp_system::telemetry::{ + self, MvpLifecycleRecord, MvpProvisionEventRecord, MvpProvisionLogRecord, +}; #[test] fn mvp_lifecycle_record_round_trips_on_owned_datastream_channel() { @@ -63,6 +64,10 @@ fn provisioning_records_round_trip_on_owned_datastream_channels() { event ); assert_eq!(MvpProvisionLogRecord::decode(&log.encode()).unwrap(), log); + assert_eq!( + telemetry::mvp_provision_log_channel(11, ProvisionLogStream::Stdout).as_str(), + "mvp.provisioning.logs.node.11.stdout" + ); } #[test] diff --git a/crates/mvp-system/src/tests/weight_shards_guarantees.rs b/crates/mvp-system/src/tests/weight_shards_guarantees.rs new file mode 100644 index 0000000..90025af --- /dev/null +++ b/crates/mvp-system/src/tests/weight_shards_guarantees.rs @@ -0,0 +1,134 @@ +use mvp_system::weight_shards as shards; + +fn model_ref() -> shards::ModelArtifactRef { + shards::ModelArtifactRef::parse("hf://org/repo@abcdef123456/model.gguf").unwrap() +} + +fn assignment() -> shards::ShardAssignment { + let model_ref = model_ref(); + let split_scheme = shards::SplitScheme::GgufLayerContiguousV1; + let split_id = shards::SplitId::derive(&model_ref, split_scheme); + shards::ShardAssignment::new( + model_ref, + split_id, + split_scheme, + 3, + 8, + shards::LayerRange::new(12, 16).unwrap(), + ) + .unwrap() +} + +fn content_hash() -> shards::ContentHash { + shards::ContentHash::literal("sha256:test-content").unwrap() +} + +#[test] +fn model_ref_canonicalization_is_stable() { + let parsed = shards::ModelArtifactRef::parse("hf://org/repo@abcdef123456/model.gguf").unwrap(); + let from_parts = + shards::ModelArtifactRef::hugging_face("/org/repo/", "abcdef123456", "/model.gguf") + .unwrap(); + + assert_eq!(parsed, from_parts); + assert_eq!(parsed.as_str(), "hf://org/repo@abcdef123456/model.gguf"); + assert_eq!(parsed.repo(), "org/repo"); + assert_eq!(parsed.revision(), "abcdef123456"); + assert_eq!(parsed.path(), "model.gguf"); +} + +#[test] +fn split_id_is_deterministic_and_model_sensitive() { + let first = model_ref(); + let same = shards::ModelArtifactRef::parse("hf://org/repo@abcdef123456/model.gguf").unwrap(); + let different = + shards::ModelArtifactRef::parse("hf://org/repo@fedcba654321/model.gguf").unwrap(); + let scheme = shards::SplitScheme::GgufLayerContiguousV1; + + assert_eq!( + shards::SplitId::derive(&first, scheme), + shards::SplitId::derive(&same, scheme) + ); + assert_ne!( + shards::SplitId::derive(&first, scheme), + shards::SplitId::derive(&different, scheme) + ); +} + +#[test] +fn assignment_rejects_invalid_stage_shape_and_ranges() { + let model_ref = model_ref(); + let scheme = shards::SplitScheme::GgufLayerContiguousV1; + let split_id = shards::SplitId::derive(&model_ref, scheme); + let range = shards::LayerRange::new(1, 2).unwrap(); + + assert_eq!( + shards::ShardAssignment::new(model_ref.clone(), split_id.clone(), scheme, 0, 0, range), + Err(shards::ShardAssignmentError::EmptyStageCount) + ); + assert_eq!( + shards::ShardAssignment::new(model_ref, split_id, scheme, 2, 2, range), + Err(shards::ShardAssignmentError::StageIndexOutOfRange) + ); + assert_eq!( + shards::LayerRange::new(4, 4), + Err(shards::LayerRangeError::EmptyOrInverted) + ); + assert_eq!( + shards::LayerRange::new(5, 4), + Err(shards::LayerRangeError::EmptyOrInverted) + ); +} + +#[test] +fn validator_accepts_matching_manifest() { + let assignment = assignment(); + let manifest = shards::ShardManifest::for_assignment(&assignment, content_hash()); + + assert_eq!( + shards::ShardValidator::validate(&assignment, &manifest), + Ok(()) + ); + assert!(shards::ValidatedShard::new(assignment, manifest, "/cache/stage-00003.gguf").is_ok()); +} + +#[test] +fn validator_rejects_mismatched_manifest() { + let assignment = assignment(); + let matching = shards::ShardManifest::for_assignment(&assignment, content_hash()); + + let mut wrong_model = matching.clone(); + wrong_model.model_digest = shards::ModelDigest::literal("wrong-model").unwrap(); + assert_eq!( + shards::ShardValidator::validate(&assignment, &wrong_model), + Err(shards::ShardValidationError::ModelDigestMismatch) + ); + + let mut wrong_split = matching.clone(); + wrong_split.split_id = shards::SplitId::literal("split-wrong").unwrap(); + assert_eq!( + shards::ShardValidator::validate(&assignment, &wrong_split), + Err(shards::ShardValidationError::SplitIdMismatch) + ); + + let mut wrong_stage = matching.clone(); + wrong_stage.stage_index = 4; + assert_eq!( + shards::ShardValidator::validate(&assignment, &wrong_stage), + Err(shards::ShardValidationError::StageIndexMismatch) + ); + + let mut wrong_count = matching.clone(); + wrong_count.stage_count = 9; + assert_eq!( + shards::ShardValidator::validate(&assignment, &wrong_count), + Err(shards::ShardValidationError::StageCountMismatch) + ); + + let mut wrong_range = matching; + wrong_range.layer_range = shards::LayerRange::new(16, 20).unwrap(); + assert_eq!( + shards::ShardValidator::validate(&assignment, &wrong_range), + Err(shards::ShardValidationError::LayerRangeMismatch) + ); +} diff --git a/crates/mvp-system/src/weight_shards.rs b/crates/mvp-system/src/weight_shards.rs new file mode 100644 index 0000000..44048d2 --- /dev/null +++ b/crates/mvp-system/src/weight_shards.rs @@ -0,0 +1,348 @@ +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct ModelArtifactRef { + canonical: String, + repo: String, + revision: String, + path: String, +} + +impl ModelArtifactRef { + pub fn parse(value: impl Into) -> Result { + let value = value.into(); + let rest = value + .strip_prefix("hf://") + .ok_or(ModelArtifactRefError::InvalidScheme)?; + let (repo, revision_and_path) = rest + .split_once('@') + .ok_or(ModelArtifactRefError::MissingRevision)?; + let (revision, path) = revision_and_path + .split_once('/') + .ok_or(ModelArtifactRefError::MissingPath)?; + Self::hugging_face(repo, revision, path) + } + + pub fn hugging_face( + repo: impl Into, + revision: impl Into, + path: impl Into, + ) -> Result { + let repo = repo.into().trim_matches('/').to_owned(); + let revision = revision.into(); + let path = path.into().trim_start_matches('/').to_owned(); + if repo.is_empty() { + return Err(ModelArtifactRefError::MissingRepo); + } + if revision.is_empty() { + return Err(ModelArtifactRefError::MissingRevision); + } + if revision.contains('/') { + return Err(ModelArtifactRefError::RevisionMustBePathSegment); + } + if path.is_empty() { + return Err(ModelArtifactRefError::MissingPath); + } + let canonical = format!("hf://{repo}@{revision}/{path}"); + Ok(Self { + canonical, + repo, + revision, + path, + }) + } + + pub fn as_str(&self) -> &str { + &self.canonical + } + + pub fn repo(&self) -> &str { + &self.repo + } + + pub fn revision(&self) -> &str { + &self.revision + } + + pub fn path(&self) -> &str { + &self.path + } + + pub fn model_digest(&self) -> ModelDigest { + ModelDigest(stable_digest_hex(&["model", self.as_str()])) + } + + pub fn shard_uri(&self, split_id: &SplitId, stage_index: u32) -> String { + format!( + "hf://{}@{}/shards/{}/stage-{stage_index:05}.gguf", + self.repo, + self.revision, + split_id.as_str(), + ) + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum ModelArtifactRefError { + InvalidScheme, + MissingRepo, + MissingRevision, + RevisionMustBePathSegment, + MissingPath, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum SplitScheme { + GgufLayerContiguousV1, +} + +impl SplitScheme { + pub fn as_str(&self) -> &'static str { + match self { + Self::GgufLayerContiguousV1 => "gguf-layer-contiguous-v1", + } + } +} + +#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct SplitId(String); + +impl SplitId { + pub fn derive(model_ref: &ModelArtifactRef, scheme: SplitScheme) -> Self { + Self(format!( + "split-{}", + stable_digest_hex(&["split", model_ref.as_str(), scheme.as_str()]) + )) + } + + pub fn literal(value: impl Into) -> Result { + let value = value.into(); + if value.is_empty() { + return Err(SplitIdError::Empty); + } + if !value + .bytes() + .all(|byte| byte.is_ascii_alphanumeric() || byte == b'-' || byte == b'_') + { + return Err(SplitIdError::InvalidCharacter); + } + Ok(Self(value)) + } + + pub fn as_str(&self) -> &str { + &self.0 + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum SplitIdError { + Empty, + InvalidCharacter, +} + +#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct ModelDigest(String); + +impl ModelDigest { + pub fn literal(value: impl Into) -> Result { + let value = value.into(); + if value.is_empty() { + return Err(ModelDigestError::Empty); + } + Ok(Self(value)) + } + + pub fn as_str(&self) -> &str { + &self.0 + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum ModelDigestError { + Empty, +} + +#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct ContentHash(String); + +impl ContentHash { + pub fn literal(value: impl Into) -> Result { + let value = value.into(); + if value.is_empty() { + return Err(ContentHashError::Empty); + } + Ok(Self(value)) + } + + pub fn as_str(&self) -> &str { + &self.0 + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum ContentHashError { + Empty, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct LayerRange { + pub start: u32, + pub end_exclusive: u32, +} + +impl LayerRange { + pub fn new(start: u32, end_exclusive: u32) -> Result { + if start >= end_exclusive { + return Err(LayerRangeError::EmptyOrInverted); + } + Ok(Self { + start, + end_exclusive, + }) + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum LayerRangeError { + EmptyOrInverted, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct ShardAssignment { + pub model_ref: ModelArtifactRef, + pub split_id: SplitId, + pub split_scheme: SplitScheme, + pub stage_index: u32, + pub stage_count: u32, + pub layer_range: LayerRange, +} + +impl ShardAssignment { + pub fn new( + model_ref: ModelArtifactRef, + split_id: SplitId, + split_scheme: SplitScheme, + stage_index: u32, + stage_count: u32, + layer_range: LayerRange, + ) -> Result { + if stage_count == 0 { + return Err(ShardAssignmentError::EmptyStageCount); + } + if stage_index >= stage_count { + return Err(ShardAssignmentError::StageIndexOutOfRange); + } + Ok(Self { + model_ref, + split_id, + split_scheme, + stage_index, + stage_count, + layer_range, + }) + } + + pub fn expected_model_digest(&self) -> ModelDigest { + self.model_ref.model_digest() + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum ShardAssignmentError { + EmptyStageCount, + StageIndexOutOfRange, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct ShardManifest { + pub model_digest: ModelDigest, + pub split_id: SplitId, + pub stage_index: u32, + pub stage_count: u32, + pub layer_range: LayerRange, + pub content_hash: ContentHash, +} + +impl ShardManifest { + pub fn for_assignment(assignment: &ShardAssignment, content_hash: ContentHash) -> Self { + Self { + model_digest: assignment.expected_model_digest(), + split_id: assignment.split_id.clone(), + stage_index: assignment.stage_index, + stage_count: assignment.stage_count, + layer_range: assignment.layer_range, + content_hash, + } + } +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct ValidatedShard { + pub assignment: ShardAssignment, + pub manifest: ShardManifest, + pub local_path: String, +} + +impl ValidatedShard { + pub fn new( + assignment: ShardAssignment, + manifest: ShardManifest, + local_path: impl Into, + ) -> Result { + ShardValidator::validate(&assignment, &manifest)?; + Ok(Self { + assignment, + manifest, + local_path: local_path.into(), + }) + } +} + +pub struct ShardValidator; + +impl ShardValidator { + pub fn validate( + assignment: &ShardAssignment, + manifest: &ShardManifest, + ) -> Result<(), ShardValidationError> { + if manifest.model_digest != assignment.expected_model_digest() { + return Err(ShardValidationError::ModelDigestMismatch); + } + if manifest.split_id != assignment.split_id { + return Err(ShardValidationError::SplitIdMismatch); + } + if manifest.stage_index != assignment.stage_index { + return Err(ShardValidationError::StageIndexMismatch); + } + if manifest.stage_count != assignment.stage_count { + return Err(ShardValidationError::StageCountMismatch); + } + if manifest.layer_range != assignment.layer_range { + return Err(ShardValidationError::LayerRangeMismatch); + } + Ok(()) + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum ShardValidationError { + ModelDigestMismatch, + SplitIdMismatch, + StageIndexMismatch, + StageCountMismatch, + LayerRangeMismatch, +} + +fn stable_digest_hex(parts: &[&str]) -> String { + const FNV_OFFSET: u64 = 0xcbf29ce484222325; + const FNV_PRIME: u64 = 0x100000001b3; + + let mut hash = FNV_OFFSET; + for part in parts { + for byte in part.as_bytes() { + hash ^= u64::from(*byte); + hash = hash.wrapping_mul(FNV_PRIME); + } + hash ^= 0xff; + hash = hash.wrapping_mul(FNV_PRIME); + } + format!("{hash:016x}") +}