From 38383e1113edc4b976595b135b7f429863a97b24 Mon Sep 17 00:00:00 2001 From: Zachery Aaron Shores-Chmielewski Date: Thu, 30 Jul 2026 12:18:02 +0400 Subject: [PATCH] refactor: fold driver pumps into `iroh-driver` and codecs into `transport` --- .gitignore | 1 + crates/datastream/src/publisher_actor.rs | 27 +- crates/distribution/src/messages.rs | 53 +--- .../src}/driver_pumps.rs | 33 ++- .../src}/endpoint_advertisement.rs | 24 +- crates/iroh-driver/src/lib.rs | 5 + .../tests/endpoint_advertisement.rs} | 4 +- crates/mvp-system/src/chat/runtime.rs | 2 +- crates/mvp-system/src/codecs.rs | 13 + crates/mvp-system/src/lib.rs | 4 +- crates/mvp-system/src/node/actor.rs | 2 +- .../src/node/worker_node_runtime.rs | 264 ++++-------------- crates/mvp-system/src/orchestration/actor.rs | 2 +- crates/mvp-system/src/orchestration/app.rs | 4 +- crates/mvp-system/src/prompt/rpc.rs | 2 +- crates/mvp-system/src/tests/mod.rs | 1 - crates/mvp-system/src/transport/mod.rs | 11 - .../transport => transport/src}/json_codec.rs | 15 +- crates/transport/src/lib.rs | 2 + 19 files changed, 156 insertions(+), 313 deletions(-) rename crates/{mvp-system/src/transport => iroh-driver/src}/driver_pumps.rs (71%) rename crates/{mvp-system/src/transport => iroh-driver/src}/endpoint_advertisement.rs (62%) rename crates/{mvp-system/src/tests/transport_guarantees.rs => iroh-driver/tests/endpoint_advertisement.rs} (89%) create mode 100644 crates/mvp-system/src/codecs.rs delete mode 100644 crates/mvp-system/src/transport/mod.rs rename crates/{mvp-system/src/transport => transport/src}/json_codec.rs (55%) diff --git a/.gitignore b/.gitignore index e25d7cc..2fdd790 100644 --- a/.gitignore +++ b/.gitignore @@ -5,6 +5,7 @@ corpus .loop/ .model-cache/ .deployment-notes/ +.omp/ # environment variables for deployment configurations .config/ \ No newline at end of file diff --git a/crates/datastream/src/publisher_actor.rs b/crates/datastream/src/publisher_actor.rs index 6dd7206..d31c540 100644 --- a/crates/datastream/src/publisher_actor.rs +++ b/crates/datastream/src/publisher_actor.rs @@ -5,17 +5,13 @@ //! injected by the runtime crate so the datastream core does not depend on the //! concrete QUIC writer implementation. -use std::marker::PhantomData; use std::sync::Arc; use iroh::EndpointAddr; -use serde::Serialize; -use serde::de::DeserializeOwned; use serde::{Deserialize, Serialize as DeriveSerialize}; -use swactor::Error; use swactor::actor::ActorInterface; use swactor::runtime::Ctx; -use swactor_transport::{Codec, CodecRegistry, NetworkMessage}; +use swactor_transport::{CodecRegistry, JsonCodec, NetworkMessage}; use crate::{DatastreamEndpoint, DatastreamSubscription, SubscriptionRequest}; @@ -80,22 +76,7 @@ impl ActorInterface for DatastreamPublisherActor { /// Register JSON encoding for remote datastream publisher messages. pub fn register_datastream_publisher_codec(registry: &mut CodecRegistry) { - registry.register::>(JsonCodec( - PhantomData, - )); -} - -struct JsonCodec(PhantomData); - -impl Codec for JsonCodec -where - M: Serialize + DeserializeOwned + Clone + Send + Sync + 'static, -{ - fn encode(&self, msg: &M) -> Result, Error> { - serde_json::to_vec(msg).map_err(|e| Error::from(format!("encode: {e}"))) - } - - fn decode(&self, bytes: &[u8]) -> Result { - serde_json::from_slice(bytes).map_err(|e| Error::from(format!("decode: {e}"))) - } + registry.register::( + JsonCodec::::default(), + ); } diff --git a/crates/distribution/src/messages.rs b/crates/distribution/src/messages.rs index ec11bb7..babfbc9 100644 --- a/crates/distribution/src/messages.rs +++ b/crates/distribution/src/messages.rs @@ -1,9 +1,9 @@ //! Protocol messages for SWIM membership and the standalone gossip frames -//! (registry, node metadata, directory), plus their JSON codec. +//! (registry, node metadata, directory), plus codec registration. use serde::{Deserialize, Serialize}; use swactor::Error; -use swactor_transport::{Codec, CodecRegistry, NetworkMessage}; +use swactor_transport::{CodecRegistry, JsonCodec, NetworkMessage}; use crate::node_metadata::NodeMetadataEntry; use crate::registry::RegistryEntry; @@ -164,52 +164,21 @@ impl NetworkMessage for DirectoryGossip { } } -// ─── JSON Codec ───────────────────────────────────────────────────────────── - -/// JSON codec for distribution protocol messages. -/// -/// Using JSON for simplicity and debuggability. Can be swapped for -/// bincode/msgpack in production via the Codec trait. -pub struct JsonCodec; - -macro_rules! impl_json_codec { - ($ty:ty) => { - impl Codec<$ty> for JsonCodec { - fn encode(&self, msg: &$ty) -> Result, Error> { - serde_json::to_vec(msg).map_err(|e| Error::from(format!("encode: {e}"))) - } - fn decode(&self, bytes: &[u8]) -> Result<$ty, Error> { - serde_json::from_slice(bytes).map_err(|e| Error::from(format!("decode: {e}"))) - } - } - }; -} - -impl_json_codec!(Ping); -impl_json_codec!(Ack); -impl_json_codec!(PingReq); -impl_json_codec!(IndirectAck); -impl_json_codec!(JoinRequest); -impl_json_codec!(JoinResponse); -impl_json_codec!(RegistryGossip); -impl_json_codec!(MetadataGossip); -impl_json_codec!(DirectoryGossip); - /// Build a `CodecRegistry` with all distribution protocol messages registered. pub fn distribution_codec_registry() -> CodecRegistry { let mut cr = CodecRegistry::new(); - cr.register::(JsonCodec); - cr.register::(JsonCodec); - cr.register::(JsonCodec); + cr.register::(JsonCodec::::default()); + cr.register::(JsonCodec::::default()); + cr.register::(JsonCodec::::default()); // §6.1 / §14.5: `IndirectAck` is folded into the shared registry so all six // SWIM message types decode through one uniform path; concrete drivers no // longer need to hand-dispatch it by tag. - cr.register::(JsonCodec); - cr.register::(JsonCodec); - cr.register::(JsonCodec); - cr.register::(JsonCodec); - cr.register::(JsonCodec); - cr.register::(JsonCodec); + cr.register::(JsonCodec::::default()); + cr.register::(JsonCodec::::default()); + cr.register::(JsonCodec::::default()); + cr.register::(JsonCodec::::default()); + cr.register::(JsonCodec::::default()); + cr.register::(JsonCodec::::default()); cr } diff --git a/crates/mvp-system/src/transport/driver_pumps.rs b/crates/iroh-driver/src/driver_pumps.rs similarity index 71% rename from crates/mvp-system/src/transport/driver_pumps.rs rename to crates/iroh-driver/src/driver_pumps.rs index bf4cb8a..90ee322 100644 --- a/crates/mvp-system/src/transport/driver_pumps.rs +++ b/crates/iroh-driver/src/driver_pumps.rs @@ -1,25 +1,31 @@ -#![allow(dead_code)] +//! Driver edge/ring/stream bookkeeping. +//! +//! The driver is the node's swactor-to-iroh boundary. It tracks which edges +//! have an established send or recv pump, maps inbound uni-streams to their +//! recv rings, and emits driver lifecycle events (edge ready, stream fault, +//! pump stopped). This module is the pure state machine; the iroh stream pumps +//! themselves live in [`crate::edge_transport`]. use std::collections::BTreeMap; #[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] -pub(crate) struct EdgeId(pub(crate) u64); +pub struct EdgeId(pub u64); #[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] -pub(crate) struct RingId(pub(crate) u64); +pub struct RingId(pub u64); #[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] -pub(crate) struct StreamId(pub(crate) u64); +pub struct StreamId(pub u64); #[derive(Clone, Debug, PartialEq, Eq)] -pub(crate) enum DriverEventOut { +pub enum DriverEventOut { DriverEdgeReady { edge_id: EdgeId }, StreamFault { edge_id: EdgeId }, PumpStopped { edge_id: EdgeId, ring_id: RingId }, } #[derive(Debug)] -pub(crate) struct Driver { +pub struct Driver { sends: BTreeMap, recv_specs: BTreeMap, pending_streams: BTreeMap, @@ -28,7 +34,7 @@ pub(crate) struct Driver { } impl Driver { - pub(crate) fn new() -> Self { + pub fn new() -> Self { Self { sends: BTreeMap::new(), recv_specs: BTreeMap::new(), @@ -38,20 +44,20 @@ impl Driver { } } - pub(crate) fn establish_send(&mut self, edge_id: EdgeId, ring_id: RingId) { + pub fn establish_send(&mut self, edge_id: EdgeId, ring_id: RingId) { self.sends.insert(edge_id, ring_id); self.events .push(DriverEventOut::DriverEdgeReady { edge_id }); } - pub(crate) fn establish_recv(&mut self, edge_id: EdgeId, ring_id: RingId) { + pub fn establish_recv(&mut self, edge_id: EdgeId, ring_id: RingId) { self.recv_specs.insert(edge_id, ring_id); if let Some(stream_id) = self.pending_streams.remove(&edge_id) { self.spawn_recv(edge_id, stream_id); } } - pub(crate) fn incoming_uni_stream(&mut self, edge_id: EdgeId, stream_id: StreamId) { + pub fn incoming_uni_stream(&mut self, edge_id: EdgeId, stream_id: StreamId) { if self.recv_specs.contains_key(&edge_id) { self.spawn_recv(edge_id, stream_id); } else { @@ -69,11 +75,11 @@ impl Driver { .push(DriverEventOut::DriverEdgeReady { edge_id }); } - pub(crate) fn read_error(&mut self, edge_id: EdgeId) { + pub fn read_error(&mut self, edge_id: EdgeId) { self.events.push(DriverEventOut::StreamFault { edge_id }); } - pub(crate) fn stop_edge(&mut self, edge_id: EdgeId) { + pub fn stop_edge(&mut self, edge_id: EdgeId) { let ring_id = self .sends .get(&edge_id) @@ -89,7 +95,8 @@ impl Driver { self.events .push(DriverEventOut::PumpStopped { edge_id, ring_id }); } - pub(crate) fn events(&self) -> &[DriverEventOut] { + + pub fn events(&self) -> &[DriverEventOut] { &self.events } } diff --git a/crates/mvp-system/src/transport/endpoint_advertisement.rs b/crates/iroh-driver/src/endpoint_advertisement.rs similarity index 62% rename from crates/mvp-system/src/transport/endpoint_advertisement.rs rename to crates/iroh-driver/src/endpoint_advertisement.rs index cf96a22..5165428 100644 --- a/crates/mvp-system/src/transport/endpoint_advertisement.rs +++ b/crates/iroh-driver/src/endpoint_advertisement.rs @@ -1,18 +1,29 @@ +//! Endpoint advertisement masking. +//! +//! Controls how much of an iroh [`EndpointAddr`] a node advertises to peers. +//! `relay-only` strips direct IP/socket addresses so peers can only reach the +//! node via its relay URL — useful for NAT-egress-only or hidden nodes. + use std::fmt; use iroh::EndpointAddr; -pub(crate) const MVP_IROH_ENDPOINT_ADDR_MASK_ENV: &str = "MVP_IROH_ENDPOINT_ADDR_MASK"; +/// Environment variable selecting the advertised endpoint address mask. +/// +/// Recognized values: `full` (default) and `relay-only`. +pub const MVP_IROH_ENDPOINT_ADDR_MASK_ENV: &str = "MVP_IROH_ENDPOINT_ADDR_MASK"; #[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] -pub(crate) enum EndpointAddrMask { +pub enum EndpointAddrMask { + /// Advertise the full endpoint address: relays and direct addresses. #[default] Full, + /// Advertise relay URLs only, omitting direct socket addresses. RelayOnly, } impl EndpointAddrMask { - pub(crate) fn parse(value: &str) -> Result { + pub fn parse(value: &str) -> Result { match value.trim().to_ascii_lowercase().as_str() { "" | "full" | "none" => Ok(Self::Full), "relay-only" | "relay_only" | "relay" => Ok(Self::RelayOnly), @@ -22,14 +33,14 @@ impl EndpointAddrMask { } } - pub(crate) fn as_str(self) -> &'static str { + pub fn as_str(self) -> &'static str { match self { Self::Full => "full", Self::RelayOnly => "relay-only", } } - pub(crate) fn requires_relay(self) -> bool { + pub fn requires_relay(self) -> bool { matches!(self, Self::RelayOnly) } } @@ -40,7 +51,8 @@ impl fmt::Display for EndpointAddrMask { } } -pub(crate) fn advertised_endpoint( +/// Apply `mask` to `endpoint`, returning the address to advertise to peers. +pub fn advertised_endpoint( endpoint: EndpointAddr, mask: EndpointAddrMask, ) -> Result { diff --git a/crates/iroh-driver/src/lib.rs b/crates/iroh-driver/src/lib.rs index 4990671..df44fe5 100644 --- a/crates/iroh-driver/src/lib.rs +++ b/crates/iroh-driver/src/lib.rs @@ -5,13 +5,18 @@ //! wire message definitions. pub mod datastream_transport; +pub mod driver_pumps; pub mod edge_transport; +pub mod endpoint_advertisement; pub mod iroh_driver; 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/tests/transport_guarantees.rs b/crates/iroh-driver/tests/endpoint_advertisement.rs similarity index 89% rename from crates/mvp-system/src/tests/transport_guarantees.rs rename to crates/iroh-driver/tests/endpoint_advertisement.rs index e31e101..b79ca06 100644 --- a/crates/mvp-system/src/tests/transport_guarantees.rs +++ b/crates/iroh-driver/tests/endpoint_advertisement.rs @@ -1,8 +1,8 @@ -//! Behavior guarantees for the `transport` module. +//! Behavior guarantees for endpoint advertisement masking. use std::net::{Ipv4Addr, SocketAddr}; -use mvp_system::transport::endpoint_advertisement::{EndpointAddrMask, advertised_endpoint}; +use iroh_driver::{EndpointAddrMask, advertised_endpoint}; #[test] fn relay_only_mask_preserves_relay_urls_and_removes_direct_addresses() { diff --git a/crates/mvp-system/src/chat/runtime.rs b/crates/mvp-system/src/chat/runtime.rs index 0de3845..fdbb120 100644 --- a/crates/mvp-system/src/chat/runtime.rs +++ b/crates/mvp-system/src/chat/runtime.rs @@ -30,7 +30,7 @@ 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 crate::transport::endpoint_advertisement::EndpointAddrMask; +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, diff --git a/crates/mvp-system/src/codecs.rs b/crates/mvp-system/src/codecs.rs new file mode 100644 index 0000000..2e4f0f5 --- /dev/null +++ b/crates/mvp-system/src/codecs.rs @@ -0,0 +1,13 @@ +//! Registers wire codecs for every MVP actor message type. +//! +//! The runtime's [`CodecRegistry`](swactor_transport::CodecRegistry) needs an +//! encoder/decoder entry for each inter-node message type. This is the single +//! aggregator that wires up the node, orchestrator, prompt, and datastream +//! publisher codecs. + +pub(crate) fn register_mvp_actor_codecs(registry: &mut swactor_transport::CodecRegistry) { + crate::node_actor::register_codecs(registry); + crate::orchestration::actor::register_codecs(registry); + datastream::register_datastream_publisher_codec(registry); + crate::prompt::rpc::register_codecs(registry); +} diff --git a/crates/mvp-system/src/lib.rs b/crates/mvp-system/src/lib.rs index 0dfef09..49fd11a 100644 --- a/crates/mvp-system/src/lib.rs +++ b/crates/mvp-system/src/lib.rs @@ -26,8 +26,6 @@ pub fn run_worker_node_from_env() -> std::process::ExitCode { node::worker_node_runtime::run_from_env() } -#[path = "transport/driver_pumps.rs"] -mod driver_pumps; #[path = "staging/gguf_common.rs"] mod gguf_common; #[path = "staging/gguf_shard.rs"] @@ -50,7 +48,7 @@ mod observability; mod orchestration; mod prompt; mod staging; -mod transport; +mod codecs; #[cfg(test)] mod tests; diff --git a/crates/mvp-system/src/node/actor.rs b/crates/mvp-system/src/node/actor.rs index 3e64b43..8a1de52 100644 --- a/crates/mvp-system/src/node/actor.rs +++ b/crates/mvp-system/src/node/actor.rs @@ -9,7 +9,7 @@ use crate::run_plan; use crate::staging as stage; use crate::orchestration::actor::OrchestratorMsg; -use crate::transport::json_codec::JsonCodec; +use swactor_transport::JsonCodec; #[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] pub(crate) enum StageEdgeKindWire { diff --git a/crates/mvp-system/src/node/worker_node_runtime.rs b/crates/mvp-system/src/node/worker_node_runtime.rs index 1f57773..5365832 100644 --- a/crates/mvp-system/src/node/worker_node_runtime.rs +++ b/crates/mvp-system/src/node/worker_node_runtime.rs @@ -22,7 +22,7 @@ use datastream::{ Lifetime, NodeId, Record, StreamDescriptor, StreamId, StreamOrigin, }; -use crate::driver_pumps as driver_model; +use iroh_driver::driver_pumps as driver_model; use crate::gguf_shard::{StageShardPlan, materialize_stage_shard_http, validate_stage_shard_cache}; use crate::node_actor::{ NodeAgentActor, NodeAgentMsg, NodeAgentReport, StageCommandWire, StageInboundEdgeWire, @@ -34,10 +34,10 @@ use crate::orchestration::provider_adapters::relay::relay_runtime_config_from_en use crate::prompt::rpc::{PromptEvent, TokenizerEvent}; use crate::run_plan::{GgufSource, TokenizerSource}; use crate::staging::control as stage; -use crate::transport::endpoint_advertisement::{ +use iroh_driver::{ EndpointAddrMask, MVP_IROH_ENDPOINT_ADDR_MASK_ENV, advertised_endpoint, }; -use crate::transport::register_mvp_actor_codecs; +use crate::codecs::register_mvp_actor_codecs; use data_plane::arena; use data_plane::edge_lifecycle as edge; use data_plane::ingress; @@ -1630,50 +1630,32 @@ fn run_stage_shard_fetcher() -> Result<(), String> { fn run() -> Result<(), String> { let config = DeploymentConfig::from_env()?; - emit_stdio_node_event( - &config, - NODE_BOOTSTRAP_CHANNEL, - "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"), - }), - )?; - emit_stdio_node_event( - &config, - NODE_BOOTSTRAP_CHANNEL, - "process", - "started", - json!({"binary":"mvp-worker-node","pid":std::process::id()}), - )?; + let boot = |phase: &str, status: &str, detail: Value| { + emit_stdio_node_event(&config, NODE_BOOTSTRAP_CHANNEL, phase, status, detail) + }; + 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()}))?; let tokio = match tokio::runtime::Runtime::new() { Ok(runtime) => { - emit_stdio_node_event( - &config, - NODE_BOOTSTRAP_CHANNEL, - "tokio_runtime", - "ready", - json!({"runtime":"tokio"}), - )?; + boot("tokio_runtime", "ready", json!({"runtime":"tokio"}))?; runtime } Err(error) => { - emit_stdio_node_event( - &config, - NODE_BOOTSTRAP_CHANNEL, - "tokio_runtime", - "failed", - json!({"error":error.to_string()}), - )?; + boot("tokio_runtime", "failed", json!({"error":error.to_string()}))?; return Err(format!("tokio runtime: {error}")); } }; @@ -1689,42 +1671,18 @@ fn run() -> Result<(), String> { ) { Ok(driver) => driver, Err(error) => { - emit_stdio_node_event( - &config, - NODE_BOOTSTRAP_CHANNEL, - "iroh_driver", - "failed", - json!({"error":error.to_string()}), - )?; + boot("iroh_driver", "failed", json!({"error":error.to_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()}), - )?; + 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)); - emit_stdio_node_event( - &config, - NODE_BOOTSTRAP_CHANNEL, - "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 { - emit_stdio_node_event( - &config, - NODE_BOOTSTRAP_CHANNEL, - "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( @@ -1735,20 +1693,8 @@ fn run() -> Result<(), String> { datastream::wire::register_datastream_codec(registry); }, ); - emit_stdio_node_event( - &config, - NODE_BOOTSTRAP_CHANNEL, - "distribution_stack", - "ready", - json!({"actors":"initialized","route_view":"initialized","swim":"initialized","outbox":"initialized"}), - )?; - emit_stdio_node_event( - &config, - NODE_BOOTSTRAP_CHANNEL, - "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(), @@ -1757,13 +1703,7 @@ fn run() -> Result<(), String> { stack.relay_mirror.clone(), stack.route_view.clone(), ); - emit_stdio_node_event( - &config, - NODE_BOOTSTRAP_CHANNEL, - "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), @@ -1771,26 +1711,14 @@ fn run() -> Result<(), String> { base_alignment: config.arena_alignment, }) { Ok(manager) => { - emit_stdio_node_event( - &config, - NODE_BOOTSTRAP_CHANNEL, - "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) => { - emit_stdio_node_event( - &config, - NODE_BOOTSTRAP_CHANNEL, - "arena_manager", - "failed", - json!({"error":format!("{error:?}")}), - )?; + boot("arena_manager", "failed", json!({"error":format!("{error:?}")}))?; return Err(format!("boot arena manager: {error:?}")); } }; @@ -1917,23 +1845,11 @@ fn run() -> Result<(), String> { let reports = match stack.runtime.new_inbox::() { Ok(inbox) => { - emit_stdio_node_event( - &config, - NODE_BOOTSTRAP_CHANNEL, - "node_report_inbox", - "ready", - json!({"actor":inbox.addr()}), - )?; + boot("node_report_inbox", "ready", json!({"actor":inbox.addr()}))?; inbox } Err(error) => { - emit_stdio_node_event( - &config, - NODE_BOOTSTRAP_CHANNEL, - "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}")); } }; @@ -1941,13 +1857,7 @@ fn run() -> Result<(), String> { "MVP_ORCHESTRATOR_ACTOR is required for runtime readiness signaling".to_owned() })?; let orchestrator_source = "env"; - emit_stdio_node_event( - &config, - NODE_BOOTSTRAP_CHANNEL, - "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, @@ -1955,59 +1865,29 @@ fn run() -> Result<(), String> { ); let node_actor = match stack.runtime.spawn(node_agent) { Ok(actor) => { - emit_stdio_node_event( - &config, - NODE_BOOTSTRAP_CHANNEL, - "node_agent", - "ready", - json!({"node_actor":actor,"source":"generated"}), - )?; + boot("node_agent", "ready", json!({"node_actor":actor,"source":"generated"}))?; actor } Err(error) => { - emit_stdio_node_event( - &config, - NODE_BOOTSTRAP_CHANNEL, - "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)); - emit_stdio_node_event( - &config, - NODE_BOOTSTRAP_CHANNEL, - "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}))?; - emit_stdio_node_event( - &config, - NODE_WORKER_CHANNEL, - "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) => { - emit_stdio_node_event( - &config, - NODE_WORKER_CHANNEL, - "worker_process", - "failed", - json!({"error":error}), - )?; + worker_evt("worker_process", "failed", json!({"error":error}))?; return Err(error); } }; @@ -2019,30 +1899,12 @@ fn run() -> Result<(), String> { sampler_health_context, vec![std::process::id(), worker.pid()], ); - emit_stdio_node_event( - &config, - NODE_WORKER_CHANNEL, - "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(()) => emit_stdio_node_event( - &config, - NODE_WORKER_CHANNEL, - "worker_initialize", - "ready", - json!({"worker_event_type":"WorkerReady"}), - )?, + Ok(()) => worker_evt("worker_initialize", "ready", json!({"worker_event_type":"WorkerReady"}))?, Err(error) => { - emit_stdio_node_event( - &config, - NODE_WORKER_CHANNEL, - "worker_initialize", - "failed", - json!({"error":error}), - )?; + worker_evt("worker_initialize", "failed", json!({"error":error}))?; return Err(error); } } @@ -2063,19 +1925,13 @@ fn run() -> Result<(), String> { "logical_node_id": config.logical_node_id, "stage_index": config.stage_index, }); - emit_stdio_node_event( - &config, - NODE_BOOTSTRAP_CHANNEL, - "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( diff --git a/crates/mvp-system/src/orchestration/actor.rs b/crates/mvp-system/src/orchestration/actor.rs index 3983fc8..c8be606 100644 --- a/crates/mvp-system/src/orchestration/actor.rs +++ b/crates/mvp-system/src/orchestration/actor.rs @@ -6,7 +6,7 @@ use swactor_transport::{CodecRegistry, NetworkMessage}; use crate::run_fsm as core; -use crate::transport::json_codec::JsonCodec; +use swactor_transport::JsonCodec; #[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] pub(crate) struct StageRefWire { diff --git a/crates/mvp-system/src/orchestration/app.rs b/crates/mvp-system/src/orchestration/app.rs index ab767db..3f15e0d 100644 --- a/crates/mvp-system/src/orchestration/app.rs +++ b/crates/mvp-system/src/orchestration/app.rs @@ -20,7 +20,7 @@ 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::transport::register_mvp_actor_codecs; +use crate::codecs::register_mvp_actor_codecs; const PROVIDER_START_MAX_ATTEMPTS: usize = 4; use crate::gguf_shard::{StageShardPlan, plan_stage_shard}; @@ -48,7 +48,7 @@ use crate::provisioning::{ }; use crate::run_fsm::{RunConfig, RunId}; use crate::run_plan::{self, GgufSource, TokenizerSource}; -use crate::transport::endpoint_advertisement::{ +use iroh_driver::{ EndpointAddrMask, MVP_IROH_ENDPOINT_ADDR_MASK_ENV, advertised_endpoint, }; use data_plane::object_record as ingress; diff --git a/crates/mvp-system/src/prompt/rpc.rs b/crates/mvp-system/src/prompt/rpc.rs index a17afdb..6af1aed 100644 --- a/crates/mvp-system/src/prompt/rpc.rs +++ b/crates/mvp-system/src/prompt/rpc.rs @@ -3,7 +3,7 @@ use std::io::{BufRead, Write}; use serde::{Deserialize, Serialize}; use swactor_transport::{CodecRegistry, NetworkMessage}; -use crate::transport::json_codec::JsonCodec; +use swactor_transport::JsonCodec; #[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] pub(crate) struct SubmitPrompt { diff --git a/crates/mvp-system/src/tests/mod.rs b/crates/mvp-system/src/tests/mod.rs index 33dff01..e757758 100644 --- a/crates/mvp-system/src/tests/mod.rs +++ b/crates/mvp-system/src/tests/mod.rs @@ -5,4 +5,3 @@ mod observability_guarantees; mod orchestration_guarantees; mod prompt_guarantees; mod staging_guarantees; -mod transport_guarantees; diff --git a/crates/mvp-system/src/transport/mod.rs b/crates/mvp-system/src/transport/mod.rs deleted file mode 100644 index acb52cd..0000000 --- a/crates/mvp-system/src/transport/mod.rs +++ /dev/null @@ -1,11 +0,0 @@ -//! MVP edge transport public surface. - -pub(crate) fn register_mvp_actor_codecs(registry: &mut swactor_transport::CodecRegistry) { - crate::node_actor::register_codecs(registry); - crate::orchestration::actor::register_codecs(registry); - datastream::register_datastream_publisher_codec(registry); - crate::prompt::rpc::register_codecs(registry); -} - -pub(crate) mod endpoint_advertisement; -pub(crate) mod json_codec; diff --git a/crates/mvp-system/src/transport/json_codec.rs b/crates/transport/src/json_codec.rs similarity index 55% rename from crates/mvp-system/src/transport/json_codec.rs rename to crates/transport/src/json_codec.rs index 7da1f23..577aa83 100644 --- a/crates/mvp-system/src/transport/json_codec.rs +++ b/crates/transport/src/json_codec.rs @@ -1,11 +1,22 @@ +//! Generic JSON codec for serde message types. +//! +//! [`JsonCodec`] is the out-of-the-box [`Codec`] for messages that are +//! [`Serialize`] + [`DeserializeOwned`]. Use a custom codec when a wire format +//! needs schema stability or compactness beyond JSON. + use std::marker::PhantomData; use serde::Serialize; use serde::de::DeserializeOwned; use swactor::Error; -use swactor_transport::Codec; -pub(crate) struct JsonCodec(PhantomData); +use crate::codec::Codec; + +/// JSON codec for message type `M`. +/// +/// Implements [`Codec`] via `serde_json`. Construct with +/// [`JsonCodec::default`] / `JsonCodec::::default()`. +pub struct JsonCodec(PhantomData); impl Default for JsonCodec { fn default() -> Self { diff --git a/crates/transport/src/lib.rs b/crates/transport/src/lib.rs index 810e6ad..20c37e0 100644 --- a/crates/transport/src/lib.rs +++ b/crates/transport/src/lib.rs @@ -8,9 +8,11 @@ pub mod codec; pub mod crypto; pub mod identity; +pub mod json_codec; pub mod transport; pub use codec::{ hex_decode, hex_encode, Codec, CodecRegistry, NetworkMessage, NodeId, WireEnvelope, }; +pub use json_codec::JsonCodec; pub use transport::{CodecRemoteSink, InMemoryTransport, Transport, TransportRouter};