Compare commits

...

2 commits

15 changed files with 1625 additions and 49 deletions

1
Cargo.lock generated
View file

@ -5736,6 +5736,7 @@ dependencies = [
name = "xtask"
version = "0.1.0"
dependencies = [
"blake3",
"libc",
"serde_json",
]

View file

@ -751,6 +751,7 @@ def decode_greedy_device_resident(
request_id: int | None,
model_id: str | None,
progress_every: int,
decode_started_at: float,
) -> list[int]:
if max_tokens <= 0:
return []
@ -797,6 +798,7 @@ def decode_greedy_device_resident(
model_id=model_id,
token_index=1,
prompt_tokens=len(prompt_tokens),
first_token_elapsed_ms=int((time.monotonic() - decode_started_at) * 1000),
)
elif progress_every > 0 and tokens_generated % progress_every == 0:
control(
@ -836,6 +838,7 @@ def infer_prompt(cmd: dict[str, Any]) -> None:
prompt_chars=len(prompt),
max_tokens=max_tokens,
)
encode_started = time.monotonic()
model_prompt, prompt_template = model_prompt_text(prompt)
control(
type="PromptEncodeStarted",
@ -852,8 +855,10 @@ def infer_prompt(cmd: dict[str, Any]) -> None:
model_prompt_bytes=len(model_prompt.encode("utf-8")),
prompt_template=prompt_template,
prompt_tokens=len(prompt_tokens),
elapsed_ms=int((time.monotonic() - encode_started) * 1000),
)
progress_every = int(os.environ.get("MVP_TOKEN_PROGRESS_EVERY", "16") or "16")
decode_started = time.monotonic()
control(
type="DecodeStarted",
request_id=request_id,
@ -874,6 +879,7 @@ def infer_prompt(cmd: dict[str, Any]) -> None:
request_id=request_id,
model_id=loaded.get("model_id"),
progress_every=progress_every,
decode_started_at=decode_started,
)
finally:
stop_cpu_line_sampler(cpu_sampler)
@ -883,7 +889,9 @@ def infer_prompt(cmd: dict[str, Any]) -> None:
model_id=loaded.get("model_id"),
prompt_tokens=len(prompt_tokens),
tokens_generated=len(generated),
elapsed_ms=int((time.monotonic() - decode_started) * 1000),
)
text_decode_started = time.monotonic()
control(type="TextDecodeStarted", request_id=request_id, model_id=loaded.get("model_id"), tokens_generated=len(generated))
raw_text = tokenizer.decode(generated) if generated else ""
text = strip_chat_stop_markers(raw_text)
@ -893,6 +901,7 @@ def infer_prompt(cmd: dict[str, Any]) -> None:
model_id=loaded.get("model_id"),
tokens_generated=len(generated),
text_bytes=len(text.encode("utf-8")),
elapsed_ms=int((time.monotonic() - text_decode_started) * 1000),
)
control(
type="PromptCompleted",
@ -930,7 +939,11 @@ def install_ring(cmd: dict[str, Any]) -> None:
type="RingInstalled",
ring_id=ring_id,
edge_id=rings[ring_id]["edge_id"],
port=rings[ring_id]["port"],
direction=rings[ring_id]["direction"],
data_capacity=rings[ring_id]["data_capacity"],
max_extent=rings[ring_id]["max_extent"],
alignment=rings[ring_id]["alignment"],
)
@ -1001,6 +1014,7 @@ def materialize_object(payload: bytes, sequence: int, flags: int) -> dict[str, A
return {
"kind": "tokens",
"tokens": tokens,
"token_count": token_count,
"tensor": TensorCls([tokens], dtype="int32").realize(),
"start_pos": object_start_pos(sequence, token_count, flags),
}
@ -1011,6 +1025,7 @@ def materialize_object(payload: bytes, sequence: int, flags: int) -> dict[str, A
return {
"kind": "tokens",
"tokens": tokens,
"token_count": token_count,
"tensor": TensorCls([tokens], dtype="int32").realize(),
"start_pos": object_start_pos(sequence, token_count, flags),
}
@ -1026,18 +1041,20 @@ def materialize_object(payload: bytes, sequence: int, flags: int) -> dict[str, A
array = np.frombuffer(payload, dtype=np.float16).copy().reshape(1, token_count, hidden_dim)
return {
"kind": "activation",
"token_count": token_count,
"hidden_dim": hidden_dim,
"tensor": TensorCls(array).realize(),
"start_pos": object_start_pos(sequence, token_count, flags),
}
def ring_readable(cmd: dict[str, Any]) -> None:
global next_handle
ring_id = int(cmd["ring_id"])
ring = rings[ring_id]
if ring["direction"] != "ingress":
fatal("WrongRingDirection", ring_id=ring_id, direction=ring["direction"])
started = time.monotonic()
object_id, sequence, extent, flags, payload = parse_record(ring)
handle = next_handle
next_handle += 1
@ -1047,7 +1064,6 @@ def ring_readable(cmd: dict[str, Any]) -> None:
sequence=sequence,
extent=extent,
flags=flags,
payload=payload,
)
device_objects[handle] = materialized
control(
@ -1059,6 +1075,11 @@ def ring_readable(cmd: dict[str, Any]) -> None:
extent=extent,
handle_generation=WORKER_GENERATION,
handle_id=handle,
kind=materialized.get("kind"),
token_count=materialized.get("token_count"),
hidden_dim=materialized.get("hidden_dim"),
start_pos=materialized.get("start_pos"),
elapsed_ms=int((time.monotonic() - started) * 1000),
)
@ -1092,6 +1113,7 @@ def write_record(ring: dict[str, Any], object_id: int, sequence: int, payload: b
def execute_step(cmd: dict[str, Any]) -> None:
if not role:
fatal("RoleNotConfigured")
step_started = time.monotonic()
handle = int(cmd["input_handle_id"])
obj = device_objects.get(handle)
if obj is None:
@ -1103,35 +1125,41 @@ 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"))
input_kind = obj.get("kind")
input_extent = int(obj.get("extent", 0))
execution_backend = "pipeline_stage"
if not isinstance(model, PipelineStageTinygradModel):
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 input_kind != "tokens":
fatal("FullTransformerInputUnsupported", step_id=int(cmd["step_id"]), kind=input_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)
output_kind = "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)))
input_tensor = model.token_hidden(obj["tensor"]) if obj.get("kind") == "tokens" else obj["tensor"]
input_tensor = model.token_hidden(obj["tensor"]) if input_kind == "tokens" else obj["tensor"]
hidden = model.forward_hidden(input_tensor, int(obj.get("start_pos", 0)))
if final_stage:
token_array = model.next_token(hidden).realize().numpy().reshape(-1)
token = int(token_array[0])
payload = struct.pack("<I", token)
output_kind = "token"
flags = 1 if token == int(loaded.get("eos_token_id", 0)) else 0
else:
import numpy as np
activation = hidden.realize().numpy().astype(np.float16, copy=False)
payload = activation.tobytes()
output_kind = "activation"
flags = 0
compute_ready = time.monotonic()
committed = write_record(
ring,
int(cmd["output_object_id"]),
@ -1139,6 +1167,7 @@ def execute_step(cmd: dict[str, Any]) -> None:
payload,
flags,
)
write_ready = time.monotonic()
control(
type="StepExecuted",
step_id=int(cmd["step_id"]),
@ -1147,6 +1176,15 @@ def execute_step(cmd: dict[str, Any]) -> None:
sequence=int(cmd["output_sequence"]),
committed_bytes=committed,
execution_backend=execution_backend,
final_stage=final_stage,
input_kind=input_kind,
input_extent=input_extent,
output_kind=output_kind,
payload_bytes=len(payload),
record_bytes=committed,
stage_execution_ms=int((compute_ready - step_started) * 1000),
record_write_ms=int((write_ready - compute_ready) * 1000),
elapsed_ms=int((write_ready - step_started) * 1000),
)
@ -1157,22 +1195,39 @@ def release_device_object(cmd: dict[str, Any]) -> None:
def encode_prompt(cmd: dict[str, Any]) -> None:
started = time.monotonic()
prompt = str(cmd.get("prompt", ""))
if tokenizer is not None:
model_prompt, _ = model_prompt_text(prompt)
tokens = [int(token) for token in tokenizer.encode(model_prompt)]
else:
model_prompt = prompt
tokens = [int(byte) for byte in prompt.encode("utf-8")] or [0]
control(type="PromptEncoded", request_id=cmd.get("request_id"), tokens=tokens)
control(
type="PromptEncoded",
request_id=cmd.get("request_id"),
tokens=tokens,
prompt_bytes=len(prompt.encode("utf-8")),
model_prompt_bytes=len(model_prompt.encode("utf-8")),
elapsed_ms=int((time.monotonic() - started) * 1000),
)
def decode_tokens(cmd: dict[str, Any]) -> None:
started = time.monotonic()
tokens = [int(token) for token in cmd.get("tokens", [])]
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)
control(type="TokensDecoded", request_id=cmd.get("request_id"), text=text)
control(
type="TokensDecoded",
request_id=cmd.get("request_id"),
text=text,
tokens=len(tokens),
text_bytes=len(text.encode("utf-8")),
elapsed_ms=int((time.monotonic() - started) * 1000),
)
def shutdown_worker(_: dict[str, Any]) -> None:

View file

@ -61,6 +61,7 @@ const CHAT_LIFECYCLE_CHANNEL: &str = "mvp.chat.lifecycle";
const CHAT_RUNTIME_CHANNEL: &str = "mvp.chat.runtime";
const CHAT_PROMPT_CHANNEL: &str = "mvp.chat.prompt";
const CHAT_COMPONENT_CHANNEL: &str = "mvp.chat.component";
const CHAT_BENCHMARK_CHANNEL: &str = "mvp.chat.benchmark";
#[derive(Debug)]
enum PromptInput {
@ -151,6 +152,7 @@ where
"gpu_run": config.gpu_run,
}),
);
progress.emit_benchmark_envelope(&config);
confirm_vastai_if_needed(&config)?;
progress.emit(
CHAT_RUNTIME_CHANNEL,
@ -331,6 +333,7 @@ struct Config {
run_id: u64,
vastai_yes: bool,
vastai: Option<ResolvedVastAiConfig>,
model: ChatModelConfig,
pipeline_stages: u32,
max_tokens: u32,
skip_rebuild: bool,
@ -379,6 +382,7 @@ impl ChatDatastream {
CHAT_RUNTIME_CHANNEL,
CHAT_PROMPT_CHANNEL,
CHAT_COMPONENT_CHANNEL,
CHAT_BENCHMARK_CHANNEL,
] {
out.channel_by_name(name);
}
@ -415,6 +419,68 @@ impl ChatDatastream {
self.flush();
}
fn emit_benchmark_envelope(&mut self, config: &Config) {
let id = self.channel_by_name(CHAT_BENCHMARK_CHANNEL);
let payload = serde_json::to_vec(&json!({
"type": "BenchmarkRunEnvelope",
"phase": "run_envelope",
"status": "ready",
"run_id": self.run_id,
"benchmark": benchmark_observability::stamp("mvp-chat"),
"detail": {
"scenario": "mvp-chat",
"detail_level": "benchmark_observability_v1",
"workload": {
"mode": "stdin_prompt_corpus",
"max_tokens": config.max_tokens,
"prompt_corpus": "external_or_stdin",
},
"model": {
"id": config.model.id.as_deref(),
"gguf_local_path": config.model.gguf_local_path.as_deref(),
"gguf_repo": config.model.gguf_repo.as_deref(),
"gguf_file": config.model.gguf_file.as_deref(),
"gguf_revision": config.model.gguf_revision.as_deref(),
"tokenizer_local_path": config.model.tokenizer_local_path.as_deref(),
"max_context": config.model.max_context,
},
"runtime": {
"provider": config.provider.as_str(),
"pipeline_stages": config.pipeline_stages,
"orchestrator_launch_mode": config.orchestrator_launch_mode(),
"gpu_run": config.gpu_run,
"relay_mode": config.relay_mode.as_deref(),
"relay_configured": config.relay_url.is_some(),
"endpoint_addr_mask": config.endpoint_addr_mask.as_str(),
},
"provider": {
"kind": config.provider.as_str(),
"node_image": &config.node_image,
"image_tag": config.image_tag.as_deref(),
"cached_model": config.cached_model.as_ref().map(|model| model.host_path.to_string_lossy().to_string()),
"vastai": config.vastai.as_ref().map(|vastai| json!({
"image": &vastai.image,
"relay_configured": !vastai.relay_url.is_empty(),
"bootstrap_command_configured": !vastai.bootstrap_command.is_empty(),
"gpu_name": vastai.gpu_name.as_deref(),
"min_gpu_ram_mb": vastai.min_gpu_ram_mb,
"min_down_mbps": vastai.min_down_mbps,
"min_up_mbps": vastai.min_up_mbps,
"max_dph_total": vastai.max_dph_total,
"min_reliability": vastai.min_reliability,
"require_verified": vastai.require_verified,
"disk_gb": vastai.disk_gb,
"has_onstart": vastai.onstart.is_some(),
"has_ssh_identity": vastai.ssh_identity.is_some(),
})),
},
},
}))
.expect("serialize mvp-chat benchmark envelope");
self.producer.submit_bytes(id, payload);
self.flush();
}
fn flush(&mut self) {
let stream = self.stream.clone();
for frame in self.endpoint.mux().drain() {
@ -521,6 +587,7 @@ struct ChatTomlConfig {
observability: ChatObservabilityConfig,
image: ChatImageConfig,
vastai: ChatVastAiConfig,
model: ChatModelConfig,
relay: ChatRelayConfig,
}
@ -568,6 +635,7 @@ struct ChatVastAiConfig {
min_gpu_ram_mb: Option<u64>,
min_down_mbps: Option<f64>,
min_up_mbps: Option<f64>,
max_dph_total: Option<f64>,
min_reliability: Option<f64>,
require_verified: Option<bool>,
disk_gb: Option<u32>,
@ -575,6 +643,18 @@ struct ChatVastAiConfig {
ssh_identity: Option<String>,
}
#[derive(Clone, Debug, Default, Deserialize)]
#[serde(default, deny_unknown_fields)]
struct ChatModelConfig {
id: Option<String>,
gguf_local_path: Option<String>,
gguf_repo: Option<String>,
gguf_file: Option<String>,
gguf_revision: Option<String>,
tokenizer_local_path: Option<String>,
max_context: Option<u32>,
}
#[derive(Clone, Debug)]
struct LoadedChatTomlConfig {
overlay: ChatTomlConfig,
@ -687,6 +767,7 @@ impl Config {
vastai_yes: args.vastai_yes,
pipeline_stages,
max_tokens,
model: toml.model,
vastai,
skip_rebuild: args.skip_rebuild,
gpu_run,
@ -714,6 +795,27 @@ impl Config {
self.pipeline_stages.to_string(),
"--no-dashboard".to_owned(),
];
if let Some(model_id) = &self.model.id {
args.extend(["--model-id".to_owned(), model_id.clone()]);
}
if let Some(path) = &self.model.gguf_local_path {
args.extend(["--gguf-local-path".to_owned(), path.clone()]);
}
if let Some(repo) = &self.model.gguf_repo {
args.extend(["--gguf-repo".to_owned(), repo.clone()]);
}
if let Some(file) = &self.model.gguf_file {
args.extend(["--gguf-file".to_owned(), file.clone()]);
}
if let Some(revision) = &self.model.gguf_revision {
args.extend(["--gguf-revision".to_owned(), revision.clone()]);
}
if let Some(path) = &self.model.tokenizer_local_path {
args.extend(["--tokenizer-local-path".to_owned(), path.clone()]);
}
if let Some(max_context) = self.model.max_context {
args.extend(["--max-context".to_owned(), max_context.to_string()]);
}
if self.provider == ProviderKind::Process {
args.extend([
"--worker-bin".to_owned(),
@ -773,6 +875,12 @@ impl Config {
if let Some(min_up_mbps) = vastai.min_up_mbps {
args.extend(["--vastai-min-up-mbps".to_owned(), min_up_mbps.to_string()]);
}
if let Some(max_dph_total) = vastai.max_dph_total {
args.extend([
"--vastai-max-dph-total".to_owned(),
max_dph_total.to_string(),
]);
}
if let Some(min_reliability) = vastai.min_reliability {
args.extend([
"--vastai-min-reliability".to_owned(),
@ -919,6 +1027,7 @@ fn resolve_vastai_config(
min_gpu_ram_mb: file.min_gpu_ram_mb,
min_down_mbps: file.min_down_mbps,
min_up_mbps: file.min_up_mbps,
max_dph_total: file.max_dph_total,
min_reliability: file.min_reliability,
require_verified: file.require_verified,
onstart: first_non_empty([file.onstart.clone()]),
@ -1362,6 +1471,23 @@ fn prepare_runtime_with_progress(
}
if config.skip_rebuild {
if config.provider == ProviderKind::VastAi {
emit_chat_progress(
&mut progress,
CHAT_RUNTIME_CHANNEL,
"ensure_worker_binary",
"skipped",
json!({"mode": binary_mode, "reason": "vastai_remote_image"}),
);
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());
}
emit_chat_progress(
&mut progress,
CHAT_RUNTIME_CHANNEL,
@ -1601,6 +1727,10 @@ fn emit_chat_progress(
}
}
fn prompt_hash_hex(prompt: &str) -> String {
blake3::hash(prompt.as_bytes()).to_hex().to_string()
}
fn run_chat_session_with_output_and_progress<R, W, O>(
writer: &mut W,
mut reader: R,
@ -1616,6 +1746,7 @@ where
{
let mut progress = progress;
let mut next_request_id = 1_u64;
let mut next_prompt_index = 1_u64;
loop {
if STOP_REQUESTED.load(Ordering::SeqCst) {
@ -1633,7 +1764,7 @@ where
CHAT_PROMPT_CHANNEL,
"waiting_for_prompt",
"started",
json!({"next_request_id": next_request_id}),
json!({"next_request_id": next_request_id, "next_prompt_index": next_prompt_index}),
);
write!(output, "prompt:> ").map_err(|e| format!("write prompt: {e}"))?;
output.flush().map_err(|e| format!("flush prompt: {e}"))?;
@ -1666,12 +1797,15 @@ where
let request_id = next_request_id;
next_request_id = next_request_id.wrapping_add(1).max(1);
let prompt_index = next_prompt_index;
next_prompt_index = next_prompt_index.wrapping_add(1).max(1);
let prompt_hash = prompt_hash_hex(&prompt);
emit_chat_progress(
&mut progress,
CHAT_PROMPT_CHANNEL,
"prompt_submitted",
"ready",
json!({"request_id": request_id, "prompt_bytes": prompt.len(), "max_tokens": max_tokens}),
json!({"request_id": request_id, "prompt_index": prompt_index, "prompt_hash": &prompt_hash, "prompt_bytes": prompt.len(), "max_tokens": max_tokens}),
);
write_json_line(
writer,
@ -1687,7 +1821,7 @@ where
CHAT_PROMPT_CHANNEL,
"decoding",
"started",
json!({"request_id": request_id}),
json!({"request_id": request_id, "prompt_index": prompt_index, "prompt_hash": &prompt_hash}),
);
let mut response_started = false;
@ -1768,7 +1902,7 @@ where
CHAT_PROMPT_CHANNEL,
"response_text",
"observed",
json!({"request_id": request_id, "text_bytes": text.len()}),
json!({"request_id": request_id, "prompt_index": prompt_index, "prompt_hash": &prompt_hash, "text_bytes": text.len()}),
);
}
PromptEvent::Done {
@ -1790,6 +1924,8 @@ where
"ready",
json!({
"request_id": request_id,
"prompt_index": prompt_index,
"prompt_hash": &prompt_hash,
"response_started": response_started,
"tokens_generated": tokens_generated,
"elapsed_ms": elapsed_ms,
@ -1806,7 +1942,7 @@ where
CHAT_PROMPT_CHANNEL,
"request_faulted",
"ready",
json!({"request_id": request_id, "error": error}),
json!({"request_id": request_id, "prompt_index": prompt_index, "prompt_hash": &prompt_hash, "error": error}),
);
break;
}
@ -2185,6 +2321,7 @@ mod tests {
run_id: 1,
vastai_yes: false,
vastai: None,
model: ChatModelConfig::default(),
pipeline_stages: 1,
max_tokens: DEFAULT_MAX_TOKENS,
skip_rebuild: true,
@ -2206,6 +2343,7 @@ mod tests {
min_gpu_ram_mb: None,
min_down_mbps: None,
min_up_mbps: None,
max_dph_total: None,
min_reliability: None,
require_verified: None,
onstart: None,
@ -2934,6 +3072,25 @@ relay_url = "https://relay.example"
assert_eq!(image_ref, "docker.io/acme/node:latest");
}
#[test]
fn vastai_skip_rebuild_uses_remote_image_without_worker_artifact() {
let temp = TempDir::new("vastai-skip-rebuild");
let orch_bin = temp.path().join("mvp-orchestrator");
let worker_bin = temp.path().join("mvp-worker-node");
let mut config = base_config(ProviderKind::VastAi);
config.skip_rebuild = true;
config.orch_bin = orch_bin.clone();
config.worker_bin = worker_bin;
config.node_image = "docker.io/acme/node:latest".to_owned();
assert!(prepare_runtime_with(&config, panic_prepare_node_image).is_err());
fs::write(&orch_bin, b"orch").expect("write orchestrator artifact");
let image_ref = prepare_runtime_with(&config, panic_prepare_node_image)
.expect("VastAI skip rebuild reuses remote image");
assert_eq!(image_ref, "docker.io/acme/node:latest");
}
#[test]
fn artifact_roots_use_current_directory() {
let temp = TempDir::new("artifact-root");

View file

@ -592,6 +592,14 @@ impl WorkerEdgeRuntime {
edge_id: driver_model::EdgeId(edge_id),
stream_id: driver_model::StreamId(stream_id),
});
emit_node_event(
datastream,
config,
NODE_STAGE_CHANNEL,
"iroh_edge_stream_arrived",
"observed",
json!({"edge_id":edge_id,"stream_id":stream_id}),
);
self.drive_edge_workflow(
stack,
node_actor,
@ -608,6 +616,15 @@ impl WorkerEdgeRuntime {
bytes,
..
} => {
let byte_count = bytes.len();
emit_node_event(
datastream,
config,
NODE_STAGE_CHANNEL,
"iroh_edge_bytes_read",
"observed",
json!({"edge_id":edge_id,"stream_id":stream_id,"bytes":byte_count}),
);
self.ingest_stream_bytes(
edge_id,
stream_id,
@ -792,6 +809,22 @@ impl WorkerEdgeRuntime {
.read_arena(lease.layout.data_offset, committed_bytes)
.map_err(|e| format!("read egress ring: {e}"))?
};
let record_bytes = record.len();
emit_node_event(
datastream,
config,
NODE_STAGE_CHANNEL,
"egress_ring_read",
"ready",
json!({
"edge_id":outbound.edge_id,
"edge_kind":format!("{:?}", outbound.kind),
"ring_id":output_ring_id,
"record_bytes":record_bytes,
"committed_bytes":committed_bytes,
"final_stage":final_stage,
}),
);
self.driver_model
.observe(driver_model::DriverEvent::EgressBytesCommitted {
edge_id: driver_model::EdgeId(outbound.edge_id),
@ -806,6 +839,19 @@ impl WorkerEdgeRuntime {
.as_ref()
.ok_or_else(|| "outbound edge sender missing".to_owned())?;
sender.send(record)?;
emit_node_event(
datastream,
config,
NODE_STAGE_CHANNEL,
"iroh_edge_bytes_sent",
"ready",
json!({
"edge_id":outbound.edge_id,
"edge_kind":format!("{:?}", outbound.kind),
"bytes":record_bytes,
"record_bytes":record_bytes,
}),
);
stack
.runtime
.send_to(node_actor, NodeAgentMsg::StepCompleted { step_id })
@ -867,6 +913,19 @@ impl WorkerEdgeRuntime {
.write_arena(lease.layout.data_offset, &record)
.map_err(|e| format!("write ingress ring: {e}"))?;
}
emit_node_event(
datastream,
config,
NODE_STAGE_CHANNEL,
"ingress_ring_write",
"ready",
json!({
"edge_id":edge_id,
"edge_kind":format!("{:?}", inbound.kind),
"ring_id":ring_id,
"record_bytes":record.len(),
}),
);
let loaded = worker.ring_readable(
ring_id,
edge_id,

View file

@ -106,6 +106,7 @@ pub struct VastAiConfig {
pub min_gpu_ram_mb: Option<u64>,
pub min_down_mbps: Option<f64>,
pub min_up_mbps: Option<f64>,
pub max_dph_total: Option<f64>,
pub min_reliability: Option<f64>,
pub require_verified: Option<bool>,
pub poll_interval_secs: Option<u64>,
@ -130,6 +131,7 @@ pub struct ResolvedVastAiConfig {
pub min_gpu_ram_mb: Option<u64>,
pub min_down_mbps: Option<f64>,
pub min_up_mbps: Option<f64>,
pub max_dph_total: Option<f64>,
pub min_reliability: Option<f64>,
pub require_verified: Option<bool>,
pub onstart: Option<String>,
@ -204,6 +206,9 @@ impl ResolvedVastAiConfig {
if let Some(min_up_mbps) = self.min_up_mbps {
policy.min_up_mbps = Some(min_up_mbps);
}
if let Some(max_dph_total) = self.max_dph_total {
policy.max_dph_total = Some(max_dph_total);
}
if let Some(min_reliability) = self.min_reliability {
policy.min_reliability = min_reliability;
}
@ -259,6 +264,7 @@ mod tests {
min_gpu_ram_mb: Some(16_000),
min_down_mbps: Some(100.0),
min_up_mbps: Some(25.0),
max_dph_total: Some(0.10),
min_reliability: Some(0.98),
require_verified: Some(true),
onstart: None,
@ -326,6 +332,7 @@ gpu_name = "RTX 4090"
min_gpu_ram_mb = 24000
min_down_mbps = 250.5
min_up_mbps = 50.25
max_dph_total = 0.10
min_reliability = 0.99
require_verified = true
onstart = "echo preparing"
@ -402,6 +409,7 @@ poll_interval_secs = 30
assert_eq!(config.vastai.min_gpu_ram_mb, Some(24_000));
assert_eq!(config.vastai.min_down_mbps, Some(250.5));
assert_eq!(config.vastai.min_up_mbps, Some(50.25));
assert_eq!(config.vastai.max_dph_total, Some(0.10));
assert_eq!(config.vastai.min_reliability, Some(0.99));
assert_eq!(config.vastai.require_verified, Some(true));
assert_eq!(config.vastai.poll_interval_secs, Some(30));

View file

@ -636,6 +636,15 @@ impl VastAiRuntimeConfig {
if let Some(min_down_mbps) = min_down_mbps {
provisioning.selection.min_down_mbps = min_down_mbps;
}
let max_dph_total = builder
.vastai_max_dph_total_raw
.as_ref()
.map(|value| ConfigBuilder::parse_value("MVP_VASTAI_MAX_DPH_TOTAL", value))
.transpose()?
.or(builder.vastai_max_dph_total);
if let Some(max_dph_total) = max_dph_total {
provisioning.selection.max_dph_total = Some(max_dph_total);
}
let min_up_mbps = builder
.vastai_min_up_mbps_raw
.as_ref()
@ -695,6 +704,7 @@ impl VastAiRuntimeConfig {
"min_gpu_ram_mb": self.provisioning.selection.min_gpu_ram_mb,
"min_down_mbps": self.provisioning.selection.min_down_mbps,
"min_up_mbps": self.provisioning.selection.min_up_mbps,
"max_dph_total": self.provisioning.selection.max_dph_total,
"min_reliability": self.provisioning.selection.min_reliability,
"require_verified": self.provisioning.selection.require_verified,
"confirm_lease": self.provisioning.confirm_lease,
@ -910,6 +920,8 @@ struct ConfigBuilder {
vastai_min_down_mbps: Option<f64>,
vastai_min_down_mbps_raw: Option<String>,
vastai_min_up_mbps: Option<f64>,
vastai_max_dph_total: Option<f64>,
vastai_max_dph_total_raw: Option<String>,
vastai_min_up_mbps_raw: Option<String>,
vastai_min_reliability: Option<f64>,
vastai_min_reliability_raw: Option<String>,
@ -965,6 +977,8 @@ impl ConfigBuilder {
vastai_min_gpu_ram_mb_raw: None,
vastai_min_down_mbps: None,
vastai_min_down_mbps_raw: None,
vastai_max_dph_total: None,
vastai_max_dph_total_raw: None,
vastai_min_up_mbps: None,
vastai_min_up_mbps_raw: None,
vastai_min_reliability: None,
@ -1083,6 +1097,9 @@ impl ConfigBuilder {
if let Some(min_down_mbps) = overlay.vastai.min_down_mbps {
self.vastai_min_down_mbps = Some(min_down_mbps);
}
if let Some(max_dph_total) = overlay.vastai.max_dph_total {
self.vastai_max_dph_total = Some(max_dph_total);
}
if let Some(min_up_mbps) = overlay.vastai.min_up_mbps {
self.vastai_min_up_mbps = Some(min_up_mbps);
}
@ -1215,6 +1232,9 @@ impl ConfigBuilder {
if let Some(min_down_mbps) = env_optional("MVP_VASTAI_MIN_DOWN_MBPS") {
self.vastai_min_down_mbps_raw = Some(min_down_mbps);
}
if let Some(max_dph_total) = env_optional("MVP_VASTAI_MAX_DPH_TOTAL") {
self.vastai_max_dph_total_raw = Some(max_dph_total);
}
if let Some(min_up_mbps) = env_optional("MVP_VASTAI_MIN_UP_MBPS") {
self.vastai_min_up_mbps_raw = Some(min_up_mbps);
}
@ -1339,6 +1359,11 @@ impl ConfigBuilder {
Some(parse_next(&mut args, "--vastai-min-down-mbps")?);
self.vastai_min_down_mbps_raw = None;
}
"--vastai-max-dph-total" => {
self.vastai_max_dph_total =
Some(parse_next(&mut args, "--vastai-max-dph-total")?);
self.vastai_max_dph_total_raw = None;
}
"--vastai-min-up-mbps" => {
self.vastai_min_up_mbps = Some(parse_next(&mut args, "--vastai-min-up-mbps")?);
self.vastai_min_up_mbps_raw = None;
@ -1380,9 +1405,9 @@ impl ConfigBuilder {
if self.pipeline_stages == 0 {
return Err("--pipeline-stages must be greater than 0".to_owned());
}
if provider == ProviderKind::VastAi && self.pipeline_stages > 1 {
if provider == ProviderKind::VastAi && self.pipeline_stages > 2 {
return Err(
"pipeline stages greater than 1 are only supported with provider=docker; provider=vastai does not support pipelined provisioning yet".to_owned(),
"provider=vastai currently supports at most 2 pipeline stages for activation-path smoke checks".to_owned(),
);
}
let mut cached_model_host_path = self.cached_model_host_path.clone();
@ -1648,6 +1673,20 @@ impl Config {
))
}
}
GgufSource::HuggingFaceGguf { repo, file, .. }
if self.provider == ProviderKind::VastAi
&& gguf_source_matches_default_pipeline_cache(&self.gguf_source) =>
{
let host_path = default_pipeline_cached_model_path();
if host_path.is_file() {
Ok(host_path)
} else {
Err(format!(
"VastAI pipeline planning requires local GGUF metadata at {}; selected remote source is {repo}/{file}",
host_path.display()
))
}
}
GgufSource::HuggingFaceGguf { repo, file, .. } => Err(format!(
"local pipeline requires a locally inspectable GGUF before provisioning; selected source {repo}/{file} is remote, so use --cached-model-host-path"
)),
@ -6618,25 +6657,91 @@ kind = "docker"
}
#[test]
fn vastai_rejects_pipeline_stages_count_above_one() {
fn vastai_rejects_pipeline_stages_count_above_two() {
let error = with_clean_env(&[], || {
match Config::from_layers_with_path_and_args(
None,
["--provider", "vastai", "--pipeline-stages", "2"]
["--provider", "vastai", "--pipeline-stages", "3"]
.into_iter()
.map(str::to_owned),
) {
Ok(_) => panic!("VastAI cannot provision more than one pipeline stage yet"),
Ok(_) => panic!("VastAI smoke runs are capped at two pipeline stages"),
Err(error) => error,
}
});
assert!(
error.contains("provider=vastai does not support pipelined provisioning yet"),
error.contains("provider=vastai currently supports at most 2 pipeline stages"),
"unexpected error: {error}"
);
}
#[test]
fn vastai_two_stage_plan_uses_remote_gguf_and_no_mounts() {
let config = with_clean_env(&[], || {
Config::from_layers_with_path_and_args(
None,
[
"--provider",
"vastai",
"--pipeline-stages",
"2",
"--model-id",
"smollm2-135m-instruct-q4",
"--gguf-repo",
"QuantFactory/SmolLM2-135M-Instruct-GGUF",
"--gguf-file",
DEFAULT_PIPELINE_CACHED_MODEL_FILE,
"--max-context",
"256",
"--relay-mode",
"default",
"--vastai-bootstrap-command",
"boot",
]
.into_iter()
.map(str::to_owned),
)
.expect("VastAI two-stage pipeline config parses")
});
let plan = config
.build_run_plan()
.expect("VastAI two-stage run plan uses local metadata only");
let coordinator = EndpointAddr::new(iroh::SecretKey::from_bytes(&[41; 32]).public());
let orchestrator_actor = ActorAddress([42; 32]);
let specs = stage_node_specs(&config, Some(&plan), coordinator, orchestrator_actor)
.expect("VastAI pipeline stage node specs build");
assert_eq!(config.provider, ProviderKind::VastAi);
assert!(
config.cached_model.is_none(),
"VastAI must not mount host caches"
);
assert_eq!(specs.len(), 2);
for (expected_stage_index, spec) in specs.iter().enumerate() {
let expected_stage_index =
u32::try_from(expected_stage_index).expect("fixture stage index fits u32");
let expected_node_id = config.node_id + 1 + u64::from(expected_stage_index);
assert_eq!(spec.node_id, expected_node_id);
assert_eq!(spec.stage_index, Some(expected_stage_index));
assert_eq!(env_value(&spec.env, "MVP_NODE_PROVIDER"), Some("vastai"));
assert_eq!(env_value(&spec.env, "MVP_PIPELINE_STAGES"), Some("2"));
assert_eq!(
env_value(&spec.env, "MVP_GGUF_REPO"),
Some("QuantFactory/SmolLM2-135M-Instruct-GGUF")
);
assert_eq!(
env_value(&spec.env, "MVP_GGUF_FILE"),
Some(DEFAULT_PIPELINE_CACHED_MODEL_FILE)
);
assert_eq!(env_value(&spec.env, "MVP_GGUF_LOCAL_PATH"), None);
assert_eq!(env_value(&spec.env, "MVP_MAX_CONTEXT"), Some("256"));
assert_eq!(spec.args, vec!["boot".to_owned()]);
assert!(spec.mounts.is_empty(), "VastAI stage specs must not mount");
}
}
#[test]
fn pipeline_cached_model_resolution_uses_toml_smollm2_gguf_file_with_cli_stage_count() {
let toml = TempTomlFile::new(

View file

@ -797,6 +797,24 @@ where
instance.contract_id
),
});
sink.observe(PluginObservation::ProviderLine {
run_id: spec.run_id,
node_id: spec.node_id,
line: serde_json::json!({
"type": "VastAiLeaseReady",
"run_id": spec.run_id,
"node_id": spec.node_id,
"label": &label,
"image": &spec.image,
"contract_id": instance.contract_id,
"offer_id": instance.offer_id,
"host_id": instance.host_id,
"gpu_name": &instance.gpu_name,
"gpu_ram": instance.gpu_ram,
"dph_total": instance.dph_total,
})
.to_string(),
});
let endpoint = match self.client.ssh_endpoint(
instance.contract_id,
@ -812,6 +830,20 @@ where
));
}
};
sink.observe(PluginObservation::ProviderLine {
run_id: spec.run_id,
node_id: spec.node_id,
line: serde_json::json!({
"type": "VastAiSshEndpointReady",
"run_id": spec.run_id,
"node_id": spec.node_id,
"contract_id": instance.contract_id,
"host": &endpoint.host,
"port": endpoint.port,
"user": &endpoint.user,
})
.to_string(),
});
let bootstrap = match self.bootstrap.start_bootstrap(
spec.clone(),
@ -843,14 +875,7 @@ where
Ok(handle)
}
fn complete_bootstrap(&mut self, handle: &PluginNodeHandle) -> Result<(), String> {
let Some(node) = self.nodes.get_mut(&handle.id) else {
return Ok(());
};
if let Some(mut bootstrap) = node.bootstrap.take() {
self.bootstrap
.stop_bootstrap(&mut bootstrap, BootstrapStopReason::RuntimeReady);
}
fn complete_bootstrap(&mut self, _handle: &PluginNodeHandle) -> Result<(), String> {
Ok(())
}
@ -937,6 +962,81 @@ mod tests {
VastAiProvisioningPlugin::new(NoopLeaseClient, NoopBootstrapLauncher, config)
}
#[derive(Clone, Default)]
struct ObservationSink {
observations: Arc<Mutex<Vec<PluginObservation>>>,
}
impl crate::provisioning::PluginObservationSink for ObservationSink {
fn observe(&self, observation: PluginObservation) {
self.observations.lock().push(observation);
}
}
#[derive(Clone)]
struct RecordingLeaseClient {
destroyed_contracts: Arc<Mutex<Vec<u64>>>,
}
impl VastAiLeaseClient for RecordingLeaseClient {
fn provision_one(
&mut self,
_request: ProvisionRequest,
) -> Result<ProvisionedInstance, String> {
Ok(ProvisionedInstance {
index: 0,
contract_id: 42,
offer_id: 7,
host_id: Some(99),
gpu_name: "RTX 4060".to_owned(),
gpu_ram: Some(8_192.0),
dph_total: 0.064,
})
}
fn ssh_endpoint(
&mut self,
_contract_id: u64,
_label: &str,
_lifecycle: &LifecyclePolicy,
ssh_user: &str,
) -> Result<VastAiSshEndpoint, String> {
Ok(VastAiSshEndpoint {
host: "ssh5.vast.ai".to_owned(),
port: 22_017,
user: ssh_user.to_owned(),
})
}
fn destroy_contract(&mut self, contract_id: u64) -> Result<(), String> {
self.destroyed_contracts.lock().push(contract_id);
Ok(())
}
}
#[derive(Clone)]
struct RecordingBootstrapLauncher {
stop_reasons: Arc<Mutex<Vec<BootstrapStopReason>>>,
}
impl VastAiBootstrapLauncher for RecordingBootstrapLauncher {
type Handle = u64;
fn start_bootstrap(
&mut self,
spec: NodeProvisionSpec,
_endpoint: VastAiSshEndpoint,
_sink: PluginSink,
_producer: Option<DatastreamProducer>,
) -> Result<Self::Handle, String> {
Ok(spec.node_id)
}
fn stop_bootstrap(&mut self, _handle: &mut Self::Handle, reason: BootstrapStopReason) {
self.stop_reasons.lock().push(reason);
}
}
#[test]
fn vastai_provisioning_build_request_keeps_bootstrap_args_out_of_onstart() {
let plugin = plugin_with_onstart(None);
@ -957,6 +1057,38 @@ mod tests {
assert_eq!(request.onstart.as_deref(), Some("echo explicit setup"));
}
#[test]
fn vastai_complete_bootstrap_keeps_log_tail_until_node_stop() {
let destroyed_contracts = Arc::new(Mutex::new(Vec::new()));
let stop_reasons = Arc::new(Mutex::new(Vec::new()));
let sink = PluginSink::new(Arc::new(ObservationSink::default()));
let mut plugin = VastAiProvisioningPlugin::new(
RecordingLeaseClient {
destroyed_contracts: destroyed_contracts.clone(),
},
RecordingBootstrapLauncher {
stop_reasons: stop_reasons.clone(),
},
VastAiProvisioningConfig::default(),
);
let handle = plugin
.start_node(node_spec_with_bootstrap_args(), sink)
.expect("VastAI node starts");
plugin
.complete_bootstrap(&handle)
.expect("runtime-ready bootstrap completion succeeds");
assert!(
stop_reasons.lock().is_empty(),
"VastAI bootstrap SSH tail must remain alive for post-ready worker logs"
);
plugin.stop_node(&handle).expect("VastAI node stops");
assert_eq!(*stop_reasons.lock(), vec![BootstrapStopReason::NodeStop]);
assert_eq!(*destroyed_contracts.lock(), vec![42]);
}
#[test]
fn vastai_provisioning_next_ssh_backoff_doubles_until_thirty_second_cap() {
for (current, expected) in [

View file

@ -681,7 +681,6 @@ impl<'a> Ctx<'a> {
self.inner.extension()
}
/// Return system-level information (worker count, actor count, uptime).
pub fn system_info(&self) -> SystemInfo {
self.inner.system_info()

View file

@ -598,7 +598,6 @@ impl Runtime {
self.stats_hook = Some(hook);
}
/// Set the sink for non-local (remote) message delivery.
///
/// The sink owns all codec/transport concerns; core only knows how to hand

View file

@ -9,6 +9,7 @@ pub const ENV_MIN_INET_UP_MBPS: &str = "PP_MIN_INET_UP_MBPS";
pub const ENV_MIN_RELIABILITY: &str = "PP_MIN_RELIABILITY";
pub const ENV_REQUIRE_VERIFIED: &str = "PP_REQUIRE_VERIFIED";
pub const ENV_DROP_CHEAP_FRAC: &str = "PP_DROP_CHEAP_FRAC";
pub const ENV_MAX_DPH_TOTAL: &str = "PP_MAX_DPH_TOTAL";
pub const ENV_LEASE_PACE_MS: &str = "PP_LEASE_PACE_MS";
pub const ENV_BLACKLIST_HOSTS: &str = "PP_BLACKLIST_HOSTS";
pub const ENV_ASSUME_YES: &str = "PP_ASSUME_YES";
@ -55,6 +56,7 @@ impl SelectionPolicy {
.unwrap_or(0.95);
policy.require_verified = truthy_env(ENV_REQUIRE_VERIFIED);
policy.min_up_mbps = env_optional_positive_f64(ENV_MIN_INET_UP_MBPS);
policy.max_dph_total = env_optional_positive_f64(ENV_MAX_DPH_TOTAL);
policy.drop_cheap_frac = std::env::var(ENV_DROP_CHEAP_FRAC)
.ok()
.and_then(|s| s.trim().parse::<f64>().ok())

View file

@ -13,5 +13,38 @@ pub(crate) fn reachable_offers(offers: Vec<Offer>, policy: &SelectionPolicy) ->
})
.filter(|o| o.host_id.map_or(true, |h| !blacklist.contains(&h)))
.filter(|o| o.verification.as_deref() != Some("deverified"))
.filter(|o| policy.max_dph_total.is_none_or(|max| o.dph_total <= max))
.collect()
}
#[cfg(test)]
mod tests {
use super::*;
fn offer(id: u64, dph_total: f64) -> Offer {
Offer {
id,
gpu_name: "Tesla T4".to_owned(),
dph_total,
gpu_ram: Some(16_000.0),
geolocation: Some("US".to_owned()),
inet_down_cost_per_tb: 0.0,
inet_up_cost_per_tb: 0.0,
host_id: Some(id),
verification: Some("unverified".to_owned()),
}
}
#[test]
fn max_price_cap_keeps_only_affordable_reachable_offers() {
let policy = SelectionPolicy {
max_dph_total: Some(0.10),
..SelectionPolicy::default()
};
let reachable = reachable_offers(vec![offer(1, 0.09), offer(2, 0.11)], &policy);
assert_eq!(reachable.len(), 1);
assert_eq!(reachable[0].id, 1);
}
}

View file

@ -47,6 +47,9 @@ pub async fn select_offer_pool_with_policy(
if let Some(min_ram) = policy.min_gpu_ram_mb {
query["gpu_ram"] = serde_json::json!({"gte": min_ram});
}
if let Some(max_dph_total) = policy.max_dph_total {
query["dph_total"] = serde_json::json!({"lte": max_dph_total});
}
if let Some(gpu_name) = policy.gpu_name.as_deref().filter(|s| !s.is_empty()) {
query["gpu_name"] = serde_json::json!({"eq": gpu_name});
}
@ -78,14 +81,19 @@ pub async fn select_offer_pool_with_policy(
let pool = rank_survivors(reachable, &cost, policy.drop_cheap_frac);
if pool.is_empty() {
return Err(
"no offers available (after quality/geo/host-blacklist filters and cheap-tail drop)"
.to_string(),
);
let cap = policy
.max_dph_total
.map_or_else(|| "uncapped".to_owned(), |max| format!("max ${max:.3}/hr"));
return Err(format!(
"no offers available ({cap}, after quality/geo/host-blacklist filters and cheap-tail drop)"
));
}
let cap = policy
.max_dph_total
.map_or_else(|| "uncapped".to_owned(), |max| format!("max ${max:.3}/hr"));
eprintln!(
"select_offer_pool: {} survivor(s) for {target_count} instance(s) after \
per-model {:.0}% cheap-drop (cheapest ${:.3}/hr eff)",
per-model {:.0}% cheap-drop ({cap}, cheapest ${:.3}/hr eff)",
pool.len(),
policy.drop_cheap_frac * 100.0,
cost.effective_price(&pool[0]),

View file

@ -62,6 +62,7 @@ pub struct SelectionPolicy {
pub require_verified: bool,
pub min_down_mbps: f64,
pub min_up_mbps: Option<f64>,
pub max_dph_total: Option<f64>,
pub blacklist_hosts: Vec<u64>,
pub drop_cheap_frac: f64,
pub image_size_gb: Option<f64>,
@ -76,6 +77,7 @@ impl Default for SelectionPolicy {
require_verified: false,
min_down_mbps: 100.0,
min_up_mbps: None,
max_dph_total: None,
blacklist_hosts: vec![59017],
drop_cheap_frac: 0.30,
image_size_gb: None,

View file

@ -5,6 +5,7 @@ edition = "2024"
[dependencies]
serde_json = "1"
blake3 = "1"
[target.'cfg(target_os = "linux")'.dependencies]
libc = "0.2"

File diff suppressed because it is too large Load diff