From 0902de639457811f882d2f45da17e42f4d2581c4 Mon Sep 17 00:00:00 2001 From: Zachery Aaron Shores-Chmielewski Date: Thu, 30 Jul 2026 15:31:23 +0400 Subject: [PATCH] refactor: more `yoke grind` --- crates/datastream/src/publisher_actor.rs | 4 +- crates/iroh-driver/src/lib.rs | 6 +- crates/mvp-system/src/chat/runtime.rs | 62 +- crates/mvp-system/src/lib.rs | 2 +- .../src/node/worker_node_runtime.rs | 576 ++++++++---------- crates/mvp-system/src/orchestration/app.rs | 543 +++++++---------- .../src/orchestration/distribution_stack.rs | 68 ++- crates/transport/src/json_codec.rs | 2 +- 8 files changed, 551 insertions(+), 712 deletions(-) diff --git a/crates/datastream/src/publisher_actor.rs b/crates/datastream/src/publisher_actor.rs index d31c540..acc96df 100644 --- a/crates/datastream/src/publisher_actor.rs +++ b/crates/datastream/src/publisher_actor.rs @@ -76,7 +76,5 @@ impl ActorInterface for DatastreamPublisherActor { /// Register JSON encoding for remote datastream publisher messages. pub fn register_datastream_publisher_codec(registry: &mut CodecRegistry) { - registry.register::( - JsonCodec::::default(), - ); + registry.register::(JsonCodec::::default()); } diff --git a/crates/iroh-driver/src/lib.rs b/crates/iroh-driver/src/lib.rs index df44fe5..c3a7138 100644 --- a/crates/iroh-driver/src/lib.rs +++ b/crates/iroh-driver/src/lib.rs @@ -10,13 +10,13 @@ pub mod edge_transport; pub mod endpoint_advertisement; pub mod iroh_driver; +pub use endpoint_advertisement::{ + EndpointAddrMask, MVP_IROH_ENDPOINT_ADDR_MASK_ENV, advertised_endpoint, +}; pub use iroh_driver::{ ConnType, DatastreamPublishHandle, IrohDriver, IrohDriverConfig, JoinPhase, JoinStatus, conn_type_of, discover_lan_ips, }; -pub use endpoint_advertisement::{ - MVP_IROH_ENDPOINT_ADDR_MASK_ENV, EndpointAddrMask, advertised_endpoint, -}; pub use edge_transport::{EDGE_ALPN, EdgeSendHandle, EdgeTransportEvent, EdgeTransportFault}; diff --git a/crates/mvp-system/src/chat/runtime.rs b/crates/mvp-system/src/chat/runtime.rs index fdbb120..c3b904e 100644 --- a/crates/mvp-system/src/chat/runtime.rs +++ b/crates/mvp-system/src/chat/runtime.rs @@ -30,11 +30,11 @@ use crate::node_provisioning::{ProviderKind, provider_kind}; use crate::observability::{benchmark, frame_archive::FrameArchive}; use crate::orchestration::config::ResolvedVastAiConfig; use crate::prompt::rpc::{PromptEvent, SubmitPrompt, write_json_line}; -use iroh_driver::EndpointAddrMask; use crate::{ DEFAULT_PIPELINE_CACHED_MODEL_FILE, DEFAULT_PIPELINE_CACHED_MODEL_ID, DEFAULT_PIPELINE_CACHED_MODEL_MAX_CONTEXT, DEFAULT_PIPELINE_CACHED_MODEL_REPO, }; +use iroh_driver::EndpointAddrMask; const DEFAULT_CONFIG_PATH: &str = ".config/config.toml"; const DEFAULT_RPC_ADDR: &str = "127.0.0.1:19777"; @@ -150,8 +150,7 @@ where ); progress.emit_benchmark_envelope(&config); progress.emit_endpoint_config_snapshot(&config); - let mut approval = StdinVastAiApproval; - confirm_vastai_if_needed_with_approval(&config, &mut approval)?; + confirm_vastai_if_needed(&config)?; let prepare_runtime_started = Instant::now(); progress.emit( CHAT_RUNTIME_CHANNEL, @@ -1264,58 +1263,39 @@ fn first_non_empty(values: [Option; N]) -> Option bool; - fn ask(&mut self) -> Result; -} - -struct StdinVastAiApproval; - -impl VastAiApproval for StdinVastAiApproval { - fn stdin_is_terminal(&self) -> bool { - io::stdin().is_terminal() - } - - fn ask(&mut self) -> Result { - #[cfg(test)] - { - let mut input = std::io::Cursor::new(Vec::::new()); - let mut output = io::sink(); - ask_vastai_approval(&mut input, &mut output) - } - #[cfg(not(test))] - { - let stdin = io::stdin(); - let mut input = stdin.lock(); - let mut output = io::stdout(); - ask_vastai_approval(&mut input, &mut output) - } - } -} - -fn confirm_vastai_if_needed_with_approval( - config: &Config, - approval: &mut A, -) -> Result<(), String> -where - A: VastAiApproval, -{ +fn confirm_vastai_if_needed(config: &Config) -> Result<(), String> { if config.vastai.is_none() { return Ok(()); } if config.vastai_yes { return Ok(()); } - if !approval.stdin_is_terminal() { + if !io::stdin().is_terminal() { return Err("Vast.ai rental requires --yes when stdin is not a terminal".to_owned()); } - if approval.ask()? { + if prompt_vastai_approval()? { Ok(()) } else { Err("Vast.ai rental declined".to_owned()) } } +fn prompt_vastai_approval() -> Result { + #[cfg(test)] + { + let mut input = std::io::Cursor::new(Vec::::new()); + let mut output = io::sink(); + ask_vastai_approval(&mut input, &mut output) + } + #[cfg(not(test))] + { + let stdin = io::stdin(); + let mut input = stdin.lock(); + let mut output = io::stdout(); + ask_vastai_approval(&mut input, &mut output) + } +} + fn ask_vastai_approval(input: &mut R, output: &mut W) -> Result where R: BufRead, diff --git a/crates/mvp-system/src/lib.rs b/crates/mvp-system/src/lib.rs index 49fd11a..95d0ed1 100644 --- a/crates/mvp-system/src/lib.rs +++ b/crates/mvp-system/src/lib.rs @@ -43,12 +43,12 @@ mod run_fsm; mod run_plan; mod chat; +mod codecs; mod node; mod observability; mod orchestration; mod prompt; mod staging; -mod codecs; #[cfg(test)] mod tests; diff --git a/crates/mvp-system/src/node/worker_node_runtime.rs b/crates/mvp-system/src/node/worker_node_runtime.rs index 5365832..9e510ff 100644 --- a/crates/mvp-system/src/node/worker_node_runtime.rs +++ b/crates/mvp-system/src/node/worker_node_runtime.rs @@ -22,34 +22,31 @@ use datastream::{ Lifetime, NodeId, Record, StreamDescriptor, StreamId, StreamOrigin, }; -use iroh_driver::driver_pumps as driver_model; +use crate::codecs::register_mvp_actor_codecs; use crate::gguf_shard::{StageShardPlan, materialize_stage_shard_http, validate_stage_shard_cache}; use crate::node_actor::{ NodeAgentActor, NodeAgentMsg, NodeAgentReport, StageCommandWire, StageInboundEdgeWire, StageObjectSpecWire, StageOutboundEdgeWire, }; use crate::observability::benchmark; -use crate::orchestration::distribution_stack::DistributionRuntimeStack; +use crate::orchestration::distribution_stack::{DistributionRuntimeStack, duration_ms_u64}; use crate::orchestration::provider_adapters::relay::relay_runtime_config_from_env; use crate::prompt::rpc::{PromptEvent, TokenizerEvent}; use crate::run_plan::{GgufSource, TokenizerSource}; use crate::staging::control as stage; -use iroh_driver::{ - EndpointAddrMask, MVP_IROH_ENDPOINT_ADDR_MASK_ENV, advertised_endpoint, -}; -use crate::codecs::register_mvp_actor_codecs; use data_plane::arena; use data_plane::edge_lifecycle as edge; use data_plane::ingress; use distribution::node::DistributedNodeConfig; -use distribution::swim::telemetry::ObservedProbeEvent; use distribution::telemetry::{MembershipTransition, SwimProbeEvent}; use distribution::types::{MemberState, NodeId as DistNodeId}; use iroh::EndpointAddr; +use iroh_driver::driver_pumps as driver_model; use iroh_driver::{ DATASTREAM_ALPN, DatastreamPublishHandle, DatastreamQuicHeader, EDGE_ALPN, EdgeSendHandle, EdgeTransportEvent, IrohDriver, IrohDriverConfig, }; +use iroh_driver::{EndpointAddrMask, MVP_IROH_ENDPOINT_ADDR_MASK_ENV, advertised_endpoint}; use parking_lot::Mutex; use serde_json::{Value, json}; use swactor::actor::{ActorAddress, ActorInterface}; @@ -805,6 +802,9 @@ impl WorkerEdgeRuntime { config: &DeploymentConfig, datastream: &mut NodeDatastream, ) -> Result<(), String> { + let node_stage = |ds: &mut NodeDatastream, phase: &str, status: &str, detail: Value| { + emit_node_event(ds, config, NODE_STAGE_CHANNEL, phase, status, detail) + }; driver.pump_edge_ingress(); for event in driver.drain_edge_events() { match event { @@ -815,10 +815,8 @@ impl WorkerEdgeRuntime { driver_model::EdgeId(edge_id), driver_model::StreamId(stream_id), ); - emit_node_event( + node_stage( datastream, - config, - NODE_STAGE_CHANNEL, "iroh_edge_stream_arrived", "observed", json!({"edge_id":edge_id,"stream_id":stream_id}), @@ -840,10 +838,8 @@ impl WorkerEdgeRuntime { .. } => { let byte_count = bytes.len(); - emit_node_event( + node_stage( datastream, - config, - NODE_STAGE_CHANNEL, "iroh_edge_bytes_read", "observed", json!({"edge_id":edge_id,"stream_id":stream_id,"bytes":byte_count}), @@ -992,6 +988,9 @@ impl WorkerEdgeRuntime { datastream: &mut NodeDatastream, driver: &mut IrohDriver, ) -> Result<(), String> { + let node_stage = |ds: &mut NodeDatastream, phase: &str, status: &str, detail: Value| { + emit_node_event(ds, config, NODE_STAGE_CHANNEL, phase, status, detail) + }; let input_key = ObjectKey { edge_id: input_edge_id, object_id, @@ -1050,10 +1049,8 @@ impl WorkerEdgeRuntime { }; let egress_read_ms = duration_ms_u64(egress_read_started.elapsed()); let record_bytes = record.len(); - emit_node_event( + node_stage( datastream, - config, - NODE_STAGE_CHANNEL, "egress_ring_read", "ready", json!({ @@ -1080,10 +1077,8 @@ impl WorkerEdgeRuntime { let edge_send_started = Instant::now(); sender.send(record)?; let edge_send_ms = duration_ms_u64(edge_send_started.elapsed()); - emit_node_event( + node_stage( datastream, - config, - NODE_STAGE_CHANNEL, "iroh_edge_bytes_sent", "ready", json!({ @@ -1130,6 +1125,9 @@ impl WorkerEdgeRuntime { datastream: &mut NodeDatastream, driver: &mut IrohDriver, ) -> Result<(), String> { + let node_stage = |ds: &mut NodeDatastream, phase: &str, status: &str, detail: Value| { + emit_node_event(ds, config, NODE_STAGE_CHANNEL, phase, status, detail) + }; let Some(inbound) = self.inbound_edge.clone() else { return Ok(()); }; @@ -1161,10 +1159,8 @@ impl WorkerEdgeRuntime { .map_err(|e| format!("write ingress ring: {e}"))?; } let ingress_ring_write_ms = duration_ms_u64(ring_write_started.elapsed()); - emit_node_event( + node_stage( datastream, - config, - NODE_STAGE_CHANNEL, "ingress_ring_write", "ready", json!({ @@ -1197,10 +1193,8 @@ impl WorkerEdgeRuntime { object_id: loaded.object_id, }; self.object_handles.insert(key, loaded.clone()); - emit_node_event( + node_stage( datastream, - config, - NODE_STAGE_CHANNEL, "object_loaded", "ready", json!({ @@ -1538,10 +1532,6 @@ impl WorkerEdgeRuntime { } } -fn duration_ms_u64(duration: Duration) -> u64 { - u64::try_from(duration.as_millis()).unwrap_or(u64::MAX) -} - struct IngressRecordBytes { bytes: Vec, object_id: u64, @@ -1636,18 +1626,26 @@ fn run() -> Result<(), String> { let worker_evt = |phase: &str, status: &str, detail: Value| { emit_stdio_node_event(&config, NODE_WORKER_CHANNEL, phase, status, detail) }; - boot("config", "ready", json!({ - "worker_script":&config.worker_script, - "device":&config.device, - "model_id":&config.model_id, - "has_coordinator_endpoint":config.coordinator_endpoint.is_some(), - "has_orchestrator_actor":config.orchestrator_actor.is_some(), - "self_test_enabled":config.self_test_prompt.is_some(), - "arena_bytes":config.arena_bytes, - "arena_alignment":config.arena_alignment, - "debug_join_socket":config.debug_join_socket.as_deref().unwrap_or("disabled"), - }))?; - boot("process", "started", json!({"binary":"mvp-worker-node","pid":std::process::id()}))?; + boot( + "config", + "ready", + json!({ + "worker_script":&config.worker_script, + "device":&config.device, + "model_id":&config.model_id, + "has_coordinator_endpoint":config.coordinator_endpoint.is_some(), + "has_orchestrator_actor":config.orchestrator_actor.is_some(), + "self_test_enabled":config.self_test_prompt.is_some(), + "arena_bytes":config.arena_bytes, + "arena_alignment":config.arena_alignment, + "debug_join_socket":config.debug_join_socket.as_deref().unwrap_or("disabled"), + }), + )?; + boot( + "process", + "started", + json!({"binary":"mvp-worker-node","pid":std::process::id()}), + )?; let tokio = match tokio::runtime::Runtime::new() { Ok(runtime) => { @@ -1655,7 +1653,11 @@ fn run() -> Result<(), String> { runtime } Err(error) => { - boot("tokio_runtime", "failed", json!({"error":error.to_string()}))?; + boot( + "tokio_runtime", + "failed", + json!({"error":error.to_string()}), + )?; return Err(format!("tokio runtime: {error}")); } }; @@ -1677,12 +1679,24 @@ fn run() -> Result<(), String> { }; let advertised_self_endpoint = advertised_endpoint(driver.endpoint_addr(), config.endpoint_addr_mask)?; - boot("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()}))?; + boot( + "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)); - boot("coordinator_join", "started", json!({"endpoint":coordinator,"has_relay":coordinator.relay_urls().next().is_some(),"direct_addr_count":coordinator.ip_addrs().count()}))?; + boot( + "coordinator_join", + "started", + json!({"endpoint":coordinator,"has_relay":coordinator.relay_urls().next().is_some(),"direct_addr_count":coordinator.ip_addrs().count()}), + )?; } else { - boot("coordinator_join", "skipped", json!({"reason":"MVP_COORDINATOR_ENDPOINT not set","mode":"standalone"}))?; + boot( + "coordinator_join", + "skipped", + json!({"reason":"MVP_COORDINATOR_ENDPOINT not set","mode":"standalone"}), + )?; } let stack = DistributionRuntimeStack::new_with_codecs( @@ -1693,8 +1707,16 @@ fn run() -> Result<(), String> { datastream::wire::register_datastream_codec(registry); }, ); - boot("distribution_stack", "ready", json!({"actors":"initialized","route_view":"initialized","swim":"initialized","outbox":"initialized"}))?; - boot("codecs", "ready", json!({"registered":["node_agent","orchestrator","provisioner","prompt_rpc","datastream"]}))?; + boot( + "distribution_stack", + "ready", + json!({"actors":"initialized","route_view":"initialized","swim":"initialized","outbox":"initialized"}), + )?; + boot( + "codecs", + "ready", + json!({"registered":["node_agent","orchestrator","provisioner","prompt_rpc","datastream"]}), + )?; driver.enable_actor_bridge( stack.runtime.clone(), stack.codec.clone(), @@ -1703,7 +1725,11 @@ fn run() -> Result<(), String> { stack.relay_mirror.clone(), stack.route_view.clone(), ); - boot("actor_bridge", "ready", json!({"transport":"iroh","routes":"attached"}))?; + boot( + "actor_bridge", + "ready", + json!({"transport":"iroh","routes":"attached"}), + )?; let arena_manager = match arena::ArenaManager::boot(arena::ArenaConfig { node_id: arena::NodeId(config.logical_node_id), @@ -1711,20 +1737,37 @@ fn run() -> Result<(), String> { base_alignment: config.arena_alignment, }) { Ok(manager) => { - boot("arena_manager", "ready", json!({ - "arena_bytes":config.arena_bytes, - "arena_alignment":config.arena_alignment, - }))?; + boot( + "arena_manager", + "ready", + json!({ + "arena_bytes":config.arena_bytes, + "arena_alignment":config.arena_alignment, + }), + )?; Arc::new(Mutex::new(manager)) } Err(error) => { - boot("arena_manager", "failed", json!({"error":format!("{error:?}")}))?; + boot( + "arena_manager", + "failed", + json!({"error":format!("{error:?}")}), + )?; return Err(format!("boot arena manager: {error:?}")); } }; let arena_fd = arena_manager.lock().arena_fd(); let mut datastream = NodeDatastream::new(&config); + let node_boot = |ds: &mut NodeDatastream, phase: &str, status: &str, detail: Value| { + emit_node_event(ds, &config, NODE_BOOTSTRAP_CHANNEL, phase, status, detail) + }; + let node_runtime = |ds: &mut NodeDatastream, phase: &str, status: &str, detail: Value| { + emit_node_event(ds, &config, NODE_RUNTIME_CHANNEL, phase, status, detail) + }; + let node_shutdown = |ds: &mut NodeDatastream, phase: &str, status: &str, detail: Value| { + emit_node_event(ds, &config, NODE_SHUTDOWN_CHANNEL, phase, status, detail) + }; let datastream_transport = driver.datastream_publish_handle(); let datastream_publisher = match stack .runtime @@ -1732,10 +1775,8 @@ fn run() -> Result<(), String> { { Ok(actor) => actor, Err(error) => { - emit_node_event( + node_boot( &mut datastream, - &config, - NODE_BOOTSTRAP_CHANNEL, "datastream_publisher", "failed", json!({"error":error.to_string()}), @@ -1766,10 +1807,8 @@ fn run() -> Result<(), String> { datastream.channels.arena, Arc::clone(&arena_manager), ); - emit_node_event( + node_boot( &mut datastream, - &config, - NODE_BOOTSTRAP_CHANNEL, "datastream_publisher", "ready", json!({"actor":datastream_publisher,"name":DATASTREAM_PUBLISHER_NAME,"subscription_transport":"iroh"}), @@ -1784,10 +1823,8 @@ fn run() -> Result<(), String> { ("DatastreamSyntheticEventSent", "sent"), ("DatastreamSyntheticEventObserved", "observed"), ] { - emit_node_event( + node_boot( &mut datastream, - &config, - NODE_BOOTSTRAP_CHANNEL, phase, status, json!({ @@ -1807,10 +1844,8 @@ fn run() -> Result<(), String> { Some(path) => { match spawn_debug_join_listener(tokio.handle().clone(), PathBuf::from(path)) { Ok(rx) => { - emit_node_event( + node_runtime( &mut datastream, - &config, - NODE_RUNTIME_CHANNEL, "debug_join_socket", "ready", json!({"socket":path}), @@ -1818,10 +1853,8 @@ fn run() -> Result<(), String> { Some(rx) } Err(error) => { - emit_node_event( + node_runtime( &mut datastream, - &config, - NODE_RUNTIME_CHANNEL, "debug_join_socket", "failed", json!({"socket":path,"error":error}), @@ -1831,10 +1864,8 @@ fn run() -> Result<(), String> { } } None => { - emit_node_event( + node_runtime( &mut datastream, - &config, - NODE_RUNTIME_CHANNEL, "debug_join_socket", "skipped", json!({"reason":"MVP_DEBUG_JOIN_SOCKET=disabled"}), @@ -1849,7 +1880,11 @@ fn run() -> Result<(), String> { inbox } Err(error) => { - boot("node_report_inbox", "failed", json!({"error":error.to_string()}))?; + boot( + "node_report_inbox", + "failed", + json!({"error":error.to_string()}), + )?; return Err(format!("node report inbox: {error}")); } }; @@ -1857,7 +1892,11 @@ fn run() -> Result<(), String> { "MVP_ORCHESTRATOR_ACTOR is required for runtime readiness signaling".to_owned() })?; let orchestrator_source = "env"; - boot("orchestrator_actor", "ready", json!({"actor":orchestrator,"source":orchestrator_source}))?; + boot( + "orchestrator_actor", + "ready", + json!({"actor":orchestrator,"source":orchestrator_source}), + )?; let node_agent = NodeAgentActor::new( stage::NodeId(config.logical_node_id), orchestrator, @@ -1865,25 +1904,41 @@ fn run() -> Result<(), String> { ); let node_actor = match stack.runtime.spawn(node_agent) { Ok(actor) => { - boot("node_agent", "ready", json!({"node_actor":actor,"source":"generated"}))?; + boot( + "node_agent", + "ready", + json!({"node_actor":actor,"source":"generated"}), + )?; actor } Err(error) => { - boot("node_agent", "failed", json!({"error":error.to_string(),"source":"generated"}))?; + boot( + "node_agent", + "failed", + json!({"error":error.to_string(),"source":"generated"}), + )?; return Err(format!("spawn node agent: {error}")); } }; stack.register_local_actor(driver.register_actor(node_actor, 1)); - boot("node_actor_registration", "ready", json!({"node_actor":node_actor,"network_reachable":true}))?; + boot( + "node_actor_registration", + "ready", + json!({"node_actor":node_actor,"network_reachable":true}), + )?; - worker_evt("worker_process", "started", json!({ - "program":"python3", - "script":&config.worker_script, - "device":&config.device, - "stdin":"piped", - "stdout":"piped", - "stderr":"piped", - }))?; + worker_evt( + "worker_process", + "started", + json!({ + "program":"python3", + "script":&config.worker_script, + "device":&config.device, + "stdin":"piped", + "stdout":"piped", + "stderr":"piped", + }), + )?; let mut worker = match TinygradWorker::spawn(&config, arena_fd) { Ok(worker) => worker, Err(error) => { @@ -1899,10 +1954,18 @@ fn run() -> Result<(), String> { sampler_health_context, vec![std::process::id(), worker.pid()], ); - worker_evt("worker_initialize", "started", json!({"command":"InitializeWorker","helper_abi_version":1,"device":&config.device}))?; + worker_evt( + "worker_initialize", + "started", + json!({"command":"InitializeWorker","helper_abi_version":1,"device":&config.device}), + )?; let mut initial_pump = || {}; match worker.initialize(&config.device, &config, &mut datastream, &mut initial_pump) { - Ok(()) => worker_evt("worker_initialize", "ready", json!({"worker_event_type":"WorkerReady"}))?, + Ok(()) => worker_evt( + "worker_initialize", + "ready", + json!({"worker_event_type":"WorkerReady"}), + )?, Err(error) => { worker_evt("worker_initialize", "failed", json!({"error":error}))?; return Err(error); @@ -1925,13 +1988,17 @@ fn run() -> Result<(), String> { "logical_node_id": config.logical_node_id, "stage_index": config.stage_index, }); - boot("runtime_ready_local", "ready", json!({ - "endpoint":advertised_self_endpoint.clone(), - "node_actor":node_actor, - "logical_node_id":config.logical_node_id, - "stage_index":config.stage_index, - "readiness_id":pending_runtime_ready.readiness_id, - }))?; + boot( + "runtime_ready_local", + "ready", + json!({ + "endpoint":advertised_self_endpoint.clone(), + "node_actor":node_actor, + "logical_node_id":config.logical_node_id, + "stage_index":config.stage_index, + "readiness_id":pending_runtime_ready.readiness_id, + }), + )?; if let Some(prompt) = &config.self_test_prompt { run_self_test( @@ -1945,18 +2012,14 @@ fn run() -> Result<(), String> { } let shutdown_rx = spawn_stdin_shutdown_listener(); - emit_node_event( + node_runtime( &mut datastream, - &config, - NODE_RUNTIME_CHANNEL, "stdin_shutdown_listener", "ready", json!({"command":"shutdown"}), ); - emit_node_event( + node_runtime( &mut datastream, - &config, - NODE_RUNTIME_CHANNEL, "main_loop", "started", json!({ @@ -2000,10 +2063,8 @@ fn run() -> Result<(), String> { } => { if pending_runtime_ready.observe_ack(run_id, node_id, stage_index, readiness_id) { - emit_node_event( + node_boot( &mut datastream, - &config, - NODE_BOOTSTRAP_CHANNEL, "runtime_ready_ack", "ready", json!({ @@ -2014,10 +2075,8 @@ fn run() -> Result<(), String> { }), ); datastream.submit_text(datastream.channels.node_ready, ready.to_string()); - emit_node_event( + node_boot( &mut datastream, - &config, - NODE_BOOTSTRAP_CHANNEL, "datastream_handoff", "ready", json!({"from":"runtime_ready_ack","to":"cluster_datastream","channel":"mvp.node.ready"}), @@ -2027,10 +2086,8 @@ fn run() -> Result<(), String> { } } if !pending_runtime_ready.swim_logged && pending_runtime_ready.swim_ready(&stack) { - emit_node_event( + node_runtime( &mut datastream, - &config, - NODE_RUNTIME_CHANNEL, "coordinator_swim", "ready", json!({ @@ -2044,10 +2101,8 @@ fn run() -> Result<(), String> { pending_runtime_ready.swim_logged = true; } if !pending_runtime_ready.acked && pending_runtime_ready.maybe_send(&stack, node_actor)? { - emit_node_event( + node_runtime( &mut datastream, - &config, - NODE_RUNTIME_CHANNEL, "runtime_ready_signal", "sent", json!({ @@ -2058,10 +2113,8 @@ fn run() -> Result<(), String> { ); } if shutdown_rx.try_recv().is_ok() { - emit_node_event( + node_shutdown( &mut datastream, - &config, - NODE_SHUTDOWN_CHANNEL, "shutdown", "started", json!({"source":"stdin","command":"shutdown"}), @@ -2069,27 +2122,21 @@ fn run() -> Result<(), String> { let mut pump = || pump_network(&mut driver, &stack); match worker.shutdown(&config, &mut datastream, &mut pump) { Ok(()) => { - emit_node_event( + node_shutdown( &mut datastream, - &config, - NODE_SHUTDOWN_CHANNEL, "worker_shutdown", "ready", json!({"worker_event_type":"WorkerStopped"}), ); - emit_node_event( + node_shutdown( &mut datastream, - &config, - NODE_SHUTDOWN_CHANNEL, "node_exit", "ready", json!({"result":"ok"}), ); } - Err(error) => emit_node_event( + Err(error) => node_shutdown( &mut datastream, - &config, - NODE_SHUTDOWN_CHANNEL, "worker_shutdown", "failed", json!({"error":error}), @@ -2098,10 +2145,8 @@ fn run() -> Result<(), String> { return Ok(()); } if let Some(status) = worker.try_wait()? { - emit_node_event( + node_shutdown( &mut datastream, - &config, - NODE_SHUTDOWN_CHANNEL, "worker_process", "failed", json!({"exit_status":status.to_string()}), @@ -2132,75 +2177,19 @@ fn emit_swim_telemetry( local_phase: &str, ) { for transition in stack.drain_swim_transitions() { - let peer = format!("{:?}", transition.peer); - let from = transition.from.map(|state| format!("{:?}", state)); - let to = format!("{:?}", transition.to); - let member_state = stack - .member_state(transition.peer) - .map(|state| format!("{:?}", state)); - let record = MembershipTransition { - peer, - from: from.unwrap_or_default(), - to, - reason: transition.reason.to_owned(), - last_ack_age_ms: transition.last_ack_age.map(duration_ms_u64), - consecutive_timeouts: transition.consecutive_timeouts, - recent_probe_targets: swim_recent_probe_targets(stack), - member_state, - }; + let record = stack.membership_transition(&transition); datastream .producer .submit_record(datastream.channels.membership, &record); } for event in stack.drain_swim_probe_events() { - let record = swim_probe_event_record(stack, event, local_phase); + let record = stack.swim_probe_event_record(event, local_phase); datastream .producer .submit_record(datastream.channels.swim_probes, &record); } } -fn swim_probe_event_record( - stack: &DistributionRuntimeStack, - event: ObservedProbeEvent, - local_phase: &str, -) -> SwimProbeEvent { - let config = &stack.swim_config; - let budget_ms = event.budget_ms; - SwimProbeEvent { - event: event.event.to_owned(), - target: format!("{:?}", event.target), - sequence: event.sequence, - kind: event.kind.to_owned(), - rtt_ms: event.rtt_ms, - budget_ms, - budget_ticks: budget_ms, - last_ack_age_ms: event.last_ack_age.map(duration_ms_u64), - consecutive_timeouts: event.consecutive_timeouts, - recent_probe_targets: swim_recent_probe_targets(stack), - member_state: stack - .member_state(event.target) - .map(|state| format!("{:?}", state)), - local_phase: local_phase.to_owned(), - probe_interval_ms: duration_ms_u64(config.probe_interval), - probe_timeout_ms: duration_ms_u64(config.probe_timeout), - indirect_probes: u32::try_from(config.indirect_probes).unwrap_or(u32::MAX), - suspicion_timeout_ms: duration_ms_u64(config.suspicion_timeout), - dead_reprobe_interval_ms: duration_ms_u64(config.dead_reprobe_interval), - probe_mode: format!("{:?}", config.probe_mode), - lifeguard_enabled: config.lifeguard.is_some(), - } -} - -fn swim_recent_probe_targets(stack: &DistributionRuntimeStack) -> Vec { - stack - .swim_telemetry - .recent_targets() - .into_iter() - .map(|node_id| format!("{:?}", node_id)) - .collect() -} - #[derive(Clone, Copy)] struct DatastreamChannelSet { node_ready: ChannelId, @@ -2566,6 +2555,9 @@ fn handle_node_report( arena_manager: &Arc>, datastream: &mut NodeDatastream, ) -> Result { + let node_stage = |ds: &mut NodeDatastream, phase: &str, status: &str, detail: Value| { + emit_node_event(ds, config, NODE_STAGE_CHANNEL, phase, status, detail) + }; let kind = match &report { NodeAgentReport::Command(_) => "Command", NodeAgentReport::Lifecycle(_) => "Lifecycle", @@ -2604,14 +2596,7 @@ fn handle_node_report( datastream.channels.node_lifecycle, json!({"type":"node_lifecycle","event":event}).to_string(), ); - emit_node_event( - datastream, - config, - NODE_STAGE_CHANNEL, - "lifecycle", - "observed", - json!({"event":event}), - ); + node_stage(datastream, "lifecycle", "observed", json!({"event":event})); Ok(NodeReportOutcome::None) } NodeAgentReport::PromptRequested { @@ -2671,19 +2656,18 @@ fn handle_prompt_request( worker: &mut TinygradWorker, datastream: &mut NodeDatastream, ) -> Result<(), String> { + let node_prompt = |ds: &mut NodeDatastream, phase: &str, status: &str, detail: Value| { + emit_node_event(ds, config, NODE_PROMPT_CHANNEL, phase, status, detail) + }; let started = Instant::now(); - emit_node_event( + node_prompt( datastream, - config, - NODE_PROMPT_CHANNEL, "prompt_requested", "started", json!({"request_id":request_id,"max_tokens":max_tokens,"reply_to":reply_to,"prompt_bytes":prompt.len()}), ); - emit_node_event( + node_prompt( datastream, - config, - NODE_PROMPT_CHANNEL, "infer_prompt", "started", json!({"request_id":request_id,"command":"InferPrompt","max_tokens":max_tokens}), @@ -2712,10 +2696,8 @@ fn handle_prompt_request( .get("elapsed_ms") .and_then(Value::as_u64) .unwrap_or_else(|| started.elapsed().as_millis() as u64); - emit_node_event( + node_prompt( datastream, - config, - NODE_PROMPT_CHANNEL, "infer_prompt", "ready", json!({"request_id":request_id,"worker_event_type":"PromptCompleted","prompt_tokens":prompt_tokens,"tokens_generated":tokens_generated,"elapsed_ms":elapsed_ms,"text_bytes":text_bytes,"worker_result_payload_bytes":worker_result_payload_bytes}), @@ -2728,19 +2710,15 @@ fn handle_prompt_request( text: text.clone(), }, ) { - Ok(()) => emit_node_event( + Ok(()) => node_prompt( datastream, - config, - NODE_PROMPT_CHANNEL, "prompt_response", "ready", json!({"request_id":request_id,"event":"TextDelta","bytes":text_bytes,"reply_to":reply_to}), ), Err(error) => { - emit_node_event( + node_prompt( datastream, - config, - NODE_PROMPT_CHANNEL, "prompt_response", "failed", json!({"request_id":request_id,"event":"TextDelta","error":error.to_string()}), @@ -2759,10 +2737,8 @@ fn handle_prompt_request( }, ) { Ok(()) => { - emit_node_event( + node_prompt( datastream, - config, - NODE_PROMPT_CHANNEL, "prompt_response", "ready", json!({"request_id":request_id,"event":"Done","tokens_generated":tokens_generated,"elapsed_ms":elapsed_ms,"final_text_bytes":text_bytes,"reply_to":reply_to}), @@ -2770,10 +2746,8 @@ fn handle_prompt_request( Ok(()) } Err(error) => { - emit_node_event( + node_prompt( datastream, - config, - NODE_PROMPT_CHANNEL, "prompt_response", "failed", json!({"request_id":request_id,"event":"Done","error":error.to_string()}), @@ -2783,10 +2757,8 @@ fn handle_prompt_request( } } Err(error) => { - emit_node_event( + node_prompt( datastream, - config, - NODE_PROMPT_CHANNEL, "infer_prompt", "failed", json!({"request_id":request_id,"error":error}), @@ -2799,10 +2771,8 @@ fn handle_prompt_request( }, ) { Ok(()) => { - emit_node_event( + node_prompt( datastream, - config, - NODE_PROMPT_CHANNEL, "prompt_response", "ready", json!({"request_id":request_id,"event":"Fault","reply_to":reply_to}), @@ -2810,10 +2780,8 @@ fn handle_prompt_request( Ok(()) } Err(send_error) => { - emit_node_event( + node_prompt( datastream, - config, - NODE_PROMPT_CHANNEL, "prompt_response", "failed", json!({"request_id":request_id,"event":"Fault","error":send_error.to_string()}), @@ -2835,10 +2803,11 @@ fn handle_encode_prompt_request( worker: &mut TinygradWorker, datastream: &mut NodeDatastream, ) -> Result<(), String> { - emit_node_event( + let node_prompt = |ds: &mut NodeDatastream, phase: &str, status: &str, detail: Value| { + emit_node_event(ds, config, NODE_PROMPT_CHANNEL, phase, status, detail) + }; + node_prompt( datastream, - config, - NODE_PROMPT_CHANNEL, "encode_prompt", "started", json!({"request_id":request_id,"prompt_bytes":prompt.len(),"reply_to":reply_to}), @@ -2846,10 +2815,8 @@ fn handle_encode_prompt_request( let mut pump = || pump_network(driver, stack); let event = match worker.encode_prompt(request_id, &prompt, config, datastream, &mut pump) { Ok(tokens) => { - emit_node_event( + node_prompt( datastream, - config, - NODE_PROMPT_CHANNEL, "encode_prompt", "ready", json!({"request_id":request_id,"tokens":tokens.len(),"reply_to":reply_to}), @@ -2857,10 +2824,8 @@ fn handle_encode_prompt_request( TokenizerEvent::PromptEncoded { request_id, tokens } } Err(error) => { - emit_node_event( + node_prompt( datastream, - config, - NODE_PROMPT_CHANNEL, "encode_prompt", "failed", json!({"request_id":request_id,"error":error,"reply_to":reply_to}), @@ -2884,10 +2849,11 @@ fn handle_decode_tokens_request( worker: &mut TinygradWorker, datastream: &mut NodeDatastream, ) -> Result<(), String> { - emit_node_event( + let node_prompt = |ds: &mut NodeDatastream, phase: &str, status: &str, detail: Value| { + emit_node_event(ds, config, NODE_PROMPT_CHANNEL, phase, status, detail) + }; + node_prompt( datastream, - config, - NODE_PROMPT_CHANNEL, "decode_tokens", "started", json!({"request_id":request_id,"tokens":tokens.len(),"reply_to":reply_to}), @@ -2895,10 +2861,8 @@ fn handle_decode_tokens_request( let mut pump = || pump_network(driver, stack); let event = match worker.decode_tokens(request_id, &tokens, config, datastream, &mut pump) { Ok(text) => { - emit_node_event( + node_prompt( datastream, - config, - NODE_PROMPT_CHANNEL, "decode_tokens", "ready", json!({"request_id":request_id,"tokens":tokens.len(),"text_bytes":text.len(),"reply_to":reply_to}), @@ -2906,10 +2870,8 @@ fn handle_decode_tokens_request( TokenizerEvent::TokensDecoded { request_id, text } } Err(error) => { - emit_node_event( + node_prompt( datastream, - config, - NODE_PROMPT_CHANNEL, "decode_tokens", "failed", json!({"request_id":request_id,"error":error,"reply_to":reply_to}), @@ -3376,6 +3338,9 @@ fn handle_stage_command( arena_manager: &Arc>, datastream: &mut NodeDatastream, ) -> Result<(), String> { + let node_stage = |ds: &mut NodeDatastream, phase: &str, status: &str, detail: Value| { + emit_node_event(ds, config, NODE_STAGE_CHANNEL, phase, status, detail) + }; match command { StageCommandWire::ConfigureWorkerRole { run_id, @@ -3383,10 +3348,8 @@ fn handle_stage_command( layer_start, layer_end_exclusive, } => { - emit_node_event( + node_stage( datastream, - config, - NODE_STAGE_CHANNEL, "configure_worker_role", "started", json!({"run_id":run_id,"stage_index":stage_index,"layer_range":{"start":layer_start,"end_exclusive":layer_end_exclusive}}), @@ -3401,19 +3364,15 @@ fn handle_stage_command( datastream, &mut pump, ) { - Ok(()) => emit_node_event( + Ok(()) => node_stage( datastream, - config, - NODE_STAGE_CHANNEL, "configure_worker_role", "ready", json!({"worker_event_type":"RoleConfigured"}), ), Err(error) => { - emit_node_event( + node_stage( datastream, - config, - NODE_STAGE_CHANNEL, "configure_worker_role", "failed", json!({"error":error}), @@ -3463,10 +3422,8 @@ fn handle_stage_command( true, ), Err(error) => { - emit_node_event( + node_stage( datastream, - config, - NODE_STAGE_CHANNEL, "load_weights", "failed", json!({"error":error,"stage_shard":true}), @@ -3484,10 +3441,8 @@ fn handle_stage_command( } else { (gguf_source, false) }; - emit_node_event( + node_stage( datastream, - config, - NODE_STAGE_CHANNEL, "load_weights", "started", json!({"model_id":&model_id,"gguf_source":gguf_source_kind,"tokenizer":tokenizer_kind,"stage_shard":using_stage_shard,"layer_range":{"start":layer_start,"end_exclusive":layer_end_exclusive}}), @@ -3503,23 +3458,14 @@ fn handle_stage_command( datastream, &mut pump, ) { - Ok(()) => emit_node_event( + Ok(()) => node_stage( datastream, - config, - NODE_STAGE_CHANNEL, "load_weights", "ready", json!({"worker_event_type":"WeightsLoaded","model_id":model_id}), ), Err(error) => { - emit_node_event( - datastream, - config, - NODE_STAGE_CHANNEL, - "load_weights", - "failed", - json!({"error":error}), - ); + node_stage(datastream, "load_weights", "failed", json!({"error":error})); let _ = stack.runtime.send_to( node_actor, NodeAgentMsg::WorkerCrashed { @@ -3551,10 +3497,8 @@ fn handle_stage_command( .runtime .send_to(node_actor, NodeAgentMsg::WorkerRingsQuiesced { run_id }) .map_err(|e| format!("mark worker rings quiesced: {e}"))?; - emit_node_event( + node_stage( datastream, - config, - NODE_STAGE_CHANNEL, "local_edges", "ready", json!({"run_id":run_id,"sent":["LocalEdgesStopped","WorkerRingsQuiesced"]}), @@ -3570,10 +3514,8 @@ fn handle_stage_command( .runtime .send_to(node_actor, NodeAgentMsg::WorkerRoleReset { run_id }) .map_err(|e| format!("mark worker role reset: {e}"))?; - emit_node_event( + node_stage( datastream, - config, - NODE_STAGE_CHANNEL, "device_objects", "ready", json!({"run_id":run_id,"sent":["DeviceObjectsReleased","WorkerRoleReset"]}), @@ -3581,10 +3523,8 @@ fn handle_stage_command( Ok(()) } StageCommandWire::EstablishInboundEdge { edge_id, edge } => { - emit_node_event( + node_stage( datastream, - config, - NODE_STAGE_CHANNEL, "inbound_edge", "started", json!({"edge_id":edge_id,"kind":format!("{:?}", edge.kind),"ring_data_capacity":edge.ring_spec.data_capacity}), @@ -3599,10 +3539,8 @@ fn handle_stage_command( datastream, driver, )?; - emit_node_event( + node_stage( datastream, - config, - NODE_STAGE_CHANNEL, "inbound_edge", "ready", json!({"edge_id":edge_id}), @@ -3610,10 +3548,8 @@ fn handle_stage_command( Ok(()) } StageCommandWire::EstablishOutboundEdge { edge_id, edge } => { - emit_node_event( + node_stage( datastream, - config, - NODE_STAGE_CHANNEL, "outbound_edge", "started", json!({"edge_id":edge_id,"kind":format!("{:?}", edge.kind),"consumer_node_id":edge.consumer_node_id,"has_consumer_endpoint":edge.consumer_endpoint.is_some()}), @@ -3628,10 +3564,8 @@ fn handle_stage_command( datastream, driver, )?; - emit_node_event( + node_stage( datastream, - config, - NODE_STAGE_CHANNEL, "outbound_edge", "ready", json!({"edge_id":edge_id}), @@ -3639,10 +3573,8 @@ fn handle_stage_command( Ok(()) } StageCommandWire::RewireEdge { .. } => { - emit_node_event( + node_stage( datastream, - config, - NODE_STAGE_CHANNEL, "rewire_edge", "skipped", json!({"reason":"not implemented in mvp-worker-node image path"}), @@ -3659,10 +3591,8 @@ fn handle_stage_command( sequence, .. } => { - emit_node_event( + node_stage( datastream, - config, - NODE_STAGE_CHANNEL, "execute_step", "started", json!({"step_id":step_id,"input_edge_id":input_edge_id,"object_id":object_id,"sequence":sequence}), @@ -3680,10 +3610,8 @@ fn handle_stage_command( datastream, driver, )?; - emit_node_event( + node_stage( datastream, - config, - NODE_STAGE_CHANNEL, "execute_step", "ready", json!({"step_id":step_id}), @@ -3932,10 +3860,11 @@ fn wait_for_helper_event( wait_config: HelperCommandWaitConfig, pump: &mut dyn FnMut(), ) -> Result { - emit_node_event( + let node_worker = |ds: &mut NodeDatastream, phase: &str, status: &str, detail: Value| { + emit_node_event(ds, config, NODE_WORKER_CHANNEL, phase, status, detail) + }; + node_worker( datastream, - config, - NODE_WORKER_CHANNEL, "worker_stdout_read", "started", json!({"command_type":command_type,"expected_event_type":expected,"channel":channel_name}), @@ -3950,18 +3879,14 @@ fn wait_for_helper_event( drain_worker_stderr(stderr_rx, config, datastream); } let line_bytes = line.len(); - emit_node_event( + node_worker( datastream, - config, - NODE_WORKER_CHANNEL, "worker_stdout_read", "ready", json!({"command_type":command_type,"expected_event_type":expected,"channel":channel_name,"line_bytes":line_bytes}), ); - emit_node_event( + node_worker( datastream, - config, - NODE_WORKER_CHANNEL, "worker_stdout_parse", "started", json!({"command_type":command_type,"expected_event_type":expected,"channel":channel_name,"line_bytes":line_bytes}), @@ -3969,10 +3894,8 @@ fn wait_for_helper_event( let value: Value = match serde_json::from_str(&line) { Ok(value) => value, Err(error) => { - emit_node_event( + node_worker( datastream, - config, - NODE_WORKER_CHANNEL, "worker_stdout_parse", "failed", json!({"command_type":command_type,"expected_event_type":expected,"channel":channel_name,"line_bytes":line_bytes,"error":error.to_string()}), @@ -3984,10 +3907,8 @@ fn wait_for_helper_event( .get("type") .and_then(Value::as_str) .unwrap_or("unknown"); - emit_node_event( + node_worker( datastream, - config, - NODE_WORKER_CHANNEL, "worker_stdout_parse", "ready", json!({"command_type":command_type,"expected_event_type":expected,"channel":channel_name,"line_bytes":line_bytes,"worker_event_type":worker_event_type}), @@ -4003,10 +3924,8 @@ fn wait_for_helper_event( if value.get("type").and_then(Value::as_str) == Some(expected) { return Ok(value); } - emit_node_event( + node_worker( datastream, - config, - NODE_WORKER_CHANNEL, "worker_event", "observed", json!({"command_type":command_type,"command_waiting_for":expected,"worker_event_type":worker_event_type,"event":value}), @@ -4016,10 +3935,8 @@ fn wait_for_helper_event( if let Some(stderr_rx) = stderr_rx { drain_worker_stderr(stderr_rx, config, datastream); } - emit_node_event( + node_worker( datastream, - config, - NODE_WORKER_CHANNEL, "worker_stdout_read", "failed", json!({"command_type":command_type,"expected_event_type":expected,"channel":channel_name,"line_bytes":0,"error":"stdout closed"}), @@ -4032,10 +3949,8 @@ fn wait_for_helper_event( if let Some(stderr_rx) = stderr_rx { drain_worker_stderr(stderr_rx, config, datastream); } - emit_node_event( + node_worker( datastream, - config, - NODE_WORKER_CHANNEL, "worker_stdout_read", "failed", json!({"command_type":command_type,"expected_event_type":expected,"channel":channel_name,"error":error}), @@ -4050,10 +3965,8 @@ fn wait_for_helper_event( } let now = Instant::now(); if now >= next_telemetry_at { - emit_node_event( + node_worker( datastream, - config, - NODE_WORKER_CHANNEL, "worker_command_wait", "waiting", json!({ @@ -4073,10 +3986,8 @@ fn wait_for_helper_event( datastream.tick(); } Err(RecvTimeoutError::Disconnected) => { - emit_node_event( + node_worker( datastream, - config, - NODE_WORKER_CHANNEL, "worker_stdout_read", "failed", json!({"command_type":command_type,"expected_event_type":expected,"channel":channel_name,"line_bytes":0,"error":"stdout reader disconnected"}), @@ -4504,6 +4415,9 @@ impl TinygradWorker { channel_name: &str, pump: &mut dyn FnMut(), ) -> Result { + let node_worker = |ds: &mut NodeDatastream, phase: &str, status: &str, detail: Value| { + emit_node_event(ds, config, NODE_WORKER_CHANNEL, phase, status, detail) + }; let channel = datastream.channel_by_name(channel_name); let command_type = command .get("type") @@ -4512,19 +4426,15 @@ impl TinygradWorker { .to_owned(); let command_text = command.to_string(); let command_bytes = command_text.len() + 1; - emit_node_event( + node_worker( datastream, - config, - NODE_WORKER_CHANNEL, "worker_command_write", "started", json!({"command_type":command_type.as_str(),"expected_event_type":expected,"command_bytes":command_bytes}), ); if let Err(error) = writeln!(self.stdin, "{command_text}") { - emit_node_event( + node_worker( datastream, - config, - NODE_WORKER_CHANNEL, "worker_command_write", "failed", json!({"command_type":command_type.as_str(),"expected_event_type":expected,"command_bytes":command_bytes,"error":error.to_string()}), @@ -4532,20 +4442,16 @@ impl TinygradWorker { return Err(format!("write helper command: {error}")); } if let Err(error) = self.stdin.flush() { - emit_node_event( + node_worker( datastream, - config, - NODE_WORKER_CHANNEL, "worker_command_write", "failed", json!({"command_type":command_type.as_str(),"expected_event_type":expected,"command_bytes":command_bytes,"error":error.to_string()}), ); return Err(format!("flush helper command: {error}")); } - emit_node_event( + node_worker( datastream, - config, - NODE_WORKER_CHANNEL, "worker_command_write", "ready", json!({"command_type":command_type.as_str(),"expected_event_type":expected,"command_bytes":command_bytes}), diff --git a/crates/mvp-system/src/orchestration/app.rs b/crates/mvp-system/src/orchestration/app.rs index 3f15e0d..b868356 100644 --- a/crates/mvp-system/src/orchestration/app.rs +++ b/crates/mvp-system/src/orchestration/app.rs @@ -11,6 +11,7 @@ use std::thread; use std::time::{Duration, Instant}; use crate::DEFAULT_PIPELINE_CACHED_MODEL_FILE; +use crate::codecs::register_mvp_actor_codecs; use crate::node_actor::{ NodeAgentMsg, StageEdgeKindWire, StageInboundEdgeWire, StageObjectSpecWire, StageOutboundEdgeWire, StageProvisionWire, StageRingSpecWire, @@ -20,7 +21,6 @@ use crate::observability::dashboard_view::MvpClusterDashboardView; use crate::observability::{benchmark, frame_archive::FrameArchive}; use crate::orchestration::actor::{OrchestratorActor, OrchestratorReport}; use crate::orchestration::config::{DEFAULT_CONFIG_PATH, TomlConfigOverlay}; -use crate::codecs::register_mvp_actor_codecs; const PROVIDER_START_MAX_ATTEMPTS: usize = 4; use crate::gguf_shard::{StageShardPlan, plan_stage_shard}; @@ -29,7 +29,7 @@ use crate::observability::telemetry::{ MVP_PROVISIONING_EVENTS, MvpProvisionEventRecord, MvpProvisionLogRecord, mvp_provision_log_channel, }; -use crate::orchestration::distribution_stack::DistributionRuntimeStack; +use crate::orchestration::distribution_stack::{DistributionRuntimeStack, duration_ms_u64}; use crate::orchestration::provider_adapters::relay::{ MVP_IROH_RELAY_URL_ENV, RelayRuntimeConfig, SWACTOR_IROH_RELAY_URL_ENV, relay_mode_env_value, relay_runtime_config_from_settings, @@ -48,9 +48,6 @@ use crate::provisioning::{ }; use crate::run_fsm::{RunConfig, RunId}; use crate::run_plan::{self, GgufSource, TokenizerSource}; -use iroh_driver::{ - EndpointAddrMask, MVP_IROH_ENDPOINT_ADDR_MASK_ENV, advertised_endpoint, -}; use data_plane::object_record as ingress; use datastream::{ ChannelContent, ChannelId, ChannelRef, DatastreamEndpoint, DatastreamEvent, DatastreamProducer, @@ -58,13 +55,13 @@ use datastream::{ StreamId, StreamOrigin, SubscriptionRequest, }; use distribution::node::DistributedNodeConfig; -use distribution::swim::telemetry::ObservedProbeEvent; use distribution::telemetry::{MembershipTransition, SwimProbeEvent}; use distribution::types::{MemberState, NodeId as DistNodeId}; use iroh::EndpointAddr; use iroh_driver::{ DATASTREAM_ALPN, EDGE_ALPN, EdgeSendHandle, EdgeTransportEvent, IrohDriver, IrohDriverConfig, }; +use iroh_driver::{EndpointAddrMask, MVP_IROH_ENDPOINT_ADDR_MASK_ENV, advertised_endpoint}; use parking_lot::Mutex; use serde_json::{Value, json}; use swactor::actor::ActorAddress; @@ -114,10 +111,18 @@ where }; let mut orch_datastream = OrchDatastream::new(config.run_id, config.datastream_frame_log.as_deref())?; - orch_datastream.emit_bootstrap( + let run_id = config.run_id; + let node_id = config.node_id; + let bootstrap = |ds: &mut OrchDatastream, + dash: Option<&DashboardSupport>, + phase: &str, + status: &str, + detail: Value| { + ds.emit_bootstrap(dash, run_id, node_id, phase, status, detail); + }; + bootstrap( + &mut orch_datastream, None, - config.run_id, - config.node_id, "config", "ready", json!({ @@ -137,10 +142,9 @@ where "provider_config":config.provider_datastream_detail(), }), ); - orch_datastream.emit_bootstrap( + bootstrap( + &mut orch_datastream, None, - config.run_id, - config.node_id, "datastream_preflight", "configured", json!({ @@ -165,10 +169,9 @@ where ("DatastreamSyntheticEventSent", "sent"), ("DatastreamSyntheticEventObserved", "observed"), ] { - orch_datastream.emit_bootstrap( + bootstrap( + &mut orch_datastream, None, - config.run_id, - config.node_id, phase, status, json!({ @@ -193,10 +196,9 @@ where ); let pipeline_plan = if config.uses_planned_execution() { let plan = config.build_run_plan()?; - orch_datastream.emit_bootstrap( + bootstrap( + &mut orch_datastream, None, - config.run_id, - config.node_id, "run_plan", "ready", json!({ @@ -215,10 +217,9 @@ where let tokio = match tokio::runtime::Runtime::new() { Ok(runtime) => { - orch_datastream.emit_bootstrap( + bootstrap( + &mut orch_datastream, None, - config.run_id, - config.node_id, "tokio_runtime", "ready", json!({"runtime":"tokio"}), @@ -226,10 +227,9 @@ where runtime } Err(error) => { - orch_datastream.emit_bootstrap( + bootstrap( + &mut orch_datastream, None, - config.run_id, - config.node_id, "tokio_runtime", "failed", json!({"error":error.to_string()}), @@ -249,10 +249,9 @@ where ) { Ok(driver) => driver, Err(error) => { - orch_datastream.emit_bootstrap( + bootstrap( + &mut orch_datastream, None, - config.run_id, - config.node_id, "iroh_driver", "failed", json!({"error":error.to_string()}), @@ -262,18 +261,16 @@ where }; let coordinator_endpoint = advertised_endpoint(driver.endpoint_addr(), config.endpoint_addr_mask)?; - orch_datastream.emit_bootstrap( + bootstrap( + &mut orch_datastream, 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()}), ); - orch_datastream.emit_bootstrap( + bootstrap( + &mut orch_datastream, None, - config.run_id, - config.node_id, "endpoint_config_snapshot", "ready", json!({ @@ -294,18 +291,16 @@ where datastream::wire::register_datastream_codec(registry); }, ); - orch_datastream.emit_bootstrap( + bootstrap( + &mut orch_datastream, None, - config.run_id, - config.node_id, "distribution_stack", "ready", json!({"actors":"initialized","route_view":"initialized","swim":"initialized"}), ); - orch_datastream.emit_bootstrap( + bootstrap( + &mut orch_datastream, None, - config.run_id, - config.node_id, "codecs", "ready", json!({"registered":["node_agent","orchestrator","provisioner","prompt_rpc","datastream"]}), @@ -318,29 +313,26 @@ where stack.relay_mirror.clone(), stack.route_view.clone(), ); - orch_datastream.emit_bootstrap( + bootstrap( + &mut orch_datastream, None, - config.run_id, - config.node_id, "actor_bridge", "ready", json!({"transport":"iroh","routes":"attached"}), ); let (frame_tx, frame_rx) = mpsc::channel::(); - orch_datastream.emit_bootstrap( + bootstrap( + &mut orch_datastream, None, - config.run_id, - config.node_id, "datastream_collector", "ready", json!({"alpn":String::from_utf8_lossy(DATASTREAM_ALPN)}), ); let dashboard = DashboardSupport::start(config.dashboard)?; - orch_datastream.emit_bootstrap( + bootstrap( + &mut orch_datastream, dashboard.as_ref(), - config.run_id, - config.node_id, "dashboard", "ready", json!({"enabled":dashboard.is_some()}), @@ -349,10 +341,9 @@ where let orchestrator_reports = match stack.runtime.new_inbox::() { Ok(inbox) => inbox, Err(error) => { - orch_datastream.emit_bootstrap( + bootstrap( + &mut orch_datastream, dashboard.as_ref(), - config.run_id, - config.node_id, "orchestrator_report_actor", "failed", json!({"error":error.to_string()}), @@ -362,10 +353,9 @@ where }; let orchestrator_report_actor = *orchestrator_reports.addr(); stack.register_local_actor(driver.register_actor(orchestrator_report_actor, 1)); - orch_datastream.emit_bootstrap( + bootstrap( + &mut orch_datastream, dashboard.as_ref(), - config.run_id, - config.node_id, "orchestrator_report_actor", "ready", json!({"actor":orchestrator_report_actor}), @@ -380,10 +370,9 @@ where )) { Ok(actor) => actor, Err(error) => { - orch_datastream.emit_bootstrap( + bootstrap( + &mut orch_datastream, dashboard.as_ref(), - config.run_id, - config.node_id, "orchestrator_actor", "failed", json!({"error":error.to_string()}), @@ -392,10 +381,9 @@ where } }; stack.register_local_actor(driver.register_actor(orchestrator_actor, 1)); - orch_datastream.emit_bootstrap( + bootstrap( + &mut orch_datastream, dashboard.as_ref(), - config.run_id, - config.node_id, "orchestrator_actor", "ready", json!({"actor":orchestrator_actor}), @@ -404,10 +392,9 @@ where let prompt_events = match stack.runtime.new_inbox::() { Ok(inbox) => inbox, Err(error) => { - orch_datastream.emit_bootstrap( + bootstrap( + &mut orch_datastream, dashboard.as_ref(), - config.run_id, - config.node_id, "prompt_reply_actor", "failed", json!({"error":error.to_string()}), @@ -417,10 +404,9 @@ where }; let prompt_reply_actor = *prompt_events.addr(); stack.register_local_actor(driver.register_actor(prompt_reply_actor, 1)); - orch_datastream.emit_bootstrap( + bootstrap( + &mut orch_datastream, dashboard.as_ref(), - config.run_id, - config.node_id, "prompt_reply_actor", "ready", json!({"actor":prompt_reply_actor}), @@ -429,10 +415,9 @@ where let tokenizer_events = match stack.runtime.new_inbox::() { Ok(inbox) => inbox, Err(error) => { - orch_datastream.emit_bootstrap( + bootstrap( + &mut orch_datastream, dashboard.as_ref(), - config.run_id, - config.node_id, "tokenizer_reply_actor", "failed", json!({"error":error.to_string()}), @@ -442,10 +427,9 @@ where }; let tokenizer_reply_actor = *tokenizer_events.addr(); stack.register_local_actor(driver.register_actor(tokenizer_reply_actor, 1)); - orch_datastream.emit_bootstrap( + bootstrap( + &mut orch_datastream, dashboard.as_ref(), - config.run_id, - config.node_id, "tokenizer_reply_actor", "ready", json!({"actor":tokenizer_reply_actor}), @@ -455,10 +439,9 @@ where let stop_rx = stop_rx.unwrap_or_else(spawn_stop_listener); let provisioner = config.build_provisioner(Arc::clone(&stack.runtime))?; - orch_datastream.emit_bootstrap( + bootstrap( + &mut orch_datastream, dashboard.as_ref(), - config.run_id, - config.node_id, "node_provisioner", "ready", json!({ @@ -497,10 +480,9 @@ where orchestrator_actor, )?; - orch_datastream.emit_bootstrap( + bootstrap( + &mut orch_datastream, dashboard.as_ref(), - config.run_id, - config.node_id, "prompt_rpc", "started", json!({ @@ -511,10 +493,9 @@ where let rpc_addr = match spawn_prompt_rpc(config.rpc_bind, work_tx, config.default_max_tokens) { Ok(addr) => { - orch_datastream.emit_bootstrap( + bootstrap( + &mut orch_datastream, dashboard.as_ref(), - config.run_id, - config.node_id, "prompt_rpc", "ready", json!({ @@ -525,10 +506,9 @@ where addr } Err(error) => { - orch_datastream.emit_bootstrap( + bootstrap( + &mut orch_datastream, dashboard.as_ref(), - config.run_id, - config.node_id, "prompt_rpc", "failed", json!({"error":error}), @@ -537,19 +517,17 @@ where } }; - orch_datastream.emit_bootstrap( + bootstrap( + &mut orch_datastream, dashboard.as_ref(), - config.run_id, - config.node_id, "prompt_loop", "ready", json!({"addr":rpc_addr.to_string(),"node_actor":ready.first_stage.node_actor}), ); - orch_datastream.emit_bootstrap( + bootstrap( + &mut orch_datastream, dashboard.as_ref(), - config.run_id, - config.node_id, "serve_prompts", "started", json!({"mode":"single_active_prompt","poll_interval_ms":PUMP_INTERVAL.as_millis()}), @@ -582,47 +560,42 @@ where ready.first_stage.endpoint.clone(), ); if let Err(error) = &result { - orch_datastream.emit_bootstrap( + bootstrap( + &mut orch_datastream, dashboard.as_ref(), - config.run_id, - config.node_id, "serve_prompts", "failed", json!({"error":error}), ); } - orch_datastream.emit_bootstrap( + bootstrap( + &mut orch_datastream, dashboard.as_ref(), - config.run_id, - config.node_id, "provider_stop", "started", json!({"provider":config.provider.as_str(),"node_id":config.node_id}), ); let stop_result = provisioned_nodes.stop(); match &stop_result { - Ok(()) => orch_datastream.emit_bootstrap( + Ok(()) => bootstrap( + &mut orch_datastream, dashboard.as_ref(), - config.run_id, - config.node_id, "provider_stop", "ready", json!({"provider":config.provider.as_str(),"node_id":config.node_id}), ), - Err(error) => orch_datastream.emit_bootstrap( + Err(error) => bootstrap( + &mut orch_datastream, dashboard.as_ref(), - config.run_id, - config.node_id, "provider_stop", "failed", json!({"provider":config.provider.as_str(),"node_id":config.node_id,"error":error}), ), } if result.is_ok() && stop_result.is_ok() { - orch_datastream.emit_bootstrap( + bootstrap( + &mut orch_datastream, dashboard.as_ref(), - config.run_id, - config.node_id, "orch_exit", "ready", json!({"result":"ok"}), @@ -2101,6 +2074,16 @@ fn wait_for_runtime_ready_acks( orchestrator_node_id, provider, } = ctx; + let bootstrap = |ds: &mut OrchDatastream, phase: &str, status: &str, detail: Value| { + ds.emit_bootstrap( + dashboard, + run_id, + orchestrator_node_id, + phase, + status, + detail, + ); + }; let mut pending = targets .iter() .cloned() @@ -2157,10 +2140,8 @@ fn wait_for_runtime_ready_acks( let Some(target) = pending.remove(&key) else { continue; }; - orch_datastream.emit_bootstrap( - dashboard, - run_id, - orchestrator_node_id, + bootstrap( + orch_datastream, "runtime_ready_ack", "ready", json!({ @@ -2186,10 +2167,8 @@ fn wait_for_runtime_ready_acks( target.node_id, ) { - orch_datastream.emit_bootstrap( - dashboard, - run_id, - orchestrator_node_id, + bootstrap( + orch_datastream, "datastream_subscribe", "failed", json!({"node_id":target.node_id,"error":error}), @@ -2199,10 +2178,8 @@ fn wait_for_runtime_ready_acks( let attempt = attempts.entry(*key).or_default(); *attempt += 1; let attempt = *attempt; - orch_datastream.emit_bootstrap( - dashboard, - run_id, - orchestrator_node_id, + bootstrap( + orch_datastream, "runtime_ready_ack", "sent", json!({ @@ -2297,6 +2274,11 @@ fn start_and_provision_workers( orch_stdio_rx, .. } = ctx; + let run_id = config.run_id; + let node_id = config.node_id; + let bootstrap = |ds: &mut OrchDatastream, phase: &str, status: &str, detail: Value| { + ds.emit_bootstrap(dashboard, run_id, node_id, phase, status, detail); + }; let stage_specs = stage_node_specs( config, pipeline_plan, @@ -2307,10 +2289,8 @@ fn start_and_provision_workers( .iter() .map(|spec| spec.node_id) .collect::>(); - orch_datastream.emit_bootstrap( - dashboard, - config.run_id, - config.node_id, + bootstrap( + orch_datastream, "node_spec", "ready", json!({ @@ -2357,10 +2337,8 @@ fn start_and_provision_workers( )), }, ); - orch_datastream.emit_bootstrap( - dashboard, - config.run_id, - config.node_id, + bootstrap( + orch_datastream, "provider_start", "started", json!({ @@ -2406,10 +2384,8 @@ fn start_and_provision_workers( for (node_spec, handle_result) in start_results { match handle_result { Ok(handle) => { - orch_datastream.emit_bootstrap( - dashboard, - config.run_id, - config.node_id, + bootstrap( + orch_datastream, "provider_start", "ready", json!({ @@ -2422,10 +2398,8 @@ fn start_and_provision_workers( handles.push(handle); } Err(error) => { - orch_datastream.emit_bootstrap( - dashboard, - config.run_id, - config.node_id, + bootstrap( + orch_datastream, "provider_start", "failed", json!({ @@ -2471,10 +2445,8 @@ fn start_and_provision_workers( config.node_id, ); - orch_datastream.emit_bootstrap( - dashboard, - config.run_id, - config.node_id, + bootstrap( + orch_datastream, "node_runtime_ready", "started", json!({"worker_count":expected_node_ids.len(),"node_ids":expected_node_ids}), @@ -2500,10 +2472,8 @@ fn start_and_provision_workers( ) { Ok(readies) => readies, Err(error) => { - orch_datastream.emit_bootstrap( - dashboard, - config.run_id, - config.node_id, + bootstrap( + orch_datastream, "node_runtime_ready", "failed", json!({"error":error}), @@ -2529,10 +2499,8 @@ fn start_and_provision_workers( }) { Ok(ready) => ready, Err(error) => { - orch_datastream.emit_bootstrap( - dashboard, - config.run_id, - config.node_id, + bootstrap( + orch_datastream, "node_runtime_ready", "failed", json!({"error":error}), @@ -2543,10 +2511,8 @@ fn start_and_provision_workers( BTreeMap::from([(config.node_id, ready)]) }; for (node_id, ready) in &readies { - orch_datastream.emit_bootstrap( - dashboard, - config.run_id, - config.node_id, + bootstrap( + orch_datastream, "node_runtime_ready", "ready", json!({"endpoint":&ready.endpoint,"node_actor":ready.node_actor,"node_id":node_id,"stage_index":ready.stage_index}), @@ -2582,10 +2548,8 @@ fn start_and_provision_workers( .complete_bootstrap() .map_err(|e| format!("complete provider bootstrap after runtime-ready: {e}"))?; - orch_datastream.emit_bootstrap( - dashboard, - config.run_id, - config.node_id, + bootstrap( + orch_datastream, "stage_provision", "started", stage_provision_detail(config, pipeline_plan), @@ -2597,10 +2561,8 @@ fn start_and_provision_workers( provision_stage(stack, ready.node_actor, config)?; } - orch_datastream.emit_bootstrap( - dashboard, - config.run_id, - config.node_id, + bootstrap( + orch_datastream, "weights_loaded", "started", json!({"model_id":&config.model_id,"expected":expected_node_ids.len()}), @@ -2649,19 +2611,15 @@ fn start_and_provision_workers( ) }; match weights_result { - Ok(()) => orch_datastream.emit_bootstrap( - dashboard, - config.run_id, - config.node_id, + Ok(()) => bootstrap( + orch_datastream, "weights_loaded", "ready", json!({"source":"actor_stage_ready","model_id":&config.model_id,"expected":expected_node_ids.len()}), ), Err(error) => { - orch_datastream.emit_bootstrap( - dashboard, - config.run_id, - config.node_id, + bootstrap( + orch_datastream, "weights_loaded", "failed", json!({"error":error}), @@ -3259,6 +3217,10 @@ fn send_pipeline_stage_provision( let datastream_route_owner = ctx.stack.route_owner(ready.datastream_publisher); let member_state = ctx.stack.member_state(ready.swim_node_id); let route_matches_ready = route_owner == Some(ready.swim_node_id); + let (dashboard, run_id, node_id) = (ctx.dashboard, ctx.run_id, ctx.node_id); + let bootstrap = |ds: &mut OrchDatastream, phase: &str, status: &str, detail: Value| { + ds.emit_bootstrap(dashboard, run_id, node_id, phase, status, detail); + }; ctx.orch_datastream.emit_bootstrap_to_channel( ctx.dashboard, MVP_STAGE_ROUTE, @@ -3294,10 +3256,8 @@ fn send_pipeline_stage_provision( route_matches_ready, "heartbeat_missed", ); - ctx.orch_datastream.emit_bootstrap( - ctx.dashboard, - ctx.run_id, - ctx.node_id, + bootstrap( + ctx.orch_datastream, "stage_provision_wait", "failed", json!({ @@ -3324,10 +3284,8 @@ fn send_pipeline_stage_provision( now, ); if !should_send { - ctx.orch_datastream.emit_bootstrap( - ctx.dashboard, - ctx.run_id, - ctx.node_id, + bootstrap( + ctx.orch_datastream, "stage_provision_wait", "observed", json!({ @@ -3368,10 +3326,8 @@ fn send_pipeline_stage_provision( *count }; ctx.stage_last_sends.insert(stage.stage_index, now); - ctx.orch_datastream.emit_bootstrap( - ctx.dashboard, - ctx.run_id, - ctx.node_id, + bootstrap( + ctx.orch_datastream, "stage_provision_send", "sent", json!({ @@ -3385,10 +3341,8 @@ fn send_pipeline_stage_provision( }), ); if stage_send_count == 1 || stage_send_count % 15 == 0 { - ctx.orch_datastream.emit_bootstrap( - ctx.dashboard, - ctx.run_id, - ctx.node_id, + bootstrap( + ctx.orch_datastream, "stage_provision_wait", "observed", json!({ @@ -4501,13 +4455,17 @@ impl PipelinePromptRuntime { run_id: u64, node_id: u64, ) -> Result<(), String> { + let emit_prompt_evt = + |ds: &mut OrchDatastream, request_id: u64, phase: &str, status: &str, detail: Value| { + ds.emit_prompt( + dashboard, run_id, node_id, request_id, phase, status, detail, + ) + }; 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, + emit_prompt_evt( + orch_datastream, request_id, "pipeline_prompt_busy", "failed", @@ -4530,10 +4488,8 @@ impl PipelinePromptRuntime { self.pending_encode = Some(PendingEncode { request_id }); self.started_at = Some(Instant::now()); self.note_progress(); - orch_datastream.emit_prompt( - dashboard, - run_id, - node_id, + emit_prompt_evt( + orch_datastream, request_id, "pipeline_tokenizer_encode", "started", @@ -4560,6 +4516,12 @@ impl PipelinePromptRuntime { run_id: u64, node_id: u64, ) { + let emit_prompt_evt = + |ds: &mut OrchDatastream, request_id: u64, phase: &str, status: &str, detail: Value| { + ds.emit_prompt( + dashboard, run_id, node_id, request_id, phase, status, detail, + ) + }; let Some(active) = self.active.as_ref() else { return; }; @@ -4574,10 +4536,8 @@ impl PipelinePromptRuntime { .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, + emit_prompt_evt( + orch_datastream, request_id, "pipeline_prompt_wait", "waiting", @@ -4640,6 +4600,12 @@ impl PipelinePromptRuntime { run_id: u64, node_id: u64, ) -> Result<(), String> { + let emit_prompt_evt = + |ds: &mut OrchDatastream, request_id: u64, phase: &str, status: &str, detail: Value| { + ds.emit_prompt( + dashboard, run_id, node_id, request_id, phase, status, detail, + ) + }; let Some(pending) = self.pending_encode.take() else { return Ok(()); }; @@ -4653,20 +4619,16 @@ impl PipelinePromptRuntime { if active.request.request_id != request_id { return Ok(()); } - orch_datastream.emit_prompt( - dashboard, - run_id, - node_id, + emit_prompt_evt( + orch_datastream, request_id, "pipeline_tokenizer_encode", "ready", json!({"node_actor":self.tokenizer_encode_actor,"reply_to":self.tokenizer_reply_to,"tokens":tokens.len()}), ); let sequence = self.next_sequence; - orch_datastream.emit_prompt( - dashboard, - run_id, - node_id, + emit_prompt_evt( + orch_datastream, request_id, "pipeline_token_in", "started", @@ -4674,10 +4636,8 @@ impl PipelinePromptRuntime { ); self.send_token_in(sequence, &tokens, true)?; self.note_progress(); - orch_datastream.emit_prompt( - dashboard, - run_id, - node_id, + emit_prompt_evt( + orch_datastream, request_id, "pipeline_token_in", "ready", @@ -4697,6 +4657,12 @@ impl PipelinePromptRuntime { run_id: u64, node_id: u64, ) -> Result<(), String> { + let emit_prompt_evt = + |ds: &mut OrchDatastream, request_id: u64, phase: &str, status: &str, detail: Value| { + ds.emit_prompt( + dashboard, run_id, node_id, request_id, phase, status, detail, + ) + }; let Some(pending) = self.pending_decode.take() else { return Ok(()); }; @@ -4711,10 +4677,8 @@ impl PipelinePromptRuntime { return Ok(()); } let events = active.events.clone(); - orch_datastream.emit_prompt( - dashboard, - run_id, - node_id, + emit_prompt_evt( + orch_datastream, request_id, "pipeline_tokenizer_decode", "ready", @@ -4732,10 +4696,8 @@ impl PipelinePromptRuntime { .unwrap_or(0); let final_text = self.final_text.clone(); let tokens_generated = self.generated_tokens.len() as u32; - orch_datastream.emit_prompt( - dashboard, - run_id, - node_id, + emit_prompt_evt( + orch_datastream, request_id, "prompt_complete", "ready", @@ -4760,10 +4722,8 @@ impl PipelinePromptRuntime { return Ok(()); } let sequence = self.next_sequence; - orch_datastream.emit_prompt( - dashboard, - run_id, - node_id, + emit_prompt_evt( + orch_datastream, request_id, "pipeline_token_in", "started", @@ -4771,10 +4731,8 @@ impl PipelinePromptRuntime { ); self.send_token_in(sequence, &[pending.token_id], false)?; self.note_progress(); - orch_datastream.emit_prompt( - dashboard, - run_id, - node_id, + emit_prompt_evt( + orch_datastream, request_id, "pipeline_token_in", "ready", @@ -4885,6 +4843,12 @@ impl PipelinePromptRuntime { run_id: u64, node_id: u64, ) -> Result<(), String> { + let emit_prompt_evt = + |ds: &mut OrchDatastream, request_id: u64, phase: &str, status: &str, detail: Value| { + ds.emit_prompt( + dashboard, run_id, node_id, request_id, phase, status, detail, + ) + }; if self.pending_decode.is_some() { return Ok(()); } @@ -4905,10 +4869,8 @@ impl PipelinePromptRuntime { continue; }; let request_id = active.request.request_id; - orch_datastream.emit_prompt( - dashboard, - run_id, - node_id, + emit_prompt_evt( + orch_datastream, request_id, "pipeline_token_out", "observed", @@ -4916,10 +4878,8 @@ impl PipelinePromptRuntime { ); self.generated_tokens.push(record.token_id); let reached_limit = self.generated_tokens.len() as u32 >= active.request.max_tokens; - orch_datastream.emit_prompt( - dashboard, - run_id, - node_id, + emit_prompt_evt( + orch_datastream, request_id, "pipeline_tokenizer_decode", "started", @@ -5009,6 +4969,12 @@ fn serve_prompts( provider, .. } = ctx; + let emit_prompt_evt = + |ds: &mut OrchDatastream, request_id: u64, phase: &str, status: &str, detail: Value| { + ds.emit_prompt( + dashboard, run_id, node_id, request_id, phase, status, detail, + ) + }; let mut pipeline_runtime = match pipeline_plan { Some(plan) => Some(PipelinePromptRuntime::new( driver, @@ -5065,10 +5031,8 @@ fn serve_prompts( { let request = work.request; let request_id = request.request_id; - orch_datastream.emit_prompt( - dashboard, - run_id, - node_id, + emit_prompt_evt( + orch_datastream, request_id, "prompt_work", "observed", @@ -5089,10 +5053,8 @@ fn serve_prompts( )?; continue; } - orch_datastream.emit_prompt( - dashboard, - run_id, - node_id, + emit_prompt_evt( + orch_datastream, request_id, "node_prompt_send", "started", @@ -5108,10 +5070,8 @@ fn serve_prompts( }, ) { Ok(()) => { - orch_datastream.emit_prompt( - dashboard, - run_id, - node_id, + emit_prompt_evt( + orch_datastream, request_id, "node_prompt_send", "ready", @@ -5123,10 +5083,8 @@ fn serve_prompts( }); } Err(error) => { - orch_datastream.emit_prompt( - dashboard, - run_id, - node_id, + emit_prompt_evt( + orch_datastream, request_id, "node_prompt_send", "failed", @@ -5140,10 +5098,8 @@ fn serve_prompts( while let Some(event) = prompt_events.try_recv() { let request_id = event.request_id(); let Some(current) = active.as_ref() else { - orch_datastream.emit_prompt( - dashboard, - run_id, - node_id, + emit_prompt_evt( + orch_datastream, request_id, "node_prompt_event", "dropped", @@ -5152,10 +5108,8 @@ fn serve_prompts( continue; }; if request_id != current.request.request_id { - orch_datastream.emit_prompt( - dashboard, - run_id, - node_id, + emit_prompt_evt( + orch_datastream, request_id, "node_prompt_event", "dropped", @@ -5175,10 +5129,8 @@ fn serve_prompts( } PromptEvent::TextDelta { .. } => None, }; - orch_datastream.emit_prompt( - dashboard, - run_id, - node_id, + emit_prompt_evt( + orch_datastream, request_id, "node_prompt_event", "observed", @@ -5186,10 +5138,8 @@ fn serve_prompts( ); let _ = current.events.send(event); if let Some((status, detail)) = completion { - orch_datastream.emit_prompt( - dashboard, - run_id, - node_id, + emit_prompt_evt( + orch_datastream, request_id, "prompt_complete", status, @@ -5437,7 +5387,7 @@ fn emit_swim_transitions( .map(|state| format!("{:?}", state)); let last_ack_age_ms = transition.last_ack_age.map(duration_ms_u64); let consecutive_timeouts = transition.consecutive_timeouts; - let recent_probe_targets = swim_recent_probe_targets(stack); + let recent_probe_targets = stack.swim_recent_probe_targets(); orch_datastream.emit_bootstrap_to_channel( dashboard, MVP_SWIM_MEMBERSHIP, @@ -5456,19 +5406,7 @@ fn emit_swim_transitions( "member_state":member_state.clone(), }), ); - orch_datastream.emit_record( - dashboard, - &MembershipTransition { - peer, - from: from.unwrap_or_default(), - to, - reason: transition.reason.to_owned(), - last_ack_age_ms, - consecutive_timeouts, - recent_probe_targets, - member_state, - }, - ); + orch_datastream.emit_record(dashboard, &stack.membership_transition(&transition)); } } @@ -5479,56 +5417,11 @@ fn emit_swim_probe_events( local_phase: &str, ) { for event in stack.drain_swim_probe_events() { - let record = swim_probe_event_record(stack, event, local_phase); + let record = stack.swim_probe_event_record(event, local_phase); orch_datastream.emit_record(dashboard, &record); } } -fn swim_probe_event_record( - stack: &DistributionRuntimeStack, - event: ObservedProbeEvent, - local_phase: &str, -) -> SwimProbeEvent { - let config = &stack.swim_config; - let budget_ms = event.budget_ms; - SwimProbeEvent { - event: event.event.to_owned(), - target: format!("{:?}", event.target), - sequence: event.sequence, - kind: event.kind.to_owned(), - rtt_ms: event.rtt_ms, - budget_ms, - budget_ticks: budget_ms, - last_ack_age_ms: event.last_ack_age.map(duration_ms_u64), - consecutive_timeouts: event.consecutive_timeouts, - recent_probe_targets: swim_recent_probe_targets(stack), - member_state: stack - .member_state(event.target) - .map(|state| format!("{:?}", state)), - local_phase: local_phase.to_owned(), - probe_interval_ms: duration_ms_u64(config.probe_interval), - probe_timeout_ms: duration_ms_u64(config.probe_timeout), - indirect_probes: u32::try_from(config.indirect_probes).unwrap_or(u32::MAX), - suspicion_timeout_ms: duration_ms_u64(config.suspicion_timeout), - dead_reprobe_interval_ms: duration_ms_u64(config.dead_reprobe_interval), - probe_mode: format!("{:?}", config.probe_mode), - lifeguard_enabled: config.lifeguard.is_some(), - } -} - -fn swim_recent_probe_targets(stack: &DistributionRuntimeStack) -> Vec { - stack - .swim_telemetry - .recent_targets() - .into_iter() - .map(|node_id| format!("{:?}", node_id)) - .collect() -} - -fn duration_ms_u64(duration: Duration) -> u64 { - u64::try_from(duration.as_millis()).unwrap_or(u64::MAX) -} - fn pump( driver: &mut IrohDriver, stack: &DistributionRuntimeStack, @@ -5700,10 +5593,7 @@ fn ensure_vastai_account_ssh_key(api_key: &str, public_key: &str) -> Result<(), if vastai_account_has_ssh_key(api_key, public_key)? { Ok(()) } else { - Err( - "VastAI SSH key registration did not make the selected key visible in vastai show ssh-keys" - .to_owned(), - ) + Err("VastAI SSH key registration did not make the selected key visible in vastai show ssh-keys".to_owned()) } } @@ -5719,8 +5609,7 @@ fn account_ssh_keys_output_contains_public_key(output: &str, public_key: &str) - fn vastai_cli_error(error: std::io::Error) -> String { if error.kind() == std::io::ErrorKind::NotFound { - "vastai CLI is required to verify/register MVP_VASTAI_SSH_IDENTITY; install with pip install vastai" - .to_owned() + "vastai CLI is required to verify/register MVP_VASTAI_SSH_IDENTITY; install with pip install vastai".to_owned() } else { format!("run vastai CLI: {error}") } diff --git a/crates/mvp-system/src/orchestration/distribution_stack.rs b/crates/mvp-system/src/orchestration/distribution_stack.rs index 8f05320..b4747db 100644 --- a/crates/mvp-system/src/orchestration/distribution_stack.rs +++ b/crates/mvp-system/src/orchestration/distribution_stack.rs @@ -9,7 +9,7 @@ use std::collections::HashMap; use std::sync::{Arc, Mutex, RwLock}; -use std::time::Instant; +use std::time::{Duration, Instant}; use swactor::actor::{ActorAddress, ActorInterface}; use swactor::config::RuntimeConfig; @@ -28,6 +28,8 @@ use distribution::swim::actor::{MembershipChanged, SwimActor, SwimIn}; use distribution::swim::member_list::MemberList; use distribution::swim::probe::SwimConfig; use distribution::swim::telemetry::{ObservedProbeEvent, ObservedTransition, SwimTelemetry}; +use distribution::telemetry::MembershipTransition; +use distribution::telemetry::SwimProbeEvent; use distribution::transport_bridge::{ Outbox, OutboxPeerDirectory, OutboxRouteBinder, RelayMirror, RouteView, RouteViewTransport, }; @@ -237,6 +239,70 @@ impl DistributionRuntimeStack { pub(crate) fn drain_swim_probe_events(&self) -> Vec { self.swim_telemetry.drain_probe_events() } + + pub(crate) fn swim_recent_probe_targets(&self) -> Vec { + self.swim_telemetry + .recent_targets() + .into_iter() + .map(|node_id| format!("{:?}", node_id)) + .collect() + } + + pub(crate) fn swim_probe_event_record( + &self, + event: ObservedProbeEvent, + local_phase: &str, + ) -> SwimProbeEvent { + let config = &self.swim_config; + let budget_ms = event.budget_ms; + SwimProbeEvent { + event: event.event.to_owned(), + target: format!("{:?}", event.target), + sequence: event.sequence, + kind: event.kind.to_owned(), + rtt_ms: event.rtt_ms, + budget_ms, + budget_ticks: budget_ms, + last_ack_age_ms: event.last_ack_age.map(duration_ms_u64), + consecutive_timeouts: event.consecutive_timeouts, + recent_probe_targets: self.swim_recent_probe_targets(), + member_state: self + .member_state(event.target) + .map(|state| format!("{:?}", state)), + local_phase: local_phase.to_owned(), + probe_interval_ms: duration_ms_u64(config.probe_interval), + probe_timeout_ms: duration_ms_u64(config.probe_timeout), + indirect_probes: u32::try_from(config.indirect_probes).unwrap_or(u32::MAX), + suspicion_timeout_ms: duration_ms_u64(config.suspicion_timeout), + dead_reprobe_interval_ms: duration_ms_u64(config.dead_reprobe_interval), + probe_mode: format!("{:?}", config.probe_mode), + lifeguard_enabled: config.lifeguard.is_some(), + } + } + pub(crate) fn membership_transition( + &self, + transition: &ObservedTransition, + ) -> MembershipTransition { + MembershipTransition { + peer: format!("{:?}", transition.peer), + from: transition + .from + .map(|state| format!("{:?}", state)) + .unwrap_or_default(), + to: format!("{:?}", transition.to), + reason: transition.reason.to_owned(), + last_ack_age_ms: transition.last_ack_age.map(duration_ms_u64), + consecutive_timeouts: transition.consecutive_timeouts, + recent_probe_targets: self.swim_recent_probe_targets(), + member_state: self + .member_state(transition.peer) + .map(|state| format!("{:?}", state)), + } + } +} + +pub(crate) fn duration_ms_u64(duration: Duration) -> u64 { + u64::try_from(duration.as_millis()).unwrap_or(u64::MAX) } struct MembershipFanout { diff --git a/crates/transport/src/json_codec.rs b/crates/transport/src/json_codec.rs index 577aa83..3b7ed0e 100644 --- a/crates/transport/src/json_codec.rs +++ b/crates/transport/src/json_codec.rs @@ -6,8 +6,8 @@ use std::marker::PhantomData; -use serde::Serialize; use serde::de::DeserializeOwned; +use serde::Serialize; use swactor::Error; use crate::codec::Codec;