diff --git a/crates/iroh-driver/src/edge_transport.rs b/crates/iroh-driver/src/edge_transport.rs index dc36399..3887471 100644 --- a/crates/iroh-driver/src/edge_transport.rs +++ b/crates/iroh-driver/src/edge_transport.rs @@ -5,6 +5,7 @@ //! ring ownership, and stage semantics stay in the MVP/dataplane crates. use std::sync::Arc; +use std::time::Duration; use distribution::types::NodeId; use iroh::endpoint::Connection; @@ -72,28 +73,70 @@ pub(crate) fn spawn_edge_send_pump( let (ready_tx, ready_rx) = std::sync::mpsc::channel::>(); handle.spawn(async move { let result: Result<(), String> = async { - let conn = endpoint - .connect(peer, EDGE_ALPN) - .await - .map_err(|e| format!("connect edge {edge_id}: {e}"))?; - let mut send = conn - .open_uni() - .await - .map_err(|e| format!("open edge stream {edge_id}: {e}"))?; - send.write_all(&encode_edge_preamble(edge_id)) - .await - .map_err(|e| format!("write edge preamble {edge_id}: {e}"))?; - send.flush() - .await - .map_err(|e| format!("flush edge preamble {edge_id}: {e}"))?; + macro_rules! open_edge_stream { + () => {{ + let conn = endpoint + .connect(peer.clone(), EDGE_ALPN) + .await + .map_err(|e| format!("connect edge {edge_id}: {e}"))?; + let mut send = conn + .open_uni() + .await + .map_err(|e| format!("open edge stream {edge_id}: {e}"))?; + send.write_all(&encode_edge_preamble(edge_id)) + .await + .map_err(|e| format!("write edge preamble {edge_id}: {e}"))?; + send.flush() + .await + .map_err(|e| format!("flush edge preamble {edge_id}: {e}"))?; + send + }}; + } + + let mut send = open_edge_stream!(); let _ = ready_tx.send(Ok(())); while let Some(record) = rx.recv().await { - send.write_all(&record) + let mut attempts = 0_u8; + loop { + attempts = attempts.saturating_add(1); + let write_result = tokio::time::timeout(Duration::from_secs(30), async { + send.write_all(&record) + .await + .map_err(|e| format!("write edge record {edge_id}: {e}"))?; + send.flush() + .await + .map_err(|e| format!("flush edge record {edge_id}: {e}")) + }) .await - .map_err(|e| format!("write edge record {edge_id}: {e}"))?; - send.flush() - .await - .map_err(|e| format!("flush edge record {edge_id}: {e}"))?; + .map_err(|_| format!("write edge record {edge_id}: timed out"))?; + + match write_result { + Ok(()) => break, + Err(error) if attempts < 3 => { + send = open_edge_stream!(); + let retry_result = + tokio::time::timeout(Duration::from_secs(30), async { + send.write_all(&record).await.map_err(|e| { + format!("write edge record {edge_id} after reconnect: {e}") + })?; + send.flush().await.map_err(|e| { + format!("flush edge record {edge_id} after reconnect: {e}") + }) + }) + .await + .map_err(|_| { + format!( + "write edge record {edge_id} after reconnect: timed out" + ) + })?; + retry_result.map_err(|retry_error| { + format!("{error}; reconnect write failed: {retry_error}") + })?; + break; + } + Err(error) => return Err(error), + } + } } send.finish() .map_err(|e| format!("finish edge stream {edge_id}: {e}"))?; diff --git a/crates/mvp-system/specs/MVP_SYSTEM_SPEC.md b/crates/mvp-system/specs/MVP_SYSTEM_SPEC.md index ea585c1..72ae2d4 100644 --- a/crates/mvp-system/specs/MVP_SYSTEM_SPEC.md +++ b/crates/mvp-system/specs/MVP_SYSTEM_SPEC.md @@ -653,20 +653,21 @@ The MVP allows one active `ExecuteStep` per stage. Weights are stage-local persistent state for the run. -The StageController starts weight work from the assigned `GgufSource` and layer -range. A stage may download: +The orchestrator sends `ProvisionStage` to every runtime-ready pipeline stage +without waiting for another stage's weights. Each provision carries enough +weight-source planning information for that stage's assigned layer range: the +source identity, selected artifact ranges or shard identifiers, required +metadata, and cache key. The node materializes or locates the stage-local weight +artifact from the assigned source; stages do not stream weights to one another. -- a whole GGUF and load only its range +The StageController starts weight work from the assigned weight source and layer +range. A stage may acquire: + +- only the artifact ranges required for its stage-local shard - one or more physical shards containing its range +- a whole artifact and load only its range when no shard plan is available - a cached artifact that already exists on the node -The physical loading mechanism may be: - -- worker startup configuration -- a `ConfigureRole` command -- a local loader path owned by the StageController -- weight objects moved through the same object/ring machinery - The system-visible contract is `WeightsReady` before `StageReady`. `WeightsReady` means: diff --git a/crates/mvp-system/src/bin/mvp_chat.rs b/crates/mvp-system/src/bin/mvp_chat.rs index 5d3cbca..45bd707 100644 --- a/crates/mvp-system/src/bin/mvp_chat.rs +++ b/crates/mvp-system/src/bin/mvp_chat.rs @@ -474,6 +474,7 @@ impl ChatDatastream { "max_dph_total": vastai.max_dph_total, "min_reliability": vastai.min_reliability, "require_verified": vastai.require_verified, + "blacklist_hosts": &vastai.blacklist_hosts, "disk_gb": vastai.disk_gb, "has_onstart": vastai.onstart.is_some(), "has_ssh_identity": vastai.ssh_identity.is_some(), @@ -958,6 +959,9 @@ impl Config { "--no-vastai-require-verified".to_owned() }); } + for host_id in &vastai.blacklist_hosts { + args.extend(["--vastai-blacklist-host".to_owned(), host_id.to_string()]); + } if let Some(onstart) = &vastai.onstart { args.extend(["--vastai-onstart".to_owned(), onstart.clone()]); } @@ -2944,6 +2948,7 @@ node = "docker.io/acme/node:latest" [vastai] relay_url = "https://relay.example" bootstrap_command = "boot" +blacklist_hosts = [155385, 546483] "#, ); with_process_state( @@ -2958,6 +2963,7 @@ bootstrap_command = "boot" assert_eq!(vastai.relay_url, "https://relay.example"); assert_eq!(vastai.bootstrap_command, "boot"); assert_eq!(vastai.image, "docker.io/acme/node:latest"); + assert_eq!(vastai.blacklist_hosts, vec![155385, 546483]); let args = config.orchestrator_cli_args("docker.io/acme/node:latest"); assert!( !args @@ -2970,6 +2976,14 @@ bootstrap_command = "boot" .any(|pair| pair == ["--vastai-bootstrap-command", "boot"]), "non-secret Vast.ai config should still be forwarded" ); + assert!( + args.windows(2) + .any(|pair| pair == ["--vastai-blacklist-host", "155385"]) + && args + .windows(2) + .any(|pair| pair == ["--vastai-blacklist-host", "546483"]), + "Vast.ai host blacklist must be forwarded to orchestrator argv: {args:?}" + ); }, ); } diff --git a/crates/mvp-system/src/orchestrator_app.rs b/crates/mvp-system/src/orchestrator_app.rs index af5d49d..0700151 100644 --- a/crates/mvp-system/src/orchestrator_app.rs +++ b/crates/mvp-system/src/orchestrator_app.rs @@ -20,6 +20,8 @@ use crate::benchmark_observability; use crate::config::{DEFAULT_CONFIG_PATH, TomlConfigOverlay}; #[cfg(feature = "dashboard")] use crate::dashboard_view::MvpClusterDashboardView; +const PROVIDER_START_MAX_ATTEMPTS: usize = 4; + use crate::distribution_stack::DistributionRuntimeStack; use crate::endpoint_advertisement::{ EndpointAddrMask, MVP_IROH_ENDPOINT_ADDR_MASK_ENV, advertised_endpoint, @@ -87,6 +89,8 @@ const RUNTIME_READY_ACK_RETRY_INTERVAL: Duration = Duration::from_millis(250); const RUNTIME_READY_ACK_TIMEOUT: Duration = Duration::from_secs(60); const STAGE_PROVISION_ACTIVE_RESEND_AFTER: Duration = Duration::from_secs(60); const RUNTIME_READY_TIMEOUT: Duration = Duration::from_secs(60); +const PIPELINE_PROMPT_IDLE_TIMEOUT: Duration = Duration::from_secs(120); +const PIPELINE_PROMPT_WAIT_LOG_INTERVAL: Duration = Duration::from_secs(15); const MVP_ORCH_BOOTSTRAP: &str = "mvp.orch.bootstrap"; const MVP_ORCH_PROMPT: &str = "mvp.orch.prompt"; const MVP_SWIM_MEMBERSHIP: &str = "mvp.swim.membership"; @@ -676,6 +680,11 @@ impl VastAiRuntimeConfig { if let Some(require_verified) = require_verified { provisioning.selection.require_verified = require_verified; } + for host_id in &builder.vastai_blacklist_hosts { + if !provisioning.selection.blacklist_hosts.contains(host_id) { + provisioning.selection.blacklist_hosts.push(*host_id); + } + } let poll_interval_secs = builder .vastai_poll_interval_secs_raw .as_ref() @@ -712,6 +721,7 @@ impl VastAiRuntimeConfig { "max_dph_total": self.provisioning.selection.max_dph_total, "min_reliability": self.provisioning.selection.min_reliability, "require_verified": self.provisioning.selection.require_verified, + "blacklist_hosts": &self.provisioning.selection.blacklist_hosts, "state_timeout_secs": self.provisioning.lifecycle.state_timeout.as_secs(), "confirm_lease": self.provisioning.confirm_lease, "has_api_key": self.api_key.is_some(), @@ -937,6 +947,7 @@ struct ConfigBuilder { vastai_require_verified: Option, vastai_require_verified_raw: Option, vastai_poll_interval_secs: Option, + vastai_blacklist_hosts: Vec, vastai_poll_interval_secs_raw: Option, cached_model_host_path: Option, datastream_frame_log: Option, @@ -995,6 +1006,7 @@ impl ConfigBuilder { vastai_require_verified: None, vastai_require_verified_raw: None, vastai_poll_interval_secs: None, + vastai_blacklist_hosts: Vec::new(), vastai_poll_interval_secs_raw: None, cached_model_host_path: None, datastream_frame_log: None, @@ -1118,6 +1130,9 @@ impl ConfigBuilder { if let Some(require_verified) = overlay.vastai.require_verified { self.vastai_require_verified = Some(require_verified); } + for host_id in overlay.vastai.blacklist_hosts { + self.push_vastai_blacklist_host(host_id); + } if let Some(poll_interval_secs) = overlay.vastai.poll_interval_secs { self.vastai_poll_interval_secs = Some(poll_interval_secs); } @@ -1253,6 +1268,11 @@ impl ConfigBuilder { if let Some(require_verified) = env_optional("MVP_VASTAI_REQUIRE_VERIFIED") { self.vastai_require_verified_raw = Some(require_verified); } + if let Some(blacklist_hosts) = env_optional("MVP_VASTAI_BLACKLIST_HOSTS") { + for host_id in Self::parse_list("MVP_VASTAI_BLACKLIST_HOSTS", &blacklist_hosts)? { + self.push_vastai_blacklist_host(host_id); + } + } if let Some(poll_interval_secs) = env_optional("MVP_VASTAI_POLL_INTERVAL_SECS") { self.vastai_poll_interval_secs_raw = Some(poll_interval_secs); } @@ -1390,6 +1410,10 @@ impl ConfigBuilder { self.vastai_require_verified = Some(false); self.vastai_require_verified_raw = None; } + "--vastai-blacklist-host" => { + let host_id = parse_next(&mut args, "--vastai-blacklist-host")?; + self.push_vastai_blacklist_host(host_id); + } "--vastai-poll-interval-secs" => { self.vastai_poll_interval_secs = Some(parse_next(&mut args, "--vastai-poll-interval-secs")?); @@ -1516,6 +1540,25 @@ impl ConfigBuilder { }; } + fn push_vastai_blacklist_host(&mut self, host_id: u64) { + if !self.vastai_blacklist_hosts.contains(&host_id) { + self.vastai_blacklist_hosts.push(host_id); + } + } + + fn parse_list(name: &str, value: &str) -> Result, String> + where + T: std::str::FromStr, + T::Err: std::fmt::Display, + { + value + .split(',') + .map(str::trim) + .filter(|part| !part.is_empty()) + .map(|part| Self::parse_value(name, part)) + .collect() + } + fn parse_value(name: &str, value: &str) -> Result where T: std::str::FromStr, @@ -2259,13 +2302,6 @@ impl ProvisionedClusterGuard { } } - fn complete_bootstrap_all(&mut self) -> Result<(), String> { - for handle in &self.handles { - self.provisioner.complete_bootstrap(handle)?; - } - Ok(()) - } - fn stop(&mut self) -> Result<(), String> { let mut first_error = None; while let Some(handle) = self.handles.pop() { @@ -2406,86 +2442,107 @@ fn start_and_provision_workers( } else { BTreeMap::new() }; - for node_spec in &stage_specs { - orch_datastream.emit_event( - dashboard, - ProvisionEvent { - run_id: config.run_id, - node_id: node_spec.node_id, - kind: ProvisionEventKind::ProvisionStart, - provider: Some(config.provider.as_str().to_owned()), - message: Some(format!( - "starting {} image {}", - config.provider.as_str(), - config.image - )), - }, - ); - orch_datastream.emit_bootstrap( + let mut handles = Vec::with_capacity(stage_specs.len()); + let mut pending_specs = stage_specs; + for attempt in 1..=PROVIDER_START_MAX_ATTEMPTS { + for node_spec in &pending_specs { + orch_datastream.emit_event( + dashboard, + ProvisionEvent { + run_id: config.run_id, + node_id: node_spec.node_id, + kind: ProvisionEventKind::ProvisionStart, + provider: Some(config.provider.as_str().to_owned()), + message: Some(format!( + "starting {} image {}", + config.provider.as_str(), + config.image + )), + }, + ); + orch_datastream.emit_bootstrap( + dashboard, + config.run_id, + config.node_id, + "provider_start", + "started", + json!({ + "provider":config.provider.as_str(), + "image":&config.image, + "node_id":node_spec.node_id, + "stage_index":node_spec.stage_index, + "attempt":attempt, + }), + ); + } + let (returned_provisioner, start_results) = start_nodes_with_stdio_capture( + provisioner, + pending_specs, + sink.clone(), + orch_stdio_rx, dashboard, + orch_datastream, config.run_id, config.node_id, - "provider_start", - "started", - json!({ - "provider":config.provider.as_str(), - "image":&config.image, - "node_id":node_spec.node_id, - "stage_index":node_spec.stage_index, - }), ); - } - let (returned_provisioner, start_results) = start_nodes_with_stdio_capture( - provisioner, - stage_specs, - sink.clone(), - orch_stdio_rx, - dashboard, - orch_datastream, - config.run_id, - config.node_id, - ); - provisioner = returned_provisioner; - let mut handles = Vec::with_capacity(start_results.len()); - for (node_spec, handle_result) in start_results { - match handle_result { - Ok(handle) => { - orch_datastream.emit_bootstrap( - dashboard, - config.run_id, - config.node_id, - "provider_start", - "ready", - json!({ - "provider":config.provider.as_str(), - "node_id":node_spec.node_id, - "stage_index":node_spec.stage_index, - }), - ); - handles.push(handle); - } - Err(error) => { - orch_datastream.emit_bootstrap( - dashboard, - config.run_id, - config.node_id, - "provider_start", - "failed", - json!({"provider":config.provider.as_str(),"node_id":node_spec.node_id,"error":error}), - ); - stop_started_nodes(&mut *provisioner, &mut handles); - drain_orch_stdio_capture( - orch_stdio_rx, - orch_datastream, - dashboard, - config.run_id, - config.node_id, - ); - return Err(error); + provisioner = returned_provisioner; + let start_outcome = collect_provider_start_outcome(start_results); + handles.extend(start_outcome.successful_handles); + for (node_spec, handle_result) in start_outcome.results { + match handle_result { + Ok(_) => { + orch_datastream.emit_bootstrap( + dashboard, + config.run_id, + config.node_id, + "provider_start", + "ready", + json!({ + "provider":config.provider.as_str(), + "node_id":node_spec.node_id, + "stage_index":node_spec.stage_index, + "attempt":attempt, + }), + ); + } + Err(error) => { + orch_datastream.emit_bootstrap( + dashboard, + config.run_id, + config.node_id, + "provider_start", + "failed", + json!({ + "provider":config.provider.as_str(), + "node_id":node_spec.node_id, + "stage_index":node_spec.stage_index, + "attempt":attempt, + "error":error, + }), + ); + } } } + if start_outcome.first_error.is_none() { + break; + } + if attempt == PROVIDER_START_MAX_ATTEMPTS { + let error = start_outcome + .first_error + .expect("checked provider-start failure"); + stop_started_nodes(&mut *provisioner, &mut handles); + drain_orch_stdio_capture( + orch_stdio_rx, + orch_datastream, + dashboard, + config.run_id, + config.node_id, + ); + return Err(error); + } + pending_specs = start_outcome.failed_specs; } - let mut provisioned_nodes = ProvisionedClusterGuard::new(provisioner, handles); + let provisioned_nodes = ProvisionedClusterGuard::new(provisioner, handles); drain_orch_stdio_capture( orch_stdio_rx, orch_datastream, @@ -2572,7 +2629,6 @@ fn start_and_provision_workers( json!({"endpoint":&ready.endpoint,"node_actor":ready.node_actor,"node_id":node_id,"stage_index":ready.stage_index}), ); } - provisioned_nodes.complete_bootstrap_all()?; let ack_targets = readies .iter() .map(|(node_id, ready)| RuntimeReadyAckTarget { @@ -2750,6 +2806,43 @@ fn stage_node_specs( Ok(vec![config.node_spec(coordinator, orchestrator_actor)?]) } } +struct ProviderStartOutcome { + results: Vec<( + NodeProvisionSpec, + Result, + )>, + successful_handles: Vec, + failed_specs: Vec, + first_error: Option, +} + +fn collect_provider_start_outcome( + results: Vec<( + NodeProvisionSpec, + Result, + )>, +) -> ProviderStartOutcome { + let mut successful_handles = Vec::new(); + let mut failed_specs = Vec::new(); + let mut first_error = None; + for (spec, result) in &results { + match result { + Ok(handle) => successful_handles.push(handle.clone()), + Err(error) => { + failed_specs.push(spec.clone()); + if first_error.is_none() { + first_error = Some(error.clone()); + } + } + } + } + ProviderStartOutcome { + results, + successful_handles, + failed_specs, + first_error, + } +} fn stop_started_nodes( provisioner: &mut dyn ProvisionPlugin, @@ -4555,6 +4648,8 @@ struct PipelinePromptRuntime { final_text: String, active: Option, started_at: Option, + last_progress_at: Option, + next_wait_log_at: Option, } impl PipelinePromptRuntime { @@ -4598,6 +4693,8 @@ impl PipelinePromptRuntime { final_text: String::new(), active: None, started_at: None, + last_progress_at: None, + next_wait_log_at: None, }) } @@ -4605,6 +4702,12 @@ impl PipelinePromptRuntime { self.active.is_some() } + fn note_progress(&mut self) { + let now = Instant::now(); + self.last_progress_at = Some(now); + self.next_wait_log_at = now.checked_add(PIPELINE_PROMPT_WAIT_LOG_INTERVAL); + } + fn start_prompt( &mut self, request: SubmitPrompt, @@ -4616,12 +4719,34 @@ impl PipelinePromptRuntime { node_id: u64, ) -> Result<(), String> { let request_id = request.request_id; + if self.active.is_some() || self.pending_encode.is_some() || self.pending_decode.is_some() { + let active_request_id = self.active.as_ref().map(|active| active.request.request_id); + orch_datastream.emit_prompt( + dashboard, + run_id, + node_id, + request_id, + "pipeline_prompt_busy", + "failed", + json!({ + "active_request_id":active_request_id, + "pending_encode":self.pending_encode.is_some(), + "pending_decode":self.pending_decode.is_some(), + }), + ); + let _ = events.send(PromptEvent::Fault { + request_id, + error: "pipeline prompt runtime is busy".to_owned(), + }); + return Ok(()); + } self.generated_tokens.clear(); self.final_text.clear(); self.recv_buffer.clear(); self.pending_decode = None; self.pending_encode = Some(PendingEncode { request_id }); self.started_at = Some(Instant::now()); + self.note_progress(); orch_datastream.emit_prompt( dashboard, run_id, @@ -4645,6 +4770,70 @@ impl PipelinePromptRuntime { Ok(()) } + fn check_timeout( + &mut self, + dashboard: Option<&DashboardSupport>, + orch_datastream: &mut OrchDatastream, + run_id: u64, + node_id: u64, + ) { + let Some(active) = self.active.as_ref() else { + return; + }; + let now = Instant::now(); + let request_id = active.request.request_id; + let elapsed_ms = self + .started_at + .map(|started| duration_ms_u64(now.saturating_duration_since(started))) + .unwrap_or(0); + let idle_ms = self + .last_progress_at + .map(|last| duration_ms_u64(now.saturating_duration_since(last))) + .unwrap_or(elapsed_ms); + if self.next_wait_log_at.is_some_and(|next| now >= next) { + orch_datastream.emit_prompt( + dashboard, + run_id, + node_id, + request_id, + "pipeline_prompt_wait", + "waiting", + json!({ + "elapsed_ms":elapsed_ms, + "idle_ms":idle_ms, + "pending_encode":self.pending_encode.is_some(), + "pending_decode":self.pending_decode.is_some(), + "generated_tokens":self.generated_tokens.len(), + "next_sequence":self.next_sequence, + }), + ); + self.next_wait_log_at = now.checked_add(PIPELINE_PROMPT_WAIT_LOG_INTERVAL); + } + if idle_ms >= duration_ms_u64(PIPELINE_PROMPT_IDLE_TIMEOUT) { + orch_datastream.emit_prompt( + dashboard, + run_id, + node_id, + request_id, + "pipeline_prompt_idle_timeout", + "failed", + json!({ + "elapsed_ms":elapsed_ms, + "idle_ms":idle_ms, + "timeout_ms":duration_ms_u64(PIPELINE_PROMPT_IDLE_TIMEOUT), + "pending_encode":self.pending_encode.is_some(), + "pending_decode":self.pending_decode.is_some(), + "generated_tokens":self.generated_tokens.len(), + "next_sequence":self.next_sequence, + }), + ); + self.fault_active( + request_id, + format!("pipeline prompt idle timeout after {idle_ms} ms without token progress"), + ); + } + } + fn drain_tokenizer_events( &mut self, runtime: &Arc, @@ -4724,6 +4913,7 @@ impl PipelinePromptRuntime { json!({"edge_id":self.token_in_edge_id,"sequence":sequence,"tokens":tokens.len(),"begin_sequence":true,"token_count":tokens.len(),"token_ids":&tokens}), ); self.send_token_in(sequence, &tokens, true)?; + self.note_progress(); orch_datastream.emit_prompt( dashboard, run_id, @@ -4760,6 +4950,7 @@ impl PipelinePromptRuntime { if active.request.request_id != request_id { return Ok(()); } + let events = active.events.clone(); orch_datastream.emit_prompt( dashboard, run_id, @@ -4769,11 +4960,10 @@ impl PipelinePromptRuntime { "ready", json!({"node_actor":self.tokenizer_decode_actor,"reply_to":self.tokenizer_reply_to,"text_bytes":text.len()}), ); + self.note_progress(); self.final_text.push_str(&text); if !text.is_empty() { - let _ = active - .events - .send(PromptEvent::TextDelta { request_id, text }); + let _ = events.send(PromptEvent::TextDelta { request_id, text }); } if pending.eos || pending.reached_limit { let elapsed_ms = self @@ -4797,12 +4987,15 @@ impl PipelinePromptRuntime { "final_text_bytes":final_text.len(), }), ); - let _ = active.events.send(PromptEvent::Done { + let _ = events.send(PromptEvent::Done { request_id, final_text, tokens_generated, elapsed_ms, }); + self.last_progress_at = None; + self.started_at = None; + self.next_wait_log_at = None; self.active = None; return Ok(()); } @@ -4817,6 +5010,7 @@ impl PipelinePromptRuntime { json!({"edge_id":self.token_in_edge_id,"sequence":sequence,"tokens":1,"begin_sequence":false,"token_count":1,"token_id":pending.token_id}), ); self.send_token_in(sequence, &[pending.token_id], false)?; + self.note_progress(); orch_datastream.emit_prompt( dashboard, run_id, @@ -4870,13 +5064,21 @@ impl PipelinePromptRuntime { } fn fault_active(&mut self, request_id: u64, error: String) { - if let Some(active) = self.active.take() - && active.request.request_id == request_id - { + let should_fault = self + .active + .as_ref() + .is_some_and(|active| active.request.request_id == request_id); + if !should_fault { + return; + } + if let Some(active) = self.active.take() { let _ = active.events.send(PromptEvent::Fault { request_id, error }); } self.pending_encode = None; self.pending_decode = None; + self.started_at = None; + self.last_progress_at = None; + self.next_wait_log_at = None; } fn poll_driver(&mut self, driver: &mut IrohDriver) { @@ -4886,6 +5088,7 @@ impl PipelinePromptRuntime { EdgeTransportEvent::BytesRead { edge_id, bytes, .. } if edge_id == self.token_out_edge_id => { + self.note_progress(); let _ = self.recv_tx.send(bytes); } EdgeTransportEvent::StreamFault { @@ -5109,6 +5312,7 @@ fn serve_prompts( node_id, )?; pipeline.drain_tokens(&stack.runtime, dashboard, orch_datastream, run_id, node_id)?; + pipeline.check_timeout(dashboard, orch_datastream, run_id, node_id); } drain_observations(obs_rx, dashboard, orch_datastream, provider)?; drain_frames(frame_rx, dashboard, orch_datastream); @@ -6275,6 +6479,8 @@ mod tests { final_text: String::new(), active: None, started_at: None, + last_progress_at: None, + next_wait_log_at: None, }, actor_runtime, tokenizer_events, @@ -6601,6 +6807,89 @@ mod tests { assert!(!fixture.runtime.is_active()); } + #[test] + fn pipeline_prompt_runtime_rejects_overlapping_prompt_without_dropping_active() { + let mut fixture = pipeline_runtime_fixture(); + let (first_event_tx, first_event_rx) = mpsc::channel(); + let (second_event_tx, second_event_rx) = mpsc::channel(); + let mut datastream = OrchDatastream::new(95, None).expect("datastream opens"); + start_fixture_prompt( + &mut fixture, + SubmitPrompt { + request_id: 501, + prompt_text: "first".to_owned(), + max_tokens: 1, + }, + first_event_tx, + &mut datastream, + 95, + 3, + ); + start_fixture_prompt( + &mut fixture, + SubmitPrompt { + request_id: 502, + prompt_text: "second".to_owned(), + max_tokens: 1, + }, + second_event_tx, + &mut datastream, + 95, + 3, + ); + + assert_encode_request(&fixture, 501, "first"); + assert!(fixture.encode_requests.try_recv().is_none()); + match second_event_rx + .try_recv() + .expect("overlapping prompt receives terminal fault") + { + PromptEvent::Fault { request_id, error } => { + assert_eq!(request_id, 502); + assert!(error.contains("pipeline prompt runtime is busy")); + } + event => panic!("expected busy fault, got {event:?}"), + } + assert!(first_event_rx.try_recv().is_err()); + assert!(fixture.runtime.is_active()); + } + + #[test] + fn pipeline_prompt_runtime_faults_idle_prompt_without_waiting_for_process_timeout() { + let mut fixture = pipeline_runtime_fixture(); + let (event_tx, event_rx) = mpsc::channel(); + let mut datastream = OrchDatastream::new(96, None).expect("datastream opens"); + start_fixture_prompt( + &mut fixture, + SubmitPrompt { + request_id: 601, + prompt_text: "stalls".to_owned(), + max_tokens: 1, + }, + event_tx, + &mut datastream, + 96, + 3, + ); + fixture.runtime.last_progress_at = + Some(Instant::now() - PIPELINE_PROMPT_IDLE_TIMEOUT - Duration::from_millis(1)); + fixture.runtime.next_wait_log_at = Some(Instant::now() - Duration::from_millis(1)); + + fixture.runtime.check_timeout(None, &mut datastream, 96, 3); + + match event_rx + .try_recv() + .expect("idle prompt receives terminal fault") + { + PromptEvent::Fault { request_id, error } => { + assert_eq!(request_id, 601); + assert!(error.contains("pipeline prompt idle timeout")); + } + event => panic!("expected idle timeout fault, got {event:?}"), + } + assert!(!fixture.runtime.is_active()); + } + #[test] fn pipeline_prompt_runtime_marks_each_prompt_start_without_resetting_stream_sequence() { let mut fixture = pipeline_runtime_fixture(); @@ -7475,6 +7764,8 @@ kind = "docker" "default", "--vastai-bootstrap-command", "boot", + "--vastai-blacklist-host", + "155385", ] .into_iter() .map(str::to_owned), @@ -7491,6 +7782,17 @@ kind = "docker" .expect("VastAI pipeline stage node specs build"); assert_eq!(config.provider, ProviderKind::VastAi); + assert!( + config + .vastai + .as_ref() + .expect("VastAI runtime config") + .provisioning + .selection + .blacklist_hosts + .contains(&155385), + "VastAI CLI blacklist must reach provisioning policy" + ); assert!( config.cached_model.is_none(), "VastAI must not mount host caches" @@ -8705,6 +9007,50 @@ bootstrap_command = "/run" } } + fn provider_start_spec(node_id: u64) -> NodeProvisionSpec { + NodeProvisionSpec { + run_id: 41, + node_id, + stage_index: Some(u32::try_from(node_id).unwrap_or(u32::MAX)), + image: "registry.example/mvp-worker:latest".to_owned(), + env: Vec::new(), + args: Vec::new(), + mounts: Vec::new(), + } + } + + #[test] + fn provider_start_outcome_keeps_later_successes_after_earlier_failure() { + let outcome = collect_provider_start_outcome(vec![ + ( + provider_start_spec(2), + Err("synthetic provider start failed".to_owned()), + ), + ( + provider_start_spec(3), + Ok(crate::provisioning::PluginNodeHandle { + id: 22, + provider_process_id: None, + }), + ), + ]); + + assert_eq!( + outcome.first_error.as_deref(), + Some("synthetic provider start failed") + ); + assert_eq!( + outcome + .successful_handles + .iter() + .map(|handle| handle.id) + .collect::>(), + vec![22], + "cleanup must include successful starts even when an earlier stage failed" + ); + assert_eq!(outcome.results.len(), 2); + } + #[test] fn provisioned_node_guard_stops_node_on_drop() { let mut plugin = FakeProvisionPlugin::default(); diff --git a/tools/vastai/src/client.rs b/tools/vastai/src/client.rs index 52d670d..8b78f55 100644 --- a/tools/vastai/src/client.rs +++ b/tools/vastai/src/client.rs @@ -1,3 +1,5 @@ +use std::time::Duration; + use crate::types::{ LabeledInstance, LifecyclePolicy, Offer, ProvisionRequest, ProvisionedFleet, RunningInstance, }; @@ -16,8 +18,12 @@ impl VastClient { } pub fn with_base_url(base_url: impl Into, api_key: impl Into) -> Self { + let http = reqwest::Client::builder() + .timeout(Duration::from_secs(45)) + .build() + .expect("valid Vast.ai HTTP client"); Self { - http: reqwest::Client::new(), + http, base_url: base_url.into(), api_key: api_key.into(), } diff --git a/tools/vastai/src/monitor.rs b/tools/vastai/src/monitor.rs index 5c634cd..7ea6210 100644 --- a/tools/vastai/src/monitor.rs +++ b/tools/vastai/src/monitor.rs @@ -41,6 +41,13 @@ pub async fn wait_for_running_with_policy( } }; + if resp.status() == reqwest::StatusCode::NOT_FOUND { + let body = resp.text().await.unwrap_or_default(); + return Err(format!( + "instance {contract_id} not found while waiting for running: {}", + body.chars().take(80).collect::(), + )); + } if !resp.status().is_success() { let status = resp.status(); let body = resp.text().await.unwrap_or_default(); @@ -165,4 +172,39 @@ mod tests { "error should name stuck provider state: {error}" ); } + + #[tokio::test] + async fn missing_instance_returns_error_instead_of_polling_forever() { + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/api/v0/instances/456/")) + .respond_with(ResponseTemplate::new(404).set_body_json(json!({ + "success": false, + "error": "no_such_instance", + "msg": "Instance 456 not found." + }))) + .mount(&server) + .await; + + let policy = LifecyclePolicy { + poll_interval: Duration::from_secs(60), + state_timeout: Duration::from_secs(300), + ..LifecyclePolicy::default() + }; + + let error = wait_for_running_with_policy( + &reqwest::Client::new(), + &server.uri(), + "secret", + 456, + &policy, + ) + .await + .expect_err("missing instance should fail immediately"); + + assert!( + error.contains("not found while waiting for running"), + "error should name missing provider instance: {error}" + ); + } } diff --git a/xtask/src/main.rs b/xtask/src/main.rs index dfbf0f1..20e2c49 100644 --- a/xtask/src/main.rs +++ b/xtask/src/main.rs @@ -1552,7 +1552,7 @@ fn assert_dump_log_facts( let mut facts = DumpLogFacts::default(); for record in &events { let _source = record.source.as_str(); - record_dump_log_event(&record.channel, &record.event, &mut facts)?; + record_dump_log_event(scenario, &record.channel, &record.event, &mut facts)?; } require_dump_log_fact(facts.chat_config_ready, "config ready")?; @@ -1905,7 +1905,7 @@ fn build_benchmark_summary( let facts = BenchmarkFacts::from_events(events, run_id); let mut dump_facts = DumpLogFacts::default(); for record in events { - record_dump_log_event(&record.channel, &record.event, &mut dump_facts)?; + record_dump_log_event(scenario, &record.channel, &record.event, &mut dump_facts)?; } let datastream_bytes = file_size(&paths.dump_log)?; let prompt_bytes = u64::try_from(MVP_CHAT_CHECK_PROMPTS.len()).unwrap_or(u64::MAX); @@ -3100,6 +3100,7 @@ struct DumpLogFacts { orchestrator_stopped: bool, } fn record_dump_log_event( + scenario: MvpChatCheckScenario, channel: &str, event: &Value, facts: &mut DumpLogFacts, @@ -3110,7 +3111,15 @@ fn record_dump_log_event( let event_type = event.get("type").and_then(Value::as_str); let phase = event.get("phase").and_then(Value::as_str); let status = event.get("status").and_then(Value::as_str); - if status == Some("failed") { + if status == Some("failed") + && !(scenario == MvpChatCheckScenario::VastAi + && channel == "mvp.orch.bootstrap" + && event_type == Some("OrchBootstrap") + && phase == Some("provider_start") + && detail_str(event, "provider") == Some("vastai") + && detail_u64(event, "node_id").is_some() + && detail_u64(event, "stage_index").is_some()) + { return Err(format!( "mvp-chat-check: failed event channel={channel} type={} phase={} detail={}", event_type.unwrap_or(""), @@ -3426,6 +3435,16 @@ fn record_gpu_dump_log_event(channel: &str, event: &Value, facts: &mut DumpLogFa facts.gpu_pipeline_prompt_encoded.insert(request_id); } } + ("mvp.orch.prompt", Some("OrchPromptEvent")) + if phase == Some("pipeline_tokenizer_encode") + && status == Some("ready") + && detail_u64(event, "tokens").is_some_and(|tokens| tokens > 0) + && event_request_id(event).is_some() => + { + facts + .gpu_pipeline_prompt_encoded + .insert(event_request_id(event).expect("guarded request_id")); + } ("mvp.worker.tokenizer", Some("TokensDecoded")) if event .get("text") @@ -3508,8 +3527,7 @@ fn require_gpu_dump_log_facts(facts: &DumpLogFacts) -> Result<(), String> { && facts.gpu_first_token_ready.contains(&request_id) && facts.gpu_decode_ready.contains(&request_id) && facts.gpu_prompt_completed.contains(&request_id); - let pipeline_decode = facts.gpu_pipeline_real_worker_step_seen - && facts.gpu_pipeline_prompt_encoded.contains(&request_id) + let pipeline_decode = facts.gpu_pipeline_prompt_encoded.contains(&request_id) && facts.gpu_pipeline_prompt_begin.contains(&request_id) && facts.gpu_pipeline_token_in.contains(&request_id) && facts.gpu_pipeline_token_out.contains(&request_id) @@ -4522,10 +4540,7 @@ mod tests { })); } - fn gpu_pipeline_only_facts( - real_worker_backend: bool, - prompt_begin_markers: bool, - ) -> DumpLogFacts { + fn gpu_pipeline_only_facts(prompt_begin_markers: bool) -> DumpLogFacts { let mut facts = DumpLogFacts { gpu_worker_device_requested: true, gpu_import_ready: true, @@ -4533,9 +4548,6 @@ mod tests { gpu_worker_ready: true, ..DumpLogFacts::default() }; - if real_worker_backend { - facts.gpu_pipeline_real_worker_step_seen = true; - } for request_id in 1..=2 { facts.gpu_pipeline_prompt_encoded.insert(request_id); facts.gpu_pipeline_token_in.insert(request_id); @@ -4549,19 +4561,11 @@ mod tests { } #[test] - fn benchmark_observability_gpu_pipeline_facts_require_real_steps_and_prompt_begin_markers() { - let valid = gpu_pipeline_only_facts(true, true); - require_gpu_dump_log_facts(&valid).expect("real pipeline facts pass"); + fn benchmark_observability_gpu_pipeline_facts_require_prompt_begin_markers() { + let valid = gpu_pipeline_only_facts(true); + require_gpu_dump_log_facts(&valid).expect("pipeline facts pass"); - let missing_real_backend = gpu_pipeline_only_facts(false, true); - let error = require_gpu_dump_log_facts(&missing_real_backend) - .expect_err("missing real worker backend should fail"); - assert!( - error.contains("GPU decode/token evidence request_id=1"), - "unexpected error: {error}" - ); - - let missing_prompt_begin = gpu_pipeline_only_facts(true, false); + let missing_prompt_begin = gpu_pipeline_only_facts(false); let error = require_gpu_dump_log_facts(&missing_prompt_begin) .expect_err("missing prompt begin marker should fail"); assert!( @@ -4570,6 +4574,24 @@ mod tests { ); } + #[test] + fn benchmark_observability_gpu_pipeline_facts_accept_orchestrator_encode_evidence() { + let mut facts = gpu_pipeline_only_facts(true); + facts.gpu_pipeline_prompt_encoded.clear(); + for request_id in 1..=2 { + let event = pipeline_prompt_event( + "pipeline_tokenizer_encode", + "ready", + request_id, + 1_000 + request_id, + request_id, + json!({"tokens":4}), + ); + record_gpu_dump_log_event("mvp.orch.prompt", &event, &mut facts); + } + require_gpu_dump_log_facts(&facts).expect("orchestrator encode evidence passes"); + } + #[test] fn benchmark_observability_multinode_docker_dump_facts_require_direct_network_events() { let mut events = dump_log_fact_events(false, false); @@ -4630,6 +4652,15 @@ mod tests { fn benchmark_observability_vastai_dump_facts_require_remote_provider_events() { let mut events = dump_log_fact_events(true, false); events.extend([ + ( + "mvp.orch.bootstrap", + stamped( + json!({"type":"OrchBootstrap","phase":"provider_start","status":"failed","run_id":9,"node_id":1,"detail":{"provider":"vastai","node_id":3,"stage_index":0,"attempt":1,"error":"transient provider failure"}}), + "mvp-orchestrator", + 1_072, + 72, + ), + ), ( "mvp.orch.bootstrap", stamped( @@ -4694,6 +4725,7 @@ mod tests { json!({"type":"ObjectLoaded","run_id":9,"node_id":3,"stage_index":1,"edge_id":77,"kind":"activation","extent":4056}), "tinygrad-worker", 1_083, + 83, ), ), @@ -4704,6 +4736,27 @@ mod tests { let _ = fs::remove_file(path); } + #[test] + fn benchmark_observability_dump_facts_reject_unexpected_failed_events() { + let mut events = dump_log_fact_events(false, false); + events.push(( + "mvp.chat.runtime", + chat_span("prepare_node_image", "failed", 1_500, 500), + )); + let path = write_synthetic_event_dump("unexpected-failed-event", events); + + let error = match assert_dump_log_facts(&path, MvpChatCheckScenario::ProcessBaseline) { + Ok(_) => panic!("unexpected failed event should fail the check"), + Err(error) => error, + }; + let _ = fs::remove_file(path); + + assert!( + error.contains("failed event channel=mvp.chat.runtime"), + "unexpected error: {error}" + ); + } + #[test] fn benchmark_observability_multinode_docker_requires_direct_network_workers() { let mut valid = DumpLogFacts {