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

feat: working two stage pipeline parallel over vastai

Browse source
This commit is contained in:
Zachery Aaron Shores-Chmielewski 2026-07-24 12:59:32 +04:00
parent 18f06e34f9
commit 3c6845b170
13 changed files with 688 additions and 31 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -930,7 +930,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 +1005,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 +1016,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,6 +1032,8 @@ 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),
}
@ -1047,7 +1055,6 @@ def ring_readable(cmd: dict[str, Any]) -> None:
sequence=sequence,
extent=extent,
flags=flags,
payload=payload,
)
device_objects[handle] = materialized
control(
@ -1059,6 +1066,10 @@ 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"),
)
@ -1103,34 +1114,39 @@ 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
committed = write_record(
ring,
@ -1147,6 +1163,12 @@ 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,
)

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

@ -331,6 +331,7 @@ struct Config {
run_id: u64,
vastai_yes: bool,
vastai: Option<ResolvedVastAiConfig>,
model: ChatModelConfig,
pipeline_stages: u32,
max_tokens: u32,
skip_rebuild: bool,
@ -521,6 +522,7 @@ struct ChatTomlConfig {
observability: ChatObservabilityConfig,
image: ChatImageConfig,
vastai: ChatVastAiConfig,
model: ChatModelConfig,
relay: ChatRelayConfig,
}
@ -568,6 +570,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 +578,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 +702,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 +730,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 +810,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 +962,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 +1406,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,
@ -2185,6 +2246,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 +2268,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 +2997,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");

59
crates/mvp-system/src/bin/worker_node.rs
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,

8
crates/mvp-system/src/config.rs
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));

117
crates/mvp-system/src/orchestrator_app.rs
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(

116
crates/mvp-system/src/vastai_provisioning.rs
View file

@ -843,14 +843,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 +930,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 +1025,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 [

1
src/actor.rs
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()

1
src/runtime.rs
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

2
tools/vastai/src/config.rs
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())

33
tools/vastai/src/filters.rs
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);
}
}

18
tools/vastai/src/search.rs
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]),

2
tools/vastai/src/types.rs
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,

250
xtask/src/main.rs
View file

