zacheryasc/swactor
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

feat: mvp-chat benchmarking

Browse source 
Add end-to-end timing instrumentation and an xtask benchmark report for mvp-chat runs.

- benchmark_observability: add a shared stamping module — stamp(component) emitting schema/pid/monotonic+wall ms from a process-global start and sequence counter, plus unix_ms_now() — stamped onto every mvp-chat/orchestrator/worker-node event and frame-archive record
- mvp-chat: thread a run_id (new --run-id, defaults to 1) through config and the orchestrator CLI, add per-phase started/ready/failed emits for ensure_orch_binary/ensure_worker_binary/prepare_node_image, and a prompt_complete record carrying tokens_generated/elapsed_ms/final_text bytes
- orchestrator/worker-node: stamp bootstrap and prompt events, add arrival_unix_ms to archived frames, propagate MVP_RUN_ID/MVP_LOGICAL_NODE_ID/MVP_STAGE_INDEX into the tinygrad worker, default the device to CPU for the process provider, emit a prompt_rpc started span, and drop the MVP_TINYGRAD_TEST_MODE passthrough
- tinygrad_worker.py: stamp every control() event and tag it with run/node/stage env, add a CPU:X86 fallback when clang is absent, and remove the test_mode() short-circuits
- xtask: replace the flat dump-log fact assertions with a benchmark report builder (build_benchmark_report) that requires named spans (prepare_runtime, ensure_*_binary, weights_loaded, prompt_rpc) and emits per-prompt first-token/decode/tokens-per-second latency; wrap the cargo run in XtaskBenchmark synthetic frames and pass a unix-ms --run-id

Signed-off-by: Zachery Aaron Shores-Chmielewski <zacheryasc@gmail.com>
This commit is contained in:
Zachery Aaron Shores-Chmielewski 2026-07-23 09:35:04 +04:00
parent 841a2de911
commit 504c2d13ad
8 changed files with 2095 additions and 506 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

195
apps/mvp-node/tinygrad_worker.py
View file

