376 lines
11 KiB
Rust
Executable file
376 lines
11 KiB
Rust
Executable file
#![cfg(target_os = "linux")]
|
|
|
|
use std::io::{BufRead, BufReader, Write};
|
|
use std::process::{Child, ChildStdin, Command, Stdio};
|
|
|
|
use mvp_system::arena_manager as arena;
|
|
use serde_json::{Value, json};
|
|
|
|
const HEADER_LEN: usize = 40;
|
|
|
|
struct WorkerProcess {
|
|
child: Child,
|
|
stdin: ChildStdin,
|
|
stdout: BufReader<std::process::ChildStdout>,
|
|
}
|
|
|
|
impl WorkerProcess {
|
|
fn spawn(arena: Option<(&arena::ArenaManager, u64)>) -> Self {
|
|
let script = worker_script();
|
|
let mut command = Command::new("python3");
|
|
command
|
|
.arg(&script)
|
|
.env("MVP_TINYGRAD_TEST_MODE", "1")
|
|
.env("DEV", "CPU")
|
|
.stdin(Stdio::piped())
|
|
.stdout(Stdio::piped())
|
|
.stderr(Stdio::piped());
|
|
if let Some((arena, arena_bytes)) = arena {
|
|
command
|
|
.env("MVP_ARENA_FD", arena.arena_fd().to_string())
|
|
.env("MVP_ARENA_BYTES", arena_bytes.to_string());
|
|
}
|
|
let mut child = command
|
|
.spawn()
|
|
.unwrap_or_else(|e| panic!("spawn {}: {e}", script.display()));
|
|
let stdin = child.stdin.take().expect("worker stdin");
|
|
let stdout = BufReader::new(child.stdout.take().expect("worker stdout"));
|
|
Self {
|
|
child,
|
|
stdin,
|
|
stdout,
|
|
}
|
|
}
|
|
|
|
fn send_expect(&mut self, command: Value, expected_type: &str) -> Value {
|
|
writeln!(self.stdin, "{command}").expect("write worker command");
|
|
self.stdin.flush().expect("flush worker command");
|
|
loop {
|
|
let mut line = String::new();
|
|
let read = self.stdout.read_line(&mut line).expect("read worker event");
|
|
assert_ne!(read, 0, "worker exited before {expected_type}");
|
|
let event: Value = serde_json::from_str(line.trim_end()).expect("worker event JSON");
|
|
let actual_type = event.get("type").and_then(Value::as_str).unwrap_or("");
|
|
assert_ne!(
|
|
actual_type, "WorkerFatal",
|
|
"worker fatal while waiting for {expected_type}: {event}"
|
|
);
|
|
if actual_type == expected_type {
|
|
return event;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
impl Drop for WorkerProcess {
|
|
fn drop(&mut self) {
|
|
let _ = self.child.kill();
|
|
let _ = self.child.wait();
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_mode_tokenizer_commands_round_trip_prompt_bytes_and_visible_tokens() {
|
|
let mut worker = WorkerProcess::spawn(None);
|
|
worker.send_expect(
|
|
json!({"type":"InitializeWorker","helper_abi_version":1,"backend":{"device":"CPU"}}),
|
|
"WorkerReady",
|
|
);
|
|
worker.send_expect(
|
|
json!({
|
|
"type":"LoadWeights",
|
|
"model_id":"test-model",
|
|
"gguf_source":{"LocalPath":"/tmp/not-used-in-test-mode.gguf"},
|
|
"tokenizer":{"EmbeddedGguf":{}},
|
|
"layer_start":0,
|
|
"layer_end_exclusive":1
|
|
}),
|
|
"WeightsLoaded",
|
|
);
|
|
|
|
let encoded = worker.send_expect(
|
|
json!({"type":"EncodePrompt","request_id":7,"prompt":"Hi!"}),
|
|
"PromptEncoded",
|
|
);
|
|
assert_eq!(encoded.get("request_id").and_then(Value::as_u64), Some(7));
|
|
assert_eq!(
|
|
encoded
|
|
.get("tokens")
|
|
.and_then(Value::as_array)
|
|
.expect("encoded tokens")
|
|
.iter()
|
|
.map(|value| value.as_u64().expect("token is u64"))
|
|
.collect::<Vec<_>>(),
|
|
vec![72, 105, 33]
|
|
);
|
|
|
|
let decoded = worker.send_expect(
|
|
json!({"type":"DecodeTokens","request_id":8,"tokens":[72,105,33,6]}),
|
|
"TokensDecoded",
|
|
);
|
|
assert_eq!(decoded.get("request_id").and_then(Value::as_u64), Some(8));
|
|
assert_eq!(
|
|
decoded.get("text").and_then(Value::as_str),
|
|
Some("Hi!<tok:6>")
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_mode_worker_executes_single_and_three_stage_mo01_flow_through_real_arena() {
|
|
let arena_bytes = 16 * 1024;
|
|
let mut arena = arena::ArenaManager::boot(arena::ArenaConfig {
|
|
node_id: arena::NodeId(1),
|
|
reservation_ceiling: arena_bytes,
|
|
base_alignment: 64,
|
|
})
|
|
.expect("arena boots");
|
|
let ingress = lease_ring(&mut arena, 1, 1024, 64);
|
|
let egress = lease_ring(&mut arena, 2, 1024, 64);
|
|
let mut worker = WorkerProcess::spawn(Some((&arena, arena_bytes)));
|
|
worker.send_expect(
|
|
json!({"type":"InitializeWorker","helper_abi_version":1,"backend":{"device":"CPU"}}),
|
|
"WorkerReady",
|
|
);
|
|
|
|
let single_token = run_stage(
|
|
&mut worker,
|
|
&arena,
|
|
&ingress,
|
|
&egress,
|
|
StageFixture {
|
|
stage_index: 0,
|
|
layer_start: 0,
|
|
layer_end_exclusive: 7,
|
|
final_stage: true,
|
|
},
|
|
900,
|
|
&[2, 3],
|
|
);
|
|
assert_eq!(
|
|
single_token,
|
|
vec![6],
|
|
"N=1 stage should emit a token record"
|
|
);
|
|
|
|
let stage0 = run_stage(
|
|
&mut worker,
|
|
&arena,
|
|
&ingress,
|
|
&egress,
|
|
StageFixture {
|
|
stage_index: 0,
|
|
layer_start: 0,
|
|
layer_end_exclusive: 3,
|
|
final_stage: false,
|
|
},
|
|
901,
|
|
&[2, 3],
|
|
);
|
|
assert_eq!(stage0, vec![8], "stage 0 activation fixture word");
|
|
|
|
let stage1 = run_stage(
|
|
&mut worker,
|
|
&arena,
|
|
&ingress,
|
|
&egress,
|
|
StageFixture {
|
|
stage_index: 1,
|
|
layer_start: 3,
|
|
layer_end_exclusive: 5,
|
|
final_stage: false,
|
|
},
|
|
902,
|
|
&stage0,
|
|
);
|
|
assert_eq!(stage1, vec![17], "stage 1 activation fixture word");
|
|
|
|
let stage2 = run_stage(
|
|
&mut worker,
|
|
&arena,
|
|
&ingress,
|
|
&egress,
|
|
StageFixture {
|
|
stage_index: 2,
|
|
layer_start: 5,
|
|
layer_end_exclusive: 7,
|
|
final_stage: true,
|
|
},
|
|
903,
|
|
&stage1,
|
|
);
|
|
assert_eq!(stage2, vec![6], "final stage should emit a token record");
|
|
}
|
|
|
|
#[derive(Clone, Copy)]
|
|
struct StageFixture {
|
|
stage_index: u32,
|
|
layer_start: u32,
|
|
layer_end_exclusive: u32,
|
|
final_stage: bool,
|
|
}
|
|
|
|
fn run_stage(
|
|
worker: &mut WorkerProcess,
|
|
arena: &arena::ArenaManager,
|
|
ingress: &arena::RingLease,
|
|
egress: &arena::RingLease,
|
|
stage: StageFixture,
|
|
output_object_id: u64,
|
|
input_words: &[u32],
|
|
) -> Vec<u32> {
|
|
worker.send_expect(
|
|
json!({
|
|
"type":"ConfigureRole",
|
|
"role_id":1,
|
|
"config":{
|
|
"run_id":1,
|
|
"stage_index":stage.stage_index,
|
|
"layer_start":stage.layer_start,
|
|
"layer_end_exclusive":stage.layer_end_exclusive
|
|
}
|
|
}),
|
|
"RoleConfigured",
|
|
);
|
|
worker.send_expect(
|
|
json!({
|
|
"type":"LoadWeights",
|
|
"model_id":"test-model",
|
|
"gguf_source":{"LocalPath":"/tmp/not-used-in-test-mode.gguf"},
|
|
"tokenizer":{"EmbeddedGguf":{}},
|
|
"layer_start":stage.layer_start,
|
|
"layer_end_exclusive":stage.layer_end_exclusive
|
|
}),
|
|
"WeightsLoaded",
|
|
);
|
|
install_ring(worker, ingress, 1, 10, "ingress", 4096, 4);
|
|
install_ring(worker, egress, 2, 11, "egress", 4096, 4);
|
|
|
|
let input_payload = words_payload(input_words);
|
|
arena
|
|
.write_arena(
|
|
ingress.layout.data_offset,
|
|
&object_record(100 + u64::from(stage.stage_index), 0, 0, &input_payload),
|
|
)
|
|
.expect("write ingress record");
|
|
let loaded = worker.send_expect(json!({"type":"RingReadable","ring_id":1}), "ObjectLoaded");
|
|
let handle_id = loaded
|
|
.get("handle_id")
|
|
.and_then(Value::as_u64)
|
|
.expect("handle id");
|
|
worker.send_expect(
|
|
json!({
|
|
"type":"ExecuteStep",
|
|
"role_id":1,
|
|
"step_id":u64::from(stage.stage_index) + 1,
|
|
"input_handle_id":handle_id,
|
|
"input_object_id":100 + u64::from(stage.stage_index),
|
|
"input_sequence":0,
|
|
"output_ring_id":2,
|
|
"output_object_id":output_object_id,
|
|
"output_sequence":0,
|
|
"final_stage":stage.final_stage
|
|
}),
|
|
"StepExecuted",
|
|
);
|
|
let output = arena
|
|
.read_arena(egress.layout.data_offset, HEADER_LEN + 4)
|
|
.expect("read egress record");
|
|
decode_record_words(&output, output_object_id)
|
|
}
|
|
|
|
fn install_ring(
|
|
worker: &mut WorkerProcess,
|
|
lease: &arena::RingLease,
|
|
ring_id: u64,
|
|
edge_id: u64,
|
|
direction: &str,
|
|
max_extent: u64,
|
|
alignment: u32,
|
|
) {
|
|
worker.send_expect(
|
|
json!({
|
|
"type":"InstallRing",
|
|
"ring_id":ring_id,
|
|
"edge_id":edge_id,
|
|
"port":direction,
|
|
"direction":direction,
|
|
"layout":{
|
|
"data_offset":lease.layout.data_offset,
|
|
"data_bytes":lease.layout.data_bytes
|
|
},
|
|
"object_spec":{
|
|
"max_extent":max_extent,
|
|
"alignment":alignment
|
|
}
|
|
}),
|
|
"RingInstalled",
|
|
);
|
|
}
|
|
|
|
fn lease_ring(
|
|
arena: &mut arena::ArenaManager,
|
|
request_id: u64,
|
|
data_bytes: u64,
|
|
alignment: u64,
|
|
) -> arena::RingLease {
|
|
let events = arena.request(arena::ArenaRequest::LeaseRing(arena::LeaseRing {
|
|
request_id: arena::LeaseRequestId(request_id),
|
|
ring_spec: arena::RingSpec {
|
|
header_bytes: 64,
|
|
data_bytes,
|
|
alignment,
|
|
},
|
|
}));
|
|
match events.into_iter().next().expect("arena event") {
|
|
arena::ArenaEvent::RingLeased { lease } => lease,
|
|
event => panic!("expected ring lease, got {event:?}"),
|
|
}
|
|
}
|
|
|
|
fn object_record(object_id: u64, sequence: u64, flags: u32, payload: &[u8]) -> Vec<u8> {
|
|
let mut bytes = vec![0_u8; HEADER_LEN];
|
|
bytes[0..4].copy_from_slice(b"MO01");
|
|
bytes[4..6].copy_from_slice(&1_u16.to_le_bytes());
|
|
bytes[6..8].copy_from_slice(&(HEADER_LEN as u16).to_le_bytes());
|
|
bytes[8..16].copy_from_slice(&object_id.to_le_bytes());
|
|
bytes[16..24].copy_from_slice(&sequence.to_le_bytes());
|
|
bytes[24..32].copy_from_slice(&(payload.len() as u64).to_le_bytes());
|
|
bytes[32..36].copy_from_slice(&flags.to_le_bytes());
|
|
bytes.extend_from_slice(payload);
|
|
bytes
|
|
}
|
|
|
|
fn words_payload(words: &[u32]) -> Vec<u8> {
|
|
words
|
|
.iter()
|
|
.flat_map(|word| word.to_le_bytes())
|
|
.collect::<Vec<_>>()
|
|
}
|
|
|
|
fn decode_record_words(record: &[u8], expected_object_id: u64) -> Vec<u32> {
|
|
assert!(record.len() >= HEADER_LEN);
|
|
assert_eq!(&record[0..4], b"MO01");
|
|
assert_eq!(u16::from_le_bytes(record[4..6].try_into().unwrap()), 1);
|
|
assert_eq!(
|
|
u16::from_le_bytes(record[6..8].try_into().unwrap()),
|
|
HEADER_LEN as u16
|
|
);
|
|
assert_eq!(
|
|
u64::from_le_bytes(record[8..16].try_into().unwrap()),
|
|
expected_object_id
|
|
);
|
|
let extent = u64::from_le_bytes(record[24..32].try_into().unwrap()) as usize;
|
|
assert_eq!(extent % 4, 0);
|
|
record[HEADER_LEN..HEADER_LEN + extent]
|
|
.chunks_exact(4)
|
|
.map(|chunk| u32::from_le_bytes(chunk.try_into().unwrap()))
|
|
.collect()
|
|
}
|
|
|
|
fn worker_script() -> std::path::PathBuf {
|
|
std::path::PathBuf::from(env!("CARGO_MANIFEST_DIR"))
|
|
.join("../..")
|
|
.join("apps/mvp-node/tinygrad_worker.py")
|
|
.canonicalize()
|
|
.expect("tinygrad worker script exists")
|
|
}
|