stash: local single-node gpu prompt loop

This commit is contained in:
Zachery Aaron Shores-Chmielewski 2026-07-23 13:39:17 +04:00
parent 92643bf399
commit 50f58432f5
8 changed files with 8618 additions and 7408 deletions

View file

@ -28,6 +28,7 @@ rings: dict[int, dict[str, Any]] = {}
device_objects: dict[int, dict[str, Any]] = {}
next_handle = 42
HEADER_LEN = 40
FLAG_BEGIN_SEQUENCE = 1 << 1
WORKER_GENERATION = 1
BENCHMARK_SCHEMA = 1
_benchmark_start = time.monotonic()
@ -575,9 +576,19 @@ 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.llm.cli import SimpleTokenizer
try:
from tinygrad.llm.cli import SimpleTokenizer
control(type="TinygradLlmImportReady", model_id=model_id)
Transformer = None
llm_backend = "tinygrad.llm"
except ModuleNotFoundError as exc:
if exc.name != "tinygrad.llm":
raise
from tinygrad.apps.llm import SimpleTokenizer, Transformer
llm_backend = "tinygrad.apps.llm"
control(type="TinygradLlmImportReady", model_id=model_id, backend=llm_backend)
max_context_raw = os.environ.get("MVP_MAX_CONTEXT", "512")
max_context = int(max_context_raw) if max_context_raw else 512
control(
@ -589,13 +600,30 @@ def load_weights(cmd: dict[str, Any]) -> None:
layer_start=layer_start,
layer_end_exclusive=layer_end_exclusive,
requested_device=os.environ.get("DEV"),
llm_backend=llm_backend,
)
model, kv = load_pipeline_stage_model(
path,
max_context=max_context,
layer_start=layer_start,
layer_end_exclusive=layer_end_exclusive,
)
if Transformer is None:
model, kv = load_pipeline_stage_model(
path,
max_context=max_context,
layer_start=layer_start,
layer_end_exclusive=layer_end_exclusive,
)
else:
TensorCls = require_tinygrad()
model, kv = Transformer.from_gguf(TensorCls(path), max_context=max_context, realize=True)
total_layers = int(kv[f"{kv['general.architecture']}.block_count"]) - int(
kv.get(f"{kv['general.architecture']}.nextn_predict_layers", 0)
)
if layer_start != 0 or layer_end_exclusive < total_layers:
fatal(
"TinygradAppsLlmPartialStageUnsupported",
layer_start=layer_start,
layer_end_exclusive=layer_end_exclusive,
total_layers=total_layers,
)
model.first_stage = True
model.final_stage = True
control(
type="PipelineStageFromGgufReady",
model_id=model_id,
@ -607,10 +635,12 @@ def load_weights(cmd: dict[str, Any]) -> None:
first_stage=model.first_stage,
final_stage=model.final_stage,
requested_device=os.environ.get("DEV"),
llm_backend=llm_backend,
)
tok_src = cmd.get("tokenizer", {"EmbeddedGguf": None})
if "EmbeddedGguf" in tok_src:
if kv.get("tokenizer.ggml.pre") == "smollm":
tokenizer_pre = str(kv.get("tokenizer.ggml.pre", "")).lower()
if tokenizer_pre == "smollm":
kv = dict(kv)
kv["tokenizer.ggml.pre"] = "qwen2"
control(type="TokenizerBuildStarted", model_id=model_id, source="EmbeddedGguf")
@ -633,6 +663,7 @@ def load_weights(cmd: dict[str, Any]) -> None:
hidden_dim=int(getattr(model, "hidden_dim", 0)),
max_context=int(getattr(model, "max_context", 0)),
eos_token_id=int(kv.get("tokenizer.ggml.eos_token_id", 0)),
tokenizer_pre=tokenizer_pre,
)
control(
type="WeightsLoaded",
@ -649,6 +680,11 @@ def prompt_template_name() -> str:
explicit = os.environ.get("MVP_PROMPT_TEMPLATE")
if explicit is not None:
return explicit.strip().lower()
tokenizer_pre = str(loaded.get("tokenizer_pre", "")).lower()
if "smollm" in tokenizer_pre or tokenizer_pre == "qwen2":
return "smollm-chat"
if "llama" in tokenizer_pre:
return "llama3-chat"
model_id = str(loaded.get("model_id", "")).lower()
if "smollm" in model_id:
return "smollm-chat"
@ -934,20 +970,26 @@ def payload_words(payload: bytes) -> list[int]:
return []
return list(struct.unpack(f"<{len(payload) // 4}I", payload))
def object_start_pos(sequence: int, token_count: int) -> int:
if sequence == 0:
def object_start_pos(sequence: int, token_count: int, flags: int) -> int:
if flags & FLAG_BEGIN_SEQUENCE or sequence == 0:
role["prompt_tokens"] = token_count
role["prompt_decode_index"] = 0
return 0
return int(role.get("prompt_tokens", 1)) + sequence - 1
decode_index = int(role.get("prompt_decode_index", max(0, sequence - 1)))
role["prompt_decode_index"] = decode_index + 1
return int(role.get("prompt_tokens", token_count)) + decode_index
def materialize_object(payload: bytes, sequence: int) -> dict[str, Any]:
def materialize_object(payload: bytes, sequence: int, flags: int) -> dict[str, Any]:
if not isinstance(model, PipelineStageTinygradModel):
TensorCls = require_tinygrad()
tokens = payload_words(payload)
token_count = len(tokens)
return {
"kind": "words",
"words": payload_words(payload),
"payload": payload,
"start_pos": object_start_pos(sequence, max(1, len(payload) // 4)),
"kind": "tokens",
"tokens": tokens,
"tensor": TensorCls([tokens], dtype="int32").realize(),
"start_pos": object_start_pos(sequence, token_count, flags),
}
TensorCls = require_tinygrad()
if bool(getattr(model, "first_stage", False)) and int(role.get("layer_start", 0)) == 0:
@ -957,7 +999,7 @@ def materialize_object(payload: bytes, sequence: int) -> dict[str, Any]:
"kind": "tokens",
"tokens": tokens,
"tensor": TensorCls([tokens], dtype="int32").realize(),
"start_pos": object_start_pos(sequence, token_count),
"start_pos": object_start_pos(sequence, token_count, flags),
}
import numpy as np
@ -972,7 +1014,7 @@ def materialize_object(payload: bytes, sequence: int) -> dict[str, Any]:
return {
"kind": "activation",
"tensor": TensorCls(array).realize(),
"start_pos": object_start_pos(sequence, token_count),
"start_pos": object_start_pos(sequence, token_count, flags),
}
@ -986,7 +1028,7 @@ def ring_readable(cmd: dict[str, Any]) -> None:
object_id, sequence, extent, flags, payload = parse_record(ring)
handle = next_handle
next_handle += 1
materialized = materialize_object(payload, sequence)
materialized = materialize_object(payload, sequence, flags)
materialized.update(
object_id=object_id,
sequence=sequence,
@ -1048,19 +1090,19 @@ 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"))
execution_backend = "pipeline_stage"
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
payload = struct.pack("<I", token)
flags = 1 if token == int(loaded.get("eos_token_id", 0)) else 0
else:
value = sum(int(word) for word in obj["words"])
value += int(role.get("layer_start", 0)) + int(role.get("layer_end_exclusive", 0)) + int(role.get("stage_index", 0))
if value <= 0:
value = 1
payload = struct.pack("<I", value)
flags = 0
execution_backend = "full_transformer"
if not final_stage:
fatal("FullTransformerNonFinalStageUnsupported", step_id=int(cmd["step_id"]))
if obj.get("kind") != "tokens":
fatal("FullTransformerInputUnsupported", step_id=int(cmd["step_id"]), kind=obj.get("kind"))
if int(obj.get("flags", 0)) & FLAG_BEGIN_SEQUENCE and hasattr(model, "forward_jit"):
model.forward_jit.reset()
token_array = model(obj["tensor"], int(obj.get("start_pos", 0))).realize().numpy().reshape(-1)
token = int(token_array[0])
payload = struct.pack("<I", token)
flags = 1 if token == int(loaded.get("eos_token_id", 0)) else 0
else:
if final_stage != bool(getattr(model, "final_stage", False)):
fatal("FinalStageMismatch", command_final_stage=final_stage, model_final_stage=bool(getattr(model, "final_stage", False)))
@ -1091,6 +1133,7 @@ def execute_step(cmd: dict[str, Any]) -> None:
object_id=int(cmd["output_object_id"]),
sequence=int(cmd["output_sequence"]),
committed_bytes=committed,
execution_backend=execution_backend,
)

View file

@ -36,6 +36,22 @@ const BASE_NODE_IMAGE: &str = "swactor-mvp-node-base:cuda12.6";
const REPO_MODEL_CACHE_DIR: &str = ".model-cache";
const DEFAULT_MAX_TOKENS: u32 = 64;
const ORCH_SHUTDOWN_GRACE_MS: u64 = 5_000;
const MVP_CHAT_GPU_RUN_ENV: &str = "MVP_CHAT_GPU_RUN";
const MVP_CHAT_USAGE: &str = "\
USAGE: cargo mvp-chat [OPTIONS]
OPTIONS:
--gpu Run the local GPU path: in-process orchestrator plus DEV=CUDA worker selection
--process | --docker | --vastai
Select the runtime provider
--config <path> Load config overlay
--pipeline-stages <count> Number of pipeline stages
--cached-model[=<path>] Use discovered or explicit cached GGUF model
--dump-logs[=<path>] Write datastream frame log
--run-id <id> Override run id
--skip-rebuild Reuse existing Cargo artifacts
--yes, -y Approve Vast.ai lease prompts
--help, -h Print this help";
const ORCH_SHUTDOWN_POLL_MS: u64 = 50;
const CHAT_LIFECYCLE_CHANNEL: &str = "mvp.chat.lifecycle";
const CHAT_RUNTIME_CHANNEL: &str = "mvp.chat.runtime";
@ -69,11 +85,54 @@ where
}
}
fn print_usage() {
println!("{MVP_CHAT_USAGE}");
}
fn is_help_request(args: &[String]) -> bool {
args.iter()
.any(|arg| matches!(arg.as_str(), "--help" | "-h" | "help"))
}
struct RuntimeEnvGuard {
name: &'static str,
original: Option<std::ffi::OsString>,
}
impl RuntimeEnvGuard {
fn apply_gpu_defaults(gpu_run: bool) -> Option<Self> {
if !gpu_run || std::env::var_os("DEV").is_some() {
return None;
}
let guard = Self {
name: "DEV",
original: None,
};
unsafe { std::env::set_var(guard.name, "CUDA") };
Some(guard)
}
}
impl Drop for RuntimeEnvGuard {
fn drop(&mut self) {
match &self.original {
Some(value) => unsafe { std::env::set_var(self.name, value) },
None => unsafe { std::env::remove_var(self.name) },
}
}
}
fn run<I>(args: I) -> Result<(), String>
where
I: IntoIterator<Item = String>,
{
let config = Config::from_args(args)?;
let provided_args = args.into_iter().collect::<Vec<_>>();
if is_help_request(&provided_args) {
print_usage();
return Ok(());
}
let config = Config::from_args(provided_args)?;
let _gpu_env = RuntimeEnvGuard::apply_gpu_defaults(config.gpu_run);
let mut progress = ChatDatastream::new(config.run_id, config.datastream_frame_log.clone())?;
progress.emit(
CHAT_LIFECYCLE_CHANNEL,
@ -85,6 +144,7 @@ where
"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()),
"gpu_run": config.gpu_run,
}),
);
confirm_vastai_if_needed(&config)?;
@ -120,21 +180,30 @@ where
CHAT_COMPONENT_CHANNEL,
"orchestrator_process_spawn",
"started",
json!({"binary": config.orch_bin.to_string_lossy()}),
json!({
"mode": config.orchestrator_launch_mode(),
"binary": config.orch_bin.to_string_lossy(),
}),
);
let mut orch = match OrchChild::spawn(&config, &image_ref) {
let mut orch = match OrchHandle::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()}),
json!({
"mode": config.orchestrator_launch_mode(),
"binary": config.orch_bin.to_string_lossy(),
}),
);
progress.emit(
CHAT_COMPONENT_CHANNEL,
"orchestrator_process",
"started",
json!({"binary": config.orch_bin.to_string_lossy()}),
json!({
"mode": config.orchestrator_launch_mode(),
"binary": config.orch_bin.to_string_lossy(),
}),
);
orch
}
@ -143,13 +212,17 @@ where
CHAT_COMPONENT_CHANNEL,
"orchestrator_process_spawn",
"failed",
json!({"binary": config.orch_bin.to_string_lossy(), "error": error}),
json!({
"mode": config.orchestrator_launch_mode(),
"binary": config.orch_bin.to_string_lossy(),
"error": error,
}),
);
progress.emit(
CHAT_COMPONENT_CHANNEL,
"orchestrator_process",
"failed",
json!({"error": error}),
json!({"mode": config.orchestrator_launch_mode(), "error": error}),
);
progress.archive_pending()?;
return Err(error);
@ -257,6 +330,7 @@ struct Config {
pipeline_stages: u32,
max_tokens: u32,
skip_rebuild: bool,
gpu_run: bool,
}
struct ChatDatastream {
@ -537,8 +611,15 @@ impl Config {
if max_tokens == 0 {
return Err("[runtime].max_tokens must be greater than 0".to_owned());
}
let cached_model = args
.cached_model
let gpu_run = args.gpu || env_flag(MVP_CHAT_GPU_RUN_ENV, false);
let cached_model_source = match args.cached_model {
Some(source) => Some(source),
None if gpu_run && provider == ProviderKind::Process => {
Some(CachedModelSource::Discover)
}
None => None,
};
let cached_model = cached_model_source
.map(CachedModelConfig::from_source)
.transpose()?;
let datastream_frame_log = if args.dump_logs {
@ -576,6 +657,7 @@ impl Config {
max_tokens,
vastai,
skip_rebuild: args.skip_rebuild,
gpu_run,
})
}
@ -666,6 +748,14 @@ impl Config {
}
args
}
fn orchestrator_launch_mode(&self) -> &'static str {
if self.gpu_run {
"in_process_actor"
} else {
"process_binary"
}
}
}
#[derive(Default, Debug)]
@ -679,6 +769,8 @@ struct ParsedArgs {
run_id: Option<u64>,
skip_rebuild: bool,
cached_model: Option<CachedModelSource>,
help: bool,
gpu: bool,
}
#[derive(Clone, Debug, PartialEq, Eq)]
@ -707,6 +799,8 @@ impl ParsedArgs {
let mut args = provided_args.into_iter().peekable();
while let Some(arg) = args.next() {
match arg.as_str() {
"--help" | "-h" | "help" => parsed.help = true,
"--gpu" => parsed.gpu = true,
"--vastai" => parsed.set_provider_selector(ProviderKind::VastAi)?,
"--process" => parsed.set_provider_selector(ProviderKind::Process)?,
"--docker" => parsed.set_provider_selector(ProviderKind::Docker)?,
@ -860,6 +954,130 @@ fn parse_approval(input: &str) -> bool {
matches!(input.trim().to_ascii_lowercase().as_str(), "y" | "yes")
}
enum OrchHandle {
Process(OrchChild),
InProcess(InProcessOrch),
}
impl OrchHandle {
fn spawn(config: &Config, image_ref: &str) -> Result<Self, String> {
if config.gpu_run {
InProcessOrch::spawn(config, image_ref).map(Self::InProcess)
} else {
OrchChild::spawn(config, image_ref).map(Self::Process)
}
}
fn wait_ready(&mut self, rpc_addr: String) -> Result<String, String> {
match self {
Self::Process(orch) => orch.wait_ready(rpc_addr),
Self::InProcess(orch) => orch.wait_ready(rpc_addr),
}
}
fn shutdown(&mut self) {
match self {
Self::Process(orch) => orch.shutdown(),
Self::InProcess(orch) => orch.shutdown(),
}
}
}
struct InProcessOrch {
stop_tx: Option<mpsc::Sender<()>>,
thread: Option<thread::JoinHandle<Result<(), String>>>,
cleaned: bool,
}
impl InProcessOrch {
fn spawn(config: &Config, image_ref: &str) -> Result<Self, String> {
let args = config.orchestrator_cli_args(image_ref);
let (stop_tx, stop_rx) = mpsc::channel();
let thread = thread::spawn(move || {
mvp_system::orchestrator_app::run_in_process_from_args(args, stop_rx)
});
Ok(Self {
stop_tx: Some(stop_tx),
thread: Some(thread),
cleaned: false,
})
}
fn wait_ready(&mut self, rpc_addr: String) -> Result<String, String> {
loop {
if STOP_REQUESTED.load(Ordering::SeqCst) {
return Err("interrupted before orchestrator became ready".to_owned());
}
match TcpStream::connect(&rpc_addr) {
Ok(stream) => {
let _ = stream.shutdown(Shutdown::Both);
return Ok(rpc_addr);
}
Err(error)
if matches!(
error.kind(),
std::io::ErrorKind::ConnectionRefused
| std::io::ErrorKind::TimedOut
| std::io::ErrorKind::AddrNotAvailable
) => {}
Err(error) => return Err(format!("connect prompt RPC {rpc_addr}: {error}")),
}
if let Some(result) = self.take_finished_result() {
return Err(format!(
"in-process orchestrator exited before prompt RPC ready: {}",
render_orch_thread_result(result)
));
}
thread::sleep(Duration::from_millis(100));
}
}
fn shutdown(&mut self) {
if self.cleaned {
return;
}
self.cleaned = true;
let _ = self.stop_tx.take().map(|tx| tx.send(()));
let grace = Duration::from_millis(ORCH_SHUTDOWN_GRACE_MS);
let poll = Duration::from_millis(ORCH_SHUTDOWN_POLL_MS);
let started = Instant::now();
while started.elapsed() < grace {
if self.take_finished_result().is_some() {
return;
}
thread::sleep(poll);
}
}
fn take_finished_result(&mut self) -> Option<Result<(), String>> {
if !self
.thread
.as_ref()
.is_some_and(|thread| thread.is_finished())
{
return None;
}
let thread = self.thread.take()?;
Some(match thread.join() {
Ok(result) => result,
Err(_) => Err("in-process orchestrator thread panicked".to_owned()),
})
}
}
impl Drop for InProcessOrch {
fn drop(&mut self) {
self.shutdown();
}
}
fn render_orch_thread_result(result: Result<(), String>) -> String {
match result {
Ok(()) => "completed successfully".to_owned(),
Err(error) => error,
}
}
struct OrchChild {
child: Child,
cleaned: bool,
@ -1007,30 +1225,47 @@ fn prepare_runtime_with_progress(
} 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(
if config.gpu_run {
emit_chat_progress(
&mut progress,
CHAT_RUNTIME_CHANNEL,
"ensure_orchestrator_actor",
"started",
json!({"mode": config.orchestrator_launch_mode()}),
);
emit_chat_progress(
&mut progress,
CHAT_RUNTIME_CHANNEL,
"ensure_orchestrator_actor",
"ready",
json!({"mode": config.orchestrator_launch_mode()}),
);
} else {
emit_chat_progress(
&mut progress,
CHAT_RUNTIME_CHANNEL,
"ensure_orch_binary",
"ready",
"started",
json!({"mode": binary_mode}),
),
Err(error) => {
emit_chat_progress(
);
match ensure_orch_binary(config) {
Ok(()) => emit_chat_progress(
&mut progress,
CHAT_RUNTIME_CHANNEL,
"ensure_orch_binary",
"failed",
json!({"mode": binary_mode, "error": error.as_str()}),
);
return Err(error);
"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);
}
}
}
@ -1709,6 +1944,16 @@ fn provider_from_sources(
Ok(ProviderKind::Process)
}
fn env_flag(name: &str, default: bool) -> bool {
match env_optional(name) {
Some(value) => !matches!(
value.to_ascii_lowercase().as_str(),
"0" | "false" | "no" | "off"
),
None => default,
}
}
fn env_optional(name: &str) -> Option<String> {
std::env::var(name)
.ok()
@ -1775,6 +2020,8 @@ mod tests {
"VASTAI_API_KEY",
"MVP_PIPELINE_STAGES",
"MVP_RUNTIME_CONFIG",
"MVP_CHAT_GPU_RUN",
"DEV",
];
struct TempDir {
@ -1886,6 +2133,7 @@ mod tests {
pipeline_stages: 1,
max_tokens: DEFAULT_MAX_TOKENS,
skip_rebuild: true,
gpu_run: false,
}
}
@ -1988,6 +2236,7 @@ mod tests {
#[test]
fn parsed_args_accepts_public_flags() {
let parsed = ParsedArgs::parse(strings(&[
"--gpu",
"--docker",
"--yes",
"--config",
@ -2008,6 +2257,34 @@ mod tests {
assert_eq!(parsed.dump_log_path, Some(PathBuf::from("logs.ndjson")));
assert_eq!(parsed.cached_model, Some(CachedModelSource::Discover));
assert!(parsed.skip_rebuild);
assert!(parsed.gpu);
let help = ParsedArgs::parse(strings(&["--help"])).expect("help parses");
assert!(help.help);
let short_help = ParsedArgs::parse(strings(&["-h"])).expect("short help parses");
assert!(short_help.help);
}
#[test]
fn config_gpu_flag_selects_in_process_gpu_run() {
let temp = TempDir::new("gpu-flag-config");
let cache = temp.path().join(REPO_MODEL_CACHE_DIR);
fs::create_dir_all(&cache).expect("create model cache");
let cached_path = cache.join("default.gguf");
fs::write(&cached_path, b"cached model").expect("write cached model");
with_process_state(&[], Some(temp.path()), || {
let config = Config::from_args(strings(&["--gpu", "--skip-rebuild"]))
.expect("gpu config resolves");
assert!(config.gpu_run);
assert_eq!(config.orchestrator_launch_mode(), "in_process_actor");
assert_eq!(
config
.cached_model
.as_ref()
.map(|model| model.host_path.clone()),
Some(cached_path.canonicalize().expect("canonical cached model"))
);
});
}
#[test]

File diff suppressed because it is too large Load diff

View file

@ -56,20 +56,25 @@ pub const OBJECT_MAGIC_BYTES: [u8; 4] = *b"MO01";
pub const OBJECT_MAGIC: u32 = u32::from_le_bytes(OBJECT_MAGIC_BYTES);
pub const OBJECT_VERSION: u16 = 1;
pub const FLAG_END_OF_SEQUENCE: u32 = 1;
pub const KNOWN_FLAGS_MASK: u32 = FLAG_END_OF_SEQUENCE;
pub const FLAG_BEGIN_SEQUENCE: u32 = 1 << 1;
pub const KNOWN_FLAGS_MASK: u32 = FLAG_END_OF_SEQUENCE | FLAG_BEGIN_SEQUENCE;
#[derive(Clone, Copy, Debug, PartialEq, Eq, Default)]
pub struct ObjectFlags {
pub end_of_sequence: bool,
pub begin_sequence: bool,
}
impl ObjectFlags {
pub fn bits(self) -> u32 {
let mut bits = 0;
if self.end_of_sequence {
FLAG_END_OF_SEQUENCE
} else {
0
bits |= FLAG_END_OF_SEQUENCE;
}
if self.begin_sequence {
bits |= FLAG_BEGIN_SEQUENCE;
}
bits
}
pub fn from_bits(bits: u32) -> Option<Self> {
@ -78,6 +83,7 @@ impl ObjectFlags {
}
Some(Self {
end_of_sequence: bits & FLAG_END_OF_SEQUENCE != 0,
begin_sequence: bits & FLAG_BEGIN_SEQUENCE != 0,
})
}
}

View file

@ -24,6 +24,7 @@ pub mod node_boot_lifecycle;
pub mod node_image;
pub mod node_provisioning;
pub mod observability_surface;
pub mod orchestrator_app;
pub mod orchestrator_run_fsm;
pub mod orchestrator_token_endpoint;
pub mod prompt_rpc;

File diff suppressed because it is too large Load diff

View file

@ -52,6 +52,7 @@ fn flagged_record(sequence: u64, extent: u64) -> Vec<u8> {
.extent(extent)
.flags(ingress::ObjectFlags {
end_of_sequence: true,
begin_sequence: true,
})
.payload(vec![7; extent as usize])
.encode()
@ -167,7 +168,8 @@ fn parser_decodes_and_propagates_object_header_flags() {
assert_eq!(
parsed.flags,
ingress::ObjectFlags {
end_of_sequence: true
end_of_sequence: true,
begin_sequence: true,
}
);
}

View file

@ -1,4 +1,4 @@
use std::collections::BTreeMap;
use std::collections::{BTreeMap, BTreeSet};
use std::fs;
use std::io::{Read, Write};
use std::path::{Path, PathBuf};
@ -37,6 +37,78 @@ struct MvpChatCheckOutput {
stdin_error: Option<String>,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum MvpChatCheckScenario {
ProcessBaseline,
Gpu,
Multinode,
MultinodeDocker,
}
impl MvpChatCheckScenario {
fn parse_args(args: Vec<String>) -> Result<Self, String> {
let mut scenario = Self::ProcessBaseline;
for arg in args {
let selected = match arg.as_str() {
"--gpu" => Self::Gpu,
"--multinode" => Self::Multinode,
"--multinode-docker" => Self::MultinodeDocker,
other => return Err(format!("unsupported mvp-chat-check argument {other:?}")),
};
if scenario != Self::ProcessBaseline {
return Err(
"mvp-chat-check accepts at most one scenario flag: --gpu, --multinode, or --multinode-docker"
.to_owned(),
);
}
scenario = selected;
}
Ok(scenario)
}
fn name(self) -> &'static str {
match self {
Self::ProcessBaseline => "process",
Self::Gpu => "gpu",
Self::Multinode => "multinode",
Self::MultinodeDocker => "multinode-docker",
}
}
fn mvp_chat_args(self, run_id: u64, dump_log: &Path) -> Vec<String> {
let mut args = Vec::new();
match self {
Self::ProcessBaseline | Self::Multinode => {
args.push("--process".to_owned());
}
Self::Gpu => {
args.extend(["--process".to_owned(), "--gpu".to_owned()]);
}
Self::MultinodeDocker => {
args.push("--docker".to_owned());
}
}
if matches!(self, Self::Multinode | Self::MultinodeDocker) {
args.extend(["--pipeline-stages".to_owned(), "2".to_owned()]);
}
if !matches!(self, Self::Gpu) {
args.push("--cached-model".to_owned());
}
args.extend([
"--run-id".to_owned(),
run_id.to_string(),
format!("--dump-logs={}", dump_log.display()),
]);
args
}
fn env_overrides(self) -> &'static [(&'static str, &'static str)] {
match self {
Self::ProcessBaseline | Self::Gpu | Self::Multinode | Self::MultinodeDocker => &[],
}
}
}
const BASIC_TESTS: &[TestStep] = &[
TestStep {
label: "root crate",
@ -92,13 +164,39 @@ fn print_usage() {
USAGE: cargo xtask <command>
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-check Run real cargo mvp-chat acceptance check.
mvp-chat [--gpu] [--process|--docker|--vastai] [--pipeline-stages n] [--cached-model] [-- args...] Run the human chat wrapper against the real orchestrator/worker bins.
mvp-chat-check [--gpu|--multinode|--multinode-docker]
Run real cargo mvp-chat acceptance check for one explicit scenario.
test Run the basic non-binding test barrier: root crate plus each
non-binding repository package with `cargo test -p`."
);
}
const MVP_CHAT_USAGE: &str = "\
USAGE: cargo mvp-chat [OPTIONS]
OPTIONS:
--gpu Run the local GPU path: in-process orchestrator plus DEV=CUDA worker selection
--process | --docker | --vastai
Select the runtime provider
--config <path> Load config overlay
--pipeline-stages <count> Number of pipeline stages
--cached-model[=<path>] Use discovered or explicit cached GGUF model
--dump-logs[=<path>] Write datastream frame log
--run-id <id> Override run id
--skip-rebuild Reuse existing Cargo artifacts
--yes, -y Approve Vast.ai lease prompts
--help, -h Print this help";
fn print_mvp_chat_usage() {
println!("{MVP_CHAT_USAGE}");
}
fn is_mvp_chat_help_request(args: &[String]) -> bool {
args.iter()
.any(|arg| matches!(arg.as_str(), "--help" | "-h" | "help"))
}
fn run_step(step: &TestStep) -> bool {
println!("\n=== {} ===", step.label);
println!(" cargo {}", step.args.join(" "));
@ -115,7 +213,7 @@ fn run_step(step: &TestStep) -> bool {
fn run_tests() -> ExitCode {
let start = Instant::now();
let check = run_mvp_chat_check();
let check = run_mvp_chat_check(Vec::new());
if check != ExitCode::SUCCESS {
return check;
}
@ -139,13 +237,17 @@ fn run_tests() -> ExitCode {
}
fn run_mvp_chat(args: Vec<String>) -> ExitCode {
let mut command = Command::new(cargo_bin());
command.args(["run", "--package", "mvp-system", "--bin", "mvp-chat", "--"]);
let forwarded = if args.first().is_some_and(|arg| arg == "--") {
args[1..].to_vec()
} else {
args
};
if is_mvp_chat_help_request(&forwarded) {
print_mvp_chat_usage();
return ExitCode::SUCCESS;
}
let mut command = Command::new(cargo_bin());
command.args(["run", "--package", "mvp-system", "--bin", "mvp-chat", "--"]);
let dump_log_path = explicit_dump_log_path_from_mvp_chat_args(&forwarded);
let run_id = run_id_from_mvp_chat_args(&forwarded);
let benchmark_target = dump_log_path.as_deref().zip(run_id);
@ -374,7 +476,15 @@ fn write_mvp_chat_check_paths(root: &Path) -> Result<MvpChatCheckPaths, String>
})
}
fn run_mvp_chat_check() -> ExitCode {
fn run_mvp_chat_check(args: Vec<String>) -> ExitCode {
let scenario = match MvpChatCheckScenario::parse_args(args) {
Ok(scenario) => scenario,
Err(error) => {
eprintln!("mvp-chat-check: failed: {error}");
print_usage();
return ExitCode::from(1);
}
};
let workspace = workspace_root();
let temp_root = unique_temp_dir("mvp-chat-check");
let paths = match write_mvp_chat_check_paths(&temp_root) {
@ -389,8 +499,9 @@ fn run_mvp_chat_check() -> ExitCode {
}
};
let run_id = mvp_chat_check_run_id();
println!("mvp-chat-check: scenario {}", scenario.name());
let output = match run_mvp_chat_check_process(&workspace, &paths, run_id) {
let output = match run_mvp_chat_check_process(&workspace, &paths, run_id, scenario) {
Ok(output) => output,
Err(error) => return fail_mvp_chat_check(&error, &paths, "", "", None),
};
@ -436,7 +547,7 @@ fn run_mvp_chat_check() -> ExitCode {
);
}
};
let events = match assert_dump_log_facts(&paths.dump_log) {
let events = match assert_dump_log_facts(&paths.dump_log, scenario) {
Ok(events) => events,
Err(error) => {
return fail_mvp_chat_check(
@ -448,7 +559,7 @@ fn run_mvp_chat_check() -> ExitCode {
);
}
};
let report = match build_benchmark_report(&events, output.child_elapsed_ms, run_id) {
let report = match build_benchmark_report(&events, output.child_elapsed_ms, run_id, scenario) {
Ok(report) => report,
Err(error) => {
return fail_mvp_chat_check(
@ -483,13 +594,17 @@ fn run_mvp_chat_check_process(
workspace: &Path,
paths: &MvpChatCheckPaths,
run_id: u64,
scenario: MvpChatCheckScenario,
) -> Result<MvpChatCheckOutput, String> {
let mut command = Command::new(cargo_bin());
command.current_dir(workspace).arg("mvp-chat").arg("--");
for arg in scenario.mvp_chat_args(run_id, &paths.dump_log) {
command.arg(arg);
}
for &(key, value) in scenario.env_overrides() {
command.env(key, value);
}
command
.current_dir(workspace)
.args(["mvp-chat", "--", "--cached-model", "--run-id"])
.arg(run_id.to_string())
.arg(format!("--dump-logs={}", paths.dump_log.display()))
.stdin(Stdio::piped())
.stdout(Stdio::piped())
.stderr(Stdio::piped());
@ -1086,7 +1201,10 @@ fn dump_log_inner_payload_text<'a>(
})
}
fn assert_dump_log_facts(path: &Path) -> Result<Vec<DumpLogEvent>, String> {
fn assert_dump_log_facts(
path: &Path,
scenario: MvpChatCheckScenario,
) -> Result<Vec<DumpLogEvent>, String> {
let events = parse_dump_log_events(path)?;
let mut facts = DumpLogFacts::default();
for record in &events {
@ -1116,6 +1234,9 @@ fn assert_dump_log_facts(path: &Path) -> Result<Vec<DumpLogEvent>, String> {
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")?;
if scenario == MvpChatCheckScenario::Gpu {
require_gpu_dump_log_facts(&facts)?;
}
Ok(events)
}
@ -1123,6 +1244,7 @@ fn build_benchmark_report(
events: &[DumpLogEvent],
child_elapsed_ms: u64,
run_id: u64,
scenario: MvpChatCheckScenario,
) -> Result<BenchmarkReport, String> {
let facts = BenchmarkFacts::from_events(events, run_id);
facts.require_span(
@ -1149,12 +1271,21 @@ fn build_benchmark_report(
"prepare_runtime",
"ready",
)?;
facts.require_span(
"mvp.chat.runtime",
"ChatProgress",
"ensure_orch_binary",
"ready",
)?;
if scenario == MvpChatCheckScenario::Gpu {
facts.require_span(
"mvp.chat.runtime",
"ChatProgress",
"ensure_orchestrator_actor",
"ready",
)?;
} else {
facts.require_span(
"mvp.chat.runtime",
"ChatProgress",
"ensure_orch_binary",
"ready",
)?;
}
facts.require_span(
"mvp.chat.runtime",
"ChatProgress",
@ -1438,6 +1569,22 @@ fn missing_benchmark_event(event: String) -> String {
#[derive(Default)]
struct DumpLogFacts {
gpu_worker_device_requested: bool,
gpu_import_ready: bool,
gpu_probe_ready: bool,
gpu_worker_ready: bool,
gpu_cpu_fallback_seen: bool,
gpu_decode_started: BTreeSet<u64>,
gpu_first_token_ready: BTreeSet<u64>,
gpu_decode_ready: BTreeSet<u64>,
gpu_prompt_completed: BTreeSet<u64>,
gpu_pipeline_prompt_encoded: BTreeSet<u64>,
gpu_pipeline_prompt_begin: BTreeSet<u64>,
gpu_pipeline_token_in: BTreeSet<u64>,
gpu_pipeline_token_out: BTreeSet<u64>,
gpu_pipeline_tokenizer_decode_ready: BTreeSet<u64>,
gpu_pipeline_tokens_decoded: BTreeSet<u64>,
gpu_pipeline_real_worker_step_seen: bool,
chat_config_ready: bool,
prepare_runtime_ready: bool,
prompt_rpc_ready: bool,
@ -1458,6 +1605,7 @@ fn record_dump_log_event(
event: &Value,
facts: &mut DumpLogFacts,
) -> Result<(), String> {
record_gpu_dump_log_event(channel, event, facts);
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);
@ -1522,12 +1670,245 @@ fn record_dump_log_event(
Ok(())
}
fn record_gpu_dump_log_event(channel: &str, event: &Value, facts: &mut DumpLogFacts) {
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 event_type == Some("TinygradCpuCompilerSelected") {
facts.gpu_cpu_fallback_seen = true;
}
match (channel, event_type) {
("mvp.worker.initialize", Some("TinygradImportStarted"))
if event_requested_device_is_cuda(event) =>
{
facts.gpu_worker_device_requested = true;
}
("mvp.worker.initialize", Some("TinygradImportReady"))
if event_requested_device_is_cuda(event) && event_env_dev_is_cuda(event) =>
{
facts.gpu_import_ready = true;
}
("mvp.worker.initialize", Some("TinygradDeviceProbeReady"))
if event_requested_device_is_cuda(event) && event_probe_result_is_one(event) =>
{
facts.gpu_probe_ready = true;
}
("mvp.worker.initialize", Some("WorkerReady")) if worker_ready_backend_is_cuda(event) => {
facts.gpu_worker_ready = true;
}
("mvp.worker.prompt", Some("DecodeStarted"))
if event_positive_u64(event, "prompt_tokens")
&& event_positive_u64(event, "max_tokens")
&& event
.get("decode_impl")
.and_then(Value::as_str)
.is_some_and(|decode_impl| !decode_impl.is_empty()) =>
{
if let Some(request_id) = event.get("request_id").and_then(Value::as_u64) {
facts.gpu_decode_started.insert(request_id);
}
}
("mvp.worker.prompt", Some("FirstTokenReady"))
if event.get("token_index").and_then(Value::as_u64) == Some(1) =>
{
if let Some(request_id) = event.get("request_id").and_then(Value::as_u64) {
facts.gpu_first_token_ready.insert(request_id);
}
}
("mvp.worker.prompt", Some("DecodeReady"))
if event_positive_u64(event, "tokens_generated") =>
{
if let Some(request_id) = event.get("request_id").and_then(Value::as_u64) {
facts.gpu_decode_ready.insert(request_id);
}
}
("mvp.worker.prompt", Some("PromptCompleted"))
if event
.get("generated_tokens")
.and_then(Value::as_array)
.is_some_and(|tokens| !tokens.is_empty()) =>
{
if let Some(request_id) = event.get("request_id").and_then(Value::as_u64) {
facts.gpu_prompt_completed.insert(request_id);
}
}
("mvp.worker.tokenizer", Some("PromptEncoded"))
if event
.get("tokens")
.and_then(Value::as_array)
.is_some_and(|tokens| !tokens.is_empty()) =>
{
if let Some(request_id) = event_request_id(event) {
facts.gpu_pipeline_prompt_encoded.insert(request_id);
}
}
("mvp.worker.tokenizer", Some("TokensDecoded"))
if event
.get("text")
.and_then(Value::as_str)
.is_some_and(|text| !text.is_empty()) =>
{
if let Some(request_id) = event_request_id(event) {
facts.gpu_pipeline_tokens_decoded.insert(request_id);
}
}
("mvp.worker.step", Some("StepExecuted"))
if event_positive_u64(event, "committed_bytes") =>
{
let backend = event.get("execution_backend").and_then(Value::as_str);
if matches!(backend, Some("pipeline_stage" | "full_transformer")) {
facts.gpu_pipeline_real_worker_step_seen = true;
}
}
("mvp.orch.prompt", Some("OrchPromptEvent"))
if phase == Some("pipeline_token_in")
&& status == Some("ready")
&& event_request_id(event).is_some() =>
{
let request_id = event_request_id(event).expect("guarded request_id");
facts.gpu_pipeline_token_in.insert(request_id);
if event
.get("detail")
.and_then(|detail| detail.get("begin_sequence"))
.and_then(Value::as_bool)
== Some(true)
{
facts.gpu_pipeline_prompt_begin.insert(request_id);
}
}
("mvp.orch.prompt", Some("OrchPromptEvent"))
if phase == Some("pipeline_token_out")
&& status == Some("observed")
&& event
.get("detail")
.and_then(|detail| detail.get("token_id"))
.and_then(Value::as_u64)
.is_some()
&& event_request_id(event).is_some() =>
{
facts
.gpu_pipeline_token_out
.insert(event_request_id(event).expect("guarded request_id"));
}
("mvp.orch.prompt", Some("OrchPromptEvent"))
if phase == Some("pipeline_tokenizer_decode")
&& status == Some("ready")
&& event
.get("detail")
.and_then(|detail| detail.get("text_bytes"))
.and_then(Value::as_u64)
.is_some_and(|bytes| bytes > 0)
&& event_request_id(event).is_some() =>
{
facts
.gpu_pipeline_tokenizer_decode_ready
.insert(event_request_id(event).expect("guarded request_id"));
}
_ => {}
}
}
fn require_gpu_dump_log_facts(facts: &DumpLogFacts) -> Result<(), String> {
if facts.gpu_cpu_fallback_seen {
return Err("mvp-chat-check: GPU run fell back to the tinygrad CPU compiler".to_owned());
}
require_dump_log_fact(
facts.gpu_worker_device_requested,
"TinygradImportStarted requested_device CUDA",
)?;
require_dump_log_fact(facts.gpu_import_ready, "TinygradImportReady env_DEV CUDA")?;
require_dump_log_fact(facts.gpu_probe_ready, "TinygradDeviceProbeReady CUDA probe")?;
require_dump_log_fact(facts.gpu_worker_ready, "WorkerReady CUDA backend")?;
for request_id in 1..=2 {
let direct_decode = facts.gpu_decode_started.contains(&request_id)
&& facts.gpu_first_token_ready.contains(&request_id)
&& facts.gpu_decode_ready.contains(&request_id)
&& facts.gpu_prompt_completed.contains(&request_id);
let pipeline_decode = facts.gpu_pipeline_real_worker_step_seen
&& facts.gpu_pipeline_prompt_encoded.contains(&request_id)
&& facts.gpu_pipeline_prompt_begin.contains(&request_id)
&& facts.gpu_pipeline_token_in.contains(&request_id)
&& facts.gpu_pipeline_token_out.contains(&request_id)
&& facts
.gpu_pipeline_tokenizer_decode_ready
.contains(&request_id)
&& facts.gpu_pipeline_tokens_decoded.contains(&request_id);
require_dump_log_fact(
direct_decode || pipeline_decode,
&format!("GPU decode/token evidence request_id={request_id}"),
)?;
}
Ok(())
}
fn event_requested_device_is_cuda(event: &Value) -> bool {
event
.get("requested_device")
.and_then(Value::as_str)
.is_some_and(is_cuda_device)
}
fn event_env_dev_is_cuda(event: &Value) -> bool {
event
.get("env_DEV")
.and_then(Value::as_str)
.is_some_and(is_cuda_device)
}
fn event_probe_result_is_one(event: &Value) -> bool {
event
.get("probe_result")
.and_then(Value::as_array)
.is_some_and(|values| values.iter().any(|value| value.as_i64() == Some(1)))
}
fn worker_ready_backend_is_cuda(event: &Value) -> bool {
let Some(backend) = event.get("backend") else {
return false;
};
backend
.get("requested_device")
.and_then(Value::as_str)
.is_some_and(is_cuda_device)
&& backend
.get("env_DEV")
.and_then(Value::as_str)
.is_some_and(is_cuda_device)
&& backend
.get("tinygrad_device")
.and_then(Value::as_str)
.is_some_and(is_cuda_device)
&& event_probe_result_is_one_from_key(event, "cuda_probe")
}
fn event_probe_result_is_one_from_key(event: &Value, key: &str) -> bool {
event
.get(key)
.and_then(Value::as_array)
.is_some_and(|values| values.iter().any(|value| value.as_i64() == Some(1)))
}
fn event_positive_u64(event: &Value, key: &str) -> bool {
event
.get(key)
.and_then(Value::as_u64)
.is_some_and(|value| value > 0)
}
fn is_cuda_device(value: &str) -> bool {
value.to_ascii_uppercase().contains("CUDA")
}
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 event_request_id(event: &Value) -> Option<u64> {
event.get("request_id").and_then(Value::as_u64)
}
fn require_dump_log_fact(found: bool, fact: &str) -> Result<(), String> {
if found {
@ -1547,6 +1928,85 @@ mod tests {
values.iter().map(|value| (*value).to_owned()).collect()
}
#[test]
fn scenario_flags_select_expected_launch_contract() {
let dump_log = Path::new("/tmp/mvp-chat-check.ndjson");
let baseline =
MvpChatCheckScenario::parse_args(Vec::new()).expect("default scenario parses");
assert_eq!(baseline, MvpChatCheckScenario::ProcessBaseline);
assert_eq!(
baseline.mvp_chat_args(42, dump_log),
strings(&[
"--process",
"--cached-model",
"--run-id",
"42",
"--dump-logs=/tmp/mvp-chat-check.ndjson",
])
);
let gpu = MvpChatCheckScenario::parse_args(strings(&["--gpu"])).expect("gpu parses");
assert_eq!(gpu, MvpChatCheckScenario::Gpu);
assert!(gpu.env_overrides().is_empty());
assert_eq!(
gpu.mvp_chat_args(42, dump_log),
strings(&[
"--process",
"--gpu",
"--run-id",
"42",
"--dump-logs=/tmp/mvp-chat-check.ndjson",
])
);
let multinode =
MvpChatCheckScenario::parse_args(strings(&["--multinode"])).expect("multinode parses");
assert_eq!(multinode, MvpChatCheckScenario::Multinode);
assert_eq!(
multinode.mvp_chat_args(42, dump_log),
strings(&[
"--process",
"--pipeline-stages",
"2",
"--cached-model",
"--run-id",
"42",
"--dump-logs=/tmp/mvp-chat-check.ndjson",
])
);
let multinode_docker = MvpChatCheckScenario::parse_args(strings(&["--multinode-docker"]))
.expect("multinode docker parses");
assert_eq!(multinode_docker, MvpChatCheckScenario::MultinodeDocker);
assert_eq!(
multinode_docker.mvp_chat_args(42, dump_log),
strings(&[
"--docker",
"--pipeline-stages",
"2",
"--cached-model",
"--run-id",
"42",
"--dump-logs=/tmp/mvp-chat-check.ndjson",
])
);
}
#[test]
fn scenario_flags_reject_unknown_or_ambiguous_invocations() {
assert!(
MvpChatCheckScenario::parse_args(strings(&["--docker"]))
.expect_err("unknown flag fails")
.contains("unsupported mvp-chat-check argument")
);
assert!(
MvpChatCheckScenario::parse_args(strings(&["--gpu", "--multinode"]))
.expect_err("multiple scenarios fail")
.contains("at most one scenario flag")
);
}
fn temp_path(label: &str) -> PathBuf {
let id = NEXT_TEST_FILE.fetch_add(1, Ordering::Relaxed);
std::env::temp_dir().join(format!(
@ -1583,6 +2043,13 @@ mod tests {
label: &str,
events: Vec<(&'static str, Value)>,
) -> Vec<DumpLogEvent> {
let path = write_synthetic_event_dump(label, events);
let parsed = parse_dump_log_events(&path).expect("parse synthetic dump log");
let _ = fs::remove_file(path);
parsed
}
fn write_synthetic_event_dump(label: &str, events: Vec<(&'static str, Value)>) -> PathBuf {
let lines = events
.into_iter()
.enumerate()
@ -1594,10 +2061,7 @@ mod tests {
)
})
.collect::<Vec<_>>();
let path = write_dump_log(label, lines);
let parsed = parse_dump_log_events(&path).expect("parse synthetic dump log");
let _ = fs::remove_file(path);
parsed
write_dump_log(label, lines)
}
fn benchmark(component: &str, wall_unix_ms: u64, mono_ms: u64) -> Value {
@ -1680,11 +2144,23 @@ mod tests {
wall_unix_ms,
mono_ms,
);
let object = event.as_object_mut().expect("event object");
match event_type {
"DecodeStarted" => {
object.insert("prompt_tokens".to_owned(), json!(4));
object.insert("max_tokens".to_owned(), json!(8));
object.insert("decode_impl".to_owned(), json!("device_resident_greedy"));
}
"FirstTokenReady" => {
object.insert("token_index".to_owned(), json!(1));
}
"DecodeReady" => {
object.insert("tokens_generated".to_owned(), json!(3));
}
_ => {}
}
if event_type == "PromptCompleted" {
event
.as_object_mut()
.expect("event object")
.insert("generated_tokens".to_owned(), json!([1, 2, 3]));
object.insert("generated_tokens".to_owned(), json!([1, 2, 3]));
}
event
}
@ -1953,6 +2429,7 @@ mod tests {
"type": "OrchBootstrap",
"phase": "weights_loaded",
"status": "ready",
"run_id": 9,
"node_id": 1,
"detail": {},
@ -1969,6 +2446,157 @@ mod tests {
]
}
fn dump_log_fact_events(
include_gpu: bool,
include_cpu_fallback: bool,
) -> Vec<(&'static str, Value)> {
let mut events = vec![
("mvp.chat.lifecycle", chat_span("config", "ready", 1_000, 0)),
(
"mvp.chat.runtime",
chat_span("prepare_runtime", "ready", 1_010, 10),
),
(
"mvp.chat.runtime",
chat_span("prompt_rpc", "ready", 1_020, 20),
),
(
"mvp.orch.bootstrap",
stamped(
json!({"type":"OrchBootstrap","phase":"iroh_driver","status":"ready","run_id":9,"node_id":1,"detail":{}}),
"mvp-orchestrator",
1_030,
30,
),
),
(
"mvp.node.bootstrap",
stamped(
json!({"type":"NodeEvent","phase":"iroh_driver","status":"ready","run_id":9,"node_id":3,"stage_index":1,"detail":{}}),
"mvp-worker-node",
1_040,
40,
),
),
(
"mvp.node.worker",
stamped(
json!({"type":"NodeEvent","phase":"worker_initialize","status":"ready","run_id":9,"node_id":3,"stage_index":1,"detail":{"device":"CUDA"}}),
"mvp-worker-node",
1_050,
50,
),
),
(
"mvp.orch.bootstrap",
stamped(
json!({"type":"OrchBootstrap","phase":"weights_loaded","status":"ready","run_id":9,"node_id":1,"detail":{}}),
"mvp-orchestrator",
1_060,
60,
),
),
(
"mvp.chat.prompt",
prompt_chat_span("response_text", "observed", 1, 1_200, 200),
),
(
"mvp.chat.prompt",
prompt_chat_span("request_completed", "ready", 1, 1_210, 210),
),
(
"mvp.chat.prompt",
prompt_chat_span("response_text", "observed", 2, 1_300, 300),
),
(
"mvp.chat.prompt",
prompt_chat_span("request_completed", "ready", 2, 1_310, 310),
),
(
"mvp.chat.lifecycle",
chat_span("shutdown", "requested", 1_400, 400),
),
(
"mvp.chat.component",
chat_span("orchestrator_process", "stopped", 1_410, 410),
),
];
if include_gpu {
events.extend([
(
"mvp.worker.initialize",
stamped(
json!({"type":"TinygradImportStarted","run_id":9,"node_id":3,"stage_index":1,"requested_device":"CUDA","env_DEV":"CUDA"}),
"tinygrad-worker",
1_070,
70,
),
),
(
"mvp.worker.initialize",
stamped(
json!({"type":"TinygradImportReady","run_id":9,"node_id":3,"stage_index":1,"requested_device":"CUDA","env_DEV":"CUDA"}),
"tinygrad-worker",
1_080,
80,
),
),
(
"mvp.worker.initialize",
stamped(
json!({"type":"TinygradDeviceProbeReady","run_id":9,"node_id":3,"stage_index":1,"requested_device":"CUDA","probe_result":[1]}),
"tinygrad-worker",
1_090,
90,
),
),
(
"mvp.worker.initialize",
stamped(
json!({"type":"WorkerReady","run_id":9,"node_id":3,"stage_index":1,"backend":{"requested_device":"CUDA","env_DEV":"CUDA","tinygrad_device":"CUDA"},"cuda_probe":[1]}),
"tinygrad-worker",
1_100,
100,
),
),
]);
if include_cpu_fallback {
events.push((
"mvp.worker.initialize",
stamped(
json!({"type":"TinygradCpuCompilerSelected","run_id":9,"node_id":3,"stage_index":1,"requested_device":"CUDA","selected_device":"CPU:X86"}),
"tinygrad-worker",
1_105,
105,
),
));
}
for request_id in 1..=2 {
let base = 1_200 + request_id * 100;
events.extend([
(
"mvp.worker.prompt",
worker_prompt_event("DecodeStarted", request_id, base + 20, base - 980),
),
(
"mvp.worker.prompt",
worker_prompt_event("FirstTokenReady", request_id, base + 25, base - 975),
),
(
"mvp.worker.prompt",
worker_prompt_event("DecodeReady", request_id, base + 40, base - 960),
),
(
"mvp.worker.prompt",
worker_prompt_event("PromptCompleted", request_id, base + 50, base - 950),
),
]);
}
}
events
}
#[test]
fn benchmark_observability_dump_log_path_parser_finds_equals_and_separate_forms() {
assert_eq!(
@ -2042,12 +2670,91 @@ mod tests {
);
}
#[test]
fn benchmark_observability_gpu_dump_facts_require_cuda_worker_and_decode_cycles() {
let path = write_synthetic_event_dump("gpu-dump-facts", dump_log_fact_events(true, false));
let events = assert_dump_log_facts(&path, MvpChatCheckScenario::Gpu)
.expect("GPU dump log facts pass");
let _ = fs::remove_file(path);
assert!(events.iter().any(|event| {
event.channel == "mvp.worker.initialize"
&& event.event.get("type").and_then(Value::as_str) == Some("WorkerReady")
}));
}
fn gpu_pipeline_only_facts(
real_worker_backend: bool,
prompt_begin_markers: bool,
) -> DumpLogFacts {
let mut facts = DumpLogFacts {
gpu_worker_device_requested: true,
gpu_import_ready: true,
gpu_probe_ready: true,
gpu_worker_ready: true,
..DumpLogFacts::default()
};
if real_worker_backend {
facts.gpu_pipeline_real_worker_step_seen = true;
}
for request_id in 1..=2 {
facts.gpu_pipeline_prompt_encoded.insert(request_id);
facts.gpu_pipeline_token_in.insert(request_id);
facts.gpu_pipeline_token_out.insert(request_id);
facts.gpu_pipeline_tokenizer_decode_ready.insert(request_id);
facts.gpu_pipeline_tokens_decoded.insert(request_id);
if prompt_begin_markers {
facts.gpu_pipeline_prompt_begin.insert(request_id);
}
}
facts
}
#[test]
fn benchmark_observability_gpu_pipeline_facts_require_real_steps_and_prompt_begin_markers() {
let valid = gpu_pipeline_only_facts(true, true);
require_gpu_dump_log_facts(&valid).expect("real pipeline facts pass");
let missing_real_backend = gpu_pipeline_only_facts(false, true);
let error = require_gpu_dump_log_facts(&missing_real_backend)
.expect_err("missing real worker backend should fail");
assert!(
error.contains("GPU decode/token evidence request_id=1"),
"unexpected error: {error}"
);
let missing_prompt_begin = gpu_pipeline_only_facts(true, false);
let error = require_gpu_dump_log_facts(&missing_prompt_begin)
.expect_err("missing prompt begin marker should fail");
assert!(
error.contains("GPU decode/token evidence request_id=1"),
"unexpected error: {error}"
);
}
#[test]
fn benchmark_observability_gpu_dump_facts_reject_cpu_fallback() {
let path = write_synthetic_event_dump("gpu-cpu-fallback", dump_log_fact_events(true, true));
let error = match assert_dump_log_facts(&path, MvpChatCheckScenario::Gpu) {
Ok(_) => panic!("CPU fallback should fail GPU check"),
Err(error) => error,
};
let _ = fs::remove_file(path);
assert!(
error.contains("fell back to the tinygrad CPU compiler"),
"{error}"
);
}
#[test]
fn benchmark_observability_report_requires_granular_decode_events() {
let events = parse_synthetic_events("missing-first-token", benchmark_report_events(false));
let error =
build_benchmark_report(&events, 80, 9).expect_err("missing first token should fail");
let error = build_benchmark_report(&events, 80, 9, MvpChatCheckScenario::ProcessBaseline)
.expect_err("missing first token should fail");
assert!(
error.starts_with("mvp-chat-check: missing benchmark event "),
@ -2064,7 +2771,8 @@ mod tests {
let events =
parse_synthetic_events("pipeline-report", pipeline_benchmark_report_events(true));
let report = build_benchmark_report(&events, 80, 9).expect("pipeline report builds");
let report = build_benchmark_report(&events, 80, 9, MvpChatCheckScenario::ProcessBaseline)
.expect("pipeline report builds");
assert!(
report
@ -2089,7 +2797,7 @@ fn main() -> ExitCode {
let mut args = std::env::args().skip(1);
match args.next().as_deref() {
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-check") => run_mvp_chat_check(args.collect()),
Some("mvp-chat") => run_mvp_chat(args.collect()),
Some("help" | "--help" | "-h") | None => {
print_usage();