@ -29,6 +29,9 @@ device_objects: dict[int, dict[str, Any]] = {}
next_handle = 42
HEADER_LEN = 40
WORKER_GENERATION = 1
BENCHMARK_SCHEMA = 1
_benchmark_start = time.monotonic()
_benchmark_seq = 0
class CpuLineSampler:
@ -142,9 +145,37 @@ def env_flag(name: str, default: bool = True) -> bool:
return raw.strip().lower() not in {"0", "false", "no", "off"}
def benchmark_stamp() -> dict[str, Any]:
global _benchmark_seq
_benchmark_seq += 1
return {
"schema": BENCHMARK_SCHEMA,
"component": "tinygrad-worker",
"pid": os.getpid(),
"seq": _benchmark_seq,
"wall_unix_ms": time.time_ns() // 1_000_000,
"mono_ms": int((time.monotonic() - _benchmark_start) * 1000),
}
def env_int(name: str) -> int | None:
raw = os.environ.get(name)
if raw is None:
return None
try:
return int(raw)
except ValueError:
return None
def control(**event: Any) -> None:
event.setdefault("benchmark", benchmark_stamp())
if (run_id := env_int("MVP_RUN_ID")) is not None:
event.setdefault("run_id", run_id)
if (node_id := env_int("MVP_LOGICAL_NODE_ID")) is not None:
event.setdefault("node_id", node_id)
if (stage_index := env_int("MVP_STAGE_INDEX")) is not None:
event.setdefault("stage_index", stage_index)
print(json.dumps(event, separators=(",", ":")), flush=True)
@ -157,8 +188,6 @@ def fatal(reason: str, **fields: Any) -> None:
raise SystemExit(1)
def test_mode() -> bool:
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":
@ -170,31 +199,31 @@ def configure_tinygrad_cuda_compiler(device: str) -> None:
os.environ["CUDA_PTX"] = "1"
control(type="TinygradCudaCompilerSelected", requested_device=device, compiler="PTX", reason="nvcc_not_found")
def select_tinygrad_device(device: str) -> str:
device_kind = device.split(":", 1)[0].upper()
if device_kind == "CPU" and ":" not in device and shutil.which("clang") is None:
selected = "CPU:X86"
os.environ["DEV"] = selected
control(type="TinygradCpuCompilerSelected", requested_device=device, selected_device=selected, compiler="X86", reason="clang_not_found")
return selected
os.environ["DEV"] = device
configure_tinygrad_cuda_compiler(device)
return device
def initialize(cmd: dict[str, Any]) -> None:
global Tensor, dtypes, arena
if int(cmd.get("helper_abi_version", 1)) != 1:
fatal("UnsupportedHelperAbi", helper_abi_version=cmd.get("helper_abi_version"))
device = str(cmd.get("backend", {}).get("device") or os.environ.get("DEV") or "CUDA")
os.environ["DEV"] = device
requested_device = str(cmd.get("backend", {}).get("device") or os.environ.get("DEV") or "CUDA")
device = select_tinygrad_device(requested_device)
arena_fd = os.environ.get("MVP_ARENA_FD")
if arena_fd is not None:
arena_bytes = int(os.environ.get("MVP_ARENA_BYTES", "0") or "0")
if arena_bytes > 0:
arena = mmap.mmap(int(arena_fd), arena_bytes)
started = time.monotonic()
if test_mode():
control(
type="WorkerReady",
pid=os.getpid(),
backend={"requested_device": device, "env_DEV": os.environ.get("DEV"), "test_mode": True},
cuda_probe=[],
test_mode=True,
elapsed_ms=int((time.monotonic() - started) * 1000),
)
return
configure_tinygrad_cuda_compiler(device)
control(type="TinygradImportStarted", requested_device=device, env_DEV=os.environ.get("DEV"))
from tinygrad import Tensor as TinyTensor, dtypes as tiny_dtypes
@ -353,31 +382,42 @@ class PipelineStageTinygradModel:
vocab_size: int,
head_dim: int,
rope_theta: float,
rope_dim: int,
v_head_dim: int,
max_context: int,
qk_norm: int,
num_experts: int,
num_experts_per_tok: int,
norm_topk_prob: bool,
qkv_bias: bool,
expert_bias: bool,
first_stage: bool,
final_stage: bool,
nn_mod: Any,
config_cls: Any,
block_cls: Any,
) -> None:
self.blk = [
block_cls(
dim,
hidden_dim,
n_heads,
n_kv_heads,
norm_eps,
head_dim,
rope_theta,
max_context,
qk_norm,
num_experts,
num_experts_per_tok,
block_config = config_cls(
num_blocks=block_count,
dim=dim,
hidden_dim=hidden_dim,
n_heads=n_heads,
n_kv_heads=n_kv_heads,
norm_eps=norm_eps,
vocab_size=vocab_size,
head_dim=head_dim,
rope_theta=rope_theta,
rope_dim=rope_dim,
v_head_dim=v_head_dim,
max_context=max_context,
qk_norm=qk_norm,
num_experts=num_experts,
num_experts_per_tok=num_experts_per_tok,
norm_topk_prob=norm_topk_prob,
qkv_bias=qkv_bias,
expert_bias=expert_bias,
)
for _ in range(block_count)
]
self.blk = [block_cls(block_config) for _ in range(block_count)]
self.max_context = max_context
self.hidden_dim = dim
self.first_stage = first_stage
@ -389,16 +429,18 @@ class PipelineStageTinygradModel:
self.output = nn_mod.Linear(dim, vocab_size, bias=False)
def token_hidden(self, tokens_tensor: Any) -> Any:
return self.token_embd(tokens_tensor)
return self.token_embd(tokens_tensor).float()
def forward_hidden(self, hidden: Any, start_pos: int) -> Any:
def forward_hidden(self, hidden: Any, start_pos: Any) -> Any:
for block in self.blk:
hidden = block(hidden, start_pos)
return hidden.contiguous()
def next_token(self, hidden: Any) -> Any:
return self.output(self.output_norm(hidden))[:, -1, :].softmax(-1, dtype="float").argmax(-1, keepdim=True)
return self.output(self.output_norm(hidden))[:, -1, :].argmax(-1, keepdim=True)
def __call__(self, tokens_tensor: Any, start_pos: Any) -> Any:
return self.next_token(self.forward_hidden(self.token_hidden(tokens_tensor), start_pos))
def remap_stage_state_dict(
state_dict: dict[str, Any],
@ -436,42 +478,66 @@ def load_pipeline_stage_model(
) -> tuple[PipelineStageTinygradModel, dict[str, Any]]:
TensorCls = require_tinygrad()
from tinygrad import nn
from tinygrad.apps.llm import TransformerBlock
from tinygrad.llm.gguf import gguf_load
from tinygrad.llm.model import TransformerBlock, TransformerConfig
kv, state_dict = nn.state.gguf_load(TensorCls(path).to(None))
kv, state_dict = gguf_load(path)
state_dict = {key: value.cast("float16") if env_flag("HALF", True) else value for key, value in state_dict.items()}
if "output.weight" not in state_dict and "token_embd.weight" in state_dict:
state_dict["output.weight"] = state_dict["token_embd.weight"]
arch = kv["general.architecture"]
max_context = min(max_context, int(kv[f"{arch}.context_length"]))
n_heads = int(kv[f"{arch}.attention.head_count"])
n_kv_heads = int(kv[f"{arch}.attention.head_count_kv"])
if arch == "llama":
dim = int(kv[f"{arch}.embedding_length"])
kv_lora_rank = int(kv.get(f"{arch}.attention.kv_lora_rank", 0))
head_dim = int(kv.get(f"{arch}.attention.key_length_mla", kv.get(f"{arch}.attention.key_length", dim // n_heads)))
rope_dim = int(kv.get(f"{arch}.rope.dimension_count", head_dim))
for name in list(state_dict):
if "attn_q.weight" in name:
state_dict[name] = state_dict[name].rearrange("(n h two) d -> (n two h) d", n=n_heads, two=2)
if "attn_k.weight" in name:
state_dict[name] = state_dict[name].rearrange("(n h two) d -> (n two h) d", n=n_kv_heads, two=2)
total_layers = int(kv[f"{arch}.block_count"])
if ("attn_q.weight" in name or "attn_q_b.weight" in name) and (arch == "llama" or kv_lora_rank):
weight = state_dict[name].reshape(n_heads, state_dict[name].shape[0] // n_heads, -1)
prefix = head_dim - rope_dim
state_dict[name] = (
weight[:, :prefix]
.cat(weight[:, prefix:].rearrange("n (h two) d -> n (two h) d", two=2), dim=1)
.reshape(-1, weight.shape[-1])
)
elif arch == "llama" and "attn_k.weight" in name:
weight = state_dict[name].reshape(n_kv_heads, state_dict[name].shape[0] // n_kv_heads, -1)
state_dict[name] = weight.rearrange("n (h two) d -> n (two h) d", two=2).reshape(-1, weight.shape[-1])
elif kv_lora_rank and "attn_kv_a_mqa.weight" in name:
state_dict[name] = state_dict[name][:kv_lora_rank].cat(
state_dict[name][kv_lora_rank:].rearrange("(h two) d -> (two h) d", two=2),
dim=0,
)
total_layers = int(kv[f"{arch}.block_count"]) - int(kv.get(f"{arch}.nextn_predict_layers", 0))
first_stage = layer_start == 0
final_stage = layer_end_exclusive >= total_layers
qk_key = f"blk.{layer_start}.attn_q_norm.weight"
qk_norm = int(state_dict[qk_key].shape[0]) if qk_key in state_dict else 0
stage_model = PipelineStageTinygradModel(
block_count=layer_end_exclusive - layer_start,
dim=int(kv[f"{arch}.embedding_length"]),
hidden_dim=int(kv.get(f"{arch}.expert_feed_forward_length", kv[f"{arch}.feed_forward_length"])),
dim=dim,
hidden_dim=int(kv.get(f"{arch}.expert_feed_forward_length", kv.get(f"{arch}.feed_forward_length", 0))),
n_heads=n_heads,
n_kv_heads=n_kv_heads,
norm_eps=float(kv[f"{arch}.attention.layer_norm_rms_epsilon"]),
vocab_size=len(kv["tokenizer.ggml.tokens"]),
head_dim=int(kv.get(f"{arch}.attention.key_length", int(kv[f"{arch}.embedding_length"]) // n_heads)),
head_dim=head_dim,
rope_theta=float(kv[f"{arch}.rope.freq_base"]),
rope_dim=rope_dim,
v_head_dim=int(kv.get(f"{arch}.attention.value_length_mla", kv.get(f"{arch}.attention.value_length", head_dim))),
max_context=max_context,
qk_norm=qk_norm,
num_experts=int(kv.get(f"{arch}.expert_count", 0)),
num_experts_per_tok=int(kv.get(f"{arch}.expert_used_count", 0)),
norm_topk_prob=bool(kv.get(f"{arch}.expert_weights_norm", arch in ("qwen3moe", "qwen35moe"))),
qkv_bias="blk.0.attn_q.bias" in state_dict,
expert_bias=f"blk.{int(kv.get(f'{arch}.leading_dense_block_count', 0))}.exp_probs_b.bias" in state_dict,
first_stage=first_stage,
final_stage=final_stage,
nn_mod=nn,
config_cls=TransformerConfig,
block_cls=TransformerBlock,
)
stage_state = remap_stage_state_dict(
@ -501,24 +567,6 @@ def load_weights(cmd: dict[str, Any]) -> None:
layer_start=int(cmd.get("layer_start", 0)),
layer_end_exclusive=int(cmd.get("layer_end_exclusive", 0)),
)
if test_mode():
model = {"test_mode": True}
tokenizer = {"test_mode": True}
loaded.clear()
loaded.update(
model_id=model_id,
path="mvp-tinygrad-test-mode",
layer_start=int(cmd.get("layer_start", 0)),
layer_end_exclusive=int(cmd.get("layer_end_exclusive", 0)),
)
control(
type="WeightsLoaded",
model_id=model_id,
path=loaded["path"],
test_mode=True,
elapsed_ms=int((time.monotonic() - started) * 1000),
)
return
control(type="GgufResolveStarted", model_id=model_id, source_kind=source_kind(source))
path = fetch_whole(source)
model_bytes = path.stat().st_size
@ -527,7 +575,7 @@ def load_weights(cmd: dict[str, Any]) -> None:
layer_end_exclusive = int(cmd.get("layer_end_exclusive", 0))
try:
control(type="TinygradLlmImportStarted", model_id=model_id)
from tinygrad.apps.llm import SimpleTokenizer
from tinygrad.llm.cli import SimpleTokenizer
control(type="TinygradLlmImportReady", model_id=model_id)
max_context_raw = os.environ.get("MVP_MAX_CONTEXT", "512")
@ -739,19 +787,6 @@ def infer_prompt(cmd: dict[str, Any]) -> None:
prompt_chars=len(prompt),
max_tokens=max_tokens,
)
if test_mode():
text = f"mvp-test response: {prompt}"
control(
type="PromptCompleted",
request_id=request_id,
model_id=loaded.get("model_id"),
prompt_tokens=[],
generated_tokens=list(range(min(max_tokens, 3))),
text=text,
test_mode=True,
elapsed_ms=int((time.monotonic() - started) * 1000),
)
return
model_prompt, prompt_template = model_prompt_text(prompt)
control(
type="PromptEncodeStarted",
@ -907,7 +942,7 @@ def object_start_pos(sequence: int, token_count: int) -> int:
def materialize_object(payload: bytes, sequence: int) -> dict[str, Any]:
if test_mode() or not isinstance(model, PipelineStageTinygradModel):
if not isinstance(model, PipelineStageTinygradModel):
return {
"kind": "words",
"words": payload_words(payload),
@ -1013,7 +1048,7 @@ def execute_step(cmd: dict[str, Any]) -> None:
if ring["direction"] != "egress":
fatal("WrongRingDirection", ring_id=output_ring_id, direction=ring["direction"])
final_stage = bool(cmd.get("final_stage"))
if test_mode() or not isinstance(model, PipelineStageTinygradModel):
if not isinstance(model, PipelineStageTinygradModel):
if final_stage:
base = sum(int(word) for word in obj["words"]) + int(role.get("stage_index", 0))
token = 6 if base % 2 else 8
@ -1067,7 +1102,7 @@ def release_device_object(cmd: dict[str, Any]) -> None:
def encode_prompt(cmd: dict[str, Any]) -> None:
prompt = str(cmd.get("prompt", ""))
if tokenizer is not None and not test_mode():
if tokenizer is not None:
model_prompt, _ = model_prompt_text(prompt)
tokens = [int(token) for token in tokenizer.encode(model_prompt)]
else:
@ -1077,7 +1112,7 @@ def encode_prompt(cmd: dict[str, Any]) -> None:
def decode_tokens(cmd: dict[str, Any]) -> None:
tokens = [int(token) for token in cmd.get("tokens", [])]
if tokenizer is not None and not test_mode():
if tokenizer is not None:
text = strip_chat_stop_markers(tokenizer.decode(tokens))
else:
text = "".join(chr(token) if 32 <= token <= 126 else f"<tok:{token}>" for token in tokens)

32
crates/mvp-system/src/benchmark_observability.rs Normal file
View file

@ -0,0 +1,32 @@
use std::sync::OnceLock;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::{Instant, SystemTime, UNIX_EPOCH};
use serde_json::{Value, json};
pub const BENCHMARK_SCHEMA: u64 = 1;
static BENCHMARK_START: OnceLock<Instant> = OnceLock::new();
static BENCHMARK_SEQ: AtomicU64 = AtomicU64::new(1);
pub fn unix_ms_now() -> u64 {
let millis = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
.as_millis();
u64::try_from(millis).unwrap_or(u64::MAX)
}
pub fn stamp(component: &'static str) -> Value {
let start = BENCHMARK_START.get_or_init(Instant::now);
let mono_ms = u64::try_from(start.elapsed().as_millis()).unwrap_or(u64::MAX);
let seq = BENCHMARK_SEQ.fetch_add(1, Ordering::Relaxed);
json!({
"schema": BENCHMARK_SCHEMA,
"component": component,
"pid": std::process::id(),
"seq": seq,
"wall_unix_ms": unix_ms_now(),
"mono_ms": mono_ms,
})
}

332
crates/mvp-system/src/bin/mvp_chat.rs
View file

@ -22,6 +22,7 @@ use signal_hook::consts::signal::{SIGINT, SIGTERM};
#[cfg(target_os = "linux")]
use signal_hook::iterator::Signals;
use mvp_system::benchmark_observability;
use mvp_system::config as chat_config;
use mvp_system::config::ResolvedVastAiConfig;
use mvp_system::node_image::{
@ -73,7 +74,7 @@ where
I: IntoIterator<Item = String>,
{
let config = Config::from_args(args)?;
let mut progress = ChatDatastream::new(1, config.datastream_frame_log.clone())?;
let mut progress = ChatDatastream::new(config.run_id, config.datastream_frame_log.clone())?;
progress.emit(
CHAT_LIFECYCLE_CHANNEL,
"config",
@ -87,7 +88,14 @@ where
}),
);
confirm_vastai_if_needed(&config)?;
let image_ref = match prepare_runtime(&config) {
progress.emit(
CHAT_RUNTIME_CHANNEL,
"prepare_runtime",
"started",
json!({"provider": config.provider.as_str()}),
);
let image_ref =
match prepare_runtime_with_progress(&config, prepare_node_image, Some(&mut progress)) {
Ok(image_ref) => {
progress.emit(
CHAT_RUNTIME_CHANNEL,
@ -108,8 +116,20 @@ where
return Err(error);
}
};
progress.emit(
CHAT_COMPONENT_CHANNEL,
"orchestrator_process_spawn",
"started",
json!({"binary": config.orch_bin.to_string_lossy()}),
);
let mut orch = match OrchChild::spawn(&config, &image_ref) {
Ok(orch) => {
progress.emit(
CHAT_COMPONENT_CHANNEL,
"orchestrator_process_spawn",
"ready",
json!({"binary": config.orch_bin.to_string_lossy(), "pid": orch.child.id()}),
);
progress.emit(
CHAT_COMPONENT_CHANNEL,
"orchestrator_process",
@ -119,6 +139,12 @@ where
orch
}
Err(error) => {
progress.emit(
CHAT_COMPONENT_CHANNEL,
"orchestrator_process_spawn",
"failed",
json!({"binary": config.orch_bin.to_string_lossy(), "error": error}),
);
progress.emit(
CHAT_COMPONENT_CHANNEL,
"orchestrator_process",
@ -129,8 +155,20 @@ where
return Err(error);
}
};
progress.emit(
CHAT_RUNTIME_CHANNEL,
"prompt_rpc_wait",
"started",
json!({"addr": config.rpc_addr}),
);
let rpc_addr = match orch.wait_ready(config.rpc_addr.clone()) {
Ok(addr) => {
progress.emit(
CHAT_RUNTIME_CHANNEL,
"prompt_rpc_wait",
"ready",
json!({"addr": addr}),
);
progress.emit(
CHAT_RUNTIME_CHANNEL,
"prompt_rpc",
@ -139,7 +177,13 @@ where
);
addr
}
Err(_) if STOP_REQUESTED.load(Ordering::SeqCst) => {
Err(error) if STOP_REQUESTED.load(Ordering::SeqCst) => {
progress.emit(
CHAT_RUNTIME_CHANNEL,
"prompt_rpc_wait",
"failed",
json!({"addr": config.rpc_addr, "error": error}),
);
progress.emit(
CHAT_LIFECYCLE_CHANNEL,
"shutdown",
@ -157,6 +201,12 @@ where
return Ok(());
}
Err(error) => {
progress.emit(
CHAT_RUNTIME_CHANNEL,
"prompt_rpc_wait",
"failed",
json!({"addr": config.rpc_addr, "error": error}),
);
progress.emit(
CHAT_RUNTIME_CHANNEL,
"prompt_rpc",
@ -201,6 +251,7 @@ struct Config {
image_tag: Option<String>,
cached_model: Option<CachedModelConfig>,
datastream_frame_log: Option<PathBuf>,
run_id: u64,
vastai_yes: bool,
vastai: Option<ResolvedVastAiConfig>,
pipeline_stages: u32,
@ -210,6 +261,7 @@ struct Config {
struct ChatDatastream {
stream: StreamId,
run_id: u64,
endpoint: DatastreamEndpoint,
producer: DatastreamProducer,
channels: BTreeMap<String, ChannelId>,
@ -233,6 +285,7 @@ impl ChatDatastream {
let producer = endpoint.producer();
let mut out = Self {
stream,
run_id,
endpoint,
producer,
channels: BTreeMap::new(),
@ -272,6 +325,8 @@ impl ChatDatastream {
"type": "ChatProgress",
"phase": phase,
"status": status,
"run_id": self.run_id,
"benchmark": benchmark_observability::stamp("mvp-chat"),
"detail": detail,
}))
.expect("serialize mvp-chat progress event");
@ -356,6 +411,7 @@ impl ChatFrameArchive {
};
let record = json!({
"arrival_seq": self.next_seq,
"arrival_unix_ms": benchmark_observability::unix_ms_now(),
"source": source,
"stream": stream.to_string(),
"channel": channel,
@ -514,6 +570,7 @@ impl Config {
image_tag: first_non_empty([toml.image.tag.clone()]),
cached_model,
datastream_frame_log,
run_id: args.run_id.unwrap_or(1),
vastai_yes: args.vastai_yes,
pipeline_stages,
max_tokens,
@ -534,6 +591,8 @@ impl Config {
self.rpc_addr.clone(),
"--max-tokens".to_owned(),
self.max_tokens.to_string(),
"--run-id".to_owned(),
self.run_id.to_string(),
"--pipeline-stages".to_owned(),
self.pipeline_stages.to_string(),
"--no-dashboard".to_owned(),
@ -617,6 +676,7 @@ struct ParsedArgs {
pipeline_stages: Option<u32>,
dump_logs: bool,
dump_log_path: Option<PathBuf>,
run_id: Option<u64>,
skip_rebuild: bool,
cached_model: Option<CachedModelSource>,
}
@ -658,6 +718,13 @@ impl ParsedArgs {
parsed.pipeline_stages =
Some(parse_pipeline_stages_value(&mut args, arg.as_str())?)
}
"--run-id" => {
let run_id: u64 = parse_next(&mut args, "--run-id")?;
if run_id == 0 {
return Err("--run-id must be greater than 0".to_owned());
}
parsed.run_id = Some(run_id);
}
"--dump-logs" => {
parsed.dump_logs = true;
}
@ -916,33 +983,193 @@ fn signal_orch_process_group(child: &Child, signal: libc::c_int) -> io::Result<(
type PrepareNodeImageFn = fn(NodeImageRequest) -> Result<PreparedNodeImage, String>;
#[allow(dead_code)]
fn prepare_runtime(config: &Config) -> Result<String, String> {
prepare_runtime_with(config, prepare_node_image)
}
#[allow(dead_code)]
fn prepare_runtime_with(
config: &Config,
prepare_node_image_fn: PrepareNodeImageFn,
) -> Result<String, String> {
ensure_orch_binary(config)?;
prepare_runtime_with_progress(config, prepare_node_image_fn, None)
}
fn prepare_runtime_with_progress(
config: &Config,
prepare_node_image_fn: PrepareNodeImageFn,
progress: Option<&mut ChatDatastream>,
) -> Result<String, String> {
let mut progress = progress;
let binary_mode = if config.skip_rebuild {
"existing_artifact"
} else {
"cargo_build"
};
emit_chat_progress(
&mut progress,
CHAT_RUNTIME_CHANNEL,
"ensure_orch_binary",
"started",
json!({"mode": binary_mode}),
);
match ensure_orch_binary(config) {
Ok(()) => emit_chat_progress(
&mut progress,
CHAT_RUNTIME_CHANNEL,
"ensure_orch_binary",
"ready",
json!({"mode": binary_mode}),
),
Err(error) => {
emit_chat_progress(
&mut progress,
CHAT_RUNTIME_CHANNEL,
"ensure_orch_binary",
"failed",
json!({"mode": binary_mode, "error": error.as_str()}),
);
return Err(error);
}
}
if config.provider == ProviderKind::Process {
ensure_worker_binary(config)?;
emit_chat_progress(
&mut progress,
CHAT_RUNTIME_CHANNEL,
"ensure_worker_binary",
"started",
json!({"mode": binary_mode}),
);
match ensure_worker_binary(config) {
Ok(()) => emit_chat_progress(
&mut progress,
CHAT_RUNTIME_CHANNEL,
"ensure_worker_binary",
"ready",
json!({"mode": binary_mode}),
),
Err(error) => {
emit_chat_progress(
&mut progress,
CHAT_RUNTIME_CHANNEL,
"ensure_worker_binary",
"failed",
json!({"mode": binary_mode, "error": error.as_str()}),
);
return Err(error);
}
}
emit_chat_progress(
&mut progress,
CHAT_RUNTIME_CHANNEL,
"prepare_node_image",
"skipped",
json!({"provider": config.provider.as_str(), "reason": "process_provider"}),
);
return Ok(config.node_image.clone());
}
if config.skip_rebuild {
ensure_worker_binary(config)?;
emit_chat_progress(
&mut progress,
CHAT_RUNTIME_CHANNEL,
"ensure_worker_binary",
"started",
json!({"mode": binary_mode}),
);
match ensure_worker_binary(config) {
Ok(()) => emit_chat_progress(
&mut progress,
CHAT_RUNTIME_CHANNEL,
"ensure_worker_binary",
"ready",
json!({"mode": binary_mode}),
),
Err(error) => {
emit_chat_progress(
&mut progress,
CHAT_RUNTIME_CHANNEL,
"ensure_worker_binary",
"failed",
json!({"mode": binary_mode, "error": error.as_str()}),
);
return Err(error);
}
}
emit_chat_progress(
&mut progress,
CHAT_RUNTIME_CHANNEL,
"prepare_node_image",
"skipped",
json!({"provider": config.provider.as_str(), "reason": "skip_rebuild"}),
);
return Ok(config.node_image.clone());
}
let prepared = prepare_node_image_fn(NodeImageRequest {
emit_chat_progress(
&mut progress,
CHAT_RUNTIME_CHANNEL,
"prepare_node_image",
"started",
json!({"provider": config.provider.as_str()}),
);
let node_bin = match node_bin_for_current_profile() {
Ok(path) => path,
Err(error) => {
emit_chat_progress(
&mut progress,
CHAT_RUNTIME_CHANNEL,
"prepare_node_image",
"failed",
json!({"provider": config.provider.as_str(), "error": error.as_str()}),
);
return Err(error);
}
};
let provider = match node_image_provider(config.provider) {
Ok(provider) => provider,
Err(error) => {
emit_chat_progress(
&mut progress,
CHAT_RUNTIME_CHANNEL,
"prepare_node_image",
"failed",
json!({"provider": config.provider.as_str(), "error": error.as_str()}),
);
return Err(error);
}
};
let prepared = match prepare_node_image_fn(NodeImageRequest {
requested_image: config.node_image.clone(),
base_image: BASE_NODE_IMAGE.to_owned(),
node_bin: node_bin_for_current_profile()?,
provider: node_image_provider(config.provider)?,
node_bin,
provider,
extra_tag: config.image_tag.clone(),
push: false,
force_refresh: false,
enabled: true,
})?;
}) {
Ok(prepared) => prepared,
Err(error) => {
emit_chat_progress(
&mut progress,
CHAT_RUNTIME_CHANNEL,
"prepare_node_image",
"failed",
json!({"provider": config.provider.as_str(), "error": error.as_str()}),
);
return Err(error);
}
};
emit_chat_progress(
&mut progress,
CHAT_RUNTIME_CHANNEL,
"prepare_node_image",
"ready",
json!({"provider": config.provider.as_str(), "image_ref": prepared.image_ref}),
);
Ok(prepared.image_ref)
}
@ -1254,7 +1481,12 @@ where
json!({"request_id": request_id, "text_bytes": text.len()}),
);
}
PromptEvent::Done { .. } => {
PromptEvent::Done {
final_text,
tokens_generated,
elapsed_ms,
..
} => {
if response_started {
writeln!(output).map_err(|e| format!("write response terminator: {e}"))?;
} else {
@ -1266,7 +1498,13 @@ where
CHAT_PROMPT_CHANNEL,
"request_completed",
"ready",
json!({"request_id": request_id, "response_started": response_started}),
json!({
"request_id": request_id,
"response_started": response_started,
"tokens_generated": tokens_generated,
"elapsed_ms": elapsed_ms,
"final_text_bytes": final_text.len(),
}),
);
break;
}
@ -1642,6 +1880,7 @@ mod tests {
image_tag: None,
cached_model: None,
datastream_frame_log: None,
run_id: 1,
vastai_yes: false,
vastai: None,
pipeline_stages: 1,
@ -1704,6 +1943,48 @@ mod tests {
.collect()
}
#[test]
fn benchmark_observability_chat_progress_records_include_run_id_and_stamp() {
let temp = TempDir::new("chat-progress-archive");
let archive_path = temp.path().join("frames.ndjson");
let mut progress = ChatDatastream::new(77, Some(archive_path.clone()))
.expect("chat datastream constructs");
progress.emit(
CHAT_RUNTIME_CHANNEL,
"unit_phase",
"ready",
serde_json::json!({"ok": true}),
);
progress.archive_pending().expect("archive pending frames");
let archive = fs::read_to_string(&archive_path).expect("read archive");
let line = archive.lines().next().expect("archive line");
let outer: serde_json::Value = serde_json::from_str(line).expect("outer archive JSON");
let inner_text = outer
.get("payload")
.and_then(|payload| payload.get("value"))
.and_then(serde_json::Value::as_str)
.expect("inner event text");
let inner: serde_json::Value = serde_json::from_str(inner_text).expect("inner event JSON");
assert_eq!(
inner.get("type").and_then(serde_json::Value::as_str),
Some("ChatProgress")
);
assert_eq!(
inner.get("run_id").and_then(serde_json::Value::as_u64),
Some(77)
);
assert_eq!(
inner
.get("benchmark")
.and_then(|benchmark| benchmark.get("schema"))
.and_then(serde_json::Value::as_u64),
Some(1)
);
}
#[test]
fn parsed_args_accepts_public_flags() {
let parsed = ParsedArgs::parse(strings(&[
@ -1729,6 +2010,32 @@ mod tests {
assert!(parsed.skip_rebuild);
}
#[test]
fn benchmark_observability_parsed_args_accepts_run_id_and_forwards_to_orchestrator() {
let parsed = ParsedArgs::parse(strings(&["--run-id", "123"])).expect("run id parses");
assert_eq!(parsed.run_id, Some(123));
let temp = TempDir::new("run-id-config");
with_process_state(&[], Some(temp.path()), || {
let config =
Config::from_args(strings(&["--run-id", "123"])).expect("config resolves run id");
assert_eq!(config.run_id, 123);
let args = config.orchestrator_cli_args("resolved-image");
let run_id_arg = args
.windows(2)
.find(|pair| pair[0] == "--run-id")
.map(|pair| pair[1].as_str());
assert_eq!(run_id_arg, Some("123"), "{args:?}");
});
}
#[test]
fn benchmark_observability_parsed_args_rejects_zero_run_id() {
let error =
ParsedArgs::parse(strings(&["--run-id", "0"])).expect_err("zero run id should fail");
assert_eq!(error, "--run-id must be greater than 0");
}
#[test]
fn parsed_args_accepts_cached_model_path() {
let parsed = ParsedArgs::parse(strings(&["--cached-model=/tmp/model.gguf"]))
@ -1795,6 +2102,7 @@ mod tests {
let defaults = Config::from_args(Vec::<String>::new()).expect("defaults resolve");
assert_eq!(defaults.provider, ProviderKind::Process);
assert_eq!(defaults.pipeline_stages, 1);
assert_eq!(defaults.run_id, 1);
assert!(defaults.datastream_frame_log.is_none());
assert!(defaults.cached_model.is_none());
assert!(defaults.vastai.is_none());

87
crates/mvp-system/src/bin/orchestrator.rs
View file

@ -27,6 +27,7 @@ use mvp_system::actors::node_agent::{
};
use mvp_system::actors::orchestrator::{OrchestratorActor, OrchestratorReport};
use mvp_system::actors::register_mvp_actor_codecs;
use mvp_system::benchmark_observability;
use mvp_system::config::{DEFAULT_CONFIG_PATH, TomlConfigOverlay};
#[cfg(feature = "dashboard")]
use mvp_system::dashboard_view::MvpClusterDashboardView;
@ -415,6 +416,18 @@ fn run() -> Result<(), String> {
orchestrator_actor,
)?;
orch_datastream.emit_bootstrap(
dashboard.as_ref(),
config.run_id,
config.node_id,
"prompt_rpc",
"started",
json!({
"bind":config.rpc_bind.to_string(),
"default_max_tokens":config.default_max_tokens,
}),
);
let rpc_addr = match spawn_prompt_rpc(config.rpc_bind, work_tx, config.default_max_tokens) {
Ok(addr) => {
orch_datastream.emit_bootstrap(
@ -1700,9 +1713,6 @@ impl Config {
if self.provider == ProviderKind::Docker {
keys.push("MVP_DOCKER_GPUS");
}
if std::env::var_os("MVP_TINYGRAD_TEST_MODE").is_some() {
keys.push("MVP_TINYGRAD_TEST_MODE");
}
if local_tinygrad_worker_env(self.provider).is_some() {
keys.push("MVP_TINYGRAD_WORKER");
}
@ -1800,7 +1810,6 @@ impl Config {
if self.provider == ProviderKind::Docker {
env.push(("MVP_DOCKER_GPUS".to_owned(), self.docker_gpus.clone()));
}
env.extend(optional_env("MVP_TINYGRAD_TEST_MODE"));
env.extend(local_tinygrad_worker_env(self.provider));
env.extend(optional_env("MVP_CPU_LINE_PROFILE"));
env.extend(optional_env("MVP_CPU_LINE_PROFILE_INTERVAL_MS"));
@ -3132,6 +3141,7 @@ impl FrameArchive {
};
let record = json!({
"arrival_seq":self.next_seq,
"arrival_unix_ms":benchmark_observability::unix_ms_now(),
"source":source,
"stream":stream.to_string(),
"channel":channel,
@ -3252,6 +3262,7 @@ impl OrchDatastream {
"status":status,
"run_id":run_id,
"node_id":node_id,
"benchmark":benchmark_observability::stamp("mvp-orchestrator"),
"detail":detail,
}))
.expect("serialize orch bootstrap event");
@ -3275,6 +3286,7 @@ impl OrchDatastream {
"run_id":run_id,
"node_id":node_id,
"request_id":request_id,
"benchmark":benchmark_observability::stamp("mvp-orchestrator"),
"detail":detail,
}))
.expect("serialize orch prompt event");
@ -4042,6 +4054,21 @@ impl PipelinePromptRuntime {
.unwrap_or(0);
let final_text = self.final_text.clone();
let tokens_generated = self.generated_tokens.len() as u32;
orch_datastream.emit_prompt(
dashboard,
run_id,
node_id,
request_id,
"prompt_complete",
"ready",
json!({
"event":"Done",
"terminal":true,
"tokens_generated":tokens_generated,
"elapsed_ms":elapsed_ms,
"final_text_bytes":final_text.len(),
}),
);
let _ = active.events.send(PromptEvent::Done {
request_id,
final_text,
@ -4985,7 +5012,6 @@ mod tests {
"MVP_PIPELINE_STAGES",
"MVP_STAGE_INDEX",
"MVP_TOKEN_PROGRESS_EVERY",
"MVP_TINYGRAD_TEST_MODE",
"MVP_TINYGRAD_WORKER",
"MVP_TOKENIZER_LOCAL_PATH",
"MVP_VASTAI_API_KEY",
@ -6070,28 +6096,37 @@ mod tests {
.collect::<Vec<_>>();
let _ = std::fs::remove_file(&path);
assert_eq!(records.len(), 2);
assert_eq!(records[0]["arrival_seq"], json!(0));
assert!(
records[0]["arrival_unix_ms"]
.as_u64()
.is_some_and(|value| value > 0)
);
assert_eq!(records[0]["source"], json!("orchestrator"));
assert_eq!(records[0]["stream"], json!("test-node#42"));
assert_eq!(records[0]["channel"], json!("stdout"));
assert_eq!(records[0]["channel_id"], json!(1));
assert_eq!(records[0]["position"], json!(7));
assert_eq!(
records,
vec![
json!({
"arrival_seq":0,
"source":"orchestrator",
"stream":"test-node#42",
"channel":"stdout",
"channel_id":1,
"position":7,
"payload":{"encoding":"utf8","value":"hello λ"},
}),
json!({
"arrival_seq":1,
"source":"orchestrator",
"stream":"test-node#42",
"channel":"stderr",
"channel_id":2,
"position":8,
"payload":{"encoding":"bytes","value":[255,0,65]},
}),
]
records[0]["payload"],
json!({"encoding":"utf8","value":"hello λ"})
);
assert_eq!(records[1]["arrival_seq"], json!(1));
assert!(
records[1]["arrival_unix_ms"]
.as_u64()
.is_some_and(|value| value > 0)
);
assert_eq!(records[1]["source"], json!("orchestrator"));
assert_eq!(records[1]["stream"], json!("test-node#42"));
assert_eq!(records[1]["channel"], json!("stderr"));
assert_eq!(records[1]["channel_id"], json!(2));
assert_eq!(records[1]["position"], json!(8));
assert_eq!(
records[1]["payload"],
json!({"encoding":"bytes","value":[255,0,65]})
);
}

61
crates/mvp-system/src/bin/worker_node.rs
View file

@ -29,6 +29,7 @@ use mvp_system::actors::node_agent::{
};
use mvp_system::actors::register_mvp_actor_codecs;
use mvp_system::arena_manager as arena;
use mvp_system::benchmark_observability;
use mvp_system::distribution_stack::DistributionRuntimeStack;
use mvp_system::driver_pumps as driver_model;
use mvp_system::edge_establisher as edge;
@ -72,10 +73,26 @@ fn node_event_payload(
"run_id":config.run_id,
"node_id":config.logical_node_id,
"stage_index":config.stage_index,
"benchmark":benchmark_observability::stamp("mvp-worker-node"),
"detail":detail,
})
}
fn emit_stdio_datastream_frame(channel: &str, payload: &Value) -> Result<(), String> {
println!(
"{}",
json!({
"mvp_stdio_event":1,
"kind":"datastream_frame",
"channel":channel,
"payload":payload,
})
);
std::io::stdout()
.flush()
.map_err(|e| format!("flush stdio datastream frame: {e}"))
}
fn emit_stdio_node_event(
config: &DeploymentConfig,
channel: &str,
@ -83,18 +100,8 @@ fn emit_stdio_node_event(
status: &str,
detail: Value,
) -> Result<(), String> {
println!(
"{}",
json!({
"mvp_stdio_event":1,
"kind":"datastream_frame",
"channel":channel,
"payload":node_event_payload(config, phase, status, detail),
})
);
std::io::stdout()
.flush()
.map_err(|e| format!("flush stdio node event: {e}"))
let payload = node_event_payload(config, phase, status, detail);
emit_stdio_datastream_frame(channel, &payload)
}
fn emit_node_event(
@ -2956,6 +2963,12 @@ impl DeploymentConfig {
.into_owned(),
),
};
let provider = env_string("MVP_NODE_PROVIDER", "process");
let default_device = if provider == "process" {
"CPU"
} else {
DEFAULT_DEVICE
};
Ok(Self {
run_id,
logical_node_id,
@ -2966,7 +2979,7 @@ impl DeploymentConfig {
debug_join_socket,
relay_mode: relay.mode,
worker_script: env_string("MVP_TINYGRAD_WORKER", DEFAULT_WORKER_SCRIPT),
device: env_string("DEV", DEFAULT_DEVICE),
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(),
@ -2991,6 +3004,9 @@ impl TinygradWorker {
let mut child = Command::new("python3")
.arg(&config.worker_script)
.env("DEV", &config.device)
.env("MVP_RUN_ID", config.run_id.to_string())
.env("MVP_LOGICAL_NODE_ID", config.logical_node_id.to_string())
.env("MVP_STAGE_INDEX", config.stage_index.to_string())
.env("MVP_ARENA_FD", arena_fd.to_string())
.env("MVP_ARENA_BYTES", config.arena_bytes.to_string())
.stdin(Stdio::piped())
@ -3518,6 +3534,8 @@ impl TinygradWorker {
json!({"expected_event_type":expected,"channel":channel_name,"line_bytes":n,"worker_event_type":worker_event_type}),
);
datastream.submit_text(channel, value.to_string());
emit_stdio_datastream_frame(channel_name, &value)
.map_err(|e| format!("emit worker stdio datastream frame: {e}"))?;
datastream.tick();
pump();
if worker_event_type == "WorkerFatal" {
@ -3651,6 +3669,23 @@ mod tests {
DistributionRuntimeStack::new(DistNodeId([1; 32]), DistributedNodeConfig::default())
}
#[test]
fn benchmark_observability_node_event_payload_includes_stamp() {
let config = test_config(None);
let payload = node_event_payload(&config, "phase", "ready", json!({"ok": true}));
assert_eq!(payload.get("run_id").and_then(Value::as_u64), Some(7));
assert_eq!(payload.get("node_id").and_then(Value::as_u64), Some(11));
assert_eq!(payload.get("stage_index").and_then(Value::as_u64), Some(3));
assert_eq!(
payload
.get("benchmark")
.and_then(|benchmark| benchmark.get("schema"))
.and_then(Value::as_u64),
Some(1)
);
}
#[test]
fn debug_join_client_serializes_endpoint_from_stdin() {
let secret = iroh::SecretKey::from_bytes(&[7; 32]);

1
crates/mvp-system/src/lib.rs
View file

@ -3,6 +3,7 @@ extern crate self as mvp_system;
pub mod actors;
pub mod arena_manager;
pub mod benchmark_observability;
pub mod bootstrap_datastream;
pub mod config;
pub mod dashboard_view;

364
crates/mvp-system/tests/python_worker_protocol.rs
View file

@ -3,11 +3,8 @@
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,
@ -15,22 +12,17 @@ struct WorkerProcess {
}
impl WorkerProcess {
fn spawn(arena: Option<(&arena::ArenaManager, u64)>) -> Self {
fn spawn() -> Self {
let script = worker_script();
let mut command = Command::new("python3");
command
let mut child = Command::new("python3")
.arg(&script)
.env("MVP_TINYGRAD_TEST_MODE", "1")
.env("DEV", "CPU")
.env("MVP_RUN_ID", "9")
.env("MVP_LOGICAL_NODE_ID", "3")
.env("MVP_STAGE_INDEX", "2")
.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
.stderr(Stdio::piped())
.spawn()
.unwrap_or_else(|e| panic!("spawn {}: {e}", script.display()));
let stdin = child.stdin.take().expect("worker stdin");
@ -42,9 +34,10 @@ impl WorkerProcess {
}
}
fn send_expect(&mut self, command: Value, expected_type: &str) -> Value {
fn send_collect_until(&mut self, command: Value, expected_type: &str) -> Vec<Value> {
writeln!(self.stdin, "{command}").expect("write worker command");
self.stdin.flush().expect("flush worker command");
let mut events = Vec::new();
loop {
let mut line = String::new();
let read = self.stdout.read_line(&mut line).expect("read worker event");
@ -55,8 +48,10 @@ impl WorkerProcess {
actual_type, "WorkerFatal",
"worker fatal while waiting for {expected_type}: {event}"
);
if actual_type == expected_type {
return event;
let done = actual_type == expected_type;
events.push(event);
if done {
return events;
}
}
}
@ -70,301 +65,62 @@ impl Drop for WorkerProcess {
}
#[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",
);
fn benchmark_observability_worker_ready_includes_stamps_and_identity() {
if !tinygrad_available() {
eprintln!("skipping Python worker protocol check: tinygrad is unavailable");
return;
}
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(
let mut worker = WorkerProcess::spawn();
let events = worker.send_collect_until(
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
assert!(
events
.iter()
.flat_map(|word| word.to_le_bytes())
.collect::<Vec<_>>()
.any(|event| event.get("type").and_then(Value::as_str) == Some("TinygradImportStarted")),
"expected tinygrad import milestone in {events:?}"
);
assert!(
events
.iter()
.any(|event| event.get("type").and_then(Value::as_str) == Some("WorkerReady")),
"expected WorkerReady in {events:?}"
);
for event in &events {
assert_eq!(
event
.get("benchmark")
.and_then(|benchmark| benchmark.get("schema"))
.and_then(Value::as_u64),
Some(1),
"{event}"
);
assert_eq!(
event.get("run_id").and_then(Value::as_u64),
Some(9),
"{event}"
);
assert_eq!(
event.get("node_id").and_then(Value::as_u64),
Some(3),
"{event}"
);
assert_eq!(
event.get("stage_index").and_then(Value::as_u64),
Some(2),
"{event}"
);
}
}
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 tinygrad_available() -> bool {
Command::new("python3")
.args(["-c", "import tinygrad"])
.status()
.is_ok_and(|status| status.success())
}
fn worker_script() -> std::path::PathBuf {

1475
xtask/src/main.rs
View file

File diff suppressed because it is too large Load diff