diff --git a/crates/mvp-system/specs/mvp_chat.md b/crates/mvp-system/specs/mvp_chat.md index 498f00c..7eefd3a 100644 --- a/crates/mvp-system/specs/mvp_chat.md +++ b/crates/mvp-system/specs/mvp_chat.md @@ -73,6 +73,10 @@ General flags: - `--cached-model` - `--cached-model=` - `--skip-rebuild` +- `--gpu` +- `--relay-mode ` +- `--relay-url ` +- `--endpoint-addr-mask ` Any process argument that is not one of the listed flags or a required or attached value for one of those flags is a configuration error. @@ -85,6 +89,11 @@ If no provider selector is supplied, the provider is `process`. `--pipeline-stages ` accepts a positive integer. Zero and invalid values are configuration errors. +`--relay-mode ` accepts `default` or `disabled`. + +`--endpoint-addr-mask ` accepts `full` or `relay-only`. `relay-only` +requires a relay URL from `--relay-url`, `[relay].url`, or `[vastai].relay_url`. + `--dump-logs` writes the consolidated log stream to the default file `mvp-chat.log` in the current working directory. @@ -279,6 +288,19 @@ Additional provider/image inputs: - registry responses when checking remote image availability; - registry responses when pushing images for remote providers. +Additional relay inputs: + +- relay mode from `--relay-mode`, `[relay].mode`, or `MVP_IROH_RELAY_MODE`; +- relay URL from `--relay-url`, `[relay].url`, `[vastai].relay_url`, or relay + environment fallbacks; +- endpoint address mask from `--endpoint-addr-mask`, `[relay].endpoint_addr_mask`, + or `MVP_IROH_ENDPOINT_ADDR_MASK`. + +When the endpoint address mask is `relay-only`, orchestrator and worker +advertisements must preserve relay URLs and strip direct socket addresses before +passing endpoints across provider/runtime boundaries. A missing relay URL is a +configuration or startup error. + Prompt engine stream records are runtime-local prompt events: - text delta; diff --git a/crates/mvp-system/src/bin/mvp_chat.rs b/crates/mvp-system/src/bin/mvp_chat.rs index 0eaefde..c619392 100644 --- a/crates/mvp-system/src/bin/mvp_chat.rs +++ b/crates/mvp-system/src/bin/mvp_chat.rs @@ -25,6 +25,7 @@ use signal_hook::iterator::Signals; use mvp_system::benchmark_observability; use mvp_system::config as chat_config; use mvp_system::config::ResolvedVastAiConfig; +use mvp_system::endpoint_advertisement::EndpointAddrMask; use mvp_system::node_image::{ NodeImageProvider, NodeImageRequest, PreparedNodeImage, prepare_node_image, }; @@ -46,6 +47,9 @@ OPTIONS: Select the runtime provider --config Load config overlay --pipeline-stages Number of pipeline stages + --relay-mode Relay mode: default or disabled + --relay-url Custom relay URL passed to mvp-orchestrator + --endpoint-addr-mask Endpoint address mask: full or relay-only --cached-model[=] Use discovered or explicit cached GGUF model --dump-logs[=] Write datastream frame log --run-id Override run id @@ -331,6 +335,9 @@ struct Config { max_tokens: u32, skip_rebuild: bool, gpu_run: bool, + relay_mode: Option, + relay_url: Option, + endpoint_addr_mask: EndpointAddrMask, } struct ChatDatastream { @@ -514,6 +521,7 @@ struct ChatTomlConfig { observability: ChatObservabilityConfig, image: ChatImageConfig, vastai: ChatVastAiConfig, + relay: ChatRelayConfig, } #[derive(Clone, Debug, Default, Deserialize)] @@ -529,6 +537,14 @@ struct ChatRuntimeConfig { max_tokens: Option, } +#[derive(Clone, Debug, Default, Deserialize)] +#[serde(default, deny_unknown_fields)] +struct ChatRelayConfig { + mode: Option, + url: Option, + endpoint_addr_mask: Option, +} + #[derive(Clone, Debug, Default, Deserialize)] #[serde(default, deny_unknown_fields)] struct ChatObservabilityConfig { @@ -612,6 +628,22 @@ impl Config { return Err("[runtime].max_tokens must be greater than 0".to_owned()); } let gpu_run = args.gpu || env_flag(MVP_CHAT_GPU_RUN_ENV, false); + let endpoint_addr_mask = match first_non_empty([ + args.endpoint_addr_mask.clone(), + toml.relay.endpoint_addr_mask.clone(), + ]) { + Some(mask) => EndpointAddrMask::parse(&mask)?, + None => EndpointAddrMask::Full, + }; + let relay_mode = first_non_empty([args.relay_mode.clone(), toml.relay.mode.clone()]); + let mut relay_url = first_non_empty([args.relay_url.clone(), toml.relay.url.clone()]); + if endpoint_addr_mask.requires_relay() && relay_url.is_none() { + relay_url = first_non_empty([toml.vastai.relay_url.clone()]); + } + if endpoint_addr_mask.requires_relay() && relay_url.is_none() { + return Err("relay-only endpoint address mask requires [relay].url, --relay-url, or [vastai].relay_url".to_owned()); + } + let relay_mode = relay_mode.or_else(|| relay_url.as_ref().map(|_| "default".to_owned())); let cached_model_source = match args.cached_model { Some(source) => Some(source), None if gpu_run && provider == ProviderKind::Process => { @@ -658,6 +690,9 @@ impl Config { vastai, skip_rebuild: args.skip_rebuild, gpu_run, + relay_mode, + relay_url, + endpoint_addr_mask, }) } @@ -697,6 +732,18 @@ impl Config { path.to_string_lossy().to_string(), ]); } + if let Some(mode) = &self.relay_mode { + args.extend(["--relay-mode".to_owned(), mode.clone()]); + } + if let Some(url) = &self.relay_url { + args.extend(["--relay-url".to_owned(), url.clone()]); + } + if self.endpoint_addr_mask != EndpointAddrMask::Full { + args.extend([ + "--endpoint-addr-mask".to_owned(), + self.endpoint_addr_mask.as_str().to_owned(), + ]); + } if let Some(vastai) = &self.vastai { args.extend([ "--vastai-api-key".to_owned(), @@ -771,6 +818,9 @@ struct ParsedArgs { cached_model: Option, help: bool, gpu: bool, + relay_mode: Option, + relay_url: Option, + endpoint_addr_mask: Option, } #[derive(Clone, Debug, PartialEq, Eq)] @@ -812,6 +862,11 @@ impl ParsedArgs { parsed.pipeline_stages = Some(parse_pipeline_stages_value(&mut args, arg.as_str())?) } + "--relay-mode" => parsed.relay_mode = Some(next_arg(&mut args, "--relay-mode")?), + "--relay-url" => parsed.relay_url = Some(next_arg(&mut args, "--relay-url")?), + "--endpoint-addr-mask" => { + parsed.endpoint_addr_mask = Some(next_arg(&mut args, "--endpoint-addr-mask")?) + } "--run-id" => { let run_id: u64 = parse_next(&mut args, "--run-id")?; if run_id == 0 { @@ -2134,6 +2189,9 @@ mod tests { max_tokens: DEFAULT_MAX_TOKENS, skip_rebuild: true, gpu_run: false, + relay_mode: None, + relay_url: None, + endpoint_addr_mask: EndpointAddrMask::Full, } } @@ -2646,6 +2704,76 @@ bootstrap_command = "boot" ); } + #[test] + fn relay_only_endpoint_mask_requires_and_forwards_relay_url() { + let missing = TempDir::new("relay-mask-missing-url"); + let missing_config = write_config( + &missing, + "chat.toml", + r#" +[provider] +kind = "docker" + +[image] +node = "docker.io/acme/node:latest" +"#, + ); + with_process_state(&[], Some(missing.path()), || { + let config_arg = missing_config.to_string_lossy().into_owned(); + let error = match Config::from_args(strings(&[ + "--config", + config_arg.as_str(), + "--endpoint-addr-mask", + "relay-only", + ])) { + Ok(_) => panic!("relay-only mask without relay URL should fail"), + Err(error) => error, + }; + assert!( + error.contains("requires [relay].url, --relay-url, or [vastai].relay_url"), + "{error}" + ); + }); + + let fallback = TempDir::new("relay-mask-vastai-fallback"); + let fallback_config = write_config( + &fallback, + "chat.toml", + r#" +[provider] +kind = "docker" + +[image] +node = "docker.io/acme/node:latest" + +[relay] +endpoint_addr_mask = "relay-only" + +[vastai] +relay_url = "https://relay.example" +"#, + ); + with_process_state(&[], Some(fallback.path()), || { + let config_arg = fallback_config.to_string_lossy().into_owned(); + let config = Config::from_args(strings(&["--config", config_arg.as_str()])) + .expect("relay-only mask uses Vast.ai relay fallback"); + + assert_eq!(config.provider, ProviderKind::Docker); + assert_eq!(config.relay_mode.as_deref(), Some("default")); + assert_eq!(config.relay_url.as_deref(), Some("https://relay.example")); + assert_eq!(config.endpoint_addr_mask, EndpointAddrMask::RelayOnly); + let args = config.orchestrator_cli_args("docker.io/acme/node:latest"); + assert!( + args.windows(2) + .any(|pair| pair == ["--relay-url", "https://relay.example"]) + ); + assert!( + args.windows(2) + .any(|pair| pair == ["--endpoint-addr-mask", "relay-only"]) + ); + }); + } + struct MockApproval { terminal: bool, answer: Result, diff --git a/crates/mvp-system/src/bin/worker_node.rs b/crates/mvp-system/src/bin/worker_node.rs index 6e03d5a..00ca822 100644 --- a/crates/mvp-system/src/bin/worker_node.rs +++ b/crates/mvp-system/src/bin/worker_node.rs @@ -33,6 +33,9 @@ use mvp_system::benchmark_observability; use mvp_system::distribution_stack::DistributionRuntimeStack; use mvp_system::driver_pumps as driver_model; use mvp_system::edge_establisher as edge; +use mvp_system::endpoint_advertisement::{ + EndpointAddrMask, MVP_IROH_ENDPOINT_ADDR_MASK_ENV, advertised_endpoint, +}; use mvp_system::gpu_worker_ingress_parser as ingress; use mvp_system::prompt_rpc::{PromptEvent, TokenizerEvent}; use mvp_system::relay_provisioning::relay_runtime_config_from_env; @@ -1326,16 +1329,7 @@ fn run() -> Result<(), String> { additional_alpns: vec![EDGE_ALPN.to_vec(), DATASTREAM_ALPN.to_vec()], }, ) { - Ok(driver) => { - emit_stdio_node_event( - &config, - NODE_BOOTSTRAP_CHANNEL, - "iroh_driver", - "ready", - json!({"endpoint":driver.endpoint_addr(),"relay_mode":format!("{:?}", config.relay_mode)}), - )?; - driver - } + Ok(driver) => driver, Err(error) => { emit_stdio_node_event( &config, @@ -1347,6 +1341,15 @@ fn run() -> Result<(), String> { return Err(format!("create iroh driver: {error}")); } }; + let advertised_self_endpoint = + advertised_endpoint(driver.endpoint_addr(), config.endpoint_addr_mask)?; + emit_stdio_node_event( + &config, + NODE_BOOTSTRAP_CHANNEL, + "iroh_driver", + "ready", + json!({"endpoint":advertised_self_endpoint.clone(),"has_relay":advertised_self_endpoint.relay_urls().next().is_some(),"direct_addr_count":advertised_self_endpoint.ip_addrs().count(),"relay_mode":format!("{:?}", config.relay_mode),"endpoint_addr_mask":config.endpoint_addr_mask.as_str()}), + )?; if let Some(coordinator) = &config.coordinator_endpoint { driver.join(std::slice::from_ref(coordinator)); emit_stdio_node_event( @@ -1354,7 +1357,7 @@ fn run() -> Result<(), String> { NODE_BOOTSTRAP_CHANNEL, "coordinator_join", "started", - json!({"endpoint":coordinator}), + json!({"endpoint":coordinator,"has_relay":coordinator.relay_urls().next().is_some(),"direct_addr_count":coordinator.ip_addrs().count()}), )?; } else { emit_stdio_node_event( @@ -1651,7 +1654,7 @@ fn run() -> Result<(), String> { let mut edge_runtime = WorkerEdgeRuntime::new(config.logical_node_id); let mut pending_runtime_ready = PendingRuntimeReady::new( &config, - driver.endpoint_addr(), + advertised_self_endpoint.clone(), node_actor, datastream_publisher, ); @@ -1659,7 +1662,7 @@ fn run() -> Result<(), String> { let ready = json!({ "type":"ready", "role":"node", - "endpoint": driver.endpoint_addr(), + "endpoint": advertised_self_endpoint.clone(), "node_actor": node_actor, "datastream_publisher": datastream_publisher, "logical_node_id": config.logical_node_id, @@ -1671,7 +1674,7 @@ fn run() -> Result<(), String> { "runtime_ready_local", "ready", json!({ - "endpoint":driver.endpoint_addr(), + "endpoint":advertised_self_endpoint.clone(), "node_actor":node_actor, "logical_node_id":config.logical_node_id, "stage_index":config.stage_index, @@ -2934,6 +2937,7 @@ struct DeploymentConfig { datastream_frame_log: Option, debug_join_socket: Option, relay_mode: iroh::RelayMode, + endpoint_addr_mask: EndpointAddrMask, worker_script: String, device: String, model_id: String, @@ -2978,6 +2982,11 @@ impl DeploymentConfig { datastream_frame_log: env_optional("MVP_DATASTREAM_FRAME_LOG"), debug_join_socket, relay_mode: relay.mode, + endpoint_addr_mask: env_optional(MVP_IROH_ENDPOINT_ADDR_MASK_ENV) + .as_deref() + .map(EndpointAddrMask::parse) + .transpose()? + .unwrap_or_default(), worker_script: env_string("MVP_TINYGRAD_WORKER", DEFAULT_WORKER_SCRIPT), device: env_string("DEV", default_device), model_id: env_string("MVP_MODEL_ID", DEFAULT_MODEL_ID), @@ -3652,6 +3661,7 @@ mod tests { datastream_frame_log: None, debug_join_socket: None, relay_mode: iroh::RelayMode::Disabled, + endpoint_addr_mask: EndpointAddrMask::Full, worker_script: DEFAULT_WORKER_SCRIPT.to_owned(), device: DEFAULT_DEVICE.to_owned(), model_id: DEFAULT_MODEL_ID.to_owned(), diff --git a/crates/mvp-system/src/endpoint_advertisement.rs b/crates/mvp-system/src/endpoint_advertisement.rs new file mode 100644 index 0000000..3a8b279 --- /dev/null +++ b/crates/mvp-system/src/endpoint_advertisement.rs @@ -0,0 +1,104 @@ +use std::fmt; + +use iroh::EndpointAddr; + +pub const MVP_IROH_ENDPOINT_ADDR_MASK_ENV: &str = "MVP_IROH_ENDPOINT_ADDR_MASK"; + +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub enum EndpointAddrMask { + #[default] + Full, + RelayOnly, +} + +impl EndpointAddrMask { + pub fn parse(value: &str) -> Result { + match value.trim().to_ascii_lowercase().as_str() { + "" | "full" | "none" => Ok(Self::Full), + "relay-only" | "relay_only" | "relay" => Ok(Self::RelayOnly), + other => Err(format!( + "unsupported {MVP_IROH_ENDPOINT_ADDR_MASK_ENV}={other:?}; use full or relay-only" + )), + } + } + + pub fn as_str(self) -> &'static str { + match self { + Self::Full => "full", + Self::RelayOnly => "relay-only", + } + } + + pub fn requires_relay(self) -> bool { + matches!(self, Self::RelayOnly) + } +} + +impl fmt::Display for EndpointAddrMask { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str(self.as_str()) + } +} + +pub fn advertised_endpoint( + endpoint: EndpointAddr, + mask: EndpointAddrMask, +) -> Result { + match mask { + EndpointAddrMask::Full => Ok(endpoint), + EndpointAddrMask::RelayOnly => relay_only_endpoint(endpoint), + } +} + +fn relay_only_endpoint(endpoint: EndpointAddr) -> Result { + let relays = endpoint.relay_urls().cloned().collect::>(); + if relays.is_empty() { + return Err("relay-only endpoint address mask requires an endpoint relay URL".to_owned()); + } + + let mut out = EndpointAddr::new(endpoint.id); + for relay in relays { + out = out.with_relay_url(relay); + } + Ok(out) +} + +#[cfg(test)] +mod tests { + use std::net::{Ipv4Addr, SocketAddr}; + + use super::*; + + #[test] + fn relay_only_mask_preserves_relay_urls_and_removes_direct_addresses() { + let relay = "http://relay.example.com" + .parse::() + .expect("relay URL parses"); + let endpoint = EndpointAddr::new(iroh::SecretKey::from_bytes(&[9; 32]).public()) + .with_relay_url(relay.clone()) + .with_ip_addr(SocketAddr::from((Ipv4Addr::LOCALHOST, 7777))); + + let masked = advertised_endpoint(endpoint, EndpointAddrMask::RelayOnly) + .expect("relay-only endpoint builds"); + + assert_eq!(masked.ip_addrs().count(), 0); + assert_eq!( + masked.relay_urls().next().map(ToString::to_string), + Some(relay.to_string()) + ); + } + + #[test] + fn relay_only_mask_rejects_endpoint_without_relay_url() { + let endpoint = EndpointAddr::new(iroh::SecretKey::from_bytes(&[7; 32]).public()) + .with_ip_addr(SocketAddr::from((Ipv4Addr::LOCALHOST, 7777))); + + let error = advertised_endpoint(endpoint, EndpointAddrMask::RelayOnly) + .expect_err("missing relay URL fails"); + + assert!( + error.contains("requires an endpoint relay URL"), + "unexpected error: {error}" + ); + } +} diff --git a/crates/mvp-system/src/lib.rs b/crates/mvp-system/src/lib.rs index 37995ea..94269c7 100644 --- a/crates/mvp-system/src/lib.rs +++ b/crates/mvp-system/src/lib.rs @@ -12,6 +12,7 @@ pub mod distribution_stack; pub mod docker_cluster_provisioning; pub mod driver_pumps; pub mod edge_establisher; +pub mod endpoint_advertisement; pub mod engine_builder; pub mod gguf_metadata; pub mod gpu_worker_ctl; diff --git a/crates/mvp-system/src/node_image.rs b/crates/mvp-system/src/node_image.rs index 95bda02..d708d3c 100644 --- a/crates/mvp-system/src/node_image.rs +++ b/crates/mvp-system/src/node_image.rs @@ -175,7 +175,7 @@ pub fn prepare_node_image(request: NodeImageRequest) -> Result Result { }) } +fn docker_build_context_path(root: &Path, path: &Path) -> Result { + let full = if path.is_absolute() { + path.to_path_buf() + } else { + root.join(path) + }; + relative_path(root, &full).map(|relative| relative.to_string_lossy().to_string()) +} + fn display_workspace_path(root: &Path, path: &Path) -> String { match path.strip_prefix(root) { Ok(relative) if relative.as_os_str().is_empty() => ".".to_owned(), @@ -778,4 +787,27 @@ mod tests { vec!["latest".to_owned(), "smoke".to_owned()] ); } + + #[test] + fn docker_build_context_path_makes_worker_binary_relative_to_workspace() { + let root = Path::new("/workspace/swactor"); + + assert_eq!( + docker_build_context_path( + root, + Path::new("/workspace/swactor/target/debug/mvp-worker-node") + ) + .expect("absolute workspace path is valid"), + "target/debug/mvp-worker-node" + ); + assert_eq!( + docker_build_context_path(root, Path::new("target/debug/mvp-worker-node")) + .expect("relative workspace path is valid"), + "target/debug/mvp-worker-node" + ); + assert!( + docker_build_context_path(root, Path::new("/tmp/mvp-worker-node")).is_err(), + "Docker COPY inputs must stay inside the build context" + ); + } } diff --git a/crates/mvp-system/src/orchestrator_app.rs b/crates/mvp-system/src/orchestrator_app.rs index cbe6b0d..56efd98 100644 --- a/crates/mvp-system/src/orchestrator_app.rs +++ b/crates/mvp-system/src/orchestrator_app.rs @@ -21,6 +21,9 @@ use crate::config::{DEFAULT_CONFIG_PATH, TomlConfigOverlay}; #[cfg(feature = "dashboard")] use crate::dashboard_view::MvpClusterDashboardView; use crate::distribution_stack::DistributionRuntimeStack; +use crate::endpoint_advertisement::{ + EndpointAddrMask, MVP_IROH_ENDPOINT_ADDR_MASK_ENV, advertised_endpoint, +}; use crate::gpu_worker_ingress_parser as ingress; use crate::node_provisioning::ProviderKind; use crate::orchestrator_run_fsm::{RunConfig, RunId}; @@ -79,6 +82,7 @@ const DEFAULT_MAX_TOKENS: u32 = 64; const PUMP_INTERVAL: Duration = Duration::from_millis(10); const RUNTIME_READY_ACK_RETRY_INTERVAL: Duration = Duration::from_millis(250); const RUNTIME_READY_ACK_TIMEOUT: Duration = Duration::from_secs(60); +const RUNTIME_READY_TIMEOUT: Duration = Duration::from_secs(60); const MVP_ORCH_BOOTSTRAP: &str = "mvp.orch.bootstrap"; const MVP_ORCH_PROMPT: &str = "mvp.orch.prompt"; const MVP_SWIM_MEMBERSHIP: &str = "mvp.swim.membership"; @@ -149,6 +153,7 @@ where "stage_index":config.stage_index, "legacy_layer_end_exclusive":config.layer_end_exclusive, "relay_mode":format!("{:?}", config.relay.mode), + "endpoint_addr_mask":config.endpoint_addr_mask.as_str(), "pipeline_stages":config.pipeline_stages, "provider_config":config.provider_datastream_detail(), }), @@ -217,17 +222,7 @@ where additional_alpns: vec![EDGE_ALPN.to_vec(), DATASTREAM_ALPN.to_vec()], }, ) { - Ok(driver) => { - orch_datastream.emit_bootstrap( - None, - config.run_id, - config.node_id, - "iroh_driver", - "ready", - json!({"relay_mode":format!("{:?}", config.relay.mode)}), - ); - driver - } + Ok(driver) => driver, Err(error) => { orch_datastream.emit_bootstrap( None, @@ -240,6 +235,16 @@ where return Err(format!("create iroh driver: {error}")); } }; + let coordinator_endpoint = + advertised_endpoint(driver.endpoint_addr(), config.endpoint_addr_mask)?; + orch_datastream.emit_bootstrap( + None, + config.run_id, + config.node_id, + "iroh_driver", + "ready", + json!({"endpoint":coordinator_endpoint.clone(),"has_relay":coordinator_endpoint.relay_urls().next().is_some(),"direct_addr_count":coordinator_endpoint.ip_addrs().count(),"relay_mode":format!("{:?}", config.relay.mode),"endpoint_addr_mask":config.endpoint_addr_mask.as_str()}), + ); let stack = DistributionRuntimeStack::new_with_codecs( driver.node_id(), DistributedNodeConfig::default(), @@ -425,7 +430,6 @@ where let sink = PluginSink::new(Arc::new(ChannelObservationSink { tx: Mutex::new(obs_tx), })); - let coordinator_endpoint = driver.endpoint_addr(); let pipeline_coordinator_endpoint = coordinator_endpoint.clone(); let (mut provisioned_nodes, ready) = start_and_provision_workers( provisioner, @@ -860,6 +864,7 @@ struct Config { dashboard: bool, max_context: Option, relay: RelayRuntimeConfig, + endpoint_addr_mask: EndpointAddrMask, vastai: Option, cached_model: Option, datastream_frame_log: Option, @@ -889,6 +894,7 @@ struct ConfigBuilder { max_context: Option, relay_mode: Option, relay_url: Option, + endpoint_addr_mask: Option, vastai_api_key: Option, vastai_bootstrap_command: Option, vastai_disk_gb: Option, @@ -944,6 +950,7 @@ impl ConfigBuilder { max_context: None, relay_mode: None, relay_url: None, + endpoint_addr_mask: None, vastai_api_key: None, vastai_bootstrap_command: None, vastai_disk_gb: None, @@ -1172,6 +1179,9 @@ impl ConfigBuilder { { self.relay_url = Some(url); } + if let Some(mask) = env_optional(MVP_IROH_ENDPOINT_ADDR_MASK_ENV) { + self.endpoint_addr_mask = Some(mask); + } if let Some(api_key) = env_optional("MVP_VASTAI_API_KEY").or_else(|| env_optional("VASTAI_API_KEY")) { @@ -1283,6 +1293,9 @@ impl ConfigBuilder { } "--relay-mode" => self.relay_mode = Some(next_arg(&mut args, "--relay-mode")?), "--relay-url" => self.relay_url = Some(next_arg(&mut args, "--relay-url")?), + "--endpoint-addr-mask" => { + self.endpoint_addr_mask = Some(next_arg(&mut args, "--endpoint-addr-mask")?) + } "--vastai-api-key" => { self.vastai_api_key = Some(next_arg(&mut args, "--vastai-api-key")?) } @@ -1392,6 +1405,10 @@ impl ConfigBuilder { self.relay_mode.as_deref(), self.relay_url.as_deref(), )?; + let endpoint_addr_mask = match self.endpoint_addr_mask.as_deref() { + Some(mask) => EndpointAddrMask::parse(mask)?, + None => EndpointAddrMask::Full, + }; let vastai = if provider == ProviderKind::VastAi { Some(VastAiRuntimeConfig::from_builder(&self)?) } else { @@ -1418,6 +1435,7 @@ impl ConfigBuilder { dashboard: self.dashboard, max_context: self.max_context, relay, + endpoint_addr_mask, vastai, cached_model, worker_bin: self.worker_bin, @@ -1736,6 +1754,7 @@ impl Config { "MVP_ORCHESTRATOR_ACTOR", "MVP_MODEL_ID", "MVP_IROH_RELAY_MODE", + MVP_IROH_ENDPOINT_ADDR_MASK_ENV, "MVP_PIPELINE_STAGES", ]; if self.relay.url.is_some() { @@ -1818,6 +1837,10 @@ impl Config { "MVP_PIPELINE_STAGES".to_owned(), self.pipeline_stages.to_string(), ), + ( + MVP_IROH_ENDPOINT_ADDR_MASK_ENV.to_owned(), + self.endpoint_addr_mask.as_str().to_owned(), + ), ( "MVP_NODE_PROVIDER".to_owned(), self.provider.as_str().to_owned(), @@ -1959,9 +1982,9 @@ fn enqueue_runtime_ready_ack( } fn enqueue_datastream_subscribe( - driver: &mut IrohDriver, stack: &DistributionRuntimeStack, ready: &RuntimeReady, + collector: &EndpointAddr, run_id: u64, node_id: u64, ) -> Result<(), String> { @@ -1973,7 +1996,7 @@ fn enqueue_datastream_subscribe( .send_to( ready.datastream_publisher, DatastreamPublisherMsg::Subscribe(DatastreamSubscribe { - collector: driver.endpoint_addr(), + collector: collector.clone(), request: SubscriptionRequest::all(), flow_id, token: Vec::new(), @@ -1998,6 +2021,7 @@ fn wait_for_runtime_ready_acks( orchestrator_node_id: u64, provider: ProviderKind, targets: &[RuntimeReadyAckTarget], + collector_endpoint: &EndpointAddr, ) -> Result<(), String> { let mut pending = targets .iter() @@ -2098,9 +2122,9 @@ fn wait_for_runtime_ready_acks( if stack.route_owner(target.ready.datastream_publisher) == Some(target.ready.swim_node_id) && let Err(error) = enqueue_datastream_subscribe( - driver, stack, &target.ready, + collector_endpoint, run_id, target.node_id, ) @@ -2300,6 +2324,7 @@ fn start_and_provision_workers( "provider":config.provider.as_str(), "image":&config.image, "relay_mode":relay_mode_env_value(&config.relay.mode), + "endpoint_addr_mask":config.endpoint_addr_mask.as_str(), "docker_gpus":if config.provider == ProviderKind::Docker { Some(config.docker_gpus.as_str()) } else { None }, "provider_config":config.provider_datastream_detail(), "env_keys":config.node_spec_env_keys(), @@ -2481,6 +2506,7 @@ fn start_and_provision_workers( config.node_id, config.provider, &ack_targets, + &pipeline_coordinator, )?; orch_datastream.emit_bootstrap( @@ -2797,6 +2823,7 @@ fn wait_for_runtime_readies( ) -> Result, String> { let expected = expected_node_ids.iter().copied().collect::>(); let mut pending = BTreeMap::::new(); + let started = Instant::now(); loop { pump(driver, stack, frame_tx); emit_swim_transitions( @@ -2865,6 +2892,21 @@ fn wait_for_runtime_readies( }) { return Ok(pending); } + if started.elapsed() >= RUNTIME_READY_TIMEOUT { + let pending_list = expected + .iter() + .filter(|node_id| { + pending + .get(node_id) + .is_none_or(|ready| !runtime_ready_barrier_met(stack, ready)) + }) + .map(u64::to_string) + .collect::>() + .join(", "); + return Err(format!( + "runtime_ready timed out for node(s): {pending_list}" + )); + } thread::sleep(PUMP_INTERVAL); } } @@ -3610,6 +3652,7 @@ fn wait_for_runtime_ready( let mut node_swim_started = false; let mut node_swim_ready = false; let mut node_route_started = false; + let started = Instant::now(); loop { pump(driver, stack, frame_tx); drain_frames(frame_rx, dashboard, orch_datastream); @@ -3706,6 +3749,18 @@ fn wait_for_runtime_ready( } } } + if started.elapsed() >= RUNTIME_READY_TIMEOUT { + let detail = pending_ready.as_ref().map(|ready| { + json!({ + "readiness_id":ready.readiness_id, + "swim_alive":stack.member_state(ready.swim_node_id) == Some(MemberState::Alive), + "route_owner":stack.route_owner(ready.node_actor).map(|node| format!("{node:?}")), + }) + }); + return Err(format!( + "runtime_ready timed out for node {node_id}: {detail:?}" + )); + } thread::sleep(PUMP_INTERVAL); } } diff --git a/xtask/src/main.rs b/xtask/src/main.rs index de9cf34..718d429 100644 --- a/xtask/src/main.rs +++ b/xtask/src/main.rs @@ -86,6 +86,12 @@ impl MvpChatCheckScenario { } Self::MultinodeDocker => { args.push("--docker".to_owned()); + args.extend([ + "--relay-mode".to_owned(), + "default".to_owned(), + "--endpoint-addr-mask".to_owned(), + "relay-only".to_owned(), + ]); } } if matches!(self, Self::Multinode | Self::MultinodeDocker) { @@ -1237,6 +1243,9 @@ fn assert_dump_log_facts( if scenario == MvpChatCheckScenario::Gpu { require_gpu_dump_log_facts(&facts)?; } + if scenario == MvpChatCheckScenario::MultinodeDocker { + require_multinode_docker_relay_facts(&facts)?; + } Ok(events) } @@ -1585,6 +1594,10 @@ struct DumpLogFacts { gpu_pipeline_tokenizer_decode_ready: BTreeSet, gpu_pipeline_tokens_decoded: BTreeSet, gpu_pipeline_real_worker_step_seen: bool, + relay_masked_orchestrator_ready: bool, + relay_masked_node_spec_worker_count: Option, + relay_masked_worker_iroh_ready: BTreeSet, + relay_masked_worker_coordinator_join: BTreeSet, chat_config_ready: bool, prepare_runtime_ready: bool, prompt_rpc_ready: bool, @@ -1630,9 +1643,32 @@ fn record_dump_log_event( } (_, Some("OrchBootstrap"), Some("iroh_driver"), Some("ready")) => { facts.orch_iroh_driver_ready = true; + if detail_relay_only_advertisement(event) { + facts.relay_masked_orchestrator_ready = true; + } } (_, Some("NodeEvent"), Some("iroh_driver"), Some("ready")) => { facts.node_iroh_driver_ready = true; + if detail_relay_only_advertisement(event) + && let Some(node_id) = event_node_id(event) + { + facts.relay_masked_worker_iroh_ready.insert(node_id); + } + } + (_, Some("OrchBootstrap"), Some("node_spec"), Some("ready")) => { + if detail_str(event, "endpoint_addr_mask") == Some("relay-only") + && detail_str(event, "relay_mode") == Some("default") + { + facts.relay_masked_node_spec_worker_count = detail_u64(event, "worker_count"); + } + } + (_, Some("NodeEvent"), Some("coordinator_join"), Some("started")) => { + if detail_bool(event, "has_relay") == Some(true) + && detail_u64(event, "direct_addr_count") == Some(0) + && let Some(node_id) = event_node_id(event) + { + facts.relay_masked_worker_coordinator_join.insert(node_id); + } } (_, Some("NodeEvent"), Some("worker_initialize"), Some("ready")) => { facts.node_worker_initialize_ready = true; @@ -1842,6 +1878,27 @@ fn require_gpu_dump_log_facts(facts: &DumpLogFacts) -> Result<(), String> { Ok(()) } +fn require_multinode_docker_relay_facts(facts: &DumpLogFacts) -> Result<(), String> { + require_dump_log_fact( + facts.relay_masked_orchestrator_ready, + "relay-masked orchestrator endpoint", + )?; + require_dump_log_fact( + facts + .relay_masked_node_spec_worker_count + .is_some_and(|count| count >= 2), + "relay-masked Docker node_spec with multiple workers", + )?; + require_dump_log_fact( + facts.relay_masked_worker_iroh_ready.len() >= 2, + "relay-masked worker iroh_driver ready for multiple nodes", + )?; + require_dump_log_fact( + facts.relay_masked_worker_coordinator_join.len() >= 2, + "relay-masked worker coordinator_join for multiple nodes", + ) +} + fn event_requested_device_is_cuda(event: &Value) -> bool { event .get("requested_device") @@ -1900,6 +1957,37 @@ fn is_cuda_device(value: &str) -> bool { value.to_ascii_uppercase().contains("CUDA") } +fn event_node_id(event: &Value) -> Option { + event.get("node_id").and_then(Value::as_u64) +} + +fn detail_bool(event: &Value, key: &str) -> Option { + event + .get("detail") + .and_then(|detail| detail.get(key)) + .and_then(Value::as_bool) +} + +fn detail_u64(event: &Value, key: &str) -> Option { + event + .get("detail") + .and_then(|detail| detail.get(key)) + .and_then(Value::as_u64) +} + +fn detail_str<'a>(event: &'a Value, key: &str) -> Option<&'a str> { + event + .get("detail") + .and_then(|detail| detail.get(key)) + .and_then(Value::as_str) +} + +fn detail_relay_only_advertisement(event: &Value) -> bool { + detail_str(event, "endpoint_addr_mask") == Some("relay-only") + && detail_bool(event, "has_relay") == Some(true) + && detail_u64(event, "direct_addr_count") == Some(0) +} + fn dump_log_request_id(event: &Value) -> Option { event .get("detail") @@ -1983,6 +2071,10 @@ mod tests { multinode_docker.mvp_chat_args(42, dump_log), strings(&[ "--docker", + "--relay-mode", + "default", + "--endpoint-addr-mask", + "relay-only", "--pipeline-stages", "2", "--cached-model", @@ -2733,6 +2825,92 @@ mod tests { ); } + #[test] + fn benchmark_observability_multinode_docker_dump_facts_require_relay_masked_events() { + let mut events = dump_log_fact_events(false, false); + events.extend([ + ( + "mvp.orch.bootstrap", + stamped( + json!({"type":"OrchBootstrap","phase":"iroh_driver","status":"ready","run_id":9,"node_id":1,"detail":{"endpoint_addr_mask":"relay-only","relay_mode":"Default","has_relay":true,"direct_addr_count":0}}), + "mvp-orchestrator", + 1_070, + 70, + ), + ), + ( + "mvp.orch.bootstrap", + stamped( + json!({"type":"OrchBootstrap","phase":"node_spec","status":"ready","run_id":9,"node_id":1,"detail":{"endpoint_addr_mask":"relay-only","relay_mode":"default","worker_count":2}}), + "mvp-orchestrator", + 1_071, + 71, + ), + ), + ( + "mvp.node.bootstrap", + stamped( + json!({"type":"NodeEvent","phase":"iroh_driver","status":"ready","run_id":9,"node_id":2,"stage_index":0,"detail":{"endpoint_addr_mask":"relay-only","has_relay":true,"direct_addr_count":0}}), + "mvp-worker-node", + 1_072, + 72, + ), + ), + ( + "mvp.node.bootstrap", + stamped( + json!({"type":"NodeEvent","phase":"iroh_driver","status":"ready","run_id":9,"node_id":3,"stage_index":1,"detail":{"endpoint_addr_mask":"relay-only","has_relay":true,"direct_addr_count":0}}), + "mvp-worker-node", + 1_073, + 73, + ), + ), + ( + "mvp.node.bootstrap", + stamped( + json!({"type":"NodeEvent","phase":"coordinator_join","status":"started","run_id":9,"node_id":2,"stage_index":0,"detail":{"has_relay":true,"direct_addr_count":0}}), + "mvp-worker-node", + 1_074, + 74, + ), + ), + ( + "mvp.node.bootstrap", + stamped( + json!({"type":"NodeEvent","phase":"coordinator_join","status":"started","run_id":9,"node_id":3,"stage_index":1,"detail":{"has_relay":true,"direct_addr_count":0}}), + "mvp-worker-node", + 1_075, + 75, + ), + ), + ]); + let path = write_synthetic_event_dump("multinode-docker-relay-mask", events); + assert_dump_log_facts(&path, MvpChatCheckScenario::MultinodeDocker) + .expect("relay-masked multinode Docker facts pass"); + let _ = fs::remove_file(path); + } + + #[test] + fn benchmark_observability_multinode_docker_requires_relay_masked_workers() { + let mut valid = DumpLogFacts { + relay_masked_orchestrator_ready: true, + relay_masked_node_spec_worker_count: Some(2), + ..DumpLogFacts::default() + }; + valid.relay_masked_worker_iroh_ready.extend([2, 3]); + valid.relay_masked_worker_coordinator_join.extend([2, 3]); + require_multinode_docker_relay_facts(&valid).expect("relay-masked Docker facts pass"); + + let mut missing_worker = valid; + missing_worker.relay_masked_worker_iroh_ready.remove(&3); + let error = require_multinode_docker_relay_facts(&missing_worker) + .expect_err("single relay-masked worker should fail"); + assert!( + error.contains("relay-masked worker iroh_driver ready for multiple nodes"), + "unexpected error: {error}" + ); + } + #[test] fn benchmark_observability_gpu_dump_facts_reject_cpu_fallback() { let path = write_synthetic_event_dump("gpu-cpu-fallback", dump_log_fact_events(true, true));