795 lines
26 KiB
Rust
795 lines
26 KiB
Rust
|
|
use std::fs::{File, OpenOptions};
|
||
|
|
use std::io::{BufRead, BufReader, Write};
|
||
|
|
use std::process::{Child, ChildStdin, ChildStdout, Command, ExitCode, Stdio};
|
||
|
|
use std::sync::{
|
||
|
|
Arc, OnceLock,
|
||
|
|
mpsc::{self, Receiver},
|
||
|
|
};
|
||
|
|
use std::thread;
|
||
|
|
use std::time::{Duration, Instant};
|
||
|
|
|
||
|
|
use datastream::emit::{ClusterFrameSink, DatastreamEmitter, EmitterConfig, FrameSink, NoopSink};
|
||
|
|
use datastream::{ChannelId, DATASTREAM_SINK_NAME};
|
||
|
|
|
||
|
|
use distribution::node::DistributedNodeConfig;
|
||
|
|
use iroh::EndpointAddr;
|
||
|
|
use iroh_driver::{IrohDriver, IrohDriverConfig};
|
||
|
|
use mvp_system::actors::node_agent::{
|
||
|
|
NodeAgentActor, NodeAgentMsg, NodeAgentReport, StageCommandWire,
|
||
|
|
};
|
||
|
|
use mvp_system::actors::register_mvp_actor_codecs;
|
||
|
|
use mvp_system::distribution_stack::DistributionRuntimeStack;
|
||
|
|
use mvp_system::prompt_rpc::PromptEvent;
|
||
|
|
use mvp_system::run_plan::{GgufSource, TokenizerSource};
|
||
|
|
use mvp_system::stage_controller as stage;
|
||
|
|
use serde_json::{Value, json};
|
||
|
|
use swactor::actor::ActorAddress;
|
||
|
|
|
||
|
|
const DEFAULT_WORKER_SCRIPT: &str = "/usr/local/share/mvp/tinygrad_worker.py";
|
||
|
|
const DEFAULT_DEVICE: &str = "CUDA";
|
||
|
|
const DEFAULT_HF_REPO: &str = "bartowski/Llama-3.2-1B-Instruct-GGUF";
|
||
|
|
const DEFAULT_HF_FILE: &str = "Llama-3.2-1B-Instruct-Q4_K_M.gguf";
|
||
|
|
const DEFAULT_MODEL_ID: &str = "llama-3.2-1b-instruct-q4";
|
||
|
|
const PUMP_INTERVAL: Duration = Duration::from_millis(10);
|
||
|
|
|
||
|
|
fn main() -> ExitCode {
|
||
|
|
match run() {
|
||
|
|
Ok(()) => ExitCode::SUCCESS,
|
||
|
|
Err(error) => {
|
||
|
|
eprintln!("mvp-node: {error}");
|
||
|
|
ExitCode::from(1)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
fn run() -> Result<(), String> {
|
||
|
|
let config = DeploymentConfig::from_env()?;
|
||
|
|
eprintln!(
|
||
|
|
"mvp-node: boot run={} logical_node={} stage={} worker={} device={}",
|
||
|
|
config.run_id,
|
||
|
|
config.logical_node_id,
|
||
|
|
config.stage_index,
|
||
|
|
config.worker_script,
|
||
|
|
config.device
|
||
|
|
);
|
||
|
|
|
||
|
|
let tokio = tokio::runtime::Runtime::new().map_err(|e| format!("tokio runtime: {e}"))?;
|
||
|
|
let mut driver = IrohDriver::with_handle(
|
||
|
|
tokio.handle().clone(),
|
||
|
|
IrohDriverConfig {
|
||
|
|
secret_key: None,
|
||
|
|
relay_mode: config.relay_mode.clone(),
|
||
|
|
node: DistributedNodeConfig::default(),
|
||
|
|
peer_auth: None,
|
||
|
|
additional_alpns: vec![],
|
||
|
|
},
|
||
|
|
)
|
||
|
|
.map_err(|e| format!("create iroh driver: {e}"))?;
|
||
|
|
if let Some(coordinator) = &config.coordinator_endpoint {
|
||
|
|
eprintln!("mvp-node: joining coordinator {coordinator:?}");
|
||
|
|
driver.join(std::slice::from_ref(coordinator));
|
||
|
|
} else {
|
||
|
|
eprintln!("mvp-node: no MVP_COORDINATOR_ENDPOINT set; running standalone until joined");
|
||
|
|
}
|
||
|
|
|
||
|
|
let stack = DistributionRuntimeStack::new_with_codecs(
|
||
|
|
driver.node_id(),
|
||
|
|
DistributedNodeConfig::default(),
|
||
|
|
|registry| {
|
||
|
|
register_mvp_actor_codecs(registry);
|
||
|
|
datastream::wire::register_datastream_codec(registry);
|
||
|
|
},
|
||
|
|
);
|
||
|
|
driver.enable_actor_bridge(
|
||
|
|
stack.runtime.clone(),
|
||
|
|
stack.codec.clone(),
|
||
|
|
stack.actor_bridge_routes(),
|
||
|
|
stack.actors.swim,
|
||
|
|
stack.relay_mirror.clone(),
|
||
|
|
stack.route_view.clone(),
|
||
|
|
);
|
||
|
|
|
||
|
|
let mut datastream = node_datastream(&config, &stack);
|
||
|
|
|
||
|
|
let reports = stack
|
||
|
|
.runtime
|
||
|
|
.new_inbox::<NodeAgentReport>()
|
||
|
|
.map_err(|e| format!("node report inbox: {e}"))?;
|
||
|
|
let orchestrator = config
|
||
|
|
.orchestrator_actor
|
||
|
|
.unwrap_or_else(ActorAddress::new_random);
|
||
|
|
let node_actor = stack
|
||
|
|
.runtime
|
||
|
|
.spawn(NodeAgentActor::new(
|
||
|
|
stage::NodeId(config.logical_node_id),
|
||
|
|
orchestrator,
|
||
|
|
Some(*reports.addr()),
|
||
|
|
))
|
||
|
|
.map_err(|e| format!("spawn node agent: {e}"))?;
|
||
|
|
stack.register_local_actor(driver.register_actor(node_actor, 1));
|
||
|
|
|
||
|
|
let mut worker = TinygradWorker::spawn(&config)?;
|
||
|
|
let mut initial_pump = || {};
|
||
|
|
worker.initialize(&config.device, &mut datastream, &mut initial_pump)?;
|
||
|
|
stack
|
||
|
|
.runtime
|
||
|
|
.send_to(node_actor, NodeAgentMsg::MarkWorkerReady)
|
||
|
|
.map_err(|e| format!("mark initialized worker ready: {e}"))?;
|
||
|
|
|
||
|
|
let ready = json!({
|
||
|
|
"type":"ready",
|
||
|
|
"role":"node",
|
||
|
|
"endpoint": driver.endpoint_addr(),
|
||
|
|
"node_actor": node_actor,
|
||
|
|
"logical_node_id": config.logical_node_id,
|
||
|
|
"stage_index": config.stage_index,
|
||
|
|
});
|
||
|
|
println!("{ready}");
|
||
|
|
datastream.submit_text(ChannelId::new("mvp.node.ready"), ready.to_string());
|
||
|
|
std::io::stdout()
|
||
|
|
.flush()
|
||
|
|
.map_err(|e| format!("flush ready line: {e}"))?;
|
||
|
|
|
||
|
|
if let Some(prompt) = &config.self_test_prompt {
|
||
|
|
run_self_test(&mut worker, &config, prompt, &mut datastream)?;
|
||
|
|
}
|
||
|
|
|
||
|
|
let shutdown_rx = spawn_stdin_shutdown_listener();
|
||
|
|
let started = Instant::now();
|
||
|
|
loop {
|
||
|
|
pump_network(&mut driver, &stack);
|
||
|
|
datastream.tick();
|
||
|
|
while let Some(report) = reports.try_recv() {
|
||
|
|
handle_node_report(
|
||
|
|
report,
|
||
|
|
&stack,
|
||
|
|
&mut driver,
|
||
|
|
node_actor,
|
||
|
|
&mut worker,
|
||
|
|
&mut datastream,
|
||
|
|
)?;
|
||
|
|
}
|
||
|
|
if shutdown_rx.try_recv().is_ok() {
|
||
|
|
eprintln!("mvp-node: shutdown requested on stdin");
|
||
|
|
let mut pump = || pump_network(&mut driver, &stack);
|
||
|
|
let _ = worker.shutdown(&mut datastream, &mut pump);
|
||
|
|
return Ok(());
|
||
|
|
}
|
||
|
|
if let Some(status) = worker.try_wait()? {
|
||
|
|
let _ = stack
|
||
|
|
.runtime
|
||
|
|
.send_to(node_actor, NodeAgentMsg::WorkerCrashed);
|
||
|
|
return Err(format!("tinygrad helper exited with {status}"));
|
||
|
|
}
|
||
|
|
if started.elapsed() > config.max_runtime && config.max_runtime != Duration::ZERO {
|
||
|
|
eprintln!("mvp-node: max runtime elapsed; shutting down cleanly");
|
||
|
|
let mut pump = || pump_network(&mut driver, &stack);
|
||
|
|
let _ = worker.shutdown(&mut datastream, &mut pump);
|
||
|
|
return Ok(());
|
||
|
|
}
|
||
|
|
thread::sleep(PUMP_INTERVAL);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
fn pump_network(driver: &mut IrohDriver, stack: &DistributionRuntimeStack) {
|
||
|
|
stack.tick_protocol_actors(Instant::now());
|
||
|
|
driver.pump_inbound_to_actors();
|
||
|
|
stack.pump_runtime_once();
|
||
|
|
driver.drain_outbox(&stack.outbox);
|
||
|
|
}
|
||
|
|
|
||
|
|
fn node_datastream(
|
||
|
|
config: &DeploymentConfig,
|
||
|
|
stack: &DistributionRuntimeStack,
|
||
|
|
) -> DatastreamEmitter {
|
||
|
|
let mut sinks: Vec<Box<dyn FrameSink>> = Vec::new();
|
||
|
|
if let Some(actor) = config.datastream_sink_actor {
|
||
|
|
let sink_addr = Arc::new(OnceLock::new());
|
||
|
|
let _ = sink_addr.set(actor);
|
||
|
|
eprintln!("mvp-node: native datastream targeting {DATASTREAM_SINK_NAME} actor {actor:?}");
|
||
|
|
sinks.push(Box::new(ClusterFrameSink::new(
|
||
|
|
stack.runtime.clone(),
|
||
|
|
sink_addr,
|
||
|
|
)));
|
||
|
|
}
|
||
|
|
if let Some(path) = &config.datastream_frame_log {
|
||
|
|
match JsonlFrameSink::open(path) {
|
||
|
|
Ok(sink) => {
|
||
|
|
eprintln!("mvp-node: native datastream frame log {path}");
|
||
|
|
sinks.push(Box::new(sink));
|
||
|
|
}
|
||
|
|
Err(error) => {
|
||
|
|
eprintln!("mvp-node: failed to open MVP_DATASTREAM_FRAME_LOG={path:?}: {error}");
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
if sinks.is_empty() {
|
||
|
|
eprintln!("mvp-node: no datastream sink configured; native datastream drops locally");
|
||
|
|
sinks.push(Box::new(NoopSink));
|
||
|
|
}
|
||
|
|
let sink: Box<dyn FrameSink> = if sinks.len() == 1 {
|
||
|
|
sinks.pop().expect("one sink")
|
||
|
|
} else {
|
||
|
|
Box::new(TeeFrameSink { sinks })
|
||
|
|
};
|
||
|
|
DatastreamEmitter::new(
|
||
|
|
EmitterConfig {
|
||
|
|
node_hex: config.logical_node_id.to_string(),
|
||
|
|
life: config.run_id,
|
||
|
|
mux_capacity: 256,
|
||
|
|
},
|
||
|
|
sink,
|
||
|
|
)
|
||
|
|
}
|
||
|
|
|
||
|
|
struct TeeFrameSink {
|
||
|
|
sinks: Vec<Box<dyn FrameSink>>,
|
||
|
|
}
|
||
|
|
|
||
|
|
impl FrameSink for TeeFrameSink {
|
||
|
|
fn ship(&mut self, stream: &datastream::StreamId, frame: &datastream::Frame) {
|
||
|
|
for sink in &mut self.sinks {
|
||
|
|
sink.ship(stream, frame);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
struct JsonlFrameSink {
|
||
|
|
file: File,
|
||
|
|
}
|
||
|
|
|
||
|
|
impl JsonlFrameSink {
|
||
|
|
fn open(path: &str) -> std::io::Result<Self> {
|
||
|
|
Ok(Self {
|
||
|
|
file: OpenOptions::new().create(true).append(true).open(path)?,
|
||
|
|
})
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
impl FrameSink for JsonlFrameSink {
|
||
|
|
fn ship(&mut self, stream: &datastream::StreamId, frame: &datastream::Frame) {
|
||
|
|
let record = json!({
|
||
|
|
"stream":stream.to_string(),
|
||
|
|
"channel":frame.channel.as_str(),
|
||
|
|
"position":frame.position.0,
|
||
|
|
"payload":String::from_utf8_lossy(&frame.payload),
|
||
|
|
});
|
||
|
|
let _ = serde_json::to_writer(&mut self.file, &record);
|
||
|
|
let _ = writeln!(self.file);
|
||
|
|
let _ = self.file.flush();
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
fn handle_node_report(
|
||
|
|
report: NodeAgentReport,
|
||
|
|
stack: &DistributionRuntimeStack,
|
||
|
|
driver: &mut IrohDriver,
|
||
|
|
node_actor: ActorAddress,
|
||
|
|
worker: &mut TinygradWorker,
|
||
|
|
datastream: &mut DatastreamEmitter,
|
||
|
|
) -> Result<(), String> {
|
||
|
|
match report {
|
||
|
|
NodeAgentReport::Command(command) => {
|
||
|
|
handle_stage_command(command, stack, driver, node_actor, worker, datastream)
|
||
|
|
}
|
||
|
|
NodeAgentReport::Lifecycle(event) => {
|
||
|
|
let record = json!({"type":"node_lifecycle","event":format!("{event:?}")});
|
||
|
|
println!("{record}");
|
||
|
|
datastream.submit_text(ChannelId::new("mvp.node.lifecycle"), record.to_string());
|
||
|
|
std::io::stdout()
|
||
|
|
.flush()
|
||
|
|
.map_err(|e| format!("flush lifecycle: {e}"))
|
||
|
|
}
|
||
|
|
NodeAgentReport::PromptRequested {
|
||
|
|
request_id,
|
||
|
|
prompt,
|
||
|
|
max_tokens,
|
||
|
|
reply_to,
|
||
|
|
} => handle_prompt_request(
|
||
|
|
request_id, prompt, max_tokens, reply_to, stack, driver, worker, datastream,
|
||
|
|
),
|
||
|
|
NodeAgentReport::Snapshot { .. } => Ok(()),
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
fn handle_prompt_request(
|
||
|
|
request_id: u64,
|
||
|
|
prompt: String,
|
||
|
|
max_tokens: u32,
|
||
|
|
reply_to: ActorAddress,
|
||
|
|
stack: &DistributionRuntimeStack,
|
||
|
|
driver: &mut IrohDriver,
|
||
|
|
worker: &mut TinygradWorker,
|
||
|
|
datastream: &mut DatastreamEmitter,
|
||
|
|
) -> Result<(), String> {
|
||
|
|
let started = Instant::now();
|
||
|
|
let mut pump = || pump_network(driver, stack);
|
||
|
|
match worker.infer_prompt(&prompt, max_tokens, datastream, &mut pump) {
|
||
|
|
Ok(result) => {
|
||
|
|
let text = result
|
||
|
|
.get("text")
|
||
|
|
.and_then(Value::as_str)
|
||
|
|
.unwrap_or_default()
|
||
|
|
.to_owned();
|
||
|
|
let tokens_generated = result
|
||
|
|
.get("generated_tokens")
|
||
|
|
.and_then(Value::as_array)
|
||
|
|
.map_or(0, |tokens| tokens.len() as u32);
|
||
|
|
if !text.is_empty() {
|
||
|
|
stack
|
||
|
|
.runtime
|
||
|
|
.send_to(
|
||
|
|
reply_to,
|
||
|
|
PromptEvent::TextDelta {
|
||
|
|
request_id,
|
||
|
|
text: text.clone(),
|
||
|
|
},
|
||
|
|
)
|
||
|
|
.map_err(|e| format!("send prompt text delta: {e}"))?;
|
||
|
|
}
|
||
|
|
stack
|
||
|
|
.runtime
|
||
|
|
.send_to(
|
||
|
|
reply_to,
|
||
|
|
PromptEvent::Done {
|
||
|
|
request_id,
|
||
|
|
final_text: text,
|
||
|
|
tokens_generated,
|
||
|
|
elapsed_ms: result
|
||
|
|
.get("elapsed_ms")
|
||
|
|
.and_then(Value::as_u64)
|
||
|
|
.unwrap_or_else(|| started.elapsed().as_millis() as u64),
|
||
|
|
},
|
||
|
|
)
|
||
|
|
.map_err(|e| format!("send prompt done: {e}"))
|
||
|
|
}
|
||
|
|
Err(error) => stack
|
||
|
|
.runtime
|
||
|
|
.send_to(reply_to, PromptEvent::Fault { request_id, error })
|
||
|
|
.map_err(|e| format!("send prompt fault: {e}")),
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
fn handle_stage_command(
|
||
|
|
command: StageCommandWire,
|
||
|
|
stack: &DistributionRuntimeStack,
|
||
|
|
driver: &mut IrohDriver,
|
||
|
|
node_actor: ActorAddress,
|
||
|
|
worker: &mut TinygradWorker,
|
||
|
|
datastream: &mut DatastreamEmitter,
|
||
|
|
) -> Result<(), String> {
|
||
|
|
match command {
|
||
|
|
StageCommandWire::ConfigureWorkerRole {
|
||
|
|
run_id,
|
||
|
|
stage_index,
|
||
|
|
layer_start,
|
||
|
|
layer_end_exclusive,
|
||
|
|
} => {
|
||
|
|
let mut pump = || pump_network(driver, stack);
|
||
|
|
worker.configure_role(
|
||
|
|
run_id,
|
||
|
|
stage_index,
|
||
|
|
layer_start,
|
||
|
|
layer_end_exclusive,
|
||
|
|
datastream,
|
||
|
|
&mut pump,
|
||
|
|
)?;
|
||
|
|
stack
|
||
|
|
.runtime
|
||
|
|
.send_to(node_actor, NodeAgentMsg::MarkWorkerReady)
|
||
|
|
.map_err(|e| format!("mark worker ready: {e}"))
|
||
|
|
}
|
||
|
|
StageCommandWire::LoadWeights {
|
||
|
|
model_id,
|
||
|
|
gguf_source,
|
||
|
|
tokenizer,
|
||
|
|
layer_start,
|
||
|
|
layer_end_exclusive,
|
||
|
|
} => {
|
||
|
|
let mut pump = || pump_network(driver, stack);
|
||
|
|
worker.load_weights(
|
||
|
|
model_id,
|
||
|
|
gguf_source,
|
||
|
|
tokenizer,
|
||
|
|
layer_start,
|
||
|
|
layer_end_exclusive,
|
||
|
|
datastream,
|
||
|
|
&mut pump,
|
||
|
|
)?;
|
||
|
|
stack
|
||
|
|
.runtime
|
||
|
|
.send_to(node_actor, NodeAgentMsg::MarkWeightsReady)
|
||
|
|
.map_err(|e| format!("mark weights ready: {e}"))
|
||
|
|
}
|
||
|
|
StageCommandWire::StopLocalEdges { run_id } => {
|
||
|
|
stack
|
||
|
|
.runtime
|
||
|
|
.send_to(node_actor, NodeAgentMsg::LocalEdgesStopped { run_id })
|
||
|
|
.map_err(|e| format!("mark local edges stopped: {e}"))?;
|
||
|
|
stack
|
||
|
|
.runtime
|
||
|
|
.send_to(node_actor, NodeAgentMsg::WorkerRingsQuiesced { run_id })
|
||
|
|
.map_err(|e| format!("mark worker rings quiesced: {e}"))
|
||
|
|
}
|
||
|
|
StageCommandWire::ReleaseRunDeviceObjects { run_id } => {
|
||
|
|
stack
|
||
|
|
.runtime
|
||
|
|
.send_to(node_actor, NodeAgentMsg::DeviceObjectsReleased { run_id })
|
||
|
|
.map_err(|e| format!("mark device objects released: {e}"))?;
|
||
|
|
stack
|
||
|
|
.runtime
|
||
|
|
.send_to(node_actor, NodeAgentMsg::WorkerRoleReset { run_id })
|
||
|
|
.map_err(|e| format!("mark worker role reset: {e}"))
|
||
|
|
}
|
||
|
|
StageCommandWire::EstablishInboundEdge { edge_id } => stack
|
||
|
|
.runtime
|
||
|
|
.send_to(node_actor, NodeAgentMsg::MarkInboundEdgeReady { edge_id })
|
||
|
|
.map_err(|e| format!("mark inbound edge ready: {e}")),
|
||
|
|
StageCommandWire::EstablishOutboundEdge { edge_id } => stack
|
||
|
|
.runtime
|
||
|
|
.send_to(node_actor, NodeAgentMsg::MarkOutboundEdgeReady { edge_id })
|
||
|
|
.map_err(|e| format!("mark outbound edge ready: {e}")),
|
||
|
|
StageCommandWire::RewireEdge { .. } | StageCommandWire::ReleaseInputHandle { .. } => Ok(()),
|
||
|
|
StageCommandWire::ExecuteStep { .. } => Err(format!(
|
||
|
|
"mvp-node image path does not carry ring payload commands yet: {command:?}"
|
||
|
|
)),
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
fn run_self_test(
|
||
|
|
worker: &mut TinygradWorker,
|
||
|
|
config: &DeploymentConfig,
|
||
|
|
prompt: &str,
|
||
|
|
datastream: &mut DatastreamEmitter,
|
||
|
|
) -> Result<(), String> {
|
||
|
|
eprintln!("mvp-node: running self-test prompt");
|
||
|
|
let mut pump = || {};
|
||
|
|
worker.configure_role(
|
||
|
|
config.run_id,
|
||
|
|
config.stage_index,
|
||
|
|
0,
|
||
|
|
config.self_test_layer_end,
|
||
|
|
datastream,
|
||
|
|
&mut pump,
|
||
|
|
)?;
|
||
|
|
worker.load_weights(
|
||
|
|
config.model_id.clone(),
|
||
|
|
config.gguf_source.clone(),
|
||
|
|
config.tokenizer.clone(),
|
||
|
|
0,
|
||
|
|
config.self_test_layer_end,
|
||
|
|
datastream,
|
||
|
|
&mut pump,
|
||
|
|
)?;
|
||
|
|
let result = worker.infer_prompt(prompt, config.self_test_max_tokens, datastream, &mut pump)?;
|
||
|
|
let record = json!({
|
||
|
|
"type":"self_test_completed",
|
||
|
|
"prompt":prompt,
|
||
|
|
"result":result,
|
||
|
|
});
|
||
|
|
println!("{record}");
|
||
|
|
datastream.submit_text(ChannelId::new("mvp.node.self_test"), record.to_string());
|
||
|
|
std::io::stdout()
|
||
|
|
.flush()
|
||
|
|
.map_err(|e| format!("flush self-test: {e}"))
|
||
|
|
}
|
||
|
|
|
||
|
|
#[derive(Clone)]
|
||
|
|
struct DeploymentConfig {
|
||
|
|
run_id: u64,
|
||
|
|
logical_node_id: u64,
|
||
|
|
stage_index: u32,
|
||
|
|
coordinator_endpoint: Option<EndpointAddr>,
|
||
|
|
orchestrator_actor: Option<ActorAddress>,
|
||
|
|
datastream_sink_actor: Option<ActorAddress>,
|
||
|
|
datastream_frame_log: Option<String>,
|
||
|
|
relay_mode: iroh::RelayMode,
|
||
|
|
worker_script: String,
|
||
|
|
device: String,
|
||
|
|
model_id: String,
|
||
|
|
gguf_source: GgufSource,
|
||
|
|
tokenizer: TokenizerSource,
|
||
|
|
self_test_prompt: Option<String>,
|
||
|
|
self_test_layer_end: u32,
|
||
|
|
self_test_max_tokens: u32,
|
||
|
|
max_runtime: Duration,
|
||
|
|
}
|
||
|
|
|
||
|
|
impl DeploymentConfig {
|
||
|
|
fn from_env() -> Result<Self, String> {
|
||
|
|
Ok(Self {
|
||
|
|
run_id: env_u64("MVP_RUN_ID", 1)?,
|
||
|
|
logical_node_id: env_u64("MVP_LOGICAL_NODE_ID", 1)?,
|
||
|
|
stage_index: env_u32("MVP_STAGE_INDEX", 0)?,
|
||
|
|
coordinator_endpoint: env_json("MVP_COORDINATOR_ENDPOINT")?,
|
||
|
|
orchestrator_actor: env_json("MVP_ORCHESTRATOR_ACTOR")?,
|
||
|
|
datastream_sink_actor: env_json("MVP_DATASTREAM_SINK_ACTOR")?,
|
||
|
|
datastream_frame_log: env_optional("MVP_DATASTREAM_FRAME_LOG"),
|
||
|
|
relay_mode: relay_mode_from_env()?,
|
||
|
|
worker_script: env_string("MVP_TINYGRAD_WORKER", DEFAULT_WORKER_SCRIPT),
|
||
|
|
device: env_string("DEV", DEFAULT_DEVICE),
|
||
|
|
model_id: env_string("MVP_MODEL_ID", DEFAULT_MODEL_ID),
|
||
|
|
gguf_source: gguf_source_from_env(),
|
||
|
|
tokenizer: tokenizer_from_env(),
|
||
|
|
self_test_prompt: env_optional("MVP_NODE_SELF_TEST_PROMPT"),
|
||
|
|
self_test_layer_end: env_u32("MVP_SELF_TEST_LAYER_END", 16)?,
|
||
|
|
self_test_max_tokens: env_u32("MVP_SELF_TEST_MAX_TOKENS", 1)?,
|
||
|
|
max_runtime: Duration::from_secs(env_u64("MVP_NODE_MAX_RUNTIME_SECS", 0)?),
|
||
|
|
})
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
struct TinygradWorker {
|
||
|
|
child: Child,
|
||
|
|
stdin: ChildStdin,
|
||
|
|
stdout: BufReader<ChildStdout>,
|
||
|
|
}
|
||
|
|
|
||
|
|
impl TinygradWorker {
|
||
|
|
fn spawn(config: &DeploymentConfig) -> Result<Self, String> {
|
||
|
|
let mut child = Command::new("python3")
|
||
|
|
.arg(&config.worker_script)
|
||
|
|
.env("DEV", &config.device)
|
||
|
|
.stdin(Stdio::piped())
|
||
|
|
.stdout(Stdio::piped())
|
||
|
|
.stderr(Stdio::inherit())
|
||
|
|
.spawn()
|
||
|
|
.map_err(|e| format!("spawn tinygrad helper {}: {e}", config.worker_script))?;
|
||
|
|
let stdin = child
|
||
|
|
.stdin
|
||
|
|
.take()
|
||
|
|
.ok_or_else(|| "tinygrad helper stdin missing".to_owned())?;
|
||
|
|
let stdout = child
|
||
|
|
.stdout
|
||
|
|
.take()
|
||
|
|
.ok_or_else(|| "tinygrad helper stdout missing".to_owned())?;
|
||
|
|
Ok(Self {
|
||
|
|
child,
|
||
|
|
stdin,
|
||
|
|
stdout: BufReader::new(stdout),
|
||
|
|
})
|
||
|
|
}
|
||
|
|
|
||
|
|
fn initialize(
|
||
|
|
&mut self,
|
||
|
|
device: &str,
|
||
|
|
datastream: &mut DatastreamEmitter,
|
||
|
|
pump: &mut dyn FnMut(),
|
||
|
|
) -> Result<(), String> {
|
||
|
|
self.command(
|
||
|
|
json!({"type":"InitializeWorker","helper_abi_version":1,"backend":{"device":device}}),
|
||
|
|
"WorkerReady",
|
||
|
|
datastream,
|
||
|
|
ChannelId::new("mvp.worker.initialize"),
|
||
|
|
pump,
|
||
|
|
)
|
||
|
|
.map(|_| ())
|
||
|
|
}
|
||
|
|
|
||
|
|
fn configure_role(
|
||
|
|
&mut self,
|
||
|
|
run_id: u64,
|
||
|
|
stage_index: u32,
|
||
|
|
layer_start: u32,
|
||
|
|
layer_end_exclusive: u32,
|
||
|
|
datastream: &mut DatastreamEmitter,
|
||
|
|
pump: &mut dyn FnMut(),
|
||
|
|
) -> Result<(), String> {
|
||
|
|
self.command(
|
||
|
|
json!({
|
||
|
|
"type":"ConfigureRole",
|
||
|
|
"role_id":stage_index + 1,
|
||
|
|
"config":{
|
||
|
|
"run_id":run_id,
|
||
|
|
"stage_index":stage_index,
|
||
|
|
"layer_start":layer_start,
|
||
|
|
"layer_end_exclusive":layer_end_exclusive,
|
||
|
|
}
|
||
|
|
}),
|
||
|
|
"RoleConfigured",
|
||
|
|
datastream,
|
||
|
|
ChannelId::new("mvp.worker.role"),
|
||
|
|
pump,
|
||
|
|
)
|
||
|
|
.map(|_| ())
|
||
|
|
}
|
||
|
|
|
||
|
|
fn load_weights(
|
||
|
|
&mut self,
|
||
|
|
model_id: String,
|
||
|
|
gguf_source: GgufSource,
|
||
|
|
tokenizer: TokenizerSource,
|
||
|
|
layer_start: u32,
|
||
|
|
layer_end_exclusive: u32,
|
||
|
|
datastream: &mut DatastreamEmitter,
|
||
|
|
pump: &mut dyn FnMut(),
|
||
|
|
) -> Result<(), String> {
|
||
|
|
self.command(
|
||
|
|
json!({
|
||
|
|
"type":"LoadWeights",
|
||
|
|
"model_id":model_id,
|
||
|
|
"gguf_source":gguf_source,
|
||
|
|
"tokenizer":tokenizer,
|
||
|
|
"layer_start":layer_start,
|
||
|
|
"layer_end_exclusive":layer_end_exclusive,
|
||
|
|
}),
|
||
|
|
"WeightsLoaded",
|
||
|
|
datastream,
|
||
|
|
ChannelId::new("mvp.worker.weights"),
|
||
|
|
pump,
|
||
|
|
)
|
||
|
|
.map(|_| ())
|
||
|
|
}
|
||
|
|
|
||
|
|
fn infer_prompt(
|
||
|
|
&mut self,
|
||
|
|
prompt: &str,
|
||
|
|
max_tokens: u32,
|
||
|
|
datastream: &mut DatastreamEmitter,
|
||
|
|
pump: &mut dyn FnMut(),
|
||
|
|
) -> Result<Value, String> {
|
||
|
|
self.command(
|
||
|
|
json!({"type":"InferPrompt","prompt":prompt,"max_tokens":max_tokens}),
|
||
|
|
"PromptCompleted",
|
||
|
|
datastream,
|
||
|
|
ChannelId::new("mvp.worker.prompt"),
|
||
|
|
pump,
|
||
|
|
)
|
||
|
|
}
|
||
|
|
|
||
|
|
fn shutdown(
|
||
|
|
&mut self,
|
||
|
|
datastream: &mut DatastreamEmitter,
|
||
|
|
pump: &mut dyn FnMut(),
|
||
|
|
) -> Result<(), String> {
|
||
|
|
self.command(
|
||
|
|
json!({"type":"ShutdownWorker"}),
|
||
|
|
"WorkerStopped",
|
||
|
|
datastream,
|
||
|
|
ChannelId::new("mvp.worker.shutdown"),
|
||
|
|
pump,
|
||
|
|
)
|
||
|
|
.map(|_| ())
|
||
|
|
}
|
||
|
|
|
||
|
|
fn try_wait(&mut self) -> Result<Option<std::process::ExitStatus>, String> {
|
||
|
|
self.child
|
||
|
|
.try_wait()
|
||
|
|
.map_err(|e| format!("poll tinygrad helper: {e}"))
|
||
|
|
}
|
||
|
|
|
||
|
|
fn command(
|
||
|
|
&mut self,
|
||
|
|
command: Value,
|
||
|
|
expected: &str,
|
||
|
|
datastream: &mut DatastreamEmitter,
|
||
|
|
channel: ChannelId,
|
||
|
|
pump: &mut dyn FnMut(),
|
||
|
|
) -> Result<Value, String> {
|
||
|
|
writeln!(self.stdin, "{command}").map_err(|e| format!("write helper command: {e}"))?;
|
||
|
|
self.stdin
|
||
|
|
.flush()
|
||
|
|
.map_err(|e| format!("flush helper command: {e}"))?;
|
||
|
|
self.expect_event(expected, datastream, channel, pump)
|
||
|
|
}
|
||
|
|
|
||
|
|
fn expect_event(
|
||
|
|
&mut self,
|
||
|
|
expected: &str,
|
||
|
|
datastream: &mut DatastreamEmitter,
|
||
|
|
channel: ChannelId,
|
||
|
|
pump: &mut dyn FnMut(),
|
||
|
|
) -> Result<Value, String> {
|
||
|
|
loop {
|
||
|
|
let mut line = String::new();
|
||
|
|
let n = self
|
||
|
|
.stdout
|
||
|
|
.read_line(&mut line)
|
||
|
|
.map_err(|e| format!("read helper stdout: {e}"))?;
|
||
|
|
if n == 0 {
|
||
|
|
return Err(format!(
|
||
|
|
"tinygrad helper stdout closed while waiting for {expected}"
|
||
|
|
));
|
||
|
|
}
|
||
|
|
let value: Value = serde_json::from_str(&line)
|
||
|
|
.map_err(|e| format!("parse helper stdout {line:?}: {e}"))?;
|
||
|
|
datastream.submit_text(channel.clone(), value.to_string());
|
||
|
|
datastream.tick();
|
||
|
|
pump();
|
||
|
|
if value.get("type").and_then(Value::as_str) == Some(expected) {
|
||
|
|
return Ok(value);
|
||
|
|
}
|
||
|
|
println!("{}", json!({"type":"tinygrad_event","event":value}));
|
||
|
|
std::io::stdout()
|
||
|
|
.flush()
|
||
|
|
.map_err(|e| format!("flush helper passthrough: {e}"))?;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
impl Drop for TinygradWorker {
|
||
|
|
fn drop(&mut self) {
|
||
|
|
let _ = self.child.kill();
|
||
|
|
let _ = self.child.wait();
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
fn spawn_stdin_shutdown_listener() -> Receiver<()> {
|
||
|
|
let (tx, rx) = mpsc::channel();
|
||
|
|
thread::spawn(move || {
|
||
|
|
let stdin = std::io::stdin();
|
||
|
|
for line in stdin.lock().lines().map_while(Result::ok) {
|
||
|
|
if line.trim().eq_ignore_ascii_case("shutdown") {
|
||
|
|
let _ = tx.send(());
|
||
|
|
break;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
});
|
||
|
|
rx
|
||
|
|
}
|
||
|
|
|
||
|
|
fn env_string(name: &str, default: &str) -> String {
|
||
|
|
env_optional(name).unwrap_or_else(|| default.to_owned())
|
||
|
|
}
|
||
|
|
|
||
|
|
fn env_optional(name: &str) -> Option<String> {
|
||
|
|
std::env::var(name)
|
||
|
|
.ok()
|
||
|
|
.map(|value| value.trim().to_owned())
|
||
|
|
.filter(|value| !value.is_empty())
|
||
|
|
}
|
||
|
|
|
||
|
|
fn env_u64(name: &str, default: u64) -> Result<u64, String> {
|
||
|
|
match env_optional(name) {
|
||
|
|
Some(value) => value
|
||
|
|
.parse::<u64>()
|
||
|
|
.map_err(|e| format!("invalid {name}={value:?}: {e}")),
|
||
|
|
None => Ok(default),
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
fn env_u32(name: &str, default: u32) -> Result<u32, String> {
|
||
|
|
match env_optional(name) {
|
||
|
|
Some(value) => value
|
||
|
|
.parse::<u32>()
|
||
|
|
.map_err(|e| format!("invalid {name}={value:?}: {e}")),
|
||
|
|
None => Ok(default),
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
fn env_json<T>(name: &str) -> Result<Option<T>, String>
|
||
|
|
where
|
||
|
|
T: serde::de::DeserializeOwned,
|
||
|
|
{
|
||
|
|
env_optional(name)
|
||
|
|
.map(|value| serde_json::from_str(&value).map_err(|e| format!("invalid {name} JSON: {e}")))
|
||
|
|
.transpose()
|
||
|
|
}
|
||
|
|
|
||
|
|
fn relay_mode_from_env() -> Result<iroh::RelayMode, String> {
|
||
|
|
match env_string("MVP_IROH_RELAY_MODE", "disabled").as_str() {
|
||
|
|
"disabled" => Ok(iroh::RelayMode::Disabled),
|
||
|
|
"default" => Ok(iroh::RelayMode::Default),
|
||
|
|
other => Err(format!(
|
||
|
|
"unsupported MVP_IROH_RELAY_MODE={other:?}; use disabled or default"
|
||
|
|
)),
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
fn gguf_source_from_env() -> GgufSource {
|
||
|
|
if let Some(path) = env_optional("MVP_GGUF_LOCAL_PATH") {
|
||
|
|
return GgufSource::LocalPath(path);
|
||
|
|
}
|
||
|
|
GgufSource::HuggingFaceGguf {
|
||
|
|
repo: env_string("MVP_GGUF_REPO", DEFAULT_HF_REPO),
|
||
|
|
file: env_string("MVP_GGUF_FILE", DEFAULT_HF_FILE),
|
||
|
|
revision: env_optional("MVP_GGUF_REVISION"),
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
fn tokenizer_from_env() -> TokenizerSource {
|
||
|
|
env_optional("MVP_TOKENIZER_LOCAL_PATH")
|
||
|
|
.map(TokenizerSource::LocalPath)
|
||
|
|
.unwrap_or(TokenizerSource::EmbeddedGguf)
|
||
|
|
}
|