zacheryasc/swactor
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

feat: mvp-chat multinode docker test with network masking

Browse source 
Add endpoint-address masking and a relay-only advertisement path so the multinode Docker mvp-chat scenario can run with direct addresses stripped.

- endpoint_advertisement: add EndpointAddrMask (Full/RelayOnly) parsed from --endpoint-addr-mask/MVP_IROH_ENDPOINT_ADDR_MASK, and advertised_endpoint that rebuilds an EndpointAddr from relay URLs only, rejecting relay-only without a relay URL
- orchestrator_app: mask the coordinator endpoint before advertising it, thread the masked collector endpoint into datastream subscribe/runtime-ready acks, surface endpoint_addr_mask/has_relay/direct_addr_count in iroh_driver and node_spec events, forward the mask env to workers, and add a 60s RUNTIME_READY_TIMEOUT to the runtime-ready barriers
- worker_node: advertise the masked self endpoint in the iroh_driver ready and coordinator_join events and propagate it through runtime_ready_local and PendingRuntimeReady
- mvp-chat: add --relay-mode/--relay-url/--endpoint-addr-mask plus a [relay] toml section, require (with a Vast.ai fallback) a relay URL when relay-only, and forward all three to the orchestrator CLI
- node_image: resolve the worker binary to a workspace-relative path for the Docker COPY via docker_build_context_path, rejecting paths outside the build context
- xtask/specs: run MultinodeDocker with --relay-mode default --endpoint-addr-mask relay-only, add dump-log fact checks for relay-masked orchestrator/node/coordinator advertisement, and document the mask/relay flags in mvp_chat.md

Signed-off-by: Zachery Aaron Shores-Chmielewski <zacheryasc@gmail.com>
This commit is contained in:
Zachery Aaron Shores-Chmielewski 2026-07-23 14:58:03 +04:00
parent f7243dbc3b
commit 42bf9cceff
8 changed files with 560 additions and 30 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

22
crates/mvp-system/specs/mvp_chat.md
View file

@ -73,6 +73,10 @@ General flags:
- `--cached-model`
- `--cached-model=<path>`
- `--skip-rebuild`
- `--gpu`
- `--relay-mode <mode>`
- `--relay-url <url>`
- `--endpoint-addr-mask <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 <count>` accepts a positive integer. Zero and invalid values
are configuration errors.
`--relay-mode <mode>` accepts `default` or `disabled`.
`--endpoint-addr-mask <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;

128
crates/mvp-system/src/bin/mvp_chat.rs
View file

@ -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 <path> Load config overlay
--pipeline-stages <count> Number of pipeline stages
--relay-mode <mode> Relay mode: default or disabled
--relay-url <url> Custom relay URL passed to mvp-orchestrator
--endpoint-addr-mask <mask> Endpoint address mask: full or relay-only
--cached-model[=<path>] Use discovered or explicit cached GGUF model
--dump-logs[=<path>] Write datastream frame log
--run-id <id> Override run id
@ -331,6 +335,9 @@ struct Config {
max_tokens: u32,
skip_rebuild: bool,
gpu_run: bool,
relay_mode: Option<String>,
relay_url: Option<String>,
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<u32>,
}
#[derive(Clone, Debug, Default, Deserialize)]
#[serde(default, deny_unknown_fields)]
struct ChatRelayConfig {
mode: Option<String>,
url: Option<String>,
endpoint_addr_mask: Option<String>,
}
#[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<CachedModelSource>,
help: bool,
gpu: bool,
relay_mode: Option<String>,
relay_url: Option<String>,
endpoint_addr_mask: Option<String>,
}
#[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<bool, String>,

38
crates/mvp-system/src/bin/worker_node.rs
View file

@ -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<String>,
debug_join_socket: Option<String>,
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(),

104
crates/mvp-system/src/endpoint_advertisement.rs Normal file
View file

@ -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<Self, String> {
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<EndpointAddr, String> {
match mask {
EndpointAddrMask::Full => Ok(endpoint),
EndpointAddrMask::RelayOnly => relay_only_endpoint(endpoint),
}
}
fn relay_only_endpoint(endpoint: EndpointAddr) -> Result<EndpointAddr, String> {
let relays = endpoint.relay_urls().cloned().collect::<Vec<_>>();
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::<iroh::RelayUrl>()
.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}"
);
}
}

1
crates/mvp-system/src/lib.rs
View file

@ -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;

34
crates/mvp-system/src/node_image.rs
View file

@ -175,7 +175,7 @@ pub fn prepare_node_image(request: NodeImageRequest) -> Result<PreparedNodeImage
)?;
}
let node_bin = request.node_bin.to_string_lossy().to_string();
let node_bin = docker_build_context_path(&root, &request.node_bin)?;
let mut build_args = vec![
"build".to_owned(),
"-f".to_owned(),
@ -349,6 +349,15 @@ fn relative_path(root: &Path, path: &Path) -> Result<PathBuf, String> {
})
}
fn docker_build_context_path(root: &Path, path: &Path) -> Result<String, String> {
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"
);
}
}

85
crates/mvp-system/src/orchestrator_app.rs
View file

@ -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<u32>,
relay: RelayRuntimeConfig,
endpoint_addr_mask: EndpointAddrMask,
vastai: Option<VastAiRuntimeConfig>,
cached_model: Option<CachedModelConfig>,
datastream_frame_log: Option<PathBuf>,
@ -889,6 +894,7 @@ struct ConfigBuilder {
max_context: Option<u32>,
relay_mode: Option<String>,
relay_url: Option<String>,
endpoint_addr_mask: Option<String>,
vastai_api_key: Option<String>,
vastai_bootstrap_command: Option<String>,
vastai_disk_gb: Option<u32>,
@ -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<BTreeMap<u64, RuntimeReady>, String> {
let expected = expected_node_ids.iter().copied().collect::<BTreeSet<_>>();
let mut pending = BTreeMap::<u64, RuntimeReady>::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::<Vec<_>>()
.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);
}
}

178
xtask/src/main.rs
View file

@ -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<u64>,
gpu_pipeline_tokens_decoded: BTreeSet<u64>,
gpu_pipeline_real_worker_step_seen: bool,
relay_masked_orchestrator_ready: bool,
relay_masked_node_spec_worker_count: Option<u64>,
relay_masked_worker_iroh_ready: BTreeSet<u64>,
relay_masked_worker_coordinator_join: BTreeSet<u64>,
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<u64> {
event.get("node_id").and_then(Value::as_u64)
}
fn detail_bool(event: &Value, key: &str) -> Option<bool> {
event
.get("detail")
.and_then(|detail| detail.get(key))
.and_then(Value::as_bool)
}
fn detail_u64(event: &Value, key: &str) -> Option<u64> {
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<u64> {
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));