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

feat: successful 8 stage pipeline parallel run, more metrics

Browse source
This commit is contained in:
Zachery Aaron Shores-Chmielewski 2026-07-25 20:04:57 +04:00
parent f8beca8c5d
commit d3a67d021f
15 changed files with 1366 additions and 79 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -1061,6 +1061,7 @@ def ring_readable(cmd: dict[str, Any]) -> None:
materialized = materialize_object(payload, sequence, flags)
materialized.update(
object_id=object_id,
edge_id=ring["edge_id"],
sequence=sequence,
extent=extent,
flags=flags,
@ -1127,6 +1128,11 @@ def execute_step(cmd: dict[str, Any]) -> None:
final_stage = bool(cmd.get("final_stage"))
input_kind = obj.get("kind")
input_extent = int(obj.get("extent", 0))
validation_ready = time.monotonic()
input_prepare_ms = 0
forward_ms = 0
realize_ms = 0
payload_pack_ms = 0
execution_backend = "pipeline_stage"
if not isinstance(model, PipelineStageTinygradModel):
execution_backend = "full_transformer"
@ -1136,29 +1142,50 @@ def execute_step(cmd: dict[str, Any]) -> None:
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()
forward_started = time.monotonic()
token_array = model(obj["tensor"], int(obj.get("start_pos", 0))).realize().numpy().reshape(-1)
forward_ready = time.monotonic()
forward_ms = int((forward_ready - forward_started) * 1000)
token = int(token_array[0])
payload_started = time.monotonic()
payload = struct.pack("<I", token)
output_kind = "token"
flags = 1 if token == int(loaded.get("eos_token_id", 0)) else 0
payload_pack_ms = int((time.monotonic() - payload_started) * 1000)
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_started = time.monotonic()
input_tensor = model.token_hidden(obj["tensor"]) if input_kind == "tokens" else obj["tensor"]
input_ready = time.monotonic()
input_prepare_ms = int((input_ready - input_started) * 1000)
forward_started = time.monotonic()
hidden = model.forward_hidden(input_tensor, int(obj.get("start_pos", 0)))
forward_ready = time.monotonic()
forward_ms = int((forward_ready - forward_started) * 1000)
if final_stage:
realize_started = time.monotonic()
token_array = model.next_token(hidden).realize().numpy().reshape(-1)
realize_ready = time.monotonic()
realize_ms = int((realize_ready - realize_started) * 1000)
token = int(token_array[0])
payload_started = time.monotonic()
payload = struct.pack("<I", token)
output_kind = "token"
flags = 1 if token == int(loaded.get("eos_token_id", 0)) else 0
payload_pack_ms = int((time.monotonic() - payload_started) * 1000)
else:
import numpy as np
realize_started = time.monotonic()
activation = hidden.realize().numpy().astype(np.float16, copy=False)
realize_ready = time.monotonic()
realize_ms = int((realize_ready - realize_started) * 1000)
payload_started = time.monotonic()
payload = activation.tobytes()
output_kind = "activation"
flags = 0
payload_pack_ms = int((time.monotonic() - payload_started) * 1000)
compute_ready = time.monotonic()
committed = write_record(
ring,
@ -1172,6 +1199,8 @@ def execute_step(cmd: dict[str, Any]) -> None:
type="StepExecuted",
step_id=int(cmd["step_id"]),
ring_id=output_ring_id,
role_id=int(cmd["role_id"]),
stage_index=max(0, int(cmd["role_id"]) - 1),
object_id=int(cmd["output_object_id"]),
sequence=int(cmd["output_sequence"]),
committed_bytes=committed,
@ -1182,6 +1211,15 @@ def execute_step(cmd: dict[str, Any]) -> None:
output_kind=output_kind,
payload_bytes=len(payload),
record_bytes=committed,
input_handle_id=handle,
input_object_id=int(cmd["input_object_id"]),
input_sequence=int(cmd["input_sequence"]),
input_edge_id=obj.get("edge_id"),
input_prepare_ms=input_prepare_ms,
model_forward_ms=forward_ms,
output_realize_ms=realize_ms,
payload_pack_ms=payload_pack_ms,
validation_ms=int((validation_ready - step_started) * 1000),
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),

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

@ -848,8 +848,6 @@ impl Config {
}
if let Some(vastai) = &self.vastai {
args.extend([
"--vastai-api-key".to_owned(),
vastai.api_key.clone(),
"--vastai-bootstrap-command".to_owned(),
vastai.bootstrap_command.clone(),
"--no-vastai-confirm-lease".to_owned(),
@ -1255,6 +1253,9 @@ impl OrchChild {
.stdin(Stdio::null())
.stdout(Stdio::null())
.stderr(Stdio::null());
if let Some(vastai) = &config.vastai {
command.env("VAST_API_KEY", &vastai.api_key);
}
#[cfg(all(target_os = "linux", not(test)))]
unsafe {
command.pre_exec(|| {
@ -2838,6 +2839,18 @@ bootstrap_command = "boot"
assert_eq!(vastai.relay_url, "https://relay.example");
assert_eq!(vastai.bootstrap_command, "boot");
assert_eq!(vastai.image, "docker.io/acme/node:latest");
let args = config.orchestrator_cli_args("docker.io/acme/node:latest");
assert!(
!args
.iter()
.any(|arg| arg == "--vastai-api-key" || arg == "secret"),
"Vast.ai API key must not be exposed in orchestrator argv: {args:?}"
);
assert!(
args.windows(2)
.any(|pair| pair == ["--vastai-bootstrap-command", "boot"]),
"non-secret Vast.ai config should still be forwarded"
);
},
);
}

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

@ -784,6 +784,7 @@ impl WorkerEdgeRuntime {
outbound.kind,
mvp_system::actors::node_agent::StageEdgeKindWire::TokenOut
);
let step_started = Instant::now();
let mut pump = || pump_network(driver, stack);
let committed_bytes = worker.execute_step(
u64::from(config.stage_index) + 1,
@ -800,6 +801,8 @@ impl WorkerEdgeRuntime {
datastream,
&mut pump,
)?;
let helper_execute_ms = duration_ms_u64(step_started.elapsed());
let egress_read_started = Instant::now();
let record = {
let arena = arena_manager.lock();
let lease = arena
@ -809,6 +812,7 @@ impl WorkerEdgeRuntime {
.read_arena(lease.layout.data_offset, committed_bytes)
.map_err(|e| format!("read egress ring: {e}"))?
};
let egress_read_ms = duration_ms_u64(egress_read_started.elapsed());
let record_bytes = record.len();
emit_node_event(
datastream,
@ -820,9 +824,17 @@ impl WorkerEdgeRuntime {
"edge_id":outbound.edge_id,
"edge_kind":format!("{:?}", outbound.kind),
"ring_id":output_ring_id,
"step_id":step_id,
"input_object_id":object_id,
"input_edge_id":input_edge_id,
"sequence":sequence,
"output_object_id":output_object_id,
"output_sequence":sequence,
"record_bytes":record_bytes,
"committed_bytes":committed_bytes,
"final_stage":final_stage,
"helper_execute_ms":helper_execute_ms,
"egress_ring_read_ms":egress_read_ms,
}),
);
self.driver_model
@ -838,7 +850,9 @@ impl WorkerEdgeRuntime {
.outbound_sender
.as_ref()
.ok_or_else(|| "outbound edge sender missing".to_owned())?;
let edge_send_started = Instant::now();
sender.send(record)?;
let edge_send_ms = duration_ms_u64(edge_send_started.elapsed());
emit_node_event(
datastream,
config,
@ -848,8 +862,12 @@ impl WorkerEdgeRuntime {
json!({
"edge_id":outbound.edge_id,
"edge_kind":format!("{:?}", outbound.kind),
"step_id":step_id,
"object_id":output_object_id,
"sequence":sequence,
"bytes":record_bytes,
"record_bytes":record_bytes,
"send_ms":edge_send_ms,
}),
);
stack
@ -891,28 +909,31 @@ impl WorkerEdgeRuntime {
if inbound.edge_id != edge_id {
return Ok(());
}
let records = {
let (records, buffered_bytes) = {
let buffer = self.ingress_streams.entry(stream_id).or_default();
buffer.extend_from_slice(&bytes);
let buffered_bytes = buffer.len();
let mut records = Vec::new();
while let Some(record) = take_complete_ingress_record(buffer, inbound.object_spec)? {
records.push(record);
}
records
(records, buffered_bytes)
};
for record in records {
let ring_id = self
.inbound_ring_id
.ok_or_else(|| "inbound ring missing".to_owned())?;
let ring_write_started = Instant::now();
{
let arena = arena_manager.lock();
let lease = arena
.lookup_lease(arena::RingId(ring_id))
.ok_or_else(|| format!("inbound ring {ring_id} lease missing"))?;
arena
.write_arena(lease.layout.data_offset, &record)
.write_arena(lease.layout.data_offset, &record.bytes)
.map_err(|e| format!("write ingress ring: {e}"))?;
}
let ingress_ring_write_ms = duration_ms_u64(ring_write_started.elapsed());
emit_node_event(
datastream,
config,
@ -923,9 +944,18 @@ impl WorkerEdgeRuntime {
"edge_id":edge_id,
"edge_kind":format!("{:?}", inbound.kind),
"ring_id":ring_id,
"record_bytes":record.len(),
"stream_id":stream_id,
"object_id":record.object_id,
"sequence":record.sequence,
"extent":record.extent,
"begin_sequence":record.begin_sequence,
"end_of_sequence":record.end_of_sequence,
"record_bytes":record.bytes.len(),
"ingress_buffer_bytes":buffered_bytes,
"ingress_ring_write_ms":ingress_ring_write_ms,
}),
);
let object_load_started = Instant::now();
let loaded = worker.ring_readable(
ring_id,
edge_id,
@ -934,11 +964,30 @@ impl WorkerEdgeRuntime {
datastream,
&mut || {},
)?;
let object_load_ms = duration_ms_u64(object_load_started.elapsed());
let key = ObjectKey {
edge_id,
object_id: loaded.object_id,
};
self.object_handles.insert(key, loaded.clone());
emit_node_event(
datastream,
config,
NODE_STAGE_CHANNEL,
"object_loaded",
"ready",
json!({
"edge_id":edge_id,
"edge_kind":format!("{:?}", inbound.kind),
"ring_id":ring_id,
"stream_id":stream_id,
"object_id":loaded.object_id,
"sequence":loaded.sequence,
"handle_generation":loaded.handle_generation,
"handle_id":loaded.handle_id,
"object_load_ms":object_load_ms,
}),
);
stack
.runtime
.send_to(
@ -1270,6 +1319,9 @@ impl WorkerEdgeRuntime {
}
}
fn duration_ms_u64(duration: Duration) -> u64 {
u64::try_from(duration.as_millis()).unwrap_or(u64::MAX)
}
fn edge_object_spec(spec: StageObjectSpecWire) -> edge::ObjectSpec {
edge::ObjectSpec {
kind: edge::ObjectKind::Activation,
@ -1294,17 +1346,34 @@ fn ingress_object_spec(spec: StageObjectSpecWire) -> ingress::ObjectSpec {
}
}
struct IngressRecordBytes {
bytes: Vec<u8>,
object_id: u64,
sequence: u64,
extent: u64,
begin_sequence: bool,
end_of_sequence: bool,
}
fn take_complete_ingress_record(
buffer: &mut Vec<u8>,
spec: StageObjectSpecWire,
) -> Result<Option<Vec<u8>>, String> {
) -> Result<Option<IngressRecordBytes>, String> {
let record = match ingress::read_object_record(buffer, ingress_object_spec(spec), false)
.map_err(|reason| format!("invalid object record: {reason:?}"))?
{
ingress::ObjectRecordRead::Incomplete => return Ok(None),
ingress::ObjectRecordRead::Complete(record) => record,
};
Ok(Some(buffer.drain(..record.total_len).collect()))
let bytes = buffer.drain(..record.total_len).collect();
Ok(Some(IngressRecordBytes {
bytes,
object_id: record.object_id.0,
sequence: record.sequence,
extent: record.extent,
begin_sequence: record.flags.begin_sequence,
end_of_sequence: record.flags.end_of_sequence,
}))
}
fn value_u64(value: &Value, field: &str) -> Result<u64, String> {

70
crates/mvp-system/src/orchestrator_app.rs
View file

@ -702,11 +702,13 @@ impl VastAiRuntimeConfig {
"ssh_user": &self.provisioning.ssh_user,
"gpu_name": &self.provisioning.selection.gpu_name,
"min_gpu_ram_mb": self.provisioning.selection.min_gpu_ram_mb,
"min_compute_cap": self.provisioning.selection.min_compute_cap,
"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,
"state_timeout_secs": self.provisioning.lifecycle.state_timeout.as_secs(),
"confirm_lease": self.provisioning.confirm_lease,
"has_api_key": self.api_key.is_some(),
"has_onstart": self.provisioning.onstart.is_some(),
@ -2975,6 +2977,7 @@ fn wait_for_weights_loaded_count(
let mut last_resend = Instant::now();
let mut active_stage = None::<u32>;
let mut resend_attempt = 0_u64;
let mut stage_resend_counts = BTreeMap::<u32, u64>::new();
loop {
pump(driver, stack, frame_tx);
emit_swim_transitions(orch_datastream, dashboard, run_id, node_id, stack);
@ -3002,6 +3005,12 @@ fn wait_for_weights_loaded_count(
));
};
let stage_node_id = stage.node_id.0;
let stage_send_count = {
let count = stage_resend_counts.entry(stage.stage_index).or_default();
*count += 1;
*count
};
let emit_wait_headline = stage_send_count == 1 || stage_send_count % 15 == 0;
let ready = readies.get(&stage_node_id).ok_or_else(|| {
format!("missing runtime-ready node for stage {}", stage.stage_index)
})?;
@ -3015,6 +3024,7 @@ fn wait_for_weights_loaded_count(
"attempt":resend_attempt,
"stage_count":pipeline_plan.stages.len(),
"stage_index":stage.stage_index,
"stage_send_count":stage_send_count,
"loaded_stage_count":loaded_stages.len()
}),
);
@ -3040,6 +3050,29 @@ fn wait_for_weights_loaded_count(
"route_matches_ready":route_owner == Some(ready.swim_node_id),
}),
);
if emit_wait_headline {
orch_datastream.emit_bootstrap(
dashboard,
run_id,
node_id,
"stage_provision_wait",
"observed",
json!({
"attempt":resend_attempt,
"stage_count":pipeline_plan.stages.len(),
"stage_index":stage.stage_index,
"stage_node_id":stage_node_id,
"stage_send_count":stage_send_count,
"loaded_stage_count":loaded_stages.len(),
"message":format!(
"loaded {} of {}; waiting on stage {}",
loaded_stages.len(),
pipeline_plan.stages.len(),
stage.stage_index
)
}),
);
}
provision_stage_from_plan(
stack,
ready.node_actor,
@ -3120,7 +3153,7 @@ fn next_pipeline_weight_load_stage<'a>(
.stages
.iter()
.filter(|stage| !loaded_stages.contains(&stage.stage_index))
.min_by_key(|stage| stage.stage_index)
.max_by_key(|stage| stage.stage_index)
}
struct FailedProvisionPlugin;
@ -4115,7 +4148,7 @@ impl PipelinePromptRuntime {
request_id,
"pipeline_token_in",
"started",
json!({"edge_id":self.token_in_edge_id,"sequence":sequence,"tokens":tokens.len(),"begin_sequence":true}),
json!({"edge_id":self.token_in_edge_id,"sequence":sequence,"tokens":tokens.len(),"begin_sequence":true,"token_count":tokens.len(),"token_ids":&tokens}),
);
self.send_token_in(sequence, &tokens, true)?;
orch_datastream.emit_prompt(
@ -4125,7 +4158,7 @@ impl PipelinePromptRuntime {
request_id,
"pipeline_token_in",
"ready",
json!({"edge_id":self.token_in_edge_id,"sequence":sequence,"begin_sequence":true}),
json!({"edge_id":self.token_in_edge_id,"sequence":sequence,"begin_sequence":true,"token_count":tokens.len()}),
);
Ok(())
}
@ -4200,7 +4233,26 @@ impl PipelinePromptRuntime {
self.active = None;
return Ok(());
}
self.send_token_in(self.next_sequence, &[pending.token_id], false)?;
let sequence = self.next_sequence;
orch_datastream.emit_prompt(
dashboard,
run_id,
node_id,
request_id,
"pipeline_token_in",
"started",
json!({"edge_id":self.token_in_edge_id,"sequence":sequence,"tokens":1,"begin_sequence":false,"token_count":1,"token_id":pending.token_id}),
);
self.send_token_in(sequence, &[pending.token_id], false)?;
orch_datastream.emit_prompt(
dashboard,
run_id,
node_id,
request_id,
"pipeline_token_in",
"ready",
json!({"edge_id":self.token_in_edge_id,"sequence":sequence,"begin_sequence":false,"token_count":1,"token_id":pending.token_id}),
);
Ok(())
}
fn send_token_in(
@ -4317,7 +4369,7 @@ impl PipelinePromptRuntime {
request_id,
"pipeline_token_out",
"observed",
json!({"edge_id":self.token_out_edge_id,"object_id":record.object_id,"sequence":record.sequence,"token_id":record.token_id,"eos":record.eos}),
json!({"edge_id":self.token_out_edge_id,"object_id":record.object_id,"sequence":record.sequence,"token_id":record.token_id,"eos":record.eos,"generated_index":self.generated_tokens.len() + 1}),
);
self.generated_tokens.push(record.token_id);
let reached_limit = self.generated_tokens.len() as u32 >= active.request.max_tokens;
@ -4328,7 +4380,7 @@ impl PipelinePromptRuntime {
request_id,
"pipeline_tokenizer_decode",
"started",
json!({"node_actor":self.tokenizer_decode_actor,"reply_to":self.tokenizer_reply_to,"token_id":record.token_id}),
json!({"node_actor":self.tokenizer_decode_actor,"reply_to":self.tokenizer_reply_to,"token_id":record.token_id,"sequence":record.sequence,"generated_index":self.generated_tokens.len()}),
);
self.request_decode(
runtime,
@ -5325,13 +5377,13 @@ mod tests {
let first =
next_pipeline_weight_load_stage(&plan, &loaded, None).expect("first stage selected");
assert_eq!(first.stage_index, 0);
assert_eq!(first.stage_index, 6);
let resent = next_pipeline_weight_load_stage(&plan, &loaded, Some(first.stage_index))
.expect("active stage is resent before it loads");
assert_eq!(resent.stage_index, 0);
assert_eq!(resent.stage_index, 6);
for expected_stage in 0..7 {
for expected_stage in (0..7).rev() {
let active = next_pipeline_weight_load_stage(&plan, &loaded, None)
.expect("next unloaded stage selected");
assert_eq!(active.stage_index, expected_stage);

46
crates/mvp-system/src/tests/vastai_provisioning_guarantees.rs
View file

@ -36,6 +36,7 @@ struct FakeLeaseClient {
destroyed: Vec<u64>,
destroy_result: Option<Result<(), String>>,
next_contract_id: u64,
host_ids: VecDeque<Option<u64>>,
}
impl FakeLeaseClient {
@ -50,11 +51,12 @@ impl VastAiLeaseClient for FakeLeaseClient {
self.requests.push(request);
let contract_id = self.next_contract_id;
self.next_contract_id = self.next_contract_id.wrapping_add(1).max(1);
let host_id = self.host_ids.pop_front().unwrap_or(Some(77));
Ok(ProvisionedInstance {
index: 0,
contract_id,
offer_id: 55,
host_id: Some(77),
host_id,
gpu_name: "RTX 4090".to_owned(),
gpu_ram: Some(24_000.0),
dph_total: 0.42,
@ -238,6 +240,38 @@ fn vastai_plugin_omits_ssh_public_key_when_unconfigured() {
assert_eq!(request.env.get("SSH_PUBLIC_KEY"), None);
}
#[test]
fn pipeline_starts_blacklist_hosts_already_leased_in_run() {
let mut client = FakeLeaseClient::default().with_contract(100);
client.host_ids.extend([Some(77), Some(88)]);
let mut plugin = VastAiProvisioningPlugin::new(client, FakeBootstrap::default(), config());
let first = spec();
let mut second = spec();
second.node_id = 12;
second.stage_index = Some(3);
let first_handle = plugin.start_node(first, sink()).unwrap();
let second_handle = plugin.start_node(second, sink()).unwrap();
let requests = &plugin.client().requests;
assert_eq!(requests.len(), 2);
assert!(
!requests[0].selection.blacklist_hosts.contains(&77),
"first node should not preemptively blacklist the host it has not leased"
);
assert!(
requests[1].selection.blacklist_hosts.contains(&77),
"second node should avoid the first node's Vast.ai host"
);
assert!(
requests[1].selection.blacklist_hosts.contains(&59017),
"existing operator blacklist must be preserved"
);
plugin.stop_node(&second_handle).unwrap();
plugin.stop_node(&first_handle).unwrap();
}
#[test]
fn stop_destroys_known_vastai_contract_exactly_once() {
let mut plugin = VastAiProvisioningPlugin::new(
@ -259,7 +293,7 @@ fn stop_destroys_known_vastai_contract_exactly_once() {
}
#[test]
fn vastai_complete_bootstrap_stops_bootstrap_without_destroying_contract() {
fn vastai_complete_bootstrap_keeps_log_tail_until_node_stop() {
let mut plugin = VastAiProvisioningPlugin::new(
FakeLeaseClient::default().with_contract(100),
FakeBootstrap::default(),
@ -269,9 +303,9 @@ fn vastai_complete_bootstrap_stops_bootstrap_without_destroying_contract() {
plugin.complete_bootstrap(&handle).unwrap();
assert_eq!(
plugin.bootstrap().stops,
vec![(1, BootstrapStopReason::RuntimeReady)]
assert!(
plugin.bootstrap().stops.is_empty(),
"runtime-ready completion should keep the SSH log tail alive"
);
assert_eq!(plugin.client().destroyed, Vec::<u64>::new());
assert_eq!(plugin.active_contract_count(), 1);
@ -281,7 +315,7 @@ fn vastai_complete_bootstrap_stops_bootstrap_without_destroying_contract() {
assert_eq!(plugin.client().destroyed, vec![100]);
assert_eq!(
plugin.bootstrap().stops,
vec![(1, BootstrapStopReason::RuntimeReady)]
vec![(1, BootstrapStopReason::NodeStop)]
);
assert_eq!(plugin.active_contract_count(), 0);
}

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

@ -57,6 +57,7 @@ mod tests {
gpu_name: "RTX 4090".to_owned(),
dph_total: 0.375,
gpu_ram: Some(23.5),
compute_cap: 890,
geolocation: Some("US".to_owned()),
inet_down_cost_per_tb: 0.0,
inet_up_cost_per_tb: 0.0,

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

@ -1,5 +1,5 @@
use parking_lot::Mutex;
use std::collections::BTreeMap;
use std::collections::{BTreeMap, BTreeSet};
use std::path::{Path, PathBuf};
use std::process::{Child, Command, Stdio};
use std::sync::{
@ -551,15 +551,14 @@ fn spawn_retrying_ssh_bootstrap(
match wait_result {
Some(Ok(status)) => {
let line = if status.success() {
format!(
"VastAI SSH bootstrap exited before runtime ready: {status}; retrying"
)
let readiness = if status.success() {
"exited before runtime ready"
} else {
format!(
"VastAI SSH bootstrap failed before runtime ready: {status}; retrying"
)
"not ready before runtime ready"
};
let line = format!(
"VastAI SSH bootstrap {readiness} (attempt {attempt}, {status}); retrying"
);
sink.observe(PluginObservation::ProviderLine {
run_id,
node_id,
@ -571,7 +570,9 @@ fn spawn_retrying_ssh_bootstrap(
sink.observe(PluginObservation::ProviderLine {
run_id,
node_id,
line: format!("wait VastAI SSH bootstrap: {error}; retrying"),
line: format!(
"wait VastAI SSH bootstrap attempt {attempt}: {error}; retrying"
),
});
break;
}
@ -583,7 +584,9 @@ fn spawn_retrying_ssh_bootstrap(
sink.observe(PluginObservation::ProviderLine {
run_id,
node_id,
line: format!("spawn VastAI SSH bootstrap failed: {error}; retrying"),
line: format!(
"spawn VastAI SSH bootstrap attempt {attempt} failed: {error}; retrying"
),
});
}
}
@ -594,7 +597,10 @@ fn spawn_retrying_ssh_bootstrap(
sink.observe(PluginObservation::ProviderLine {
run_id,
node_id,
line: format!("VastAI SSH bootstrap retrying in {}s", backoff.as_secs()),
line: format!(
"VastAI SSH bootstrap retrying in {}s after attempt {attempt}",
backoff.as_secs()
),
});
std::thread::sleep(backoff);
backoff = next_ssh_backoff(backoff);
@ -670,6 +676,7 @@ where
bootstrap: B,
config: VastAiProvisioningConfig,
bootstrap_producer: Option<DatastreamProducer>,
leased_host_ids: BTreeSet<u64>,
next_handle_id: u64,
nodes: BTreeMap<u64, VastAiNode<B::Handle>>,
}
@ -677,6 +684,7 @@ where
struct VastAiNode<H> {
contract_id: u64,
bootstrap: Option<H>,
host_id: Option<u64>,
}
impl<C, B> VastAiProvisioningPlugin<C, B>
@ -691,6 +699,7 @@ where
config,
bootstrap_producer: None,
next_handle_id: 1,
leased_host_ids: BTreeSet::new(),
nodes: BTreeMap::new(),
}
}
@ -741,6 +750,13 @@ where
{
env.insert("SSH_PUBLIC_KEY".to_owned(), key.to_owned());
}
let mut selection = self.config.selection.clone();
for host_id in &self.leased_host_ids {
if !selection.blacklist_hosts.contains(host_id) {
selection.blacklist_hosts.push(*host_id);
}
}
ProvisionRequest {
count: 1,
image: spec.image.clone(),
@ -749,7 +765,7 @@ where
env,
per_instance_env: vec![BTreeMap::new()],
onstart: self.config.onstart.clone(),
selection: self.config.selection.clone(),
selection,
lifecycle: self.config.lifecycle.clone(),
confirm_lease: self.config.confirm_lease,
}
@ -860,6 +876,11 @@ where
}
};
let host_id = instance.host_id;
if let Some(host_id) = host_id {
self.leased_host_ids.insert(host_id);
}
let handle = PluginNodeHandle {
id: self.next_handle_id,
provider_process_id: None,
@ -870,6 +891,7 @@ where
VastAiNode {
contract_id: instance.contract_id,
bootstrap: Some(bootstrap),
host_id,
},
);
Ok(handle)
@ -883,6 +905,9 @@ where
let Some(mut node) = self.nodes.remove(&handle.id) else {
return Ok(());
};
if let Some(host_id) = node.host_id {
self.leased_host_ids.remove(&host_id);
}
if let Some(mut bootstrap) = node.bootstrap.take() {
self.bootstrap
.stop_bootstrap(&mut bootstrap, BootstrapStopReason::NodeStop);

2
tools/vastai/Cargo.toml
View file

@ -8,7 +8,7 @@ publish = false
reqwest = { version = "0.12", features = ["json"] }
serde = { version = "1", features = ["derive"] }
serde_json = "1"
tokio = { version = "1", features = ["time"] }
tokio = { version = "1", features = ["macros", "rt", "time"] }
urlencoding = "2"
[dev-dependencies]

8
tools/vastai/src/config.rs
View file

@ -4,6 +4,7 @@ use crate::types::{LifecyclePolicy, SelectionPolicy};
pub const ENV_IMAGE_SIZE_GB: &str = "PP_IMAGE_SIZE_GB";
pub const ENV_GPU_MIN_RAM_MB: &str = "PP_GPU_MIN_RAM_MB";
pub const ENV_MIN_COMPUTE_CAP: &str = "PP_MIN_COMPUTE_CAP";
pub const ENV_MIN_INET_DOWN_MBPS: &str = "PP_MIN_INET_DOWN_MBPS";
pub const ENV_MIN_INET_UP_MBPS: &str = "PP_MIN_INET_UP_MBPS";
pub const ENV_MIN_RELIABILITY: &str = "PP_MIN_RELIABILITY";
@ -11,6 +12,7 @@ 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_STATE_TIMEOUT_SECS: &str = "PP_STATE_TIMEOUT_SECS";
pub const ENV_BLACKLIST_HOSTS: &str = "PP_BLACKLIST_HOSTS";
pub const ENV_ASSUME_YES: &str = "PP_ASSUME_YES";
@ -48,6 +50,9 @@ impl SelectionPolicy {
pub fn from_env() -> Self {
let mut policy = Self::default();
policy.min_gpu_ram_mb = env_positive_u64(ENV_GPU_MIN_RAM_MB);
if let Some(min_compute_cap) = env_positive_u64(ENV_MIN_COMPUTE_CAP) {
policy.min_compute_cap = Some(min_compute_cap);
}
policy.min_down_mbps = env_nonnegative_f64(ENV_MIN_INET_DOWN_MBPS, 100.0);
policy.min_reliability = std::env::var(ENV_MIN_RELIABILITY)
.ok()
@ -84,6 +89,9 @@ impl LifecyclePolicy {
.unwrap_or(600),
);
policy.poll_interval = poll_interval;
if let Some(state_timeout_secs) = env_positive_u64(ENV_STATE_TIMEOUT_SECS) {
policy.state_timeout = Duration::from_secs(state_timeout_secs);
}
policy
}
}

19
tools/vastai/src/filters.rs
View file

@ -14,6 +14,11 @@ 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))
.filter(|o| {
policy
.min_compute_cap
.is_none_or(|min_compute_cap| o.compute_cap >= min_compute_cap)
})
.collect()
}
@ -27,6 +32,7 @@ mod tests {
gpu_name: "Tesla T4".to_owned(),
dph_total,
gpu_ram: Some(16_000.0),
compute_cap: 750,
geolocation: Some("US".to_owned()),
inet_down_cost_per_tb: 0.0,
inet_up_cost_per_tb: 0.0,
@ -47,4 +53,17 @@ mod tests {
assert_eq!(reachable.len(), 1);
assert_eq!(reachable[0].id, 1);
}
#[test]
fn minimum_compute_cap_keeps_only_modern_cuda_offers() {
let mut legacy = offer(1, 0.03);
legacy.compute_cap = 520;
let mut modern = offer(2, 0.04);
modern.compute_cap = 750;
let reachable = reachable_offers(vec![legacy, modern], &SelectionPolicy::default());
assert_eq!(reachable.len(), 1);
assert_eq!(reachable[0].id, 2);
}
}

54
tools/vastai/src/monitor.rs
View file

@ -92,6 +92,12 @@ pub async fn wait_for_running_with_policy(
msg.unwrap_or_default()
));
}
if !policy.state_timeout.is_zero() && state_since.elapsed() >= policy.state_timeout {
return Err(format!(
"instance {contract_id} stuck in status {actual} for {}s",
policy.state_timeout.as_secs()
));
}
match actual {
"running" => {
@ -112,3 +118,51 @@ pub async fn wait_for_running_with_policy(
}
}
}
#[cfg(test)]
mod tests {
use std::time::Duration;
use serde_json::json;
use wiremock::matchers::{method, path};
use wiremock::{Mock, MockServer, ResponseTemplate};
use super::*;
#[tokio::test]
async fn stuck_loading_state_returns_error_instead_of_polling_forever() {
let server = MockServer::start().await;
Mock::given(method("GET"))
.and(path("/api/v0/instances/123/"))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
"instances": {
"actual_status": "loading",
"intended_status": "running",
"status_msg": "afad30e59d72: Already exists"
}
})))
.mount(&server)
.await;
let policy = LifecyclePolicy {
poll_interval: Duration::from_millis(1),
state_timeout: Duration::from_millis(5),
..LifecyclePolicy::default()
};
let error = wait_for_running_with_policy(
&reqwest::Client::new(),
&server.uri(),
"secret",
123,
&policy,
)
.await
.expect_err("stuck loading should be replaceable");
assert!(
error.contains("stuck in status loading"),
"error should name stuck provider state: {error}"
);
}
}

3
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(min_compute_cap) = policy.min_compute_cap {
query["compute_cap"] = serde_json::json!({"gte": min_compute_cap});
}
if let Some(max_dph_total) = policy.max_dph_total {
query["dph_total"] = serde_json::json!({"lte": max_dph_total});
}

5
tools/vastai/src/teardown.rs
View file

@ -9,7 +9,10 @@ pub async fn destroy_instance(
api_key: &str,
contract_id: u64,
) -> Result<(), String> {
let url = format!("{base_url}/api/v0/instances/{contract_id}/");
let url = format!(
"{base_url}/api/v0/instances/{contract_id}/?api_key={}",
urlencoding::encode(api_key)
);
let resp = client
.delete(&url)
.header("Authorization", format!("Bearer {api_key}"))

9
tools/vastai/src/types.rs
View file

@ -17,6 +17,9 @@ pub struct Offer {
pub dph_total: f64,
#[serde(default)]
pub gpu_ram: Option<f64>,
/// CUDA compute capability encoded as vast.ai reports it (`890` = 8.9).
#[serde(default)]
pub compute_cap: u64,
#[serde(default)]
pub geolocation: Option<String>,
/// Inbound bandwidth price ($/TB). vast.ai bills Docker image pulls here.
@ -58,6 +61,8 @@ pub struct LabeledInstance {
pub struct SelectionPolicy {
pub gpu_name: Option<String>,
pub min_gpu_ram_mb: Option<u64>,
/// Minimum CUDA compute capability encoded as vast.ai reports it (`700` = 7.0).
pub min_compute_cap: Option<u64>,
pub min_reliability: f64,
pub require_verified: bool,
pub min_down_mbps: f64,
@ -73,6 +78,7 @@ impl Default for SelectionPolicy {
Self {
gpu_name: None,
min_gpu_ram_mb: None,
min_compute_cap: Some(700),
min_reliability: 0.95,
require_verified: false,
min_down_mbps: 100.0,
@ -90,6 +96,8 @@ impl Default for SelectionPolicy {
pub struct LifecyclePolicy {
pub lease_pace: Duration,
pub poll_interval: Duration,
/// Maximum time to stay in one non-running provider state before replacing the lease.
pub state_timeout: Duration,
}
impl Default for LifecyclePolicy {
@ -97,6 +105,7 @@ impl Default for LifecyclePolicy {
Self {
lease_pace: Duration::from_millis(600),
poll_interval: Duration::from_secs(10),
state_timeout: Duration::from_secs(300),
}
}
}

1043
xtask/src/main.rs
View file

File diff suppressed because it is too large Load diff