Compare commits

...

2 commits

26 changed files with 2988 additions and 240 deletions

1
Cargo.lock generated
View file

@ -2477,6 +2477,7 @@ dependencies = [
"distribution", "distribution",
"iroh", "iroh",
"iroh-driver", "iroh-driver",
"iroh-relay",
"libc", "libc",
"parking_lot", "parking_lot",
"serde", "serde",

View file

@ -49,6 +49,7 @@ mod datastream_records {
fn distribution_emits_owned_channel_through_datastream_mux() { fn distribution_emits_owned_channel_through_datastream_mux() {
let stream = StreamId::new(NodeId::new("dist-node"), Lifetime(1)); let stream = StreamId::new(NodeId::new("dist-node"), Lifetime(1));
let mux = Mux::unbounded(stream); let mux = Mux::unbounded(stream);
mux.set_frame_timing_enabled(false);
let state = DistributionState { let state = DistributionState {
registry_size: 9, registry_size: 9,
..Default::default() ..Default::default()

View file

@ -237,8 +237,8 @@ pub struct IrohDriver {
peer_relay_urls: HashMap<NodeId, iroh::RelayUrl>, peer_relay_urls: HashMap<NodeId, iroh::RelayUrl>,
/// Real-time join status for each peer being joined. /// Real-time join status for each peer being joined.
join_statuses: Arc<Mutex<HashMap<NodeId, JoinStatus>>>, join_statuses: Arc<Mutex<HashMap<NodeId, JoinStatus>>>,
/// Relay URL exposed by the bound endpoint, if any. /// Relay URL configured or exposed by the bound endpoint, if any.
relay_url: Option<String>, relay_url: Option<iroh::RelayUrl>,
/// Actor-bridge wiring, installed via [`Self::enable_actor_bridge`]. When /// Actor-bridge wiring, installed via [`Self::enable_actor_bridge`]. When
/// present, the driver decodes inbound frames into actor messages /// present, the driver decodes inbound frames into actor messages
/// ([`Self::pump_inbound_to_actors`]) and writes actor-produced outbound /// ([`Self::pump_inbound_to_actors`]) and writes actor-produced outbound
@ -330,6 +330,11 @@ impl IrohDriver {
config: IrohDriverConfig, config: IrohDriverConfig,
) -> Result<Self, Box<dyn std::error::Error>> { ) -> Result<Self, Box<dyn std::error::Error>> {
let effective_relay_mode = config.relay_mode; let effective_relay_mode = config.relay_mode;
let configured_relay_url = match &effective_relay_mode {
RelayMode::Custom(relay_map) => relay_map.urls::<Vec<_>>().into_iter().next(),
_ => None,
};
let custom_relay = matches!(&effective_relay_mode, RelayMode::Custom(_));
// A custom relay is operator-controlled (typically `iroh-driver-relay` // A custom relay is operator-controlled (typically `iroh-driver-relay`
// on a VPS, serving QUIC Address Discovery with a self-signed cert). // on a VPS, serving QUIC Address Discovery with a self-signed cert).
@ -337,8 +342,6 @@ impl IrohDriver {
// that, address discovery fails and every connection stays // that, address discovery fails and every connection stays
// `conn_type=Relay`, which defeats hole-punching and makes a NAT'd peer // `conn_type=Relay`, which defeats hole-punching and makes a NAT'd peer
// (e.g. a locally-run orchestrator) reachable only over the relay. // (e.g. a locally-run orchestrator) reachable only over the relay.
let custom_relay = matches!(effective_relay_mode, RelayMode::Custom(_));
let endpoint = rt.block_on(async { let endpoint = rt.block_on(async {
let mut alpns = vec![ALPN.to_vec()]; let mut alpns = vec![ALPN.to_vec()];
alpns.extend(config.additional_alpns.iter().cloned()); alpns.extend(config.additional_alpns.iter().cloned());
@ -362,7 +365,8 @@ impl IrohDriver {
.addr() .addr()
.relay_urls() .relay_urls()
.next() .next()
.map(|url| url.to_string()); .cloned()
.or(configured_relay_url);
// The driver's signing identity matches the iroh endpoint: both use // The driver's signing identity matches the iroh endpoint: both use
// ed25519-dalek, so we reconstruct our Keypair from iroh's secret key. // ed25519-dalek, so we reconstruct our Keypair from iroh's secret key.
@ -458,16 +462,20 @@ impl IrohDriver {
self.keypair.node_id() self.keypair.node_id()
} }
/// The endpoint's full address (public key + direct socket addresses). /// The endpoint's current relay/direct advertised address.
/// ///
/// Constructs the address from the endpoint's public key and bound /// Starts from Iroh's live endpoint address, which includes the current
/// sockets. For sockets bound to `0.0.0.0`, emits one address per /// home relay when one is available, then merges normalized direct socket
/// addresses. For sockets bound to `0.0.0.0`, emits one address per
/// discovered LAN IP so that peers on the same network can connect /// discovered LAN IP so that peers on the same network can connect
/// directly. IPv6 unspecified is mapped to localhost. /// directly. IPv6 unspecified is mapped to localhost.
pub fn endpoint_addr(&self) -> EndpointAddr { pub fn endpoint_addr(&self) -> EndpointAddr {
let key = PublicKey::from_bytes(&self.keypair.node_id().0) let mut addr = self.endpoint.addr();
.expect("node_id is a valid public key"); if addr.relay_urls().next().is_none() {
let mut addr = EndpointAddr::new(key); if let Some(relay) = self.relay_url.clone() {
addr = addr.with_relay_url(relay);
}
}
for sa in self.direct_addresses() { for sa in self.direct_addresses() {
addr = addr.with_ip_addr(sa); addr = addr.with_ip_addr(sa);
} }
@ -622,7 +630,7 @@ impl IrohDriver {
.peer_relay_urls .peer_relay_urls
.get(&seed_node_id) .get(&seed_node_id)
.cloned() .cloned()
.or_else(|| self.endpoint.addr().relay_urls().next().cloned()) .or_else(|| self.home_relay_url())
{ {
seed_addr.clone().with_relay_url(relay) seed_addr.clone().with_relay_url(relay)
} else { } else {
@ -983,7 +991,7 @@ impl IrohDriver {
.and_then(|s| s.parse::<iroh::RelayUrl>().ok()) .and_then(|s| s.parse::<iroh::RelayUrl>().ok())
{ {
Some(r) Some(r)
} else if let Some(r) = self.endpoint.addr().relay_urls().next().cloned() { } else if let Some(r) = self.home_relay_url() {
Some(r) Some(r)
} else { } else {
None None
@ -1051,14 +1059,19 @@ impl IrohDriver {
} }
} }
/// Relay URL exposed by the bound endpoint, if any. /// Relay URL configured or exposed by the bound endpoint, if any.
pub fn relay_url(&self) -> Option<&str> { pub fn relay_url(&self) -> Option<&str> {
self.relay_url.as_deref() self.relay_url.as_ref().map(|url| url.as_str())
} }
/// The endpoint's home relay URL (from RelayMode::Custom), if connected. /// The endpoint's live or configured home relay URL, if any.
pub fn home_relay_url(&self) -> Option<iroh::RelayUrl> { pub fn home_relay_url(&self) -> Option<iroh::RelayUrl> {
self.endpoint.addr().relay_urls().next().cloned() self.endpoint
.addr()
.relay_urls()
.next()
.cloned()
.or_else(|| self.relay_url.clone())
} }
/// Async teardown for the unified driver loop, which runs on a tokio worker /// Async teardown for the unified driver loop, which runs on a tokio worker

View file

@ -28,6 +28,7 @@ async fn iroh_datastream_alpn_carries_endpoint_subscription() {
let endpoint = DatastreamEndpoint::with_capacity(stream.clone(), 8, 8); let endpoint = DatastreamEndpoint::with_capacity(stream.clone(), 8, 8);
let subscription = endpoint.subscribe_all("iroh"); let subscription = endpoint.subscribe_all("iroh");
let producer = endpoint.producer(); let producer = endpoint.producer();
producer.set_frame_timing_enabled(false);
producer.submit_text("runtime.log", "alpha"); producer.submit_text("runtime.log", "alpha");
producer.submit_text("runtime.log", "beta"); producer.submit_text("runtime.log", "beta");

View file

@ -11,7 +11,7 @@
mod common; mod common;
use std::sync::{Arc, Mutex}; use std::sync::{Arc, Mutex};
use std::time::Duration; use std::time::{Duration, Instant};
use common::iroh::*; use common::iroh::*;
use distribution::peer_auth::PeerAllowList; use distribution::peer_auth::PeerAllowList;
@ -38,6 +38,33 @@ fn iroh_driver_reports_listen_addr_and_no_routes() {
driver.shutdown(); driver.shutdown();
} }
#[test]
fn endpoint_addr_includes_home_relay() {
let (relay_url, _relay_guard) = spawn_test_relay();
let expected_relay_url = relay_url.to_string();
let mut node = make_driver_with_relay(relay_url.clone());
let start = Instant::now();
let (relay_advertised, observed_relay_url) = loop {
let endpoint = node.endpoint_addr();
let current_relay_url = endpoint.relay_urls().next().map(|url| url.to_string());
if current_relay_url.as_deref() == Some(expected_relay_url.as_str()) {
break (true, current_relay_url);
}
if start.elapsed() >= Duration::from_secs(5) {
break (false, current_relay_url);
}
pump_one(&mut node);
std::thread::sleep(Duration::from_millis(10));
};
node.shutdown();
assert!(
relay_advertised,
"advertised endpoint did not include home relay {expected_relay_url} within timeout; last relay URL: {observed_relay_url:?}"
);
}
// ─── Join integration tests ───────────────────────────────────────────── // ─── Join integration tests ─────────────────────────────────────────────
#[test] #[test]

View file

@ -0,0 +1,75 @@
# Actor Control Audit Ideas
**status**: early draft
Grounding from `crates/mvp-system`: the specs already give a useful audit line. `MVP_SYSTEM_SPEC.md` says the orchestrator is run authority, swactor owns the control plane, actors establish/observe/tear down components, and tensor bytes are explicitly not actor-mailbox traffic. `MVP_NODE_PROVISIONING_SPEC.md` also gives a key exception: provider I/O and SSH bootstrap are temporary pre-swactor paths; after convergence, swactor is the live control path.
Suggested somewhat-deterministic identification passes:
1. **Execution-boundary denylist scan** -- Yes, and the inverse, any code not called from an Actor::handle(...) needs inspection.
AST-scan for `std::thread::spawn`, `tokio::spawn`, `Handle::spawn`, `spawn_blocking`, `Command::new(...).spawn`, `Runtime::new`, and `block_on`. Anything not inside an actor, runtime bootstrap, hot-path byte pump, or pre-swactor bootstrap allowlist is a candidate.
2. **Process ownership audit**
Find every `std::process::Child`, `ChildStdin`, `ChildStdout`, `ChildStderr`, and `Command::new`. Require each long-lived child to have an actor owner, stop message, exit observation path, and teardown report; otherwise it is likely imperative supervision.
3. **Network listener audit**
Scan for `TcpListener`, `UnixListener`, `UnixStream`, `UnixDatagram`, `accept`, and per-connection threads/tasks. A listener is acceptable if it immediately decodes ingress into actor messages; if it owns request state or invokes domain operations directly, flag it.
4. **Channel-as-shadow-mailbox audit**
Scan for `std::sync::mpsc`, `tokio::sync::mpsc`, `oneshot`, `watch`, `broadcast`, and custom queues. Channels outside actor shells often mean a parallel control surface; classify each as actor ingress adapter, data hot-path helper, test harness, or suspect.
7. **Actor reachability taint analysis** -- Yes see my comments on 1
Treat `impl ActorInterface::handle` and actor constructors as roots, then build a call graph. Side-effectful functions reachable only from bins/tests/background threads but not actor roots become candidates for migration.
8. **Side-effect import layering rule**
Flag `std::process`, `std::net`, `tokio::net`, `std::fs`, Docker/VastAI/SSH clients, driver joins, and datastream emitters in modules that are supposed to be pure domain state machines. Pure cores should emit commands/events, not perform effects.
10. **Runtime creation inventory** -- If this happens at all, massive red flag.
Enumerate every `tokio::runtime::Runtime::new` and `swactor::runtime::Runtime::new`. Runtime creation should cluster at process/runtime-stack boundaries and tests; nested or ad-hoc runtimes usually indicate imperative islands.
11. **Post-handoff control-path check** -- All bootstrap monitoring should be owned by an actor, no exceptions.
Encode the provisioning spec as an audit rule: after `swactor` convergence/handoff, SSH/bootstrap/provider code may not remain the live node control path. Scan for SSH or bootstrap-session methods that can act after convergence without going through a node actor.
13. **Datastream emission provenance check**
Find direct calls that emit provisioning/readiness/fault/teardown telemetry. Control-plane telemetry should be derived from actor-observed events or actor-owned adapters; direct emission from random loops can hide imperative authority.
18. **External API client audit**
Identify VastAI, Docker, SSH, git, and filesystem operations. Provider plugins can perform provider I/O, but they should be stateless with respect to run authority; any retained run/node state inside the client/plugin is suspect.
19. **Ownership matrix by resource** -- Yes, but let us be careful about resource definition to catch these.
Build a table: resource type -> owning actor -> allowed non-actor adapter -> teardown message. Missing owner for processes, sockets, rings, leases, workers, or node records is a concrete migration target.
20. **Control-plane exception registry** -- How about a critical section boundary, so that any unactorized code gets flagged
Maintain a small checked-in allowlist: pure core, hot tensor byte path, startup bootstrap, pre-swactor SSH bootstrap, provider I/O adapter, test harness. Every denylist hit must match one exception or be filed as non-actor control code.
22. **Backtrace-based audit mode** -- Yes, but not with a 'registry', and only certain critical datastructures
Wrap side-effect APIs behind crate-local helpers and, in audit builds, record a lightweight backtrace/source tag. During e2e runs, fail or report when control-plane effects happen without an actor frame or registered bootstrap exception.
24. **Spawn wrapper migration** -- Interesting idea, consider later. Eventually want to migrate task/thread behavior to swactor runtime, but that is currently deferred to post-alpha.
Replace direct `thread::spawn`, `tokio::spawn`, and `Command::spawn` with crate-local wrappers like `spawn_actor_adapter`, `spawn_byte_pump`, `spawn_pre_swactor_bootstrap`, `spawn_test_helper`. The wrapper name forces classification and makes unclassified spawns easy to detect.
25. **Shadow-runtime detector** -- Multiple runtimes should be considered always wrong until future notice.
Flag ad-hoc Tokio runtimes or swactor runtimes not created by the runtime stack/binary bootstrap. Multiple runtimes are not always wrong, but they often correlate with code escaping the actor scheduler/control surface.
28. **Readiness/fault/teardown vocabulary scan**
Search emitted JSON/log labels and enum variants containing `ready`, `live`, `failed`, `fault`, `stopped`, `exited`, `teardown`, `destroyed`. These are control-plane facts; require actor observation/provenance.
33. **Test-harness exclusion rule**
Keep tests out of the main migration signal unless they define production-like support code reused by binaries. The crate has many e2e helpers with threads/processes; classify those separately to avoid noisy false positives.

View file

@ -19,7 +19,7 @@ swactor-transport = { path = "../transport" }
distribution = { path = "../distribution" } distribution = { path = "../distribution" }
iroh-driver = { path = "../iroh-driver" } iroh-driver = { path = "../iroh-driver" }
iroh = "0.98" iroh = "0.98"
tokio = { version = "1", features = ["rt-multi-thread", "macros", "process", "io-util", "sync", "time"] } tokio = { version = "1", features = ["rt-multi-thread", "macros", "process", "io-util", "sync", "time", "net"] }
swactor-vastai = { path = "../../tools/vastai" } swactor-vastai = { path = "../../tools/vastai" }
parking_lot = "0.12" parking_lot = "0.12"
blake3 = "1" blake3 = "1"
@ -28,6 +28,9 @@ toml = "0.8"
[target.'cfg(target_os = "linux")'.dependencies] [target.'cfg(target_os = "linux")'.dependencies]
libc = "0.2" libc = "0.2"
[dev-dependencies]
iroh-relay = { version = "0.98", features = ["server", "test-utils"] }
[[bin]] [[bin]]
name = "mvp-worker-node" name = "mvp-worker-node"
path = "src/bin/worker_node.rs" path = "src/bin/worker_node.rs"

View file

@ -1,6 +1,7 @@
# MVP Node Provisioning Specification # MVP Node Provisioning Specification
***STALE! FOR HISTORICAL REFERENCE ONLY***
**Status:** draft node-provisioning specification. **Status:**draft node-provisioning specification.
This document defines the MVP path from a static runplan node requirement to a This document defines the MVP path from a static runplan node requirement to a
remote swactor runtime joined to the orchestrator-side swarm. It covers provider remote swactor runtime joined to the orchestrator-side swarm. It covers provider

View file

@ -1,4 +1,5 @@
# MVP System Specification # MVP System Specification
***STALE! FOR HISTORICAL REFERENCE ONLY***
**Status:** draft consolidated system specification. **Status:** draft consolidated system specification.

View file

@ -58,6 +58,13 @@ pub enum NodeAgentMsg {
stage_index: u32, stage_index: u32,
endpoint: EndpointAddr, endpoint: EndpointAddr,
node_actor: ActorAddress, node_actor: ActorAddress,
readiness_id: u64,
},
RuntimeReadyAck {
run_id: u64,
node_id: u64,
stage_index: u32,
readiness_id: u64,
}, },
MarkWeightsReady, MarkWeightsReady,
MarkInboundEdgeReady { MarkInboundEdgeReady {
@ -184,6 +191,12 @@ pub enum NodeAgentReport {
max_tokens: u32, max_tokens: u32,
reply_to: ActorAddress, reply_to: ActorAddress,
}, },
RuntimeReadyAck {
run_id: u64,
node_id: u64,
stage_index: u32,
readiness_id: u64,
},
Snapshot { Snapshot {
commands: Vec<StageCommandWire>, commands: Vec<StageCommandWire>,
events: Vec<StageLifecycleWire>, events: Vec<StageLifecycleWire>,
@ -234,6 +247,7 @@ impl NodeAgentActor {
stage_index, stage_index,
endpoint, endpoint,
node_actor, node_actor,
readiness_id,
} => { } => {
self.core.observe(stage::StageEvent::WorkerReady); self.core.observe(stage::StageEvent::WorkerReady);
let _ = ctx.send( let _ = ctx.send(
@ -244,9 +258,29 @@ impl NodeAgentActor {
stage_index, stage_index,
endpoint, endpoint,
node_actor, node_actor,
readiness_id,
}, },
); );
} }
NodeAgentMsg::RuntimeReadyAck {
run_id,
node_id,
stage_index,
readiness_id,
} => {
if let Some(report_to) = self.report_to {
let _ = ctx.send(
report_to,
NodeAgentReport::RuntimeReadyAck {
run_id,
node_id,
stage_index,
readiness_id,
},
);
}
return;
}
NodeAgentMsg::MarkWeightsReady => self.core.observe(stage::StageEvent::WeightsReady), NodeAgentMsg::MarkWeightsReady => self.core.observe(stage::StageEvent::WeightsReady),
NodeAgentMsg::MarkInboundEdgeReady { edge_id } => { NodeAgentMsg::MarkInboundEdgeReady { edge_id } => {
self.core.observe(stage::StageEvent::InboundEdgeReady { self.core.observe(stage::StageEvent::InboundEdgeReady {
@ -534,6 +568,7 @@ mod tests {
stage_index: 3, stage_index: 3,
endpoint: endpoint.clone(), endpoint: endpoint.clone(),
node_actor, node_actor,
readiness_id: 99,
}, },
) )
.expect("send runtime loaded"); .expect("send runtime loaded");
@ -547,7 +582,50 @@ mod tests {
stage_index: 3, stage_index: 3,
endpoint, endpoint,
node_actor, node_actor,
readiness_id: 99,
}) })
); );
} }
#[test]
fn node_agent_runtime_ready_ack_reports_worker_loop() {
let runtime = Runtime::new(RuntimeConfig::default());
let orchestrator_inbox = runtime
.new_inbox::<OrchestratorMsg>()
.expect("orchestrator inbox");
let reports = runtime
.new_inbox::<NodeAgentReport>()
.expect("node report inbox");
let actor = runtime
.spawn(NodeAgentActor::new(
stage::NodeId(11),
*orchestrator_inbox.addr(),
Some(*reports.addr()),
))
.expect("spawn node agent");
runtime
.send_to(
actor,
NodeAgentMsg::RuntimeReadyAck {
run_id: 7,
node_id: 11,
stage_index: 3,
readiness_id: 99,
},
)
.expect("send runtime ready ack");
runtime.tick();
assert_eq!(
reports.try_recv(),
Some(NodeAgentReport::RuntimeReadyAck {
run_id: 7,
node_id: 11,
stage_index: 3,
readiness_id: 99,
})
);
assert_eq!(reports.try_recv(), None);
}
} }

View file

@ -33,6 +33,7 @@ pub enum OrchestratorMsg {
stage_index: u32, stage_index: u32,
endpoint: EndpointAddr, endpoint: EndpointAddr,
node_actor: ActorAddress, node_actor: ActorAddress,
readiness_id: u64,
}, },
ObserveTokenInEndpointReady, ObserveTokenInEndpointReady,
ObserveTokenOutEndpointReady, ObserveTokenOutEndpointReady,
@ -134,6 +135,7 @@ pub enum OrchestratorReport {
stage_index: u32, stage_index: u32,
endpoint: EndpointAddr, endpoint: EndpointAddr,
node_actor: ActorAddress, node_actor: ActorAddress,
readiness_id: u64,
}, },
Snapshot { Snapshot {
commands: Vec<RunCommandWire>, commands: Vec<RunCommandWire>,
@ -269,6 +271,7 @@ impl ActorInterface for OrchestratorActor {
stage_index, stage_index,
endpoint, endpoint,
node_actor, node_actor,
readiness_id,
} = msg.clone() } = msg.clone()
{ {
if let Some(report_to) = self.report_to { if let Some(report_to) = self.report_to {
@ -280,6 +283,7 @@ impl ActorInterface for OrchestratorActor {
stage_index, stage_index,
endpoint, endpoint,
node_actor, node_actor,
readiness_id,
}, },
); );
} }
@ -421,6 +425,7 @@ mod tests {
stage_index: 3, stage_index: 3,
endpoint: endpoint.clone(), endpoint: endpoint.clone(),
node_actor, node_actor,
readiness_id: 99,
}, },
) )
.expect("send runtime ready"); .expect("send runtime ready");
@ -434,6 +439,7 @@ mod tests {
stage_index: 3, stage_index: 3,
endpoint, endpoint,
node_actor, node_actor,
readiness_id: 99,
}) })
); );
assert_eq!(reports.try_recv(), None); assert_eq!(reports.try_recv(), None);

View file

@ -1109,6 +1109,22 @@ fn orch_local_e2e_marker(bin: &Path) -> PathBuf {
marker marker
} }
fn orch_binary_fingerprint(bin: &Path, root: &Path) -> Result<String, String> {
let display = display_workspace_path(root, bin);
let metadata = fs::metadata(bin).map_err(|e| format!("stat {display}: {e}"))?;
let modified = metadata
.modified()
.map_err(|e| format!("modified time {display}: {e}"))?;
let modified_ns = modified
.duration_since(std::time::UNIX_EPOCH)
.map_err(|e| format!("modified time before Unix epoch for {display}: {e}"))?
.as_nanos();
Ok(format!(
"local-e2e\nlen={}\nmodified_ns={modified_ns}\n",
metadata.len()
))
}
fn orch_local_e2e_marker_stale(bin: &Path, root: &Path) -> Result<bool, String> { fn orch_local_e2e_marker_stale(bin: &Path, root: &Path) -> Result<bool, String> {
if !bin.is_file() { if !bin.is_file() {
return Ok(true); return Ok(true);
@ -1117,13 +1133,16 @@ fn orch_local_e2e_marker_stale(bin: &Path, root: &Path) -> Result<bool, String>
if !marker.is_file() { if !marker.is_file() {
return Ok(true); return Ok(true);
} }
Ok(modified_time(root, &marker)? < modified_time(root, bin)?) let expected = orch_binary_fingerprint(bin, root)?;
let actual = fs::read_to_string(&marker).unwrap_or_default();
Ok(actual != expected)
} }
fn write_orch_local_e2e_marker(bin: &Path, root: &Path) -> Result<(), String> { fn write_orch_local_e2e_marker(bin: &Path, root: &Path) -> Result<(), String> {
let marker = orch_local_e2e_marker(bin); let marker = orch_local_e2e_marker(bin);
let display = display_workspace_path(root, &marker); let display = display_workspace_path(root, &marker);
fs::write(&marker, b"local-e2e\n").map_err(|e| format!("write {display}: {e}")) let fingerprint = orch_binary_fingerprint(bin, root)?;
fs::write(&marker, fingerprint).map_err(|e| format!("write {display}: {e}"))
} }
fn latest_mtime(root: &Path, path: &Path) -> Result<SystemTime, String> { fn latest_mtime(root: &Path, path: &Path) -> Result<SystemTime, String> {

View file

@ -11,6 +11,7 @@ use std::time::{Duration, Instant};
use datastream::{ChannelId, DatastreamSink, Frame, Lifetime, Mux, NodeId, StreamId}; use datastream::{ChannelId, DatastreamSink, Frame, Lifetime, Mux, NodeId, StreamId};
use distribution::node::DistributedNodeConfig; use distribution::node::DistributedNodeConfig;
use distribution::types::{MemberState, NodeId as DistNodeId};
use iroh::EndpointAddr; use iroh::EndpointAddr;
use iroh_driver::{IrohDriver, IrohDriverConfig}; use iroh_driver::{IrohDriver, IrohDriverConfig};
use mvp_system::actors::node_agent::{NodeAgentMsg, StageProvisionWire}; use mvp_system::actors::node_agent::{NodeAgentMsg, StageProvisionWire};
@ -60,6 +61,8 @@ const PUMP_INTERVAL: Duration = Duration::from_millis(10);
const MVP_ORCH_BOOTSTRAP: &str = "mvp.orch.bootstrap"; const MVP_ORCH_BOOTSTRAP: &str = "mvp.orch.bootstrap";
const MVP_ORCH_PROMPT: &str = "mvp.orch.prompt"; const MVP_ORCH_PROMPT: &str = "mvp.orch.prompt";
const DATASTREAM_FRAME_LOG_ENV: &str = "MVP_DATASTREAM_FRAME_LOG"; const DATASTREAM_FRAME_LOG_ENV: &str = "MVP_DATASTREAM_FRAME_LOG";
const DEFAULT_DOCKER_CONTAINER_PREFIX: &str = "mvp-orchestrator";
const MVP_DOCKER_CONTAINER_PREFIX_ENV: &str = "MVP_DOCKER_CONTAINER_PREFIX";
fn main() -> ExitCode { fn main() -> ExitCode {
match run() { match run() {
@ -484,36 +487,38 @@ fn run() -> Result<(), String> {
} }
}; };
provisioned_node.complete_bootstrap()?; provisioned_node.complete_bootstrap()?;
driver.join(std::slice::from_ref(&ready.endpoint)); match enqueue_runtime_ready_ack(&stack, &ready, config.run_id, config.node_id) {
Ok(()) => {
driver.drain_outbox(&stack.outbox);
orch_datastream.emit_bootstrap(
dashboard.as_ref(),
config.run_id,
config.node_id,
"runtime_ready_ack",
"ready",
json!({"node_actor":ready.node_actor,"readiness_id":ready.readiness_id}),
);
}
Err(error) => {
orch_datastream.emit_bootstrap(
dashboard.as_ref(),
config.run_id,
config.node_id,
"runtime_ready_ack",
"failed",
json!({"node_actor":ready.node_actor,"readiness_id":ready.readiness_id,"error":error}),
);
return Err(error);
}
}
orch_datastream.emit_bootstrap( orch_datastream.emit_bootstrap(
dashboard.as_ref(), dashboard.as_ref(),
config.run_id, config.run_id,
config.node_id, config.node_id,
"node_join", "node_join",
"ready", "ready",
json!({"endpoint":&ready.endpoint}), json!({"endpoint":&ready.endpoint,"source":"runtime_ready_barrier"}),
); );
match wait_for_route(&mut driver, &stack, ready.node_actor, &stop_rx) {
Ok(()) => orch_datastream.emit_bootstrap(
dashboard.as_ref(),
config.run_id,
config.node_id,
"node_route",
"ready",
json!({"node_actor":ready.node_actor}),
),
Err(error) => {
orch_datastream.emit_bootstrap(
dashboard.as_ref(),
config.run_id,
config.node_id,
"node_route",
"failed",
json!({"node_actor":ready.node_actor,"error":error}),
);
return Err(error);
}
}
orch_datastream.emit_bootstrap( orch_datastream.emit_bootstrap(
dashboard.as_ref(), dashboard.as_ref(),
config.run_id, config.run_id,
@ -1609,7 +1614,7 @@ impl Config {
bootstrap_runtime: Arc<swactor::runtime::Runtime>, bootstrap_runtime: Arc<swactor::runtime::Runtime>,
) -> Result<Box<dyn ProvisionPlugin>, String> { ) -> Result<Box<dyn ProvisionPlugin>, String> {
match self.provider { match self.provider {
ProviderKind::Docker => Ok(Box::new(LocalDockerPlugin::new("mvp-orchestrator"))), ProviderKind::Docker => Ok(Box::new(LocalDockerPlugin::new(docker_container_prefix()))),
ProviderKind::VastAi => { ProviderKind::VastAi => {
let vastai = self.vastai.as_ref().ok_or_else(|| { let vastai = self.vastai.as_ref().ok_or_else(|| {
"VastAI config was not resolved for provider vastai".to_owned() "VastAI config was not resolved for provider vastai".to_owned()
@ -1805,6 +1810,33 @@ struct RuntimeReady {
endpoint: EndpointAddr, endpoint: EndpointAddr,
node_actor: ActorAddress, node_actor: ActorAddress,
stage_index: u32, stage_index: u32,
readiness_id: u64,
swim_node_id: DistNodeId,
}
fn runtime_ready_barrier_met(stack: &DistributionRuntimeStack, ready: &RuntimeReady) -> bool {
stack.member_state(ready.swim_node_id) == Some(MemberState::Alive)
&& stack.route_owner(ready.node_actor) == Some(ready.swim_node_id)
}
fn enqueue_runtime_ready_ack(
stack: &DistributionRuntimeStack,
ready: &RuntimeReady,
run_id: u64,
node_id: u64,
) -> Result<(), String> {
stack
.runtime
.send_to(
ready.node_actor,
NodeAgentMsg::RuntimeReadyAck {
run_id,
node_id,
stage_index: ready.stage_index,
readiness_id: ready.readiness_id,
},
)
.map_err(|e| format!("send runtime ready ack: {e}"))
} }
struct ProvisionedNodeGuard<'a> { struct ProvisionedNodeGuard<'a> {
@ -2294,6 +2326,10 @@ fn wait_for_runtime_ready(
node_id: u64, node_id: u64,
provider: ProviderKind, provider: ProviderKind,
) -> Result<RuntimeReady, String> { ) -> Result<RuntimeReady, String> {
let mut pending_ready: Option<RuntimeReady> = None;
let mut node_swim_started = false;
let mut node_swim_ready = false;
let mut node_route_started = false;
loop { loop {
pump(driver, stack); pump(driver, stack);
drain_frames(frame_rx, dashboard, orch_datastream); drain_frames(frame_rx, dashboard, orch_datastream);
@ -2321,39 +2357,72 @@ fn wait_for_runtime_ready(
stage_index, stage_index,
endpoint, endpoint,
node_actor, node_actor,
readiness_id,
} = report } = report
{ {
if report_run_id == run_id && report_node_id == node_id { if report_run_id == run_id && report_node_id == node_id {
return Ok(RuntimeReady { let reset_progress = pending_ready
.as_ref()
.map(|ready| ready.readiness_id != readiness_id)
.unwrap_or(true);
if reset_progress {
node_swim_started = false;
node_swim_ready = false;
node_route_started = false;
}
let swim_node_id = DistNodeId(*endpoint.id.as_bytes());
pending_ready = Some(RuntimeReady {
endpoint, endpoint,
node_actor, node_actor,
stage_index, stage_index,
readiness_id,
swim_node_id,
}); });
} }
} }
} }
thread::sleep(PUMP_INTERVAL); if let Some(ready) = pending_ready.as_ref() {
if runtime_ready_barrier_met(stack, ready) {
return Ok(ready.clone());
}
let swim_ready = stack.member_state(ready.swim_node_id) == Some(MemberState::Alive);
let route_ready = stack.route_owner(ready.node_actor) == Some(ready.swim_node_id);
if !swim_ready {
if !node_swim_started {
orch_datastream.emit_bootstrap(
dashboard,
run_id,
node_id,
"node_swim",
"started",
json!({"node":format!("{:?}", ready.swim_node_id),"readiness_id":ready.readiness_id}),
);
node_swim_started = true;
}
} else if !route_ready {
if !node_swim_ready {
orch_datastream.emit_bootstrap(
dashboard,
run_id,
node_id,
"node_swim",
"ready",
json!({"node":format!("{:?}", ready.swim_node_id),"readiness_id":ready.readiness_id}),
);
node_swim_ready = true;
}
if !node_route_started {
orch_datastream.emit_bootstrap(
dashboard,
run_id,
node_id,
"node_route",
"started",
json!({"node_actor":ready.node_actor,"node":format!("{:?}", ready.swim_node_id),"readiness_id":ready.readiness_id}),
);
node_route_started = true;
} }
}
fn wait_for_route(
driver: &mut IrohDriver,
stack: &DistributionRuntimeStack,
actor: ActorAddress,
stop_rx: &mpsc::Receiver<()>,
) -> Result<(), String> {
loop {
pump(driver, stack);
if stop_requested(stop_rx) {
return Err("shutdown requested while waiting for node route".to_owned());
} }
let ready = stack
.route_view
.read()
.map(|view| view.contains_key(&actor))
.unwrap_or(false);
if ready {
return Ok(());
} }
thread::sleep(PUMP_INTERVAL); thread::sleep(PUMP_INTERVAL);
} }
@ -2812,6 +2881,11 @@ fn env_optional(name: &str) -> Option<String> {
.filter(|value| !value.is_empty()) .filter(|value| !value.is_empty())
} }
fn docker_container_prefix() -> String {
env_optional(MVP_DOCKER_CONTAINER_PREFIX_ENV)
.unwrap_or_else(|| DEFAULT_DOCKER_CONTAINER_PREFIX.to_owned())
}
fn optional_env(name: &str) -> Option<(String, String)> { fn optional_env(name: &str) -> Option<(String, String)> {
env_optional(name).map(|value| (name.to_owned(), value)) env_optional(name).map(|value| (name.to_owned(), value))
} }
@ -3045,6 +3119,7 @@ mod tests {
"MVP_VASTAI_SSH_IDENTITY", "MVP_VASTAI_SSH_IDENTITY",
"VASTAI_API_KEY", "VASTAI_API_KEY",
SWACTOR_IROH_RELAY_URL_ENV, SWACTOR_IROH_RELAY_URL_ENV,
MVP_DOCKER_CONTAINER_PREFIX_ENV,
]; ];
struct RestoreEnv { struct RestoreEnv {
@ -3126,6 +3201,16 @@ mod tests {
.map(|(_, value)| value.as_str()) .map(|(_, value)| value.as_str())
} }
#[test]
fn docker_container_prefix_defaults_and_trims_env_override() {
with_clean_env(&[], || {
assert_eq!(docker_container_prefix(), DEFAULT_DOCKER_CONTAINER_PREFIX);
});
with_clean_env(&[(MVP_DOCKER_CONTAINER_PREFIX_ENV, " custom-prefix ")], || {
assert_eq!(docker_container_prefix(), "custom-prefix");
});
}
#[test] #[test]
fn expand_home_path_expands_leading_home_segment() { fn expand_home_path_expands_leading_home_segment() {
with_clean_env(&[("HOME", "/tmp/mvp-vastai-home")], || { with_clean_env(&[("HOME", "/tmp/mvp-vastai-home")], || {
@ -3574,6 +3659,39 @@ bootstrap_command = "/run"
assert_eq!(env_value(&disabled_env, MVP_IROH_RELAY_URL_ENV), None); assert_eq!(env_value(&disabled_env, MVP_IROH_RELAY_URL_ENV), None);
} }
#[test]
fn node_spec_preserves_relay_transport_in_coordinator_endpoint() {
with_clean_env(&[], || {
let config = Config::from_layers_with_path_and_args(None, std::iter::empty::<String>())
.expect("config parses");
let secret = iroh::SecretKey::from_bytes(&[10; 32]);
let coordinator = EndpointAddr::new(secret.public()).with_relay_url(
"http://relay.example.com"
.parse::<iroh::RelayUrl>()
.unwrap(),
);
let datastream_sink = ActorAddress([19; 32]);
let orchestrator_actor = ActorAddress([20; 32]);
let spec = config
.node_spec(coordinator, datastream_sink, orchestrator_actor)
.expect("node spec builds");
let coordinator_endpoint_json = env_value(&spec.env, "MVP_COORDINATOR_ENDPOINT")
.expect("coordinator endpoint env is present");
let coordinator_endpoint =
serde_json::from_str::<EndpointAddr>(coordinator_endpoint_json)
.expect("coordinator endpoint env deserializes");
assert_eq!(
coordinator_endpoint
.relay_urls()
.next()
.map(|url| url.to_string()),
Some("http://relay.example.com/".to_owned())
);
});
}
#[test] #[test]
fn docker_config_construction_ignores_malformed_vastai_environment() { fn docker_config_construction_ignores_malformed_vastai_environment() {
let config = with_clean_env( let config = with_clean_env(
@ -3735,4 +3853,110 @@ bootstrap_command = "/run"
assert!(stop_requested(&rx)); assert!(stop_requested(&rx));
assert!(!stop_requested(&rx)); assert!(!stop_requested(&rx));
} }
fn endpoint(seed: u8) -> EndpointAddr {
EndpointAddr::new(iroh::SecretKey::from_bytes(&[seed; 32]).public())
}
fn insert_route(stack: &DistributionRuntimeStack, actor: ActorAddress, owner: DistNodeId) {
let mut route_view = match stack.route_view.write() {
Ok(route_view) => route_view,
Err(poisoned) => poisoned.into_inner(),
};
route_view.insert(actor, owner);
}
fn mark_alive(stack: &DistributionRuntimeStack, node_id: DistNodeId) {
stack
.runtime
.send_to(
stack.actors.membership_fanout,
distribution::swim::actor::MembershipChanged {
node_id,
state: MemberState::Alive,
incarnation: 1,
},
)
.expect("send membership change");
stack.pump_runtime_once();
}
#[test]
fn runtime_ready_barrier_waits_for_specific_swim_and_route() {
let remote = DistNodeId([2; 32]);
let node_actor = ActorAddress::new_random();
let ready = RuntimeReady {
endpoint: endpoint(2),
node_actor,
stage_index: 3,
readiness_id: 99,
swim_node_id: remote,
};
let stack =
DistributionRuntimeStack::new(DistNodeId([1; 32]), DistributedNodeConfig::default());
assert!(!runtime_ready_barrier_met(&stack, &ready));
let stack =
DistributionRuntimeStack::new(DistNodeId([1; 32]), DistributedNodeConfig::default());
mark_alive(&stack, remote);
assert!(!runtime_ready_barrier_met(&stack, &ready));
let stack =
DistributionRuntimeStack::new(DistNodeId([1; 32]), DistributedNodeConfig::default());
insert_route(&stack, node_actor, remote);
assert!(!runtime_ready_barrier_met(&stack, &ready));
let stack =
DistributionRuntimeStack::new(DistNodeId([1; 32]), DistributedNodeConfig::default());
mark_alive(&stack, remote);
insert_route(&stack, node_actor, remote);
assert!(runtime_ready_barrier_met(&stack, &ready));
}
#[test]
fn enqueue_runtime_ready_ack_reports_to_node_agent() {
use mvp_system::actors::node_agent::{NodeAgentActor, NodeAgentReport};
use mvp_system::actors::orchestrator::OrchestratorMsg;
use mvp_system::stage_controller as stage;
let stack =
DistributionRuntimeStack::new(DistNodeId([1; 32]), DistributedNodeConfig::default());
let orchestrator_inbox = stack
.runtime
.new_inbox::<OrchestratorMsg>()
.expect("orchestrator inbox");
let reports = stack
.runtime
.new_inbox::<NodeAgentReport>()
.expect("node report inbox");
let node_actor = stack
.runtime
.spawn(NodeAgentActor::new(
stage::NodeId(11),
*orchestrator_inbox.addr(),
Some(*reports.addr()),
))
.expect("spawn node agent");
let ready = RuntimeReady {
endpoint: endpoint(9),
node_actor,
stage_index: 3,
readiness_id: 99,
swim_node_id: DistNodeId([2; 32]),
};
enqueue_runtime_ready_ack(&stack, &ready, 7, 11).expect("enqueue runtime ready ack");
stack.pump_runtime_once();
assert_eq!(
reports.try_recv(),
Some(NodeAgentReport::RuntimeReadyAck {
run_id: 7,
node_id: 11,
stage_index: ready.stage_index,
readiness_id: 99,
})
);
}
} }

View file

@ -1,5 +1,6 @@
use std::fs::{File, OpenOptions}; use std::fs::{self, File, OpenOptions};
use std::io::{BufRead, BufReader, Write}; use std::io::{BufRead, BufReader, Read, Write};
use std::path::PathBuf;
use std::process::{Child, ChildStdin, ChildStdout, Command, ExitCode, Stdio}; use std::process::{Child, ChildStdin, ChildStdout, Command, ExitCode, Stdio};
use std::sync::{ use std::sync::{
Arc, OnceLock, Arc, OnceLock,
@ -14,6 +15,7 @@ use datastream::emit::{
}; };
use distribution::node::DistributedNodeConfig; use distribution::node::DistributedNodeConfig;
use distribution::types::{MemberState, NodeId as DistNodeId};
use iroh::EndpointAddr; use iroh::EndpointAddr;
use iroh_driver::{IrohDriver, IrohDriverConfig}; use iroh_driver::{IrohDriver, IrohDriverConfig};
use mvp_system::actors::node_agent::{ use mvp_system::actors::node_agent::{
@ -29,6 +31,7 @@ use mvp_system::stage_controller as stage;
use parking_lot::Mutex; use parking_lot::Mutex;
use serde_json::{Value, json}; use serde_json::{Value, json};
use swactor::actor::ActorAddress; use swactor::actor::ActorAddress;
use tokio::io::{AsyncBufReadExt, AsyncWriteExt};
const DEFAULT_WORKER_SCRIPT: &str = "/usr/local/share/mvp/tinygrad_worker.py"; const DEFAULT_WORKER_SCRIPT: &str = "/usr/local/share/mvp/tinygrad_worker.py";
const DEFAULT_DEVICE: &str = "CUDA"; const DEFAULT_DEVICE: &str = "CUDA";
@ -38,6 +41,8 @@ const DEFAULT_MODEL_ID: &str = "llama-3.2-1b-instruct-q4";
const DEFAULT_ARENA_BYTES: u64 = 64 * 1024 * 1024; const DEFAULT_ARENA_BYTES: u64 = 64 * 1024 * 1024;
const DEFAULT_ARENA_ALIGNMENT: u64 = 64; const DEFAULT_ARENA_ALIGNMENT: u64 = 64;
const PUMP_INTERVAL: Duration = Duration::from_millis(10); const PUMP_INTERVAL: Duration = Duration::from_millis(10);
const RUNTIME_READY_RETRY_INITIAL: Duration = Duration::from_millis(100);
const RUNTIME_READY_RETRY_MAX: Duration = Duration::from_secs(2);
const NODE_BOOTSTRAP_CHANNEL: &str = "mvp.node.bootstrap"; const NODE_BOOTSTRAP_CHANNEL: &str = "mvp.node.bootstrap";
const NODE_RUNTIME_CHANNEL: &str = "mvp.node.runtime"; const NODE_RUNTIME_CHANNEL: &str = "mvp.node.runtime";
const NODE_STAGE_CHANNEL: &str = "mvp.node.stage"; const NODE_STAGE_CHANNEL: &str = "mvp.node.stage";
@ -98,6 +103,298 @@ fn emit_node_event(
datastream.tick(); datastream.tick();
} }
#[derive(serde::Deserialize, serde::Serialize)]
#[serde(tag = "type")]
enum DebugJoinRequestWire {
JoinEndpoint { endpoint: EndpointAddr },
}
#[derive(serde::Deserialize, serde::Serialize)]
#[serde(tag = "type")]
enum DebugJoinResponseWire {
JoinQueued {
peer_node_id: String,
has_relay: bool,
direct_addr_count: usize,
},
JoinRejected {
error: String,
detail: String,
},
}
enum DebugJoinCommand {
JoinEndpoint {
endpoint: EndpointAddr,
reply: tokio::sync::oneshot::Sender<DebugJoinResponseWire>,
},
}
enum DebugJoinClientError {
Cli(String),
Runtime(String),
}
fn debug_join_client_main(args: Vec<String>) -> ExitCode {
match run_debug_join_client(args) {
Ok(response) => {
let queued = matches!(response, DebugJoinResponseWire::JoinQueued { .. });
match serde_json::to_string(&response) {
Ok(line) => println!("{line}"),
Err(error) => {
eprintln!("mvp-worker-node debug-join: serialize response: {error}");
return ExitCode::from(1);
}
}
if queued {
ExitCode::SUCCESS
} else {
ExitCode::from(1)
}
}
Err(DebugJoinClientError::Cli(error)) => {
eprintln!("mvp-worker-node debug-join: {error}");
ExitCode::from(2)
}
Err(DebugJoinClientError::Runtime(error)) => {
eprintln!("mvp-worker-node debug-join: {error}");
ExitCode::from(1)
}
}
}
fn run_debug_join_client(args: Vec<String>) -> Result<DebugJoinResponseWire, DebugJoinClientError> {
let mut socket = None;
let mut endpoint_json = None;
let mut read_endpoint_stdin = false;
let mut iter = args.into_iter();
while let Some(arg) = iter.next() {
match arg.as_str() {
"--socket" => {
socket = Some(PathBuf::from(iter.next().ok_or_else(|| {
DebugJoinClientError::Cli("--socket requires a path".to_owned())
})?));
}
"--endpoint-json" => {
endpoint_json = Some(iter.next().ok_or_else(|| {
DebugJoinClientError::Cli("--endpoint-json requires JSON".to_owned())
})?);
}
"--endpoint-json-stdin" => read_endpoint_stdin = true,
other => {
return Err(DebugJoinClientError::Cli(format!(
"unknown argument {other:?}; usage: debug-join --socket <path> (--endpoint-json <json> | --endpoint-json-stdin)"
)));
}
}
}
let socket = socket.ok_or_else(|| {
DebugJoinClientError::Cli(
"missing --socket <path>; usage: debug-join --socket <path> (--endpoint-json <json> | --endpoint-json-stdin)".to_owned(),
)
})?;
let endpoint_json = match (endpoint_json, read_endpoint_stdin) {
(Some(_), true) => {
return Err(DebugJoinClientError::Cli(
"use either --endpoint-json or --endpoint-json-stdin, not both".to_owned(),
));
}
(Some(json), false) => json,
(None, true) => {
let mut json = String::new();
std::io::stdin()
.read_to_string(&mut json)
.map_err(|e| DebugJoinClientError::Runtime(format!("read endpoint stdin: {e}")))?;
json
}
(None, false) => {
return Err(DebugJoinClientError::Cli(
"missing endpoint JSON; use --endpoint-json <json> or --endpoint-json-stdin"
.to_owned(),
));
}
};
let endpoint = serde_json::from_str::<EndpointAddr>(&endpoint_json)
.map_err(|e| DebugJoinClientError::Cli(format!("parse endpoint JSON: {e}")))?;
let request = debug_join_request_line(endpoint).map_err(DebugJoinClientError::Runtime)?;
let mut stream = std::os::unix::net::UnixStream::connect(&socket)
.map_err(|e| DebugJoinClientError::Runtime(format!("connect {}: {e}", socket.display())))?;
stream
.write_all(request.as_bytes())
.map_err(|e| DebugJoinClientError::Runtime(format!("write request: {e}")))?;
stream
.flush()
.map_err(|e| DebugJoinClientError::Runtime(format!("flush request: {e}")))?;
let mut response_line = String::new();
BufReader::new(stream)
.read_line(&mut response_line)
.map_err(|e| DebugJoinClientError::Runtime(format!("read response: {e}")))?;
if response_line.trim().is_empty() {
return Err(DebugJoinClientError::Runtime(
"debug join socket closed without response".to_owned(),
));
}
serde_json::from_str::<DebugJoinResponseWire>(&response_line)
.map_err(|e| DebugJoinClientError::Runtime(format!("parse response JSON: {e}")))
}
fn debug_join_request_line(endpoint: EndpointAddr) -> Result<String, String> {
serde_json::to_string(&DebugJoinRequestWire::JoinEndpoint { endpoint })
.map(|mut line| {
line.push('\n');
line
})
.map_err(|e| format!("serialize debug join request: {e}"))
}
fn spawn_debug_join_listener(
handle: tokio::runtime::Handle,
path: PathBuf,
) -> Result<tokio::sync::mpsc::UnboundedReceiver<DebugJoinCommand>, String> {
use std::os::unix::fs::PermissionsExt;
match fs::remove_file(&path) {
Ok(()) => {}
Err(error) if error.kind() == std::io::ErrorKind::NotFound => {}
Err(error) => {
return Err(format!(
"remove stale debug join socket {}: {error}",
path.display()
));
}
}
let listener = {
let _guard = handle.enter();
tokio::net::UnixListener::bind(&path)
.map_err(|e| format!("bind debug join socket {}: {e}", path.display()))?
};
fs::set_permissions(&path, fs::Permissions::from_mode(0o600))
.map_err(|e| format!("chmod debug join socket {}: {e}", path.display()))?;
let (command_tx, command_rx) = tokio::sync::mpsc::unbounded_channel::<DebugJoinCommand>();
handle.spawn(async move {
loop {
match listener.accept().await {
Ok((stream, _addr)) => {
let command_tx = command_tx.clone();
tokio::spawn(async move {
handle_debug_join_stream(stream, command_tx).await;
});
}
Err(error) => {
eprintln!("mvp-worker-node debug join listener stopped: {error}");
break;
}
}
}
});
Ok(command_rx)
}
async fn handle_debug_join_stream(
stream: tokio::net::UnixStream,
command_tx: tokio::sync::mpsc::UnboundedSender<DebugJoinCommand>,
) {
let mut reader = tokio::io::BufReader::new(stream);
let mut line = String::new();
let response = match reader.read_line(&mut line).await {
Ok(0) => DebugJoinResponseWire::JoinRejected {
error: "MalformedCommand".to_owned(),
detail: "empty request".to_owned(),
},
Ok(_) => match parse_debug_join_request(&line) {
Ok(DebugJoinRequestWire::JoinEndpoint { endpoint }) => {
let (reply, response_rx) = tokio::sync::oneshot::channel();
if command_tx
.send(DebugJoinCommand::JoinEndpoint { endpoint, reply })
.is_err()
{
DebugJoinResponseWire::JoinRejected {
error: "CommandQueueClosed".to_owned(),
detail: "worker main loop is not accepting debug join commands".to_owned(),
}
} else {
response_rx
.await
.unwrap_or_else(|error| DebugJoinResponseWire::JoinRejected {
error: "CommandCancelled".to_owned(),
detail: error.to_string(),
})
}
}
Err(response) => response,
},
Err(error) => DebugJoinResponseWire::JoinRejected {
error: "MalformedCommand".to_owned(),
detail: format!("read request: {error}"),
},
};
let mut stream = reader.into_inner();
if let Ok(line) = serde_json::to_string(&response) {
let _ = stream.write_all(line.as_bytes()).await;
let _ = stream.write_all(b"\n").await;
let _ = stream.flush().await;
}
}
fn parse_debug_join_request(raw: &str) -> Result<DebugJoinRequestWire, DebugJoinResponseWire> {
let value =
serde_json::from_str::<Value>(raw).map_err(|e| DebugJoinResponseWire::JoinRejected {
error: "MalformedCommand".to_owned(),
detail: e.to_string(),
})?;
let endpoint_decode_error = value.get("type").and_then(Value::as_str) == Some("JoinEndpoint")
&& value.get("endpoint").is_some();
serde_json::from_value::<DebugJoinRequestWire>(value).map_err(|e| {
DebugJoinResponseWire::JoinRejected {
error: if endpoint_decode_error {
"MalformedEndpoint"
} else {
"MalformedCommand"
}
.to_owned(),
detail: e.to_string(),
}
})
}
fn drain_debug_join_commands(
debug_join_rx: &mut Option<tokio::sync::mpsc::UnboundedReceiver<DebugJoinCommand>>,
driver: &mut IrohDriver,
config: &DeploymentConfig,
datastream: &mut DatastreamEmitter,
) {
let Some(rx) = debug_join_rx else {
return;
};
while let Ok(command) = rx.try_recv() {
match command {
DebugJoinCommand::JoinEndpoint { endpoint, reply } => {
let peer_node_id = endpoint.id.to_string();
let has_relay = endpoint.relay_urls().next().is_some();
let direct_addr_count = endpoint.ip_addrs().count();
driver.join(std::slice::from_ref(&endpoint));
emit_node_event(
datastream,
config,
NODE_RUNTIME_CHANNEL,
"debug_join",
"queued",
json!({
"peer_node_id":peer_node_id,
"has_relay":has_relay,
"direct_addr_count":direct_addr_count,
}),
);
let _ = reply.send(DebugJoinResponseWire::JoinQueued {
peer_node_id,
has_relay,
direct_addr_count,
});
}
}
}
}
fn spawn_host_gpu_sampler(handle: tokio::runtime::Handle, sink: DatastreamEventSink) { fn spawn_host_gpu_sampler(handle: tokio::runtime::Handle, sink: DatastreamEventSink) {
handle.spawn(async move { handle.spawn(async move {
let mut seq = 0_u64; let mut seq = 0_u64;
@ -192,6 +489,11 @@ fn spawn_arena_sampler(
} }
fn main() -> ExitCode { fn main() -> ExitCode {
let mut args = std::env::args().skip(1).collect::<Vec<_>>();
if args.first().map(String::as_str) == Some("debug-join") {
args.remove(0);
return debug_join_client_main(args);
}
match run() { match run() {
Ok(()) => ExitCode::SUCCESS, Ok(()) => ExitCode::SUCCESS,
Err(error) => { Err(error) => {
@ -218,6 +520,7 @@ fn run() -> Result<(), String> {
"self_test_enabled":config.self_test_prompt.is_some(), "self_test_enabled":config.self_test_prompt.is_some(),
"arena_bytes":config.arena_bytes, "arena_bytes":config.arena_bytes,
"arena_alignment":config.arena_alignment, "arena_alignment":config.arena_alignment,
"debug_join_socket":config.debug_join_socket.as_deref().unwrap_or("disabled"),
}), }),
)?; )?;
emit_stdio_node_event( emit_stdio_node_event(
@ -383,6 +686,45 @@ fn run() -> Result<(), String> {
"ready", "ready",
config.datastream_sink_detail(), config.datastream_sink_detail(),
)?; )?;
let mut debug_join_rx = match &config.debug_join_socket {
Some(path) => {
match spawn_debug_join_listener(tokio.handle().clone(), PathBuf::from(path)) {
Ok(rx) => {
emit_node_event(
&mut datastream,
&config,
NODE_RUNTIME_CHANNEL,
"debug_join_socket",
"ready",
json!({"socket":path}),
);
Some(rx)
}
Err(error) => {
emit_node_event(
&mut datastream,
&config,
NODE_RUNTIME_CHANNEL,
"debug_join_socket",
"failed",
json!({"socket":path,"error":error}),
);
return Err(format!("bind debug join socket {}: {error}", path));
}
}
}
None => {
emit_node_event(
&mut datastream,
&config,
NODE_RUNTIME_CHANNEL,
"debug_join_socket",
"skipped",
json!({"reason":"MVP_DEBUG_JOIN_SOCKET=disabled"}),
);
None
}
};
let reports = match stack.runtime.new_inbox::<NodeAgentReport>() { let reports = match stack.runtime.new_inbox::<NodeAgentReport>() {
Ok(inbox) => { Ok(inbox) => {
@ -512,34 +854,8 @@ fn run() -> Result<(), String> {
return Err(error); return Err(error);
} }
} }
match stack.runtime.send_to( let mut pending_runtime_ready =
node_actor, PendingRuntimeReady::new(&config, driver.endpoint_addr(), node_actor);
NodeAgentMsg::RuntimeLoaded {
run_id: config.run_id,
node_id: config.logical_node_id,
stage_index: config.stage_index,
endpoint: driver.endpoint_addr(),
node_actor,
},
) {
Ok(()) => emit_stdio_node_event(
&config,
NODE_RUNTIME_CHANNEL,
"runtime_loaded",
"ready",
json!({"sent":"NodeAgentMsg::RuntimeLoaded","node_actor":node_actor}),
)?,
Err(error) => {
emit_stdio_node_event(
&config,
NODE_RUNTIME_CHANNEL,
"runtime_loaded",
"failed",
json!({"error":error.to_string()}),
)?;
return Err(format!("signal runtime loaded: {error}"));
}
}
let ready = json!({ let ready = json!({
"type":"ready", "type":"ready",
@ -552,23 +868,16 @@ fn run() -> Result<(), String> {
emit_stdio_node_event( emit_stdio_node_event(
&config, &config,
NODE_BOOTSTRAP_CHANNEL, NODE_BOOTSTRAP_CHANNEL,
"runtime_ready", "runtime_ready_local",
"ready", "ready",
json!({ json!({
"endpoint":driver.endpoint_addr(), "endpoint":driver.endpoint_addr(),
"node_actor":node_actor, "node_actor":node_actor,
"logical_node_id":config.logical_node_id, "logical_node_id":config.logical_node_id,
"stage_index":config.stage_index, "stage_index":config.stage_index,
"readiness_id":pending_runtime_ready.readiness_id,
}), }),
)?; )?;
datastream.submit_text(ChannelId::new("mvp.node.ready"), ready.to_string());
emit_stdio_node_event(
&config,
NODE_BOOTSTRAP_CHANNEL,
"datastream_handoff",
"ready",
json!({"from":"stdio_envelope","to":"cluster_datastream","channel":NODE_BOOTSTRAP_CHANNEL}),
)?;
if let Some(prompt) = &config.self_test_prompt { if let Some(prompt) = &config.self_test_prompt {
run_self_test( run_self_test(
@ -603,10 +912,11 @@ fn run() -> Result<(), String> {
); );
loop { loop {
pump_network(&mut driver, &stack); pump_network(&mut driver, &stack);
drain_debug_join_commands(&mut debug_join_rx, &mut driver, &config, &mut datastream);
datastream.tick(); datastream.tick();
worker.drain_stderr(&config, &mut datastream); worker.drain_stderr(&config, &mut datastream);
while let Some(report) = reports.try_recv() { while let Some(report) = reports.try_recv() {
handle_node_report( match handle_node_report(
report, report,
&config, &config,
&stack, &stack,
@ -614,7 +924,72 @@ fn run() -> Result<(), String> {
node_actor, node_actor,
&mut worker, &mut worker,
&mut datastream, &mut datastream,
)?; )? {
NodeReportOutcome::None => {}
NodeReportOutcome::RuntimeReadyAck {
run_id,
node_id,
stage_index,
readiness_id,
} => {
if pending_runtime_ready.observe_ack(run_id, node_id, stage_index, readiness_id)
{
emit_node_event(
&mut datastream,
&config,
NODE_BOOTSTRAP_CHANNEL,
"runtime_ready_ack",
"ready",
json!({
"readiness_id":readiness_id,
"attempts":pending_runtime_ready.attempts,
"endpoint":&pending_runtime_ready.endpoint,
"node_actor":pending_runtime_ready.node_actor,
}),
);
datastream.submit_text(ChannelId::new("mvp.node.ready"), ready.to_string());
emit_node_event(
&mut datastream,
&config,
NODE_BOOTSTRAP_CHANNEL,
"datastream_handoff",
"ready",
json!({"from":"runtime_ready_ack","to":"cluster_datastream","channel":"mvp.node.ready"}),
);
}
}
}
}
if !pending_runtime_ready.swim_logged && pending_runtime_ready.swim_ready(&stack) {
emit_node_event(
&mut datastream,
&config,
NODE_RUNTIME_CHANNEL,
"coordinator_swim",
"ready",
json!({
"coordinator":pending_runtime_ready
.coordinator
.map(|node| format!("{node:?}"))
.unwrap_or_else(|| "standalone".to_owned()),
"readiness_id":pending_runtime_ready.readiness_id,
}),
);
pending_runtime_ready.swim_logged = true;
}
if !pending_runtime_ready.acked && pending_runtime_ready.maybe_send(&stack, node_actor)? {
emit_node_event(
&mut datastream,
&config,
NODE_RUNTIME_CHANNEL,
"runtime_ready_signal",
"sent",
json!({
"readiness_id":pending_runtime_ready.readiness_id,
"attempts":pending_runtime_ready.attempts,
"next_backoff_ms":pending_runtime_ready.backoff.as_millis(),
}),
);
} }
if shutdown_rx.try_recv().is_ok() { if shutdown_rx.try_recv().is_ok() {
emit_node_event( emit_node_event(
@ -758,6 +1133,115 @@ impl FrameSink for JsonlFrameSink {
} }
} }
enum NodeReportOutcome {
None,
RuntimeReadyAck {
run_id: u64,
node_id: u64,
stage_index: u32,
readiness_id: u64,
},
}
struct PendingRuntimeReady {
run_id: u64,
node_id: u64,
stage_index: u32,
endpoint: EndpointAddr,
node_actor: ActorAddress,
coordinator: Option<DistNodeId>,
readiness_id: u64,
attempts: u32,
next_attempt_at: Instant,
backoff: Duration,
acked: bool,
swim_logged: bool,
}
impl PendingRuntimeReady {
fn new(config: &DeploymentConfig, endpoint: EndpointAddr, node_actor: ActorAddress) -> Self {
Self {
run_id: config.run_id,
node_id: config.logical_node_id,
stage_index: config.stage_index,
endpoint,
node_actor,
coordinator: config
.coordinator_endpoint
.as_ref()
.map(|endpoint| DistNodeId(*endpoint.id.as_bytes())),
readiness_id: 1,
attempts: 0,
next_attempt_at: Instant::now(),
backoff: RUNTIME_READY_RETRY_INITIAL,
acked: false,
swim_logged: false,
}
}
fn swim_ready(&self, stack: &DistributionRuntimeStack) -> bool {
let Some(coordinator) = self.coordinator else {
return true;
};
stack.member_state(coordinator) == Some(MemberState::Alive)
}
fn observe_ack(
&mut self,
run_id: u64,
node_id: u64,
stage_index: u32,
readiness_id: u64,
) -> bool {
if self.acked
|| self.run_id != run_id
|| self.node_id != node_id
|| self.stage_index != stage_index
|| self.readiness_id != readiness_id
{
return false;
}
self.acked = true;
true
}
fn maybe_send(
&mut self,
stack: &DistributionRuntimeStack,
node_actor: ActorAddress,
) -> Result<bool, String> {
if self.acked || !self.swim_ready(stack) {
return Ok(false);
}
let now = Instant::now();
if now < self.next_attempt_at {
return Ok(false);
}
stack
.runtime
.send_to(
node_actor,
NodeAgentMsg::RuntimeLoaded {
run_id: self.run_id,
node_id: self.node_id,
stage_index: self.stage_index,
endpoint: self.endpoint.clone(),
node_actor: self.node_actor,
readiness_id: self.readiness_id,
},
)
.map_err(|error| format!("signal runtime loaded: {error}"))?;
self.attempts = self.attempts.saturating_add(1);
self.next_attempt_at = now + self.backoff;
self.backoff = self
.backoff
.checked_mul(2)
.unwrap_or(RUNTIME_READY_RETRY_MAX)
.min(RUNTIME_READY_RETRY_MAX);
Ok(true)
}
}
fn handle_node_report( fn handle_node_report(
report: NodeAgentReport, report: NodeAgentReport,
config: &DeploymentConfig, config: &DeploymentConfig,
@ -766,11 +1250,12 @@ fn handle_node_report(
node_actor: ActorAddress, node_actor: ActorAddress,
worker: &mut TinygradWorker, worker: &mut TinygradWorker,
datastream: &mut DatastreamEmitter, datastream: &mut DatastreamEmitter,
) -> Result<(), String> { ) -> Result<NodeReportOutcome, String> {
let kind = match &report { let kind = match &report {
NodeAgentReport::Command(_) => "Command", NodeAgentReport::Command(_) => "Command",
NodeAgentReport::Lifecycle(_) => "Lifecycle", NodeAgentReport::Lifecycle(_) => "Lifecycle",
NodeAgentReport::PromptRequested { .. } => "PromptRequested", NodeAgentReport::PromptRequested { .. } => "PromptRequested",
NodeAgentReport::RuntimeReadyAck { .. } => "RuntimeReadyAck",
NodeAgentReport::Snapshot { .. } => "Snapshot", NodeAgentReport::Snapshot { .. } => "Snapshot",
}; };
emit_node_event( emit_node_event(
@ -782,9 +1267,12 @@ fn handle_node_report(
json!({"kind":kind}), json!({"kind":kind}),
); );
match report { match report {
NodeAgentReport::Command(command) => handle_stage_command( NodeAgentReport::Command(command) => {
handle_stage_command(
command, config, stack, driver, node_actor, worker, datastream, command, config, stack, driver, node_actor, worker, datastream,
), )?;
Ok(NodeReportOutcome::None)
}
NodeAgentReport::Lifecycle(event) => { NodeAgentReport::Lifecycle(event) => {
let event = format!("{event:?}"); let event = format!("{event:?}");
datastream.submit_text( datastream.submit_text(
@ -799,17 +1287,31 @@ fn handle_node_report(
"observed", "observed",
json!({"event":event}), json!({"event":event}),
); );
Ok(()) Ok(NodeReportOutcome::None)
} }
NodeAgentReport::PromptRequested { NodeAgentReport::PromptRequested {
request_id, request_id,
prompt, prompt,
max_tokens, max_tokens,
reply_to, reply_to,
} => handle_prompt_request( } => {
handle_prompt_request(
request_id, prompt, max_tokens, reply_to, config, stack, driver, worker, datastream, request_id, prompt, max_tokens, reply_to, config, stack, driver, worker, datastream,
), )?;
NodeAgentReport::Snapshot { .. } => Ok(()), Ok(NodeReportOutcome::None)
}
NodeAgentReport::RuntimeReadyAck {
run_id,
node_id,
stage_index,
readiness_id,
} => Ok(NodeReportOutcome::RuntimeReadyAck {
run_id,
node_id,
stage_index,
readiness_id,
}),
NodeAgentReport::Snapshot { .. } => Ok(NodeReportOutcome::None),
} }
} }
@ -1267,6 +1769,7 @@ struct DeploymentConfig {
orchestrator_actor: Option<ActorAddress>, orchestrator_actor: Option<ActorAddress>,
datastream_sink_actor: Option<ActorAddress>, datastream_sink_actor: Option<ActorAddress>,
datastream_frame_log: Option<String>, datastream_frame_log: Option<String>,
debug_join_socket: Option<String>,
relay_mode: iroh::RelayMode, relay_mode: iroh::RelayMode,
worker_script: String, worker_script: String,
device: String, device: String,
@ -1283,15 +1786,29 @@ struct DeploymentConfig {
impl DeploymentConfig { impl DeploymentConfig {
fn from_env() -> Result<Self, String> { fn from_env() -> Result<Self, String> {
let run_id = env_u64("MVP_RUN_ID", 1)?; let run_id = env_u64("MVP_RUN_ID", 1)?;
let logical_node_id = env_u64("MVP_LOGICAL_NODE_ID", 1)?;
let relay = relay_runtime_config_from_env(run_id)?; let relay = relay_runtime_config_from_env(run_id)?;
let debug_join_socket = match env_optional("MVP_DEBUG_JOIN_SOCKET").as_deref() {
Some("disabled") => None,
Some(path) => Some(path.to_owned()),
None => Some(
std::env::temp_dir()
.join(format!(
"mvp-node-debug-join-{run_id}-{logical_node_id}.sock"
))
.to_string_lossy()
.into_owned(),
),
};
Ok(Self { Ok(Self {
run_id, run_id,
logical_node_id: env_u64("MVP_LOGICAL_NODE_ID", 1)?, logical_node_id,
stage_index: env_u32("MVP_STAGE_INDEX", 0)?, stage_index: env_u32("MVP_STAGE_INDEX", 0)?,
coordinator_endpoint: env_json("MVP_COORDINATOR_ENDPOINT")?, coordinator_endpoint: env_json("MVP_COORDINATOR_ENDPOINT")?,
orchestrator_actor: env_json("MVP_ORCHESTRATOR_ACTOR")?, orchestrator_actor: env_json("MVP_ORCHESTRATOR_ACTOR")?,
datastream_sink_actor: env_json("MVP_DATASTREAM_SINK_ACTOR")?, datastream_sink_actor: env_json("MVP_DATASTREAM_SINK_ACTOR")?,
datastream_frame_log: env_optional("MVP_DATASTREAM_FRAME_LOG"), datastream_frame_log: env_optional("MVP_DATASTREAM_FRAME_LOG"),
debug_join_socket,
relay_mode: relay.mode, relay_mode: relay.mode,
worker_script: env_string("MVP_TINYGRAD_WORKER", DEFAULT_WORKER_SCRIPT), worker_script: env_string("MVP_TINYGRAD_WORKER", DEFAULT_WORKER_SCRIPT),
device: env_string("DEV", DEFAULT_DEVICE), device: env_string("DEV", DEFAULT_DEVICE),
@ -1743,3 +2260,212 @@ fn tokenizer_from_env() -> TokenizerSource {
.map(TokenizerSource::LocalPath) .map(TokenizerSource::LocalPath)
.unwrap_or(TokenizerSource::EmbeddedGguf) .unwrap_or(TokenizerSource::EmbeddedGguf)
} }
#[cfg(test)]
mod tests {
use super::*;
use distribution::swim::actor::MembershipChanged;
use mvp_system::actors::orchestrator::OrchestratorMsg;
fn endpoint(seed: u8) -> EndpointAddr {
EndpointAddr::new(iroh::SecretKey::from_bytes(&[seed; 32]).public())
}
fn test_config(coordinator_endpoint: Option<EndpointAddr>) -> DeploymentConfig {
DeploymentConfig {
run_id: 7,
logical_node_id: 11,
stage_index: 3,
coordinator_endpoint,
orchestrator_actor: Some(ActorAddress::new_random()),
datastream_sink_actor: None,
datastream_frame_log: None,
debug_join_socket: None,
relay_mode: iroh::RelayMode::Disabled,
worker_script: DEFAULT_WORKER_SCRIPT.to_owned(),
device: DEFAULT_DEVICE.to_owned(),
model_id: DEFAULT_MODEL_ID.to_owned(),
gguf_source: GgufSource::LocalPath("/tmp/model.gguf".to_owned()),
tokenizer: TokenizerSource::EmbeddedGguf,
self_test_prompt: None,
self_test_layer_end: 16,
self_test_max_tokens: 1,
arena_bytes: DEFAULT_ARENA_BYTES,
arena_alignment: DEFAULT_ARENA_ALIGNMENT,
}
}
fn test_stack() -> DistributionRuntimeStack {
DistributionRuntimeStack::new(DistNodeId([1; 32]), DistributedNodeConfig::default())
}
#[test]
fn debug_join_client_serializes_endpoint_from_stdin() {
let secret = iroh::SecretKey::from_bytes(&[7; 32]);
let endpoint = EndpointAddr::new(secret.public()).with_relay_url(
"http://relay.example.com"
.parse::<iroh::RelayUrl>()
.unwrap(),
);
let line = debug_join_request_line(endpoint).expect("serialize debug join request");
let request: DebugJoinRequestWire =
serde_json::from_str(&line).expect("deserialize debug join request");
match request {
DebugJoinRequestWire::JoinEndpoint { endpoint } => {
assert_eq!(
endpoint.relay_urls().next().map(ToString::to_string),
Some("http://relay.example.com/".to_owned())
);
}
}
}
#[test]
fn debug_join_listener_queues_join_endpoint() {
let root = std::env::temp_dir().join(format!(
"mvp-worker-debug-join-test-{}-{}",
std::process::id(),
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.expect("system time after epoch")
.as_nanos()
));
std::fs::create_dir(&root).expect("create debug join test temp dir");
let socket_path = root.join("debug-join.sock");
let runtime = tokio::runtime::Runtime::new().expect("create tokio runtime");
let mut commands = spawn_debug_join_listener(runtime.handle().clone(), socket_path.clone())
.expect("spawn debug join listener");
let secret = iroh::SecretKey::from_bytes(&[8; 32]);
let endpoint = EndpointAddr::new(secret.public()).with_relay_url(
"http://relay.example.com"
.parse::<iroh::RelayUrl>()
.unwrap(),
);
let request = DebugJoinRequestWire::JoinEndpoint { endpoint };
let mut request_line =
serde_json::to_string(&request).expect("serialize debug join request");
request_line.push('\n');
runtime.block_on(async {
let mut stream = tokio::net::UnixStream::connect(&socket_path)
.await
.expect("connect to debug join listener");
stream
.write_all(request_line.as_bytes())
.await
.expect("write debug join request");
stream.flush().await.expect("flush debug join request");
let DebugJoinCommand::JoinEndpoint { endpoint, reply } =
commands.recv().await.expect("receive debug join command");
assert_eq!(
endpoint.relay_urls().next().map(ToString::to_string),
Some("http://relay.example.com/".to_owned())
);
let peer_node_id = endpoint.id.to_string();
assert!(
reply
.send(DebugJoinResponseWire::JoinQueued {
peer_node_id,
has_relay: true,
direct_addr_count: 0,
})
.is_ok()
);
let mut reader = tokio::io::BufReader::new(stream);
let mut response_line = String::new();
reader
.read_line(&mut response_line)
.await
.expect("read debug join response");
let response: DebugJoinResponseWire =
serde_json::from_str(&response_line).expect("deserialize debug join response");
match response {
DebugJoinResponseWire::JoinQueued { has_relay, .. } => {
assert!(has_relay);
}
DebugJoinResponseWire::JoinRejected { error, detail } => {
panic!("debug join was rejected: {error}: {detail}");
}
}
});
drop(commands);
let _ = std::fs::remove_file(&socket_path);
std::fs::remove_dir(&root).expect("remove debug join test temp dir");
}
#[test]
fn runtime_ready_retry_waits_for_swim() {
let stack = test_stack();
let node_actor = ActorAddress::new_random();
let mut pending = PendingRuntimeReady::new(&test_config(None), endpoint(3), node_actor);
pending.coordinator = Some(DistNodeId([2; 32]));
assert_eq!(pending.maybe_send(&stack, node_actor), Ok(false));
assert_eq!(pending.attempts, 0);
}
#[test]
fn runtime_ready_retry_stops_after_matching_ack() {
let stack = test_stack();
let node_actor = ActorAddress::new_random();
let mut pending = PendingRuntimeReady::new(&test_config(None), endpoint(4), node_actor);
assert!(pending.observe_ack(
pending.run_id,
pending.node_id,
pending.stage_index,
pending.readiness_id,
));
assert!(pending.acked);
assert_eq!(pending.maybe_send(&stack, node_actor), Ok(false));
}
#[test]
fn runtime_ready_retry_backoff_caps() {
let stack = test_stack();
let orchestrator_inbox = stack
.runtime
.new_inbox::<OrchestratorMsg>()
.expect("orchestrator inbox");
let node_actor = stack
.runtime
.spawn(NodeAgentActor::new(
stage::NodeId(11),
*orchestrator_inbox.addr(),
None,
))
.expect("spawn node agent");
let coordinator = DistNodeId([2; 32]);
stack
.runtime
.send_to(
stack.actors.membership_fanout,
MembershipChanged {
node_id: coordinator,
state: MemberState::Alive,
incarnation: 1,
},
)
.expect("send membership change");
stack.pump_runtime_once();
let mut pending = PendingRuntimeReady::new(&test_config(None), endpoint(5), node_actor);
pending.coordinator = Some(coordinator);
for expected_attempts in 1..=4 {
pending.next_attempt_at = Instant::now();
assert!(
pending
.maybe_send(&stack, node_actor)
.expect("runtime ready send")
);
assert_eq!(pending.attempts, expected_attempts);
assert!(pending.backoff <= RUNTIME_READY_RETRY_MAX);
}
}
}

View file

@ -204,6 +204,18 @@ impl DistributionRuntimeStack {
.filter(|entry| entry.state == MemberState::Alive) .filter(|entry| entry.state == MemberState::Alive)
.count() .count()
} }
pub fn member_state(&self, node_id: NodeId) -> Option<MemberState> {
self.membership_mirror
.lock()
.ok()?
.get(&node_id)
.map(|entry| entry.state)
}
pub fn route_owner(&self, actor: ActorAddress) -> Option<NodeId> {
self.route_view.read().ok()?.get(&actor).copied()
}
} }
struct MembershipFanout { struct MembershipFanout {
@ -227,3 +239,41 @@ impl ActorInterface for MembershipFanout {
let _ = ctx.send(self.directory, DirectoryIn::Membership(change)); let _ = ctx.send(self.directory, DirectoryIn::Membership(change));
} }
} }
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn distribution_stack_reports_member_state_and_route_owner() {
let stack = DistributionRuntimeStack::new(
distribution::types::NodeId([1; 32]),
DistributedNodeConfig::default(),
);
let remote = distribution::types::NodeId([2; 32]);
stack
.runtime
.send_to(
stack.actors.membership_fanout,
distribution::swim::actor::MembershipChanged {
node_id: remote,
state: MemberState::Alive,
incarnation: 1,
},
)
.expect("send membership change");
stack.pump_runtime_once();
assert_eq!(stack.member_state(remote), Some(MemberState::Alive));
let actor = ActorAddress::new_random();
{
let mut route_view = match stack.route_view.write() {
Ok(route_view) => route_view,
Err(poisoned) => poisoned.into_inner(),
};
route_view.insert(actor, remote);
}
assert_eq!(stack.route_owner(actor), Some(remote));
}
}

View file

@ -387,7 +387,7 @@ pub trait VastAiBootstrapLauncher: Send {
#[derive(Clone)] #[derive(Clone)]
enum SshBootstrapMsg { enum SshBootstrapMsg {
Stop { reason: BootstrapStopReason }, Stop,
} }
struct SshBootstrapActor { struct SshBootstrapActor {
@ -407,7 +407,7 @@ impl ActorInterface for SshBootstrapActor {
fn handle(&mut self, _ctx: &Ctx, msg: Self::Incoming) { fn handle(&mut self, _ctx: &Ctx, msg: Self::Incoming) {
match msg { match msg {
SshBootstrapMsg::Stop { reason: _ } => { SshBootstrapMsg::Stop => {
self.stopping.store(true, Ordering::SeqCst); self.stopping.store(true, Ordering::SeqCst);
stop_ssh_child(&self.child); stop_ssh_child(&self.child);
} }
@ -481,10 +481,8 @@ impl VastAiBootstrapLauncher for SshCommandBootstrapLauncher {
}) })
} }
fn stop_bootstrap(&mut self, handle: &mut Self::Handle, reason: BootstrapStopReason) { fn stop_bootstrap(&mut self, handle: &mut Self::Handle, _reason: BootstrapStopReason) {
let _ = handle let _ = handle.runtime.send_to(handle.actor, SshBootstrapMsg::Stop);
.runtime
.send_to(handle.actor, SshBootstrapMsg::Stop { reason });
handle.runtime.tick(); handle.runtime.tick();
} }
} }
@ -1051,12 +1049,7 @@ mod tests {
.expect("spawn ssh bootstrap actor"); .expect("spawn ssh bootstrap actor");
runtime runtime
.send_to( .send_to(actor, SshBootstrapMsg::Stop)
actor,
SshBootstrapMsg::Stop {
reason: BootstrapStopReason::RuntimeReady,
},
)
.expect("send stop"); .expect("send stop");
runtime.tick(); runtime.tick();

View file

@ -1,15 +1,20 @@
use std::path::{Path, PathBuf}; #![recursion_limit = "256"]
use std::path::Path;
use std::process::{Command, ExitCode}; use std::process::{Command, ExitCode};
use std::time::{Duration, Instant};
#[path = "support/local_e2e_cluster.rs"] #[path = "support/local_e2e_cluster.rs"]
mod local_e2e_cluster; mod local_e2e_cluster;
const IMAGE: &str = "swactor-mvp-local-e2e-cluster:latest"; const IMAGE: &str = "swactor-mvp-local-e2e-cluster:latest";
const SKIP_BUILD_ENV: &str = "MVP_LOCAL_E2E_CLUSTER_SKIP_BUILD";
const BUILD_ONLY_ENV: &str = "MVP_LOCAL_E2E_CLUSTER_BUILD_ONLY";
fn main() -> ExitCode { fn main() -> ExitCode {
let args = std::env::args().collect::<Vec<_>>(); let args = std::env::args().collect::<Vec<_>>();
match std::env::var("MVP_TEST_ROLE").ok().as_deref() { match std::env::var("MVP_TEST_ROLE").ok().as_deref() {
Some("cluster-supervisor") => return local_e2e_cluster::run_main(), Some("cluster-supervisor" | "cluster-relay") => return local_e2e_cluster::run_main(),
Some(role) => { Some(role) => {
eprintln!("unknown MVP_TEST_ROLE={role}"); eprintln!("unknown MVP_TEST_ROLE={role}");
return ExitCode::from(2); return ExitCode::from(2);
@ -30,16 +35,20 @@ fn local_e2e_cluster_docker_cpu_pipeline_prompt() {
eprintln!("skipping; set MVP_SYSTEM_LOCAL_E2E_CLUSTER=1 to run Docker CPU cluster e2e"); eprintln!("skipping; set MVP_SYSTEM_LOCAL_E2E_CLUSTER=1 to run Docker CPU cluster e2e");
return; return;
} }
if !Path::new("/var/run/docker.sock").exists() {
eprintln!(
"skipping; /var/run/docker.sock is required for the relay-only Docker cluster e2e"
);
return;
}
build_docker_fixture(); build_docker_fixture();
if std::env::var_os(BUILD_ONLY_ENV).is_some() {
return;
}
let output = Command::new(current_test_exe()) let docker = DockerRelayFixture::start();
.env("MVP_TEST_ROLE", "cluster-supervisor") let output = docker.run_supervisor("ping");
.arg("--prompt")
.arg("ping")
.env("MVP_LOCAL_E2E_CLUSTER_IMAGE", IMAGE)
.output()
.expect("run local e2e cluster supervisor");
assert!( assert!(
output.status.success(), output.status.success(),
@ -95,6 +104,17 @@ fn local_e2e_cluster_docker_cpu_pipeline_prompt() {
"{value}" "{value}"
); );
assert_eq!(value["provision_nodes_stopped"], true); assert_eq!(value["provision_nodes_stopped"], true);
assert_eq!(value["relay_only"], true, "{value}");
assert_eq!(value["relay_url"], "http://relay:7843/");
assert_eq!(value["orchestrator_endpoint_has_relay"], true, "{value}");
assert_eq!(
value["orchestrator_endpoint_relay_url"],
"http://relay:7843/"
);
assert_eq!(value["node0_endpoint_has_relay"], true, "{value}");
assert_eq!(value["node0_endpoint_relay_url"], "http://relay:7843/");
assert_eq!(value["node1_endpoint_has_relay"], true, "{value}");
assert_eq!(value["node1_endpoint_relay_url"], "http://relay:7843/");
assert!( assert!(
value["node0_endpoint"]["addrs"] value["node0_endpoint"]["addrs"]
.as_array() .as_array()
@ -108,8 +128,166 @@ fn local_e2e_cluster_docker_cpu_pipeline_prompt() {
"{value}" "{value}"
); );
} }
struct DockerRelayFixture {
relay_container: String,
supervisor_network: String,
node0_network: String,
node1_network: String,
}
impl DockerRelayFixture {
fn start() -> Self {
let suffix = format!("{}-{}", std::process::id(), unique_nanos());
let relay_container = format!("mvp-local-e2e-relay-{suffix}");
let supervisor_network = format!("mvp-local-e2e-supervisor-{suffix}");
let node0_network = format!("mvp-local-e2e-node0-{suffix}");
let node1_network = format!("mvp-local-e2e-node1-{suffix}");
for network in [&supervisor_network, &node0_network, &node1_network] {
docker_status(
["network", "create", network],
"create relay-only Docker network",
);
}
docker_status(
[
"run",
"-d",
"--rm",
"--name",
&relay_container,
"--network",
&supervisor_network,
"--network-alias",
"relay",
"-e",
"MVP_TEST_ROLE=cluster-relay",
"-e",
"MVP_LOCAL_E2E_RELAY_LISTEN=0.0.0.0:7843",
IMAGE,
],
"start relay sidecar",
);
docker_status(
[
"network",
"connect",
"--alias",
"relay",
&node0_network,
&relay_container,
],
"attach relay to node0 network",
);
docker_status(
[
"network",
"connect",
"--alias",
"relay",
&node1_network,
&relay_container,
],
"attach relay to node1 network",
);
let fixture = Self {
relay_container,
supervisor_network,
node0_network,
node1_network,
};
fixture.wait_for_relay();
fixture
}
fn run_supervisor(&self, prompt: &str) -> std::process::Output {
Command::new("docker")
.args([
"run",
"--rm",
"--name",
&format!("mvp-local-e2e-supervisor-{}", unique_nanos()),
"--network",
&self.supervisor_network,
"--network-alias",
"supervisor",
"-v",
"/var/run/docker.sock:/var/run/docker.sock",
"-e",
"MVP_TEST_ROLE=cluster-supervisor",
"-e",
&format!("MVP_LOCAL_E2E_CLUSTER_IMAGE={IMAGE}"),
"-e",
&format!("MVP_LOCAL_E2E_DOCKER_NETWORK_NODE0={}", self.node0_network),
"-e",
&format!("MVP_LOCAL_E2E_DOCKER_NETWORK_NODE1={}", self.node1_network),
"-e",
"MVP_IROH_RELAY_MODE=default",
"-e",
"MVP_IROH_RELAY_URL=http://relay:7843/",
IMAGE,
"--prompt",
prompt,
])
.output()
.expect("run relay-only local e2e cluster supervisor")
}
fn wait_for_relay(&self) {
let started = Instant::now();
while started.elapsed() < Duration::from_secs(20) {
let logs = Command::new("docker")
.args(["logs", &self.relay_container])
.output()
.expect("read relay sidecar logs");
let stdout = String::from_utf8_lossy(&logs.stdout);
let stderr = String::from_utf8_lossy(&logs.stderr);
if stdout.contains("relay ready") || stderr.contains("relay ready") {
return;
}
std::thread::sleep(Duration::from_millis(100));
}
panic!("relay sidecar did not report ready within 20s");
}
}
impl Drop for DockerRelayFixture {
fn drop(&mut self) {
let _ = Command::new("docker")
.args(["stop", "-t", "2", &self.relay_container])
.status();
for network in [
&self.node1_network,
&self.node0_network,
&self.supervisor_network,
] {
let _ = Command::new("docker")
.args(["network", "rm", network])
.status();
}
}
}
fn docker_status<const N: usize>(args: [&str; N], action: &str) {
let status = Command::new("docker")
.args(args)
.status()
.unwrap_or_else(|error| panic!("{action}: {error}"));
assert!(status.success(), "{action} failed with status {status}");
}
fn unique_nanos() -> u128 {
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.expect("system time after epoch")
.as_nanos()
}
fn build_docker_fixture() { fn build_docker_fixture() {
if std::env::var_os(SKIP_BUILD_ENV).is_some() {
phase("using existing Docker CPU cluster fixture image");
return;
}
let crate_dir = Path::new(env!("CARGO_MANIFEST_DIR")); let crate_dir = Path::new(env!("CARGO_MANIFEST_DIR"));
let workspace = crate_dir let workspace = crate_dir
.parent() .parent()
@ -176,7 +354,3 @@ fn copy_context_entry(source: &Path, dest: &Path) {
fn phase(message: &str) { fn phase(message: &str) {
eprintln!("local-e2e-cluster: {message}"); eprintln!("local-e2e-cluster: {message}");
} }
fn current_test_exe() -> PathBuf {
std::env::current_exe().expect("current test exe")
}

View file

@ -3,6 +3,7 @@ FROM rust:1-bookworm
RUN apt-get update && \ RUN apt-get update && \
DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \ DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \
ca-certificates \ ca-certificates \
docker.io \
pkg-config \ pkg-config \
python3 \ python3 \
python3-pip && \ python3-pip && \
@ -12,9 +13,12 @@ RUN apt-get update && \
ENV DEV=CPU ENV DEV=CPU
ENV PYTHONDONTWRITEBYTECODE=1 ENV PYTHONDONTWRITEBYTECODE=1
ENV CARGO_TARGET_DIR=/workspace/target ENV CARGO_TARGET_DIR=/tmp/mvp-local-e2e-target
COPY . /workspace COPY . /workspace
WORKDIR /workspace WORKDIR /workspace
RUN cargo test -p mvp-system --features local-e2e --test local-e2e-cluster --no-run RUN cargo test -p mvp-system --features local-e2e --test local-e2e-cluster --no-run && \
test_bin="$(find /tmp/mvp-local-e2e-target/debug/deps -maxdepth 1 -type f -perm /111 \( -name 'local_e2e_cluster-*' -o -name 'local-e2e-cluster-*' \) | head -n 1)" && \
cp "$test_bin" /usr/local/bin/local-e2e-cluster && \
rm -rf /tmp/mvp-local-e2e-target
ENTRYPOINT ["sh", "-c", "test_bin=$(find /workspace/target/debug/deps -maxdepth 1 -type f -perm /111 \\( -name 'local_e2e_cluster-*' -o -name 'local-e2e-cluster-*' \\) | head -n 1); exec \"$test_bin\" \"$@\"", "local-e2e-cluster"] ENTRYPOINT ["/usr/local/bin/local-e2e-cluster"]

View file

@ -14,12 +14,13 @@ const TEST_WATCHDOG: Duration = Duration::from_secs(1_800);
const PROMPT_WATCHDOG: Duration = Duration::from_secs(600); const PROMPT_WATCHDOG: Duration = Duration::from_secs(600);
const SHUTDOWN_WATCHDOG: Duration = Duration::from_secs(60); const SHUTDOWN_WATCHDOG: Duration = Duration::from_secs(60);
const DASHBOARD_ADDR: &str = "127.0.0.1:9090"; const DASHBOARD_ADDR: &str = "127.0.0.1:9090";
const DEFAULT_CONTAINER: &str = "mvp-orchestrator-1-1"; const DOCKER_CONTAINER_PREFIX_ENV: &str = "MVP_DOCKER_CONTAINER_PREFIX";
#[test] #[test]
fn one_node_chat_docker_cuda_e2e() { fn one_node_chat_docker_cuda_e2e() {
let root = workspace_root(); let root = workspace_root();
require_docker(&root); require_docker(&root);
let container_prefix = format!("mvp-orchestrator-e2e-{}", std::process::id());
let mut command = Command::new("cargo"); let mut command = Command::new("cargo");
command command
@ -27,6 +28,9 @@ fn one_node_chat_docker_cuda_e2e() {
.args(["mvp-chat"]) .args(["mvp-chat"])
.env("MVP_RUNTIME_CONFIG", "local") .env("MVP_RUNTIME_CONFIG", "local")
.env("MVP_IROH_RELAY_MODE", "disabled") .env("MVP_IROH_RELAY_MODE", "disabled")
.env(DOCKER_CONTAINER_PREFIX_ENV, &container_prefix)
.env("MVP_TINYGRAD_TEST_MODE", "1")
.env("MVP_PROMPT_MAX_TOKENS", "3")
.stdin(Stdio::piped()) .stdin(Stdio::piped())
.stdout(Stdio::piped()) .stdout(Stdio::piped())
.stderr(Stdio::piped()); .stderr(Stdio::piped());
@ -60,7 +64,7 @@ fn one_node_chat_docker_cuda_e2e() {
if result.is_ok() { if result.is_ok() {
result = assert_no_lower_layer_terminal_leaks(&stdout, &stderr); result = assert_no_lower_layer_terminal_leaks(&stdout, &stderr);
} }
assert_container_removed(&root, DEFAULT_CONTAINER); assert_containers_with_prefix_removed(&root, &container_prefix);
if let Err(error) = result { if let Err(error) = result {
panic!( panic!(
@ -84,9 +88,9 @@ fn run_full_flow(
}) })
.map_err(|e| format!("provisioning frames not visible in dashboard: {e}"))?; .map_err(|e| format!("provisioning frames not visible in dashboard: {e}"))?;
wait_for_child_or(TEST_WATCHDOG, child, || { wait_for_child_or(TEST_WATCHDOG, child, || {
dashboard_has_frame("mvp.worker.weights", "GgufDownloadProgress") dashboard_has_frame("mvp.worker.weights", "LoadWeightsStarted")
}) })
.map_err(|e| format!("GGUF download progress not visible in dashboard: {e}"))?; .map_err(|e| format!("weight loading start not visible in dashboard: {e}"))?;
wait_for_child_or(TEST_WATCHDOG, child, || { wait_for_child_or(TEST_WATCHDOG, child, || {
dashboard_has_frame("mvp.worker.weights", "WeightsLoaded") dashboard_has_frame("mvp.worker.weights", "WeightsLoaded")
}) })
@ -111,10 +115,13 @@ fn run_full_flow(
.map_err(|e| format!("orchestrator prompt lifecycle not visible in dashboard: {e}"))?; .map_err(|e| format!("orchestrator prompt lifecycle not visible in dashboard: {e}"))?;
wait_for_child_or(PROMPT_WATCHDOG, child, || prompt_count(stdout) >= 2) wait_for_child_or(PROMPT_WATCHDOG, child, || prompt_count(stdout) >= 2)
.map_err(|e| format!("chat prompt did not return after response: {e}"))?; .map_err(|e| format!("chat prompt did not return after response: {e}"))?;
request_child_interrupt(child); writeln!(stdin, "/exit").map_err(|e| format!("write exit command: {e}"))?;
stdin
.flush()
.map_err(|e| format!("flush exit command: {e}"))?;
let status = wait_child(child, SHUTDOWN_WATCHDOG) let status = wait_child(child, SHUTDOWN_WATCHDOG)
.ok_or_else(|| "cargo mvp-chat did not exit after Ctrl-C".to_owned())?; .ok_or_else(|| "cargo mvp-chat did not exit after /exit".to_owned())?;
if status.success() || status.code() == Some(130) || status.signal_name() == Some("SIGINT") { if status.success() {
Ok(()) Ok(())
} else { } else {
Err(format!("cargo mvp-chat exited with {status}")) Err(format!("cargo mvp-chat exited with {status}"))
@ -400,28 +407,6 @@ fn request_child_interrupt(child: &Child) {
} }
} }
trait ExitStatusSignalName {
fn signal_name(&self) -> Option<&'static str>;
}
impl ExitStatusSignalName for std::process::ExitStatus {
fn signal_name(&self) -> Option<&'static str> {
#[cfg(target_os = "linux")]
{
use std::os::unix::process::ExitStatusExt;
match self.signal() {
Some(libc::SIGINT) => Some("SIGINT"),
Some(libc::SIGTERM) => Some("SIGTERM"),
_ => None,
}
}
#[cfg(not(target_os = "linux"))]
{
let _ = self;
None
}
}
}
fn require_docker(root: &std::path::Path) { fn require_docker(root: &std::path::Path) {
let version = Command::new("docker") let version = Command::new("docker")
@ -437,14 +422,17 @@ fn require_docker(root: &std::path::Path) {
); );
} }
fn assert_container_removed(root: &std::path::Path, container: &str) { fn assert_containers_with_prefix_removed(root: &std::path::Path, prefix: &str) {
let start = Instant::now();
let mut containers = String::new();
while start.elapsed() < SHUTDOWN_WATCHDOG {
let output = Command::new("docker") let output = Command::new("docker")
.current_dir(root) .current_dir(root)
.args([ .args([
"ps", "ps",
"-a", "-a",
"--filter", "--filter",
&format!("name=^{container}$"), &format!("name=^{prefix}-"),
"--format", "--format",
"{{.Names}}", "{{.Names}}",
]) ])
@ -456,10 +444,13 @@ fn assert_container_removed(root: &std::path::Path, container: &str) {
String::from_utf8_lossy(&output.stdout), String::from_utf8_lossy(&output.stdout),
String::from_utf8_lossy(&output.stderr) String::from_utf8_lossy(&output.stderr)
); );
assert!( containers = String::from_utf8_lossy(&output.stdout).into_owned();
String::from_utf8_lossy(&output.stdout).trim().is_empty(), if containers.trim().is_empty() {
"container {container} still exists" return;
); }
thread::sleep(Duration::from_millis(100));
}
panic!("containers with prefix {prefix} still exist:\n{containers}");
} }
fn workspace_root() -> std::path::PathBuf { fn workspace_root() -> std::path::PathBuf {

View file

@ -1,5 +1,6 @@
use std::collections::{BTreeMap, HashMap, VecDeque}; use std::collections::{BTreeMap, HashMap, VecDeque};
use std::io::{BufRead, BufReader, Write}; use std::io::{BufRead, BufReader, Write};
use std::net::SocketAddr;
use std::path::PathBuf; use std::path::PathBuf;
use std::process::{Child, ChildStdin, ChildStdout, Command, ExitCode, Stdio}; use std::process::{Child, ChildStdin, ChildStdout, Command, ExitCode, Stdio};
use std::sync::{ use std::sync::{
@ -40,7 +41,8 @@ use mvp_system::observability_surface as obs;
use mvp_system::orchestrator_run_fsm as fsm; use mvp_system::orchestrator_run_fsm as fsm;
use mvp_system::provisioning::{NodeProvisionSpec, ProvisionLogStream}; use mvp_system::provisioning::{NodeProvisionSpec, ProvisionLogStream};
use mvp_system::relay_provisioning::{ use mvp_system::relay_provisioning::{
LocalShimRelayProvider, RelayProvider, RelayProvisionRequest, RelayPurpose, LocalShimRelayProvider, MVP_IROH_RELAY_URL_ENV, RelayProvider, RelayProvisionRequest,
RelayPurpose, relay_runtime_config_from_env,
}; };
use mvp_system::run_plan as plan; use mvp_system::run_plan as plan;
use mvp_system::stage_controller as stage; use mvp_system::stage_controller as stage;
@ -63,6 +65,11 @@ const OBJECT_ALIGNMENT: u64 = 4;
const ARENA_BYTES: usize = 16 * 1024; const ARENA_BYTES: usize = 16 * 1024;
const RING_BYTES: usize = 4096; const RING_BYTES: usize = 4096;
const DEFAULT_RUNTIME_SNAPSHOT_INTERVAL: Duration = Duration::from_millis(500); const DEFAULT_RUNTIME_SNAPSHOT_INTERVAL: Duration = Duration::from_millis(500);
const LOCAL_E2E_ROUTE_TIMEOUT: Duration = Duration::from_secs(60);
const LOCAL_E2E_WORKFLOW_TIMEOUT: Duration = Duration::from_secs(120);
const LOCAL_E2E_RELAY_LISTEN_ENV: &str = "MVP_LOCAL_E2E_RELAY_LISTEN";
const LOCAL_E2E_DOCKER_NETWORK_NODE0_ENV: &str = "MVP_LOCAL_E2E_DOCKER_NETWORK_NODE0";
const LOCAL_E2E_DOCKER_NETWORK_NODE1_ENV: &str = "MVP_LOCAL_E2E_DOCKER_NETWORK_NODE1";
static STOP_REQUESTED: AtomicBool = AtomicBool::new(false); static STOP_REQUESTED: AtomicBool = AtomicBool::new(false);
struct MvpDashboard { struct MvpDashboard {
@ -216,7 +223,9 @@ fn runtime_snapshot_interval_from_env() -> Result<Duration, String> {
pub fn run_main() -> ExitCode { pub fn run_main() -> ExitCode {
install_signal_handlers(); install_signal_handlers();
let args = std::env::args().collect::<Vec<_>>(); let args = std::env::args().collect::<Vec<_>>();
let result = if args.iter().any(|arg| arg == "--role=node") { let result = if std::env::var("MVP_TEST_ROLE").ok().as_deref() == Some("cluster-relay") {
run_relay_role()
} else if args.iter().any(|arg| arg == "--role=node") {
run_node_role(&args) run_node_role(&args)
} else if std::env::var_os("MVP_DASHBOARD").is_some() { } else if std::env::var_os("MVP_DASHBOARD").is_some() {
run_supervisor_dashboard_loop() run_supervisor_dashboard_loop()
@ -236,6 +245,53 @@ pub fn run_main() -> ExitCode {
} }
} }
struct LocalRelayGuard {
_server: iroh_relay::server::Server,
_rt: tokio::runtime::Runtime,
}
fn run_relay_role() -> Result<(), String> {
let listen =
std::env::var(LOCAL_E2E_RELAY_LISTEN_ENV).unwrap_or_else(|_| "0.0.0.0:7843".to_owned());
let _relay = spawn_local_relay(&listen)?;
eprintln!("mvp-local-e2e-cluster: relay ready on {listen}");
while !STOP_REQUESTED.load(Ordering::SeqCst) {
thread::sleep(Duration::from_millis(100));
}
Ok(())
}
fn spawn_local_relay(listen: &str) -> Result<LocalRelayGuard, String> {
let bind_addr = listen
.parse::<SocketAddr>()
.map_err(|e| format!("invalid {LOCAL_E2E_RELAY_LISTEN_ENV}={listen:?}: {e}"))?;
let rt = tokio::runtime::Builder::new_multi_thread()
.enable_all()
.worker_threads(2)
.build()
.map_err(|e| format!("relay tokio runtime: {e}"))?;
let server = rt
.block_on(async {
iroh_relay::server::Server::spawn(iroh_relay::server::ServerConfig::<(), ()> {
relay: Some(iroh_relay::server::RelayConfig {
http_bind_addr: bind_addr,
tls: None,
limits: Default::default(),
key_cache_capacity: Some(256),
access: iroh_relay::server::AccessConfig::Everyone,
}),
quic: None,
metrics_addr: None,
})
.await
})
.map_err(|e| format!("spawn relay server: {e}"))?;
Ok(LocalRelayGuard {
_server: server,
_rt: rt,
})
}
#[derive(Clone, Debug, Deserialize)] #[derive(Clone, Debug, Deserialize)]
struct NodeStdoutLine { struct NodeStdoutLine {
#[serde(rename = "type")] #[serde(rename = "type")]
@ -532,7 +588,8 @@ fn run_supervisor_once(
.cloned() .cloned()
.ok_or_else(|| "engine builder did not assign stage 1".to_owned())?; .ok_or_else(|| "engine builder did not assign stage 1".to_owned())?;
let self_endpoint_json = serde_json::to_string(&driver.endpoint_addr()) let self_endpoint = driver.endpoint_addr();
let self_endpoint_json = serde_json::to_string(&self_endpoint)
.map_err(|e| format!("serialize endpoint addr: {e}"))?; .map_err(|e| format!("serialize endpoint addr: {e}"))?;
let orchestrator_actor_json = serde_json::to_string(&orchestrator_addr) let orchestrator_actor_json = serde_json::to_string(&orchestrator_addr)
.map_err(|e| format!("serialize orchestrator actor: {e}"))?; .map_err(|e| format!("serialize orchestrator actor: {e}"))?;
@ -660,8 +717,22 @@ fn run_supervisor_once(
let mut response_tokens = Vec::<u32>::new(); let mut response_tokens = Vec::<u32>::new();
let mut edge_stream_count = 0usize; let mut edge_stream_count = 0usize;
let mut token_out_streams = HashMap::<u64, Vec<u8>>::new(); let mut token_out_streams = HashMap::<u64, Vec<u8>>::new();
let workflow_started = Instant::now();
while !STOP_REQUESTED.load(Ordering::SeqCst) { while !STOP_REQUESTED.load(Ordering::SeqCst) {
if workflow_started.elapsed() >= LOCAL_E2E_WORKFLOW_TIMEOUT {
let _ = stop_provisioned_nodes(
&mut [&mut node0, &mut node1],
&mut driver,
&stack,
&mut dashboard,
);
return Err(format!(
"workflow timed out after {:?}: injected={injected}, token_received={token_received}, completed={completed}, torn_down={torn_down}, stop_node0={sent_stop_to_node0}, stop_node1={sent_stop_to_node1}, stage_ready_count={stage_ready_count}, edge_stream_count={edge_stream_count}",
LOCAL_E2E_WORKFLOW_TIMEOUT
));
}
pump_network(&mut driver, &stack); pump_network(&mut driver, &stack);
driver_runtime.poll_iroh(&driver); driver_runtime.poll_iroh(&driver);
while let Some(event) = driver_runtime.try_recv() { while let Some(event) = driver_runtime.try_recv() {
@ -827,7 +898,8 @@ fn run_supervisor_once(
return Err(format!("run failed: {event:?}")); return Err(format!("run failed: {event:?}"));
} }
}, },
OrchestratorReport::Snapshot { .. } => {} OrchestratorReport::NodeRuntimeReady { .. }
| OrchestratorReport::Snapshot { .. } => {}
} }
} }
@ -858,6 +930,10 @@ fn run_supervisor_once(
}) })
.count(); .count();
let response_text = detokenize_response(&response_tokens); let response_text = detokenize_response(&response_tokens);
let relay_url = relay_url_from_env();
let orchestrator_endpoint_relay_url = endpoint_relay_url(&self_endpoint);
let node0_endpoint_relay_url = endpoint_relay_url(&node0.endpoint);
let node1_endpoint_relay_url = endpoint_relay_url(&node1.endpoint);
let summary = json!({ let summary = json!({
"ok": true, "ok": true,
"actor_plane": "iroh-swactor", "actor_plane": "iroh-swactor",
@ -871,6 +947,15 @@ fn run_supervisor_once(
}, },
"node0_endpoint": node0.endpoint, "node0_endpoint": node0.endpoint,
"node1_endpoint": node1.endpoint, "node1_endpoint": node1.endpoint,
"orchestrator_endpoint": self_endpoint,
"relay_only": relay_url.is_some(),
"relay_url": relay_url,
"orchestrator_endpoint_has_relay": orchestrator_endpoint_relay_url.is_some(),
"orchestrator_endpoint_relay_url": orchestrator_endpoint_relay_url,
"node0_endpoint_has_relay": node0_endpoint_relay_url.is_some(),
"node0_endpoint_relay_url": node0_endpoint_relay_url,
"node1_endpoint_has_relay": node1_endpoint_relay_url.is_some(),
"node1_endpoint_relay_url": node1_endpoint_relay_url,
"node0_logical_id": node0.node_id, "node0_logical_id": node0.node_id,
"node1_logical_id": node1.node_id, "node1_logical_id": node1.node_id,
"node0_stage_index": node0.stage_index, "node0_stage_index": node0.stage_index,
@ -1132,7 +1217,9 @@ fn run_node_role(args: &[String]) -> Result<(), String> {
.flush() .flush()
.map_err(|e| format!("flush lifecycle stdout: {e}"))?; .map_err(|e| format!("flush lifecycle stdout: {e}"))?;
} }
NodeAgentReport::PromptRequested { .. } | NodeAgentReport::Snapshot { .. } => {} NodeAgentReport::PromptRequested { .. }
| NodeAgentReport::RuntimeReadyAck { .. }
| NodeAgentReport::Snapshot { .. } => {}
} }
} }
@ -1203,12 +1290,16 @@ fn run_node_role(args: &[String]) -> Result<(), String> {
} }
fn new_driver(handle: tokio::runtime::Handle) -> Result<IrohDriver, String> { fn new_driver(handle: tokio::runtime::Handle) -> Result<IrohDriver, String> {
let relay_mode = if relay_url_from_env().is_some() {
relay_runtime_config_from_env(RUN_ID)?.mode
} else {
let mut relay_provider = LocalShimRelayProvider; let mut relay_provider = LocalShimRelayProvider;
let relay = relay_provider.provision_relay(RelayProvisionRequest { let relay = relay_provider.provision_relay(RelayProvisionRequest {
run_id: RUN_ID, run_id: RUN_ID,
purpose: RelayPurpose::Combined, purpose: RelayPurpose::Combined,
})?; })?;
let relay_mode = relay_provider.relay_mode(&relay)?; relay_provider.relay_mode(&relay)?
};
IrohDriver::with_handle( IrohDriver::with_handle(
handle, handle,
IrohDriverConfig { IrohDriverConfig {
@ -1234,8 +1325,14 @@ fn wait_for_routes(
stack: &DistributionRuntimeStack, stack: &DistributionRuntimeStack,
actors: &[ActorAddress], actors: &[ActorAddress],
) -> Result<(), String> { ) -> Result<(), String> {
let started = Instant::now();
loop { loop {
pump_network(driver, stack); pump_network(driver, stack);
let route_count = stack
.route_view
.read()
.map(|view| view.len())
.unwrap_or_default();
let ready = stack let ready = stack
.route_view .route_view
.read() .read()
@ -1244,6 +1341,13 @@ fn wait_for_routes(
if ready { if ready {
return Ok(()); return Ok(());
} }
if started.elapsed() >= LOCAL_E2E_ROUTE_TIMEOUT {
return Err(format!(
"routes not ready after {:?}: expected {} actor routes, observed {route_count}",
LOCAL_E2E_ROUTE_TIMEOUT,
actors.len()
));
}
thread::sleep(Duration::from_millis(20)); thread::sleep(Duration::from_millis(20));
} }
} }
@ -2239,12 +2343,7 @@ fn local_docker_spec(
coordinator_endpoint_json: &str, coordinator_endpoint_json: &str,
orchestrator_actor_json: &str, orchestrator_actor_json: &str,
) -> NodeProvisionSpec { ) -> NodeProvisionSpec {
NodeProvisionSpec { let mut env = vec![
run_id,
node_id,
stage_index: Some(stage_index),
image: docker_image(),
env: vec![
("DEV".to_owned(), "CPU".to_owned()), ("DEV".to_owned(), "CPU".to_owned()),
("PYTHONDONTWRITEBYTECODE".to_owned(), "1".to_owned()), ("PYTHONDONTWRITEBYTECODE".to_owned(), "1".to_owned()),
( (
@ -2252,7 +2351,18 @@ fn local_docker_spec(
"/workspace/crates/mvp-system/tests/local_e2e_cluster/tinygrad_cpu_worker.py" "/workspace/crates/mvp-system/tests/local_e2e_cluster/tinygrad_cpu_worker.py"
.to_owned(), .to_owned(),
), ),
], ];
if let Some(relay_url) = relay_url_from_env() {
env.push(("MVP_IROH_RELAY_MODE".to_owned(), "default".to_owned()));
env.push((MVP_IROH_RELAY_URL_ENV.to_owned(), relay_url));
}
NodeProvisionSpec {
run_id,
node_id,
stage_index: Some(stage_index),
image: docker_image(),
env,
args: vec![ args: vec![
"--role=node".to_owned(), "--role=node".to_owned(),
"--logical-node-id".to_owned(), "--logical-node-id".to_owned(),
@ -2308,14 +2418,15 @@ impl docker_provision::DockerCli for LocalE2eDockerCli {
} }
let mut command = Command::new("docker"); let mut command = Command::new("docker");
command.arg("run").arg("--rm");
if let Some(network) = docker_network_for_node(self.spec.node_id) {
command.arg("--network").arg(network);
} else {
command command
.arg("run")
.arg("--rm")
.arg("--add-host") .arg("--add-host")
.arg("host.docker.internal:host-gateway") .arg("host.docker.internal:host-gateway");
.arg("--name") }
.arg(&request.container_name) command.arg("--name").arg(&request.container_name).arg("-i");
.arg("-i");
for (key, value) in &request.labels { for (key, value) in &request.labels {
command.arg("--label").arg(format!("{key}={value}")); command.arg("--label").arg(format!("{key}={value}"));
} }
@ -2939,6 +3050,29 @@ fn spawn_shutdown_listener() -> Receiver<()> {
rx rx
} }
fn relay_url_from_env() -> Option<String> {
std::env::var(MVP_IROH_RELAY_URL_ENV)
.ok()
.map(|value| value.trim().to_owned())
.filter(|value| !value.is_empty())
}
fn endpoint_relay_url(endpoint: &EndpointAddr) -> Option<String> {
endpoint.relay_urls().next().map(|url| url.to_string())
}
fn docker_network_for_node(node_id: u64) -> Option<String> {
let env_name = match node_id {
NODE0_LOGICAL_ID => LOCAL_E2E_DOCKER_NETWORK_NODE0_ENV,
NODE1_LOGICAL_ID => LOCAL_E2E_DOCKER_NETWORK_NODE1_ENV,
_ => return None,
};
std::env::var(env_name)
.ok()
.map(|value| value.trim().to_owned())
.filter(|value| !value.is_empty())
}
fn parse_arg<'a>(args: &'a [String], name: &str) -> Result<&'a str, String> { fn parse_arg<'a>(args: &'a [String], name: &str) -> Result<&'a str, String> {
let index = args let index = args
.iter() .iter()

View file

@ -443,6 +443,22 @@ impl<A: ActorInterface> Actor<A> {
pub fn new(inner: A) -> Self { pub fn new(inner: A) -> Self {
Self { inner } Self { inner }
} }
pub(crate) fn inner(&self) -> &A {
&self.inner
}
pub(crate) fn replace_inner(&mut self, inner: A) {
self.inner = inner;
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct ActorTypeMetadata {
pub actor_type_id: TypeId,
pub actor_type_name: &'static str,
pub message_type_id: TypeId,
pub message_type_name: &'static str,
} }
/// Trait for type-erased actors — single-message handler. /// Trait for type-erased actors — single-message handler.
@ -456,6 +472,12 @@ pub trait AnyActor: Send {
/// Called on graceful stop, before removal. See [`ActorInterface::on_stop`]. /// Called on graceful stop, before removal. See [`ActorInterface::on_stop`].
fn on_stop(&mut self, _ctx: &Ctx) {} fn on_stop(&mut self, _ctx: &Ctx) {}
fn metadata(&self) -> ActorTypeMetadata;
fn as_any(&self) -> &dyn Any;
fn as_any_mut(&mut self) -> &mut dyn Any;
} }
impl<A> AnyActor for Actor<A> impl<A> AnyActor for Actor<A>
@ -493,6 +515,23 @@ where
fn on_stop(&mut self, ctx: &Ctx) { fn on_stop(&mut self, ctx: &Ctx) {
self.inner.on_stop(ctx); self.inner.on_stop(ctx);
} }
fn metadata(&self) -> ActorTypeMetadata {
ActorTypeMetadata {
actor_type_id: TypeId::of::<A>(),
actor_type_name: std::any::type_name::<A>(),
message_type_id: TypeId::of::<A::Incoming>(),
message_type_name: std::any::type_name::<A::Incoming>(),
}
}
fn as_any(&self) -> &dyn Any {
self
}
fn as_any_mut(&mut self) -> &mut dyn Any {
self
}
} }
/// Unique token identifying a monitor subscription. /// Unique token identifying a monitor subscription.

163
src/admin.rs Normal file
View file

@ -0,0 +1,163 @@
use crate::actor::{ActorAddress, ActorInterface, ActorTypeMetadata, AnyActor, Message};
use crate::runtime::{Inbox, Runtime};
use parking_lot::Mutex;
use std::any::Any;
use std::sync::Arc;
use std::sync::atomic::AtomicUsize;
pub type AdminResult<T> = Result<T, AdminError>;
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum AdminError {
ActorNotFound {
actor: ActorAddress,
},
AddressMismatch {
requested: ActorAddress,
snapshot: ActorAddress,
},
TypeMismatch {
expected_actor_type: &'static str,
expected_message_type: &'static str,
actual_actor_type: &'static str,
actual_message_type: &'static str,
},
Timeout,
}
pub struct RuntimeAdmin<'a> {
pub(crate) runtime: &'a Runtime,
}
pub struct Admin<T: Message> {
pub(crate) inbox: Inbox<AdminResult<T>>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct OperationResult {
pub applied: bool,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ActorStatus {
pub started: bool,
pub suspended: bool,
pub stopping: bool,
pub poisoned: bool,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ActorSummary {
pub address: ActorAddress,
pub actor_type: &'static str,
pub message_type: &'static str,
pub worker_id: usize,
pub parent: Option<ActorAddress>,
pub mailbox_depth: usize,
pub status: ActorStatus,
pub last_message_type: Option<&'static str>,
pub messages_handled: u64,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ListActorsResponse {
pub actors: Vec<ActorSummary>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct InspectActorResponse {
pub summary: ActorSummary,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ActorStateSnapshot<A> {
pub actor: ActorAddress,
pub actor_type: &'static str,
pub message_type: &'static str,
pub actor_instance: A,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct GetActorStateResponse<A> {
pub state: ActorStateSnapshot<A>,
}
impl<A: ActorInterface> ActorStateSnapshot<A> {
pub fn new(actor: ActorAddress, actor_instance: A) -> Self {
Self {
actor,
actor_type: std::any::type_name::<A>(),
message_type: std::any::type_name::<A::Incoming>(),
actor_instance,
}
}
}
impl<T: Message> Admin<T> {
pub(crate) fn new(inbox: Inbox<AdminResult<T>>) -> Self {
Self { inbox }
}
pub fn try_recv(&self) -> Option<AdminResult<T>> {
self.inbox.try_recv()
}
pub fn recv_ticking(&self, rt: &Runtime, max_ticks: usize) -> AdminResult<T> {
for _ in 0..max_ticks {
rt.tick();
if let Some(resp) = self.inbox.try_recv() {
return resp;
}
}
Err(AdminError::Timeout)
}
pub fn reply_addr(&self) -> &ActorAddress {
self.inbox.addr()
}
}
pub(crate) type AdminBoxedReply = Box<dyn Any + Send>;
pub(crate) enum AdminCommand {
ListActors {
acc: Arc<ListActorsAccumulator>,
},
InspectActor {
actor: ActorAddress,
reply_to: ActorAddress,
},
GetActorState {
actor: ActorAddress,
reply_to: ActorAddress,
get: Box<
dyn FnOnce(ActorAddress, &dyn AnyActor, ActorTypeMetadata) -> AdminBoxedReply + Send,
>,
not_found: Box<dyn FnOnce(ActorAddress) -> AdminBoxedReply + Send>,
},
ReplaceActorState {
actor: ActorAddress,
reply_to: ActorAddress,
replace: Box<
dyn FnOnce(&mut dyn AnyActor, ActorTypeMetadata) -> AdminResult<OperationResult> + Send,
>,
},
StopActor {
actor: ActorAddress,
reply_to: ActorAddress,
},
SuspendActor {
actor: ActorAddress,
reply_to: ActorAddress,
},
ResumeActor {
actor: ActorAddress,
reply_to: ActorAddress,
},
}
pub(crate) struct ListActorsAccumulator {
pub(crate) remaining: AtomicUsize,
pub(crate) summaries: Mutex<Vec<ActorSummary>>,
pub(crate) reply_to: ActorAddress,
}

View file

@ -1,4 +1,5 @@
pub mod actor; pub mod actor;
pub mod admin;
pub mod extension; pub mod extension;
pub mod process_observer; pub mod process_observer;
pub mod worker; pub mod worker;

View file

@ -1,15 +1,19 @@
use crate::Instant; use crate::Instant;
use std::any::Any; use std::any::{Any, TypeId};
use std::cell::RefCell; use std::cell::RefCell;
use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::{Arc, OnceLock}; use std::sync::{Arc, OnceLock};
use std::thread::Thread; use std::thread::Thread;
#[cfg(not(target_arch = "wasm32"))] #[cfg(not(target_arch = "wasm32"))]
use std::thread::{self, JoinHandle}; use std::thread::{self, JoinHandle};
use crate::actor::{ use crate::actor::{
Actor, ActorAddress, ActorInterface, AnyActor, Environment, ExitValue, Message, ResumeSignal, Actor, ActorAddress, ActorInterface, ActorTypeMetadata, AnyActor, Environment, ExitValue,
SpawnRequest, StopSignal, StopWithSignal, SystemInfo, Message, ResumeSignal, SpawnRequest, StopSignal, StopWithSignal, SystemInfo,
};
use crate::admin::{
ActorStateSnapshot, Admin, AdminCommand, AdminError, AdminResult, GetActorStateResponse,
InspectActorResponse, ListActorsAccumulator, ListActorsResponse, OperationResult, RuntimeAdmin,
}; };
use crate::channel::{Receiver, Sender}; use crate::channel::{Receiver, Sender};
// Re-export config types so existing code using `runtime::RuntimeConfig` still works // Re-export config types so existing code using `runtime::RuntimeConfig` still works
@ -106,6 +110,7 @@ pub struct Runtime {
extension: Option<Arc<dyn RuntimeExtension>>, extension: Option<Arc<dyn RuntimeExtension>>,
transfer_txs: Vec<Sender<Envelope>>, transfer_txs: Vec<Sender<Envelope>>,
spawn_txs: Vec<Sender<SpawnRequest>>, spawn_txs: Vec<Sender<SpawnRequest>>,
admin_txs: Vec<Sender<AdminCommand>>,
placement: Placement, placement: Placement,
is_running: AtomicBool, is_running: AtomicBool,
worker_stats: Vec<Arc<WorkerStats>>, worker_stats: Vec<Arc<WorkerStats>>,
@ -221,6 +226,7 @@ impl Runtime {
let mut transfer_txs = Vec::with_capacity(num_workers); let mut transfer_txs = Vec::with_capacity(num_workers);
let mut spawn_txs = Vec::with_capacity(num_workers); let mut spawn_txs = Vec::with_capacity(num_workers);
let mut admin_txs = Vec::with_capacity(num_workers);
let mut worker_stats = Vec::with_capacity(num_workers); let mut worker_stats = Vec::with_capacity(num_workers);
let mut workers = Vec::with_capacity(num_workers); let mut workers = Vec::with_capacity(num_workers);
@ -233,9 +239,19 @@ impl Runtime {
let spawn_tx = spawn_rx.new_sender(); let spawn_tx = spawn_rx.new_sender();
spawn_txs.push(spawn_tx); spawn_txs.push(spawn_tx);
let admin_rx = Receiver::<AdminCommand>::new(config.channel_buffer_size);
let admin_tx = admin_rx.new_sender();
admin_txs.push(admin_tx);
let stats = Arc::new(WorkerStats::new()); let stats = Arc::new(WorkerStats::new());
worker_stats.push(stats.clone()); worker_stats.push(stats.clone());
workers.push(Worker::new(WorkerId(i), transfer_rx, spawn_rx, stats)); workers.push(Worker::new(
WorkerId(i),
transfer_rx,
spawn_rx,
admin_rx,
stats,
));
} }
let placement = Placement::new(num_workers, worker_stats.clone()); let placement = Placement::new(num_workers, worker_stats.clone());
@ -250,6 +266,7 @@ impl Runtime {
extension: None, extension: None,
transfer_txs, transfer_txs,
spawn_txs, spawn_txs,
admin_txs,
placement, placement,
is_running: AtomicBool::new(false), is_running: AtomicBool::new(false),
worker_stats, worker_stats,
@ -342,6 +359,10 @@ impl Runtime {
self.extension.as_deref() self.extension.as_deref()
} }
pub fn admin(&self) -> RuntimeAdmin<'_> {
RuntimeAdmin { runtime: self }
}
/// Send a request and get a handle for the response. /// Send a request and get a handle for the response.
/// ///
/// Creates a temporary inbox, calls `msg_builder` with the inbox's address /// Creates a temporary inbox, calls `msg_builder` with the inbox's address
@ -590,6 +611,218 @@ impl Runtime {
} }
} }
impl RuntimeAdmin<'_> {
fn new_admin<T: Message>(&self) -> Result<(Admin<T>, ActorAddress), Error> {
let inbox = self.runtime.new_inbox::<AdminResult<T>>()?;
let reply_to = *inbox.addr();
Ok((Admin::new(inbox), reply_to))
}
fn ready<T: Message>(&self, result: AdminResult<T>) -> Result<Admin<T>, Error> {
let (admin, reply_to) = self.new_admin::<T>()?;
let _ = self
.runtime
.inbox_registry
.try_deliver(reply_to, Box::new(result));
Ok(admin)
}
pub fn list_actors(&self) -> Result<Admin<ListActorsResponse>, Error> {
let (admin, reply_to) = self.new_admin::<ListActorsResponse>()?;
let acc = Arc::new(ListActorsAccumulator {
remaining: AtomicUsize::new(self.runtime.admin_txs.len()),
summaries: parking_lot::Mutex::new(Vec::new()),
reply_to,
});
for (idx, tx) in self.runtime.admin_txs.iter().enumerate() {
tx.send(AdminCommand::ListActors { acc: acc.clone() });
notify_worker(&self.runtime.worker_threads, idx);
}
Ok(admin)
}
pub fn inspect_actor(&self, actor: ActorAddress) -> Result<Admin<InspectActorResponse>, Error> {
let Some(wid) = self.runtime.address_map.lookup(&actor) else {
return self.ready::<InspectActorResponse>(Err(AdminError::ActorNotFound { actor }));
};
let (admin, reply_to) = self.new_admin::<InspectActorResponse>()?;
let worker_idx = wid.as_usize();
self.runtime.admin_txs[worker_idx].send(AdminCommand::InspectActor { actor, reply_to });
notify_worker(&self.runtime.worker_threads, worker_idx);
Ok(admin)
}
pub fn get_actor_state<A>(
&self,
actor: ActorAddress,
) -> Result<Admin<GetActorStateResponse<A>>, Error>
where
A: ActorInterface + Clone + Sync,
{
let Some(wid) = self.runtime.address_map.lookup(&actor) else {
return self
.ready::<GetActorStateResponse<A>>(Err(AdminError::ActorNotFound { actor }));
};
let (admin, reply_to) = self.new_admin::<GetActorStateResponse<A>>()?;
let get = Box::new(
|actor: ActorAddress,
erased: &dyn AnyActor,
metadata: ActorTypeMetadata|
-> Box<dyn Any + Send> {
let expected_actor_type = std::any::type_name::<A>();
let expected_message_type = std::any::type_name::<A::Incoming>();
if metadata.actor_type_id != TypeId::of::<A>()
|| metadata.message_type_id != TypeId::of::<A::Incoming>()
{
return Box::new(Err::<GetActorStateResponse<A>, AdminError>(
AdminError::TypeMismatch {
expected_actor_type,
expected_message_type,
actual_actor_type: metadata.actor_type_name,
actual_message_type: metadata.message_type_name,
},
));
}
let Some(typed) = erased.as_any().downcast_ref::<Actor<A>>() else {
return Box::new(Err::<GetActorStateResponse<A>, AdminError>(
AdminError::TypeMismatch {
expected_actor_type,
expected_message_type,
actual_actor_type: metadata.actor_type_name,
actual_message_type: metadata.message_type_name,
},
));
};
Box::new(Ok::<GetActorStateResponse<A>, AdminError>(
GetActorStateResponse {
state: ActorStateSnapshot {
actor,
actor_type: metadata.actor_type_name,
message_type: metadata.message_type_name,
actor_instance: typed.inner().clone(),
},
},
))
},
);
let not_found = Box::new(|actor| {
Box::new(Err::<GetActorStateResponse<A>, AdminError>(
AdminError::ActorNotFound { actor },
)) as Box<dyn Any + Send>
});
let worker_idx = wid.as_usize();
self.runtime.admin_txs[worker_idx].send(AdminCommand::GetActorState {
actor,
reply_to,
get,
not_found,
});
notify_worker(&self.runtime.worker_threads, worker_idx);
Ok(admin)
}
pub fn replace_actor_state<A>(
&self,
actor: ActorAddress,
state: ActorStateSnapshot<A>,
) -> Result<Admin<OperationResult>, Error>
where
A: ActorInterface,
{
if state.actor != actor {
return self.ready::<OperationResult>(Err(AdminError::AddressMismatch {
requested: actor,
snapshot: state.actor,
}));
}
let Some(wid) = self.runtime.address_map.lookup(&actor) else {
return self.ready::<OperationResult>(Err(AdminError::ActorNotFound { actor }));
};
let (admin, reply_to) = self.new_admin::<OperationResult>()?;
let actor_instance = state.actor_instance;
let replace = Box::new(
move |erased: &mut dyn AnyActor,
metadata: ActorTypeMetadata|
-> AdminResult<OperationResult> {
let expected_actor_type = std::any::type_name::<A>();
let expected_message_type = std::any::type_name::<A::Incoming>();
if metadata.actor_type_id != TypeId::of::<A>()
|| metadata.message_type_id != TypeId::of::<A::Incoming>()
{
return Err(AdminError::TypeMismatch {
expected_actor_type,
expected_message_type,
actual_actor_type: metadata.actor_type_name,
actual_message_type: metadata.message_type_name,
});
}
let Some(typed) = erased.as_any_mut().downcast_mut::<Actor<A>>() else {
return Err(AdminError::TypeMismatch {
expected_actor_type,
expected_message_type,
actual_actor_type: metadata.actor_type_name,
actual_message_type: metadata.message_type_name,
});
};
typed.replace_inner(actor_instance);
Ok(OperationResult { applied: true })
},
);
let worker_idx = wid.as_usize();
self.runtime.admin_txs[worker_idx].send(AdminCommand::ReplaceActorState {
actor,
reply_to,
replace,
});
notify_worker(&self.runtime.worker_threads, worker_idx);
Ok(admin)
}
pub fn stop_actor(&self, actor: ActorAddress) -> Result<Admin<OperationResult>, Error> {
let Some(wid) = self.runtime.address_map.lookup(&actor) else {
return self.ready::<OperationResult>(Err(AdminError::ActorNotFound { actor }));
};
let (admin, reply_to) = self.new_admin::<OperationResult>()?;
let worker_idx = wid.as_usize();
self.runtime.admin_txs[worker_idx].send(AdminCommand::StopActor { actor, reply_to });
notify_worker(&self.runtime.worker_threads, worker_idx);
Ok(admin)
}
pub fn suspend_actor(&self, actor: ActorAddress) -> Result<Admin<OperationResult>, Error> {
let Some(wid) = self.runtime.address_map.lookup(&actor) else {
return self.ready::<OperationResult>(Err(AdminError::ActorNotFound { actor }));
};
let (admin, reply_to) = self.new_admin::<OperationResult>()?;
let worker_idx = wid.as_usize();
self.runtime.admin_txs[worker_idx].send(AdminCommand::SuspendActor { actor, reply_to });
notify_worker(&self.runtime.worker_threads, worker_idx);
Ok(admin)
}
pub fn resume_actor(&self, actor: ActorAddress) -> Result<Admin<OperationResult>, Error> {
let Some(wid) = self.runtime.address_map.lookup(&actor) else {
return self.ready::<OperationResult>(Err(AdminError::ActorNotFound { actor }));
};
let (admin, reply_to) = self.new_admin::<OperationResult>()?;
let worker_idx = wid.as_usize();
self.runtime.admin_txs[worker_idx].send(AdminCommand::ResumeActor { actor, reply_to });
notify_worker(&self.runtime.worker_threads, worker_idx);
Ok(admin)
}
}
/// Wake a parked worker thread so it can process new work. /// Wake a parked worker thread so it can process new work.
/// No-op if the thread handle hasn't been registered yet (single-threaded tick mode). /// No-op if the thread handle hasn't been registered yet (single-threaded tick mode).
#[inline] #[inline]

View file

@ -11,6 +11,10 @@ use crate::actor::{
ActorAddress, AnyActor, ContextInner, Ctx, Environment, ExitValue, ResumeSignal, SpawnRequest, ActorAddress, AnyActor, ContextInner, Ctx, Environment, ExitValue, ResumeSignal, SpawnRequest,
StopReason, StopSignal, StopWithSignal, SystemInfo, StopReason, StopSignal, StopWithSignal, SystemInfo,
}; };
use crate::admin::{
ActorStatus, ActorSummary, AdminCommand, AdminError, AdminResult, InspectActorResponse,
ListActorsResponse, OperationResult,
};
use crate::channel::Receiver; use crate::channel::Receiver;
use crate::delivery::{AddrBuildHasher, AddrMap, Envelope, TickContext, WorkerId}; use crate::delivery::{AddrBuildHasher, AddrMap, Envelope, TickContext, WorkerId};
use crate::stats::{ActorSnapshot, TickTiming, WorkerStats}; use crate::stats::{ActorSnapshot, TickTiming, WorkerStats};
@ -69,6 +73,7 @@ pub(crate) struct Worker {
pub(crate) pool: ActorPool, pub(crate) pool: ActorPool,
transfer_rx: Receiver<Envelope>, transfer_rx: Receiver<Envelope>,
spawn_rx: Receiver<SpawnRequest>, spawn_rx: Receiver<SpawnRequest>,
admin_rx: Receiver<AdminCommand>,
stats: Arc<WorkerStats>, stats: Arc<WorkerStats>,
/// Reusable scratch buffer for building per-actor snapshots. /// Reusable scratch buffer for building per-actor snapshots.
snapshot_buf: Vec<ActorSnapshot>, snapshot_buf: Vec<ActorSnapshot>,
@ -84,6 +89,7 @@ impl Worker {
id: WorkerId, id: WorkerId,
transfer_rx: Receiver<Envelope>, transfer_rx: Receiver<Envelope>,
spawn_rx: Receiver<SpawnRequest>, spawn_rx: Receiver<SpawnRequest>,
admin_rx: Receiver<AdminCommand>,
stats: Arc<WorkerStats>, stats: Arc<WorkerStats>,
) -> Self { ) -> Self {
Self { Self {
@ -91,6 +97,7 @@ impl Worker {
pool: ActorPool::new(), pool: ActorPool::new(),
transfer_rx, transfer_rx,
spawn_rx, spawn_rx,
admin_rx,
stats, stats,
snapshot_buf: Vec::new(), snapshot_buf: Vec::new(),
worker_ext: None, worker_ext: None,
@ -132,6 +139,88 @@ impl Worker {
did_work did_work
} }
fn drain_admin(&mut self, tc: &TickContext) -> bool {
let mut did_work = false;
while let Some(cmd) = self.admin_rx.try_recv() {
did_work = true;
self.apply_admin_command(tc, cmd);
}
did_work
}
fn send_admin_reply<T: crate::actor::Message>(
tc: &TickContext,
reply_to: ActorAddress,
result: AdminResult<T>,
) {
let _ = tc.inbox_registry.try_deliver(reply_to, Box::new(result));
}
fn apply_admin_command(&mut self, tc: &TickContext, cmd: AdminCommand) {
match cmd {
AdminCommand::ListActors { acc } => {
let mut local = Vec::new();
self.pool.actor_summaries_into(self.id, &mut local);
{
let mut summaries = acc.summaries.lock();
summaries.extend(local);
}
if acc.remaining.fetch_sub(1, Ordering::AcqRel) == 1 {
let actors = {
let mut summaries = acc.summaries.lock();
std::mem::take(&mut *summaries)
};
Self::send_admin_reply(tc, acc.reply_to, Ok(ListActorsResponse { actors }));
}
}
AdminCommand::InspectActor { actor, reply_to } => {
let result = self
.pool
.actor_summary(self.id, actor)
.map(|summary| InspectActorResponse { summary });
Self::send_admin_reply(tc, reply_to, result);
}
AdminCommand::GetActorState {
actor,
reply_to,
get,
not_found,
} => {
let boxed = match self.pool.get_actor_erased(actor) {
Some(erased) => get(actor, erased, erased.metadata()),
None => not_found(actor),
};
let _ = tc.inbox_registry.try_deliver(reply_to, boxed);
}
AdminCommand::ReplaceActorState {
actor,
reply_to,
replace,
} => {
let result = match self.pool.get_actor_erased_mut(actor) {
Some(erased) => {
let metadata = erased.metadata();
replace(erased, metadata)
}
None => Err(AdminError::ActorNotFound { actor }),
};
Self::send_admin_reply(tc, reply_to, result);
}
AdminCommand::StopActor { actor, reply_to } => {
let result = self.pool.stop_actor_admin(actor, &self.stats);
Self::send_admin_reply(tc, reply_to, result);
}
AdminCommand::SuspendActor { actor, reply_to } => {
let result = self.pool.suspend_actor_admin(actor);
Self::send_admin_reply(tc, reply_to, result);
}
AdminCommand::ResumeActor { actor, reply_to } => {
let result = self.pool.resume_actor_admin(actor);
Self::send_admin_reply(tc, reply_to, result);
}
}
}
/// Phase 7: clean up dead actors, deliver death notifications, GC extension state. /// Phase 7: clean up dead actors, deliver death notifications, GC extension state.
fn cleanup_dead_actors(&mut self, tc: &TickContext) -> bool { fn cleanup_dead_actors(&mut self, tc: &TickContext) -> bool {
let cleanup_pending: RefCell<Vec<(ActorAddress, Box<dyn Any + Send>)>> = let cleanup_pending: RefCell<Vec<(ActorAddress, Box<dyn Any + Send>)>> =
@ -197,6 +286,7 @@ impl Worker {
if !self.has_backlog if !self.has_backlog
&& self.spawn_rx.is_empty() && self.spawn_rx.is_empty()
&& self.transfer_rx.is_empty() && self.transfer_rx.is_empty()
&& self.admin_rx.is_empty()
&& !self && !self
.worker_ext .worker_ext
.as_ref() .as_ref()
@ -221,7 +311,10 @@ impl Worker {
} }
let t2 = Instant::now(); let t2 = Instant::now();
// 2.5. Fire per-worker extension (e.g., timers) → deliver before tick_all // 3. Drain admin queue → inspect or mutate worker-owned slots before handlers
did_work |= self.drain_admin(tc);
// 4. Fire per-worker extension (e.g., timers) → deliver before tick_all
let ext_msgs: Vec<_> = self let ext_msgs: Vec<_> = self
.worker_ext .worker_ext
.as_mut() .as_mut()
@ -232,7 +325,7 @@ impl Worker {
did_work = true; did_work = true;
} }
// 3. Tick all actors with WorkerContext // 5. Tick all actors with WorkerContext
let pending_local: RefCell<Vec<(ActorAddress, Box<dyn Any + Send>)>> = let pending_local: RefCell<Vec<(ActorAddress, Box<dyn Any + Send>)>> =
RefCell::new(Vec::new()); RefCell::new(Vec::new());
let stop_requests: RefCell<Vec<ActorAddress>> = RefCell::new(Vec::new()); let stop_requests: RefCell<Vec<ActorAddress>> = RefCell::new(Vec::new());
@ -275,12 +368,12 @@ impl Worker {
); );
} }
// 4. Drain spawn queue again — actors spawned during step 3 // 6. Drain spawn queue again — actors spawned during step 5
// must be in the pool before pending_local delivery. // must be in the pool before pending_local delivery.
did_work |= self.drain_spawns(tc); did_work |= self.drain_spawns(tc);
let t4 = Instant::now(); let t4 = Instant::now();
// 5. Drain pending_local buffer → deliver to local actors // 7. Drain pending_local buffer → deliver to local actors
let pending = pending_local.into_inner(); let pending = pending_local.into_inner();
if !pending.is_empty() { if !pending.is_empty() {
did_work = true; did_work = true;
@ -289,7 +382,7 @@ impl Worker {
self.pool.deliver(&addr, msg); self.pool.deliver(&addr, msg);
} }
// 5.5. Process worker extension requests from handlers (e.g., timer scheduling) // 7.5. Process worker extension requests from handlers (e.g., timer scheduling)
if let Some(ext) = &mut self.worker_ext { if let Some(ext) = &mut self.worker_ext {
for request in worker_requests.into_inner() { for request in worker_requests.into_inner() {
ext.handle_request(request); ext.handle_request(request);
@ -298,7 +391,7 @@ impl Worker {
let t5 = Instant::now(); let t5 = Instant::now();
// 6. Publish stats (skip entirely when idle to avoid allocation + mutex) // 8. Publish stats (skip entirely when idle to avoid allocation + mutex)
if did_work { if did_work {
self.stats self.stats
.num_actors .num_actors
@ -344,7 +437,7 @@ impl Worker {
); );
} }
// 7. Clean up poisoned and stopping actors // 9. Clean up poisoned and stopping actors
did_work |= self.cleanup_dead_actors(tc); did_work |= self.cleanup_dead_actors(tc);
self.has_backlog = did_work; self.has_backlog = did_work;
@ -545,6 +638,95 @@ impl ActorPool {
} }
} }
fn get_actor_erased(&self, addr: ActorAddress) -> Option<&dyn AnyActor> {
self.actors.get(&addr).map(|slot| slot.actor.as_ref())
}
fn get_actor_erased_mut(
&mut self,
addr: ActorAddress,
) -> Option<&mut (dyn AnyActor + 'static)> {
self.actors.get_mut(&addr).map(|slot| slot.actor.as_mut())
}
fn actor_summary_from_slot(
worker_id: WorkerId,
address: ActorAddress,
slot: &ActorSlot,
) -> ActorSummary {
let metadata = slot.actor.metadata();
ActorSummary {
address,
actor_type: metadata.actor_type_name,
message_type: metadata.message_type_name,
worker_id: worker_id.as_usize(),
parent: slot.parent_addr,
mailbox_depth: slot.mailbox.len(),
status: ActorStatus {
started: slot.started,
suspended: slot.suspended,
stopping: slot.stopping,
poisoned: slot.poisoned,
},
last_message_type: slot.last_msg_type,
messages_handled: slot.messages_processed,
}
}
fn actor_summary(&self, worker_id: WorkerId, addr: ActorAddress) -> AdminResult<ActorSummary> {
self.actors
.get(&addr)
.map(|slot| Self::actor_summary_from_slot(worker_id, addr, slot))
.ok_or(AdminError::ActorNotFound { actor: addr })
}
fn actor_summaries_into(&self, worker_id: WorkerId, out: &mut Vec<ActorSummary>) {
out.clear();
out.extend(
self.actors
.iter()
.map(|(&addr, slot)| Self::actor_summary_from_slot(worker_id, addr, slot)),
);
}
fn suspend_actor_admin(&mut self, addr: ActorAddress) -> AdminResult<OperationResult> {
match self.actors.get_mut(&addr) {
Some(slot) => {
slot.suspended = true;
Ok(OperationResult { applied: true })
}
None => Err(AdminError::ActorNotFound { actor: addr }),
}
}
fn resume_actor_admin(&mut self, addr: ActorAddress) -> AdminResult<OperationResult> {
match self.actors.get_mut(&addr) {
Some(slot) => {
slot.suspended = false;
Ok(OperationResult { applied: true })
}
None => Err(AdminError::ActorNotFound { actor: addr }),
}
}
fn stop_actor_admin(
&mut self,
addr: ActorAddress,
stats: &WorkerStats,
) -> AdminResult<OperationResult> {
match self.actors.get_mut(&addr) {
Some(slot) => {
if !slot.stopping {
stats.stops.fetch_add(1, Ordering::Relaxed);
}
slot.stopping = true;
slot.mailbox.clear();
Ok(OperationResult { applied: true })
}
None => Err(AdminError::ActorNotFound { actor: addr }),
}
}
/// Tick all actors in the pool. Returns the number of messages processed. /// Tick all actors in the pool. Returns the number of messages processed.
/// ///
/// Each actor processes up to `budget` messages per tick (0 = unlimited). /// Each actor processes up to `budget` messages per tick (0 = unlimited).

608
tests/runtime_admin.rs Normal file
View file

@ -0,0 +1,608 @@
//! Runtime Admin API tests — inventory, typed actor state, lifecycle control, and scheduling.
mod common;
use common::*;
use std::collections::HashSet;
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::time::{Duration, Instant};
use swactor::admin::{ActorStateSnapshot, AdminError, OperationResult};
use swactor::config::RuntimeConfig;
#[derive(Clone)]
struct AddAndReport {
delta: usize,
reply_to: ActorAddress,
}
#[derive(Clone)]
struct ReplaceProbe {
value: usize,
started: Arc<AtomicUsize>,
stopped: Arc<AtomicUsize>,
}
impl ActorInterface for ReplaceProbe {
type Incoming = AddAndReport;
type Response = Count;
fn on_start(&mut self, _ctx: &Ctx) {
self.started.fetch_add(1, Ordering::SeqCst);
}
fn on_stop(&mut self, _ctx: &Ctx) {
self.stopped.fetch_add(1, Ordering::SeqCst);
}
fn handle(&mut self, ctx: &Ctx, msg: AddAndReport) {
self.value += msg.delta;
let _ = ctx.send(msg.reply_to, Count(self.value));
}
}
struct WrongProbe;
impl ActorInterface for WrongProbe {
type Incoming = Ping;
type Response = Pong;
fn handle(&mut self, _ctx: &Ctx, _msg: Ping) {}
}
struct StopProbe {
started: Arc<AtomicUsize>,
handled: Arc<AtomicUsize>,
stopped: Arc<AtomicUsize>,
}
impl ActorInterface for StopProbe {
type Incoming = Ping;
type Response = Pong;
fn on_start(&mut self, _ctx: &Ctx) {
self.started.fetch_add(1, Ordering::SeqCst);
}
fn on_stop(&mut self, _ctx: &Ctx) {
self.stopped.fetch_add(1, Ordering::SeqCst);
}
fn handle(&mut self, ctx: &Ctx, msg: Ping) {
self.handled.fetch_add(1, Ordering::SeqCst);
let _ = ctx.send(msg.reply_to, Pong);
}
}
fn poll_admin<T: swactor::actor::Message>(
admin: &swactor::admin::Admin<T>,
timeout: Duration,
) -> Option<swactor::admin::AdminResult<T>> {
let start = Instant::now();
while start.elapsed() < timeout {
if let Some(value) = admin.try_recv() {
return Some(value);
}
std::thread::sleep(Duration::from_millis(5));
}
None
}
fn poll_inbox<M: swactor::actor::Message>(inbox: &Inbox<M>, timeout: Duration) -> Option<M> {
let start = Instant::now();
while start.elapsed() < timeout {
if let Some(value) = inbox.try_recv() {
return Some(value);
}
std::thread::sleep(Duration::from_millis(5));
}
None
}
fn operation_applied() -> OperationResult {
OperationResult { applied: true }
}
#[test]
fn ask_recv_ticking_delivers_reply_through_runtime_inbox() {
let rt = std_runtime(RuntimeConfig::default());
let actor = rt.spawn(SelfAddrActor).unwrap();
let ask = rt
.ask::<WhoAreYou, MyAddr>(actor, |reply_to| WhoAreYou { reply_to })
.unwrap();
assert_eq!(
ask.try_recv(),
None,
"ask reply is not available before ticking"
);
assert_eq!(
ask.recv_ticking(&rt, 5).unwrap(),
MyAddr(actor),
"recv_ticking drives the runtime inbox reply path"
);
}
#[test]
fn admin_list_and_inspect_report_actor_slot_metadata() {
let rt = std_runtime(RuntimeConfig::default());
let ping_pong = rt.spawn(PingPongActor).unwrap();
let counter = rt.spawn(CounterActor { count: 0 }).unwrap();
rt.tick();
let count_inbox = rt.new_inbox::<Count>().unwrap();
rt.send_to(
counter,
Increment {
reply_to: *count_inbox.addr(),
},
)
.unwrap();
rt.send_to(
counter,
Increment {
reply_to: *count_inbox.addr(),
},
)
.unwrap();
assert_eq!(tick_until_recv(&rt, &count_inbox, 5), Some(Count(1)));
assert_eq!(tick_until_recv(&rt, &count_inbox, 5), Some(Count(2)));
let response = rt
.admin()
.list_actors()
.unwrap()
.recv_ticking(&rt, 5)
.unwrap();
assert_eq!(
response
.actors
.iter()
.filter(|summary| summary.address == ping_pong)
.count(),
1,
"ping-pong actor appears exactly once in inventory"
);
assert_eq!(
response
.actors
.iter()
.filter(|summary| summary.address == counter)
.count(),
1,
"counter actor appears exactly once in inventory"
);
let counter_summary = response
.actors
.iter()
.find(|summary| summary.address == counter)
.expect("counter summary missing");
assert_eq!(counter_summary.worker_id, 0);
assert_eq!(counter_summary.parent, None);
assert_eq!(counter_summary.mailbox_depth, 0);
assert!(counter_summary.status.started);
assert!(!counter_summary.status.suspended);
assert!(!counter_summary.status.stopping);
assert!(!counter_summary.status.poisoned);
assert_eq!(counter_summary.messages_handled, 2);
assert!(counter_summary.actor_type.ends_with("CounterActor"));
assert!(counter_summary.message_type.ends_with("Increment"));
let inspect = rt
.admin()
.inspect_actor(counter)
.unwrap()
.recv_ticking(&rt, 5)
.unwrap();
assert_eq!(inspect.summary, *counter_summary);
let missing = ActorAddress::new_random();
let missing_result = rt
.admin()
.inspect_actor(missing)
.unwrap()
.recv_ticking(&rt, 5);
assert!(
matches!(missing_result, Err(AdminError::ActorNotFound { actor }) if actor == missing),
"missing actor is reported through AdminResult"
);
}
#[test]
fn admin_get_and_replace_actor_state_preserves_slot_metadata() {
let rt = std_runtime(RuntimeConfig::default());
let started = Arc::new(AtomicUsize::new(0));
let stopped = Arc::new(AtomicUsize::new(0));
let addr = rt
.spawn(ReplaceProbe {
value: 1,
started: started.clone(),
stopped: stopped.clone(),
})
.unwrap();
rt.tick();
assert_eq!(started.load(Ordering::SeqCst), 1);
assert_eq!(stopped.load(Ordering::SeqCst), 0);
let count_inbox = rt.new_inbox::<Count>().unwrap();
rt.send_to(
addr,
AddAndReport {
delta: 1,
reply_to: *count_inbox.addr(),
},
)
.unwrap();
assert_eq!(tick_until_recv(&rt, &count_inbox, 5), Some(Count(2)));
let state = rt
.admin()
.get_actor_state::<ReplaceProbe>(addr)
.unwrap()
.recv_ticking(&rt, 5)
.unwrap()
.state;
assert_eq!(state.actor, addr);
assert!(state.actor_type.ends_with("ReplaceProbe"));
assert!(state.message_type.ends_with("AddAndReport"));
assert_eq!(state.actor_instance.value, 2);
let replacement = ActorStateSnapshot::new(
addr,
ReplaceProbe {
value: 100,
started: started.clone(),
stopped: stopped.clone(),
},
);
let replace_result = rt
.admin()
.replace_actor_state::<ReplaceProbe>(addr, replacement)
.unwrap()
.recv_ticking(&rt, 5)
.unwrap();
assert_eq!(replace_result, operation_applied());
assert_eq!(
started.load(Ordering::SeqCst),
1,
"replacement does not call on_start"
);
assert_eq!(
stopped.load(Ordering::SeqCst),
0,
"replacement does not call on_stop"
);
rt.send_to(
addr,
AddAndReport {
delta: 1,
reply_to: *count_inbox.addr(),
},
)
.unwrap();
assert_eq!(tick_until_recv(&rt, &count_inbox, 5), Some(Count(101)));
let summary = rt
.admin()
.inspect_actor(addr)
.unwrap()
.recv_ticking(&rt, 5)
.unwrap()
.summary;
assert_eq!(summary.address, addr);
assert_eq!(summary.worker_id, 0);
assert_eq!(
summary.messages_handled, 2,
"state replacement preserves slot-owned message counters"
);
let stop_result = rt
.admin()
.stop_actor(addr)
.unwrap()
.recv_ticking(&rt, 5)
.unwrap();
assert_eq!(stop_result, operation_applied());
rt.tick();
assert_eq!(stopped.load(Ordering::SeqCst), 1);
}
#[test]
fn admin_replace_rejects_wrong_actor_type_and_wrong_snapshot_address() {
let rt = std_runtime(RuntimeConfig::default());
let started = Arc::new(AtomicUsize::new(0));
let stopped = Arc::new(AtomicUsize::new(0));
let addr = rt
.spawn(ReplaceProbe {
value: 10,
started: started.clone(),
stopped: stopped.clone(),
})
.unwrap();
rt.tick();
let wrong_type_snapshot = ActorStateSnapshot::new(addr, WrongProbe);
let wrong_type = rt
.admin()
.replace_actor_state::<WrongProbe>(addr, wrong_type_snapshot)
.unwrap()
.recv_ticking(&rt, 5);
assert!(
matches!(wrong_type, Err(AdminError::TypeMismatch { .. })),
"wrong concrete actor type is rejected"
);
let wrong_addr = ActorAddress::new_random();
let wrong_addr_snapshot = ActorStateSnapshot::new(
wrong_addr,
ReplaceProbe {
value: 50,
started: started.clone(),
stopped: stopped.clone(),
},
);
let wrong_address = rt
.admin()
.replace_actor_state::<ReplaceProbe>(addr, wrong_addr_snapshot)
.unwrap()
.recv_ticking(&rt, 5);
assert!(
matches!(wrong_address, Err(AdminError::AddressMismatch { requested, snapshot }) if requested == addr && snapshot == wrong_addr),
"snapshot address must match the target address"
);
let count_inbox = rt.new_inbox::<Count>().unwrap();
rt.send_to(
addr,
AddAndReport {
delta: 1,
reply_to: *count_inbox.addr(),
},
)
.unwrap();
assert_eq!(
tick_until_recv(&rt, &count_inbox, 5),
Some(Count(11)),
"failed replacements do not mutate the original actor state"
);
}
#[test]
fn admin_suspend_queues_messages_until_resume() {
let rt = std_runtime(RuntimeConfig::default());
let counter = Arc::new(AtomicUsize::new(0));
let addr = rt
.spawn(CountingPingActor {
counter: counter.clone(),
})
.unwrap();
rt.tick();
let suspend_result = rt
.admin()
.suspend_actor(addr)
.unwrap()
.recv_ticking(&rt, 5)
.unwrap();
assert_eq!(suspend_result, operation_applied());
let pong_inbox = rt.new_inbox::<Pong>().unwrap();
for _ in 0..3 {
rt.send_to(
addr,
Ping {
reply_to: *pong_inbox.addr(),
},
)
.unwrap();
}
tick_n(&rt, 5);
assert_eq!(counter.load(Ordering::SeqCst), 0);
assert_eq!(pong_inbox.try_recv(), None);
let suspended = rt
.admin()
.inspect_actor(addr)
.unwrap()
.recv_ticking(&rt, 5)
.unwrap()
.summary;
assert!(suspended.status.suspended);
assert_eq!(suspended.mailbox_depth, 3);
let resume_result = rt
.admin()
.resume_actor(addr)
.unwrap()
.recv_ticking(&rt, 5)
.unwrap();
assert_eq!(resume_result, operation_applied());
for _ in 0..3 {
assert_eq!(tick_until_recv(&rt, &pong_inbox, 5), Some(Pong));
}
assert_eq!(counter.load(Ordering::SeqCst), 3);
let resumed = rt
.admin()
.inspect_actor(addr)
.unwrap()
.recv_ticking(&rt, 5)
.unwrap()
.summary;
assert!(!resumed.status.suspended);
assert_eq!(resumed.mailbox_depth, 0);
assert_eq!(resumed.messages_handled, 3);
}
#[test]
fn admin_stop_clears_pending_mailbox_without_calling_handle() {
let rt = std_runtime(RuntimeConfig::default());
let started = Arc::new(AtomicUsize::new(0));
let handled = Arc::new(AtomicUsize::new(0));
let stopped = Arc::new(AtomicUsize::new(0));
let addr = rt
.spawn(StopProbe {
started: started.clone(),
handled: handled.clone(),
stopped: stopped.clone(),
})
.unwrap();
rt.tick();
assert_eq!(started.load(Ordering::SeqCst), 1);
let pong_inbox = rt.new_inbox::<Pong>().unwrap();
for _ in 0..5 {
rt.send_to(
addr,
Ping {
reply_to: *pong_inbox.addr(),
},
)
.unwrap();
}
let stop_result = rt
.admin()
.stop_actor(addr)
.unwrap()
.recv_ticking(&rt, 5)
.unwrap();
assert_eq!(stop_result, operation_applied());
rt.tick();
assert_eq!(handled.load(Ordering::SeqCst), 0);
assert_eq!(stopped.load(Ordering::SeqCst), 1);
assert_eq!(pong_inbox.try_recv(), None);
assert!(
rt.send_to(
addr,
Ping {
reply_to: *pong_inbox.addr(),
},
)
.is_err(),
"admin-stopped actor is removed from normal send routing"
);
let inspect = rt.admin().inspect_actor(addr).unwrap().recv_ticking(&rt, 5);
assert!(
matches!(inspect, Err(AdminError::ActorNotFound { actor }) if actor == addr),
"admin-stopped actor is no longer inspectable"
);
}
#[test]
fn threaded_admin_suspend_resume_wakes_parked_worker() {
let rt = std_runtime(RuntimeConfig {
num_threads: 2,
..Default::default()
});
let counter = Arc::new(AtomicUsize::new(0));
let addr = rt
.spawn(CountingPingActor {
counter: counter.clone(),
})
.unwrap();
let pong_inbox = rt.new_inbox::<Pong>().unwrap();
let handle = rt.run().unwrap();
std::thread::sleep(Duration::from_millis(50));
let suspended = handle.runtime.admin().suspend_actor(addr).unwrap();
let suspended = poll_admin(&suspended, Duration::from_secs(1));
handle
.runtime
.send_to(
addr,
Ping {
reply_to: *pong_inbox.addr(),
},
)
.unwrap();
let pong_while_suspended = poll_inbox(&pong_inbox, Duration::from_millis(100));
let count_while_suspended = counter.load(Ordering::SeqCst);
let resumed = handle.runtime.admin().resume_actor(addr).unwrap();
let resumed = poll_admin(&resumed, Duration::from_secs(1));
let pong_after_resume = poll_inbox(&pong_inbox, Duration::from_secs(1));
let final_count = counter.load(Ordering::SeqCst);
handle.shutdown();
handle.join();
assert_eq!(suspended, Some(Ok(operation_applied())));
assert_eq!(pong_while_suspended, None);
assert_eq!(count_while_suspended, 0);
assert_eq!(resumed, Some(Ok(operation_applied())));
assert_eq!(pong_after_resume, Some(Pong));
assert_eq!(final_count, 1);
}
#[test]
fn admin_list_actors_aggregates_all_workers() {
let rt = std_runtime(RuntimeConfig {
num_threads: 4,
max_actors: 100,
..Default::default()
});
let mut addrs = Vec::new();
for _ in 0..16 {
addrs.push(rt.spawn(CounterActor { count: 0 }).unwrap());
}
let handle = rt.run().unwrap();
let start = Instant::now();
let mut response = None;
while start.elapsed() < Duration::from_secs(1) {
let admin = handle.runtime.admin().list_actors().unwrap();
if let Some(Ok(list)) = poll_admin(&admin, Duration::from_millis(100)) {
if list.actors.len() == addrs.len() {
response = Some(list);
break;
}
}
std::thread::sleep(Duration::from_millis(5));
}
handle.shutdown();
handle.join();
let response = response.expect("admin list did not observe all spawned actors within timeout");
let expected: HashSet<_> = addrs.iter().copied().collect();
let actual: HashSet<_> = response
.actors
.iter()
.map(|summary| summary.address)
.collect();
assert_eq!(actual, expected);
for addr in &addrs {
assert_eq!(
response
.actors
.iter()
.filter(|summary| summary.address == *addr)
.count(),
1,
"actor {addr} appears exactly once in aggregated list"
);
}
let worker_ids: HashSet<_> = response
.actors
.iter()
.map(|summary| summary.worker_id)
.collect();
assert!(
worker_ids.len() >= 2,
"aggregation should include actors from at least two workers, got {worker_ids:?}"
);
}