@ -22,6 +22,7 @@ const MVP_CHAT_CHECK_TIMEOUT_SECS: u64 = 900;
const MVP_CHAT_CHECK_POLL_MS: u64 = 100;
const MVP_CHAT_CHECK_TERM_GRACE_MS: u64 = 2_000;
const MVP_CHAT_CHECK_PROMPTS: &[u8] = b"ping\nsecond prompt\n";
const DATA_PATH_MIN_PAYLOAD_BYTES: u64 = 512;
struct MvpChatCheckPaths {
root: PathBuf,
@ -105,6 +106,7 @@ impl MvpChatCheckScenario {
"--yes".to_owned(),
"--endpoint-addr-mask".to_owned(),
"relay-only".to_owned(),
"--skip-rebuild".to_owned(),
]);
}
args.extend([
@ -1258,7 +1260,9 @@ fn assert_dump_log_facts(
require_multinode_docker_network_facts(&facts)?;
}
if scenario == MvpChatCheckScenario::VastAi {
require_gpu_dump_log_facts(&facts)?;
require_vastai_network_facts(&facts)?;
require_vastai_data_path_facts(&facts)?;
}
Ok(events)
}
@ -1613,6 +1617,21 @@ struct DumpLogFacts {
gpu_pipeline_tokenizer_decode_ready: BTreeSet<u64>,
gpu_pipeline_tokens_decoded: BTreeSet<u64>,
gpu_pipeline_real_worker_step_seen: bool,
ring_installed_ingress: bool,
ring_installed_egress: bool,
ring_installed_ingress_edge_ids: BTreeSet<u64>,
ring_installed_egress_edge_ids: BTreeSet<u64>,
activation_object_loaded: bool,
activation_step_executed: bool,
activation_egress_ring_read: bool,
activation_ingress_ring_write: bool,
activation_iroh_edge_sent: bool,
activation_iroh_edge_read: bool,
activation_interstage_handoff: bool,
activation_edge_ids: BTreeSet<u64>,
iroh_read_edge_ids: BTreeSet<u64>,
max_activation_record_bytes: u64,
max_worker_command_bytes: u64,
docker_node_spec_worker_count: Option<u64>,
worker_iroh_ready: BTreeSet<u64>,
docker_worker_coordinator_join: BTreeSet<u64>,
@ -1640,6 +1659,7 @@ fn record_dump_log_event(
) -> Result<(), String> {
record_vastai_provision_dump_log_event(channel, event, facts);
record_gpu_dump_log_event(channel, event, facts);
record_data_path_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);
@ -1746,6 +1766,146 @@ fn record_vastai_provision_dump_log_event(channel: &str, event: &Value, facts: &
}
}
fn record_data_path_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);
match (channel, event_type) {
("mvp.worker.ring", Some("RingInstalled")) => {
let edge_id = event.get("edge_id").and_then(Value::as_u64);
match event.get("direction").and_then(Value::as_str) {
Some("ingress") => {
facts.ring_installed_ingress = true;
if let Some(edge_id) = edge_id {
facts.ring_installed_ingress_edge_ids.insert(edge_id);
}
}
Some("egress") => {
facts.ring_installed_egress = true;
if let Some(edge_id) = edge_id {
facts.ring_installed_egress_edge_ids.insert(edge_id);
}
}
_ => {}
}
}
("mvp.worker.ingress", Some("ObjectLoaded")) => {
let extent = event.get("extent").and_then(Value::as_u64).unwrap_or(0);
if event.get("kind").and_then(Value::as_str) == Some("activation")
|| extent >= DATA_PATH_MIN_PAYLOAD_BYTES
{
facts.activation_object_loaded = true;
facts.max_activation_record_bytes = facts.max_activation_record_bytes.max(extent);
if let Some(edge_id) = event.get("edge_id").and_then(Value::as_u64) {
record_activation_edge_id(facts, edge_id);
record_interstage_activation_handoff(
facts,
edge_id,
event.get("stage_index").and_then(Value::as_u64),
);
}
}
}
("mvp.worker.step", Some("StepExecuted")) => {
let payload_bytes = event
.get("payload_bytes")
.or_else(|| event.get("committed_bytes"))
.and_then(Value::as_u64)
.unwrap_or(0);
let output_kind = event.get("output_kind").and_then(Value::as_str);
let legacy_activation_sized_pipeline_output = output_kind.is_none()
&& event.get("execution_backend").and_then(Value::as_str) == Some("pipeline_stage")
&& payload_bytes >= DATA_PATH_MIN_PAYLOAD_BYTES;
if (output_kind == Some("activation") || legacy_activation_sized_pipeline_output)
&& payload_bytes >= DATA_PATH_MIN_PAYLOAD_BYTES
{
facts.activation_step_executed = true;
facts.max_activation_record_bytes =
facts.max_activation_record_bytes.max(payload_bytes);
}
}
(_, Some("NodeEvent")) => match (phase, status) {
(Some("worker_command_write"), Some("ready")) => {
if detail_str(event, "command_type").is_some_and(|command| {
matches!(command, "InstallRing" | "RingReadable" | "ExecuteStep")
}) && let Some(command_bytes) = detail_u64(event, "command_bytes")
{
facts.max_worker_command_bytes =
facts.max_worker_command_bytes.max(command_bytes);
}
}
(Some("egress_ring_read"), Some("ready")) => {
if detail_str(event, "edge_kind") == Some("Activation")
&& let Some(record_bytes) = detail_u64(event, "record_bytes")
&& record_bytes >= DATA_PATH_MIN_PAYLOAD_BYTES
{
facts.activation_egress_ring_read = true;
facts.max_activation_record_bytes =
facts.max_activation_record_bytes.max(record_bytes);
if let Some(edge_id) = detail_u64(event, "edge_id") {
record_activation_edge_id(facts, edge_id);
}
}
}
(Some("ingress_ring_write"), Some("ready")) => {
if detail_str(event, "edge_kind") == Some("Activation")
&& let Some(record_bytes) = detail_u64(event, "record_bytes")
&& record_bytes >= DATA_PATH_MIN_PAYLOAD_BYTES
{
facts.activation_ingress_ring_write = true;
facts.max_activation_record_bytes =
facts.max_activation_record_bytes.max(record_bytes);
if let Some(edge_id) = detail_u64(event, "edge_id") {
record_activation_edge_id(facts, edge_id);
}
}
}
(Some("iroh_edge_bytes_sent"), Some("ready")) => {
if detail_str(event, "edge_kind") == Some("Activation")
&& detail_u64(event, "bytes").is_some_and(|bytes| bytes > 0)
{
facts.activation_iroh_edge_sent = true;
if let Some(edge_id) = detail_u64(event, "edge_id") {
record_activation_edge_id(facts, edge_id);
}
}
}
(Some("iroh_edge_bytes_read"), Some("observed")) => {
if detail_u64(event, "bytes").is_some_and(|bytes| bytes > 0)
&& let Some(edge_id) = detail_u64(event, "edge_id")
{
facts.iroh_read_edge_ids.insert(edge_id);
if facts.activation_edge_ids.contains(&edge_id) {
facts.activation_iroh_edge_read = true;
}
}
}
_ => {}
},
_ => {}
}
}
fn record_activation_edge_id(facts: &mut DumpLogFacts, edge_id: u64) {
facts.activation_edge_ids.insert(edge_id);
if facts.iroh_read_edge_ids.contains(&edge_id) {
facts.activation_iroh_edge_read = true;
}
}
fn record_interstage_activation_handoff(
facts: &mut DumpLogFacts,
edge_id: u64,
stage_index: Option<u64>,
) {
if stage_index.is_some_and(|stage_index| stage_index > 0)
&& facts.ring_installed_ingress_edge_ids.contains(&edge_id)
&& facts.ring_installed_egress_edge_ids.contains(&edge_id)
{
facts.activation_interstage_handoff = true;
}
}
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);
@ -1939,8 +2099,8 @@ fn require_vastai_network_facts(facts: &DumpLogFacts) -> Result<(), String> {
require_dump_log_fact(
facts
.vastai_node_spec_worker_count
.is_some_and(|count| count >= 1),
"VastAI node_spec with workers",
.is_some_and(|count| count >= 2),
"VastAI node_spec with multiple workers",
)?;
require_dump_log_fact(
!facts.vastai_provision_start_nodes.is_empty(),
@ -1951,8 +2111,40 @@ fn require_vastai_network_facts(facts: &DumpLogFacts) -> Result<(), String> {
"VastAI provider_start",
)?;
require_dump_log_fact(
!facts.worker_iroh_ready.is_empty(),
"VastAI worker iroh_driver ready",
facts.worker_iroh_ready.len() >= 2,
"VastAI worker iroh_driver ready for multiple nodes",
)
}
fn require_vastai_data_path_facts(facts: &DumpLogFacts) -> Result<(), String> {
require_dump_log_fact(
facts.ring_installed_ingress,
"worker ingress ring installed",
)?;
require_dump_log_fact(facts.ring_installed_egress, "worker egress ring installed")?;
require_dump_log_fact(
facts.activation_object_loaded,
"activation object loaded from ingress ring",
)?;
require_dump_log_fact(
facts.activation_step_executed,
"activation-producing worker step executed",
)?;
let explicit_transport = facts.activation_egress_ring_read
&& facts.activation_iroh_edge_sent
&& facts.activation_iroh_edge_read
&& facts.activation_ingress_ring_write;
require_dump_log_fact(
explicit_transport || facts.activation_interstage_handoff,
"activation inter-stage transport evidence",
)?;
let payload_outsizes_observed_command = facts.max_worker_command_bytes > 0
&& facts.max_activation_record_bytes > facts.max_worker_command_bytes;
let activation_sized_interstage_handoff = facts.activation_interstage_handoff
&& facts.max_activation_record_bytes >= DATA_PATH_MIN_PAYLOAD_BYTES;
require_dump_log_fact(
payload_outsizes_observed_command || activation_sized_interstage_handoff,
"activation payload not carried as worker JSON command",
)
}
@ -2134,6 +2326,7 @@ mod tests {
"--yes",
"--endpoint-addr-mask",
"relay-only",
"--skip-rebuild",
"--run-id",
"42",
"--dump-logs=/tmp/mvp-chat-check.ndjson",
@ -2939,12 +3132,12 @@ mod tests {
#[test]
fn benchmark_observability_vastai_dump_facts_require_remote_provider_events() {
let mut events = dump_log_fact_events(false, false);
let mut events = dump_log_fact_events(true, false);
events.extend([
(
"mvp.orch.bootstrap",
stamped(
json!({"type":"OrchBootstrap","phase":"node_spec","status":"ready","run_id":9,"node_id":1,"detail":{"endpoint_addr_mask":"relay-only","provider":"vastai","worker_count":1}}),
json!({"type":"OrchBootstrap","phase":"node_spec","status":"ready","run_id":9,"node_id":1,"detail":{"endpoint_addr_mask":"relay-only","provider":"vastai","worker_count":2}}),
"mvp-orchestrator",
1_071,
71,
@ -2963,6 +3156,51 @@ mod tests {
73,
),
),
(
"mvp.node.bootstrap",
stamped(
json!({"type":"NodeEvent","phase":"iroh_driver","status":"ready","run_id":9,"node_id":2,"stage_index":0,"detail":{"endpoint_addr_mask":"relay-only","has_relay":true,"direct_addr_count":0}}),
"mvp-worker-node",
1_074,
74,
),
),
(
"mvp.worker.ring",
stamped(
json!({"type":"RingInstalled","run_id":9,"node_id":2,"stage_index":0,"ring_id":1,"direction":"egress","edge_id":77,"kind":"activation","max_extent":4096}),
"tinygrad-worker",
1_075,
75,
),
),
(
"mvp.worker.ring",
stamped(
json!({"type":"RingInstalled","run_id":9,"node_id":3,"stage_index":1,"ring_id":2,"direction":"ingress","edge_id":77,"kind":"activation","max_extent":4096}),
"tinygrad-worker",
1_076,
76,
),
),
(
"mvp.worker.step",
stamped(
json!({"type":"StepExecuted","run_id":9,"node_id":2,"stage_index":0,"execution_backend":"pipeline_stage","committed_bytes":4096}),
"tinygrad-worker",
1_078,
78,
),
),
(
"mvp.worker.ingress",
stamped(
json!({"type":"ObjectLoaded","run_id":9,"node_id":3,"stage_index":1,"edge_id":77,"kind":"activation","extent":4056}),
"tinygrad-worker",
1_083,
83,
),
),
]);
let path = write_synthetic_event_dump("vastai-remote-provider", events);
assert_dump_log_facts(&path, MvpChatCheckScenario::VastAi)