slow but working mvp-chat
This commit is contained in:
parent
7a7283ab85
commit
9e0a26d99d
29 changed files with 1603 additions and 8373 deletions
|
|
@ -1,4 +1,3 @@
|
||||||
[alias]
|
[alias]
|
||||||
xtask = "run --package xtask --"
|
xtask = "run --package xtask --"
|
||||||
mvp-chat = "run --package xtask -- mvp-chat"
|
mvp-chat = "run --package xtask -- mvp-chat"
|
||||||
mvp-chat-test = "test -p mvp-system --features local-e2e --test one_node_chat_e2e -- --nocapture"
|
|
||||||
|
|
|
||||||
5
Cargo.lock
generated
5
Cargo.lock
generated
|
|
@ -2481,7 +2481,6 @@ dependencies = [
|
||||||
"distribution",
|
"distribution",
|
||||||
"iroh",
|
"iroh",
|
||||||
"iroh-driver",
|
"iroh-driver",
|
||||||
"iroh-relay",
|
|
||||||
"libc",
|
"libc",
|
||||||
"parking_lot",
|
"parking_lot",
|
||||||
"serde",
|
"serde",
|
||||||
|
|
@ -5736,6 +5735,10 @@ dependencies = [
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "xtask"
|
name = "xtask"
|
||||||
version = "0.1.0"
|
version = "0.1.0"
|
||||||
|
dependencies = [
|
||||||
|
"libc",
|
||||||
|
"serde_json",
|
||||||
|
]
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "yasna"
|
name = "yasna"
|
||||||
|
|
|
||||||
|
|
@ -11,6 +11,7 @@ import threading
|
||||||
import time
|
import time
|
||||||
import struct
|
import struct
|
||||||
import traceback
|
import traceback
|
||||||
|
import shutil
|
||||||
import urllib.parse
|
import urllib.parse
|
||||||
import urllib.request
|
import urllib.request
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
|
|
@ -159,6 +160,17 @@ def fatal(reason: str, **fields: Any) -> None:
|
||||||
def test_mode() -> bool:
|
def test_mode() -> bool:
|
||||||
return os.environ.get("MVP_TINYGRAD_TEST_MODE", "").strip().lower() in {"1", "true", "yes", "on"}
|
return os.environ.get("MVP_TINYGRAD_TEST_MODE", "").strip().lower() in {"1", "true", "yes", "on"}
|
||||||
|
|
||||||
|
def configure_tinygrad_cuda_compiler(device: str) -> None:
|
||||||
|
if device.split(":", 1)[0].upper() != "CUDA":
|
||||||
|
return
|
||||||
|
if os.environ.get("CUDA_PTX") or os.environ.get("CUDA_CC"):
|
||||||
|
return
|
||||||
|
if shutil.which("nvcc") is not None:
|
||||||
|
return
|
||||||
|
os.environ["CUDA_PTX"] = "1"
|
||||||
|
control(type="TinygradCudaCompilerSelected", requested_device=device, compiler="PTX", reason="nvcc_not_found")
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
def initialize(cmd: dict[str, Any]) -> None:
|
def initialize(cmd: dict[str, Any]) -> None:
|
||||||
global Tensor, dtypes, arena
|
global Tensor, dtypes, arena
|
||||||
|
|
@ -182,6 +194,7 @@ def initialize(cmd: dict[str, Any]) -> None:
|
||||||
elapsed_ms=int((time.monotonic() - started) * 1000),
|
elapsed_ms=int((time.monotonic() - started) * 1000),
|
||||||
)
|
)
|
||||||
return
|
return
|
||||||
|
configure_tinygrad_cuda_compiler(device)
|
||||||
control(type="TinygradImportStarted", requested_device=device, env_DEV=os.environ.get("DEV"))
|
control(type="TinygradImportStarted", requested_device=device, env_DEV=os.environ.get("DEV"))
|
||||||
from tinygrad import Tensor as TinyTensor, dtypes as tiny_dtypes
|
from tinygrad import Tensor as TinyTensor, dtypes as tiny_dtypes
|
||||||
|
|
||||||
|
|
|
||||||
177
crates/iroh-driver/src/edge_transport.rs
Normal file
177
crates/iroh-driver/src/edge_transport.rs
Normal file
|
|
@ -0,0 +1,177 @@
|
||||||
|
//! Driver-owned byte transport for ring-backed MVP edge protocols.
|
||||||
|
//!
|
||||||
|
//! This module deliberately owns only transport framing: one edge id preamble per
|
||||||
|
//! unidirectional stream, followed by opaque byte chunks. Object-record parsing,
|
||||||
|
//! ring ownership, and stage semantics stay in the MVP/dataplane crates.
|
||||||
|
|
||||||
|
use std::sync::Arc;
|
||||||
|
|
||||||
|
use distribution::types::NodeId;
|
||||||
|
use iroh::endpoint::Connection;
|
||||||
|
use iroh::{Endpoint, EndpointAddr};
|
||||||
|
use parking_lot::Mutex;
|
||||||
|
use tokio::io::AsyncWriteExt;
|
||||||
|
use tokio::runtime::Handle;
|
||||||
|
use tokio::sync::mpsc as tokio_mpsc;
|
||||||
|
|
||||||
|
pub const EDGE_ALPN: &[u8] = b"mvp/pipeline-edge/0";
|
||||||
|
|
||||||
|
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||||
|
pub enum EdgeTransportEvent {
|
||||||
|
StreamArrived {
|
||||||
|
peer: NodeId,
|
||||||
|
edge_id: u64,
|
||||||
|
stream_id: u64,
|
||||||
|
},
|
||||||
|
BytesRead {
|
||||||
|
peer: NodeId,
|
||||||
|
edge_id: u64,
|
||||||
|
stream_id: u64,
|
||||||
|
bytes: Vec<u8>,
|
||||||
|
},
|
||||||
|
StreamEnded {
|
||||||
|
peer: NodeId,
|
||||||
|
edge_id: u64,
|
||||||
|
stream_id: u64,
|
||||||
|
},
|
||||||
|
StreamFault {
|
||||||
|
peer: NodeId,
|
||||||
|
edge_id: Option<u64>,
|
||||||
|
stream_id: Option<u64>,
|
||||||
|
reason: EdgeTransportFault,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||||
|
pub enum EdgeTransportFault {
|
||||||
|
ReadError,
|
||||||
|
WriteError,
|
||||||
|
ProtocolError,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Clone)]
|
||||||
|
pub struct EdgeSendHandle {
|
||||||
|
tx: tokio_mpsc::UnboundedSender<Vec<u8>>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl EdgeSendHandle {
|
||||||
|
pub fn send(&self, bytes: Vec<u8>) -> Result<(), String> {
|
||||||
|
self.tx
|
||||||
|
.send(bytes)
|
||||||
|
.map_err(|_| "edge sender task stopped".to_owned())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(crate) fn spawn_edge_send_pump(
|
||||||
|
handle: Handle,
|
||||||
|
endpoint: Endpoint,
|
||||||
|
peer: EndpointAddr,
|
||||||
|
edge_id: u64,
|
||||||
|
) -> Result<EdgeSendHandle, String> {
|
||||||
|
let (tx, mut rx) = tokio_mpsc::unbounded_channel::<Vec<u8>>();
|
||||||
|
let (ready_tx, ready_rx) = std::sync::mpsc::channel::<Result<(), String>>();
|
||||||
|
handle.spawn(async move {
|
||||||
|
let result: Result<(), String> = async {
|
||||||
|
let conn = endpoint
|
||||||
|
.connect(peer, EDGE_ALPN)
|
||||||
|
.await
|
||||||
|
.map_err(|e| format!("connect edge {edge_id}: {e}"))?;
|
||||||
|
let mut send = conn
|
||||||
|
.open_uni()
|
||||||
|
.await
|
||||||
|
.map_err(|e| format!("open edge stream {edge_id}: {e}"))?;
|
||||||
|
send.write_all(&encode_edge_preamble(edge_id))
|
||||||
|
.await
|
||||||
|
.map_err(|e| format!("write edge preamble {edge_id}: {e}"))?;
|
||||||
|
send.flush()
|
||||||
|
.await
|
||||||
|
.map_err(|e| format!("flush edge preamble {edge_id}: {e}"))?;
|
||||||
|
let _ = ready_tx.send(Ok(()));
|
||||||
|
while let Some(record) = rx.recv().await {
|
||||||
|
send.write_all(&record)
|
||||||
|
.await
|
||||||
|
.map_err(|e| format!("write edge record {edge_id}: {e}"))?;
|
||||||
|
send.flush()
|
||||||
|
.await
|
||||||
|
.map_err(|e| format!("flush edge record {edge_id}: {e}"))?;
|
||||||
|
}
|
||||||
|
send.finish()
|
||||||
|
.map_err(|e| format!("finish edge stream {edge_id}: {e}"))?;
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
.await;
|
||||||
|
if let Err(error) = result {
|
||||||
|
let _ = ready_tx.send(Err(error));
|
||||||
|
}
|
||||||
|
});
|
||||||
|
ready_rx
|
||||||
|
.recv()
|
||||||
|
.map_err(|e| format!("edge {edge_id} sender startup channel closed: {e}"))??;
|
||||||
|
Ok(EdgeSendHandle { tx })
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(crate) fn spawn_edge_recv_pump(
|
||||||
|
handle: Handle,
|
||||||
|
conn: Connection,
|
||||||
|
peer: NodeId,
|
||||||
|
events: Arc<Mutex<Vec<EdgeTransportEvent>>>,
|
||||||
|
stream_group: u64,
|
||||||
|
) {
|
||||||
|
handle.spawn(async move {
|
||||||
|
let mut next_uni_stream_id = stream_group << 32;
|
||||||
|
while let Ok(mut recv) = conn.accept_uni().await {
|
||||||
|
next_uni_stream_id = next_uni_stream_id.saturating_add(1);
|
||||||
|
let current_stream_id = next_uni_stream_id;
|
||||||
|
let mut preamble = [0u8; 8];
|
||||||
|
if recv.read_exact(&mut preamble).await.is_err() {
|
||||||
|
events.lock().push(EdgeTransportEvent::StreamFault {
|
||||||
|
peer,
|
||||||
|
edge_id: None,
|
||||||
|
stream_id: Some(current_stream_id),
|
||||||
|
reason: EdgeTransportFault::ProtocolError,
|
||||||
|
});
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
let edge_id = u64::from_le_bytes(preamble);
|
||||||
|
events.lock().push(EdgeTransportEvent::StreamArrived {
|
||||||
|
peer,
|
||||||
|
edge_id,
|
||||||
|
stream_id: current_stream_id,
|
||||||
|
});
|
||||||
|
let mut chunk = vec![0u8; 4096];
|
||||||
|
loop {
|
||||||
|
match recv.read(&mut chunk).await {
|
||||||
|
Ok(Some(0)) | Ok(None) => {
|
||||||
|
events.lock().push(EdgeTransportEvent::StreamEnded {
|
||||||
|
peer,
|
||||||
|
edge_id,
|
||||||
|
stream_id: current_stream_id,
|
||||||
|
});
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
Ok(Some(n)) => {
|
||||||
|
events.lock().push(EdgeTransportEvent::BytesRead {
|
||||||
|
peer,
|
||||||
|
edge_id,
|
||||||
|
stream_id: current_stream_id,
|
||||||
|
bytes: chunk[..n].to_vec(),
|
||||||
|
});
|
||||||
|
}
|
||||||
|
Err(_) => {
|
||||||
|
events.lock().push(EdgeTransportEvent::StreamFault {
|
||||||
|
peer,
|
||||||
|
edge_id: Some(edge_id),
|
||||||
|
stream_id: Some(current_stream_id),
|
||||||
|
reason: EdgeTransportFault::ReadError,
|
||||||
|
});
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
fn encode_edge_preamble(edge_id: u64) -> [u8; 8] {
|
||||||
|
edge_id.to_le_bytes()
|
||||||
|
}
|
||||||
|
|
@ -30,7 +30,14 @@ use distribution::swim::actor::SwimIn;
|
||||||
use distribution::transport_bridge::{OutFrame, Outbox, RelayMirror, RouteView, peer_addr};
|
use distribution::transport_bridge::{OutFrame, Outbox, RelayMirror, RouteView, peer_addr};
|
||||||
use distribution::types::NodeId;
|
use distribution::types::NodeId;
|
||||||
|
|
||||||
use crate::datastream_transport::DATASTREAM_ALPN;
|
use crate::datastream_transport::{
|
||||||
|
DATASTREAM_ALPN, DatastreamQuicHeader, DatastreamQuicRead, read_events_from_stream,
|
||||||
|
spawn_subscription_writer,
|
||||||
|
};
|
||||||
|
use crate::edge_transport::{
|
||||||
|
EDGE_ALPN, EdgeSendHandle, EdgeTransportEvent, spawn_edge_recv_pump,
|
||||||
|
spawn_edge_send_pump as spawn_edge_sender_task,
|
||||||
|
};
|
||||||
use swactor::actor::ActorAddress;
|
use swactor::actor::ActorAddress;
|
||||||
use swactor::runtime::Runtime;
|
use swactor::runtime::Runtime;
|
||||||
use swactor_transport::CodecRegistry;
|
use swactor_transport::CodecRegistry;
|
||||||
|
|
@ -207,6 +214,34 @@ pub struct JoinStatus {
|
||||||
|
|
||||||
// ─── Driver ─────────────────────────────────────────────────────────────────
|
// ─── Driver ─────────────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
/// Cloneable logical datastream publisher transport. It hides the raw iroh
|
||||||
|
/// endpoint and Tokio task handle from callers while leaving datastream
|
||||||
|
/// subscription/catalog semantics in the datastream crate.
|
||||||
|
#[derive(Clone)]
|
||||||
|
pub struct DatastreamPublishHandle {
|
||||||
|
rt: Handle,
|
||||||
|
endpoint: Endpoint,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl DatastreamPublishHandle {
|
||||||
|
pub fn publish_subscription(
|
||||||
|
&self,
|
||||||
|
peer: EndpointAddr,
|
||||||
|
header: DatastreamQuicHeader,
|
||||||
|
subscription: datastream::DatastreamSubscription,
|
||||||
|
idle_sleep: Duration,
|
||||||
|
) {
|
||||||
|
let _ = spawn_subscription_writer(
|
||||||
|
&self.rt,
|
||||||
|
self.endpoint.clone(),
|
||||||
|
peer,
|
||||||
|
header,
|
||||||
|
subscription,
|
||||||
|
idle_sleep,
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
/// iroh P2P network transport bridge.
|
/// iroh P2P network transport bridge.
|
||||||
///
|
///
|
||||||
/// Bridges the actorized distribution protocol (running on a swactor runtime)
|
/// Bridges the actorized distribution protocol (running on a swactor runtime)
|
||||||
|
|
@ -236,8 +271,13 @@ pub struct IrohDriver {
|
||||||
dialing: Arc<Mutex<HashSet<NodeId>>>,
|
dialing: Arc<Mutex<HashSet<NodeId>>>,
|
||||||
/// Connections accepted by the background accept loop (SWIM ALPN).
|
/// Connections accepted by the background accept loop (SWIM ALPN).
|
||||||
accepted_conns: Arc<Mutex<Vec<(NodeId, Connection)>>>,
|
accepted_conns: Arc<Mutex<Vec<(NodeId, Connection)>>>,
|
||||||
/// Connections accepted on non-SWIM ALPNs (streams, datastream, etc.).
|
/// Connections accepted on non-SWIM ALPNs before driver-owned adapters claim them.
|
||||||
other_accepted_conns: Arc<Mutex<Vec<(NodeId, Vec<u8>, Connection)>>>,
|
other_accepted_conns: Arc<Mutex<Vec<(NodeId, Vec<u8>, Connection)>>>,
|
||||||
|
/// Completed datastream QUIC reads from driver-owned DATASTREAM_ALPN adapters.
|
||||||
|
datastream_reads: Arc<Mutex<Vec<DatastreamQuicRead>>>,
|
||||||
|
/// Logical edge events emitted by driver-owned EDGE_ALPN byte pumps.
|
||||||
|
edge_events: Arc<Mutex<Vec<EdgeTransportEvent>>>,
|
||||||
|
next_edge_stream_group: u64,
|
||||||
/// Frames read by per-connection reader tasks, drained synchronously by
|
/// Frames read by per-connection reader tasks, drained synchronously by
|
||||||
/// `recv()` / `pump_inbound_to_actors()`. This decouples network reads from the
|
/// `recv()` / `pump_inbound_to_actors()`. This decouples network reads from the
|
||||||
/// state machine so `recv()`/`tick()` are pure-sync (no `block_on`) and can run
|
/// state machine so `recv()`/`tick()` are pure-sync (no `block_on`) and can run
|
||||||
|
|
@ -395,6 +435,9 @@ impl IrohDriver {
|
||||||
Arc::new(Mutex::new(Vec::new()));
|
Arc::new(Mutex::new(Vec::new()));
|
||||||
let other_accepted_conns: Arc<Mutex<Vec<(NodeId, Vec<u8>, Connection)>>> =
|
let other_accepted_conns: Arc<Mutex<Vec<(NodeId, Vec<u8>, Connection)>>> =
|
||||||
Arc::new(Mutex::new(Vec::new()));
|
Arc::new(Mutex::new(Vec::new()));
|
||||||
|
let datastream_reads: Arc<Mutex<Vec<DatastreamQuicRead>>> =
|
||||||
|
Arc::new(Mutex::new(Vec::new()));
|
||||||
|
let edge_events: Arc<Mutex<Vec<EdgeTransportEvent>>> = Arc::new(Mutex::new(Vec::new()));
|
||||||
{
|
{
|
||||||
let ep = endpoint.clone();
|
let ep = endpoint.clone();
|
||||||
let peer_auth = config.peer_auth.clone();
|
let peer_auth = config.peer_auth.clone();
|
||||||
|
|
@ -443,6 +486,9 @@ impl IrohDriver {
|
||||||
dialing: Arc::new(Mutex::new(HashSet::new())),
|
dialing: Arc::new(Mutex::new(HashSet::new())),
|
||||||
accepted_conns,
|
accepted_conns,
|
||||||
other_accepted_conns,
|
other_accepted_conns,
|
||||||
|
datastream_reads,
|
||||||
|
edge_events,
|
||||||
|
next_edge_stream_group: 1,
|
||||||
incoming: Arc::new(Mutex::new(Vec::new())),
|
incoming: Arc::new(Mutex::new(Vec::new())),
|
||||||
evict: Arc::new(Mutex::new(Vec::new())),
|
evict: Arc::new(Mutex::new(Vec::new())),
|
||||||
peer_relay_urls: HashMap::new(),
|
peer_relay_urls: HashMap::new(),
|
||||||
|
|
@ -501,6 +547,81 @@ impl IrohDriver {
|
||||||
drained
|
drained
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Claim accepted datastream connections and read them inside driver-owned tasks.
|
||||||
|
pub fn pump_datastream_ingress(&mut self) {
|
||||||
|
for (_node, conn) in self.drain_accepted_for_alpn(DATASTREAM_ALPN) {
|
||||||
|
let reads = Arc::clone(&self.datastream_reads);
|
||||||
|
self.rt.spawn(async move {
|
||||||
|
while let Ok(recv) = conn.accept_uni().await {
|
||||||
|
match read_events_from_stream(recv).await {
|
||||||
|
Ok(read) => reads.lock().push(read),
|
||||||
|
Err(_) => break,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
});
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Drain decoded datastream QUIC reads emitted by driver-owned adapter tasks.
|
||||||
|
pub fn drain_datastream_reads(&self) -> Vec<DatastreamQuicRead> {
|
||||||
|
self.datastream_reads.lock().drain(..).collect()
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Start a driver-owned datastream subscription writer task.
|
||||||
|
pub fn publish_datastream_subscription(
|
||||||
|
&self,
|
||||||
|
peer: EndpointAddr,
|
||||||
|
header: DatastreamQuicHeader,
|
||||||
|
subscription: datastream::DatastreamSubscription,
|
||||||
|
idle_sleep: Duration,
|
||||||
|
) {
|
||||||
|
let _ = spawn_subscription_writer(
|
||||||
|
&self.rt,
|
||||||
|
self.endpoint.clone(),
|
||||||
|
peer,
|
||||||
|
header,
|
||||||
|
subscription,
|
||||||
|
idle_sleep,
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Return a cloneable logical datastream transport handle for publisher actors.
|
||||||
|
pub fn datastream_publish_handle(&self) -> DatastreamPublishHandle {
|
||||||
|
DatastreamPublishHandle {
|
||||||
|
rt: self.rt.clone(),
|
||||||
|
endpoint: self.endpoint.clone(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Claim accepted MVP edge connections and read opaque edge bytes inside the driver.
|
||||||
|
pub fn pump_edge_ingress(&mut self) {
|
||||||
|
for (node, conn) in self.drain_accepted_for_alpn(EDGE_ALPN) {
|
||||||
|
let stream_group = self.next_edge_stream_group;
|
||||||
|
self.next_edge_stream_group = self.next_edge_stream_group.saturating_add(1).max(1);
|
||||||
|
spawn_edge_recv_pump(
|
||||||
|
self.rt.clone(),
|
||||||
|
conn,
|
||||||
|
node,
|
||||||
|
Arc::clone(&self.edge_events),
|
||||||
|
stream_group,
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Drain logical edge transport events emitted by driver-owned byte pumps.
|
||||||
|
pub fn drain_edge_events(&self) -> Vec<EdgeTransportEvent> {
|
||||||
|
self.edge_events.lock().drain(..).collect()
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Start a driver-owned EDGE_ALPN send pump and return its logical byte input handle.
|
||||||
|
pub fn spawn_edge_send_pump(
|
||||||
|
&self,
|
||||||
|
peer: EndpointAddr,
|
||||||
|
edge_id: u64,
|
||||||
|
) -> Result<EdgeSendHandle, String> {
|
||||||
|
spawn_edge_sender_task(self.rt.clone(), self.endpoint.clone(), peer, edge_id)
|
||||||
|
}
|
||||||
|
|
||||||
/// The node's identity.
|
/// The node's identity.
|
||||||
pub fn node_id(&self) -> NodeId {
|
pub fn node_id(&self) -> NodeId {
|
||||||
self.keypair.node_id()
|
self.keypair.node_id()
|
||||||
|
|
|
||||||
|
|
@ -5,12 +5,16 @@
|
||||||
//! wire message definitions.
|
//! wire message definitions.
|
||||||
|
|
||||||
pub mod datastream_transport;
|
pub mod datastream_transport;
|
||||||
|
pub mod edge_transport;
|
||||||
pub mod iroh_driver;
|
pub mod iroh_driver;
|
||||||
|
|
||||||
pub use iroh_driver::{
|
pub use iroh_driver::{
|
||||||
ConnType, IrohDriver, IrohDriverConfig, JoinPhase, JoinStatus, conn_type_of, discover_lan_ips,
|
ConnType, DatastreamPublishHandle, IrohDriver, IrohDriverConfig, JoinPhase, JoinStatus,
|
||||||
|
conn_type_of, discover_lan_ips,
|
||||||
};
|
};
|
||||||
|
|
||||||
|
pub use edge_transport::{EDGE_ALPN, EdgeSendHandle, EdgeTransportEvent, EdgeTransportFault};
|
||||||
|
|
||||||
pub use datastream_transport::{
|
pub use datastream_transport::{
|
||||||
DATASTREAM_ALPN, DatastreamQuicHeader, DatastreamQuicRead, DatastreamQuicWriteStats,
|
DATASTREAM_ALPN, DatastreamQuicHeader, DatastreamQuicRead, DatastreamQuicWriteStats,
|
||||||
read_events_from_stream, read_next_event, read_next_uni_from_connection, read_stream_header,
|
read_events_from_stream, read_next_event, read_next_uni_from_connection, read_stream_header,
|
||||||
|
|
|
||||||
|
|
@ -1,75 +0,0 @@
|
||||||
# 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.
|
|
||||||
|
|
||||||
|
|
@ -7,6 +7,7 @@ autobins = false
|
||||||
|
|
||||||
[features]
|
[features]
|
||||||
default = []
|
default = []
|
||||||
|
dashboard = []
|
||||||
|
|
||||||
[dependencies]
|
[dependencies]
|
||||||
datastream = { path = "../datastream" }
|
datastream = { path = "../datastream" }
|
||||||
|
|
@ -28,9 +29,6 @@ toml = "0.8"
|
||||||
libc = "0.2"
|
libc = "0.2"
|
||||||
signal-hook = "0.3"
|
signal-hook = "0.3"
|
||||||
|
|
||||||
[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"
|
||||||
|
|
@ -43,23 +41,6 @@ path = "src/bin/orchestrator.rs"
|
||||||
name = "mvp-chat"
|
name = "mvp-chat"
|
||||||
path = "src/bin/mvp_chat.rs"
|
path = "src/bin/mvp_chat.rs"
|
||||||
|
|
||||||
[[test]]
|
|
||||||
name = "local_unmocked_mvp_e2e"
|
|
||||||
path = "tests/local_unmocked_mvp_e2e.rs"
|
|
||||||
harness = false
|
|
||||||
required-features = ["local-e2e"]
|
|
||||||
|
|
||||||
[[test]]
|
|
||||||
name = "gpu_worker_node_e2e"
|
|
||||||
path = "tests/gpu_worker_node_e2e.rs"
|
|
||||||
required-features = ["local-e2e"]
|
|
||||||
|
|
||||||
[[test]]
|
|
||||||
name = "local-e2e-cluster"
|
|
||||||
path = "tests/local_e2e_cluster.rs"
|
|
||||||
harness = false
|
|
||||||
required-features = ["local-e2e"]
|
|
||||||
|
|
||||||
[[test]]
|
[[test]]
|
||||||
name = "mvp_chat_mock"
|
name = "mvp_chat_mock"
|
||||||
path = "tests/mvp_chat_mock.rs"
|
path = "tests/mvp_chat_mock.rs"
|
||||||
|
|
|
||||||
|
|
@ -1,882 +0,0 @@
|
||||||
# MVP Node Provisioning Specification
|
|
||||||
***STALE! FOR HISTORICAL REFERENCE ONLY***
|
|
||||||
|
|
||||||
**Status:**draft node-provisioning specification.
|
|
||||||
|
|
||||||
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
|
|
||||||
leasing, SSH bootstrap, stdout/stderr collection, handoff, and known-lease
|
|
||||||
teardown.
|
|
||||||
|
|
||||||
It intentionally does not define general run management, automatic replacement,
|
|
||||||
provider recovery, or a second post-handoff health system.
|
|
||||||
|
|
||||||
---
|
|
||||||
|
|
||||||
## 1. Purpose
|
|
||||||
|
|
||||||
The MVP needs to rent GPU nodes, bring them to the point where swactor can manage
|
|
||||||
them, and then stop managing them through SSH.
|
|
||||||
|
|
||||||
The intended path is:
|
|
||||||
|
|
||||||
```text
|
|
||||||
static runplan
|
|
||||||
-> logical node specs
|
|
||||||
-> one NodeManager actor per logical node
|
|
||||||
-> provider plugin creates a lease
|
|
||||||
-> BootstrapSession holds SSH until swactor convergence
|
|
||||||
-> stdout/stderr flows into datastream
|
|
||||||
-> remote swactor joins
|
|
||||||
-> BootstrapSession closes SSH and exits
|
|
||||||
-> NodeManager becomes dormant and keeps lease state for teardown
|
|
||||||
```
|
|
||||||
|
|
||||||
The actor system owns state transitions. Provider plugins perform provider I/O.
|
|
||||||
SSH bootstrap is a temporary pre-swactor transport, not long-term node
|
|
||||||
management.
|
|
||||||
|
|
||||||
---
|
|
||||||
|
|
||||||
## 2. Scope
|
|
||||||
|
|
||||||
In scope:
|
|
||||||
|
|
||||||
- static runplan node-group shape
|
|
||||||
- expansion of node groups into logical node specs
|
|
||||||
- one `NodeManager` actor per logical node
|
|
||||||
- node-local inventory/ledger owned by `NodeManager`
|
|
||||||
- readiness as a `NodeManager` state flag
|
|
||||||
- stateless provider plugin boundary for Vast.ai
|
|
||||||
- transient `BootstrapSession` for SSH, remote boot observation, and swactor
|
|
||||||
startup
|
|
||||||
- stdout/stderr forwarding from bootstrap SSH to datastream
|
|
||||||
- handoff from SSH bootstrap to swactor control
|
|
||||||
- teardown of leases already known to `NodeManager`
|
|
||||||
|
|
||||||
Out of scope:
|
|
||||||
|
|
||||||
- automatic replacement after lease, bootstrap, or runtime failure
|
|
||||||
- provider-label recovery or hidden provider scans
|
|
||||||
- post-handoff heartbeat layer outside swactor
|
|
||||||
- bidding/account/billing policy beyond selecting and destroying leases
|
|
||||||
- repairing a node after it disappears from swactor
|
|
||||||
|
|
||||||
---
|
|
||||||
|
|
||||||
## 3. Design Commitments
|
|
||||||
|
|
||||||
`NodeManager` owns one node. It owns both the node finite-state machine and that
|
|
||||||
node's inventory record.
|
|
||||||
|
|
||||||
There is no central actor that owns the run. A short-lived bootstrap procedure may
|
|
||||||
expand a runplan and spawn node actors, but it does not retain authority over
|
|
||||||
node state.
|
|
||||||
|
|
||||||
Readiness is node-local. A node is ready only when its `NodeManager` has recorded
|
|
||||||
successful swactor handoff.
|
|
||||||
|
|
||||||
The provider plugin is stateless with respect to the run. It maps desired node
|
|
||||||
shape to provider API calls and maps known lease handles to destroy calls.
|
|
||||||
|
|
||||||
`BootstrapSession` owns SSH and early stdout/stderr. It exists only between
|
|
||||||
provider endpoint availability and swactor convergence.
|
|
||||||
|
|
||||||
After swactor convergence, swactor is the live control path. The MVP does not add
|
|
||||||
another liveness or heartbeat system.
|
|
||||||
|
|
||||||
Known lease state is explicit. Teardown uses only lease handles already recorded
|
|
||||||
by `NodeManager`.
|
|
||||||
|
|
||||||
---
|
|
||||||
|
|
||||||
## 4. Identifiers
|
|
||||||
|
|
||||||
Identifier types are schematic. Concrete Rust APIs may wrap these as newtypes.
|
|
||||||
|
|
||||||
```rust
|
|
||||||
struct RunId(u64);
|
|
||||||
struct LogicalNodeId(String); // e.g. "workers-0"
|
|
||||||
struct NodeGroupId(String); // e.g. "workers"
|
|
||||||
struct RoleId(String); // e.g. "worker"
|
|
||||||
struct ProviderLeaseId(String); // e.g. "vastai:123456"
|
|
||||||
struct SwactorId(String);
|
|
||||||
struct BootstrapSessionId(u64);
|
|
||||||
struct DatastreamStreamId(String);
|
|
||||||
```
|
|
||||||
|
|
||||||
`LogicalNodeId` is stable for the run. It is assigned before provisioning and is
|
|
||||||
used to correlate provider lease, SSH bootstrap logs, and swactor identity.
|
|
||||||
|
|
||||||
`ProviderLeaseId` names the external billing/lease resource. For Vast.ai it wraps
|
|
||||||
the contract id.
|
|
||||||
|
|
||||||
`SwactorId` is not known until the remote runtime joins.
|
|
||||||
|
|
||||||
---
|
|
||||||
|
|
||||||
## 5. Static Runplan Node Shape
|
|
||||||
|
|
||||||
The runplan describes desired node groups. It does not describe provider API
|
|
||||||
steps or SSH polling details.
|
|
||||||
|
|
||||||
```rust
|
|
||||||
struct RunNodeGroupSpec {
|
|
||||||
run_id: RunId,
|
|
||||||
group_id: NodeGroupId,
|
|
||||||
role: RoleId,
|
|
||||||
count: u32,
|
|
||||||
provider: ProviderKind,
|
|
||||||
shape: DesiredNodeShape,
|
|
||||||
boot: BootSpec,
|
|
||||||
swarm_join: SwarmJoinSpec,
|
|
||||||
}
|
|
||||||
```
|
|
||||||
|
|
||||||
Provider-neutral desired shape:
|
|
||||||
|
|
||||||
```rust
|
|
||||||
struct DesiredNodeShape {
|
|
||||||
image: String,
|
|
||||||
disk_gb: u32,
|
|
||||||
gpu_name: Option<String>,
|
|
||||||
min_gpu_ram_mb: Option<u64>,
|
|
||||||
min_down_mbps: Option<f64>,
|
|
||||||
min_up_mbps: Option<f64>,
|
|
||||||
min_reliability: Option<f64>,
|
|
||||||
require_verified: bool,
|
|
||||||
provider_labels: BTreeMap<String, String>,
|
|
||||||
}
|
|
||||||
```
|
|
||||||
|
|
||||||
Remote boot specification:
|
|
||||||
|
|
||||||
```rust
|
|
||||||
struct BootSpec {
|
|
||||||
ssh_user: String,
|
|
||||||
verify_commands: Vec<String>,
|
|
||||||
start_swactor_command: String,
|
|
||||||
stdout_sources: Vec<String>,
|
|
||||||
stderr_sources: Vec<String>,
|
|
||||||
timeout_policy: BootstrapTimeoutPolicy,
|
|
||||||
}
|
|
||||||
```
|
|
||||||
|
|
||||||
Swarm join material:
|
|
||||||
|
|
||||||
```rust
|
|
||||||
struct SwarmJoinSpec {
|
|
||||||
orch_swactor_addr: String,
|
|
||||||
join_token_ref: String,
|
|
||||||
expected_logical_node_id: LogicalNodeId,
|
|
||||||
}
|
|
||||||
```
|
|
||||||
|
|
||||||
A bootstrap procedure expands each group into logical specs:
|
|
||||||
|
|
||||||
```text
|
|
||||||
workers count=3
|
|
||||||
-> workers-0
|
|
||||||
-> workers-1
|
|
||||||
-> workers-2
|
|
||||||
```
|
|
||||||
|
|
||||||
Each expanded logical spec starts one `NodeManager` actor.
|
|
||||||
|
|
||||||
---
|
|
||||||
|
|
||||||
## 6. Runtime Topology
|
|
||||||
|
|
||||||
|
|
||||||
The orchestrator host runs the swactor actor runtime and a datastream producer.
|
|
||||||
|
|
||||||
```text
|
|
||||||
orchestrator host
|
|
||||||
Swactor actor runtime
|
|
||||||
NodeManager(workers-0)
|
|
||||||
BootstrapSession(workers-0) while pre-handoff
|
|
||||||
NodeManager(workers-1)
|
|
||||||
BootstrapSession(workers-1) while pre-handoff
|
|
||||||
|
|
||||||
Orchestrator control endpoint
|
|
||||||
receives remote runtime joins
|
|
||||||
owns post-handoff actor communication
|
|
||||||
|
|
||||||
Provider plugins
|
|
||||||
VastAiPlugin, called by NodeManager
|
|
||||||
|
|
||||||
Datastream
|
|
||||||
receives bootstrap stdout/stderr records
|
|
||||||
```
|
|
||||||
|
|
||||||
---
|
|
||||||
|
|
||||||
## 7. NodeManager Actor
|
|
||||||
|
|
||||||
### 7.1 Responsibility
|
|
||||||
|
|
||||||
`NodeManager` owns one logical node's state and lifecycle.
|
|
||||||
|
|
||||||
It:
|
|
||||||
|
|
||||||
- stores desired node spec
|
|
||||||
- requests a provider lease
|
|
||||||
- records lease facts
|
|
||||||
- waits for provider endpoint facts when needed
|
|
||||||
- starts a `BootstrapSession`
|
|
||||||
- records compact bootstrap observations
|
|
||||||
- records swactor identity on join
|
|
||||||
- sets `ready = true` after handoff
|
|
||||||
- keeps known lease state while dormant
|
|
||||||
- releases its known lease on `Destroy`
|
|
||||||
|
|
||||||
It does not:
|
|
||||||
|
|
||||||
- aggregate run readiness
|
|
||||||
- replace failed nodes
|
|
||||||
- poll post-handoff liveness
|
|
||||||
- own provider search state after a plugin call returns
|
|
||||||
- store full stdout/stderr logs
|
|
||||||
|
|
||||||
### 7.2 Node Record
|
|
||||||
|
|
||||||
```rust
|
|
||||||
struct NodeRecord {
|
|
||||||
logical_node_id: LogicalNodeId,
|
|
||||||
run_id: RunId,
|
|
||||||
group_id: NodeGroupId,
|
|
||||||
role: RoleId,
|
|
||||||
|
|
||||||
desired: LogicalNodeSpec,
|
|
||||||
stage: NodeStage,
|
|
||||||
ready: bool,
|
|
||||||
|
|
||||||
lease: Option<LeaseFacts>,
|
|
||||||
connection: Option<SshEndpoint>,
|
|
||||||
bootstrap: Option<BootstrapFacts>,
|
|
||||||
swactor: Option<SwactorFacts>,
|
|
||||||
|
|
||||||
failed_reason: Option<String>,
|
|
||||||
destroyed_at: Option<SystemTime>,
|
|
||||||
}
|
|
||||||
```
|
|
||||||
|
|
||||||
Provider lease facts:
|
|
||||||
|
|
||||||
```rust
|
|
||||||
struct LeaseFacts {
|
|
||||||
provider: ProviderKind,
|
|
||||||
lease_id: ProviderLeaseId,
|
|
||||||
provider_contract_id: String,
|
|
||||||
offer_id: Option<String>,
|
|
||||||
destroy_handle: DestroyHandle,
|
|
||||||
provider_metadata: BTreeMap<String, String>,
|
|
||||||
}
|
|
||||||
```
|
|
||||||
|
|
||||||
SSH endpoint:
|
|
||||||
|
|
||||||
```rust
|
|
||||||
struct SshEndpoint {
|
|
||||||
host: String,
|
|
||||||
port: u16,
|
|
||||||
user: String,
|
|
||||||
auth_ref: String,
|
|
||||||
}
|
|
||||||
```
|
|
||||||
|
|
||||||
Bootstrap facts are compact. Full logs belong to datastream.
|
|
||||||
|
|
||||||
```rust
|
|
||||||
struct BootstrapFacts {
|
|
||||||
session_id: BootstrapSessionId,
|
|
||||||
last_stage: BootstrapStage,
|
|
||||||
last_stdout_seq: Option<u64>,
|
|
||||||
last_stderr_seq: Option<u64>,
|
|
||||||
last_observed_at: SystemTime,
|
|
||||||
}
|
|
||||||
```
|
|
||||||
|
|
||||||
Swactor facts:
|
|
||||||
|
|
||||||
```rust
|
|
||||||
struct SwactorFacts {
|
|
||||||
swactor_id: SwactorId,
|
|
||||||
joined_at: SystemTime,
|
|
||||||
handed_off_at: Option<SystemTime>,
|
|
||||||
}
|
|
||||||
```
|
|
||||||
|
|
||||||
### 7.3 Node Stages
|
|
||||||
|
|
||||||
```text
|
|
||||||
New
|
|
||||||
-> LeaseRequested
|
|
||||||
-> LeaseCreated
|
|
||||||
-> EndpointKnown
|
|
||||||
-> BootstrapRunning
|
|
||||||
-> SwactorJoined
|
|
||||||
-> HandedOff
|
|
||||||
-> Dormant
|
|
||||||
```
|
|
||||||
|
|
||||||
Terminal stages:
|
|
||||||
|
|
||||||
```text
|
|
||||||
Failed
|
|
||||||
Destroyed
|
|
||||||
```
|
|
||||||
|
|
||||||
`ready = true` only after handoff has completed. `Dormant` means the actor keeps
|
|
||||||
state for query and teardown but performs no polling, heartbeating, or repair.
|
|
||||||
|
|
||||||
### 7.4 Inbound Messages
|
|
||||||
|
|
||||||
Messages are logical actor signals. Some implementations may deliver provider
|
|
||||||
results as awaited futures and then enqueue the equivalent event to the actor FSM.
|
|
||||||
|
|
||||||
```rust
|
|
||||||
enum NodeManagerMsg {
|
|
||||||
Start(LogicalNodeSpec),
|
|
||||||
|
|
||||||
LeaseCreated(LeaseFacts, Option<SshEndpoint>),
|
|
||||||
LeaseFailed(String),
|
|
||||||
EndpointKnown(SshEndpoint),
|
|
||||||
EndpointFailed(String),
|
|
||||||
|
|
||||||
BootstrapObserved(BootstrapObservation),
|
|
||||||
BootstrapFailed(String),
|
|
||||||
BootstrapClosed,
|
|
||||||
|
|
||||||
SwactorJoined { swactor_id: SwactorId },
|
|
||||||
HandoffComplete { swactor_id: SwactorId },
|
|
||||||
|
|
||||||
Destroy,
|
|
||||||
GetStatus { reply_to: ActorAddress },
|
|
||||||
GetRecord { reply_to: ActorAddress },
|
|
||||||
}
|
|
||||||
```
|
|
||||||
|
|
||||||
### 7.5 Outbound Effects
|
|
||||||
|
|
||||||
`NodeManager` may perform these effects:
|
|
||||||
|
|
||||||
```text
|
|
||||||
ProviderPlugin.create_lease(shape)
|
|
||||||
ProviderPlugin.lookup_endpoint(lease)
|
|
||||||
spawn BootstrapSession(spec)
|
|
||||||
BootstrapSession.ConvergenceObserved(swactor_id)
|
|
||||||
ProviderPlugin.destroy_lease(destroy_handle)
|
|
||||||
reply with node status or record
|
|
||||||
```
|
|
||||||
|
|
||||||
It does not send full logs. `BootstrapSession` writes logs directly to
|
|
||||||
datastream.
|
|
||||||
|
|
||||||
---
|
|
||||||
|
|
||||||
## 8. NodeManager FSM Behavior
|
|
||||||
|
|
||||||
### 8.1 Start
|
|
||||||
|
|
||||||
On `Start(spec)`:
|
|
||||||
|
|
||||||
```text
|
|
||||||
record.desired = spec
|
|
||||||
record.stage = New
|
|
||||||
record.ready = false
|
|
||||||
record.failed_reason = None
|
|
||||||
```
|
|
||||||
|
|
||||||
Then:
|
|
||||||
|
|
||||||
```text
|
|
||||||
record.stage = LeaseRequested
|
|
||||||
call ProviderPlugin.create_lease(spec.shape)
|
|
||||||
```
|
|
||||||
|
|
||||||
### 8.2 Lease Result
|
|
||||||
|
|
||||||
On `LeaseCreated(lease, endpoint)`:
|
|
||||||
|
|
||||||
```text
|
|
||||||
record.lease = lease
|
|
||||||
record.stage = LeaseCreated
|
|
||||||
```
|
|
||||||
|
|
||||||
If `endpoint` is present:
|
|
||||||
|
|
||||||
```text
|
|
||||||
record.connection = endpoint
|
|
||||||
record.stage = EndpointKnown
|
|
||||||
spawn BootstrapSession
|
|
||||||
record.stage = BootstrapRunning
|
|
||||||
```
|
|
||||||
|
|
||||||
If `endpoint` is absent:
|
|
||||||
|
|
||||||
```text
|
|
||||||
call ProviderPlugin.lookup_endpoint(lease) until endpoint timeout or success
|
|
||||||
```
|
|
||||||
|
|
||||||
On `LeaseFailed(reason)`:
|
|
||||||
|
|
||||||
```text
|
|
||||||
record.stage = Failed
|
|
||||||
record.ready = false
|
|
||||||
record.failed_reason = reason
|
|
||||||
```
|
|
||||||
|
|
||||||
### 8.3 Endpoint Result
|
|
||||||
|
|
||||||
On `EndpointKnown(endpoint)`:
|
|
||||||
|
|
||||||
```text
|
|
||||||
record.connection = endpoint
|
|
||||||
record.stage = EndpointKnown
|
|
||||||
spawn BootstrapSession
|
|
||||||
record.stage = BootstrapRunning
|
|
||||||
```
|
|
||||||
|
|
||||||
On `EndpointFailed(reason)`:
|
|
||||||
|
|
||||||
```text
|
|
||||||
record.stage = Failed
|
|
||||||
record.ready = false
|
|
||||||
record.failed_reason = reason
|
|
||||||
```
|
|
||||||
|
|
||||||
A lease may still exist after endpoint failure. It is destroyed only when the
|
|
||||||
actor later receives `Destroy`.
|
|
||||||
|
|
||||||
### 8.4 Bootstrap Observations
|
|
||||||
|
|
||||||
On `BootstrapObserved(obs)`:
|
|
||||||
|
|
||||||
```text
|
|
||||||
record.bootstrap.last_stage = obs.stage
|
|
||||||
record.bootstrap.last_observed_at = now
|
|
||||||
record.bootstrap.last_stdout_seq = obs.last_stdout_seq if present
|
|
||||||
record.bootstrap.last_stderr_seq = obs.last_stderr_seq if present
|
|
||||||
```
|
|
||||||
|
|
||||||
The node stage remains `BootstrapRunning` until swactor join. Optional UI views
|
|
||||||
may display the finer bootstrap stage from `record.bootstrap.last_stage`.
|
|
||||||
|
|
||||||
On `BootstrapFailed(reason)`:
|
|
||||||
|
|
||||||
```text
|
|
||||||
record.stage = Failed
|
|
||||||
record.ready = false
|
|
||||||
record.failed_reason = reason
|
|
||||||
```
|
|
||||||
|
|
||||||
### 8.5 Swactor Join And Handoff
|
|
||||||
|
|
||||||
On `SwactorJoined { swactor_id }`:
|
|
||||||
|
|
||||||
```text
|
|
||||||
record.swactor.swactor_id = swactor_id
|
|
||||||
record.swactor.joined_at = now
|
|
||||||
record.stage = SwactorJoined
|
|
||||||
send BootstrapSession.ConvergenceObserved(swactor_id)
|
|
||||||
```
|
|
||||||
|
|
||||||
On `BootstrapClosed` after swactor join:
|
|
||||||
|
|
||||||
```text
|
|
||||||
record.swactor.handed_off_at = now
|
|
||||||
record.stage = HandedOff
|
|
||||||
record.ready = true
|
|
||||||
record.stage = Dormant
|
|
||||||
```
|
|
||||||
|
|
||||||
The SSH handle must be closed before `ready` becomes true.
|
|
||||||
|
|
||||||
### 8.6 Destroy
|
|
||||||
|
|
||||||
On `Destroy`:
|
|
||||||
|
|
||||||
```text
|
|
||||||
if BootstrapSession active:
|
|
||||||
cancel BootstrapSession
|
|
||||||
|
|
||||||
if lease exists and not destroyed:
|
|
||||||
call ProviderPlugin.destroy_lease(lease.destroy_handle)
|
|
||||||
|
|
||||||
record.stage = Destroyed on success
|
|
||||||
record.destroyed_at = now
|
|
||||||
record.ready = false
|
|
||||||
```
|
|
||||||
|
|
||||||
Destroy uses only the lease stored in `NodeRecord`. There is no provider scan.
|
|
||||||
|
|
||||||
---
|
|
||||||
|
|
||||||
## 9. Provider Plugin Boundary
|
|
||||||
|
|
||||||
Provider plugins are adapters. They are not run supervisors.
|
|
||||||
|
|
||||||
```rust
|
|
||||||
trait ProviderPlugin {
|
|
||||||
async fn create_lease(&self, request: CreateLeaseRequest)
|
|
||||||
-> Result<CreateLeaseResult, ProviderError>;
|
|
||||||
|
|
||||||
async fn lookup_endpoint(&self, lease: &LeaseFacts)
|
|
||||||
-> Result<Option<SshEndpoint>, ProviderError>;
|
|
||||||
|
|
||||||
async fn destroy_lease(&self, handle: &DestroyHandle)
|
|
||||||
-> Result<(), ProviderError>;
|
|
||||||
}
|
|
||||||
```
|
|
||||||
|
|
||||||
`create_lease` may search, filter, rank, and create a provider lease. For Vast.ai
|
|
||||||
this maps to offer search and instance creation.
|
|
||||||
|
|
||||||
`lookup_endpoint` may poll provider APIs until SSH endpoint facts are known. It
|
|
||||||
must not open SSH or inspect remote boot.
|
|
||||||
|
|
||||||
`destroy_lease` destroys a known provider lease.
|
|
||||||
|
|
||||||
Provider plugin output must include enough facts for teardown:
|
|
||||||
|
|
||||||
```rust
|
|
||||||
struct CreateLeaseResult {
|
|
||||||
lease: LeaseFacts,
|
|
||||||
endpoint: Option<SshEndpoint>,
|
|
||||||
}
|
|
||||||
```
|
|
||||||
|
|
||||||
The plugin must not:
|
|
||||||
|
|
||||||
- own `NodeRecord`
|
|
||||||
- stream stdout/stderr
|
|
||||||
- start swactor
|
|
||||||
- infer run readiness
|
|
||||||
- replace failed nodes
|
|
||||||
- recover unknown leases by provider label
|
|
||||||
|
|
||||||
---
|
|
||||||
|
|
||||||
## 10. Vast.ai Plugin Mapping
|
|
||||||
|
|
||||||
For Vast.ai, `CreateLeaseRequest` is derived from `DesiredNodeShape`:
|
|
||||||
|
|
||||||
```text
|
|
||||||
gpu_name -> SelectionPolicy.gpu_name
|
|
||||||
min_gpu_ram_mb -> SelectionPolicy.min_gpu_ram_mb
|
|
||||||
min_down_mbps -> SelectionPolicy.min_down_mbps
|
|
||||||
min_up_mbps -> SelectionPolicy.min_up_mbps
|
|
||||||
min_reliability -> SelectionPolicy.min_reliability
|
|
||||||
require_verified -> SelectionPolicy.require_verified
|
|
||||||
image -> CreateInstanceRequest.image
|
|
||||||
disk_gb -> CreateInstanceRequest.disk_gb
|
|
||||||
provider labels -> CreateInstanceRequest.label / env labels as needed
|
|
||||||
```
|
|
||||||
|
|
||||||
The plugin may wait until Vast.ai exposes a usable SSH endpoint. Once that
|
|
||||||
endpoint is returned, provider provisioning is complete from the plugin's
|
|
||||||
perspective.
|
|
||||||
|
|
||||||
The plugin does not determine whether the remote image booted correctly. That is
|
|
||||||
`BootstrapSession` work.
|
|
||||||
|
|
||||||
---
|
|
||||||
|
|
||||||
## 11. BootstrapSession
|
|
||||||
|
|
||||||
### 11.1 Responsibility
|
|
||||||
|
|
||||||
`BootstrapSession` is a transient child of `NodeManager`.
|
|
||||||
|
|
||||||
It owns:
|
|
||||||
|
|
||||||
- SSH connection attempts
|
|
||||||
- SSH handle
|
|
||||||
- remote bootstrap command handles
|
|
||||||
- stdout/stderr collection before swactor handoff
|
|
||||||
- boot verification commands
|
|
||||||
- swactor start command
|
|
||||||
- waiting for convergence acknowledgement
|
|
||||||
- closing SSH after handoff
|
|
||||||
|
|
||||||
It exits after either convergence or failure.
|
|
||||||
|
|
||||||
### 11.2 Input
|
|
||||||
|
|
||||||
```rust
|
|
||||||
struct BootstrapSessionSpec {
|
|
||||||
run_id: RunId,
|
|
||||||
logical_node_id: LogicalNodeId,
|
|
||||||
lease_id: ProviderLeaseId,
|
|
||||||
ssh: SshEndpoint,
|
|
||||||
boot: BootSpec,
|
|
||||||
swarm_join: SwarmJoinSpec,
|
|
||||||
datastream: DatastreamStreamId,
|
|
||||||
timeout_policy: BootstrapTimeoutPolicy,
|
|
||||||
}
|
|
||||||
```
|
|
||||||
|
|
||||||
### 11.3 Stages
|
|
||||||
|
|
||||||
```text
|
|
||||||
Created
|
|
||||||
-> SshConnecting
|
|
||||||
-> SshReady
|
|
||||||
-> StdoutStreaming
|
|
||||||
-> BootChecking
|
|
||||||
-> SwactorStarting
|
|
||||||
-> WaitingForSwactorJoin
|
|
||||||
-> Converged
|
|
||||||
-> Closed
|
|
||||||
```
|
|
||||||
|
|
||||||
Failure stages:
|
|
||||||
|
|
||||||
```text
|
|
||||||
SshTimeout
|
|
||||||
BootCheckFailed
|
|
||||||
StartFailed
|
|
||||||
JoinTimeout
|
|
||||||
StreamError
|
|
||||||
Cancelled
|
|
||||||
```
|
|
||||||
|
|
||||||
### 11.4 Behavior
|
|
||||||
|
|
||||||
1. Connect SSH until timeout.
|
|
||||||
2. Prove the machine is touchable by running a small command and reading output.
|
|
||||||
3. Start stdout/stderr capture for configured sources.
|
|
||||||
4. Emit full log records to datastream.
|
|
||||||
5. Emit compact observations to `NodeManager`.
|
|
||||||
6. Run boot verification commands.
|
|
||||||
7. Run or verify the swactor start command with the join spec.
|
|
||||||
8. Wait for convergence acknowledgement.
|
|
||||||
9. Flush datastream writes.
|
|
||||||
10. Close SSH.
|
|
||||||
11. Notify `NodeManager` with `BootstrapClosed`.
|
|
||||||
|
|
||||||
### 11.5 Datastream Records
|
|
||||||
|
|
||||||
Bootstrap logs use a stable stream per logical node.
|
|
||||||
|
|
||||||
```rust
|
|
||||||
struct BootstrapLogRecord {
|
|
||||||
run_id: RunId,
|
|
||||||
logical_node_id: LogicalNodeId,
|
|
||||||
lease_id: ProviderLeaseId,
|
|
||||||
source: BootstrapLogSource, // ssh-bootstrap
|
|
||||||
stream: BootstrapLogStream, // stdout | stderr
|
|
||||||
seq: u64,
|
|
||||||
timestamp: SystemTime,
|
|
||||||
line: String,
|
|
||||||
}
|
|
||||||
```
|
|
||||||
|
|
||||||
`NodeManager` stores only sequence numbers and the latest compact observation.
|
|
||||||
It does not retain log bodies.
|
|
||||||
|
|
||||||
### 11.6 Convergence Signal
|
|
||||||
|
|
||||||
Preferred signal path:
|
|
||||||
|
|
||||||
```text
|
|
||||||
remote swactor runtime -> orchestrator control endpoint -> NodeManager.SwactorJoined(swactor_id)
|
|
||||||
NodeManager -> BootstrapSession.ConvergenceObserved(swactor_id)
|
|
||||||
BootstrapSession -> NodeManager.BootstrapClosed
|
|
||||||
```
|
|
||||||
|
|
||||||
This keeps swactor membership authoritative while still letting
|
|
||||||
`BootstrapSession` close the SSH transport.
|
|
||||||
|
|
||||||
---
|
|
||||||
|
|
||||||
## 12. Handoff Contract
|
|
||||||
|
|
||||||
Handoff is complete only when all are true:
|
|
||||||
|
|
||||||
- the remote swactor runtime has joined the orchestrator-side actor runtime
|
|
||||||
- the actor runtime can address the remote by `SwactorId`
|
|
||||||
- `NodeManager` has recorded `SwactorFacts`
|
|
||||||
- bootstrap stdout/stderr records have been flushed
|
|
||||||
- SSH has been closed
|
|
||||||
- `NodeManager.ready == true`
|
|
||||||
- `NodeManager.stage == Dormant`
|
|
||||||
|
|
||||||
After handoff:
|
|
||||||
|
|
||||||
- `BootstrapSession` is gone
|
|
||||||
- `NodeManager` does not poll the node
|
|
||||||
- the swactor actor runtime owns live communication
|
|
||||||
- the node is considered usable by the MVP run
|
|
||||||
|
|
||||||
---
|
|
||||||
|
|
||||||
## 13. Failure Behavior
|
|
||||||
|
|
||||||
Failures before handoff are terminal for the logical node.
|
|
||||||
|
|
||||||
```text
|
|
||||||
LeaseFailed
|
|
||||||
EndpointFailed
|
|
||||||
BootstrapFailed
|
|
||||||
JoinTimeout
|
|
||||||
```
|
|
||||||
|
|
||||||
Terminal behavior:
|
|
||||||
|
|
||||||
```text
|
|
||||||
record.stage = Failed
|
|
||||||
record.ready = false
|
|
||||||
record.failed_reason = reason
|
|
||||||
```
|
|
||||||
|
|
||||||
The MVP does not create a replacement lease.
|
|
||||||
|
|
||||||
A failed node with a known lease is still eligible for explicit teardown through
|
|
||||||
`Destroy`.
|
|
||||||
|
|
||||||
Failures after handoff are handled by actor-runtime behavior. `NodeManager` is
|
|
||||||
dormant and does not repair the node. If the runtime loses the remote node, the
|
|
||||||
run stalls or fails according to existing swactor behavior.
|
|
||||||
|
|
||||||
---
|
|
||||||
|
|
||||||
## 14. Teardown Behavior
|
|
||||||
|
|
||||||
Teardown targets `NodeManager` actors.
|
|
||||||
|
|
||||||
```text
|
|
||||||
Teardown caller -> NodeManager.Destroy
|
|
||||||
NodeManager -> BootstrapSession.Cancel if active
|
|
||||||
NodeManager -> ProviderPlugin.destroy_lease if lease known
|
|
||||||
NodeManager records Destroyed
|
|
||||||
```
|
|
||||||
|
|
||||||
Rules:
|
|
||||||
|
|
||||||
- only known leases are destroyed
|
|
||||||
- no provider label sweep
|
|
||||||
- no hidden recovery of missing state
|
|
||||||
- destroy failure is recorded as node-local failure state
|
|
||||||
|
|
||||||
If the process lost all `NodeManager` state, this MVP spec does not define an
|
|
||||||
automatic cleanup path. Operator/provider-side cleanup remains manual for that
|
|
||||||
case.
|
|
||||||
|
|
||||||
---
|
|
||||||
|
|
||||||
## 15. Single End-To-End Worked Example
|
|
||||||
|
|
||||||
Input runplan node group:
|
|
||||||
|
|
||||||
```text
|
|
||||||
run_id: 42
|
|
||||||
group: workers
|
|
||||||
role: worker
|
|
||||||
count: 2
|
|
||||||
provider: vastai
|
|
||||||
shape:
|
|
||||||
image: ghcr.io/acme/mvp-worker:sha123
|
|
||||||
disk_gb: 80
|
|
||||||
gpu_name: RTX 4090
|
|
||||||
min_gpu_ram_mb: 20000
|
|
||||||
boot:
|
|
||||||
ssh_user: root
|
|
||||||
verify_commands:
|
|
||||||
- test -x /opt/mvp/swactor
|
|
||||||
start_swactor_command:
|
|
||||||
/opt/mvp/swactor-node --join ${ORCH_ADDR} --node ${LOGICAL_NODE_ID}
|
|
||||||
swarm_join:
|
|
||||||
orch_swactor_addr: quic://orch.example:9443
|
|
||||||
join_token_ref: secret://run-42-join-token
|
|
||||||
```
|
|
||||||
|
|
||||||
Expansion:
|
|
||||||
|
|
||||||
```text
|
|
||||||
workers-0
|
|
||||||
workers-1
|
|
||||||
```
|
|
||||||
|
|
||||||
For `workers-0`:
|
|
||||||
|
|
||||||
1. Bootstrap procedure spawns `NodeManager(workers-0)` with its logical spec.
|
|
||||||
2. `NodeManager` records `New`, then `LeaseRequested`.
|
|
||||||
3. `NodeManager` calls `VastAiPlugin.create_lease`.
|
|
||||||
4. `VastAiPlugin` searches offers, creates a Vast.ai instance, and returns:
|
|
||||||
|
|
||||||
```text
|
|
||||||
lease_id: vastai:123
|
|
||||||
contract_id: 123
|
|
||||||
offer_id: 9001
|
|
||||||
ssh: root@203.0.113.10:22001
|
|
||||||
```
|
|
||||||
|
|
||||||
5. `NodeManager` records `LeaseCreated`, `EndpointKnown`, then spawns
|
|
||||||
`BootstrapSession(workers-0)` and records `BootstrapRunning`.
|
|
||||||
6. `BootstrapSession` connects over SSH, runs a probe command, and sends:
|
|
||||||
|
|
||||||
```text
|
|
||||||
BootstrapObserved(stage=SshReady)
|
|
||||||
```
|
|
||||||
|
|
||||||
7. `BootstrapSession` streams bootstrap stdout/stderr into datastream:
|
|
||||||
|
|
||||||
```text
|
|
||||||
run=42 node=workers-0 stream=stdout seq=1 line="container boot entered"
|
|
||||||
run=42 node=workers-0 stream=stdout seq=2 line="swactor binary found"
|
|
||||||
```
|
|
||||||
|
|
||||||
8. `BootstrapSession` runs `test -x /opt/mvp/swactor`, then runs the swactor
|
|
||||||
start command with `LOGICAL_NODE_ID=workers-0` and the run join material.
|
|
||||||
9. Remote swactor runtime joins the orchestrator-side actor runtime.
|
|
||||||
10. The orchestrator control endpoint sends:
|
|
||||||
|
|
||||||
```text
|
|
||||||
NodeManager(workers-0).SwactorJoined(swactor_id=swactor-a7)
|
|
||||||
```
|
|
||||||
|
|
||||||
11. `NodeManager` records `SwactorJoined` and tells the bootstrap session that
|
|
||||||
convergence was observed.
|
|
||||||
12. `BootstrapSession` flushes datastream writes, closes SSH, and sends
|
|
||||||
`BootstrapClosed`.
|
|
||||||
13. `NodeManager` records:
|
|
||||||
|
|
||||||
```text
|
|
||||||
stage = Dormant
|
|
||||||
ready = true
|
|
||||||
swactor_id = swactor-a7
|
|
||||||
lease_id = vastai:123
|
|
||||||
```
|
|
||||||
|
|
||||||
The same sequence runs independently for `workers-1`.
|
|
||||||
|
|
||||||
The run bootstrap caller can determine node readiness by querying both
|
|
||||||
`NodeManager` actors:
|
|
||||||
|
|
||||||
```text
|
|
||||||
workers-0.ready == true
|
|
||||||
workers-1.ready == true
|
|
||||||
```
|
|
||||||
|
|
||||||
At run completion, teardown sends `Destroy` to both node managers. Each manager
|
|
||||||
destroys only its recorded Vast.ai contract and records `Destroyed`.
|
|
||||||
|
|
||||||
---
|
|
||||||
|
|
||||||
## 16. Implementation Boundaries
|
|
||||||
|
|
||||||
The code should preserve these boundaries even if local test plugins combine
|
|
||||||
steps for convenience:
|
|
||||||
|
|
||||||
- provider lease creation is not SSH bootstrap
|
|
||||||
- SSH bootstrap is not post-handoff supervision
|
|
||||||
- `NodeManager` state is the per-node ledger
|
|
||||||
- datastream owns log bodies
|
|
||||||
- the swactor actor runtime owns live communication after handoff
|
|
||||||
- teardown uses known lease handles only
|
|
||||||
|
|
||||||
A local Docker test provider may emit lease, endpoint, bootstrap, and swactor
|
|
||||||
join observations quickly, but the observations should still map onto the same
|
|
||||||
FSM stages. This keeps local tests aligned with Vast.ai behavior.
|
|
||||||
|
|
@ -1,4 +1,4 @@
|
||||||
//! swactor actor shells for the MVP system local E2E stack.
|
//! swactor actor shells for the MVP system runtime.
|
||||||
//!
|
//!
|
||||||
//! Each actor module owns its message type. Pure state machines remain in the
|
//! Each actor module owns its message type. Pure state machines remain in the
|
||||||
//! existing domain modules; actors translate mailbox messages into those cores and
|
//! existing domain modules; actors translate mailbox messages into those cores and
|
||||||
|
|
|
||||||
|
|
@ -1,4 +1,5 @@
|
||||||
use std::fs;
|
use std::collections::BTreeMap;
|
||||||
|
use std::fs::{self, File, OpenOptions};
|
||||||
use std::io::{self, BufRead, BufReader, IsTerminal, Write};
|
use std::io::{self, BufRead, BufReader, IsTerminal, Write};
|
||||||
use std::net::{Shutdown, TcpStream};
|
use std::net::{Shutdown, TcpStream};
|
||||||
#[cfg(all(target_os = "linux", not(test)))]
|
#[cfg(all(target_os = "linux", not(test)))]
|
||||||
|
|
@ -10,7 +11,12 @@ use std::sync::{Mutex, mpsc};
|
||||||
use std::thread;
|
use std::thread;
|
||||||
use std::time::{Duration, Instant};
|
use std::time::{Duration, Instant};
|
||||||
|
|
||||||
|
use datastream::{
|
||||||
|
ChannelContent, ChannelId, DatastreamEndpoint, DatastreamProducer, Frame, Lifetime, NodeId,
|
||||||
|
StreamDescriptor, StreamId, StreamOrigin,
|
||||||
|
};
|
||||||
use serde::Deserialize;
|
use serde::Deserialize;
|
||||||
|
use serde_json::{Value, json};
|
||||||
#[cfg(target_os = "linux")]
|
#[cfg(target_os = "linux")]
|
||||||
use signal_hook::consts::signal::{SIGINT, SIGTERM};
|
use signal_hook::consts::signal::{SIGINT, SIGTERM};
|
||||||
#[cfg(target_os = "linux")]
|
#[cfg(target_os = "linux")]
|
||||||
|
|
@ -30,6 +36,10 @@ const REPO_MODEL_CACHE_DIR: &str = ".model-cache";
|
||||||
const DEFAULT_MAX_TOKENS: u32 = 64;
|
const DEFAULT_MAX_TOKENS: u32 = 64;
|
||||||
const ORCH_SHUTDOWN_GRACE_MS: u64 = 5_000;
|
const ORCH_SHUTDOWN_GRACE_MS: u64 = 5_000;
|
||||||
const ORCH_SHUTDOWN_POLL_MS: u64 = 50;
|
const ORCH_SHUTDOWN_POLL_MS: u64 = 50;
|
||||||
|
const CHAT_LIFECYCLE_CHANNEL: &str = "mvp.chat.lifecycle";
|
||||||
|
const CHAT_RUNTIME_CHANNEL: &str = "mvp.chat.runtime";
|
||||||
|
const CHAT_PROMPT_CHANNEL: &str = "mvp.chat.prompt";
|
||||||
|
const CHAT_COMPONENT_CHANNEL: &str = "mvp.chat.component";
|
||||||
|
|
||||||
#[derive(Debug)]
|
#[derive(Debug)]
|
||||||
enum PromptInput {
|
enum PromptInput {
|
||||||
|
|
@ -63,19 +73,122 @@ where
|
||||||
I: IntoIterator<Item = String>,
|
I: IntoIterator<Item = String>,
|
||||||
{
|
{
|
||||||
let config = Config::from_args(args)?;
|
let config = Config::from_args(args)?;
|
||||||
|
let mut progress = ChatDatastream::new(1, config.datastream_frame_log.clone())?;
|
||||||
|
progress.emit(
|
||||||
|
CHAT_LIFECYCLE_CHANNEL,
|
||||||
|
"config",
|
||||||
|
"ready",
|
||||||
|
json!({
|
||||||
|
"provider": config.provider.as_str(),
|
||||||
|
"pipeline_stages": config.pipeline_stages,
|
||||||
|
"max_tokens": config.max_tokens,
|
||||||
|
"cached_model": config.cached_model.as_ref().map(|model| model.host_path.to_string_lossy().to_string()),
|
||||||
|
"dump_logs": config.datastream_frame_log.as_ref().map(|path| path.to_string_lossy().to_string()),
|
||||||
|
}),
|
||||||
|
);
|
||||||
confirm_vastai_if_needed(&config)?;
|
confirm_vastai_if_needed(&config)?;
|
||||||
let image_ref = prepare_runtime(&config)?;
|
let image_ref = match prepare_runtime(&config) {
|
||||||
let mut orch = OrchChild::spawn(&config, &image_ref)?;
|
Ok(image_ref) => {
|
||||||
|
progress.emit(
|
||||||
|
CHAT_RUNTIME_CHANNEL,
|
||||||
|
"prepare_runtime",
|
||||||
|
"ready",
|
||||||
|
json!({"image_ref": image_ref}),
|
||||||
|
);
|
||||||
|
image_ref
|
||||||
|
}
|
||||||
|
Err(error) => {
|
||||||
|
progress.emit(
|
||||||
|
CHAT_RUNTIME_CHANNEL,
|
||||||
|
"prepare_runtime",
|
||||||
|
"failed",
|
||||||
|
json!({"error": error}),
|
||||||
|
);
|
||||||
|
progress.archive_pending()?;
|
||||||
|
return Err(error);
|
||||||
|
}
|
||||||
|
};
|
||||||
|
let mut orch = match OrchChild::spawn(&config, &image_ref) {
|
||||||
|
Ok(orch) => {
|
||||||
|
progress.emit(
|
||||||
|
CHAT_COMPONENT_CHANNEL,
|
||||||
|
"orchestrator_process",
|
||||||
|
"started",
|
||||||
|
json!({"binary": config.orch_bin.to_string_lossy()}),
|
||||||
|
);
|
||||||
|
orch
|
||||||
|
}
|
||||||
|
Err(error) => {
|
||||||
|
progress.emit(
|
||||||
|
CHAT_COMPONENT_CHANNEL,
|
||||||
|
"orchestrator_process",
|
||||||
|
"failed",
|
||||||
|
json!({"error": error}),
|
||||||
|
);
|
||||||
|
progress.archive_pending()?;
|
||||||
|
return Err(error);
|
||||||
|
}
|
||||||
|
};
|
||||||
let rpc_addr = match orch.wait_ready(config.rpc_addr.clone()) {
|
let rpc_addr = match orch.wait_ready(config.rpc_addr.clone()) {
|
||||||
Ok(addr) => addr,
|
Ok(addr) => {
|
||||||
|
progress.emit(
|
||||||
|
CHAT_RUNTIME_CHANNEL,
|
||||||
|
"prompt_rpc",
|
||||||
|
"ready",
|
||||||
|
json!({"addr": addr}),
|
||||||
|
);
|
||||||
|
addr
|
||||||
|
}
|
||||||
Err(_) if STOP_REQUESTED.load(Ordering::SeqCst) => {
|
Err(_) if STOP_REQUESTED.load(Ordering::SeqCst) => {
|
||||||
|
progress.emit(
|
||||||
|
CHAT_LIFECYCLE_CHANNEL,
|
||||||
|
"shutdown",
|
||||||
|
"requested",
|
||||||
|
json!({"reason": "interrupted_before_ready"}),
|
||||||
|
);
|
||||||
orch.shutdown();
|
orch.shutdown();
|
||||||
|
progress.emit(
|
||||||
|
CHAT_COMPONENT_CHANNEL,
|
||||||
|
"orchestrator_process",
|
||||||
|
"stopped",
|
||||||
|
json!({"reason": "interrupted_before_ready"}),
|
||||||
|
);
|
||||||
|
progress.archive_pending()?;
|
||||||
return Ok(());
|
return Ok(());
|
||||||
}
|
}
|
||||||
Err(error) => return Err(error),
|
Err(error) => {
|
||||||
|
progress.emit(
|
||||||
|
CHAT_RUNTIME_CHANNEL,
|
||||||
|
"prompt_rpc",
|
||||||
|
"failed",
|
||||||
|
json!({"error": error}),
|
||||||
|
);
|
||||||
|
orch.shutdown();
|
||||||
|
progress.emit(
|
||||||
|
CHAT_COMPONENT_CHANNEL,
|
||||||
|
"orchestrator_process",
|
||||||
|
"stopped",
|
||||||
|
json!({"reason": "startup_failed"}),
|
||||||
|
);
|
||||||
|
progress.archive_pending()?;
|
||||||
|
return Err(error);
|
||||||
|
}
|
||||||
};
|
};
|
||||||
let result = run_chat_loop(&rpc_addr, config.max_tokens);
|
let result = run_chat_loop_with_progress(&rpc_addr, config.max_tokens, Some(&mut progress));
|
||||||
|
progress.emit(
|
||||||
|
CHAT_LIFECYCLE_CHANNEL,
|
||||||
|
"shutdown",
|
||||||
|
"requested",
|
||||||
|
json!({"reason": "prompt_loop_exited", "ok": result.is_ok()}),
|
||||||
|
);
|
||||||
orch.shutdown();
|
orch.shutdown();
|
||||||
|
progress.emit(
|
||||||
|
CHAT_COMPONENT_CHANNEL,
|
||||||
|
"orchestrator_process",
|
||||||
|
"stopped",
|
||||||
|
json!({"reason": "shutdown_requested"}),
|
||||||
|
);
|
||||||
|
progress.archive_pending()?;
|
||||||
result
|
result
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
@ -95,6 +208,174 @@ struct Config {
|
||||||
skip_rebuild: bool,
|
skip_rebuild: bool,
|
||||||
}
|
}
|
||||||
|
|
||||||
|
struct ChatDatastream {
|
||||||
|
stream: StreamId,
|
||||||
|
endpoint: DatastreamEndpoint,
|
||||||
|
producer: DatastreamProducer,
|
||||||
|
channels: BTreeMap<String, ChannelId>,
|
||||||
|
channel_names: BTreeMap<ChannelId, String>,
|
||||||
|
archive_path: Option<PathBuf>,
|
||||||
|
pending: Vec<(String, StreamId, String, Frame)>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl ChatDatastream {
|
||||||
|
fn new(run_id: u64, archive_path: Option<PathBuf>) -> Result<Self, String> {
|
||||||
|
let stream = StreamId::new(NodeId::new("mvp-chat"), Lifetime(run_id));
|
||||||
|
let endpoint = DatastreamEndpoint::with_descriptor(
|
||||||
|
StreamDescriptor {
|
||||||
|
stream: stream.clone(),
|
||||||
|
label: Some("mvp chat".to_owned()),
|
||||||
|
origin: StreamOrigin::Orchestrator,
|
||||||
|
},
|
||||||
|
1024,
|
||||||
|
256,
|
||||||
|
);
|
||||||
|
let producer = endpoint.producer();
|
||||||
|
let mut out = Self {
|
||||||
|
stream,
|
||||||
|
endpoint,
|
||||||
|
producer,
|
||||||
|
channels: BTreeMap::new(),
|
||||||
|
channel_names: BTreeMap::new(),
|
||||||
|
archive_path,
|
||||||
|
pending: Vec::new(),
|
||||||
|
};
|
||||||
|
for name in [
|
||||||
|
CHAT_LIFECYCLE_CHANNEL,
|
||||||
|
CHAT_RUNTIME_CHANNEL,
|
||||||
|
CHAT_PROMPT_CHANNEL,
|
||||||
|
CHAT_COMPONENT_CHANNEL,
|
||||||
|
] {
|
||||||
|
out.channel_by_name(name);
|
||||||
|
}
|
||||||
|
Ok(out)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn channel_by_name(&mut self, name: &str) -> ChannelId {
|
||||||
|
if let Some(id) = self.channels.get(name).copied() {
|
||||||
|
return id;
|
||||||
|
}
|
||||||
|
let id = self.producer.register_channel(
|
||||||
|
name,
|
||||||
|
ChannelContent::JsonRecord {
|
||||||
|
schema: Some(name.to_owned()),
|
||||||
|
},
|
||||||
|
);
|
||||||
|
self.channels.insert(name.to_owned(), id);
|
||||||
|
self.channel_names.insert(id, name.to_owned());
|
||||||
|
id
|
||||||
|
}
|
||||||
|
|
||||||
|
fn emit(&mut self, channel: &str, phase: &str, status: &str, detail: Value) {
|
||||||
|
let id = self.channel_by_name(channel);
|
||||||
|
let payload = serde_json::to_vec(&json!({
|
||||||
|
"type": "ChatProgress",
|
||||||
|
"phase": phase,
|
||||||
|
"status": status,
|
||||||
|
"detail": detail,
|
||||||
|
}))
|
||||||
|
.expect("serialize mvp-chat progress event");
|
||||||
|
self.producer.submit_bytes(id, payload);
|
||||||
|
self.flush();
|
||||||
|
}
|
||||||
|
|
||||||
|
fn flush(&mut self) {
|
||||||
|
let stream = self.stream.clone();
|
||||||
|
for frame in self.endpoint.mux().drain() {
|
||||||
|
let channel = self
|
||||||
|
.channel_names
|
||||||
|
.get(&frame.channel)
|
||||||
|
.cloned()
|
||||||
|
.unwrap_or_else(|| format!("channel#{}", frame.channel.0));
|
||||||
|
self.pending
|
||||||
|
.push(("mvp-chat".to_owned(), stream.clone(), channel, frame));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn archive_pending(&mut self) -> Result<(), String> {
|
||||||
|
let Some(path) = self.archive_path.as_deref() else {
|
||||||
|
self.pending.clear();
|
||||||
|
return Ok(());
|
||||||
|
};
|
||||||
|
if self.pending.is_empty() {
|
||||||
|
return Ok(());
|
||||||
|
}
|
||||||
|
let mut archive = ChatFrameArchive::open(path)?;
|
||||||
|
for (source, stream, channel, frame) in self.pending.drain(..) {
|
||||||
|
archive.record(&source, &stream, &channel, &frame)?;
|
||||||
|
}
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
struct ChatFrameArchive {
|
||||||
|
file: File,
|
||||||
|
next_seq: u64,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl ChatFrameArchive {
|
||||||
|
fn open(path: &Path) -> Result<Self, String> {
|
||||||
|
if let Some(parent) = path.parent()
|
||||||
|
&& !parent.as_os_str().is_empty()
|
||||||
|
{
|
||||||
|
fs::create_dir_all(parent).map_err(|e| {
|
||||||
|
format!(
|
||||||
|
"create mvp-chat datastream frame log dir {}: {e}",
|
||||||
|
parent.display()
|
||||||
|
)
|
||||||
|
})?;
|
||||||
|
}
|
||||||
|
let next_seq = match File::open(path) {
|
||||||
|
Ok(file) => BufReader::new(file).lines().count() as u64,
|
||||||
|
Err(error) if error.kind() == std::io::ErrorKind::NotFound => 0,
|
||||||
|
Err(error) => {
|
||||||
|
return Err(format!(
|
||||||
|
"read mvp-chat datastream frame log {}: {error}",
|
||||||
|
path.display()
|
||||||
|
));
|
||||||
|
}
|
||||||
|
};
|
||||||
|
let file = OpenOptions::new()
|
||||||
|
.create(true)
|
||||||
|
.append(true)
|
||||||
|
.open(path)
|
||||||
|
.map_err(|e| format!("open mvp-chat datastream frame log {}: {e}", path.display()))?;
|
||||||
|
Ok(Self { file, next_seq })
|
||||||
|
}
|
||||||
|
|
||||||
|
fn record(
|
||||||
|
&mut self,
|
||||||
|
source: &str,
|
||||||
|
stream: &StreamId,
|
||||||
|
channel: &str,
|
||||||
|
frame: &Frame,
|
||||||
|
) -> Result<(), String> {
|
||||||
|
let payload = match std::str::from_utf8(&frame.payload) {
|
||||||
|
Ok(text) => json!({"encoding": "utf8", "value": text}),
|
||||||
|
Err(_) => json!({"encoding": "bytes", "value": frame.payload}),
|
||||||
|
};
|
||||||
|
let record = json!({
|
||||||
|
"arrival_seq": self.next_seq,
|
||||||
|
"source": source,
|
||||||
|
"stream": stream.to_string(),
|
||||||
|
"channel": channel,
|
||||||
|
"channel_id": frame.channel.0,
|
||||||
|
"position": frame.position.0,
|
||||||
|
"payload": payload,
|
||||||
|
});
|
||||||
|
self.next_seq += 1;
|
||||||
|
let mut line = serde_json::to_vec(&record)
|
||||||
|
.map_err(|e| format!("serialize mvp-chat frame log: {e}"))?;
|
||||||
|
line.push(b'\n');
|
||||||
|
self.file
|
||||||
|
.write_all(&line)
|
||||||
|
.map_err(|e| format!("write mvp-chat frame log: {e}"))?;
|
||||||
|
self.file
|
||||||
|
.flush()
|
||||||
|
.map_err(|e| format!("flush mvp-chat frame log: {e}"))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
#[derive(Clone, Debug, Default, Deserialize)]
|
#[derive(Clone, Debug, Default, Deserialize)]
|
||||||
#[serde(default, deny_unknown_fields)]
|
#[serde(default, deny_unknown_fields)]
|
||||||
struct ChatTomlConfig {
|
struct ChatTomlConfig {
|
||||||
|
|
@ -255,6 +536,7 @@ impl Config {
|
||||||
self.max_tokens.to_string(),
|
self.max_tokens.to_string(),
|
||||||
"--pipeline-stages".to_owned(),
|
"--pipeline-stages".to_owned(),
|
||||||
self.pipeline_stages.to_string(),
|
self.pipeline_stages.to_string(),
|
||||||
|
"--no-dashboard".to_owned(),
|
||||||
];
|
];
|
||||||
if self.provider == ProviderKind::Process {
|
if self.provider == ProviderKind::Process {
|
||||||
args.extend([
|
args.extend([
|
||||||
|
|
@ -412,7 +694,7 @@ fn resolve_vastai_config(
|
||||||
node_image: &str,
|
node_image: &str,
|
||||||
) -> Result<ResolvedVastAiConfig, String> {
|
) -> Result<ResolvedVastAiConfig, String> {
|
||||||
ResolvedVastAiConfig {
|
ResolvedVastAiConfig {
|
||||||
api_key: first_non_empty([env_optional("VAST_API_KEY")]).unwrap_or_default(),
|
api_key: first_non_empty([env_optional("VASTAI_API_KEY")]).unwrap_or_default(),
|
||||||
relay_url: first_non_empty([file.relay_url.clone()]).unwrap_or_default(),
|
relay_url: first_non_empty([file.relay_url.clone()]).unwrap_or_default(),
|
||||||
image: node_image.to_owned(),
|
image: node_image.to_owned(),
|
||||||
bootstrap_command: first_non_empty([file.bootstrap_command.clone()]).unwrap_or_default(),
|
bootstrap_command: first_non_empty([file.bootstrap_command.clone()]).unwrap_or_default(),
|
||||||
|
|
@ -692,42 +974,70 @@ fn stdin_prompt_events() -> mpsc::Receiver<PromptInput> {
|
||||||
rx
|
rx
|
||||||
}
|
}
|
||||||
|
|
||||||
fn run_chat_loop(addr: &str, max_tokens: u32) -> Result<(), String> {
|
fn run_chat_loop_with_progress(
|
||||||
run_chat_loop_with_input(addr, max_tokens, stdin_prompt_events())
|
addr: &str,
|
||||||
|
max_tokens: u32,
|
||||||
|
progress: Option<&mut ChatDatastream>,
|
||||||
|
) -> Result<(), String> {
|
||||||
|
run_chat_loop_with_input_and_progress(addr, max_tokens, stdin_prompt_events(), progress)
|
||||||
}
|
}
|
||||||
|
|
||||||
fn run_chat_loop_with_input(
|
fn run_chat_loop_with_input_and_progress(
|
||||||
addr: &str,
|
addr: &str,
|
||||||
max_tokens: u32,
|
max_tokens: u32,
|
||||||
input_rx: mpsc::Receiver<PromptInput>,
|
input_rx: mpsc::Receiver<PromptInput>,
|
||||||
|
progress: Option<&mut ChatDatastream>,
|
||||||
) -> Result<(), String> {
|
) -> Result<(), String> {
|
||||||
let mut stream =
|
let mut progress = progress;
|
||||||
TcpStream::connect(addr).map_err(|e| format!("connect prompt RPC {addr}: {e}"))?;
|
emit_chat_progress(
|
||||||
let reader = BufReader::new(
|
&mut progress,
|
||||||
stream
|
CHAT_RUNTIME_CHANNEL,
|
||||||
.try_clone()
|
"prompt_rpc",
|
||||||
.map_err(|e| format!("clone prompt RPC stream: {e}"))?,
|
"connecting",
|
||||||
|
json!({"addr": addr}),
|
||||||
);
|
);
|
||||||
run_chat_session(&mut stream, reader, input_rx, max_tokens)
|
let mut stream = match TcpStream::connect(addr) {
|
||||||
}
|
Ok(stream) => {
|
||||||
|
emit_chat_progress(
|
||||||
fn run_chat_session<R, W>(
|
&mut progress,
|
||||||
writer: &mut W,
|
CHAT_RUNTIME_CHANNEL,
|
||||||
reader: R,
|
"prompt_rpc",
|
||||||
input_rx: mpsc::Receiver<PromptInput>,
|
"connected",
|
||||||
max_tokens: u32,
|
json!({"addr": addr}),
|
||||||
) -> Result<(), String>
|
);
|
||||||
where
|
stream
|
||||||
R: BufRead,
|
}
|
||||||
W: Write,
|
Err(error) => {
|
||||||
{
|
emit_chat_progress(
|
||||||
let mut output = io::stdout();
|
&mut progress,
|
||||||
run_chat_session_with_output(writer, reader, input_rx, max_tokens, &mut output)
|
CHAT_RUNTIME_CHANNEL,
|
||||||
|
"prompt_rpc",
|
||||||
|
"failed",
|
||||||
|
json!({"addr": addr, "error": error.to_string()}),
|
||||||
|
);
|
||||||
|
return Err(format!("connect prompt RPC {addr}: {error}"));
|
||||||
|
}
|
||||||
|
};
|
||||||
|
let reader = match stream.try_clone() {
|
||||||
|
Ok(stream) => BufReader::new(stream),
|
||||||
|
Err(error) => {
|
||||||
|
emit_chat_progress(
|
||||||
|
&mut progress,
|
||||||
|
CHAT_RUNTIME_CHANNEL,
|
||||||
|
"prompt_rpc_clone",
|
||||||
|
"failed",
|
||||||
|
json!({"error": error.to_string()}),
|
||||||
|
);
|
||||||
|
return Err(format!("clone prompt RPC stream: {error}"));
|
||||||
|
}
|
||||||
|
};
|
||||||
|
run_chat_session_with_progress(&mut stream, reader, input_rx, max_tokens, progress)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
fn run_chat_session_with_output<R, W, O>(
|
fn run_chat_session_with_output<R, W, O>(
|
||||||
writer: &mut W,
|
writer: &mut W,
|
||||||
mut reader: R,
|
reader: R,
|
||||||
input_rx: mpsc::Receiver<PromptInput>,
|
input_rx: mpsc::Receiver<PromptInput>,
|
||||||
max_tokens: u32,
|
max_tokens: u32,
|
||||||
output: &mut O,
|
output: &mut O,
|
||||||
|
|
@ -737,17 +1047,101 @@ where
|
||||||
W: Write,
|
W: Write,
|
||||||
O: Write,
|
O: Write,
|
||||||
{
|
{
|
||||||
|
run_chat_session_with_output_and_progress(writer, reader, input_rx, max_tokens, output, None)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn run_chat_session_with_progress<R, W>(
|
||||||
|
writer: &mut W,
|
||||||
|
reader: R,
|
||||||
|
input_rx: mpsc::Receiver<PromptInput>,
|
||||||
|
max_tokens: u32,
|
||||||
|
progress: Option<&mut ChatDatastream>,
|
||||||
|
) -> Result<(), String>
|
||||||
|
where
|
||||||
|
R: BufRead,
|
||||||
|
W: Write,
|
||||||
|
{
|
||||||
|
let mut output = io::stdout();
|
||||||
|
run_chat_session_with_output_and_progress(
|
||||||
|
writer,
|
||||||
|
reader,
|
||||||
|
input_rx,
|
||||||
|
max_tokens,
|
||||||
|
&mut output,
|
||||||
|
progress,
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn emit_chat_progress(
|
||||||
|
progress: &mut Option<&mut ChatDatastream>,
|
||||||
|
channel: &str,
|
||||||
|
phase: &str,
|
||||||
|
status: &str,
|
||||||
|
detail: Value,
|
||||||
|
) {
|
||||||
|
if let Some(progress) = progress.as_deref_mut() {
|
||||||
|
progress.emit(channel, phase, status, detail);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn run_chat_session_with_output_and_progress<R, W, O>(
|
||||||
|
writer: &mut W,
|
||||||
|
mut reader: R,
|
||||||
|
input_rx: mpsc::Receiver<PromptInput>,
|
||||||
|
max_tokens: u32,
|
||||||
|
output: &mut O,
|
||||||
|
progress: Option<&mut ChatDatastream>,
|
||||||
|
) -> Result<(), String>
|
||||||
|
where
|
||||||
|
R: BufRead,
|
||||||
|
W: Write,
|
||||||
|
O: Write,
|
||||||
|
{
|
||||||
|
let mut progress = progress;
|
||||||
let mut next_request_id = 1_u64;
|
let mut next_request_id = 1_u64;
|
||||||
|
|
||||||
loop {
|
loop {
|
||||||
if STOP_REQUESTED.load(Ordering::SeqCst) {
|
if STOP_REQUESTED.load(Ordering::SeqCst) {
|
||||||
|
emit_chat_progress(
|
||||||
|
&mut progress,
|
||||||
|
CHAT_PROMPT_CHANNEL,
|
||||||
|
"prompt_loop",
|
||||||
|
"exited",
|
||||||
|
json!({"reason": "stop_requested"}),
|
||||||
|
);
|
||||||
return Ok(());
|
return Ok(());
|
||||||
}
|
}
|
||||||
|
emit_chat_progress(
|
||||||
|
&mut progress,
|
||||||
|
CHAT_PROMPT_CHANNEL,
|
||||||
|
"waiting_for_prompt",
|
||||||
|
"started",
|
||||||
|
json!({"next_request_id": next_request_id}),
|
||||||
|
);
|
||||||
write!(output, "prompt:> ").map_err(|e| format!("write prompt: {e}"))?;
|
write!(output, "prompt:> ").map_err(|e| format!("write prompt: {e}"))?;
|
||||||
output.flush().map_err(|e| format!("flush prompt: {e}"))?;
|
output.flush().map_err(|e| format!("flush prompt: {e}"))?;
|
||||||
let prompt = match input_rx.recv() {
|
let prompt = match input_rx.recv() {
|
||||||
Ok(PromptInput::Line(line)) => line.trim_end().to_owned(),
|
Ok(PromptInput::Line(line)) => line.trim_end().to_owned(),
|
||||||
Ok(PromptInput::Closed | PromptInput::StopRequested) | Err(_) => return Ok(()),
|
Ok(PromptInput::Closed) | Err(_) => {
|
||||||
|
emit_chat_progress(
|
||||||
|
&mut progress,
|
||||||
|
CHAT_PROMPT_CHANNEL,
|
||||||
|
"prompt_loop",
|
||||||
|
"exited",
|
||||||
|
json!({"reason": "input_closed"}),
|
||||||
|
);
|
||||||
|
return Ok(());
|
||||||
|
}
|
||||||
|
Ok(PromptInput::StopRequested) => {
|
||||||
|
emit_chat_progress(
|
||||||
|
&mut progress,
|
||||||
|
CHAT_PROMPT_CHANNEL,
|
||||||
|
"prompt_loop",
|
||||||
|
"exited",
|
||||||
|
json!({"reason": "stop_requested"}),
|
||||||
|
);
|
||||||
|
return Ok(());
|
||||||
|
}
|
||||||
};
|
};
|
||||||
if prompt.trim().is_empty() {
|
if prompt.trim().is_empty() {
|
||||||
continue;
|
continue;
|
||||||
|
|
@ -755,6 +1149,13 @@ where
|
||||||
|
|
||||||
let request_id = next_request_id;
|
let request_id = next_request_id;
|
||||||
next_request_id = next_request_id.wrapping_add(1).max(1);
|
next_request_id = next_request_id.wrapping_add(1).max(1);
|
||||||
|
emit_chat_progress(
|
||||||
|
&mut progress,
|
||||||
|
CHAT_PROMPT_CHANNEL,
|
||||||
|
"prompt_submitted",
|
||||||
|
"ready",
|
||||||
|
json!({"request_id": request_id, "prompt_bytes": prompt.len(), "max_tokens": max_tokens}),
|
||||||
|
);
|
||||||
write_json_line(
|
write_json_line(
|
||||||
writer,
|
writer,
|
||||||
&SubmitPrompt {
|
&SubmitPrompt {
|
||||||
|
|
@ -764,22 +1165,72 @@ where
|
||||||
},
|
},
|
||||||
)?;
|
)?;
|
||||||
writeln!(output, "decoding...").map_err(|e| format!("write decoding marker: {e}"))?;
|
writeln!(output, "decoding...").map_err(|e| format!("write decoding marker: {e}"))?;
|
||||||
|
emit_chat_progress(
|
||||||
|
&mut progress,
|
||||||
|
CHAT_PROMPT_CHANNEL,
|
||||||
|
"decoding",
|
||||||
|
"started",
|
||||||
|
json!({"request_id": request_id}),
|
||||||
|
);
|
||||||
let mut response_started = false;
|
let mut response_started = false;
|
||||||
|
|
||||||
loop {
|
loop {
|
||||||
if STOP_REQUESTED.load(Ordering::SeqCst) {
|
if STOP_REQUESTED.load(Ordering::SeqCst) {
|
||||||
|
emit_chat_progress(
|
||||||
|
&mut progress,
|
||||||
|
CHAT_PROMPT_CHANNEL,
|
||||||
|
"prompt_loop",
|
||||||
|
"exited",
|
||||||
|
json!({"reason": "stop_requested"}),
|
||||||
|
);
|
||||||
return Ok(());
|
return Ok(());
|
||||||
}
|
}
|
||||||
let mut line = String::new();
|
let mut line = String::new();
|
||||||
match reader.read_line(&mut line) {
|
match reader.read_line(&mut line) {
|
||||||
Ok(0) => return Err("prompt RPC closed".to_owned()),
|
Ok(0) => {
|
||||||
|
emit_chat_progress(
|
||||||
|
&mut progress,
|
||||||
|
CHAT_PROMPT_CHANNEL,
|
||||||
|
"prompt_rpc",
|
||||||
|
"failed",
|
||||||
|
json!({"request_id": request_id, "error": "prompt RPC closed"}),
|
||||||
|
);
|
||||||
|
return Err("prompt RPC closed".to_owned());
|
||||||
|
}
|
||||||
Ok(_) => {}
|
Ok(_) => {}
|
||||||
Err(error) => return Err(format!("read prompt RPC event: {error}")),
|
Err(error) => {
|
||||||
|
emit_chat_progress(
|
||||||
|
&mut progress,
|
||||||
|
CHAT_PROMPT_CHANNEL,
|
||||||
|
"prompt_rpc",
|
||||||
|
"failed",
|
||||||
|
json!({"request_id": request_id, "error": error.to_string()}),
|
||||||
|
);
|
||||||
|
return Err(format!("read prompt RPC event: {error}"));
|
||||||
|
}
|
||||||
}
|
}
|
||||||
let event = serde_json::from_str::<PromptEvent>(&line)
|
let event = match serde_json::from_str::<PromptEvent>(&line) {
|
||||||
.map_err(|e| format!("parse prompt RPC event: {e}"))?;
|
Ok(event) => event,
|
||||||
|
Err(error) => {
|
||||||
|
emit_chat_progress(
|
||||||
|
&mut progress,
|
||||||
|
CHAT_PROMPT_CHANNEL,
|
||||||
|
"prompt_event_parse",
|
||||||
|
"failed",
|
||||||
|
json!({"request_id": request_id, "error": error.to_string()}),
|
||||||
|
);
|
||||||
|
return Err(format!("parse prompt RPC event: {error}"));
|
||||||
|
}
|
||||||
|
};
|
||||||
let seen = event.request_id();
|
let seen = event.request_id();
|
||||||
if seen != request_id {
|
if seen != request_id {
|
||||||
|
emit_chat_progress(
|
||||||
|
&mut progress,
|
||||||
|
CHAT_PROMPT_CHANNEL,
|
||||||
|
"prompt_request_id",
|
||||||
|
"failed",
|
||||||
|
json!({"expected": request_id, "observed": seen}),
|
||||||
|
);
|
||||||
return Err(format!(
|
return Err(format!(
|
||||||
"prompt RPC protocol error: response request_id {seen} does not match active request_id {request_id}"
|
"prompt RPC protocol error: response request_id {seen} does not match active request_id {request_id}"
|
||||||
));
|
));
|
||||||
|
|
@ -795,6 +1246,13 @@ where
|
||||||
output
|
output
|
||||||
.flush()
|
.flush()
|
||||||
.map_err(|e| format!("flush response text: {e}"))?;
|
.map_err(|e| format!("flush response text: {e}"))?;
|
||||||
|
emit_chat_progress(
|
||||||
|
&mut progress,
|
||||||
|
CHAT_PROMPT_CHANNEL,
|
||||||
|
"response_text",
|
||||||
|
"observed",
|
||||||
|
json!({"request_id": request_id, "text_bytes": text.len()}),
|
||||||
|
);
|
||||||
}
|
}
|
||||||
PromptEvent::Done { .. } => {
|
PromptEvent::Done { .. } => {
|
||||||
if response_started {
|
if response_started {
|
||||||
|
|
@ -803,11 +1261,25 @@ where
|
||||||
writeln!(output, "Response: ")
|
writeln!(output, "Response: ")
|
||||||
.map_err(|e| format!("write empty response: {e}"))?;
|
.map_err(|e| format!("write empty response: {e}"))?;
|
||||||
}
|
}
|
||||||
|
emit_chat_progress(
|
||||||
|
&mut progress,
|
||||||
|
CHAT_PROMPT_CHANNEL,
|
||||||
|
"request_completed",
|
||||||
|
"ready",
|
||||||
|
json!({"request_id": request_id, "response_started": response_started}),
|
||||||
|
);
|
||||||
break;
|
break;
|
||||||
}
|
}
|
||||||
PromptEvent::Fault { error, .. } => {
|
PromptEvent::Fault { error, .. } => {
|
||||||
writeln!(output, "error: {error}")
|
writeln!(output, "error: {error}")
|
||||||
.map_err(|e| format!("write prompt fault: {e}"))?;
|
.map_err(|e| format!("write prompt fault: {e}"))?;
|
||||||
|
emit_chat_progress(
|
||||||
|
&mut progress,
|
||||||
|
CHAT_PROMPT_CHANNEL,
|
||||||
|
"request_faulted",
|
||||||
|
"ready",
|
||||||
|
json!({"request_id": request_id, "error": error}),
|
||||||
|
);
|
||||||
break;
|
break;
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
@ -1061,8 +1533,11 @@ mod tests {
|
||||||
static PROCESS_STATE_LOCK: Mutex<()> = Mutex::new(());
|
static PROCESS_STATE_LOCK: Mutex<()> = Mutex::new(());
|
||||||
static NEXT_TEMP_ID: AtomicU64 = AtomicU64::new(1);
|
static NEXT_TEMP_ID: AtomicU64 = AtomicU64::new(1);
|
||||||
|
|
||||||
const PROCESS_ENV_KEYS: &[&str] =
|
const PROCESS_ENV_KEYS: &[&str] = &[
|
||||||
&["VAST_API_KEY", "MVP_PIPELINE_STAGES", "MVP_RUNTIME_CONFIG"];
|
"VASTAI_API_KEY",
|
||||||
|
"MVP_PIPELINE_STAGES",
|
||||||
|
"MVP_RUNTIME_CONFIG",
|
||||||
|
];
|
||||||
|
|
||||||
struct TempDir {
|
struct TempDir {
|
||||||
path: PathBuf,
|
path: PathBuf,
|
||||||
|
|
@ -1521,7 +1996,7 @@ bootstrap_command = "boot"
|
||||||
"#,
|
"#,
|
||||||
);
|
);
|
||||||
with_process_state(
|
with_process_state(
|
||||||
&[("VAST_API_KEY", Some("secret"))],
|
&[("VASTAI_API_KEY", Some("secret"))],
|
||||||
Some(missing_relay.path()),
|
Some(missing_relay.path()),
|
||||||
|| {
|
|| {
|
||||||
let config_arg = missing_relay_config.to_string_lossy().into_owned();
|
let config_arg = missing_relay_config.to_string_lossy().into_owned();
|
||||||
|
|
@ -1546,7 +2021,7 @@ bootstrap_command = "boot"
|
||||||
"#,
|
"#,
|
||||||
);
|
);
|
||||||
with_process_state(
|
with_process_state(
|
||||||
&[("VAST_API_KEY", Some("secret"))],
|
&[("VASTAI_API_KEY", Some("secret"))],
|
||||||
Some(local_image.path()),
|
Some(local_image.path()),
|
||||||
|| {
|
|| {
|
||||||
let config_arg = local_image_config.to_string_lossy().into_owned();
|
let config_arg = local_image_config.to_string_lossy().into_owned();
|
||||||
|
|
@ -1571,7 +2046,7 @@ bootstrap_command = "boot"
|
||||||
"#,
|
"#,
|
||||||
);
|
);
|
||||||
with_process_state(
|
with_process_state(
|
||||||
&[("VAST_API_KEY", Some("secret"))],
|
&[("VASTAI_API_KEY", Some("secret"))],
|
||||||
Some(valid.path()),
|
Some(valid.path()),
|
||||||
|| {
|
|| {
|
||||||
let config_arg = valid_config.to_string_lossy().into_owned();
|
let config_arg = valid_config.to_string_lossy().into_owned();
|
||||||
|
|
|
||||||
|
|
@ -17,9 +17,9 @@ use datastream::{
|
||||||
};
|
};
|
||||||
use distribution::node::DistributedNodeConfig;
|
use distribution::node::DistributedNodeConfig;
|
||||||
use distribution::types::{MemberState, NodeId as DistNodeId};
|
use distribution::types::{MemberState, NodeId as DistNodeId};
|
||||||
use iroh::{Endpoint, EndpointAddr};
|
use iroh::EndpointAddr;
|
||||||
use iroh_driver::{
|
use iroh_driver::{
|
||||||
DATASTREAM_ALPN, IrohDriver, IrohDriverConfig, read_next_event, read_stream_header,
|
DATASTREAM_ALPN, EDGE_ALPN, EdgeSendHandle, EdgeTransportEvent, IrohDriver, IrohDriverConfig,
|
||||||
};
|
};
|
||||||
use mvp_system::actors::node_agent::{
|
use mvp_system::actors::node_agent::{
|
||||||
NodeAgentMsg, StageEdgeKindWire, StageInboundEdgeWire, StageObjectSpecWire,
|
NodeAgentMsg, StageEdgeKindWire, StageInboundEdgeWire, StageObjectSpecWire,
|
||||||
|
|
@ -28,7 +28,7 @@ use mvp_system::actors::node_agent::{
|
||||||
use mvp_system::actors::orchestrator::{OrchestratorActor, OrchestratorReport};
|
use mvp_system::actors::orchestrator::{OrchestratorActor, OrchestratorReport};
|
||||||
use mvp_system::actors::register_mvp_actor_codecs;
|
use mvp_system::actors::register_mvp_actor_codecs;
|
||||||
use mvp_system::config::{DEFAULT_CONFIG_PATH, TomlConfigOverlay};
|
use mvp_system::config::{DEFAULT_CONFIG_PATH, TomlConfigOverlay};
|
||||||
#[cfg(feature = "local-e2e")]
|
#[cfg(feature = "dashboard")]
|
||||||
use mvp_system::dashboard_view::MvpClusterDashboardView;
|
use mvp_system::dashboard_view::MvpClusterDashboardView;
|
||||||
use mvp_system::distribution_stack::DistributionRuntimeStack;
|
use mvp_system::distribution_stack::DistributionRuntimeStack;
|
||||||
use mvp_system::gpu_worker_ingress_parser as ingress;
|
use mvp_system::gpu_worker_ingress_parser as ingress;
|
||||||
|
|
@ -60,7 +60,7 @@ use mvp_system::vastai_provisioning::{
|
||||||
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::AsyncWriteExt;
|
#[cfg(test)]
|
||||||
use tokio::sync::mpsc as tokio_mpsc;
|
use tokio::sync::mpsc as tokio_mpsc;
|
||||||
|
|
||||||
const DEFAULT_IMAGE: &str = "swactor-mvp-node:latest";
|
const DEFAULT_IMAGE: &str = "swactor-mvp-node:latest";
|
||||||
|
|
@ -85,7 +85,6 @@ const MVP_STAGE_ROUTE: &str = "mvp.orch.stage_route";
|
||||||
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 DEFAULT_DOCKER_CONTAINER_PREFIX: &str = "mvp-orchestrator";
|
||||||
const MVP_DOCKER_CONTAINER_PREFIX_ENV: &str = "MVP_DOCKER_CONTAINER_PREFIX";
|
const MVP_DOCKER_CONTAINER_PREFIX_ENV: &str = "MVP_DOCKER_CONTAINER_PREFIX";
|
||||||
const EDGE_ALPN: &[u8] = b"mvp/pipeline-edge/0";
|
|
||||||
|
|
||||||
fn main() -> ExitCode {
|
fn main() -> ExitCode {
|
||||||
match run() {
|
match run() {
|
||||||
|
|
@ -97,41 +96,6 @@ fn main() -> ExitCode {
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
fn build_pipeline_edge_endpoint(
|
|
||||||
handle: &tokio::runtime::Handle,
|
|
||||||
relay: &RelayRuntimeConfig,
|
|
||||||
) -> Result<Endpoint, String> {
|
|
||||||
let relay_mode = relay.mode.clone();
|
|
||||||
let custom_relay = matches!(&relay_mode, iroh::RelayMode::Custom(_));
|
|
||||||
handle
|
|
||||||
.block_on(async move {
|
|
||||||
let mut builder = Endpoint::builder(iroh::endpoint::presets::Minimal)
|
|
||||||
.relay_mode(relay_mode)
|
|
||||||
.alpns(vec![EDGE_ALPN.to_vec(), DATASTREAM_ALPN.to_vec()]);
|
|
||||||
if custom_relay {
|
|
||||||
builder = builder.ca_roots_config(iroh::tls::CaRootsConfig::insecure_skip_verify());
|
|
||||||
}
|
|
||||||
builder.bind().await
|
|
||||||
})
|
|
||||||
.map_err(|e| format!("create pipeline edge endpoint: {e}"))
|
|
||||||
}
|
|
||||||
|
|
||||||
fn pipeline_edge_endpoint_addr(
|
|
||||||
endpoint: &Endpoint,
|
|
||||||
relay: &RelayRuntimeConfig,
|
|
||||||
) -> Result<EndpointAddr, String> {
|
|
||||||
let mut addr = endpoint.addr();
|
|
||||||
if addr.relay_urls().next().is_none() {
|
|
||||||
if let Some(url) = &relay.url {
|
|
||||||
addr = addr.with_relay_url(
|
|
||||||
url.parse()
|
|
||||||
.map_err(|e| format!("parse pipeline edge relay URL {url:?}: {e}"))?,
|
|
||||||
);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
Ok(addr)
|
|
||||||
}
|
|
||||||
|
|
||||||
fn run() -> Result<(), String> {
|
fn run() -> Result<(), String> {
|
||||||
let mut config = Config::from_defaults_toml_env_args(std::env::args().skip(1))?;
|
let mut config = Config::from_defaults_toml_env_args(std::env::args().skip(1))?;
|
||||||
config.prepare_vastai_ssh_key()?;
|
config.prepare_vastai_ssh_key()?;
|
||||||
|
|
@ -429,29 +393,8 @@ fn run() -> Result<(), String> {
|
||||||
let sink = PluginSink::new(Arc::new(ChannelObservationSink {
|
let sink = PluginSink::new(Arc::new(ChannelObservationSink {
|
||||||
tx: Mutex::new(obs_tx),
|
tx: Mutex::new(obs_tx),
|
||||||
}));
|
}));
|
||||||
let pipeline_edge_endpoint = if pipeline_plan.is_some() {
|
|
||||||
let endpoint = build_pipeline_edge_endpoint(tokio.handle(), &config.relay)?;
|
|
||||||
let addr = pipeline_edge_endpoint_addr(&endpoint, &config.relay)?;
|
|
||||||
orch_datastream.emit_bootstrap(
|
|
||||||
dashboard.as_ref(),
|
|
||||||
config.run_id,
|
|
||||||
config.node_id,
|
|
||||||
"pipeline_edge_endpoint",
|
|
||||||
"ready",
|
|
||||||
json!({"endpoint":addr}),
|
|
||||||
);
|
|
||||||
Some(endpoint)
|
|
||||||
} else {
|
|
||||||
None
|
|
||||||
};
|
|
||||||
let pipeline_token_ingress = pipeline_edge_endpoint
|
|
||||||
.as_ref()
|
|
||||||
.map(|endpoint| PipelineTokenIngress::start(tokio.handle().clone(), endpoint.clone()));
|
|
||||||
let coordinator_endpoint = driver.endpoint_addr();
|
let coordinator_endpoint = driver.endpoint_addr();
|
||||||
let pipeline_coordinator_endpoint = match &pipeline_edge_endpoint {
|
let pipeline_coordinator_endpoint = coordinator_endpoint.clone();
|
||||||
Some(endpoint) => pipeline_edge_endpoint_addr(endpoint, &config.relay)?,
|
|
||||||
None => coordinator_endpoint.clone(),
|
|
||||||
};
|
|
||||||
let (mut provisioned_nodes, ready) = start_and_provision_workers(
|
let (mut provisioned_nodes, ready) = start_and_provision_workers(
|
||||||
provisioner,
|
provisioner,
|
||||||
&config,
|
&config,
|
||||||
|
|
@ -539,8 +482,6 @@ fn run() -> Result<(), String> {
|
||||||
tokenizer_reply_actor,
|
tokenizer_reply_actor,
|
||||||
config.provider,
|
config.provider,
|
||||||
pipeline_plan.as_ref(),
|
pipeline_plan.as_ref(),
|
||||||
pipeline_edge_endpoint.as_ref(),
|
|
||||||
pipeline_token_ingress,
|
|
||||||
ready.first_stage.endpoint.clone(),
|
ready.first_stage.endpoint.clone(),
|
||||||
);
|
);
|
||||||
if let Err(error) = &result {
|
if let Err(error) = &result {
|
||||||
|
|
@ -3105,71 +3046,60 @@ struct CollectedDatastreamFrame {
|
||||||
}
|
}
|
||||||
|
|
||||||
fn drain_datastream_connections(
|
fn drain_datastream_connections(
|
||||||
driver: &IrohDriver,
|
driver: &mut IrohDriver,
|
||||||
frame_tx: &mpsc::Sender<CollectedDatastreamFrame>,
|
frame_tx: &mpsc::Sender<CollectedDatastreamFrame>,
|
||||||
) {
|
) {
|
||||||
for (_node, conn) in driver.drain_accepted_for_alpn(DATASTREAM_ALPN) {
|
driver.pump_datastream_ingress();
|
||||||
let tx = frame_tx.clone();
|
for read in driver.drain_datastream_reads() {
|
||||||
driver.runtime_handle().spawn(async move {
|
let mut channels = read
|
||||||
while let Ok(mut recv) = conn.accept_uni().await {
|
.header
|
||||||
let Ok(header) = read_stream_header(&mut recv).await else {
|
.channels
|
||||||
break;
|
.iter()
|
||||||
};
|
.map(|descriptor| {
|
||||||
let mut channels = header
|
(
|
||||||
.channels
|
ChannelRef {
|
||||||
.iter()
|
stream: descriptor.stream.clone(),
|
||||||
.map(|descriptor| {
|
channel: descriptor.id,
|
||||||
(
|
},
|
||||||
ChannelRef {
|
descriptor.name.clone(),
|
||||||
stream: descriptor.stream.clone(),
|
)
|
||||||
channel: descriptor.id,
|
})
|
||||||
},
|
.collect::<BTreeMap<_, _>>();
|
||||||
descriptor.name.clone(),
|
for event in read.events {
|
||||||
)
|
match event {
|
||||||
})
|
DatastreamEvent::ChannelDeclared(descriptor) => {
|
||||||
.collect::<BTreeMap<_, _>>();
|
channels.insert(
|
||||||
loop {
|
ChannelRef {
|
||||||
let event = match read_next_event(&mut recv, &header.stream).await {
|
stream: descriptor.stream.clone(),
|
||||||
Ok(Some(event)) => event,
|
channel: descriptor.id,
|
||||||
Ok(None) => break,
|
},
|
||||||
Err(_) => break,
|
descriptor.name,
|
||||||
};
|
);
|
||||||
match event {
|
}
|
||||||
DatastreamEvent::ChannelDeclared(descriptor) => {
|
DatastreamEvent::Frame(delivery) => {
|
||||||
channels.insert(
|
let channel_name = channels
|
||||||
ChannelRef {
|
.get(&delivery.channel)
|
||||||
stream: descriptor.stream.clone(),
|
.cloned()
|
||||||
channel: descriptor.id,
|
.unwrap_or_else(|| format!("channel#{}", delivery.channel.channel.0));
|
||||||
},
|
let frame = Frame::new(
|
||||||
descriptor.name,
|
delivery.channel.channel,
|
||||||
);
|
delivery.position,
|
||||||
}
|
delivery.payload,
|
||||||
DatastreamEvent::Frame(delivery) => {
|
);
|
||||||
let channel_name =
|
if frame_tx
|
||||||
channels.get(&delivery.channel).cloned().unwrap_or_else(|| {
|
.send(CollectedDatastreamFrame {
|
||||||
format!("channel#{}", delivery.channel.channel.0)
|
stream: delivery.channel.stream,
|
||||||
});
|
channel_name,
|
||||||
let frame = Frame::new(
|
frame,
|
||||||
delivery.channel.channel,
|
})
|
||||||
delivery.position,
|
.is_err()
|
||||||
delivery.payload,
|
{
|
||||||
);
|
return;
|
||||||
if tx
|
|
||||||
.send(CollectedDatastreamFrame {
|
|
||||||
stream: delivery.channel.stream,
|
|
||||||
channel_name,
|
|
||||||
frame,
|
|
||||||
})
|
|
||||||
.is_err()
|
|
||||||
{
|
|
||||||
break;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
DatastreamEvent::StreamDeclared(_) | DatastreamEvent::StreamEnded(_) => {}
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
DatastreamEvent::StreamDeclared(_) | DatastreamEvent::StreamEnded(_) => {}
|
||||||
}
|
}
|
||||||
});
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
@ -3485,12 +3415,13 @@ fn drain_orch_stdio_capture(
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
#[cfg(feature = "local-e2e")]
|
#[cfg(feature = "dashboard")]
|
||||||
struct DashboardSupport {
|
struct DashboardSupport {
|
||||||
handle: dashboard::DashboardHandle,
|
handle: dashboard::DashboardHandle,
|
||||||
|
_runtime: tokio::runtime::Runtime,
|
||||||
}
|
}
|
||||||
|
|
||||||
#[cfg(feature = "local-e2e")]
|
#[cfg(feature = "dashboard")]
|
||||||
impl DashboardSupport {
|
impl DashboardSupport {
|
||||||
fn start(enabled: bool) -> Result<Option<Self>, String> {
|
fn start(enabled: bool) -> Result<Option<Self>, String> {
|
||||||
if !enabled {
|
if !enabled {
|
||||||
|
|
@ -3502,10 +3433,17 @@ impl DashboardSupport {
|
||||||
.parse::<u16>()
|
.parse::<u16>()
|
||||||
.map_err(|e| format!("invalid MVP_DASHBOARD_PORT={port:?}: {e}"))?;
|
.map_err(|e| format!("invalid MVP_DASHBOARD_PORT={port:?}: {e}"))?;
|
||||||
}
|
}
|
||||||
let handle = dashboard::start_dashboard(config);
|
let runtime = tokio::runtime::Builder::new_multi_thread()
|
||||||
|
.enable_all()
|
||||||
|
.build()
|
||||||
|
.map_err(|e| format!("dashboard runtime: {e}"))?;
|
||||||
|
let handle = dashboard::DashboardHandle::new(config);
|
||||||
handle.register_view(Arc::new(MvpClusterDashboardView::new()));
|
handle.register_view(Arc::new(MvpClusterDashboardView::new()));
|
||||||
handle.start_http_standalone();
|
handle.spawn_http(runtime.handle());
|
||||||
Ok(Some(Self { handle }))
|
Ok(Some(Self {
|
||||||
|
handle,
|
||||||
|
_runtime: runtime,
|
||||||
|
}))
|
||||||
}
|
}
|
||||||
|
|
||||||
fn publish_frame(&self, stream: &StreamId, channel: &str, frame: &Frame) {
|
fn publish_frame(&self, stream: &StreamId, channel: &str, frame: &Frame) {
|
||||||
|
|
@ -3521,15 +3459,15 @@ impl DashboardSupport {
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
#[cfg(not(feature = "local-e2e"))]
|
#[cfg(not(feature = "dashboard"))]
|
||||||
struct DashboardSupport;
|
struct DashboardSupport;
|
||||||
|
|
||||||
#[cfg(not(feature = "local-e2e"))]
|
#[cfg(not(feature = "dashboard"))]
|
||||||
impl DashboardSupport {
|
impl DashboardSupport {
|
||||||
fn start(enabled: bool) -> Result<Option<Self>, String> {
|
fn start(enabled: bool) -> Result<Option<Self>, String> {
|
||||||
if enabled {
|
if enabled {
|
||||||
return Err(
|
return Err(
|
||||||
"MVP_DASHBOARD requires building mvp-system with feature local-e2e".to_owned(),
|
"MVP_DASHBOARD requires building mvp-system with feature dashboard".to_owned(),
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
Ok(None)
|
Ok(None)
|
||||||
|
|
@ -3830,31 +3768,23 @@ struct PipelineTokenRecord {
|
||||||
eos: bool,
|
eos: bool,
|
||||||
}
|
}
|
||||||
|
|
||||||
struct PipelineSendHandle {
|
enum PipelineSendHandle {
|
||||||
tx: tokio_mpsc::UnboundedSender<Vec<u8>>,
|
Driver(EdgeSendHandle),
|
||||||
|
#[cfg(test)]
|
||||||
|
Channel(tokio_mpsc::UnboundedSender<Vec<u8>>),
|
||||||
}
|
}
|
||||||
|
|
||||||
impl PipelineSendHandle {
|
impl PipelineSendHandle {
|
||||||
fn send(&self, bytes: Vec<u8>) -> Result<(), String> {
|
fn send(&self, bytes: Vec<u8>) -> Result<(), String> {
|
||||||
self.tx
|
match self {
|
||||||
.send(bytes)
|
Self::Driver(handle) => handle.send(bytes),
|
||||||
.map_err(|_| "pipeline token-in sender stopped".to_owned())
|
#[cfg(test)]
|
||||||
|
Self::Channel(tx) => tx
|
||||||
|
.send(bytes)
|
||||||
|
.map_err(|_| "pipeline token-in sender stopped".to_owned()),
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
struct PipelineTokenIngress {
|
|
||||||
recv_rx: mpsc::Receiver<Vec<u8>>,
|
|
||||||
recv_tx: mpsc::Sender<Vec<u8>>,
|
|
||||||
}
|
|
||||||
|
|
||||||
impl PipelineTokenIngress {
|
|
||||||
fn start(handle: tokio::runtime::Handle, endpoint: Endpoint) -> Self {
|
|
||||||
let (recv_tx, recv_rx) = mpsc::channel();
|
|
||||||
spawn_pipeline_token_acceptor(handle, endpoint, recv_tx.clone());
|
|
||||||
Self { recv_rx, recv_tx }
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
struct PendingEncode {
|
struct PendingEncode {
|
||||||
request_id: u64,
|
request_id: u64,
|
||||||
}
|
}
|
||||||
|
|
@ -3889,11 +3819,9 @@ struct PipelinePromptRuntime {
|
||||||
|
|
||||||
impl PipelinePromptRuntime {
|
impl PipelinePromptRuntime {
|
||||||
fn new(
|
fn new(
|
||||||
handle: tokio::runtime::Handle,
|
driver: &IrohDriver,
|
||||||
endpoint: Endpoint,
|
|
||||||
plan: &run_plan::RunPlan,
|
plan: &run_plan::RunPlan,
|
||||||
first_stage_endpoint: EndpointAddr,
|
first_stage_endpoint: EndpointAddr,
|
||||||
ingress: PipelineTokenIngress,
|
|
||||||
tokenizer_encode_actor: ActorAddress,
|
tokenizer_encode_actor: ActorAddress,
|
||||||
tokenizer_decode_actor: ActorAddress,
|
tokenizer_decode_actor: ActorAddress,
|
||||||
tokenizer_reply_to: ActorAddress,
|
tokenizer_reply_to: ActorAddress,
|
||||||
|
|
@ -3908,19 +3836,17 @@ impl PipelinePromptRuntime {
|
||||||
.iter()
|
.iter()
|
||||||
.find(|edge| edge.kind == run_plan::EdgeKind::TokenOut)
|
.find(|edge| edge.kind == run_plan::EdgeKind::TokenOut)
|
||||||
.ok_or_else(|| "pipeline plan missing token-out edge".to_owned())?;
|
.ok_or_else(|| "pipeline plan missing token-out edge".to_owned())?;
|
||||||
|
let (recv_tx, recv_rx) = mpsc::channel();
|
||||||
Ok(Self {
|
Ok(Self {
|
||||||
token_in_edge_id: token_in_edge.edge_id.0,
|
token_in_edge_id: token_in_edge.edge_id.0,
|
||||||
token_out_edge_id: token_out_edge.edge_id.0,
|
token_out_edge_id: token_out_edge.edge_id.0,
|
||||||
token_spec: token_in_edge.object_spec,
|
token_spec: token_in_edge.object_spec,
|
||||||
token_out_spec: token_out_edge.object_spec,
|
token_out_spec: token_out_edge.object_spec,
|
||||||
token_in_sender: spawn_pipeline_token_sender(
|
token_in_sender: PipelineSendHandle::Driver(
|
||||||
handle,
|
driver.spawn_edge_send_pump(first_stage_endpoint, token_in_edge.edge_id.0)?,
|
||||||
endpoint,
|
),
|
||||||
first_stage_endpoint,
|
recv_rx,
|
||||||
token_in_edge.edge_id.0,
|
recv_tx,
|
||||||
)?,
|
|
||||||
recv_rx: ingress.recv_rx,
|
|
||||||
recv_tx: ingress.recv_tx,
|
|
||||||
tokenizer_encode_actor,
|
tokenizer_encode_actor,
|
||||||
tokenizer_decode_actor,
|
tokenizer_decode_actor,
|
||||||
tokenizer_reply_to,
|
tokenizer_reply_to,
|
||||||
|
|
@ -3950,7 +3876,6 @@ impl PipelinePromptRuntime {
|
||||||
node_id: u64,
|
node_id: u64,
|
||||||
) -> Result<(), String> {
|
) -> Result<(), String> {
|
||||||
let request_id = request.request_id;
|
let request_id = request.request_id;
|
||||||
self.next_sequence = 0;
|
|
||||||
self.generated_tokens.clear();
|
self.generated_tokens.clear();
|
||||||
self.final_text.clear();
|
self.final_text.clear();
|
||||||
self.recv_buffer.clear();
|
self.recv_buffer.clear();
|
||||||
|
|
@ -4048,6 +3973,7 @@ impl PipelinePromptRuntime {
|
||||||
"ready",
|
"ready",
|
||||||
json!({"node_actor":self.tokenizer_encode_actor,"reply_to":self.tokenizer_reply_to,"tokens":tokens.len()}),
|
json!({"node_actor":self.tokenizer_encode_actor,"reply_to":self.tokenizer_reply_to,"tokens":tokens.len()}),
|
||||||
);
|
);
|
||||||
|
let sequence = self.next_sequence;
|
||||||
orch_datastream.emit_prompt(
|
orch_datastream.emit_prompt(
|
||||||
dashboard,
|
dashboard,
|
||||||
run_id,
|
run_id,
|
||||||
|
|
@ -4055,9 +3981,9 @@ impl PipelinePromptRuntime {
|
||||||
request_id,
|
request_id,
|
||||||
"pipeline_token_in",
|
"pipeline_token_in",
|
||||||
"started",
|
"started",
|
||||||
json!({"edge_id":self.token_in_edge_id,"sequence":0,"tokens":tokens.len()}),
|
json!({"edge_id":self.token_in_edge_id,"sequence":sequence,"tokens":tokens.len()}),
|
||||||
);
|
);
|
||||||
self.send_token_in(0, &tokens)?;
|
self.send_token_in(sequence, &tokens)?;
|
||||||
orch_datastream.emit_prompt(
|
orch_datastream.emit_prompt(
|
||||||
dashboard,
|
dashboard,
|
||||||
run_id,
|
run_id,
|
||||||
|
|
@ -4065,7 +3991,7 @@ impl PipelinePromptRuntime {
|
||||||
request_id,
|
request_id,
|
||||||
"pipeline_token_in",
|
"pipeline_token_in",
|
||||||
"ready",
|
"ready",
|
||||||
json!({"edge_id":self.token_in_edge_id,"sequence":0}),
|
json!({"edge_id":self.token_in_edge_id,"sequence":sequence}),
|
||||||
);
|
);
|
||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|
@ -4175,9 +4101,31 @@ impl PipelinePromptRuntime {
|
||||||
self.pending_decode = None;
|
self.pending_decode = None;
|
||||||
}
|
}
|
||||||
|
|
||||||
fn poll_driver(&mut self, driver: &IrohDriver) {
|
fn poll_driver(&mut self, driver: &mut IrohDriver) {
|
||||||
for (_node, conn) in driver.drain_other_connections() {
|
driver.pump_edge_ingress();
|
||||||
spawn_pipeline_token_receiver(driver.tokio_handle(), conn, self.recv_tx.clone());
|
for event in driver.drain_edge_events() {
|
||||||
|
match event {
|
||||||
|
EdgeTransportEvent::BytesRead { edge_id, bytes, .. }
|
||||||
|
if edge_id == self.token_out_edge_id =>
|
||||||
|
{
|
||||||
|
let _ = self.recv_tx.send(bytes);
|
||||||
|
}
|
||||||
|
EdgeTransportEvent::StreamFault {
|
||||||
|
edge_id: Some(edge_id),
|
||||||
|
reason,
|
||||||
|
..
|
||||||
|
} if edge_id == self.token_out_edge_id => {
|
||||||
|
if let Some(request_id) =
|
||||||
|
self.active.as_ref().map(|active| active.request.request_id)
|
||||||
|
{
|
||||||
|
self.fault_active(
|
||||||
|
request_id,
|
||||||
|
format!("pipeline token-out stream fault: {reason:?}"),
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
_ => {}
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
@ -4308,96 +4256,6 @@ fn take_pipeline_token_record(
|
||||||
Ok(Some(out))
|
Ok(Some(out))
|
||||||
}
|
}
|
||||||
|
|
||||||
fn spawn_pipeline_token_sender(
|
|
||||||
handle: tokio::runtime::Handle,
|
|
||||||
endpoint: iroh::Endpoint,
|
|
||||||
peer: EndpointAddr,
|
|
||||||
edge_id: u64,
|
|
||||||
) -> Result<PipelineSendHandle, String> {
|
|
||||||
let (tx, mut rx) = tokio_mpsc::unbounded_channel::<Vec<u8>>();
|
|
||||||
let (ready_tx, ready_rx) = mpsc::channel::<Result<(), String>>();
|
|
||||||
handle.spawn(async move {
|
|
||||||
let result: Result<(), String> = async {
|
|
||||||
let conn = endpoint
|
|
||||||
.connect(peer, EDGE_ALPN)
|
|
||||||
.await
|
|
||||||
.map_err(|e| format!("connect token-in edge {edge_id}: {e}"))?;
|
|
||||||
let mut send = conn
|
|
||||||
.open_uni()
|
|
||||||
.await
|
|
||||||
.map_err(|e| format!("open token-in stream {edge_id}: {e}"))?;
|
|
||||||
send.write_all(&edge_id.to_le_bytes())
|
|
||||||
.await
|
|
||||||
.map_err(|e| format!("write token-in preamble {edge_id}: {e}"))?;
|
|
||||||
send.flush()
|
|
||||||
.await
|
|
||||||
.map_err(|e| format!("flush token-in preamble {edge_id}: {e}"))?;
|
|
||||||
let _ = ready_tx.send(Ok(()));
|
|
||||||
while let Some(record) = rx.recv().await {
|
|
||||||
send.write_all(&record)
|
|
||||||
.await
|
|
||||||
.map_err(|e| format!("write token-in record {edge_id}: {e}"))?;
|
|
||||||
send.flush()
|
|
||||||
.await
|
|
||||||
.map_err(|e| format!("flush token-in record {edge_id}: {e}"))?;
|
|
||||||
}
|
|
||||||
send.finish()
|
|
||||||
.map_err(|e| format!("finish token-in stream {edge_id}: {e}"))?;
|
|
||||||
Ok(())
|
|
||||||
}
|
|
||||||
.await;
|
|
||||||
if let Err(error) = result {
|
|
||||||
let _ = ready_tx.send(Err(error));
|
|
||||||
}
|
|
||||||
});
|
|
||||||
ready_rx
|
|
||||||
.recv()
|
|
||||||
.map_err(|e| format!("token-in sender startup channel closed: {e}"))??;
|
|
||||||
Ok(PipelineSendHandle { tx })
|
|
||||||
}
|
|
||||||
|
|
||||||
fn spawn_pipeline_token_receiver(
|
|
||||||
handle: tokio::runtime::Handle,
|
|
||||||
conn: iroh::endpoint::Connection,
|
|
||||||
tx: mpsc::Sender<Vec<u8>>,
|
|
||||||
) {
|
|
||||||
handle.spawn(async move {
|
|
||||||
while let Ok(mut recv) = conn.accept_uni().await {
|
|
||||||
let mut preamble = [0u8; 8];
|
|
||||||
if recv.read_exact(&mut preamble).await.is_err() {
|
|
||||||
continue;
|
|
||||||
}
|
|
||||||
let mut chunk = vec![0u8; 4096];
|
|
||||||
loop {
|
|
||||||
match recv.read(&mut chunk).await {
|
|
||||||
Ok(Some(0)) | Ok(None) => break,
|
|
||||||
Ok(Some(n)) => {
|
|
||||||
if tx.send(chunk[..n].to_vec()).is_err() {
|
|
||||||
break;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
Err(_) => break,
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
});
|
|
||||||
}
|
|
||||||
|
|
||||||
fn spawn_pipeline_token_acceptor(
|
|
||||||
handle: tokio::runtime::Handle,
|
|
||||||
endpoint: Endpoint,
|
|
||||||
tx: mpsc::Sender<Vec<u8>>,
|
|
||||||
) {
|
|
||||||
let accept_handle = handle.clone();
|
|
||||||
handle.spawn(async move {
|
|
||||||
while let Some(incoming) = endpoint.accept().await {
|
|
||||||
if let Ok(conn) = incoming.await {
|
|
||||||
spawn_pipeline_token_receiver(accept_handle.clone(), conn, tx.clone());
|
|
||||||
}
|
|
||||||
}
|
|
||||||
});
|
|
||||||
}
|
|
||||||
|
|
||||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||||
enum PromptRuntimeMode {
|
enum PromptRuntimeMode {
|
||||||
DirectInferPrompt,
|
DirectInferPrompt,
|
||||||
|
|
@ -4434,28 +4292,17 @@ fn serve_prompts(
|
||||||
tokenizer_reply_to: ActorAddress,
|
tokenizer_reply_to: ActorAddress,
|
||||||
provider: ProviderKind,
|
provider: ProviderKind,
|
||||||
pipeline_plan: Option<&run_plan::RunPlan>,
|
pipeline_plan: Option<&run_plan::RunPlan>,
|
||||||
pipeline_edge_endpoint: Option<&Endpoint>,
|
|
||||||
pipeline_token_ingress: Option<PipelineTokenIngress>,
|
|
||||||
prompt_endpoint: EndpointAddr,
|
prompt_endpoint: EndpointAddr,
|
||||||
) -> Result<(), String> {
|
) -> Result<(), String> {
|
||||||
let mut pipeline_runtime = match prompt_runtime_mode(pipeline_plan) {
|
let mut pipeline_runtime = match prompt_runtime_mode(pipeline_plan) {
|
||||||
PromptRuntimeMode::PipelineTokenEdges => {
|
PromptRuntimeMode::PipelineTokenEdges => Some(PipelinePromptRuntime::new(
|
||||||
let endpoint = pipeline_edge_endpoint
|
driver,
|
||||||
.ok_or_else(|| "pipeline mode requires an edge endpoint".to_owned())?
|
pipeline_plan.expect("pipeline mode requires plan"),
|
||||||
.clone();
|
prompt_endpoint,
|
||||||
let ingress = pipeline_token_ingress
|
tokenizer_encode_actor,
|
||||||
.ok_or_else(|| "pipeline mode requires edge ingress".to_owned())?;
|
tokenizer_decode_actor,
|
||||||
Some(PipelinePromptRuntime::new(
|
tokenizer_reply_to,
|
||||||
driver.tokio_handle(),
|
)?),
|
||||||
endpoint,
|
|
||||||
pipeline_plan.expect("pipeline mode requires plan"),
|
|
||||||
prompt_endpoint,
|
|
||||||
ingress,
|
|
||||||
tokenizer_encode_actor,
|
|
||||||
tokenizer_decode_actor,
|
|
||||||
tokenizer_reply_to,
|
|
||||||
)?)
|
|
||||||
}
|
|
||||||
PromptRuntimeMode::DirectInferPrompt => None,
|
PromptRuntimeMode::DirectInferPrompt => None,
|
||||||
};
|
};
|
||||||
let mut active: Option<ActivePrompt> = None;
|
let mut active: Option<ActivePrompt> = None;
|
||||||
|
|
@ -5469,7 +5316,7 @@ mod tests {
|
||||||
token_out_edge_id,
|
token_out_edge_id,
|
||||||
token_spec,
|
token_spec,
|
||||||
token_out_spec,
|
token_out_spec,
|
||||||
token_in_sender: PipelineSendHandle { tx: token_in_tx },
|
token_in_sender: PipelineSendHandle::Channel(token_in_tx),
|
||||||
recv_rx,
|
recv_rx,
|
||||||
recv_tx,
|
recv_tx,
|
||||||
recv_buffer: Vec::new(),
|
recv_buffer: Vec::new(),
|
||||||
|
|
|
||||||
|
|
@ -20,7 +20,8 @@ use distribution::node::DistributedNodeConfig;
|
||||||
use distribution::types::{MemberState, NodeId as DistNodeId};
|
use distribution::types::{MemberState, NodeId as DistNodeId};
|
||||||
use iroh::EndpointAddr;
|
use iroh::EndpointAddr;
|
||||||
use iroh_driver::{
|
use iroh_driver::{
|
||||||
DATASTREAM_ALPN, DatastreamQuicHeader, IrohDriver, IrohDriverConfig, spawn_subscription_writer,
|
DATASTREAM_ALPN, DatastreamPublishHandle, DatastreamQuicHeader, EDGE_ALPN, EdgeSendHandle,
|
||||||
|
EdgeTransportEvent, IrohDriver, IrohDriverConfig,
|
||||||
};
|
};
|
||||||
use mvp_system::actors::node_agent::{
|
use mvp_system::actors::node_agent::{
|
||||||
NodeAgentActor, NodeAgentMsg, NodeAgentReport, StageCommandWire, StageInboundEdgeWire,
|
NodeAgentActor, NodeAgentMsg, NodeAgentReport, StageCommandWire, StageInboundEdgeWire,
|
||||||
|
|
@ -40,7 +41,6 @@ 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};
|
use tokio::io::{AsyncBufReadExt, AsyncWriteExt};
|
||||||
use tokio::sync::mpsc as tokio_mpsc;
|
|
||||||
|
|
||||||
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";
|
||||||
|
|
@ -58,7 +58,6 @@ const NODE_STAGE_CHANNEL: &str = "mvp.node.stage";
|
||||||
const NODE_WORKER_CHANNEL: &str = "mvp.node.worker";
|
const NODE_WORKER_CHANNEL: &str = "mvp.node.worker";
|
||||||
const NODE_PROMPT_CHANNEL: &str = "mvp.node.prompt";
|
const NODE_PROMPT_CHANNEL: &str = "mvp.node.prompt";
|
||||||
const NODE_SHUTDOWN_CHANNEL: &str = "mvp.node.shutdown";
|
const NODE_SHUTDOWN_CHANNEL: &str = "mvp.node.shutdown";
|
||||||
const EDGE_ALPN: &[u8] = b"mvp/pipeline-edge/0";
|
|
||||||
|
|
||||||
fn node_event_payload(
|
fn node_event_payload(
|
||||||
config: &DeploymentConfig,
|
config: &DeploymentConfig,
|
||||||
|
|
@ -509,65 +508,6 @@ fn spawn_arena_sampler(
|
||||||
});
|
});
|
||||||
}
|
}
|
||||||
|
|
||||||
#[derive(Clone, Debug)]
|
|
||||||
enum DriverIngressEvent {
|
|
||||||
StreamArrived {
|
|
||||||
edge_id: u64,
|
|
||||||
stream_id: u64,
|
|
||||||
},
|
|
||||||
BytesRead {
|
|
||||||
edge_id: u64,
|
|
||||||
stream_id: u64,
|
|
||||||
bytes: Vec<u8>,
|
|
||||||
},
|
|
||||||
}
|
|
||||||
|
|
||||||
#[derive(Clone)]
|
|
||||||
struct SendPumpHandle {
|
|
||||||
tx: tokio_mpsc::UnboundedSender<Vec<u8>>,
|
|
||||||
}
|
|
||||||
|
|
||||||
impl SendPumpHandle {
|
|
||||||
fn send(&self, record: Vec<u8>) -> Result<(), String> {
|
|
||||||
self.tx
|
|
||||||
.send(record)
|
|
||||||
.map_err(|_| "edge sender task stopped".to_owned())
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
struct DriverRuntime {
|
|
||||||
tx: mpsc::Sender<DriverIngressEvent>,
|
|
||||||
rx: mpsc::Receiver<DriverIngressEvent>,
|
|
||||||
next_stream_id: u64,
|
|
||||||
}
|
|
||||||
|
|
||||||
impl DriverRuntime {
|
|
||||||
fn new() -> Self {
|
|
||||||
let (tx, rx) = mpsc::channel();
|
|
||||||
Self {
|
|
||||||
tx,
|
|
||||||
rx,
|
|
||||||
next_stream_id: 1,
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
fn poll_iroh(&mut self, driver: &IrohDriver) {
|
|
||||||
for (_node, conn) in driver.drain_other_connections() {
|
|
||||||
spawn_recv_pump(
|
|
||||||
driver.tokio_handle(),
|
|
||||||
conn,
|
|
||||||
self.tx.clone(),
|
|
||||||
self.next_stream_id,
|
|
||||||
);
|
|
||||||
self.next_stream_id = self.next_stream_id.saturating_add(1);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
fn try_recv(&self) -> Option<DriverIngressEvent> {
|
|
||||||
self.rx.try_recv().ok()
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)]
|
#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)]
|
||||||
struct ObjectKey {
|
struct ObjectKey {
|
||||||
edge_id: u64,
|
edge_id: u64,
|
||||||
|
|
@ -585,7 +525,6 @@ struct LoadedObject {
|
||||||
struct WorkerEdgeRuntime {
|
struct WorkerEdgeRuntime {
|
||||||
establisher: edge::EdgeEstablisher,
|
establisher: edge::EdgeEstablisher,
|
||||||
driver_model: driver_model::Driver,
|
driver_model: driver_model::Driver,
|
||||||
driver_runtime: DriverRuntime,
|
|
||||||
edge_command_cursor: usize,
|
edge_command_cursor: usize,
|
||||||
edge_event_cursor: usize,
|
edge_event_cursor: usize,
|
||||||
driver_event_cursor: usize,
|
driver_event_cursor: usize,
|
||||||
|
|
@ -593,7 +532,7 @@ struct WorkerEdgeRuntime {
|
||||||
outbound_edge: Option<StageOutboundEdgeWire>,
|
outbound_edge: Option<StageOutboundEdgeWire>,
|
||||||
inbound_ring_id: Option<u64>,
|
inbound_ring_id: Option<u64>,
|
||||||
outbound_ring_id: Option<u64>,
|
outbound_ring_id: Option<u64>,
|
||||||
outbound_sender: Option<SendPumpHandle>,
|
outbound_sender: Option<EdgeSendHandle>,
|
||||||
next_output_object_id: u64,
|
next_output_object_id: u64,
|
||||||
object_handles: BTreeMap<ObjectKey, LoadedObject>,
|
object_handles: BTreeMap<ObjectKey, LoadedObject>,
|
||||||
ingress_streams: BTreeMap<u64, Vec<u8>>,
|
ingress_streams: BTreeMap<u64, Vec<u8>>,
|
||||||
|
|
@ -607,7 +546,6 @@ impl WorkerEdgeRuntime {
|
||||||
local_node_id: driver_model::NodeId(local_node_id),
|
local_node_id: driver_model::NodeId(local_node_id),
|
||||||
alpn: driver_model::Alpn(String::from_utf8_lossy(EDGE_ALPN).into_owned()),
|
alpn: driver_model::Alpn(String::from_utf8_lossy(EDGE_ALPN).into_owned()),
|
||||||
}),
|
}),
|
||||||
driver_runtime: DriverRuntime::new(),
|
|
||||||
edge_command_cursor: 0,
|
edge_command_cursor: 0,
|
||||||
edge_event_cursor: 0,
|
edge_event_cursor: 0,
|
||||||
driver_event_cursor: 0,
|
driver_event_cursor: 0,
|
||||||
|
|
@ -625,7 +563,7 @@ impl WorkerEdgeRuntime {
|
||||||
#[allow(clippy::too_many_arguments)]
|
#[allow(clippy::too_many_arguments)]
|
||||||
fn poll_iroh(
|
fn poll_iroh(
|
||||||
&mut self,
|
&mut self,
|
||||||
driver: &IrohDriver,
|
driver: &mut IrohDriver,
|
||||||
stack: &DistributionRuntimeStack,
|
stack: &DistributionRuntimeStack,
|
||||||
node_actor: ActorAddress,
|
node_actor: ActorAddress,
|
||||||
worker: &mut TinygradWorker,
|
worker: &mut TinygradWorker,
|
||||||
|
|
@ -633,10 +571,12 @@ impl WorkerEdgeRuntime {
|
||||||
config: &DeploymentConfig,
|
config: &DeploymentConfig,
|
||||||
datastream: &mut NodeDatastream,
|
datastream: &mut NodeDatastream,
|
||||||
) -> Result<(), String> {
|
) -> Result<(), String> {
|
||||||
self.driver_runtime.poll_iroh(driver);
|
driver.pump_edge_ingress();
|
||||||
while let Some(event) = self.driver_runtime.try_recv() {
|
for event in driver.drain_edge_events() {
|
||||||
match event {
|
match event {
|
||||||
DriverIngressEvent::StreamArrived { edge_id, stream_id } => {
|
EdgeTransportEvent::StreamArrived {
|
||||||
|
edge_id, stream_id, ..
|
||||||
|
} => {
|
||||||
self.driver_model
|
self.driver_model
|
||||||
.observe(driver_model::DriverEvent::IncomingUniStream {
|
.observe(driver_model::DriverEvent::IncomingUniStream {
|
||||||
edge_id: driver_model::EdgeId(edge_id),
|
edge_id: driver_model::EdgeId(edge_id),
|
||||||
|
|
@ -649,14 +589,14 @@ impl WorkerEdgeRuntime {
|
||||||
arena_manager,
|
arena_manager,
|
||||||
config,
|
config,
|
||||||
datastream,
|
datastream,
|
||||||
driver.tokio_handle(),
|
driver,
|
||||||
driver.endpoint().clone(),
|
|
||||||
)?;
|
)?;
|
||||||
}
|
}
|
||||||
DriverIngressEvent::BytesRead {
|
EdgeTransportEvent::BytesRead {
|
||||||
edge_id,
|
edge_id,
|
||||||
stream_id,
|
stream_id,
|
||||||
bytes,
|
bytes,
|
||||||
|
..
|
||||||
} => {
|
} => {
|
||||||
self.ingest_stream_bytes(
|
self.ingest_stream_bytes(
|
||||||
edge_id,
|
edge_id,
|
||||||
|
|
@ -668,10 +608,29 @@ impl WorkerEdgeRuntime {
|
||||||
arena_manager,
|
arena_manager,
|
||||||
config,
|
config,
|
||||||
datastream,
|
datastream,
|
||||||
driver.tokio_handle(),
|
driver,
|
||||||
driver.endpoint().clone(),
|
|
||||||
)?;
|
)?;
|
||||||
}
|
}
|
||||||
|
EdgeTransportEvent::StreamEnded { .. } => {}
|
||||||
|
EdgeTransportEvent::StreamFault {
|
||||||
|
edge_id: Some(edge_id),
|
||||||
|
..
|
||||||
|
} => {
|
||||||
|
self.driver_model
|
||||||
|
.observe(driver_model::DriverEvent::ReadError {
|
||||||
|
edge_id: driver_model::EdgeId(edge_id),
|
||||||
|
});
|
||||||
|
self.drive_edge_workflow(
|
||||||
|
stack,
|
||||||
|
node_actor,
|
||||||
|
worker,
|
||||||
|
arena_manager,
|
||||||
|
config,
|
||||||
|
datastream,
|
||||||
|
driver,
|
||||||
|
)?;
|
||||||
|
}
|
||||||
|
EdgeTransportEvent::StreamFault { edge_id: None, .. } => {}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
Ok(())
|
Ok(())
|
||||||
|
|
@ -687,8 +646,7 @@ impl WorkerEdgeRuntime {
|
||||||
arena_manager: &Arc<Mutex<arena::ArenaManager>>,
|
arena_manager: &Arc<Mutex<arena::ArenaManager>>,
|
||||||
config: &DeploymentConfig,
|
config: &DeploymentConfig,
|
||||||
datastream: &mut NodeDatastream,
|
datastream: &mut NodeDatastream,
|
||||||
handle: tokio::runtime::Handle,
|
driver: &mut IrohDriver,
|
||||||
endpoint: iroh::Endpoint,
|
|
||||||
) -> Result<(), String> {
|
) -> Result<(), String> {
|
||||||
self.inbound_edge = Some(edge.clone());
|
self.inbound_edge = Some(edge.clone());
|
||||||
self.establisher
|
self.establisher
|
||||||
|
|
@ -706,8 +664,7 @@ impl WorkerEdgeRuntime {
|
||||||
arena_manager,
|
arena_manager,
|
||||||
config,
|
config,
|
||||||
datastream,
|
datastream,
|
||||||
handle,
|
driver,
|
||||||
endpoint,
|
|
||||||
)
|
)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
@ -721,8 +678,7 @@ impl WorkerEdgeRuntime {
|
||||||
arena_manager: &Arc<Mutex<arena::ArenaManager>>,
|
arena_manager: &Arc<Mutex<arena::ArenaManager>>,
|
||||||
config: &DeploymentConfig,
|
config: &DeploymentConfig,
|
||||||
datastream: &mut NodeDatastream,
|
datastream: &mut NodeDatastream,
|
||||||
handle: tokio::runtime::Handle,
|
driver: &mut IrohDriver,
|
||||||
endpoint: iroh::Endpoint,
|
|
||||||
) -> Result<(), String> {
|
) -> Result<(), String> {
|
||||||
if edge.consumer_endpoint.is_none() {
|
if edge.consumer_endpoint.is_none() {
|
||||||
stack
|
stack
|
||||||
|
|
@ -754,8 +710,7 @@ impl WorkerEdgeRuntime {
|
||||||
arena_manager,
|
arena_manager,
|
||||||
config,
|
config,
|
||||||
datastream,
|
datastream,
|
||||||
handle,
|
driver,
|
||||||
endpoint,
|
|
||||||
)
|
)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
@ -872,8 +827,7 @@ impl WorkerEdgeRuntime {
|
||||||
arena_manager: &Arc<Mutex<arena::ArenaManager>>,
|
arena_manager: &Arc<Mutex<arena::ArenaManager>>,
|
||||||
config: &DeploymentConfig,
|
config: &DeploymentConfig,
|
||||||
datastream: &mut NodeDatastream,
|
datastream: &mut NodeDatastream,
|
||||||
handle: tokio::runtime::Handle,
|
driver: &mut IrohDriver,
|
||||||
endpoint: iroh::Endpoint,
|
|
||||||
) -> Result<(), String> {
|
) -> Result<(), String> {
|
||||||
let Some(inbound) = self.inbound_edge.clone() else {
|
let Some(inbound) = self.inbound_edge.clone() else {
|
||||||
return Ok(());
|
return Ok(());
|
||||||
|
|
@ -937,8 +891,7 @@ impl WorkerEdgeRuntime {
|
||||||
arena_manager,
|
arena_manager,
|
||||||
config,
|
config,
|
||||||
datastream,
|
datastream,
|
||||||
handle,
|
driver,
|
||||||
endpoint,
|
|
||||||
)
|
)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
@ -951,8 +904,7 @@ impl WorkerEdgeRuntime {
|
||||||
arena_manager: &Arc<Mutex<arena::ArenaManager>>,
|
arena_manager: &Arc<Mutex<arena::ArenaManager>>,
|
||||||
config: &DeploymentConfig,
|
config: &DeploymentConfig,
|
||||||
datastream: &mut NodeDatastream,
|
datastream: &mut NodeDatastream,
|
||||||
handle: tokio::runtime::Handle,
|
driver: &mut IrohDriver,
|
||||||
endpoint: iroh::Endpoint,
|
|
||||||
) -> Result<(), String> {
|
) -> Result<(), String> {
|
||||||
loop {
|
loop {
|
||||||
let mut progressed = false;
|
let mut progressed = false;
|
||||||
|
|
@ -1099,12 +1051,7 @@ impl WorkerEdgeRuntime {
|
||||||
},
|
},
|
||||||
},
|
},
|
||||||
));
|
));
|
||||||
self.outbound_sender = Some(spawn_send_pump(
|
self.outbound_sender = Some(driver.spawn_edge_send_pump(peer, edge_id.0)?);
|
||||||
handle.clone(),
|
|
||||||
endpoint.clone(),
|
|
||||||
peer,
|
|
||||||
edge_id.0,
|
|
||||||
)?);
|
|
||||||
}
|
}
|
||||||
edge::EdgeCommand::EstablishRecv { edge_id, .. } => {
|
edge::EdgeCommand::EstablishRecv { edge_id, .. } => {
|
||||||
let record = self
|
let record = self
|
||||||
|
|
@ -1291,104 +1238,6 @@ fn take_complete_ingress_record(
|
||||||
Ok(Some(buffer.drain(..record.total_len).collect()))
|
Ok(Some(buffer.drain(..record.total_len).collect()))
|
||||||
}
|
}
|
||||||
|
|
||||||
fn spawn_send_pump(
|
|
||||||
handle: tokio::runtime::Handle,
|
|
||||||
endpoint: iroh::Endpoint,
|
|
||||||
peer: EndpointAddr,
|
|
||||||
edge_id: u64,
|
|
||||||
) -> Result<SendPumpHandle, String> {
|
|
||||||
let (tx, mut rx) = tokio_mpsc::unbounded_channel::<Vec<u8>>();
|
|
||||||
let (ready_tx, ready_rx) = mpsc::channel::<Result<(), String>>();
|
|
||||||
handle.spawn(async move {
|
|
||||||
let result: Result<(), String> = async {
|
|
||||||
let conn = endpoint
|
|
||||||
.connect(peer, EDGE_ALPN)
|
|
||||||
.await
|
|
||||||
.map_err(|e| format!("connect edge {edge_id}: {e}"))?;
|
|
||||||
let mut send = conn
|
|
||||||
.open_uni()
|
|
||||||
.await
|
|
||||||
.map_err(|e| format!("open edge stream {edge_id}: {e}"))?;
|
|
||||||
send.write_all(&driver_model::encode_edge_preamble(driver_model::EdgeId(
|
|
||||||
edge_id,
|
|
||||||
)))
|
|
||||||
.await
|
|
||||||
.map_err(|e| format!("write edge preamble {edge_id}: {e}"))?;
|
|
||||||
send.flush()
|
|
||||||
.await
|
|
||||||
.map_err(|e| format!("flush edge preamble {edge_id}: {e}"))?;
|
|
||||||
let _ = ready_tx.send(Ok(()));
|
|
||||||
while let Some(record) = rx.recv().await {
|
|
||||||
send.write_all(&record)
|
|
||||||
.await
|
|
||||||
.map_err(|e| format!("write edge record {edge_id}: {e}"))?;
|
|
||||||
send.flush()
|
|
||||||
.await
|
|
||||||
.map_err(|e| format!("flush edge record {edge_id}: {e}"))?;
|
|
||||||
}
|
|
||||||
send.finish()
|
|
||||||
.map_err(|e| format!("finish edge stream {edge_id}: {e}"))?;
|
|
||||||
Ok(())
|
|
||||||
}
|
|
||||||
.await;
|
|
||||||
if let Err(error) = result {
|
|
||||||
let _ = ready_tx.send(Err(error));
|
|
||||||
}
|
|
||||||
});
|
|
||||||
ready_rx
|
|
||||||
.recv()
|
|
||||||
.map_err(|e| format!("edge {edge_id} sender startup channel closed: {e}"))??;
|
|
||||||
Ok(SendPumpHandle { tx })
|
|
||||||
}
|
|
||||||
|
|
||||||
fn spawn_recv_pump(
|
|
||||||
handle: tokio::runtime::Handle,
|
|
||||||
conn: iroh::endpoint::Connection,
|
|
||||||
tx: mpsc::Sender<DriverIngressEvent>,
|
|
||||||
stream_id: u64,
|
|
||||||
) {
|
|
||||||
handle.spawn(async move {
|
|
||||||
let mut next_uni_stream_id = stream_id << 32;
|
|
||||||
while let Ok(mut recv) = conn.accept_uni().await {
|
|
||||||
next_uni_stream_id = next_uni_stream_id.saturating_add(1);
|
|
||||||
let current_stream_id = next_uni_stream_id;
|
|
||||||
let mut preamble = [0u8; 8];
|
|
||||||
if recv.read_exact(&mut preamble).await.is_err() {
|
|
||||||
continue;
|
|
||||||
}
|
|
||||||
let edge_id = u64::from_le_bytes(preamble);
|
|
||||||
if tx
|
|
||||||
.send(DriverIngressEvent::StreamArrived {
|
|
||||||
edge_id,
|
|
||||||
stream_id: current_stream_id,
|
|
||||||
})
|
|
||||||
.is_err()
|
|
||||||
{
|
|
||||||
break;
|
|
||||||
}
|
|
||||||
let mut chunk = vec![0u8; 4096];
|
|
||||||
loop {
|
|
||||||
match recv.read(&mut chunk).await {
|
|
||||||
Ok(Some(0)) | Ok(None) => break,
|
|
||||||
Ok(Some(n)) => {
|
|
||||||
if tx
|
|
||||||
.send(DriverIngressEvent::BytesRead {
|
|
||||||
edge_id,
|
|
||||||
stream_id: current_stream_id,
|
|
||||||
bytes: chunk[..n].to_vec(),
|
|
||||||
})
|
|
||||||
.is_err()
|
|
||||||
{
|
|
||||||
break;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
Err(_) => break,
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
});
|
|
||||||
}
|
|
||||||
|
|
||||||
fn value_u64(value: &Value, field: &str) -> Result<u64, String> {
|
fn value_u64(value: &Value, field: &str) -> Result<u64, String> {
|
||||||
value
|
value
|
||||||
.get(field)
|
.get(field)
|
||||||
|
|
@ -1580,9 +1429,10 @@ fn run() -> Result<(), String> {
|
||||||
let arena_fd = arena_manager.lock().arena_fd();
|
let arena_fd = arena_manager.lock().arena_fd();
|
||||||
|
|
||||||
let mut datastream = node_datastream(&config);
|
let mut datastream = node_datastream(&config);
|
||||||
|
let datastream_transport = driver.datastream_publish_handle();
|
||||||
let datastream_publisher = match stack
|
let datastream_publisher = match stack
|
||||||
.runtime
|
.runtime
|
||||||
.spawn(datastream.publisher_actor(tokio.handle().clone(), driver.endpoint()))
|
.spawn(datastream.publisher_actor(datastream_transport))
|
||||||
{
|
{
|
||||||
Ok(actor) => actor,
|
Ok(actor) => actor,
|
||||||
Err(error) => {
|
Err(error) => {
|
||||||
|
|
@ -1859,7 +1709,7 @@ fn run() -> Result<(), String> {
|
||||||
datastream.tick();
|
datastream.tick();
|
||||||
worker.drain_stderr(&config, &mut datastream);
|
worker.drain_stderr(&config, &mut datastream);
|
||||||
edge_runtime.poll_iroh(
|
edge_runtime.poll_iroh(
|
||||||
&driver,
|
&mut driver,
|
||||||
&stack,
|
&stack,
|
||||||
node_actor,
|
node_actor,
|
||||||
&mut worker,
|
&mut worker,
|
||||||
|
|
@ -2154,11 +2004,7 @@ impl NodeDatastream {
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
fn publisher_actor(
|
fn publisher_actor(&self, transport: DatastreamPublishHandle) -> DatastreamPublisherActor {
|
||||||
&self,
|
|
||||||
tokio: tokio::runtime::Handle,
|
|
||||||
iroh_endpoint: iroh::Endpoint,
|
|
||||||
) -> DatastreamPublisherActor {
|
|
||||||
DatastreamPublisherActor::new(
|
DatastreamPublisherActor::new(
|
||||||
Arc::clone(&self.endpoint),
|
Arc::clone(&self.endpoint),
|
||||||
move |subscribe: DatastreamSubscribe, subscription: DatastreamSubscription| {
|
move |subscribe: DatastreamSubscribe, subscription: DatastreamSubscription| {
|
||||||
|
|
@ -2169,9 +2015,7 @@ impl NodeDatastream {
|
||||||
) else {
|
) else {
|
||||||
return;
|
return;
|
||||||
};
|
};
|
||||||
let _ = spawn_subscription_writer(
|
transport.publish_subscription(
|
||||||
&tokio,
|
|
||||||
iroh_endpoint.clone(),
|
|
||||||
subscribe.collector,
|
subscribe.collector,
|
||||||
header,
|
header,
|
||||||
subscription,
|
subscription,
|
||||||
|
|
@ -2920,8 +2764,7 @@ fn handle_stage_command(
|
||||||
arena_manager,
|
arena_manager,
|
||||||
config,
|
config,
|
||||||
datastream,
|
datastream,
|
||||||
driver.tokio_handle(),
|
driver,
|
||||||
driver.endpoint().clone(),
|
|
||||||
)?;
|
)?;
|
||||||
emit_node_event(
|
emit_node_event(
|
||||||
datastream,
|
datastream,
|
||||||
|
|
@ -2950,8 +2793,7 @@ fn handle_stage_command(
|
||||||
arena_manager,
|
arena_manager,
|
||||||
config,
|
config,
|
||||||
datastream,
|
datastream,
|
||||||
driver.tokio_handle(),
|
driver,
|
||||||
driver.endpoint().clone(),
|
|
||||||
)?;
|
)?;
|
||||||
emit_node_event(
|
emit_node_event(
|
||||||
datastream,
|
datastream,
|
||||||
|
|
|
||||||
|
|
@ -174,7 +174,7 @@ impl TomlConfigOverlay {
|
||||||
|
|
||||||
impl ResolvedVastAiConfig {
|
impl ResolvedVastAiConfig {
|
||||||
pub fn validate(self) -> Result<Self, String> {
|
pub fn validate(self) -> Result<Self, String> {
|
||||||
require_non_empty("VAST_API_KEY", &self.api_key)?;
|
require_non_empty("VASTAI_API_KEY", &self.api_key)?;
|
||||||
require_non_empty("relay.url", &self.relay_url)?;
|
require_non_empty("relay.url", &self.relay_url)?;
|
||||||
require_non_empty("vastai.image", &self.image)?;
|
require_non_empty("vastai.image", &self.image)?;
|
||||||
require_non_empty("vastai.bootstrap_command", &self.bootstrap_command)?;
|
require_non_empty("vastai.bootstrap_command", &self.bootstrap_command)?;
|
||||||
|
|
|
||||||
|
|
@ -5,6 +5,7 @@ pub mod actors;
|
||||||
pub mod arena_manager;
|
pub mod arena_manager;
|
||||||
pub mod bootstrap_datastream;
|
pub mod bootstrap_datastream;
|
||||||
pub mod config;
|
pub mod config;
|
||||||
|
pub mod dashboard_view;
|
||||||
pub mod device_bridge;
|
pub mod device_bridge;
|
||||||
pub mod distribution_stack;
|
pub mod distribution_stack;
|
||||||
pub mod docker_cluster_provisioning;
|
pub mod docker_cluster_provisioning;
|
||||||
|
|
|
||||||
|
|
@ -1,841 +0,0 @@
|
||||||
use std::ffi::CString;
|
|
||||||
use std::io::{BufRead, BufReader, Write};
|
|
||||||
use std::os::fd::RawFd;
|
|
||||||
use std::os::unix::process::CommandExt;
|
|
||||||
use std::path::{Path, PathBuf};
|
|
||||||
use std::process::{Command, Stdio};
|
|
||||||
use std::sync::Arc;
|
|
||||||
use std::sync::atomic::{AtomicBool, Ordering};
|
|
||||||
use std::sync::mpsc;
|
|
||||||
use std::thread;
|
|
||||||
use std::time::{Duration, Instant};
|
|
||||||
|
|
||||||
use datastream::emit::{DatastreamEmitter, EmitterConfig, FrameSink};
|
|
||||||
use datastream::{Frame, StreamId};
|
|
||||||
use serde_json::{Value, json};
|
|
||||||
use swactor::actor::{ActorAddress, ActorInterface};
|
|
||||||
use swactor::runtime::{Ctx, ExternalSender, Runtime, RuntimeConfig};
|
|
||||||
|
|
||||||
const IMAGE: &str = "swactor-mvp-gpu-worker-node-e2e:latest";
|
|
||||||
const WORKER_EVENTS_CHANNEL: &str = "mvp.worker.events";
|
|
||||||
const ARENA_BYTES: usize = 8192;
|
|
||||||
const INGRESS_RING_ID: u64 = 8001;
|
|
||||||
const EGRESS_RING_ID: u64 = 8002;
|
|
||||||
const INGRESS_EDGE_ID: u64 = 7001;
|
|
||||||
const EGRESS_EDGE_ID: u64 = 7002;
|
|
||||||
const INGRESS_BASE: usize = 0;
|
|
||||||
const EGRESS_BASE: usize = 4096;
|
|
||||||
const RING_BYTES: usize = 1024;
|
|
||||||
const HEADER_LEN: usize = 48;
|
|
||||||
const PREFLIGHT_WATCHDOG: Duration = Duration::from_secs(45);
|
|
||||||
const EVENT_WATCHDOG: Duration = Duration::from_secs(60);
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn gpu_worker_node_e2e_cuda() {
|
|
||||||
if std::env::var_os("MVP_SYSTEM_CUDA_E2E_IN_CONTAINER").is_some() {
|
|
||||||
run_integrated_node_harness();
|
|
||||||
} else if std::env::var_os("MVP_SYSTEM_CUDA_E2E").is_some() {
|
|
||||||
build_and_run_docker_fixture();
|
|
||||||
} else {
|
|
||||||
eprintln!("skipping; set MVP_SYSTEM_CUDA_E2E=1 to run docker CUDA e2e");
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
fn build_and_run_docker_fixture() {
|
|
||||||
let crate_dir = Path::new(env!("CARGO_MANIFEST_DIR"));
|
|
||||||
let workspace = crate_dir
|
|
||||||
.parent()
|
|
||||||
.and_then(Path::parent)
|
|
||||||
.expect("workspace root")
|
|
||||||
.canonicalize()
|
|
||||||
.expect("canonical workspace root");
|
|
||||||
let context =
|
|
||||||
std::env::temp_dir().join(format!("mvp-system-docker-context-{}", std::process::id()));
|
|
||||||
let _ = std::fs::remove_dir_all(&context);
|
|
||||||
copy_workspace_context(&workspace, &context);
|
|
||||||
let dockerfile = context.join("crates/mvp-system/tests/gpu_worker_node_e2e/Dockerfile");
|
|
||||||
|
|
||||||
phase("building CUDA Docker fixture image");
|
|
||||||
let build = Command::new("docker")
|
|
||||||
.args(["build", "-f"])
|
|
||||||
.arg(&dockerfile)
|
|
||||||
.args(["-t", IMAGE])
|
|
||||||
.arg(&context)
|
|
||||||
.status()
|
|
||||||
.expect("run docker build");
|
|
||||||
assert!(build.success(), "docker build failed with status {build}");
|
|
||||||
|
|
||||||
phase("running CUDA Docker fixture");
|
|
||||||
let run = Command::new("docker")
|
|
||||||
.args(["run", "--rm", "--gpus"])
|
|
||||||
.arg(std::env::var("MVP_CUDA_GPUS").unwrap_or_else(|_| "all".to_owned()))
|
|
||||||
.args(["-e", "MVP_SYSTEM_CUDA_E2E_IN_CONTAINER=1"])
|
|
||||||
.args(["-e", "CARGO_TARGET_DIR=/tmp/mvp-system-target"])
|
|
||||||
.arg(IMAGE)
|
|
||||||
.status()
|
|
||||||
.expect("run docker CUDA fixture");
|
|
||||||
assert!(
|
|
||||||
run.success(),
|
|
||||||
"docker CUDA fixture failed with status {run}"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
fn copy_workspace_context(source: &Path, dest: &Path) {
|
|
||||||
std::fs::create_dir_all(dest).expect("create docker context");
|
|
||||||
for entry in std::fs::read_dir(source).expect("read workspace") {
|
|
||||||
let entry = entry.expect("read workspace entry");
|
|
||||||
let name = entry.file_name();
|
|
||||||
let name = name.to_string_lossy();
|
|
||||||
if matches!(name.as_ref(), ".git" | "target" | ".dockerignore") {
|
|
||||||
continue;
|
|
||||||
}
|
|
||||||
copy_context_entry(&entry.path(), &dest.join(name.as_ref()));
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
fn copy_context_entry(source: &Path, dest: &Path) {
|
|
||||||
let metadata = std::fs::symlink_metadata(source).expect("context metadata");
|
|
||||||
if metadata.file_type().is_symlink() {
|
|
||||||
return;
|
|
||||||
}
|
|
||||||
if metadata.is_dir() {
|
|
||||||
std::fs::create_dir_all(dest).expect("create context dir");
|
|
||||||
for entry in std::fs::read_dir(source).expect("read context dir") {
|
|
||||||
let entry = entry.expect("read context entry");
|
|
||||||
let name = entry.file_name();
|
|
||||||
let name = name.to_string_lossy();
|
|
||||||
if matches!(name.as_ref(), ".git" | "target" | ".dockerignore") {
|
|
||||||
continue;
|
|
||||||
}
|
|
||||||
copy_context_entry(&entry.path(), &dest.join(name.as_ref()));
|
|
||||||
}
|
|
||||||
} else if metadata.is_file() {
|
|
||||||
if let Some(parent) = dest.parent() {
|
|
||||||
std::fs::create_dir_all(parent).expect("create context parent");
|
|
||||||
}
|
|
||||||
std::fs::copy(source, dest).expect("copy context file");
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
fn phase(message: &str) {
|
|
||||||
eprintln!("gpu-worker-node-e2e: {message}");
|
|
||||||
}
|
|
||||||
|
|
||||||
fn run_cuda_preflight() {
|
|
||||||
let script = r#"
|
|
||||||
import os
|
|
||||||
print("cuda preflight: importing tinygrad", flush=True)
|
|
||||||
from tinygrad import Tensor, dtypes
|
|
||||||
print(f"cuda preflight: DEV={os.environ.get('DEV')}", flush=True)
|
|
||||||
print("cuda preflight: realizing Tensor([1])", flush=True)
|
|
||||||
value = Tensor([1], dtype=dtypes.int32).realize().numpy().tolist()
|
|
||||||
print(f"cuda preflight: ok {value}", flush=True)
|
|
||||||
"#;
|
|
||||||
let child = Command::new("python3")
|
|
||||||
.arg("-c")
|
|
||||||
.arg(script)
|
|
||||||
.env("DEV", "CUDA")
|
|
||||||
.stdout(Stdio::piped())
|
|
||||||
.stderr(Stdio::piped())
|
|
||||||
.spawn()
|
|
||||||
.expect("spawn CUDA preflight");
|
|
||||||
let pid = child.id();
|
|
||||||
let (tx, rx) = mpsc::channel();
|
|
||||||
thread::spawn(move || {
|
|
||||||
let _ = tx.send(child.wait_with_output());
|
|
||||||
});
|
|
||||||
|
|
||||||
match rx.recv_timeout(PREFLIGHT_WATCHDOG) {
|
|
||||||
Ok(output) => {
|
|
||||||
let output = output.expect("wait CUDA preflight");
|
|
||||||
eprintln!(
|
|
||||||
"gpu-worker-node-e2e: CUDA preflight stdout:\n{}",
|
|
||||||
String::from_utf8_lossy(&output.stdout)
|
|
||||||
);
|
|
||||||
assert!(
|
|
||||||
output.status.success(),
|
|
||||||
"CUDA preflight failed with status {}\nstdout:\n{}\nstderr:\n{}",
|
|
||||||
output.status,
|
|
||||||
String::from_utf8_lossy(&output.stdout),
|
|
||||||
String::from_utf8_lossy(&output.stderr)
|
|
||||||
);
|
|
||||||
}
|
|
||||||
Err(mpsc::RecvTimeoutError::Timeout) => {
|
|
||||||
unsafe {
|
|
||||||
libc::kill(pid as libc::pid_t, libc::SIGKILL);
|
|
||||||
}
|
|
||||||
let output = rx
|
|
||||||
.recv_timeout(Duration::from_secs(5))
|
|
||||||
.ok()
|
|
||||||
.and_then(Result::ok);
|
|
||||||
let (stdout, stderr) = output
|
|
||||||
.as_ref()
|
|
||||||
.map(|output| {
|
|
||||||
(
|
|
||||||
String::from_utf8_lossy(&output.stdout).into_owned(),
|
|
||||||
String::from_utf8_lossy(&output.stderr).into_owned(),
|
|
||||||
)
|
|
||||||
})
|
|
||||||
.unwrap_or_else(|| ("<unavailable>".to_owned(), "<unavailable>".to_owned()));
|
|
||||||
panic!(
|
|
||||||
"CUDA preflight test watchdog after {:?}; killed pid {pid}\nstdout:\n{stdout}\nstderr:\n{stderr}",
|
|
||||||
PREFLIGHT_WATCHDOG
|
|
||||||
);
|
|
||||||
}
|
|
||||||
Err(mpsc::RecvTimeoutError::Disconnected) => panic!("CUDA preflight waiter disconnected"),
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
fn run_integrated_node_harness() {
|
|
||||||
phase("running CUDA tinygrad preflight");
|
|
||||||
run_cuda_preflight();
|
|
||||||
|
|
||||||
phase("creating shared arena and telemetry socket");
|
|
||||||
let arena_fd = create_arena(ARENA_BYTES);
|
|
||||||
let worker_path = Path::new(env!("CARGO_MANIFEST_DIR"))
|
|
||||||
.join("tests/gpu_worker_node_e2e/mvp_tinygrad_worker.py");
|
|
||||||
let socket_path =
|
|
||||||
std::env::temp_dir().join(format!("mvp-worker-events-{}.sock", std::process::id()));
|
|
||||||
let _ = std::fs::remove_file(&socket_path);
|
|
||||||
|
|
||||||
let (frame_tx, frame_rx) = mpsc::channel();
|
|
||||||
let mut emitter = DatastreamEmitter::new(
|
|
||||||
EmitterConfig {
|
|
||||||
node_hex: "node-11".to_owned(),
|
|
||||||
life: 1,
|
|
||||||
mux_capacity: 256,
|
|
||||||
},
|
|
||||||
Box::new(ChannelFrameSink { tx: frame_tx }),
|
|
||||||
);
|
|
||||||
let ingest_alive = Arc::new(AtomicBool::new(true));
|
|
||||||
let ingest_thread = spawn_worker_event_ingest(
|
|
||||||
socket_path.clone(),
|
|
||||||
emitter.event_sink(),
|
|
||||||
Arc::clone(&ingest_alive),
|
|
||||||
);
|
|
||||||
|
|
||||||
phase("spawning Rust worker node actor");
|
|
||||||
let rt = Runtime::new(RuntimeConfig::default());
|
|
||||||
let reports = rt.new_inbox::<HarnessReport>().expect("report inbox");
|
|
||||||
let node = rt
|
|
||||||
.spawn(GpuWorkerNodeActor::new(rt.create_sender(), *reports.addr()))
|
|
||||||
.expect("spawn gpu worker node actor");
|
|
||||||
|
|
||||||
phase("spawning tinygrad worker process");
|
|
||||||
rt.send_to(
|
|
||||||
node,
|
|
||||||
NodeMsg::Start(StartWorker {
|
|
||||||
worker_path,
|
|
||||||
socket_path: socket_path.clone(),
|
|
||||||
arena_fd,
|
|
||||||
arena_bytes: ARENA_BYTES as u64,
|
|
||||||
}),
|
|
||||||
)
|
|
||||||
.expect("send start");
|
|
||||||
wait_for_report(
|
|
||||||
&rt,
|
|
||||||
node,
|
|
||||||
&mut emitter,
|
|
||||||
&frame_rx,
|
|
||||||
&reports,
|
|
||||||
"worker process start",
|
|
||||||
|report| matches!(report, HarnessReport::ProcessStarted),
|
|
||||||
);
|
|
||||||
|
|
||||||
phase("initializing CUDA backend");
|
|
||||||
send_command(
|
|
||||||
&rt,
|
|
||||||
node,
|
|
||||||
json!({"type":"InitializeWorker","helper_abi_version":1}),
|
|
||||||
);
|
|
||||||
wait_for_control(
|
|
||||||
&rt,
|
|
||||||
node,
|
|
||||||
&mut emitter,
|
|
||||||
&frame_rx,
|
|
||||||
&reports,
|
|
||||||
"WorkerReady",
|
|
||||||
"worker ready",
|
|
||||||
);
|
|
||||||
|
|
||||||
phase("installing ingress and egress rings");
|
|
||||||
send_command(
|
|
||||||
&rt,
|
|
||||||
node,
|
|
||||||
install_ring_command(
|
|
||||||
INGRESS_RING_ID,
|
|
||||||
INGRESS_EDGE_ID,
|
|
||||||
"in",
|
|
||||||
"ingress",
|
|
||||||
INGRESS_BASE,
|
|
||||||
),
|
|
||||||
);
|
|
||||||
send_command(
|
|
||||||
&rt,
|
|
||||||
node,
|
|
||||||
install_ring_command(EGRESS_RING_ID, EGRESS_EDGE_ID, "out", "egress", EGRESS_BASE),
|
|
||||||
);
|
|
||||||
wait_for_control(
|
|
||||||
&rt,
|
|
||||||
node,
|
|
||||||
&mut emitter,
|
|
||||||
&frame_rx,
|
|
||||||
&reports,
|
|
||||||
"RingInstalled",
|
|
||||||
"ingress ring installed",
|
|
||||||
);
|
|
||||||
wait_for_control(
|
|
||||||
&rt,
|
|
||||||
node,
|
|
||||||
&mut emitter,
|
|
||||||
&frame_rx,
|
|
||||||
&reports,
|
|
||||||
"RingInstalled",
|
|
||||||
"egress ring installed",
|
|
||||||
);
|
|
||||||
|
|
||||||
phase("configuring worker role");
|
|
||||||
send_command(&rt, node, json!({"type":"ConfigureRole","role_id":1}));
|
|
||||||
wait_for_control(
|
|
||||||
&rt,
|
|
||||||
node,
|
|
||||||
&mut emitter,
|
|
||||||
&frame_rx,
|
|
||||||
&reports,
|
|
||||||
"RoleLoaded",
|
|
||||||
"role loaded",
|
|
||||||
);
|
|
||||||
|
|
||||||
phase("copying ingress object into CUDA tensor");
|
|
||||||
let input_record = object_record(9000, 0, &[1, 2, 3, 4]);
|
|
||||||
pwrite_all(arena_fd, INGRESS_BASE, &input_record);
|
|
||||||
send_command(
|
|
||||||
&rt,
|
|
||||||
node,
|
|
||||||
json!({"type":"RingReadable","ring_id":INGRESS_RING_ID,"committed_bytes":input_record.len()}),
|
|
||||||
);
|
|
||||||
let loaded = wait_for_control(
|
|
||||||
&rt,
|
|
||||||
node,
|
|
||||||
&mut emitter,
|
|
||||||
&frame_rx,
|
|
||||||
&reports,
|
|
||||||
"ObjectLoaded",
|
|
||||||
"object loaded",
|
|
||||||
);
|
|
||||||
let handle = loaded["handle"]["id"].as_u64().expect("device handle id");
|
|
||||||
|
|
||||||
phase("executing CUDA step and writing egress object");
|
|
||||||
send_command(
|
|
||||||
&rt,
|
|
||||||
node,
|
|
||||||
json!({
|
|
||||||
"type":"ExecuteStep",
|
|
||||||
"step_id":9001,
|
|
||||||
"input_handle":handle,
|
|
||||||
"egress_ring_id":EGRESS_RING_ID,
|
|
||||||
"output_object_id":9001
|
|
||||||
}),
|
|
||||||
);
|
|
||||||
let produced = wait_for_control(
|
|
||||||
&rt,
|
|
||||||
node,
|
|
||||||
&mut emitter,
|
|
||||||
&frame_rx,
|
|
||||||
&reports,
|
|
||||||
"ObjectProduced",
|
|
||||||
"object produced",
|
|
||||||
);
|
|
||||||
wait_for_control(
|
|
||||||
&rt,
|
|
||||||
node,
|
|
||||||
&mut emitter,
|
|
||||||
&frame_rx,
|
|
||||||
&reports,
|
|
||||||
"StepCompleted",
|
|
||||||
"step completed",
|
|
||||||
);
|
|
||||||
let committed = produced["committed_bytes"]
|
|
||||||
.as_u64()
|
|
||||||
.expect("committed bytes") as usize;
|
|
||||||
let mut egress = vec![0u8; committed];
|
|
||||||
pread_exact(arena_fd, EGRESS_BASE, &mut egress);
|
|
||||||
assert_eq!(decode_payload_words(&egress), vec![2, 4, 6, 8]);
|
|
||||||
|
|
||||||
phase("releasing device object");
|
|
||||||
send_command(
|
|
||||||
&rt,
|
|
||||||
node,
|
|
||||||
json!({"type":"ReleaseDeviceObject","handle":handle}),
|
|
||||||
);
|
|
||||||
wait_for_control(
|
|
||||||
&rt,
|
|
||||||
node,
|
|
||||||
&mut emitter,
|
|
||||||
&frame_rx,
|
|
||||||
&reports,
|
|
||||||
"DeviceObjectReleased",
|
|
||||||
"device object released",
|
|
||||||
);
|
|
||||||
|
|
||||||
phase("shutting down worker process");
|
|
||||||
send_command(&rt, node, json!({"type":"ShutdownWorker"}));
|
|
||||||
wait_for_control(
|
|
||||||
&rt,
|
|
||||||
node,
|
|
||||||
&mut emitter,
|
|
||||||
&frame_rx,
|
|
||||||
&reports,
|
|
||||||
"WorkerStopped",
|
|
||||||
"worker stopped",
|
|
||||||
);
|
|
||||||
wait_for_report(
|
|
||||||
&rt,
|
|
||||||
node,
|
|
||||||
&mut emitter,
|
|
||||||
&frame_rx,
|
|
||||||
&reports,
|
|
||||||
"worker process exit",
|
|
||||||
|report| matches!(report, HarnessReport::ProcessExited(0)),
|
|
||||||
);
|
|
||||||
|
|
||||||
phase("asserting worker telemetry");
|
|
||||||
let telemetry = collect_telemetry(&mut emitter, &frame_rx, Duration::from_secs(1));
|
|
||||||
assert_has_worker_event(&telemetry, "importing_tinygrad");
|
|
||||||
assert_has_worker_event(&telemetry, "tinygrad_imported");
|
|
||||||
assert_has_worker_event(&telemetry, "realizing_cuda_probe");
|
|
||||||
assert_has_worker_event(&telemetry, "backend_initialized");
|
|
||||||
assert_has_worker_event(&telemetry, "worker_ready");
|
|
||||||
assert_has_worker_event(&telemetry, "ring_installed");
|
|
||||||
assert_has_worker_event(&telemetry, "role_loaded");
|
|
||||||
assert_has_worker_event(&telemetry, "object_copy_started");
|
|
||||||
assert_has_worker_event_with(&telemetry, "object_loaded", |event| {
|
|
||||||
event["object_id"] == 9000 && event["sequence"] == 0 && event["device_sum"] == 10
|
|
||||||
});
|
|
||||||
assert_has_worker_event(&telemetry, "execute_step_started");
|
|
||||||
assert_has_worker_event_with(&telemetry, "object_produced", |event| {
|
|
||||||
event["object_id"] == 9001 && event["sequence"] == 0 && event["device_sum"] == 20
|
|
||||||
});
|
|
||||||
assert_has_worker_event(&telemetry, "step_completed");
|
|
||||||
assert_has_worker_event(&telemetry, "device_object_released");
|
|
||||||
assert_has_worker_event(&telemetry, "worker_stopped");
|
|
||||||
|
|
||||||
ingest_alive.store(false, Ordering::SeqCst);
|
|
||||||
let _ = ingest_thread.join();
|
|
||||||
let _ = std::fs::remove_file(&socket_path);
|
|
||||||
unsafe {
|
|
||||||
libc::close(arena_fd);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
struct ChannelFrameSink {
|
|
||||||
tx: mpsc::Sender<Frame>,
|
|
||||||
}
|
|
||||||
|
|
||||||
impl FrameSink for ChannelFrameSink {
|
|
||||||
fn ship(&mut self, _stream: &StreamId, frame: &Frame) {
|
|
||||||
self.tx.send(frame.clone()).expect("ship datastream frame");
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
#[derive(Clone)]
|
|
||||||
struct StartWorker {
|
|
||||||
worker_path: PathBuf,
|
|
||||||
socket_path: PathBuf,
|
|
||||||
arena_fd: RawFd,
|
|
||||||
arena_bytes: u64,
|
|
||||||
}
|
|
||||||
|
|
||||||
#[derive(Clone)]
|
|
||||||
enum NodeMsg {
|
|
||||||
Start(StartWorker),
|
|
||||||
SendCommand(Value),
|
|
||||||
ControlLine(String),
|
|
||||||
StderrLine(String),
|
|
||||||
ProcessExited(i32),
|
|
||||||
KillWorker(String),
|
|
||||||
}
|
|
||||||
|
|
||||||
#[derive(Clone, Debug)]
|
|
||||||
enum HarnessReport {
|
|
||||||
ProcessStarted,
|
|
||||||
ControlEvent(Value),
|
|
||||||
StderrLine(String),
|
|
||||||
ProcessExited(i32),
|
|
||||||
}
|
|
||||||
|
|
||||||
struct GpuWorkerNodeActor {
|
|
||||||
sender: ExternalSender,
|
|
||||||
report_to: ActorAddress,
|
|
||||||
stdin: Option<std::process::ChildStdin>,
|
|
||||||
child_pid: Option<u32>,
|
|
||||||
}
|
|
||||||
|
|
||||||
impl GpuWorkerNodeActor {
|
|
||||||
fn new(sender: ExternalSender, report_to: ActorAddress) -> Self {
|
|
||||||
Self {
|
|
||||||
sender,
|
|
||||||
report_to,
|
|
||||||
stdin: None,
|
|
||||||
child_pid: None,
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
impl ActorInterface for GpuWorkerNodeActor {
|
|
||||||
type Incoming = NodeMsg;
|
|
||||||
type Response = ();
|
|
||||||
|
|
||||||
fn handle(&mut self, ctx: &Ctx, msg: NodeMsg) {
|
|
||||||
match msg {
|
|
||||||
NodeMsg::Start(start) => self.start_worker(ctx, start),
|
|
||||||
NodeMsg::SendCommand(value) => {
|
|
||||||
let stdin = self.stdin.as_mut().expect("worker stdin");
|
|
||||||
writeln!(stdin, "{}", value).expect("write worker command");
|
|
||||||
stdin.flush().expect("flush worker command");
|
|
||||||
}
|
|
||||||
NodeMsg::ControlLine(line) => {
|
|
||||||
let value: Value = serde_json::from_str(&line).expect("control JSON");
|
|
||||||
ctx.send(self.report_to, HarnessReport::ControlEvent(value))
|
|
||||||
.expect("send control report");
|
|
||||||
}
|
|
||||||
NodeMsg::StderrLine(line) => {
|
|
||||||
ctx.send(self.report_to, HarnessReport::StderrLine(line))
|
|
||||||
.expect("send stderr report");
|
|
||||||
}
|
|
||||||
NodeMsg::ProcessExited(code) => {
|
|
||||||
self.child_pid = None;
|
|
||||||
ctx.send(self.report_to, HarnessReport::ProcessExited(code))
|
|
||||||
.expect("send exit report");
|
|
||||||
}
|
|
||||||
NodeMsg::KillWorker(reason) => {
|
|
||||||
self.kill_worker(&reason);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
impl GpuWorkerNodeActor {
|
|
||||||
fn start_worker(&mut self, ctx: &Ctx, start: StartWorker) {
|
|
||||||
let mut command = Command::new(start.worker_path);
|
|
||||||
command
|
|
||||||
.stdin(Stdio::piped())
|
|
||||||
.stdout(Stdio::piped())
|
|
||||||
.stderr(Stdio::piped())
|
|
||||||
.env("SWACTOR_ARENA_FD", start.arena_fd.to_string())
|
|
||||||
.env("SWACTOR_ARENA_BYTES", start.arena_bytes.to_string())
|
|
||||||
.env("SWACTOR_WORKER_EVENT_SOCK", &start.socket_path)
|
|
||||||
.env("SWACTOR_NODE_ID", "11")
|
|
||||||
.env("SWACTOR_RUN_ID", "77")
|
|
||||||
.env("SWACTOR_STAGE_INDEX", "0")
|
|
||||||
.env("DEV", "CUDA");
|
|
||||||
unsafe {
|
|
||||||
command.pre_exec(|| Ok(()));
|
|
||||||
}
|
|
||||||
let mut child = command.spawn().expect("spawn tinygrad worker");
|
|
||||||
let stdout = child.stdout.take().expect("worker stdout");
|
|
||||||
let stderr = child.stderr.take().expect("worker stderr");
|
|
||||||
self.stdin = Some(child.stdin.take().expect("worker stdin"));
|
|
||||||
self.child_pid = Some(child.id());
|
|
||||||
|
|
||||||
let target = ctx.self_addr();
|
|
||||||
let sender = self.sender.clone();
|
|
||||||
thread::spawn(move || {
|
|
||||||
for line in BufReader::new(stdout).lines() {
|
|
||||||
match line {
|
|
||||||
Ok(line) => {
|
|
||||||
if sender.send_to(target, NodeMsg::ControlLine(line)).is_err() {
|
|
||||||
break;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
Err(_) => break,
|
|
||||||
}
|
|
||||||
}
|
|
||||||
});
|
|
||||||
|
|
||||||
let target = ctx.self_addr();
|
|
||||||
let sender = self.sender.clone();
|
|
||||||
thread::spawn(move || {
|
|
||||||
for line in BufReader::new(stderr).lines() {
|
|
||||||
match line {
|
|
||||||
Ok(line) => {
|
|
||||||
eprintln!("gpu-worker-node-e2e worker stderr: {line}");
|
|
||||||
if sender.send_to(target, NodeMsg::StderrLine(line)).is_err() {
|
|
||||||
break;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
Err(_) => break,
|
|
||||||
}
|
|
||||||
}
|
|
||||||
});
|
|
||||||
|
|
||||||
let target = ctx.self_addr();
|
|
||||||
let sender = self.sender.clone();
|
|
||||||
thread::spawn(move || {
|
|
||||||
let code = child
|
|
||||||
.wait()
|
|
||||||
.ok()
|
|
||||||
.and_then(|status| status.code())
|
|
||||||
.unwrap_or(-1);
|
|
||||||
let _ = sender.send_to(target, NodeMsg::ProcessExited(code));
|
|
||||||
});
|
|
||||||
|
|
||||||
ctx.send(self.report_to, HarnessReport::ProcessStarted)
|
|
||||||
.expect("send started report");
|
|
||||||
}
|
|
||||||
|
|
||||||
fn kill_worker(&mut self, reason: &str) {
|
|
||||||
if let Some(pid) = self.child_pid.take() {
|
|
||||||
eprintln!("gpu-worker-node-e2e: killing worker pid {pid}: {reason}");
|
|
||||||
unsafe {
|
|
||||||
libc::kill(pid as libc::pid_t, libc::SIGKILL);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
fn spawn_worker_event_ingest(
|
|
||||||
socket_path: PathBuf,
|
|
||||||
sink: datastream::emit::DatastreamEventSink,
|
|
||||||
alive: Arc<AtomicBool>,
|
|
||||||
) -> thread::JoinHandle<()> {
|
|
||||||
let (ready_tx, ready_rx) = mpsc::channel();
|
|
||||||
let handle = thread::spawn(move || {
|
|
||||||
let socket = std::os::unix::net::UnixDatagram::bind(&socket_path).expect("bind UDS ingest");
|
|
||||||
ready_tx.send(()).expect("signal UDS ingest ready");
|
|
||||||
socket
|
|
||||||
.set_read_timeout(Some(Duration::from_millis(50)))
|
|
||||||
.expect("set UDS read watchdog");
|
|
||||||
let mut buf = vec![0u8; 8192];
|
|
||||||
while alive.load(Ordering::SeqCst) {
|
|
||||||
match socket.recv(&mut buf) {
|
|
||||||
Ok(len) => {
|
|
||||||
sink.submit_bytes(WORKER_EVENTS_CHANNEL, buf[..len].to_vec());
|
|
||||||
}
|
|
||||||
Err(error)
|
|
||||||
if error.kind() == std::io::ErrorKind::WouldBlock
|
|
||||||
|| error.kind() == std::io::ErrorKind::TimedOut => {}
|
|
||||||
Err(_) => break,
|
|
||||||
}
|
|
||||||
}
|
|
||||||
});
|
|
||||||
ready_rx
|
|
||||||
.recv_timeout(Duration::from_secs(5))
|
|
||||||
.expect("UDS ingest socket ready");
|
|
||||||
handle
|
|
||||||
}
|
|
||||||
|
|
||||||
fn send_command(rt: &Runtime, node: ActorAddress, command: Value) {
|
|
||||||
rt.send_to(node, NodeMsg::SendCommand(command))
|
|
||||||
.expect("send node command");
|
|
||||||
}
|
|
||||||
|
|
||||||
fn wait_for_control(
|
|
||||||
rt: &Runtime,
|
|
||||||
node: ActorAddress,
|
|
||||||
emitter: &mut DatastreamEmitter,
|
|
||||||
frame_rx: &mpsc::Receiver<Frame>,
|
|
||||||
reports: &swactor::runtime::Inbox<HarnessReport>,
|
|
||||||
kind: &str,
|
|
||||||
phase_name: &str,
|
|
||||||
) -> Value {
|
|
||||||
match wait_for_report(
|
|
||||||
rt,
|
|
||||||
node,
|
|
||||||
emitter,
|
|
||||||
frame_rx,
|
|
||||||
reports,
|
|
||||||
phase_name,
|
|
||||||
|report| matches!(report, HarnessReport::ControlEvent(value) if value["type"] == kind),
|
|
||||||
) {
|
|
||||||
HarnessReport::ControlEvent(value) => value,
|
|
||||||
other => panic!("unexpected report for {kind}: {other:?}"),
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
fn wait_for_report(
|
|
||||||
rt: &Runtime,
|
|
||||||
node: ActorAddress,
|
|
||||||
emitter: &mut DatastreamEmitter,
|
|
||||||
_frame_rx: &mpsc::Receiver<Frame>,
|
|
||||||
reports: &swactor::runtime::Inbox<HarnessReport>,
|
|
||||||
phase_name: &str,
|
|
||||||
mut predicate: impl FnMut(&HarnessReport) -> bool,
|
|
||||||
) -> HarnessReport {
|
|
||||||
let started = Instant::now();
|
|
||||||
let mut stderr_lines = Vec::new();
|
|
||||||
while started.elapsed() < EVENT_WATCHDOG {
|
|
||||||
rt.tick();
|
|
||||||
emitter.tick();
|
|
||||||
while let Some(report) = reports.try_recv() {
|
|
||||||
if predicate(&report) {
|
|
||||||
return report;
|
|
||||||
}
|
|
||||||
match &report {
|
|
||||||
HarnessReport::StderrLine(line) => stderr_lines.push(line.clone()),
|
|
||||||
HarnessReport::ControlEvent(value) if value["type"] == "WorkerFatal" => {
|
|
||||||
rt.send_to(
|
|
||||||
node,
|
|
||||||
NodeMsg::KillWorker(format!("fatal while waiting for {phase_name}")),
|
|
||||||
)
|
|
||||||
.expect("send kill after fatal");
|
|
||||||
rt.tick();
|
|
||||||
panic!(
|
|
||||||
"worker fatal while waiting for {phase_name}: {value}\nstderr={stderr_lines:?}"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
HarnessReport::ProcessExited(code) if *code != 0 => {
|
|
||||||
panic!(
|
|
||||||
"worker exited with {code} while waiting for {phase_name}; stderr={stderr_lines:?}"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
_ => {}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
thread::sleep(Duration::from_millis(5));
|
|
||||||
}
|
|
||||||
panic!(
|
|
||||||
"test watchdog after {:?} waiting for {phase_name}; stderr={stderr_lines:?}",
|
|
||||||
EVENT_WATCHDOG
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
fn collect_telemetry(
|
|
||||||
emitter: &mut DatastreamEmitter,
|
|
||||||
frame_rx: &mpsc::Receiver<Frame>,
|
|
||||||
duration: Duration,
|
|
||||||
) -> Vec<Value> {
|
|
||||||
let started = Instant::now();
|
|
||||||
let mut frames = Vec::new();
|
|
||||||
while started.elapsed() < duration {
|
|
||||||
emitter.tick();
|
|
||||||
while let Ok(frame) = frame_rx.try_recv() {
|
|
||||||
frames.push(frame);
|
|
||||||
}
|
|
||||||
thread::sleep(Duration::from_millis(10));
|
|
||||||
}
|
|
||||||
frames
|
|
||||||
.into_iter()
|
|
||||||
.filter(|frame| frame.channel.as_str() == WORKER_EVENTS_CHANNEL)
|
|
||||||
.map(|frame| serde_json::from_slice::<Value>(&frame.payload).expect("telemetry JSON"))
|
|
||||||
.collect()
|
|
||||||
}
|
|
||||||
|
|
||||||
fn assert_has_worker_event(events: &[Value], kind: &str) {
|
|
||||||
assert_has_worker_event_with(events, kind, |_| true);
|
|
||||||
}
|
|
||||||
|
|
||||||
fn assert_has_worker_event_with(events: &[Value], kind: &str, extra: impl Fn(&Value) -> bool) {
|
|
||||||
assert!(
|
|
||||||
events.iter().any(|event| {
|
|
||||||
event["schema"] == "mvp.worker.event.v1"
|
|
||||||
&& event["kind"] == kind
|
|
||||||
&& event["node_id"] == 11
|
|
||||||
&& event["run_id"] == 77
|
|
||||||
&& event["stage_index"] == 0
|
|
||||||
&& event["worker_generation"] == 1
|
|
||||||
&& extra(event)
|
|
||||||
}),
|
|
||||||
"missing worker event {kind}; events={events:#?}"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
fn install_ring_command(
|
|
||||||
ring_id: u64,
|
|
||||||
edge_id: u64,
|
|
||||||
port_id: &str,
|
|
||||||
direction: &str,
|
|
||||||
base: usize,
|
|
||||||
) -> Value {
|
|
||||||
json!({
|
|
||||||
"type":"InstallRing",
|
|
||||||
"ring_id": ring_id,
|
|
||||||
"edge_id": edge_id,
|
|
||||||
"port_id": port_id,
|
|
||||||
"direction": direction,
|
|
||||||
"base": base,
|
|
||||||
"bytes": RING_BYTES,
|
|
||||||
"object_spec": {
|
|
||||||
"max_extent": 16,
|
|
||||||
"alignment": 4,
|
|
||||||
"layout": "token"
|
|
||||||
}
|
|
||||||
})
|
|
||||||
}
|
|
||||||
|
|
||||||
fn object_record(object_id: u64, sequence: u64, words: &[i32]) -> Vec<u8> {
|
|
||||||
let mut record = vec![0u8; HEADER_LEN];
|
|
||||||
record[0..4].copy_from_slice(b"MO01");
|
|
||||||
record[4] = 1;
|
|
||||||
record[5] = HEADER_LEN as u8;
|
|
||||||
record[8..16].copy_from_slice(&object_id.to_le_bytes());
|
|
||||||
record[16..24].copy_from_slice(&sequence.to_le_bytes());
|
|
||||||
record[24..32].copy_from_slice(&((words.len() * 4) as u64).to_le_bytes());
|
|
||||||
record[32..40].copy_from_slice(&16u64.to_le_bytes());
|
|
||||||
record[40..48].copy_from_slice(&4u64.to_le_bytes());
|
|
||||||
for word in words {
|
|
||||||
record.extend_from_slice(&word.to_le_bytes());
|
|
||||||
}
|
|
||||||
record
|
|
||||||
}
|
|
||||||
|
|
||||||
fn decode_payload_words(record: &[u8]) -> Vec<i32> {
|
|
||||||
assert_eq!(&record[0..4], b"MO01");
|
|
||||||
let extent = u64::from_le_bytes(record[24..32].try_into().unwrap()) as usize;
|
|
||||||
record[HEADER_LEN..HEADER_LEN + extent]
|
|
||||||
.chunks_exact(4)
|
|
||||||
.map(|chunk| i32::from_le_bytes(chunk.try_into().unwrap()))
|
|
||||||
.collect()
|
|
||||||
}
|
|
||||||
|
|
||||||
fn create_arena(bytes: usize) -> RawFd {
|
|
||||||
let name = CString::new("mvp-system-gpu-worker-node-e2e").expect("memfd name");
|
|
||||||
let fd = unsafe { libc::memfd_create(name.as_ptr(), 0) };
|
|
||||||
assert!(
|
|
||||||
fd >= 0,
|
|
||||||
"memfd_create failed: {}",
|
|
||||||
std::io::Error::last_os_error()
|
|
||||||
);
|
|
||||||
let truncate = unsafe { libc::ftruncate(fd, bytes as libc::off_t) };
|
|
||||||
assert_eq!(
|
|
||||||
truncate,
|
|
||||||
0,
|
|
||||||
"ftruncate failed: {}",
|
|
||||||
std::io::Error::last_os_error()
|
|
||||||
);
|
|
||||||
fd
|
|
||||||
}
|
|
||||||
|
|
||||||
fn pwrite_all(fd: RawFd, offset: usize, bytes: &[u8]) {
|
|
||||||
let written = unsafe {
|
|
||||||
libc::pwrite(
|
|
||||||
fd,
|
|
||||||
bytes.as_ptr().cast(),
|
|
||||||
bytes.len(),
|
|
||||||
offset as libc::off_t,
|
|
||||||
)
|
|
||||||
};
|
|
||||||
assert_eq!(
|
|
||||||
written,
|
|
||||||
bytes.len() as isize,
|
|
||||||
"pwrite failed: {}",
|
|
||||||
std::io::Error::last_os_error()
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
fn pread_exact(fd: RawFd, offset: usize, bytes: &mut [u8]) {
|
|
||||||
let read = unsafe {
|
|
||||||
libc::pread(
|
|
||||||
fd,
|
|
||||||
bytes.as_mut_ptr().cast(),
|
|
||||||
bytes.len(),
|
|
||||||
offset as libc::off_t,
|
|
||||||
)
|
|
||||||
};
|
|
||||||
assert_eq!(
|
|
||||||
read,
|
|
||||||
bytes.len() as isize,
|
|
||||||
"pread failed: {}",
|
|
||||||
std::io::Error::last_os_error()
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
@ -1,23 +0,0 @@
|
||||||
FROM nvidia/cuda:12.6.3-runtime-ubuntu24.04
|
|
||||||
|
|
||||||
RUN apt-get update && \
|
|
||||||
DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \
|
|
||||||
ca-certificates \
|
|
||||||
curl \
|
|
||||||
build-essential \
|
|
||||||
pkg-config \
|
|
||||||
python3 \
|
|
||||||
python3-pip \
|
|
||||||
cuda-cudart-dev-12-6 \
|
|
||||||
cuda-nvrtc-12-6 && \
|
|
||||||
python3 -m pip install --no-cache-dir --break-system-packages tinygrad==0.12.0 numpy && \
|
|
||||||
curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y --profile minimal && \
|
|
||||||
apt-get clean && \
|
|
||||||
rm -rf /var/lib/apt/lists/*
|
|
||||||
|
|
||||||
ENV PATH=/root/.cargo/bin:$PATH
|
|
||||||
ENV DEV=CUDA
|
|
||||||
ENV PYTHONDONTWRITEBYTECODE=1
|
|
||||||
COPY . /workspace
|
|
||||||
WORKDIR /workspace
|
|
||||||
CMD ["sh", "-lc", "python3 -c 'import os; print(\"cuda preflight: importing tinygrad\", flush=True); from tinygrad import Tensor,dtypes; print(\"cuda preflight: DEV=\" + str(os.environ.get(\"DEV\")), flush=True); print(\"cuda preflight: realizing Tensor([1])\", flush=True); print(Tensor([1], dtype=dtypes.int32).realize().numpy().tolist(), flush=True)' && cargo test -p mvp-system --features local-e2e --test gpu_worker_node_e2e -- --nocapture"]
|
|
||||||
|
|
@ -1,228 +0,0 @@
|
||||||
#!/usr/bin/env python3
|
|
||||||
from __future__ import annotations
|
|
||||||
|
|
||||||
import json
|
|
||||||
import mmap
|
|
||||||
import os
|
|
||||||
import socket
|
|
||||||
import struct
|
|
||||||
import sys
|
|
||||||
import time
|
|
||||||
from typing import Any
|
|
||||||
|
|
||||||
|
|
||||||
HEADER_LEN = 48
|
|
||||||
GENERATION = 1
|
|
||||||
|
|
||||||
arena: mmap.mmap | None = None
|
|
||||||
telemetry: socket.socket | None = None
|
|
||||||
telemetry_path = os.environ["SWACTOR_WORKER_EVENT_SOCK"]
|
|
||||||
rings: dict[int, dict[str, Any]] = {}
|
|
||||||
objects: dict[int, dict[str, Any]] = {}
|
|
||||||
role_configured = False
|
|
||||||
Tensor: Any = None
|
|
||||||
dtypes: Any = None
|
|
||||||
|
|
||||||
next_handle = 42
|
|
||||||
|
|
||||||
|
|
||||||
def control(**event: Any) -> None:
|
|
||||||
print(json.dumps(event, separators=(",", ":")), flush=True)
|
|
||||||
|
|
||||||
|
|
||||||
def log(message: str) -> None:
|
|
||||||
print(f"mvp_tinygrad_worker: {message}", file=sys.stderr, flush=True)
|
|
||||||
|
|
||||||
|
|
||||||
def observe(kind: str, **fields: Any) -> None:
|
|
||||||
global telemetry
|
|
||||||
if telemetry is None:
|
|
||||||
telemetry = socket.socket(socket.AF_UNIX, socket.SOCK_DGRAM)
|
|
||||||
event = {
|
|
||||||
"schema": "mvp.worker.event.v1",
|
|
||||||
"kind": kind,
|
|
||||||
"node_id": int(os.environ["SWACTOR_NODE_ID"]),
|
|
||||||
"run_id": int(os.environ["SWACTOR_RUN_ID"]),
|
|
||||||
"stage_index": int(os.environ["SWACTOR_STAGE_INDEX"]),
|
|
||||||
"worker_generation": GENERATION,
|
|
||||||
"ts_ns": time.time_ns(),
|
|
||||||
}
|
|
||||||
event.update(fields)
|
|
||||||
telemetry.sendto(json.dumps(event, separators=(",", ":")).encode(), telemetry_path)
|
|
||||||
|
|
||||||
|
|
||||||
def fatal(reason: str, **fields: Any) -> None:
|
|
||||||
observe("worker_fatal", reason=reason, **fields)
|
|
||||||
control(type="WorkerFatal", reason=reason, **fields)
|
|
||||||
raise SystemExit(1)
|
|
||||||
|
|
||||||
|
|
||||||
def require_arena() -> mmap.mmap:
|
|
||||||
if arena is None:
|
|
||||||
fatal("ArenaNotMapped")
|
|
||||||
return arena
|
|
||||||
|
|
||||||
|
|
||||||
def require_tinygrad() -> tuple[Any, Any]:
|
|
||||||
if Tensor is None or dtypes is None:
|
|
||||||
fatal("BackendNotInitialized")
|
|
||||||
return Tensor, dtypes
|
|
||||||
|
|
||||||
|
|
||||||
def initialize(cmd: dict[str, Any]) -> None:
|
|
||||||
global arena, Tensor, dtypes
|
|
||||||
if int(cmd["helper_abi_version"]) != 1:
|
|
||||||
fatal("UnsupportedHelperAbi", helper_abi_version=cmd["helper_abi_version"])
|
|
||||||
fd = int(os.environ["SWACTOR_ARENA_FD"])
|
|
||||||
size = int(os.environ["SWACTOR_ARENA_BYTES"])
|
|
||||||
arena = mmap.mmap(fd, size)
|
|
||||||
|
|
||||||
log("importing tinygrad")
|
|
||||||
observe("importing_tinygrad")
|
|
||||||
import_start = time.monotonic()
|
|
||||||
from tinygrad import Tensor as TinyTensor, dtypes as tiny_dtypes
|
|
||||||
|
|
||||||
Tensor = TinyTensor
|
|
||||||
dtypes = tiny_dtypes
|
|
||||||
observe("tinygrad_imported", elapsed_ms=int((time.monotonic() - import_start) * 1000))
|
|
||||||
|
|
||||||
log(f"realizing CUDA probe with DEV={os.environ.get('DEV')}")
|
|
||||||
observe("realizing_cuda_probe", dev=os.environ.get("DEV"))
|
|
||||||
probe_start = time.monotonic()
|
|
||||||
Tensor([1], dtype=dtypes.int32).realize().numpy().tolist()
|
|
||||||
observe("backend_initialized", elapsed_ms=int((time.monotonic() - probe_start) * 1000))
|
|
||||||
observe("worker_ready")
|
|
||||||
control(type="WorkerReady", generation=GENERATION)
|
|
||||||
|
|
||||||
|
|
||||||
def install_ring(cmd: dict[str, Any]) -> None:
|
|
||||||
ring_id = int(cmd["ring_id"])
|
|
||||||
rings[ring_id] = {
|
|
||||||
"ring_id": ring_id,
|
|
||||||
"edge_id": int(cmd["edge_id"]),
|
|
||||||
"port_id": cmd["port_id"],
|
|
||||||
"direction": cmd["direction"],
|
|
||||||
"base": int(cmd["base"]),
|
|
||||||
"bytes": int(cmd["bytes"]),
|
|
||||||
"max_extent": int(cmd["object_spec"]["max_extent"]),
|
|
||||||
"alignment": int(cmd["object_spec"]["alignment"]),
|
|
||||||
}
|
|
||||||
observe("ring_installed", ring_id=ring_id, edge_id=rings[ring_id]["edge_id"], direction=rings[ring_id]["direction"])
|
|
||||||
control(type="RingInstalled", ring_id=ring_id)
|
|
||||||
|
|
||||||
|
|
||||||
def configure_role(cmd: dict[str, Any]) -> None:
|
|
||||||
global role_configured
|
|
||||||
role_configured = True
|
|
||||||
observe("role_loaded", role_id=int(cmd["role_id"]))
|
|
||||||
control(type="RoleLoaded", role_id=int(cmd["role_id"]))
|
|
||||||
|
|
||||||
|
|
||||||
def parse_record(base: int, committed_bytes: int, ring: dict[str, Any]) -> tuple[int, int, int, bytes]:
|
|
||||||
view = require_arena()
|
|
||||||
if committed_bytes < HEADER_LEN:
|
|
||||||
fatal("MalformedHeaderLength", ring_id=ring["ring_id"])
|
|
||||||
header = view[base : base + HEADER_LEN]
|
|
||||||
if header[0:4] != b"MO01" or header[4] != 1 or header[5] != HEADER_LEN:
|
|
||||||
fatal("InvalidObjectHeader", ring_id=ring["ring_id"])
|
|
||||||
object_id = struct.unpack_from("<Q", header, 8)[0]
|
|
||||||
sequence = struct.unpack_from("<Q", header, 16)[0]
|
|
||||||
extent = struct.unpack_from("<Q", header, 24)[0]
|
|
||||||
encoded_max = struct.unpack_from("<Q", header, 32)[0]
|
|
||||||
encoded_alignment = struct.unpack_from("<Q", header, 40)[0]
|
|
||||||
if encoded_max != ring["max_extent"] or encoded_alignment != ring["alignment"]:
|
|
||||||
fatal("ObjectSpecMismatch", ring_id=ring["ring_id"], object_id=object_id)
|
|
||||||
if extent > ring["max_extent"] or (ring["alignment"] and extent % ring["alignment"]):
|
|
||||||
fatal("ObjectExtentInvalid", ring_id=ring["ring_id"], object_id=object_id, extent=extent)
|
|
||||||
total = HEADER_LEN + extent
|
|
||||||
if committed_bytes < total:
|
|
||||||
fatal("EofBeforeFullPayload", ring_id=ring["ring_id"], object_id=object_id)
|
|
||||||
payload = bytes(view[base + HEADER_LEN : base + total])
|
|
||||||
return object_id, sequence, extent, payload
|
|
||||||
|
|
||||||
|
|
||||||
def ring_readable(cmd: dict[str, Any]) -> None:
|
|
||||||
global next_handle
|
|
||||||
Tensor, dtypes = require_tinygrad()
|
|
||||||
ring_id = int(cmd["ring_id"])
|
|
||||||
ring = rings[ring_id]
|
|
||||||
if ring["direction"] != "ingress":
|
|
||||||
fatal("WrongRingDirection", ring_id=ring_id)
|
|
||||||
object_id, sequence, extent, payload = parse_record(ring["base"], int(cmd["committed_bytes"]), ring)
|
|
||||||
observe("object_copy_started", ring_id=ring_id, edge_id=ring["edge_id"], object_id=object_id, sequence=sequence, extent=extent)
|
|
||||||
values = list(struct.unpack(f"<{extent // 4}i", payload))
|
|
||||||
tensor = Tensor(values, dtype=dtypes.int32).realize()
|
|
||||||
handle = next_handle
|
|
||||||
next_handle += 1
|
|
||||||
device_sum = int(tensor.sum().item())
|
|
||||||
objects[handle] = {"object_id": object_id, "sequence": sequence, "tensor": tensor, "extent": extent}
|
|
||||||
observe("object_loaded", ring_id=ring_id, edge_id=ring["edge_id"], object_id=object_id, sequence=sequence, extent=extent, handle=handle, device_sum=device_sum)
|
|
||||||
control(type="ObjectLoaded", ring_id=ring_id, edge_id=ring["edge_id"], object_id=object_id, sequence=sequence, extent=extent, handle={"generation": GENERATION, "id": handle})
|
|
||||||
|
|
||||||
|
|
||||||
def write_record(ring: dict[str, Any], object_id: int, sequence: int, words: list[int]) -> int:
|
|
||||||
payload = b"".join(struct.pack("<i", word) for word in words)
|
|
||||||
extent = len(payload)
|
|
||||||
if extent > ring["max_extent"]:
|
|
||||||
fatal("OutputExtentInvalid", ring_id=ring["ring_id"], extent=extent)
|
|
||||||
header = bytearray(HEADER_LEN)
|
|
||||||
header[0:4] = b"MO01"
|
|
||||||
header[4] = 1
|
|
||||||
header[5] = HEADER_LEN
|
|
||||||
struct.pack_into("<Q", header, 8, object_id)
|
|
||||||
struct.pack_into("<Q", header, 16, sequence)
|
|
||||||
struct.pack_into("<Q", header, 24, extent)
|
|
||||||
struct.pack_into("<Q", header, 32, ring["max_extent"])
|
|
||||||
struct.pack_into("<Q", header, 40, ring["alignment"])
|
|
||||||
view = require_arena()
|
|
||||||
base = ring["base"]
|
|
||||||
view[base : base + HEADER_LEN] = header
|
|
||||||
view[base + HEADER_LEN : base + HEADER_LEN + extent] = payload
|
|
||||||
return HEADER_LEN + extent
|
|
||||||
|
|
||||||
|
|
||||||
def execute_step(cmd: dict[str, Any]) -> None:
|
|
||||||
if not role_configured:
|
|
||||||
fatal("RoleNotConfigured")
|
|
||||||
handle = int(cmd["input_handle"])
|
|
||||||
step_id = int(cmd["step_id"])
|
|
||||||
output_object_id = int(cmd["output_object_id"])
|
|
||||||
egress_ring_id = int(cmd["egress_ring_id"])
|
|
||||||
obj = objects[handle]
|
|
||||||
observe("execute_step_started", step_id=step_id, object_id=obj["object_id"], sequence=obj["sequence"], handle=handle)
|
|
||||||
output = (obj["tensor"] * 2).realize()
|
|
||||||
words = [int(value) for value in output.numpy().tolist()]
|
|
||||||
ring = rings[egress_ring_id]
|
|
||||||
committed = write_record(ring, output_object_id, int(obj["sequence"]), words)
|
|
||||||
output_sum = sum(words)
|
|
||||||
observe("object_produced", ring_id=egress_ring_id, edge_id=ring["edge_id"], object_id=output_object_id, sequence=obj["sequence"], extent=len(words) * 4, device_sum=output_sum, committed_bytes=committed)
|
|
||||||
observe("step_completed", step_id=step_id)
|
|
||||||
control(type="ObjectProduced", ring_id=egress_ring_id, object_id=output_object_id, sequence=obj["sequence"], committed_bytes=committed)
|
|
||||||
control(type="StepCompleted", step_id=step_id)
|
|
||||||
|
|
||||||
|
|
||||||
def release_device_object(cmd: dict[str, Any]) -> None:
|
|
||||||
handle = int(cmd["handle"])
|
|
||||||
objects.pop(handle, None)
|
|
||||||
observe("device_object_released", handle=handle)
|
|
||||||
control(type="DeviceObjectReleased", handle=handle)
|
|
||||||
|
|
||||||
|
|
||||||
handlers = {
|
|
||||||
"InitializeWorker": initialize,
|
|
||||||
"InstallRing": install_ring,
|
|
||||||
"ConfigureRole": configure_role,
|
|
||||||
"RingReadable": ring_readable,
|
|
||||||
"ExecuteStep": execute_step,
|
|
||||||
"ReleaseDeviceObject": release_device_object,
|
|
||||||
}
|
|
||||||
|
|
||||||
for raw in sys.stdin:
|
|
||||||
if not raw.strip():
|
|
||||||
continue
|
|
||||||
command = json.loads(raw)
|
|
||||||
if command["type"] == "ShutdownWorker":
|
|
||||||
observe("worker_stopped")
|
|
||||||
control(type="WorkerStopped", generation=GENERATION)
|
|
||||||
break
|
|
||||||
handlers[command["type"]](command)
|
|
||||||
|
|
@ -1,356 +0,0 @@
|
||||||
#![recursion_limit = "256"]
|
|
||||||
|
|
||||||
use std::path::Path;
|
|
||||||
use std::process::{Command, ExitCode};
|
|
||||||
use std::time::{Duration, Instant};
|
|
||||||
|
|
||||||
#[path = "support/local_e2e_cluster.rs"]
|
|
||||||
mod local_e2e_cluster;
|
|
||||||
|
|
||||||
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 {
|
|
||||||
let args = std::env::args().collect::<Vec<_>>();
|
|
||||||
match std::env::var("MVP_TEST_ROLE").ok().as_deref() {
|
|
||||||
Some("cluster-supervisor" | "cluster-relay") => return local_e2e_cluster::run_main(),
|
|
||||||
Some(role) => {
|
|
||||||
eprintln!("unknown MVP_TEST_ROLE={role}");
|
|
||||||
return ExitCode::from(2);
|
|
||||||
}
|
|
||||||
None => {}
|
|
||||||
}
|
|
||||||
|
|
||||||
if args.iter().any(|arg| arg == "--role=node") {
|
|
||||||
return local_e2e_cluster::run_main();
|
|
||||||
}
|
|
||||||
|
|
||||||
local_e2e_cluster_docker_cpu_pipeline_prompt();
|
|
||||||
ExitCode::SUCCESS
|
|
||||||
}
|
|
||||||
|
|
||||||
fn local_e2e_cluster_docker_cpu_pipeline_prompt() {
|
|
||||||
if std::env::var_os("MVP_SYSTEM_LOCAL_E2E_CLUSTER").is_none() {
|
|
||||||
eprintln!("skipping; set MVP_SYSTEM_LOCAL_E2E_CLUSTER=1 to run Docker CPU cluster e2e");
|
|
||||||
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();
|
|
||||||
if std::env::var_os(BUILD_ONLY_ENV).is_some() {
|
|
||||||
return;
|
|
||||||
}
|
|
||||||
|
|
||||||
let docker = DockerRelayFixture::start();
|
|
||||||
let output = docker.run_supervisor("ping");
|
|
||||||
|
|
||||||
assert!(
|
|
||||||
output.status.success(),
|
|
||||||
"mvp-local-e2e-cluster failed\nstdout:\n{}\nstderr:\n{}",
|
|
||||||
String::from_utf8_lossy(&output.stdout),
|
|
||||||
String::from_utf8_lossy(&output.stderr)
|
|
||||||
);
|
|
||||||
|
|
||||||
let value: serde_json::Value = serde_json::from_slice(&output.stdout).expect("json stdout");
|
|
||||||
assert_eq!(value["ok"], true);
|
|
||||||
assert_eq!(value["actor_plane"], "iroh-swactor");
|
|
||||||
assert_eq!(value["data_plane"], "iroh-quic-persistent-edge-streams");
|
|
||||||
assert_eq!(
|
|
||||||
value["edge_protocol"],
|
|
||||||
"edge-id-preamble-mo01-object-records"
|
|
||||||
);
|
|
||||||
assert_eq!(
|
|
||||||
value["node_local_data_plane"],
|
|
||||||
"arena-backed-rings-json-metadata-only"
|
|
||||||
);
|
|
||||||
assert_eq!(
|
|
||||||
value["worker_processes"],
|
|
||||||
"docker-tinygrad-cpu-worker-per-node"
|
|
||||||
);
|
|
||||||
assert_eq!(value["tinygrad_device"], "CPU");
|
|
||||||
assert_eq!(value["prompt_text"], "ping");
|
|
||||||
assert_eq!(value["response_text"], "pong");
|
|
||||||
assert_eq!(value["response_tokens"].as_array().map(Vec::len), Some(1));
|
|
||||||
assert_eq!(
|
|
||||||
value["engine_builder_pattern"],
|
|
||||||
"host-coordinator-static-topology-docker-workers"
|
|
||||||
);
|
|
||||||
assert_eq!(value["engine_builder_node_count"], 3);
|
|
||||||
assert_eq!(value["engine_builder_stage_assignments"], 2);
|
|
||||||
assert_eq!(value["injected_prompt_observed"], true);
|
|
||||||
assert_eq!(value["token_received_observed"], true);
|
|
||||||
assert_eq!(value["run_completed_observed"], true);
|
|
||||||
assert_eq!(value["run_torn_down_observed"], true);
|
|
||||||
assert_eq!(value["stop_sent_to_all_nodes"], true);
|
|
||||||
assert_eq!(value["stage_ready_stdout_count"], 2);
|
|
||||||
assert_eq!(value["provisioned_node_count"], 2);
|
|
||||||
assert_eq!(value["provision_node_live_count"], 2);
|
|
||||||
assert!(
|
|
||||||
value["provision_stdout_line_count"]
|
|
||||||
.as_u64()
|
|
||||||
.is_some_and(|count| count >= 2),
|
|
||||||
"{value}"
|
|
||||||
);
|
|
||||||
assert!(
|
|
||||||
value["provision_stderr_line_count"]
|
|
||||||
.as_u64()
|
|
||||||
.is_some_and(|count| count >= 2),
|
|
||||||
"{value}"
|
|
||||||
);
|
|
||||||
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!(
|
|
||||||
value["node0_endpoint"]["addrs"]
|
|
||||||
.as_array()
|
|
||||||
.is_some_and(|addrs| !addrs.is_empty()),
|
|
||||||
"{value}"
|
|
||||||
);
|
|
||||||
assert!(
|
|
||||||
value["node1_endpoint"]["addrs"]
|
|
||||||
.as_array()
|
|
||||||
.is_some_and(|addrs| !addrs.is_empty()),
|
|
||||||
"{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() {
|
|
||||||
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 workspace = crate_dir
|
|
||||||
.parent()
|
|
||||||
.and_then(Path::parent)
|
|
||||||
.expect("workspace root")
|
|
||||||
.canonicalize()
|
|
||||||
.expect("canonical workspace root");
|
|
||||||
let context = std::env::temp_dir().join(format!(
|
|
||||||
"mvp-system-local-e2e-cluster-context-{}",
|
|
||||||
std::process::id()
|
|
||||||
));
|
|
||||||
let _ = std::fs::remove_dir_all(&context);
|
|
||||||
copy_workspace_context(&workspace, &context);
|
|
||||||
let dockerfile = context.join("crates/mvp-system/tests/local_e2e_cluster/Dockerfile");
|
|
||||||
|
|
||||||
phase("building Docker CPU cluster fixture image");
|
|
||||||
let build = Command::new("docker")
|
|
||||||
.args(["build", "-f"])
|
|
||||||
.arg(&dockerfile)
|
|
||||||
.args(["-t", IMAGE])
|
|
||||||
.arg(&context)
|
|
||||||
.status()
|
|
||||||
.expect("run docker build");
|
|
||||||
assert!(build.success(), "docker build failed with status {build}");
|
|
||||||
}
|
|
||||||
|
|
||||||
fn copy_workspace_context(source: &Path, dest: &Path) {
|
|
||||||
std::fs::create_dir_all(dest).expect("create docker context");
|
|
||||||
for entry in std::fs::read_dir(source).expect("read workspace") {
|
|
||||||
let entry = entry.expect("read workspace entry");
|
|
||||||
let name = entry.file_name();
|
|
||||||
let name = name.to_string_lossy();
|
|
||||||
if matches!(name.as_ref(), ".git" | "target" | ".dockerignore") {
|
|
||||||
continue;
|
|
||||||
}
|
|
||||||
copy_context_entry(&entry.path(), &dest.join(name.as_ref()));
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
fn copy_context_entry(source: &Path, dest: &Path) {
|
|
||||||
let metadata = std::fs::symlink_metadata(source).expect("context metadata");
|
|
||||||
if metadata.file_type().is_symlink() {
|
|
||||||
return;
|
|
||||||
}
|
|
||||||
if metadata.is_dir() {
|
|
||||||
std::fs::create_dir_all(dest).expect("create context dir");
|
|
||||||
for entry in std::fs::read_dir(source).expect("read context dir") {
|
|
||||||
let entry = entry.expect("read context entry");
|
|
||||||
let name = entry.file_name();
|
|
||||||
let name = name.to_string_lossy();
|
|
||||||
if matches!(name.as_ref(), ".git" | "target" | ".dockerignore") {
|
|
||||||
continue;
|
|
||||||
}
|
|
||||||
copy_context_entry(&entry.path(), &dest.join(name.as_ref()));
|
|
||||||
}
|
|
||||||
} else if metadata.is_file() {
|
|
||||||
if let Some(parent) = dest.parent() {
|
|
||||||
std::fs::create_dir_all(parent).expect("create context parent");
|
|
||||||
}
|
|
||||||
std::fs::copy(source, dest).expect("copy context file");
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
fn phase(message: &str) {
|
|
||||||
eprintln!("local-e2e-cluster: {message}");
|
|
||||||
}
|
|
||||||
|
|
@ -1,24 +0,0 @@
|
||||||
FROM rust:1-bookworm
|
|
||||||
|
|
||||||
RUN apt-get update && \
|
|
||||||
DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \
|
|
||||||
ca-certificates \
|
|
||||||
docker.io \
|
|
||||||
pkg-config \
|
|
||||||
python3 \
|
|
||||||
python3-pip && \
|
|
||||||
python3 -m pip install --no-cache-dir --break-system-packages tinygrad==0.12.0 numpy && \
|
|
||||||
apt-get clean && \
|
|
||||||
rm -rf /var/lib/apt/lists/*
|
|
||||||
|
|
||||||
ENV DEV=CPU
|
|
||||||
ENV PYTHONDONTWRITEBYTECODE=1
|
|
||||||
ENV CARGO_TARGET_DIR=/tmp/mvp-local-e2e-target
|
|
||||||
COPY . /workspace
|
|
||||||
WORKDIR /workspace
|
|
||||||
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 ["/usr/local/bin/local-e2e-cluster"]
|
|
||||||
|
|
@ -1,273 +0,0 @@
|
||||||
#!/usr/bin/env python3
|
|
||||||
from __future__ import annotations
|
|
||||||
|
|
||||||
import json
|
|
||||||
import mmap
|
|
||||||
import os
|
|
||||||
import struct
|
|
||||||
import sys
|
|
||||||
from typing import Any
|
|
||||||
|
|
||||||
HEADER_LEN = 40
|
|
||||||
worker_generation = 0
|
|
||||||
arena: mmap.mmap | None = None
|
|
||||||
rings: dict[int, dict[str, Any]] = {}
|
|
||||||
objects: dict[int, dict[str, Any]] = {}
|
|
||||||
role: dict[str, Any] = {}
|
|
||||||
Tensor: Any = None
|
|
||||||
dtypes: Any = None
|
|
||||||
next_handle = 42
|
|
||||||
|
|
||||||
|
|
||||||
def control(**event: Any) -> None:
|
|
||||||
print(json.dumps(event, separators=(",", ":")), flush=True)
|
|
||||||
|
|
||||||
|
|
||||||
def fatal(reason: str, **fields: Any) -> None:
|
|
||||||
control(type="WorkerFatal", reason=reason, **fields)
|
|
||||||
raise SystemExit(1)
|
|
||||||
|
|
||||||
|
|
||||||
def require_arena() -> mmap.mmap:
|
|
||||||
if arena is None:
|
|
||||||
fatal("ArenaNotMapped")
|
|
||||||
return arena
|
|
||||||
|
|
||||||
|
|
||||||
def require_tinygrad() -> tuple[Any, Any]:
|
|
||||||
if Tensor is None or dtypes is None:
|
|
||||||
fatal("BackendNotInitialized")
|
|
||||||
return Tensor, dtypes
|
|
||||||
|
|
||||||
|
|
||||||
def initialize(cmd: dict[str, Any]) -> None:
|
|
||||||
global arena, Tensor, dtypes, worker_generation
|
|
||||||
if int(cmd["required_ring_helper_abi"]) != 1:
|
|
||||||
fatal("UnsupportedHelperAbi", required_ring_helper_abi=cmd["required_ring_helper_abi"])
|
|
||||||
worker_generation = int(cmd["worker_generation"])
|
|
||||||
fd = int(os.environ["SWACTOR_ARENA_FD"])
|
|
||||||
size = int(cmd.get("arena_ceiling", os.environ["SWACTOR_ARENA_BYTES"]))
|
|
||||||
arena = mmap.mmap(fd, size)
|
|
||||||
os.environ.setdefault("DEV", cmd.get("backend", {}).get("device", "CPU"))
|
|
||||||
from tinygrad import Tensor as TinyTensor, dtypes as tiny_dtypes
|
|
||||||
|
|
||||||
Tensor = TinyTensor
|
|
||||||
dtypes = tiny_dtypes
|
|
||||||
Tensor([1], dtype=dtypes.int32).realize().numpy().tolist()
|
|
||||||
control(
|
|
||||||
type="WorkerReady",
|
|
||||||
pid=os.getpid(),
|
|
||||||
worker_generation=worker_generation,
|
|
||||||
ring_helper_abi=1,
|
|
||||||
backend={"device": os.environ.get("DEV", "CPU")},
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
def install_ring(cmd: dict[str, Any]) -> None:
|
|
||||||
ring_id = int(cmd["ring_id"])
|
|
||||||
layout = cmd["layout"]
|
|
||||||
spec = cmd["object_spec"]
|
|
||||||
rings[ring_id] = {
|
|
||||||
"ring_id": ring_id,
|
|
||||||
"edge_id": int(cmd["edge_id"]),
|
|
||||||
"port_id": cmd["port_id"],
|
|
||||||
"direction": cmd["direction"],
|
|
||||||
"data_offset": int(layout["data_offset"]),
|
|
||||||
"data_capacity": int(layout["data_capacity"]),
|
|
||||||
"max_extent": int(spec["max_extent"]),
|
|
||||||
"alignment": int(spec["alignment"]),
|
|
||||||
"next_sequence": 0,
|
|
||||||
}
|
|
||||||
control(type="RingInstalled", ring_id=ring_id, edge_id=rings[ring_id]["edge_id"], port_id=cmd["port_id"])
|
|
||||||
|
|
||||||
|
|
||||||
def configure_role(cmd: dict[str, Any]) -> None:
|
|
||||||
config = cmd["config"]
|
|
||||||
role.clear()
|
|
||||||
role.update(
|
|
||||||
role_id=int(cmd["role_id"]),
|
|
||||||
run_id=int(config["run_id"]),
|
|
||||||
stage_index=int(config["stage_index"]),
|
|
||||||
layer_start=int(config["layer_start"]),
|
|
||||||
layer_end_exclusive=int(config["layer_end_exclusive"]),
|
|
||||||
)
|
|
||||||
control(type="RoleConfigured", role_id=role["role_id"])
|
|
||||||
|
|
||||||
|
|
||||||
def load_weights(cmd: dict[str, Any]) -> None:
|
|
||||||
if not role:
|
|
||||||
fatal("RoleNotConfigured")
|
|
||||||
role["model_id"] = cmd["model_id"]
|
|
||||||
role["gguf_source"] = cmd["gguf_source"]
|
|
||||||
role["tokenizer"] = cmd["tokenizer"]
|
|
||||||
role["weight_layer_start"] = int(cmd["layer_start"])
|
|
||||||
role["weight_layer_end_exclusive"] = int(cmd["layer_end_exclusive"])
|
|
||||||
control(
|
|
||||||
type="WeightsLoaded",
|
|
||||||
model_id=role["model_id"],
|
|
||||||
layer_start=role["weight_layer_start"],
|
|
||||||
layer_end_exclusive=role["weight_layer_end_exclusive"],
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
def parse_record(ring: dict[str, Any]) -> tuple[int, int, int, bytes]:
|
|
||||||
view = require_arena()
|
|
||||||
base = ring["data_offset"]
|
|
||||||
header = view[base : base + HEADER_LEN]
|
|
||||||
version = struct.unpack_from("<H", header, 4)[0]
|
|
||||||
header_len = struct.unpack_from("<H", header, 6)[0]
|
|
||||||
if header[0:4] != b"MO01" or version != 1 or header_len != HEADER_LEN:
|
|
||||||
fatal("InvalidObjectHeader", ring_id=ring["ring_id"])
|
|
||||||
object_id = struct.unpack_from("<Q", header, 8)[0]
|
|
||||||
sequence = struct.unpack_from("<Q", header, 16)[0]
|
|
||||||
extent = struct.unpack_from("<Q", header, 24)[0]
|
|
||||||
flags = struct.unpack_from("<I", header, 32)[0]
|
|
||||||
reserved = struct.unpack_from("<I", header, 36)[0]
|
|
||||||
del flags, reserved
|
|
||||||
if extent > ring["max_extent"] or (ring["alignment"] and extent % ring["alignment"]):
|
|
||||||
fatal("ObjectExtentInvalid", ring_id=ring["ring_id"], object_id=object_id, extent=extent)
|
|
||||||
if sequence != ring["next_sequence"]:
|
|
||||||
fatal("SequenceViolation", ring_id=ring["ring_id"], expected=ring["next_sequence"], actual=sequence)
|
|
||||||
payload = bytes(view[base + HEADER_LEN : base + HEADER_LEN + extent])
|
|
||||||
ring["next_sequence"] += 1
|
|
||||||
return object_id, sequence, extent, payload
|
|
||||||
|
|
||||||
|
|
||||||
def ring_readable(cmd: dict[str, Any]) -> None:
|
|
||||||
global next_handle
|
|
||||||
tensor, dtype_mod = require_tinygrad()
|
|
||||||
ring_id = int(cmd["ring_id"])
|
|
||||||
ring = rings[ring_id]
|
|
||||||
if ring["direction"] != "ingress":
|
|
||||||
fatal("WrongRingDirection", ring_id=ring_id)
|
|
||||||
object_id, sequence, extent, payload = parse_record(ring)
|
|
||||||
values = list(struct.unpack(f"<{extent // 4}i", payload))
|
|
||||||
loaded = tensor(values, dtype=dtype_mod.int32).realize()
|
|
||||||
handle = next_handle
|
|
||||||
next_handle += 1
|
|
||||||
objects[handle] = {
|
|
||||||
"object_id": object_id,
|
|
||||||
"sequence": sequence,
|
|
||||||
"tensor": loaded,
|
|
||||||
"extent": extent,
|
|
||||||
}
|
|
||||||
control(
|
|
||||||
type="ObjectLoaded",
|
|
||||||
ring_id=ring_id,
|
|
||||||
edge_id=ring["edge_id"],
|
|
||||||
port_id=ring["port_id"],
|
|
||||||
object_id=object_id,
|
|
||||||
sequence=sequence,
|
|
||||||
extent=extent,
|
|
||||||
device_handle={"worker_generation": worker_generation, "id": handle},
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
def write_record(ring: dict[str, Any], object_id: int, sequence: int, words: list[int], flags: int) -> int:
|
|
||||||
payload = b"".join(struct.pack("<i", word) for word in words)
|
|
||||||
extent = len(payload)
|
|
||||||
if extent > ring["max_extent"]:
|
|
||||||
fatal("OutputExtentInvalid", ring_id=ring["ring_id"], extent=extent)
|
|
||||||
header = bytearray(HEADER_LEN)
|
|
||||||
header[0:4] = b"MO01"
|
|
||||||
struct.pack_into("<H", header, 4, 1)
|
|
||||||
struct.pack_into("<H", header, 6, HEADER_LEN)
|
|
||||||
struct.pack_into("<Q", header, 8, object_id)
|
|
||||||
struct.pack_into("<Q", header, 16, sequence)
|
|
||||||
struct.pack_into("<Q", header, 24, extent)
|
|
||||||
struct.pack_into("<I", header, 32, flags)
|
|
||||||
struct.pack_into("<I", header, 36, 0)
|
|
||||||
view = require_arena()
|
|
||||||
base = ring["data_offset"]
|
|
||||||
view[base : base + HEADER_LEN] = header
|
|
||||||
view[base + HEADER_LEN : base + HEADER_LEN + extent] = payload
|
|
||||||
return HEADER_LEN + extent
|
|
||||||
|
|
||||||
|
|
||||||
def execute_step(cmd: dict[str, Any]) -> None:
|
|
||||||
if not role:
|
|
||||||
fatal("RoleNotConfigured")
|
|
||||||
tensor, _ = require_tinygrad()
|
|
||||||
del tensor
|
|
||||||
role_id = int(cmd["role_id"])
|
|
||||||
if role_id != role["role_id"]:
|
|
||||||
fatal("RoleMismatch", expected=role["role_id"], actual=role_id)
|
|
||||||
step_id = int(cmd["step_id"])
|
|
||||||
input_binding = cmd["inputs"][0]
|
|
||||||
output_binding = cmd["outputs"][0]
|
|
||||||
device_handle = input_binding["device_handle"]
|
|
||||||
if int(device_handle["worker_generation"]) != worker_generation:
|
|
||||||
fatal("OldGenerationHandle", handle=device_handle)
|
|
||||||
handle = int(device_handle["id"])
|
|
||||||
obj = objects[handle]
|
|
||||||
if int(input_binding["object_id"]) != obj["object_id"] or int(input_binding["sequence"]) != obj["sequence"]:
|
|
||||||
fatal("InputBindingMismatch", step_id=step_id)
|
|
||||||
transformed = int(obj["tensor"].sum().item()) + role["layer_start"] + role["layer_end_exclusive"] + role["stage_index"]
|
|
||||||
if bool(cmd.get("runtime", {}).get("final_stage")):
|
|
||||||
words = [6 if transformed % 2 == 1 else 8]
|
|
||||||
else:
|
|
||||||
words = [transformed if transformed > 0 else 1]
|
|
||||||
ring = rings[int(output_binding["ring_id"])]
|
|
||||||
if ring["direction"] != "egress":
|
|
||||||
fatal("WrongRingDirection", ring_id=ring["ring_id"])
|
|
||||||
committed = write_record(
|
|
||||||
ring,
|
|
||||||
int(output_binding["object_id"]),
|
|
||||||
int(output_binding["sequence"]),
|
|
||||||
words,
|
|
||||||
int(output_binding.get("flags", 0)),
|
|
||||||
)
|
|
||||||
control(
|
|
||||||
type="ObjectProduced",
|
|
||||||
ring_id=ring["ring_id"],
|
|
||||||
edge_id=ring["edge_id"],
|
|
||||||
port_id=ring["port_id"],
|
|
||||||
object_id=int(output_binding["object_id"]),
|
|
||||||
sequence=int(output_binding["sequence"]),
|
|
||||||
committed_bytes=committed,
|
|
||||||
)
|
|
||||||
if bool(cmd.get("release_inputs_after")):
|
|
||||||
objects.pop(handle, None)
|
|
||||||
control(type="StepCompleted", role_id=role_id, step_id=step_id)
|
|
||||||
|
|
||||||
|
|
||||||
def release_device_object(cmd: dict[str, Any]) -> None:
|
|
||||||
handle = int(cmd["device_handle"]["id"])
|
|
||||||
objects.pop(handle, None)
|
|
||||||
control(type="DeviceObjectReleased", device_handle=cmd["device_handle"])
|
|
||||||
|
|
||||||
|
|
||||||
def uninstall_ring(cmd: dict[str, Any]) -> None:
|
|
||||||
ring_id = int(cmd["ring_id"])
|
|
||||||
rings.pop(ring_id, None)
|
|
||||||
control(type="RingQuiesced", ring_id=ring_id)
|
|
||||||
|
|
||||||
|
|
||||||
def shutdown_worker(_: dict[str, Any]) -> None:
|
|
||||||
control(type="WorkerStopped", reason="Graceful")
|
|
||||||
raise SystemExit(0)
|
|
||||||
|
|
||||||
|
|
||||||
HANDLERS = {
|
|
||||||
"InitializeWorker": initialize,
|
|
||||||
"InstallRing": install_ring,
|
|
||||||
"ConfigureRole": configure_role,
|
|
||||||
"LoadWeights": load_weights,
|
|
||||||
"RingReadable": ring_readable,
|
|
||||||
"ExecuteStep": execute_step,
|
|
||||||
"ReleaseDeviceObject": release_device_object,
|
|
||||||
"UninstallRing": uninstall_ring,
|
|
||||||
"ShutdownWorker": shutdown_worker,
|
|
||||||
}
|
|
||||||
|
|
||||||
for raw in sys.stdin:
|
|
||||||
if not raw.strip():
|
|
||||||
continue
|
|
||||||
try:
|
|
||||||
command = json.loads(raw)
|
|
||||||
except json.JSONDecodeError as exc:
|
|
||||||
fatal("InvalidJson", error=str(exc))
|
|
||||||
handler = HANDLERS.get(command.get("type"))
|
|
||||||
if handler is None:
|
|
||||||
fatal("UnknownCommand", command=command.get("type"))
|
|
||||||
handler(command)
|
|
||||||
|
|
@ -1,119 +0,0 @@
|
||||||
use std::io::Write;
|
|
||||||
use std::path::PathBuf;
|
|
||||||
use std::process::{Command, ExitCode, Stdio};
|
|
||||||
|
|
||||||
#[path = "support/dumb_worker.rs"]
|
|
||||||
mod dumb_worker;
|
|
||||||
#[path = "support/local_e2e.rs"]
|
|
||||||
mod local_e2e;
|
|
||||||
|
|
||||||
fn main() -> ExitCode {
|
|
||||||
let args = std::env::args().collect::<Vec<_>>();
|
|
||||||
match std::env::var("MVP_TEST_ROLE").ok().as_deref() {
|
|
||||||
Some("dumb-worker") => return dumb_worker::run_main(),
|
|
||||||
Some("local-e2e") => return local_e2e::run_main(),
|
|
||||||
Some(role) => {
|
|
||||||
eprintln!("unknown MVP_TEST_ROLE={role}");
|
|
||||||
return ExitCode::from(2);
|
|
||||||
}
|
|
||||||
None => {}
|
|
||||||
}
|
|
||||||
|
|
||||||
if args.iter().any(|arg| arg == "--role=node") {
|
|
||||||
return local_e2e::run_main();
|
|
||||||
}
|
|
||||||
|
|
||||||
local_e2e_binary_drives_real_local_process_deployment();
|
|
||||||
dumb_worker_is_a_real_child_process_protocol_endpoint();
|
|
||||||
ExitCode::SUCCESS
|
|
||||||
}
|
|
||||||
|
|
||||||
fn local_e2e_binary_drives_real_local_process_deployment() {
|
|
||||||
let output = Command::new(current_test_exe())
|
|
||||||
.env("MVP_TEST_ROLE", "local-e2e")
|
|
||||||
.output()
|
|
||||||
.expect("run local e2e supervisor");
|
|
||||||
|
|
||||||
assert!(
|
|
||||||
output.status.success(),
|
|
||||||
"mvp-local-e2e failed\nstdout:\n{}\nstderr:\n{}",
|
|
||||||
String::from_utf8_lossy(&output.stdout),
|
|
||||||
String::from_utf8_lossy(&output.stderr)
|
|
||||||
);
|
|
||||||
|
|
||||||
let value: serde_json::Value = serde_json::from_slice(&output.stdout).expect("json stdout");
|
|
||||||
assert_eq!(value["ok"], true);
|
|
||||||
assert_eq!(value["actor_plane"], "iroh-swactor");
|
|
||||||
assert_eq!(value["data_plane"], "tcp-loopback-streams");
|
|
||||||
assert_eq!(value["worker_processes"], "mvp-dumb-worker-per-node");
|
|
||||||
assert_eq!(
|
|
||||||
value["engine_builder_pattern"],
|
|
||||||
"pool-first-static-launcher"
|
|
||||||
);
|
|
||||||
assert_eq!(value["engine_builder_node_count"], 3);
|
|
||||||
assert_eq!(value["engine_builder_stage_assignments"], 2);
|
|
||||||
assert!(
|
|
||||||
value["engine_builder_event_count"]
|
|
||||||
.as_u64()
|
|
||||||
.is_some_and(|count| count >= 10),
|
|
||||||
"{value}"
|
|
||||||
);
|
|
||||||
assert_eq!(value["injected_prompt_observed"], true);
|
|
||||||
assert_eq!(value["token_received_observed"], true);
|
|
||||||
assert_eq!(value["run_completed_observed"], true);
|
|
||||||
assert_eq!(value["run_torn_down_observed"], true);
|
|
||||||
assert_eq!(value["stop_sent_to_all_nodes"], true);
|
|
||||||
assert_eq!(value["stage_ready_stdout_count"], 2);
|
|
||||||
assert!(value["processes"]["node0"].as_u64().is_some(), "{value}");
|
|
||||||
assert!(value["processes"]["node1"].as_u64().is_some(), "{value}");
|
|
||||||
assert!(
|
|
||||||
value["node0_endpoint"]["addrs"]
|
|
||||||
.as_array()
|
|
||||||
.is_some_and(|addrs| !addrs.is_empty()),
|
|
||||||
"{value}"
|
|
||||||
);
|
|
||||||
assert!(
|
|
||||||
value["node1_endpoint"]["addrs"]
|
|
||||||
.as_array()
|
|
||||||
.is_some_and(|addrs| !addrs.is_empty()),
|
|
||||||
"{value}"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
fn dumb_worker_is_a_real_child_process_protocol_endpoint() {
|
|
||||||
let mut child = Command::new(current_test_exe())
|
|
||||||
.env("MVP_TEST_ROLE", "dumb-worker")
|
|
||||||
.stdin(Stdio::piped())
|
|
||||||
.stdout(Stdio::piped())
|
|
||||||
.stderr(Stdio::piped())
|
|
||||||
.spawn()
|
|
||||||
.expect("spawn dumb worker test role");
|
|
||||||
|
|
||||||
{
|
|
||||||
let stdin = child.stdin.as_mut().expect("worker stdin");
|
|
||||||
writeln!(
|
|
||||||
stdin,
|
|
||||||
"{{\"type\":\"InitializeWorker\",\"helper_abi_version\":1}}"
|
|
||||||
)
|
|
||||||
.expect("write initialize");
|
|
||||||
writeln!(stdin, "{{\"type\":\"ExecuteStep\",\"step_id\":7}}").expect("write execute");
|
|
||||||
writeln!(stdin, "{{\"type\":\"ShutdownWorker\"}}").expect("write shutdown");
|
|
||||||
}
|
|
||||||
|
|
||||||
let output = child.wait_with_output().expect("worker output");
|
|
||||||
assert!(
|
|
||||||
output.status.success(),
|
|
||||||
"worker failed\nstdout:\n{}\nstderr:\n{}",
|
|
||||||
String::from_utf8_lossy(&output.stdout),
|
|
||||||
String::from_utf8_lossy(&output.stderr)
|
|
||||||
);
|
|
||||||
let stdout = String::from_utf8(output.stdout).expect("utf8 stdout");
|
|
||||||
assert!(stdout.contains("\"type\":\"WorkerReady\""), "{stdout}");
|
|
||||||
assert!(stdout.contains("\"type\":\"StepCompleted\""), "{stdout}");
|
|
||||||
assert!(stdout.contains("\"step_id\":7"), "{stdout}");
|
|
||||||
assert!(stdout.contains("\"type\":\"WorkerStopped\""), "{stdout}");
|
|
||||||
}
|
|
||||||
|
|
||||||
fn current_test_exe() -> PathBuf {
|
|
||||||
std::env::current_exe().expect("current test exe")
|
|
||||||
}
|
|
||||||
|
|
@ -1,594 +0,0 @@
|
||||||
use std::io::{Read, Write};
|
|
||||||
use std::net::TcpStream;
|
|
||||||
use std::process::{Child, Command, Stdio};
|
|
||||||
use std::sync::{Arc, Mutex};
|
|
||||||
use std::thread;
|
|
||||||
use std::time::{Duration, Instant};
|
|
||||||
|
|
||||||
use serde_json::Value;
|
|
||||||
|
|
||||||
#[cfg(target_os = "linux")]
|
|
||||||
use std::os::unix::process::CommandExt;
|
|
||||||
|
|
||||||
const TEST_WATCHDOG: Duration = Duration::from_secs(1_800);
|
|
||||||
const PROMPT_WATCHDOG: Duration = Duration::from_secs(600);
|
|
||||||
const SHUTDOWN_WATCHDOG: Duration = Duration::from_secs(60);
|
|
||||||
const DASHBOARD_ADDR: &str = "127.0.0.1:9090";
|
|
||||||
const DOCKER_CONTAINER_PREFIX_ENV: &str = "MVP_DOCKER_CONTAINER_PREFIX";
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn one_node_chat_docker_cuda_e2e() {
|
|
||||||
let root = workspace_root();
|
|
||||||
require_docker(&root);
|
|
||||||
let container_prefix = format!("mvp-orchestrator-e2e-{}", std::process::id());
|
|
||||||
|
|
||||||
let mut command = Command::new("cargo");
|
|
||||||
command
|
|
||||||
.current_dir(&root)
|
|
||||||
.args(["mvp-chat", "--docker"])
|
|
||||||
.env("MVP_RUNTIME_CONFIG", "local")
|
|
||||||
.env("MVP_IROH_RELAY_MODE", "disabled")
|
|
||||||
.env(DOCKER_CONTAINER_PREFIX_ENV, &container_prefix)
|
|
||||||
.env("MVP_TINYGRAD_TEST_MODE", "1")
|
|
||||||
.env("MVP_LAYER_END_EXCLUSIVE", "1")
|
|
||||||
.env("MVP_PROMPT_MAX_TOKENS", "3")
|
|
||||||
.stdin(Stdio::piped())
|
|
||||||
.stdout(Stdio::piped())
|
|
||||||
.stderr(Stdio::piped());
|
|
||||||
#[cfg(target_os = "linux")]
|
|
||||||
unsafe {
|
|
||||||
command.pre_exec(|| {
|
|
||||||
if libc::setpgid(0, 0) == 0 {
|
|
||||||
Ok(())
|
|
||||||
} else {
|
|
||||||
Err(std::io::Error::last_os_error())
|
|
||||||
}
|
|
||||||
});
|
|
||||||
}
|
|
||||||
|
|
||||||
let mut child = command.spawn().expect("spawn cargo mvp-chat");
|
|
||||||
let mut stdin = child.stdin.take().expect("cargo mvp-chat stdin");
|
|
||||||
let stdout = Arc::new(Mutex::new(String::new()));
|
|
||||||
let stderr = Arc::new(Mutex::new(String::new()));
|
|
||||||
let stdout_reader = spawn_capture(child.stdout.take().expect("stdout"), Arc::clone(&stdout));
|
|
||||||
let stderr_reader = spawn_capture(child.stderr.take().expect("stderr"), Arc::clone(&stderr));
|
|
||||||
|
|
||||||
let mut result = run_full_flow(
|
|
||||||
&mut child,
|
|
||||||
&mut stdin,
|
|
||||||
&stdout,
|
|
||||||
"mvp.provisioning.logs",
|
|
||||||
"mvp-entrypoint",
|
|
||||||
dashboard_has_worker_prompt_completed,
|
|
||||||
);
|
|
||||||
if result.is_err() {
|
|
||||||
request_child_interrupt(&child);
|
|
||||||
let _ = wait_child(&mut child, SHUTDOWN_WATCHDOG);
|
|
||||||
let _ = child.kill();
|
|
||||||
let _ = child.wait();
|
|
||||||
}
|
|
||||||
let _ = stdout_reader.join();
|
|
||||||
let _ = stderr_reader.join();
|
|
||||||
if result.is_ok() {
|
|
||||||
result = assert_no_lower_layer_terminal_leaks(&stdout, &stderr);
|
|
||||||
}
|
|
||||||
assert_containers_with_prefix_removed(&root, &container_prefix);
|
|
||||||
|
|
||||||
if let Err(error) = result {
|
|
||||||
panic!(
|
|
||||||
"{error}\nstdout:\n{}\nstderr:\n{}\ndashboard frames:\n{}",
|
|
||||||
snapshot(&stdout),
|
|
||||||
snapshot(&stderr),
|
|
||||||
dashboard_snapshot().unwrap_or_else(|err| format!("<dashboard unavailable: {err}>"))
|
|
||||||
);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn one_node_chat_process_cached_model_e2e() {
|
|
||||||
let root = workspace_root();
|
|
||||||
let cached_model = root
|
|
||||||
.join(".model-cache")
|
|
||||||
.join("SmolLM2-135M-Instruct.Q4_0.gguf");
|
|
||||||
assert!(
|
|
||||||
cached_model.is_file(),
|
|
||||||
"cached model fixture is required: {}",
|
|
||||||
cached_model.display()
|
|
||||||
);
|
|
||||||
let container_prefix = format!("mvp-orchestrator-process-e2e-{}", std::process::id());
|
|
||||||
|
|
||||||
let mut command = Command::new("cargo");
|
|
||||||
command
|
|
||||||
.current_dir(&root)
|
|
||||||
.args(["mvp-chat", "--cached-model"])
|
|
||||||
.args(["--pipeline-stages", "2", "--dump-logs"])
|
|
||||||
.env("MVP_RUNTIME_CONFIG", "local")
|
|
||||||
.env("MVP_IROH_RELAY_MODE", "disabled")
|
|
||||||
.env(DOCKER_CONTAINER_PREFIX_ENV, &container_prefix)
|
|
||||||
.env("MVP_TINYGRAD_TEST_MODE", "1")
|
|
||||||
.env("MVP_LAYER_END_EXCLUSIVE", "1")
|
|
||||||
.env("MVP_PROMPT_MAX_TOKENS", "3")
|
|
||||||
.stdin(Stdio::piped())
|
|
||||||
.stdout(Stdio::piped())
|
|
||||||
.stderr(Stdio::piped());
|
|
||||||
#[cfg(target_os = "linux")]
|
|
||||||
unsafe {
|
|
||||||
command.pre_exec(|| {
|
|
||||||
if libc::setpgid(0, 0) == 0 {
|
|
||||||
Ok(())
|
|
||||||
} else {
|
|
||||||
Err(std::io::Error::last_os_error())
|
|
||||||
}
|
|
||||||
});
|
|
||||||
}
|
|
||||||
|
|
||||||
let mut child = command.spawn().expect("spawn cargo mvp-chat");
|
|
||||||
let mut stdin = child.stdin.take().expect("cargo mvp-chat stdin");
|
|
||||||
let stdout = Arc::new(Mutex::new(String::new()));
|
|
||||||
let stderr = Arc::new(Mutex::new(String::new()));
|
|
||||||
let stdout_reader = spawn_capture(child.stdout.take().expect("stdout"), Arc::clone(&stdout));
|
|
||||||
let stderr_reader = spawn_capture(child.stderr.take().expect("stderr"), Arc::clone(&stderr));
|
|
||||||
|
|
||||||
let mut result = run_full_flow(
|
|
||||||
&mut child,
|
|
||||||
&mut stdin,
|
|
||||||
&stdout,
|
|
||||||
"mvp.node.bootstrap",
|
|
||||||
"runtime_ready_local",
|
|
||||||
dashboard_has_pipeline_prompt_output,
|
|
||||||
);
|
|
||||||
if result.is_err() {
|
|
||||||
request_child_interrupt(&child);
|
|
||||||
let _ = wait_child(&mut child, SHUTDOWN_WATCHDOG);
|
|
||||||
let _ = child.kill();
|
|
||||||
let _ = child.wait();
|
|
||||||
}
|
|
||||||
let _ = stdout_reader.join();
|
|
||||||
let _ = stderr_reader.join();
|
|
||||||
if result.is_ok() {
|
|
||||||
result = assert_no_lower_layer_terminal_leaks(&stdout, &stderr);
|
|
||||||
}
|
|
||||||
assert_no_containers_with_prefix_if_docker_available(&root, &container_prefix);
|
|
||||||
|
|
||||||
if let Err(error) = result {
|
|
||||||
panic!(
|
|
||||||
"{error}\nstdout:\n{}\nstderr:\n{}\ndashboard frames:\n{}",
|
|
||||||
snapshot(&stdout),
|
|
||||||
snapshot(&stderr),
|
|
||||||
dashboard_snapshot().unwrap_or_else(|err| format!("<dashboard unavailable: {err}>"))
|
|
||||||
);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
fn run_full_flow(
|
|
||||||
child: &mut Child,
|
|
||||||
stdin: &mut impl Write,
|
|
||||||
stdout: &Arc<Mutex<String>>,
|
|
||||||
provisioning_channel_substr: &str,
|
|
||||||
provisioning_payload_substr: &str,
|
|
||||||
prompt_dashboard: fn() -> bool,
|
|
||||||
) -> Result<(), String> {
|
|
||||||
wait_for_child_or(TEST_WATCHDOG, child, dashboard_responding)
|
|
||||||
.map_err(|e| format!("dashboard API not live: {e}"))?;
|
|
||||||
wait_for_child_or(TEST_WATCHDOG, child, || {
|
|
||||||
dashboard_has_channel_or_payload(provisioning_channel_substr, provisioning_payload_substr)
|
|
||||||
})
|
|
||||||
.map_err(|e| format!("provisioning frames not visible in dashboard: {e}"))?;
|
|
||||||
wait_for_child_or(TEST_WATCHDOG, child, || {
|
|
||||||
dashboard_has_frame("mvp.worker.weights", "LoadWeightsStarted")
|
|
||||||
})
|
|
||||||
.map_err(|e| format!("weight loading start not visible in dashboard: {e}"))?;
|
|
||||||
wait_for_child_or(TEST_WATCHDOG, child, || {
|
|
||||||
dashboard_has_frame("mvp.worker.weights", "WeightsLoaded")
|
|
||||||
})
|
|
||||||
.map_err(|e| format!("WeightsLoaded not visible in dashboard: {e}"))?;
|
|
||||||
wait_for_child_or(TEST_WATCHDOG, child, || prompt_visible(stdout))
|
|
||||||
.map_err(|e| format!("chat prompt not visible: {e}"))?;
|
|
||||||
|
|
||||||
writeln!(stdin, "hello from full cargo mvp-chat e2e")
|
|
||||||
.map_err(|e| format!("write prompt: {e}"))?;
|
|
||||||
stdin.flush().map_err(|e| format!("flush prompt: {e}"))?;
|
|
||||||
wait_for_child_or(PROMPT_WATCHDOG, child, || {
|
|
||||||
stdout_contains(stdout, "decoding...")
|
|
||||||
})
|
|
||||||
.map_err(|e| format!("prompt was not submitted to chat loop: {e}"))?;
|
|
||||||
wait_for_child_or(PROMPT_WATCHDOG, child, || response_text_visible(stdout))
|
|
||||||
.map_err(|e| format!("decoded response text not visible: {e}"))?;
|
|
||||||
wait_for_child_or(PROMPT_WATCHDOG, child, prompt_dashboard)
|
|
||||||
.map_err(|e| format!("prompt result not visible in dashboard: {e}"))?;
|
|
||||||
wait_for_child_or(PROMPT_WATCHDOG, child, dashboard_has_orch_prompt_lifecycle)
|
|
||||||
.map_err(|e| format!("orchestrator prompt lifecycle not visible in dashboard: {e}"))?;
|
|
||||||
wait_for_child_or(PROMPT_WATCHDOG, child, || prompt_count(stdout) >= 2)
|
|
||||||
.map_err(|e| format!("chat prompt did not return after response: {e}"))?;
|
|
||||||
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)
|
|
||||||
.ok_or_else(|| "cargo mvp-chat did not exit after /exit".to_owned())?;
|
|
||||||
if status.success() {
|
|
||||||
Ok(())
|
|
||||||
} else {
|
|
||||||
Err(format!("cargo mvp-chat exited with {status}"))
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
fn wait_for_child_or(
|
|
||||||
watchdog: Duration,
|
|
||||||
child: &mut Child,
|
|
||||||
mut predicate: impl FnMut() -> bool,
|
|
||||||
) -> Result<(), String> {
|
|
||||||
let start = Instant::now();
|
|
||||||
while start.elapsed() < watchdog {
|
|
||||||
if predicate() {
|
|
||||||
return Ok(());
|
|
||||||
}
|
|
||||||
if let Some(status) = child.try_wait().map_err(|e| format!("poll child: {e}"))? {
|
|
||||||
return Err(format!("child exited early with {status}"));
|
|
||||||
}
|
|
||||||
thread::sleep(Duration::from_millis(250));
|
|
||||||
}
|
|
||||||
Err("test watchdog".to_owned())
|
|
||||||
}
|
|
||||||
|
|
||||||
fn spawn_capture(
|
|
||||||
mut reader: impl Read + Send + 'static,
|
|
||||||
out: Arc<Mutex<String>>,
|
|
||||||
) -> thread::JoinHandle<()> {
|
|
||||||
thread::spawn(move || {
|
|
||||||
let mut buf = [0_u8; 8192];
|
|
||||||
loop {
|
|
||||||
match reader.read(&mut buf) {
|
|
||||||
Ok(0) => break,
|
|
||||||
Ok(n) => out
|
|
||||||
.lock()
|
|
||||||
.expect("capture mutex")
|
|
||||||
.push_str(&String::from_utf8_lossy(&buf[..n])),
|
|
||||||
Err(_) => break,
|
|
||||||
}
|
|
||||||
}
|
|
||||||
})
|
|
||||||
}
|
|
||||||
|
|
||||||
fn dashboard_responding() -> bool {
|
|
||||||
dashboard_frames().is_ok()
|
|
||||||
}
|
|
||||||
|
|
||||||
fn dashboard_has_channel_or_payload(channel_substr: &str, payload_substr: &str) -> bool {
|
|
||||||
dashboard_frames()
|
|
||||||
.map(|frames| {
|
|
||||||
frames.iter().any(|frame| {
|
|
||||||
frame.channel.contains(channel_substr) || frame.payload.contains(payload_substr)
|
|
||||||
})
|
|
||||||
})
|
|
||||||
.unwrap_or(false)
|
|
||||||
}
|
|
||||||
|
|
||||||
fn dashboard_has_frame(channel_substr: &str, payload_substr: &str) -> bool {
|
|
||||||
dashboard_frames()
|
|
||||||
.map(|frames| {
|
|
||||||
frames.iter().any(|frame| {
|
|
||||||
frame.channel.contains(channel_substr) && frame.payload.contains(payload_substr)
|
|
||||||
})
|
|
||||||
})
|
|
||||||
.unwrap_or(false)
|
|
||||||
}
|
|
||||||
|
|
||||||
fn dashboard_has_worker_prompt_completed() -> bool {
|
|
||||||
dashboard_has_frame("mvp.worker.prompt", "PromptCompleted")
|
|
||||||
}
|
|
||||||
|
|
||||||
fn dashboard_has_pipeline_prompt_output() -> bool {
|
|
||||||
dashboard_has_frame("mvp.orch.prompt", "pipeline_token_out")
|
|
||||||
}
|
|
||||||
|
|
||||||
fn dashboard_has_orch_prompt_lifecycle() -> bool {
|
|
||||||
dashboard_frames()
|
|
||||||
.map(|frames| {
|
|
||||||
let events = frames
|
|
||||||
.iter()
|
|
||||||
.filter_map(orch_prompt_observation)
|
|
||||||
.collect::<Vec<_>>();
|
|
||||||
|
|
||||||
events.iter().any(|prompt_work| {
|
|
||||||
prompt_work.phase == "prompt_work"
|
|
||||||
&& prompt_work.status == "observed"
|
|
||||||
&& (direct_orch_prompt_lifecycle(&events, prompt_work)
|
|
||||||
|| pipeline_orch_prompt_lifecycle(&events, prompt_work))
|
|
||||||
})
|
|
||||||
})
|
|
||||||
.unwrap_or(false)
|
|
||||||
}
|
|
||||||
|
|
||||||
fn direct_orch_prompt_lifecycle(
|
|
||||||
events: &[OrchPromptObservation],
|
|
||||||
prompt_work: &OrchPromptObservation,
|
|
||||||
) -> bool {
|
|
||||||
events.iter().any(|event| {
|
|
||||||
same_prompt(prompt_work, event)
|
|
||||||
&& event.phase == "node_prompt_send"
|
|
||||||
&& event.status == "ready"
|
|
||||||
}) && events.iter().any(|event| {
|
|
||||||
same_prompt(prompt_work, event)
|
|
||||||
&& event.phase == "node_prompt_event"
|
|
||||||
&& event.status == "observed"
|
|
||||||
&& event.detail_event.as_deref() == Some("Done")
|
|
||||||
}) && events.iter().any(|event| {
|
|
||||||
same_prompt(prompt_work, event)
|
|
||||||
&& event.phase == "prompt_complete"
|
|
||||||
&& event.status == "ready"
|
|
||||||
&& event.detail_event.as_deref() == Some("Done")
|
|
||||||
})
|
|
||||||
}
|
|
||||||
|
|
||||||
fn pipeline_orch_prompt_lifecycle(
|
|
||||||
events: &[OrchPromptObservation],
|
|
||||||
prompt_work: &OrchPromptObservation,
|
|
||||||
) -> bool {
|
|
||||||
events.iter().any(|event| {
|
|
||||||
same_prompt(prompt_work, event)
|
|
||||||
&& event.phase == "pipeline_tokenizer_encode"
|
|
||||||
&& event.status == "ready"
|
|
||||||
}) && events.iter().any(|event| {
|
|
||||||
same_prompt(prompt_work, event)
|
|
||||||
&& event.phase == "pipeline_token_in"
|
|
||||||
&& event.status == "ready"
|
|
||||||
}) && events.iter().any(|event| {
|
|
||||||
same_prompt(prompt_work, event)
|
|
||||||
&& event.phase == "pipeline_token_out"
|
|
||||||
&& event.status == "observed"
|
|
||||||
}) && events.iter().any(|event| {
|
|
||||||
same_prompt(prompt_work, event)
|
|
||||||
&& event.phase == "pipeline_tokenizer_decode"
|
|
||||||
&& event.status == "ready"
|
|
||||||
})
|
|
||||||
}
|
|
||||||
|
|
||||||
fn orch_prompt_observation(frame: &SeenFrame) -> Option<OrchPromptObservation> {
|
|
||||||
if frame.channel != "mvp.orch.prompt" {
|
|
||||||
return None;
|
|
||||||
}
|
|
||||||
|
|
||||||
let value = serde_json::from_str::<Value>(&frame.payload).ok()?;
|
|
||||||
if value.get("type").and_then(Value::as_str)? != "OrchPromptEvent" {
|
|
||||||
return None;
|
|
||||||
}
|
|
||||||
let detail = value.get("detail")?;
|
|
||||||
|
|
||||||
Some(OrchPromptObservation {
|
|
||||||
phase: value.get("phase").and_then(Value::as_str)?.to_owned(),
|
|
||||||
status: value.get("status").and_then(Value::as_str)?.to_owned(),
|
|
||||||
run_id: value.get("run_id").and_then(Value::as_u64)?,
|
|
||||||
node_id: value.get("node_id").and_then(Value::as_u64)?,
|
|
||||||
request_id: value.get("request_id").and_then(Value::as_u64)?,
|
|
||||||
detail_event: detail
|
|
||||||
.get("event")
|
|
||||||
.and_then(Value::as_str)
|
|
||||||
.map(str::to_owned),
|
|
||||||
})
|
|
||||||
}
|
|
||||||
|
|
||||||
fn same_prompt(left: &OrchPromptObservation, right: &OrchPromptObservation) -> bool {
|
|
||||||
left.run_id == right.run_id
|
|
||||||
&& left.node_id == right.node_id
|
|
||||||
&& left.request_id == right.request_id
|
|
||||||
}
|
|
||||||
|
|
||||||
#[derive(Debug)]
|
|
||||||
struct OrchPromptObservation {
|
|
||||||
phase: String,
|
|
||||||
status: String,
|
|
||||||
run_id: u64,
|
|
||||||
node_id: u64,
|
|
||||||
request_id: u64,
|
|
||||||
detail_event: Option<String>,
|
|
||||||
}
|
|
||||||
|
|
||||||
#[derive(Debug)]
|
|
||||||
struct SeenFrame {
|
|
||||||
channel: String,
|
|
||||||
payload: String,
|
|
||||||
}
|
|
||||||
|
|
||||||
fn dashboard_frames() -> Result<Vec<SeenFrame>, String> {
|
|
||||||
let response = http_get("/api/frames")?;
|
|
||||||
let (_, body) = response
|
|
||||||
.split_once("\r\n\r\n")
|
|
||||||
.ok_or_else(|| "HTTP response missing body".to_owned())?;
|
|
||||||
let values = serde_json::from_str::<Vec<Value>>(body)
|
|
||||||
.map_err(|e| format!("parse dashboard frames JSON: {e}; body={body:?}"))?;
|
|
||||||
Ok(values
|
|
||||||
.into_iter()
|
|
||||||
.map(|value| SeenFrame {
|
|
||||||
channel: value
|
|
||||||
.get("channel")
|
|
||||||
.and_then(Value::as_str)
|
|
||||||
.unwrap_or_default()
|
|
||||||
.to_owned(),
|
|
||||||
payload: decode_payload(value.get("payload")).unwrap_or_default(),
|
|
||||||
})
|
|
||||||
.collect())
|
|
||||||
}
|
|
||||||
|
|
||||||
fn dashboard_snapshot() -> Result<String, String> {
|
|
||||||
let mut frames = dashboard_frames()?;
|
|
||||||
let keep = frames.len().saturating_sub(40);
|
|
||||||
frames.drain(0..keep);
|
|
||||||
Ok(frames
|
|
||||||
.into_iter()
|
|
||||||
.map(|frame| format!("{} {}", frame.channel, frame.payload))
|
|
||||||
.collect::<Vec<_>>()
|
|
||||||
.join("\n"))
|
|
||||||
}
|
|
||||||
|
|
||||||
fn decode_payload(value: Option<&Value>) -> Option<String> {
|
|
||||||
let bytes = value?
|
|
||||||
.as_array()?
|
|
||||||
.iter()
|
|
||||||
.map(|byte| byte.as_u64().map(|n| n as u8))
|
|
||||||
.collect::<Option<Vec<_>>>()?;
|
|
||||||
Some(String::from_utf8_lossy(&bytes).to_string())
|
|
||||||
}
|
|
||||||
|
|
||||||
fn http_get(path: &str) -> Result<String, String> {
|
|
||||||
let mut stream = TcpStream::connect(DASHBOARD_ADDR)
|
|
||||||
.map_err(|e| format!("connect dashboard {DASHBOARD_ADDR}: {e}"))?;
|
|
||||||
stream
|
|
||||||
.set_read_timeout(Some(Duration::from_secs(2)))
|
|
||||||
.map_err(|e| format!("set read watchdog: {e}"))?;
|
|
||||||
write!(
|
|
||||||
stream,
|
|
||||||
"GET {path} HTTP/1.1\r\nHost: localhost\r\nConnection: close\r\n\r\n"
|
|
||||||
)
|
|
||||||
.map_err(|e| format!("write HTTP request: {e}"))?;
|
|
||||||
let mut response = String::new();
|
|
||||||
stream
|
|
||||||
.read_to_string(&mut response)
|
|
||||||
.map_err(|e| format!("read HTTP response: {e}"))?;
|
|
||||||
if response.starts_with("HTTP/1.1 200") {
|
|
||||||
Ok(response)
|
|
||||||
} else {
|
|
||||||
Err(format!("non-200 dashboard response: {response:?}"))
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
fn stdout_contains(stdout: &Arc<Mutex<String>>, needle: &str) -> bool {
|
|
||||||
snapshot(stdout).contains(needle)
|
|
||||||
}
|
|
||||||
|
|
||||||
fn prompt_visible(stdout: &Arc<Mutex<String>>) -> bool {
|
|
||||||
snapshot(stdout).contains("prompt:> ")
|
|
||||||
}
|
|
||||||
|
|
||||||
fn prompt_count(stdout: &Arc<Mutex<String>>) -> usize {
|
|
||||||
snapshot(stdout).matches("prompt:> ").count()
|
|
||||||
}
|
|
||||||
|
|
||||||
fn response_text_visible(stdout: &Arc<Mutex<String>>) -> bool {
|
|
||||||
snapshot(stdout)
|
|
||||||
.split("Response: ")
|
|
||||||
.skip(1)
|
|
||||||
.any(|text| !text.lines().next().unwrap_or_default().trim().is_empty())
|
|
||||||
}
|
|
||||||
|
|
||||||
fn assert_no_lower_layer_terminal_leaks(
|
|
||||||
stdout: &Arc<Mutex<String>>,
|
|
||||||
stderr: &Arc<Mutex<String>>,
|
|
||||||
) -> Result<(), String> {
|
|
||||||
let leaks = lower_layer_leak_lines("stdout", &snapshot(stdout))
|
|
||||||
.into_iter()
|
|
||||||
.chain(lower_layer_leak_lines("stderr", &snapshot(stderr)))
|
|
||||||
.collect::<Vec<_>>();
|
|
||||||
if leaks.is_empty() {
|
|
||||||
Ok(())
|
|
||||||
} else {
|
|
||||||
Err(format!(
|
|
||||||
"lower-layer runtime output leaked to terminal:\n{}",
|
|
||||||
leaks.join("\n")
|
|
||||||
))
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
fn lower_layer_leak_lines(stream: &str, output: &str) -> Vec<String> {
|
|
||||||
output
|
|
||||||
.lines()
|
|
||||||
.filter_map(|line| {
|
|
||||||
let trimmed = line.trim_start();
|
|
||||||
let leaked = trimmed.contains("prompt_loop_ready")
|
|
||||||
|| trimmed.contains("dashboard_ready")
|
|
||||||
|| trimmed.starts_with("mvp-orchestrator:")
|
|
||||||
|| trimmed.starts_with("mvp-worker-node:")
|
|
||||||
|| trimmed.starts_with("mvp_tinygrad_worker:");
|
|
||||||
leaked.then(|| format!("{stream}: {line}"))
|
|
||||||
})
|
|
||||||
.collect()
|
|
||||||
}
|
|
||||||
fn snapshot(buf: &Arc<Mutex<String>>) -> String {
|
|
||||||
buf.lock().expect("capture mutex").clone()
|
|
||||||
}
|
|
||||||
|
|
||||||
fn wait_child(child: &mut Child, watchdog: Duration) -> Option<std::process::ExitStatus> {
|
|
||||||
let start = Instant::now();
|
|
||||||
while start.elapsed() < watchdog {
|
|
||||||
if let Some(status) = child.try_wait().expect("poll child") {
|
|
||||||
return Some(status);
|
|
||||||
}
|
|
||||||
thread::sleep(Duration::from_millis(100));
|
|
||||||
}
|
|
||||||
None
|
|
||||||
}
|
|
||||||
|
|
||||||
fn request_child_interrupt(child: &Child) {
|
|
||||||
#[cfg(target_os = "linux")]
|
|
||||||
unsafe {
|
|
||||||
let _ = libc::kill(-(child.id() as libc::pid_t), libc::SIGINT);
|
|
||||||
}
|
|
||||||
#[cfg(not(target_os = "linux"))]
|
|
||||||
{
|
|
||||||
let _ = child;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
fn require_docker(root: &std::path::Path) {
|
|
||||||
let version = Command::new("docker")
|
|
||||||
.current_dir(root)
|
|
||||||
.arg("version")
|
|
||||||
.output()
|
|
||||||
.expect("run docker version");
|
|
||||||
assert!(
|
|
||||||
version.status.success(),
|
|
||||||
"docker is not available\nstdout:\n{}\nstderr:\n{}",
|
|
||||||
String::from_utf8_lossy(&version.stdout),
|
|
||||||
String::from_utf8_lossy(&version.stderr)
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
fn assert_no_containers_with_prefix_if_docker_available(root: &std::path::Path, prefix: &str) {
|
|
||||||
let docker_available = Command::new("docker")
|
|
||||||
.current_dir(root)
|
|
||||||
.arg("version")
|
|
||||||
.output()
|
|
||||||
.map(|output| output.status.success())
|
|
||||||
.unwrap_or(false);
|
|
||||||
if docker_available {
|
|
||||||
assert_containers_with_prefix_removed(root, prefix);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
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")
|
|
||||||
.current_dir(root)
|
|
||||||
.args([
|
|
||||||
"ps",
|
|
||||||
"-a",
|
|
||||||
"--filter",
|
|
||||||
&format!("name=^{prefix}-"),
|
|
||||||
"--format",
|
|
||||||
"{{.Names}}",
|
|
||||||
])
|
|
||||||
.output()
|
|
||||||
.expect("run docker ps");
|
|
||||||
assert!(
|
|
||||||
output.status.success(),
|
|
||||||
"docker ps failed\nstdout:\n{}\nstderr:\n{}",
|
|
||||||
String::from_utf8_lossy(&output.stdout),
|
|
||||||
String::from_utf8_lossy(&output.stderr)
|
|
||||||
);
|
|
||||||
containers = String::from_utf8_lossy(&output.stdout).into_owned();
|
|
||||||
if containers.trim().is_empty() {
|
|
||||||
return;
|
|
||||||
}
|
|
||||||
thread::sleep(Duration::from_millis(100));
|
|
||||||
}
|
|
||||||
panic!("containers with prefix {prefix} still exist:\n{containers}");
|
|
||||||
}
|
|
||||||
|
|
||||||
fn workspace_root() -> std::path::PathBuf {
|
|
||||||
std::path::PathBuf::from(env!("CARGO_MANIFEST_DIR"))
|
|
||||||
.parent()
|
|
||||||
.and_then(std::path::Path::parent)
|
|
||||||
.expect("workspace root")
|
|
||||||
.to_path_buf()
|
|
||||||
}
|
|
||||||
|
|
@ -1,109 +0,0 @@
|
||||||
use std::io::{self, BufRead, Write};
|
|
||||||
use std::process::ExitCode;
|
|
||||||
|
|
||||||
use serde_json::json;
|
|
||||||
|
|
||||||
pub fn run_main() -> ExitCode {
|
|
||||||
let stdin = io::stdin();
|
|
||||||
let mut stdout = io::stdout();
|
|
||||||
|
|
||||||
for line in stdin.lock().lines() {
|
|
||||||
let line = match line {
|
|
||||||
Ok(line) => line,
|
|
||||||
Err(error) => {
|
|
||||||
let _ = writeln!(
|
|
||||||
stdout,
|
|
||||||
"{}",
|
|
||||||
json!({"type":"WorkerFatal","reason":format!("stdin:{error}")})
|
|
||||||
);
|
|
||||||
return ExitCode::from(1);
|
|
||||||
}
|
|
||||||
};
|
|
||||||
if line.trim().is_empty() {
|
|
||||||
continue;
|
|
||||||
}
|
|
||||||
let value: serde_json::Value = match serde_json::from_str(&line) {
|
|
||||||
Ok(value) => value,
|
|
||||||
Err(error) => {
|
|
||||||
let _ = writeln!(
|
|
||||||
stdout,
|
|
||||||
"{}",
|
|
||||||
json!({"type":"ProcessFault","reason":format!("invalid-json:{error}")})
|
|
||||||
);
|
|
||||||
continue;
|
|
||||||
}
|
|
||||||
};
|
|
||||||
let Some(kind) = value.get("type").and_then(|value| value.as_str()) else {
|
|
||||||
let _ = writeln!(
|
|
||||||
stdout,
|
|
||||||
"{}",
|
|
||||||
json!({"type":"ProcessFault","reason":"missing-type"})
|
|
||||||
);
|
|
||||||
continue;
|
|
||||||
};
|
|
||||||
match kind {
|
|
||||||
"InitializeWorker" => {
|
|
||||||
let _ = writeln!(stdout, "{}", json!({"type":"WorkerReady","generation":1}));
|
|
||||||
}
|
|
||||||
"InstallRing" => {
|
|
||||||
let ring_id = value
|
|
||||||
.get("ring_id")
|
|
||||||
.and_then(|value| value.as_u64())
|
|
||||||
.unwrap_or(0);
|
|
||||||
let _ = writeln!(
|
|
||||||
stdout,
|
|
||||||
"{}",
|
|
||||||
json!({"type":"RingInstalled","ring_id":ring_id})
|
|
||||||
);
|
|
||||||
}
|
|
||||||
"ExecuteStep" => {
|
|
||||||
let step_id = value
|
|
||||||
.get("step_id")
|
|
||||||
.and_then(|value| value.as_u64())
|
|
||||||
.unwrap_or(0);
|
|
||||||
let _ = writeln!(
|
|
||||||
stdout,
|
|
||||||
"{}",
|
|
||||||
json!({"type":"StepCompleted","step_id":step_id})
|
|
||||||
);
|
|
||||||
}
|
|
||||||
"ReleaseDeviceObject" => {
|
|
||||||
let handle = value
|
|
||||||
.get("handle")
|
|
||||||
.and_then(|value| value.as_u64())
|
|
||||||
.unwrap_or(0);
|
|
||||||
let _ = writeln!(
|
|
||||||
stdout,
|
|
||||||
"{}",
|
|
||||||
json!({"type":"DeviceObjectReleased","handle":handle})
|
|
||||||
);
|
|
||||||
}
|
|
||||||
"RingReadable" => {
|
|
||||||
let ring_id = value
|
|
||||||
.get("ring_id")
|
|
||||||
.and_then(|value| value.as_u64())
|
|
||||||
.unwrap_or(0);
|
|
||||||
let _ = writeln!(
|
|
||||||
stdout,
|
|
||||||
"{}",
|
|
||||||
json!({"type":"RingReadableAck","ring_id":ring_id})
|
|
||||||
);
|
|
||||||
}
|
|
||||||
"ShutdownWorker" => {
|
|
||||||
let _ = writeln!(stdout, "{}", json!({"type":"WorkerStopped"}));
|
|
||||||
let _ = stdout.flush();
|
|
||||||
return ExitCode::SUCCESS;
|
|
||||||
}
|
|
||||||
_ => {
|
|
||||||
let _ = writeln!(
|
|
||||||
stdout,
|
|
||||||
"{}",
|
|
||||||
json!({"type":"ProcessFault","reason":"unknown-command","command":kind})
|
|
||||||
);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
let _ = stdout.flush();
|
|
||||||
}
|
|
||||||
|
|
||||||
ExitCode::SUCCESS
|
|
||||||
}
|
|
||||||
File diff suppressed because it is too large
Load diff
File diff suppressed because it is too large
Load diff
|
|
@ -4,3 +4,7 @@ version = "0.1.0"
|
||||||
edition = "2024"
|
edition = "2024"
|
||||||
|
|
||||||
[dependencies]
|
[dependencies]
|
||||||
|
serde_json = "1"
|
||||||
|
|
||||||
|
[target.'cfg(target_os = "linux")'.dependencies]
|
||||||
|
libc = "0.2"
|
||||||
|
|
|
||||||
|
|
@ -1,11 +1,38 @@
|
||||||
use std::process::{Command, ExitCode};
|
use std::fs;
|
||||||
use std::time::Instant;
|
use std::io::{Read, Write};
|
||||||
|
use std::path::{Path, PathBuf};
|
||||||
|
use std::process::{Child, Command, ExitCode, ExitStatus, Stdio};
|
||||||
|
use std::thread;
|
||||||
|
use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
|
||||||
|
|
||||||
|
#[cfg(target_os = "linux")]
|
||||||
|
use std::os::unix::process::CommandExt;
|
||||||
|
|
||||||
|
use serde_json::Value;
|
||||||
|
|
||||||
struct TestStep {
|
struct TestStep {
|
||||||
label: &'static str,
|
label: &'static str,
|
||||||
args: &'static [&'static str],
|
args: &'static [&'static str],
|
||||||
}
|
}
|
||||||
|
|
||||||
|
const MVP_CHAT_CHECK_TIMEOUT_SECS: u64 = 900;
|
||||||
|
const MVP_CHAT_CHECK_POLL_MS: u64 = 100;
|
||||||
|
const MVP_CHAT_CHECK_TERM_GRACE_MS: u64 = 2_000;
|
||||||
|
const MVP_CHAT_CHECK_PROMPTS: &[u8] = b"ping\nsecond prompt\n";
|
||||||
|
|
||||||
|
struct MvpChatCheckPaths {
|
||||||
|
root: PathBuf,
|
||||||
|
dump_log: PathBuf,
|
||||||
|
}
|
||||||
|
|
||||||
|
struct MvpChatCheckOutput {
|
||||||
|
status: ExitStatus,
|
||||||
|
stdout: String,
|
||||||
|
stderr: String,
|
||||||
|
timed_out: bool,
|
||||||
|
stdin_error: Option<String>,
|
||||||
|
}
|
||||||
|
|
||||||
const BASIC_TESTS: &[TestStep] = &[
|
const BASIC_TESTS: &[TestStep] = &[
|
||||||
TestStep {
|
TestStep {
|
||||||
label: "root crate",
|
label: "root crate",
|
||||||
|
|
@ -62,9 +89,9 @@ USAGE: cargo xtask <command>
|
||||||
|
|
||||||
COMMANDS:
|
COMMANDS:
|
||||||
mvp-chat [--process|--docker|--vastai] [--pipeline-stages n] [--cached-model] [-- args...] Run the human chat wrapper against the real orchestrator/worker bins.
|
mvp-chat [--process|--docker|--vastai] [--pipeline-stages n] [--cached-model] [-- args...] Run the human chat wrapper against the real orchestrator/worker bins.
|
||||||
test Run all basic non-binding tests. This includes the root crate with
|
mvp-chat-check Run real cargo mvp-chat acceptance check.
|
||||||
`cargo test` plus each non-binding repository package with `cargo test -p`.
|
test Run the basic non-binding test barrier: root crate plus each
|
||||||
Feature-gated E2E/bin tests are intentionally excluded."
|
non-binding repository package with `cargo test -p`."
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
@ -84,6 +111,10 @@ fn run_step(step: &TestStep) -> bool {
|
||||||
|
|
||||||
fn run_tests() -> ExitCode {
|
fn run_tests() -> ExitCode {
|
||||||
let start = Instant::now();
|
let start = Instant::now();
|
||||||
|
let check = run_mvp_chat_check();
|
||||||
|
if check != ExitCode::SUCCESS {
|
||||||
|
return check;
|
||||||
|
}
|
||||||
|
|
||||||
for (index, step) in BASIC_TESTS.iter().enumerate() {
|
for (index, step) in BASIC_TESTS.iter().enumerate() {
|
||||||
if !run_step(step) {
|
if !run_step(step) {
|
||||||
|
|
@ -106,7 +137,12 @@ fn run_tests() -> ExitCode {
|
||||||
fn run_mvp_chat(args: Vec<String>) -> ExitCode {
|
fn run_mvp_chat(args: Vec<String>) -> ExitCode {
|
||||||
let mut command = Command::new(cargo_bin());
|
let mut command = Command::new(cargo_bin());
|
||||||
command.args(["run", "--package", "mvp-system", "--bin", "mvp-chat", "--"]);
|
command.args(["run", "--package", "mvp-system", "--bin", "mvp-chat", "--"]);
|
||||||
command.args(args);
|
let forwarded = if args.first().is_some_and(|arg| arg == "--") {
|
||||||
|
args[1..].to_vec()
|
||||||
|
} else {
|
||||||
|
args
|
||||||
|
};
|
||||||
|
command.args(forwarded);
|
||||||
|
|
||||||
match command.status() {
|
match command.status() {
|
||||||
Ok(status) if status.success() => ExitCode::SUCCESS,
|
Ok(status) if status.success() => ExitCode::SUCCESS,
|
||||||
|
|
@ -123,10 +159,550 @@ fn run_mvp_chat(args: Vec<String>) -> ExitCode {
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn workspace_root() -> PathBuf {
|
||||||
|
PathBuf::from(env!("CARGO_MANIFEST_DIR"))
|
||||||
|
.parent()
|
||||||
|
.expect("xtask manifest dir has a parent")
|
||||||
|
.to_path_buf()
|
||||||
|
}
|
||||||
|
|
||||||
|
fn unique_temp_dir(prefix: &str) -> PathBuf {
|
||||||
|
let pid = std::process::id();
|
||||||
|
for attempt in 0..100 {
|
||||||
|
let timestamp_nanos = SystemTime::now()
|
||||||
|
.duration_since(UNIX_EPOCH)
|
||||||
|
.unwrap_or_default()
|
||||||
|
.as_nanos();
|
||||||
|
let root =
|
||||||
|
std::env::temp_dir().join(format!("{prefix}-{pid}-{timestamp_nanos}-{attempt}"));
|
||||||
|
match fs::create_dir(&root) {
|
||||||
|
Ok(()) => return root,
|
||||||
|
Err(error) if error.kind() == std::io::ErrorKind::AlreadyExists => continue,
|
||||||
|
Err(error) => panic!("mvp-chat-check: create temp dir {}: {error}", root.display()),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
panic!("mvp-chat-check: could not allocate unique temp dir for prefix {prefix}");
|
||||||
|
}
|
||||||
|
|
||||||
|
fn write_mvp_chat_check_paths(root: &Path) -> Result<MvpChatCheckPaths, String> {
|
||||||
|
if !root.is_dir() {
|
||||||
|
return Err(format!(
|
||||||
|
"mvp-chat-check: temp root {} is not a directory",
|
||||||
|
root.display()
|
||||||
|
));
|
||||||
|
}
|
||||||
|
let dump_log = root.join("mvp-chat.ndjson");
|
||||||
|
if dump_log.exists() {
|
||||||
|
return Err(format!(
|
||||||
|
"mvp-chat-check: dump log path already exists: {}",
|
||||||
|
dump_log.display()
|
||||||
|
));
|
||||||
|
}
|
||||||
|
Ok(MvpChatCheckPaths {
|
||||||
|
root: root.to_path_buf(),
|
||||||
|
dump_log,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
fn run_mvp_chat_check() -> ExitCode {
|
||||||
|
let workspace = workspace_root();
|
||||||
|
let temp_root = unique_temp_dir("mvp-chat-check");
|
||||||
|
let paths = match write_mvp_chat_check_paths(&temp_root) {
|
||||||
|
Ok(paths) => paths,
|
||||||
|
Err(error) => {
|
||||||
|
eprintln!("{error}");
|
||||||
|
eprintln!(
|
||||||
|
"mvp-chat-check: temp directory kept at {}",
|
||||||
|
temp_root.display()
|
||||||
|
);
|
||||||
|
return ExitCode::from(1);
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
let output = match run_mvp_chat_check_process(&workspace, &paths) {
|
||||||
|
Ok(output) => output,
|
||||||
|
Err(error) => return fail_mvp_chat_check(&error, &paths, "", "", None),
|
||||||
|
};
|
||||||
|
|
||||||
|
if let Some(error) = &output.stdin_error {
|
||||||
|
return fail_mvp_chat_check(
|
||||||
|
error,
|
||||||
|
&paths,
|
||||||
|
&output.stdout,
|
||||||
|
&output.stderr,
|
||||||
|
Some(&output.status),
|
||||||
|
);
|
||||||
|
}
|
||||||
|
if output.timed_out {
|
||||||
|
let reason = format!("timeout after {MVP_CHAT_CHECK_TIMEOUT_SECS} seconds");
|
||||||
|
return fail_mvp_chat_check(
|
||||||
|
&reason,
|
||||||
|
&paths,
|
||||||
|
&output.stdout,
|
||||||
|
&output.stderr,
|
||||||
|
Some(&output.status),
|
||||||
|
);
|
||||||
|
}
|
||||||
|
if !output.status.success() {
|
||||||
|
return fail_mvp_chat_check(
|
||||||
|
"child exited nonzero",
|
||||||
|
&paths,
|
||||||
|
&output.stdout,
|
||||||
|
&output.stderr,
|
||||||
|
Some(&output.status),
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
let responses = match assert_stdout_contains_two_prompt_cycles(&output.stdout) {
|
||||||
|
Ok(responses) => responses,
|
||||||
|
Err(error) => {
|
||||||
|
return fail_mvp_chat_check(
|
||||||
|
&error,
|
||||||
|
&paths,
|
||||||
|
&output.stdout,
|
||||||
|
&output.stderr,
|
||||||
|
Some(&output.status),
|
||||||
|
);
|
||||||
|
}
|
||||||
|
};
|
||||||
|
if let Err(error) = assert_dump_log_facts(&paths.dump_log) {
|
||||||
|
return fail_mvp_chat_check(
|
||||||
|
&error,
|
||||||
|
&paths,
|
||||||
|
&output.stdout,
|
||||||
|
&output.stderr,
|
||||||
|
Some(&output.status),
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
if let Err(error) = fs::remove_dir_all(&paths.root) {
|
||||||
|
eprintln!(
|
||||||
|
"mvp-chat-check: remove temp directory {}: {error}",
|
||||||
|
paths.root.display()
|
||||||
|
);
|
||||||
|
return ExitCode::from(1);
|
||||||
|
}
|
||||||
|
|
||||||
|
println!("mvp-chat-check: ok");
|
||||||
|
for (index, response) in responses.iter().enumerate() {
|
||||||
|
println!("mvp-chat-check: response {}: {}", index + 1, response);
|
||||||
|
}
|
||||||
|
ExitCode::SUCCESS
|
||||||
|
}
|
||||||
|
|
||||||
|
fn run_mvp_chat_check_process(
|
||||||
|
workspace: &Path,
|
||||||
|
paths: &MvpChatCheckPaths,
|
||||||
|
) -> Result<MvpChatCheckOutput, String> {
|
||||||
|
let mut command = Command::new(cargo_bin());
|
||||||
|
command
|
||||||
|
.current_dir(workspace)
|
||||||
|
.args(["mvp-chat", "--", "--cached-model"])
|
||||||
|
.arg(format!("--dump-logs={}", paths.dump_log.display()))
|
||||||
|
.stdin(Stdio::piped())
|
||||||
|
.stdout(Stdio::piped())
|
||||||
|
.stderr(Stdio::piped());
|
||||||
|
|
||||||
|
#[cfg(target_os = "linux")]
|
||||||
|
unsafe {
|
||||||
|
command.pre_exec(|| {
|
||||||
|
let result = libc::setpgid(0, 0);
|
||||||
|
if result == 0 {
|
||||||
|
Ok(())
|
||||||
|
} else {
|
||||||
|
Err(std::io::Error::last_os_error())
|
||||||
|
}
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
let mut child = command
|
||||||
|
.spawn()
|
||||||
|
.map_err(|e| format!("mvp-chat-check: spawn cargo mvp-chat: {e}"))?;
|
||||||
|
let stdout = child
|
||||||
|
.stdout
|
||||||
|
.take()
|
||||||
|
.ok_or_else(|| "mvp-chat-check: child stdout was not piped".to_owned())?;
|
||||||
|
let stderr = child
|
||||||
|
.stderr
|
||||||
|
.take()
|
||||||
|
.ok_or_else(|| "mvp-chat-check: child stderr was not piped".to_owned())?;
|
||||||
|
let stdout_reader = thread::spawn(move || read_pipe_to_string(stdout, "stdout"));
|
||||||
|
let stderr_reader = thread::spawn(move || read_pipe_to_string(stderr, "stderr"));
|
||||||
|
|
||||||
|
let stdin_error = match child.stdin.take() {
|
||||||
|
Some(mut stdin) => {
|
||||||
|
let result = stdin.write_all(MVP_CHAT_CHECK_PROMPTS);
|
||||||
|
drop(stdin);
|
||||||
|
result
|
||||||
|
.err()
|
||||||
|
.map(|error| format!("mvp-chat-check: write child stdin: {error}"))
|
||||||
|
}
|
||||||
|
None => Some("mvp-chat-check: child stdin was not piped".to_owned()),
|
||||||
|
};
|
||||||
|
|
||||||
|
let (status, timed_out) = if stdin_error.is_some() {
|
||||||
|
(terminate_mvp_chat_child(&mut child)?, false)
|
||||||
|
} else {
|
||||||
|
wait_mvp_chat_check_child(&mut child)?
|
||||||
|
};
|
||||||
|
|
||||||
|
let stdout = join_reader(stdout_reader, "stdout")?;
|
||||||
|
let stderr = join_reader(stderr_reader, "stderr")?;
|
||||||
|
Ok(MvpChatCheckOutput {
|
||||||
|
status,
|
||||||
|
stdout,
|
||||||
|
stderr,
|
||||||
|
timed_out,
|
||||||
|
stdin_error,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
fn wait_mvp_chat_check_child(child: &mut Child) -> Result<(ExitStatus, bool), String> {
|
||||||
|
let timeout = Duration::from_secs(MVP_CHAT_CHECK_TIMEOUT_SECS);
|
||||||
|
let poll = Duration::from_millis(MVP_CHAT_CHECK_POLL_MS);
|
||||||
|
let deadline = Instant::now() + timeout;
|
||||||
|
loop {
|
||||||
|
match child.try_wait() {
|
||||||
|
Ok(Some(status)) => return Ok((status, false)),
|
||||||
|
Ok(None) if Instant::now() >= deadline => {
|
||||||
|
return terminate_mvp_chat_child(child).map(|status| (status, true));
|
||||||
|
}
|
||||||
|
Ok(None) => thread::sleep(poll),
|
||||||
|
Err(error) => return Err(format!("mvp-chat-check: poll child status: {error}")),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn terminate_mvp_chat_child(child: &mut Child) -> Result<ExitStatus, String> {
|
||||||
|
#[cfg(target_os = "linux")]
|
||||||
|
{
|
||||||
|
signal_mvp_chat_process_group(child, libc::SIGTERM);
|
||||||
|
let grace_polls = MVP_CHAT_CHECK_TERM_GRACE_MS / MVP_CHAT_CHECK_POLL_MS;
|
||||||
|
for _ in 0..grace_polls {
|
||||||
|
match child.try_wait() {
|
||||||
|
Ok(Some(status)) => return Ok(status),
|
||||||
|
Ok(None) => thread::sleep(Duration::from_millis(MVP_CHAT_CHECK_POLL_MS)),
|
||||||
|
Err(error) => {
|
||||||
|
return Err(format!(
|
||||||
|
"mvp-chat-check: poll child after SIGTERM: {error}"
|
||||||
|
));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
signal_mvp_chat_process_group(child, libc::SIGKILL);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(not(target_os = "linux"))]
|
||||||
|
{
|
||||||
|
let _ = child.kill();
|
||||||
|
}
|
||||||
|
|
||||||
|
child
|
||||||
|
.wait()
|
||||||
|
.map_err(|e| format!("mvp-chat-check: wait for terminated child: {e}"))
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(target_os = "linux")]
|
||||||
|
fn signal_mvp_chat_process_group(child: &Child, signal: libc::c_int) {
|
||||||
|
let process_group = -(child.id() as libc::pid_t);
|
||||||
|
let _ = unsafe { libc::kill(process_group, signal) };
|
||||||
|
}
|
||||||
|
|
||||||
|
fn read_pipe_to_string<R: Read>(mut reader: R, label: &'static str) -> Result<String, String> {
|
||||||
|
let mut text = String::new();
|
||||||
|
reader
|
||||||
|
.read_to_string(&mut text)
|
||||||
|
.map_err(|e| format!("mvp-chat-check: read child {label}: {e}"))?;
|
||||||
|
Ok(text)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn join_reader(
|
||||||
|
handle: thread::JoinHandle<Result<String, String>>,
|
||||||
|
label: &str,
|
||||||
|
) -> Result<String, String> {
|
||||||
|
handle
|
||||||
|
.join()
|
||||||
|
.map_err(|_| format!("mvp-chat-check: child {label} reader panicked"))?
|
||||||
|
}
|
||||||
|
|
||||||
|
fn fail_mvp_chat_check(
|
||||||
|
reason: &str,
|
||||||
|
paths: &MvpChatCheckPaths,
|
||||||
|
stdout: &str,
|
||||||
|
stderr: &str,
|
||||||
|
status: Option<&ExitStatus>,
|
||||||
|
) -> ExitCode {
|
||||||
|
eprintln!("mvp-chat-check: failed: {reason}");
|
||||||
|
if let Some(status) = status {
|
||||||
|
eprintln!("mvp-chat-check: child exit status: {status}");
|
||||||
|
}
|
||||||
|
eprintln!(
|
||||||
|
"mvp-chat-check: temp directory kept at {}",
|
||||||
|
paths.root.display()
|
||||||
|
);
|
||||||
|
eprintln!("--- captured stdout ---");
|
||||||
|
if stdout.is_empty() {
|
||||||
|
eprintln!("<empty>");
|
||||||
|
} else {
|
||||||
|
eprint!("{stdout}");
|
||||||
|
if !stdout.ends_with('\n') {
|
||||||
|
eprintln!();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
eprintln!("--- captured stderr ---");
|
||||||
|
if stderr.is_empty() {
|
||||||
|
eprintln!("<empty>");
|
||||||
|
} else {
|
||||||
|
eprint!("{stderr}");
|
||||||
|
if !stderr.ends_with('\n') {
|
||||||
|
eprintln!();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
ExitCode::from(1)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn assert_stdout_contains_two_prompt_cycles(stdout: &str) -> Result<Vec<String>, String> {
|
||||||
|
let decoding_count = stdout.matches("decoding...").count();
|
||||||
|
if decoding_count < 2 {
|
||||||
|
return Err(format!(
|
||||||
|
"mvp-chat-check: expected at least two decoding... markers, found {decoding_count}"
|
||||||
|
));
|
||||||
|
}
|
||||||
|
let response_count = stdout.matches("Response: ").count();
|
||||||
|
if response_count < 2 {
|
||||||
|
return Err(format!(
|
||||||
|
"mvp-chat-check: expected at least two Response: prefixes, found {response_count}"
|
||||||
|
));
|
||||||
|
}
|
||||||
|
|
||||||
|
let mut cursor = 0;
|
||||||
|
let mut responses = Vec::with_capacity(2);
|
||||||
|
for cycle in 1..=2 {
|
||||||
|
let prompt_at = find_stdout_marker(stdout, "prompt:>", cursor, cycle, "prompt")?;
|
||||||
|
let decoding_at = find_stdout_marker(
|
||||||
|
stdout,
|
||||||
|
"decoding...",
|
||||||
|
prompt_at + "prompt:>".len(),
|
||||||
|
cycle,
|
||||||
|
"decoding",
|
||||||
|
)?;
|
||||||
|
let response_at = find_stdout_marker(
|
||||||
|
stdout,
|
||||||
|
"Response: ",
|
||||||
|
decoding_at + "decoding...".len(),
|
||||||
|
cycle,
|
||||||
|
"response",
|
||||||
|
)?;
|
||||||
|
let response_start = response_at + "Response: ".len();
|
||||||
|
let response_end = stdout[response_start..]
|
||||||
|
.find('\n')
|
||||||
|
.map_or(stdout.len(), |offset| response_start + offset);
|
||||||
|
let response = &stdout[response_start..response_end];
|
||||||
|
if !response.chars().any(|ch| !ch.is_whitespace()) {
|
||||||
|
return Err(format!(
|
||||||
|
"mvp-chat-check: empty Response text for prompt cycle {cycle}"
|
||||||
|
));
|
||||||
|
}
|
||||||
|
responses.push(response.to_owned());
|
||||||
|
cursor = response_end;
|
||||||
|
}
|
||||||
|
Ok(responses)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn find_stdout_marker(
|
||||||
|
stdout: &str,
|
||||||
|
marker: &str,
|
||||||
|
start: usize,
|
||||||
|
cycle: usize,
|
||||||
|
label: &str,
|
||||||
|
) -> Result<usize, String> {
|
||||||
|
stdout[start..]
|
||||||
|
.find(marker)
|
||||||
|
.map(|offset| start + offset)
|
||||||
|
.ok_or_else(|| {
|
||||||
|
format!("mvp-chat-check: missing {label} marker for prompt cycle {cycle}")
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Default)]
|
||||||
|
struct DumpLogFacts {
|
||||||
|
chat_config_ready: bool,
|
||||||
|
prepare_runtime_ready: bool,
|
||||||
|
prompt_rpc_ready: bool,
|
||||||
|
orch_iroh_driver_ready: bool,
|
||||||
|
node_iroh_driver_ready: bool,
|
||||||
|
node_worker_initialize_ready: bool,
|
||||||
|
orch_weights_loaded_ready: bool,
|
||||||
|
response_text_1: bool,
|
||||||
|
request_completed_1: bool,
|
||||||
|
response_text_2: bool,
|
||||||
|
request_completed_2: bool,
|
||||||
|
shutdown_requested: bool,
|
||||||
|
orchestrator_stopped: bool,
|
||||||
|
}
|
||||||
|
|
||||||
|
fn assert_dump_log_facts(path: &Path) -> Result<(), String> {
|
||||||
|
let content = fs::read_to_string(path)
|
||||||
|
.map_err(|e| format!("mvp-chat-check: read dump log {}: {e}", path.display()))?;
|
||||||
|
if content.lines().next().is_none() {
|
||||||
|
return Err(format!(
|
||||||
|
"mvp-chat-check: dump log {} is empty",
|
||||||
|
path.display()
|
||||||
|
));
|
||||||
|
}
|
||||||
|
|
||||||
|
let mut facts = DumpLogFacts::default();
|
||||||
|
for (line_index, line) in content.lines().enumerate() {
|
||||||
|
if line.trim().is_empty() {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
let outer: Value = serde_json::from_str(line).map_err(|e| {
|
||||||
|
format!(
|
||||||
|
"mvp-chat-check: parse dump log {} line {}: {e}",
|
||||||
|
path.display(),
|
||||||
|
line_index + 1
|
||||||
|
)
|
||||||
|
})?;
|
||||||
|
let channel = outer
|
||||||
|
.get("channel")
|
||||||
|
.and_then(Value::as_str)
|
||||||
|
.ok_or_else(|| {
|
||||||
|
format!(
|
||||||
|
"mvp-chat-check: dump log line {} missing channel",
|
||||||
|
line_index + 1
|
||||||
|
)
|
||||||
|
})?;
|
||||||
|
let payload = outer.get("payload").ok_or_else(|| {
|
||||||
|
format!(
|
||||||
|
"mvp-chat-check: dump log line {} missing payload",
|
||||||
|
line_index + 1
|
||||||
|
)
|
||||||
|
})?;
|
||||||
|
if payload.get("encoding").and_then(Value::as_str) != Some("utf8") {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
let inner_text = payload.get("value").and_then(Value::as_str).ok_or_else(|| {
|
||||||
|
format!(
|
||||||
|
"mvp-chat-check: dump log line {} missing utf8 payload value",
|
||||||
|
line_index + 1
|
||||||
|
)
|
||||||
|
})?;
|
||||||
|
let event: Value = serde_json::from_str(inner_text).map_err(|e| {
|
||||||
|
format!(
|
||||||
|
"mvp-chat-check: parse inner event on dump log line {}: {e}",
|
||||||
|
line_index + 1
|
||||||
|
)
|
||||||
|
})?;
|
||||||
|
record_dump_log_event(channel, &event, &mut facts)?;
|
||||||
|
}
|
||||||
|
|
||||||
|
require_dump_log_fact(facts.chat_config_ready, "config ready")?;
|
||||||
|
require_dump_log_fact(facts.prepare_runtime_ready, "prepare_runtime ready")?;
|
||||||
|
require_dump_log_fact(facts.prompt_rpc_ready, "prompt_rpc ready")?;
|
||||||
|
require_dump_log_fact(facts.orch_iroh_driver_ready, "OrchBootstrap iroh_driver ready")?;
|
||||||
|
require_dump_log_fact(facts.node_iroh_driver_ready, "NodeEvent iroh_driver ready")?;
|
||||||
|
require_dump_log_fact(
|
||||||
|
facts.node_worker_initialize_ready,
|
||||||
|
"NodeEvent worker_initialize ready",
|
||||||
|
)?;
|
||||||
|
require_dump_log_fact(facts.orch_weights_loaded_ready, "OrchBootstrap weights_loaded ready")?;
|
||||||
|
require_dump_log_fact(facts.response_text_1, "response_text request_id=1")?;
|
||||||
|
require_dump_log_fact(facts.request_completed_1, "request_completed request_id=1")?;
|
||||||
|
require_dump_log_fact(facts.response_text_2, "response_text request_id=2")?;
|
||||||
|
require_dump_log_fact(facts.request_completed_2, "request_completed request_id=2")?;
|
||||||
|
require_dump_log_fact(facts.shutdown_requested, "shutdown requested")?;
|
||||||
|
require_dump_log_fact(facts.orchestrator_stopped, "orchestrator_process stopped")
|
||||||
|
}
|
||||||
|
|
||||||
|
fn record_dump_log_event(
|
||||||
|
channel: &str,
|
||||||
|
event: &Value,
|
||||||
|
facts: &mut DumpLogFacts,
|
||||||
|
) -> Result<(), String> {
|
||||||
|
let event_type = event.get("type").and_then(Value::as_str);
|
||||||
|
let phase = event.get("phase").and_then(Value::as_str);
|
||||||
|
let status = event.get("status").and_then(Value::as_str);
|
||||||
|
if status == Some("failed") {
|
||||||
|
return Err(format!(
|
||||||
|
"mvp-chat-check: failed event channel={channel} type={} phase={} detail={}",
|
||||||
|
event_type.unwrap_or("<missing>"),
|
||||||
|
phase.unwrap_or("<missing>"),
|
||||||
|
event.get("detail").unwrap_or(&Value::Null)
|
||||||
|
));
|
||||||
|
}
|
||||||
|
|
||||||
|
match (channel, event_type, phase, status) {
|
||||||
|
("mvp.chat.lifecycle", Some("ChatProgress"), Some("config"), Some("ready")) => {
|
||||||
|
facts.chat_config_ready = true;
|
||||||
|
}
|
||||||
|
("mvp.chat.runtime", Some("ChatProgress"), Some("prepare_runtime"), Some("ready")) => {
|
||||||
|
facts.prepare_runtime_ready = true;
|
||||||
|
}
|
||||||
|
("mvp.chat.runtime", Some("ChatProgress"), Some("prompt_rpc"), Some("ready")) => {
|
||||||
|
facts.prompt_rpc_ready = true;
|
||||||
|
}
|
||||||
|
(_, Some("OrchBootstrap"), Some("iroh_driver"), Some("ready")) => {
|
||||||
|
facts.orch_iroh_driver_ready = true;
|
||||||
|
}
|
||||||
|
(_, Some("NodeEvent"), Some("iroh_driver"), Some("ready")) => {
|
||||||
|
facts.node_iroh_driver_ready = true;
|
||||||
|
}
|
||||||
|
(_, Some("NodeEvent"), Some("worker_initialize"), Some("ready")) => {
|
||||||
|
facts.node_worker_initialize_ready = true;
|
||||||
|
}
|
||||||
|
(_, Some("OrchBootstrap"), Some("weights_loaded"), Some("ready")) => {
|
||||||
|
facts.orch_weights_loaded_ready = true;
|
||||||
|
}
|
||||||
|
("mvp.chat.prompt", Some("ChatProgress"), Some("response_text"), Some("observed")) => {
|
||||||
|
match dump_log_request_id(event) {
|
||||||
|
Some(1) => facts.response_text_1 = true,
|
||||||
|
Some(2) => facts.response_text_2 = true,
|
||||||
|
_ => {}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
("mvp.chat.prompt", Some("ChatProgress"), Some("request_completed"), Some("ready")) => {
|
||||||
|
match dump_log_request_id(event) {
|
||||||
|
Some(1) => facts.request_completed_1 = true,
|
||||||
|
Some(2) => facts.request_completed_2 = true,
|
||||||
|
_ => {}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
("mvp.chat.lifecycle", Some("ChatProgress"), Some("shutdown"), Some("requested")) => {
|
||||||
|
facts.shutdown_requested = true;
|
||||||
|
}
|
||||||
|
(
|
||||||
|
"mvp.chat.component",
|
||||||
|
Some("ChatProgress"),
|
||||||
|
Some("orchestrator_process"),
|
||||||
|
Some("stopped"),
|
||||||
|
) => {
|
||||||
|
facts.orchestrator_stopped = true;
|
||||||
|
}
|
||||||
|
_ => {}
|
||||||
|
}
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn dump_log_request_id(event: &Value) -> Option<u64> {
|
||||||
|
event
|
||||||
|
.get("detail")
|
||||||
|
.and_then(|detail| detail.get("request_id"))
|
||||||
|
.and_then(Value::as_u64)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn require_dump_log_fact(found: bool, fact: &str) -> Result<(), String> {
|
||||||
|
if found {
|
||||||
|
Ok(())
|
||||||
|
} else {
|
||||||
|
Err(format!("mvp-chat-check: missing {fact}"))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
fn main() -> ExitCode {
|
fn main() -> ExitCode {
|
||||||
let mut args = std::env::args().skip(1);
|
let mut args = std::env::args().skip(1);
|
||||||
match args.next().as_deref() {
|
match args.next().as_deref() {
|
||||||
Some("test") if args.next().is_none() => run_tests(),
|
Some("test") if args.next().is_none() => run_tests(),
|
||||||
|
Some("mvp-chat-check") if args.next().is_none() => run_mvp_chat_check(),
|
||||||
Some("mvp-chat") => run_mvp_chat(args.collect()),
|
Some("mvp-chat") => run_mvp_chat(args.collect()),
|
||||||
Some("help" | "--help" | "-h") | None => {
|
Some("help" | "--help" | "-h") | None => {
|
||||||
print_usage();
|
print_usage();
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue