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

Compare commits

base: zacheryasc:a5736730b5b172a30495f287b6187bd5a3bc9d1f
Branches Tags
zacheryasc:master
zacheryasc:development
zacheryasc:backup/pre-rewrite-msg
...
compare: zacheryasc:28a7292de2a798872386c2eb84bfc97417b93a24
Branches Tags
zacheryasc:development
zacheryasc:master
zacheryasc:backup/pre-rewrite-msg

2 commits
a5736730b5 ... 28a7292de2

Author SHA1 Message Date
Zachery Aaron Shores-Chmielewski
28a7292de2 fix: faster provisioning, better ssh checks 2026-07-27 12:30:49 +04:00
Zachery Aaron Shores-Chmielewski
155741bedb feat: working 7B inferenced over 4 pipeline stages 2026-07-27 10:47:17 +04:00
14 changed files with 2898 additions and 248 deletions
Show stats Expand all files Collapse all files

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

@ -149,13 +149,30 @@ def env_flag(name: str, default: bool = True) -> bool:
def benchmark_stamp() -> dict[str, Any]:
global _benchmark_seq
_benchmark_seq += 1
pid = os.getpid()
wall_ms = time.time_ns() // 1_000_000
mono_ms = int((time.monotonic() - _benchmark_start) * 1000)
return {
"schema": BENCHMARK_SCHEMA,
"schema_version": BENCHMARK_SCHEMA,
"component": "tinygrad-worker",
"pid": os.getpid(),
"producer_component": "tinygrad-worker",
"producer_instance_id": os.environ.get(
"MVP_BENCHMARK_PRODUCER_INSTANCE",
f"tinygrad-worker:{pid}",
),
"producer_process_id": pid,
"pid": pid,
"seq": _benchmark_seq,
"wall_unix_ms": time.time_ns() // 1_000_000,
"mono_ms": int((time.monotonic() - _benchmark_start) * 1000),
"producer_sequence": _benchmark_seq,
"wall_unix_ms": wall_ms,
"wall_clock_unix_ms": wall_ms,
"mono_ms": mono_ms,
"monotonic_ms": mono_ms,
"clock_source": {
"wall": "time.time_ns_unix_ms",
"monotonic": "time.monotonic_process_elapsed_ms",
},
}
@ -169,8 +186,45 @@ def env_int(name: str) -> int | None:
return None
def datastream_endpoint_snapshot() -> dict[str, Any]:
return {
"role": "python-worker-stdio-json-bridge",
"transport": "stdout-json-lines",
"endpoint_identity": os.environ.get("MVP_DATASTREAM_ENDPOINT_ID", "worker-stdio-bridge"),
"configured_source": "worker-node-env",
"resolved_source": "TinygradWorker::spawn environment",
"authentication_present": False,
"tls_present": False,
"relay_mode": os.environ.get("MVP_IROH_RELAY_MODE"),
"endpoint_addr_mask": os.environ.get("MVP_IROH_ENDPOINT_ADDR_MASK"),
"connectivity_result": "configured",
}
def apply_canonical_envelope(event: dict[str, Any]) -> None:
benchmark = event.setdefault("benchmark", benchmark_stamp())
event.setdefault("schema_version", BENCHMARK_SCHEMA)
event.setdefault("event_type", event.get("type"))
event.setdefault("event_name", event.get("phase", event.get("type")))
event.setdefault("producer_component", benchmark.get("producer_component", "tinygrad-worker"))
event.setdefault("producer_instance_id", benchmark.get("producer_instance_id"))
event.setdefault("producer_process_id", benchmark.get("producer_process_id", os.getpid()))
event.setdefault("producer_sequence", benchmark.get("producer_sequence", benchmark.get("seq")))
event.setdefault("wall_clock_unix_ms", benchmark.get("wall_clock_unix_ms", benchmark.get("wall_unix_ms")))
event.setdefault("monotonic_ms", benchmark.get("monotonic_ms", benchmark.get("mono_ms")))
event.setdefault("clock_source", benchmark.get("clock_source"))
event.setdefault("datastream_endpoint", datastream_endpoint_snapshot())
event.setdefault(
"span_id",
f"{event.get('producer_instance_id')}:{event.get('producer_sequence')}:{event.get('event_name')}",
)
if "parent_span_id" not in event:
request_id = event.get("request_id")
event["parent_span_id"] = f"request:{request_id}" if request_id is not None else None
def control(**event: Any) -> None:
event.setdefault("benchmark", benchmark_stamp())
apply_canonical_envelope(event)
if (run_id := env_int("MVP_RUN_ID")) is not None:
event.setdefault("run_id", run_id)
if (node_id := env_int("MVP_LOGICAL_NODE_ID")) is not None:
@ -1273,6 +1327,37 @@ def shutdown_worker(_: dict[str, Any]) -> None:
raise SystemExit(0)
def emit_python_datastream_preflight() -> None:
endpoint = datastream_endpoint_snapshot()
control(
type="PythonDatastreamConfigured",
phase="PythonDatastreamConfigured",
status="configured",
endpoint=endpoint,
)
control(
type="PythonDatastreamConnected",
phase="PythonDatastreamConnected",
status="ready",
endpoint=endpoint,
)
synthetic_id = f"python-{os.getpid()}-{_benchmark_seq + 1}"
control(
type="PythonDatastreamSyntheticEventSent",
phase="PythonDatastreamSyntheticEventSent",
status="sent",
endpoint=endpoint,
synthetic_id=synthetic_id,
)
control(
type="PythonDatastreamSyntheticEventObserved",
phase="PythonDatastreamSyntheticEventObserved",
status="observed",
endpoint=endpoint,
synthetic_id=synthetic_id,
)
HANDLERS = {
"InitializeWorker": initialize,
"ConfigureRole": configure_role,
@ -1288,6 +1373,8 @@ HANDLERS = {
"ShutdownWorker": shutdown_worker,
}
emit_python_datastream_preflight()
for raw in sys.stdin:
if not raw.strip():
continue

17
crates/mvp-system/src/benchmark_observability.rs
View file

@ -21,12 +21,25 @@ pub fn stamp(component: &'static str) -> Value {
let start = BENCHMARK_START.get_or_init(Instant::now);
let mono_ms = u64::try_from(start.elapsed().as_millis()).unwrap_or(u64::MAX);
let seq = BENCHMARK_SEQ.fetch_add(1, Ordering::Relaxed);
let pid = std::process::id();
let wall_ms = unix_ms_now();
json!({
"schema": BENCHMARK_SCHEMA,
"schema_version": BENCHMARK_SCHEMA,
"component": component,
"pid": std::process::id(),
"producer_component": component,
"producer_instance_id": format!("{component}:{pid}"),
"producer_process_id": pid,
"pid": pid,
"seq": seq,
"wall_unix_ms": unix_ms_now(),
"producer_sequence": seq,
"wall_unix_ms": wall_ms,
"wall_clock_unix_ms": wall_ms,
"mono_ms": mono_ms,
"monotonic_ms": mono_ms,
"clock_source": {
"wall": "system_unix_ms",
"monotonic": "process_elapsed_ms"
},
})
}

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

@ -1,3 +1,5 @@
#![recursion_limit = "256"]
use std::collections::BTreeMap;
use std::fs::{self, File, OpenOptions};
use std::io::{self, BufRead, BufReader, IsTerminal, Write};
@ -47,7 +49,7 @@ OPTIONS:
--process | --docker | --vastai
Select the runtime provider
--config <path> Load config overlay
--pipeline-stages <count> Number of pipeline stages
--pipeline-stages|--pipeline-parallel <count>
--relay-mode <mode> Relay mode: default or disabled
--relay-url <url> Custom relay URL passed to mvp-orchestrator
--endpoint-addr-mask <mask> Endpoint address mask: full or relay-only
@ -154,6 +156,7 @@ where
}),
);
progress.emit_benchmark_envelope(&config);
progress.emit_endpoint_config_snapshot(&config);
confirm_vastai_if_needed(&config)?;
let prepare_runtime_started = Instant::now();
progress.emit(
@ -411,12 +414,25 @@ impl ChatDatastream {
fn emit(&mut self, channel: &str, phase: &str, status: &str, detail: Value) {
let id = self.channel_by_name(channel);
let benchmark = benchmark_observability::stamp("mvp-chat");
let payload = serde_json::to_vec(&json!({
"schema_version": benchmark["schema_version"].clone(),
"type": "ChatProgress",
"event_type": "ChatProgress",
"event_name": phase,
"phase": phase,
"status": status,
"run_id": self.run_id,
"benchmark": benchmark_observability::stamp("mvp-chat"),
"producer_component": benchmark["producer_component"].clone(),
"producer_instance_id": benchmark["producer_instance_id"].clone(),
"producer_process_id": benchmark["producer_process_id"].clone(),
"producer_sequence": benchmark["producer_sequence"].clone(),
"wall_clock_unix_ms": benchmark["wall_clock_unix_ms"].clone(),
"monotonic_ms": benchmark["monotonic_ms"].clone(),
"clock_source": benchmark["clock_source"].clone(),
"span_id": format!("mvp-chat:{}:{}:{phase}", self.run_id, benchmark["producer_sequence"]),
"parent_span_id": Value::Null,
"benchmark": benchmark,
"detail": detail,
}))
.expect("serialize mvp-chat progress event");
@ -426,12 +442,25 @@ impl ChatDatastream {
fn emit_benchmark_envelope(&mut self, config: &Config) {
let id = self.channel_by_name(CHAT_BENCHMARK_CHANNEL);
let benchmark = benchmark_observability::stamp("mvp-chat");
let payload = serde_json::to_vec(&json!({
"schema_version": benchmark["schema_version"].clone(),
"type": "BenchmarkRunEnvelope",
"event_type": "BenchmarkRunEnvelope",
"event_name": "run_envelope",
"phase": "run_envelope",
"status": "ready",
"run_id": self.run_id,
"benchmark": benchmark_observability::stamp("mvp-chat"),
"producer_component": benchmark["producer_component"].clone(),
"producer_instance_id": benchmark["producer_instance_id"].clone(),
"producer_process_id": benchmark["producer_process_id"].clone(),
"producer_sequence": benchmark["producer_sequence"].clone(),
"wall_clock_unix_ms": benchmark["wall_clock_unix_ms"].clone(),
"monotonic_ms": benchmark["monotonic_ms"].clone(),
"clock_source": benchmark["clock_source"].clone(),
"span_id": format!("mvp-chat:{}:{}:run_envelope", self.run_id, benchmark["producer_sequence"]),
"parent_span_id": Value::Null,
"benchmark": benchmark,
"detail": {
"scenario": "mvp-chat",
"detail_level": "benchmark_observability_v1",
@ -487,6 +516,56 @@ impl ChatDatastream {
self.flush();
}
fn emit_endpoint_config_snapshot(&mut self, config: &Config) {
let endpoint = json!({
"role": "chat-frame-archive",
"transport": "datastream-frame-log",
"configured": config.datastream_frame_log.is_some(),
"archive_path": config.datastream_frame_log.as_ref().map(|path| path.to_string_lossy().to_string()),
});
let runtime_endpoint = json!({
"provider": config.provider.as_str(),
"relay_mode": config.relay_mode.as_deref(),
"relay_configured": config.relay_url.is_some(),
"endpoint_addr_mask": config.endpoint_addr_mask.as_str(),
});
let synthetic_id = format!("mvp-chat-{}-datastream-preflight", self.run_id);
for (phase, status) in [
("DatastreamProducerConfigured", "configured"),
("DatastreamProducerConnected", "ready"),
("DatastreamSyntheticEventSent", "sent"),
("DatastreamSyntheticEventObserved", "observed"),
] {
self.emit(
CHAT_BENCHMARK_CHANNEL,
phase,
status,
json!({
"producer": "mvp-chat",
"producer_class": "rust-chat",
"synthetic_id": synthetic_id,
"datastream_endpoint": endpoint,
"runtime_endpoint": runtime_endpoint,
}),
);
}
self.emit(
CHAT_BENCHMARK_CHANNEL,
"endpoint_config_snapshot",
"ready",
json!({
"producer": "mvp-chat",
"expected_producers": ["mvp-chat", "mvp-orchestrator", "mvp-worker-node", "tinygrad-worker"],
"datastream_endpoint": endpoint,
"runtime_endpoint": runtime_endpoint,
"connectivity_preflight": {
"status": "configured",
"canonical_datastream_required": true,
},
}),
);
}
fn flush(&mut self) {
let stream = self.stream.clone();
for frame in self.endpoint.mux().drain() {
@ -1034,7 +1113,10 @@ impl ParsedArgs {
"--config" => {
parsed.config_path = Some(PathBuf::from(next_arg(&mut args, "--config")?))
}
"--pipeline-stages" => {
"--pipeline-stages" | "--pipeline-parallel" => {
if parsed.pipeline_stages.is_some() {
return Err("pipeline stage count was provided more than once".to_owned());
}
parsed.pipeline_stages =
Some(parse_pipeline_stages_value(&mut args, arg.as_str())?)
}
@ -2583,6 +2665,11 @@ mod tests {
assert!(help.help);
let short_help = ParsedArgs::parse(strings(&["-h"])).expect("short help parses");
assert!(short_help.help);
let alias = ParsedArgs::parse(strings(&["--vastai", "--pipeline-parallel", "4"]))
.expect("pipeline-parallel alias parses");
assert_eq!(alias.provider, Some(ProviderKind::VastAi));
assert_eq!(alias.pipeline_stages, Some(4));
}
#[test]

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

@ -74,14 +74,27 @@ fn node_event_payload(
status: &str,
detail: Value,
) -> Value {
let benchmark = benchmark_observability::stamp("mvp-worker-node");
json!({
"schema_version": benchmark["schema_version"].clone(),
"type":"NodeEvent",
"event_type":"NodeEvent",
"event_name":phase,
"phase":phase,
"status":status,
"run_id":config.run_id,
"node_id":config.logical_node_id,
"stage_index":config.stage_index,
"benchmark":benchmark_observability::stamp("mvp-worker-node"),
"producer_component":benchmark["producer_component"].clone(),
"producer_instance_id":benchmark["producer_instance_id"].clone(),
"producer_process_id":benchmark["producer_process_id"].clone(),
"producer_sequence":benchmark["producer_sequence"].clone(),
"wall_clock_unix_ms":benchmark["wall_clock_unix_ms"].clone(),
"monotonic_ms":benchmark["monotonic_ms"].clone(),
"clock_source":benchmark["clock_source"].clone(),
"span_id":format!("mvp-worker-node:{}:{}:{}", config.run_id, benchmark["producer_sequence"], phase),
"parent_span_id":Value::Null,
"benchmark":benchmark,
"detail":detail,
})
}
@ -444,18 +457,31 @@ fn sampler_health_payload(
status: &str,
detail: Value,
) -> Value {
let benchmark = benchmark_observability::stamp("mvp-worker-node");
json!({
"schema_version":benchmark["schema_version"].clone(),
"type":"SamplerHealth",
"event_type":"SamplerHealth",
"event_name":"host_sampler_health",
"schema":"mvp.node.sampler.health.v1",
"run_id":context.run_id,
"node_id":context.node_id,
"stage_index":context.stage_index,
"producer_component":benchmark["producer_component"].clone(),
"producer_instance_id":benchmark["producer_instance_id"].clone(),
"producer_process_id":benchmark["producer_process_id"].clone(),
"producer_sequence":benchmark["producer_sequence"].clone(),
"wall_clock_unix_ms":benchmark["wall_clock_unix_ms"].clone(),
"monotonic_ms":benchmark["monotonic_ms"].clone(),
"clock_source":benchmark["clock_source"].clone(),
"span_id":format!("mvp-worker-node:{}:{}:host_sampler_health", context.run_id, benchmark["producer_sequence"]),
"parent_span_id":Value::Null,
"phase":"host_sampler_health",
"status":status,
"sampler":sampler,
"sample_channel":sample_channel,
"detail":detail,
"benchmark":benchmark_observability::stamp("mvp-worker-node"),
"benchmark":benchmark,
})
}
@ -1827,6 +1853,35 @@ fn run() -> Result<(), String> {
"ready",
json!({"actor":datastream_publisher,"name":DATASTREAM_PUBLISHER_NAME,"subscription_transport":"iroh"}),
);
let worker_synthetic_id = format!(
"mvp-worker-node-{}-{}-datastream-preflight",
config.logical_node_id, config.stage_index
);
for (phase, status) in [
("DatastreamProducerConfigured", "configured"),
("DatastreamProducerConnected", "ready"),
("DatastreamSyntheticEventSent", "sent"),
("DatastreamSyntheticEventObserved", "observed"),
] {
emit_node_event(
&mut datastream,
&config,
NODE_BOOTSTRAP_CHANNEL,
phase,
status,
json!({
"producer":"mvp-worker-node",
"producer_class":"rust-worker-node",
"synthetic_id":worker_synthetic_id,
"datastream_endpoint":{
"role":"worker-node-iroh-publisher",
"transport":"iroh-datastream",
"endpoint_addr_mask":config.endpoint_addr_mask.as_str(),
"relay_mode":format!("{:?}", config.relay_mode),
},
}),
);
}
let mut debug_join_rx = match &config.debug_join_socket {
Some(path) => {
match spawn_debug_join_listener(tokio.handle().clone(), PathBuf::from(path)) {
@ -3890,6 +3945,25 @@ impl TinygradWorker {
.env("MVP_STAGE_INDEX", config.stage_index.to_string())
.env("MVP_ARENA_FD", arena_fd.to_string())
.env("MVP_ARENA_BYTES", config.arena_bytes.to_string())
.env(
"MVP_DATASTREAM_ENDPOINT_ID",
format!(
"worker-node-{}-stage-{}-stdio-bridge",
config.logical_node_id, config.stage_index
),
)
.env(
"MVP_BENCHMARK_PRODUCER_INSTANCE",
format!(
"tinygrad-worker:{}:{}",
config.logical_node_id, config.stage_index
),
)
.env(
"MVP_IROH_ENDPOINT_ADDR_MASK",
config.endpoint_addr_mask.as_str(),
)
.env("MVP_IROH_RELAY_MODE", format!("{:?}", config.relay_mode))
.stdin(Stdio::piped())
.stdout(Stdio::piped())
.stderr(Stdio::piped())

138
crates/mvp-system/src/node_image.rs
View file

@ -7,25 +7,8 @@ use std::sync::mpsc;
use std::thread;
use std::time::{Duration, Instant};
const IMAGE_SOURCE_INPUTS: &[&str] = &[
"Cargo.lock",
"Cargo.toml",
"src",
"crates/datastream/Cargo.toml",
"crates/datastream/src",
"crates/distribution/Cargo.toml",
"crates/distribution/src",
"crates/iroh-driver/Cargo.toml",
"crates/iroh-driver/src",
"crates/mvp-system/Cargo.toml",
"crates/mvp-system/src",
"crates/transport/Cargo.toml",
"crates/transport/src",
"tools/vastai/Cargo.toml",
"tools/vastai/src",
const NODE_IMAGE_CONTENT_INPUTS: &[&str] = &[
"apps/mvp-node/Dockerfile",
"apps/mvp-node/Dockerfile.base",
"apps/mvp-node/mvp_entrypoint.sh",
"apps/mvp-node/tinygrad_worker.py",
];
@ -192,13 +175,31 @@ fn prepare_node_image_inner(
));
}
let tag = image_version_tag(&root)?;
run_status(
runner,
progress,
&root,
"cargo",
&[
"build",
"--quiet",
"-p",
"mvp-system",
"--bin",
"mvp-worker-node",
],
"build mvp-worker-node",
None,
)?;
let base_hash = content_hash_for_inputs(&root, BASE_IMAGE_SOURCE_INPUTS)?;
let image_content_hash = node_image_content_hash(&root, &request.node_bin, &base_hash)?;
let tag = image_version_tag(&root, &image_content_hash)?;
let image_ref = image.ref_for_tag(&tag);
emit_image_reference(progress, "resolved", &image_ref);
let source_hash = source_content_hash(&root)?;
let worker_hash = file_content_hash(&root, Path::new("apps/mvp-node/tinygrad_worker.py"))?;
let base_hash = content_hash_for_inputs(&root, BASE_IMAGE_SOURCE_INPUTS)?;
let expected_node_labels = node_image_labels(&tag, &source_hash, &worker_hash, &base_hash);
let expected_node_labels =
node_image_labels(&tag, &image_content_hash, &worker_hash, &base_hash);
let expected_base_labels = base_image_labels(&base_hash);
let alias_tags = alias_tags(&image, request.extra_tag.as_deref(), &tag)?;
for alias in alias_refs(&image, &alias_tags) {
@ -247,23 +248,6 @@ fn prepare_node_image_inner(
pushed: false,
});
}
run_status(
runner,
progress,
&root,
"cargo",
&[
"build",
"--quiet",
"-p",
"mvp-system",
"--bin",
"mvp-worker-node",
],
"build mvp-worker-node",
None,
)?;
let base_image_matches =
docker_image_labels_match(runner, &root, &request.base_image, &expected_base_labels)?;
if !base_image_matches {
@ -350,12 +334,12 @@ fn workspace_root() -> Result<PathBuf, String> {
))
}
fn image_version_tag(root: &Path) -> Result<String, String> {
fn image_version_tag(root: &Path, image_content_hash: &str) -> Result<String, String> {
if git_worktree_clean(root)? {
let sha = git_capture(root, &["rev-parse", "--short=12", "HEAD"])?;
Ok(format!("git-{}", sha.trim()))
} else {
Ok(format!("dirty-{}", dirty_content_hash(root)?))
Ok(format!("dirty-{image_content_hash}"))
}
}
@ -384,12 +368,25 @@ fn git_capture(root: &Path, args: &[&str]) -> Result<String, String> {
}
}
fn dirty_content_hash(root: &Path) -> Result<String, String> {
source_content_hash(root)
}
fn source_content_hash(root: &Path) -> Result<String, String> {
content_hash_for_inputs(root, IMAGE_SOURCE_INPUTS)
fn node_image_content_hash(
root: &Path,
node_bin: &Path,
base_hash: &str,
) -> Result<String, String> {
let mut files = Vec::new();
for input in NODE_IMAGE_CONTENT_INPUTS {
let path = root.join(input);
collect_hash_inputs(root, &path, &mut files)?;
}
let node_bin = if node_bin.is_absolute() {
node_bin.to_path_buf()
} else {
root.join(node_bin)
};
files.push(relative_path(root, &node_bin)?);
files.sort();
files.dedup();
hash_relative_files_with_salts(root, files, &[("base", base_hash)])
}
fn content_hash_for_inputs(root: &Path, inputs: &[&str]) -> Result<String, String> {
@ -408,7 +405,21 @@ fn file_content_hash(root: &Path, path: &Path) -> Result<String, String> {
}
fn hash_relative_files(root: &Path, files: Vec<PathBuf>) -> Result<String, String> {
hash_relative_files_with_salts(root, files, &[])
}
fn hash_relative_files_with_salts(
root: &Path,
files: Vec<PathBuf>,
salts: &[(&str, &str)],
) -> Result<String, String> {
let mut hasher = blake3::Hasher::new();
for (key, value) in salts {
hasher.update(key.as_bytes());
hasher.update(b"\0");
hasher.update(value.as_bytes());
hasher.update(b"\0");
}
for relative in files {
let full = root.join(&relative);
hasher.update(relative.to_string_lossy().as_bytes());
@ -1348,7 +1359,7 @@ mod tests {
#[test]
fn prepare_node_image_with_dry_runner_returns_expected_image_and_progress() {
let root = workspace_root().expect("workspace root resolves");
let node_bin = std::env::current_exe().expect("test binary path resolves");
let mut runner = DryImageCommandRunner::default();
let mut progress = CollectProgress::default();
let mut sink: Option<&mut dyn NodeImageProgressSink> = Some(&mut progress);
@ -1357,7 +1368,7 @@ mod tests {
NodeImageRequest {
requested_image: "docker.io/acme/mvp-node:latest".to_owned(),
base_image: "swactor-mvp-node-base:cuda12.6".to_owned(),
node_bin: root.join("target/debug/mvp-worker-node"),
node_bin,
provider: NodeImageProvider::Docker,
extra_tag: Some("smoke".to_owned()),
push: false,
@ -1391,6 +1402,35 @@ mod tests {
)));
}
#[test]
fn node_image_content_hash_tracks_node_payload_not_unrelated_files() {
let root = workspace_root().expect("workspace root resolves");
let scratch = root
.join("target/node-image-hash-test")
.join(std::process::id().to_string());
fs::create_dir_all(&scratch).expect("scratch dir is writable");
let node_bin = scratch.join("mvp-worker-node");
fs::write(&node_bin, b"worker binary v1").expect("node bin fixture is writable");
let initial =
node_image_content_hash(&root, &node_bin, "base-v1").expect("initial hash succeeds");
fs::write(scratch.join("unrelated.txt"), b"not part of the image")
.expect("unrelated fixture is writable");
let after_unrelated =
node_image_content_hash(&root, &node_bin, "base-v1").expect("unrelated hash succeeds");
fs::write(&node_bin, b"worker binary v2").expect("node bin fixture update is writable");
let after_node_bin =
node_image_content_hash(&root, &node_bin, "base-v1").expect("node bin hash succeeds");
let after_base =
node_image_content_hash(&root, &node_bin, "base-v2").expect("base hash succeeds");
let _ = fs::remove_dir_all(&scratch);
assert_eq!(initial, after_unrelated);
assert_ne!(initial, after_node_bin);
assert_ne!(after_node_bin, after_base);
}
#[test]
fn image_name_splits_tag_after_last_slash() {
let image = ImageName::parse("localhost:5000/team/mvp-node:trial").unwrap();

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

@ -166,6 +166,53 @@ where
"provider_config":config.provider_datastream_detail(),
}),
);
orch_datastream.emit_bootstrap(
None,
config.run_id,
config.node_id,
"datastream_preflight",
"configured",
json!({
"producer":"mvp-orchestrator",
"datastream_endpoint":{
"role":"orchestrator-frame-archive",
"transport":"datastream-frame-log",
"configured":config.datastream_frame_log.is_some(),
"archive_path":config.datastream_frame_log.as_ref().map(|path| path.to_string_lossy().to_string()),
},
"expected_worker_producers":["mvp-worker-node","tinygrad-worker"],
"provider":config.provider.as_str(),
"pipeline_stages":config.pipeline_stages,
"endpoint_addr_mask":config.endpoint_addr_mask.as_str(),
"provider_config":config.provider_datastream_detail(),
}),
);
let orch_synthetic_id = format!("mvp-orchestrator-{}-datastream-preflight", config.run_id);
for (phase, status) in [
("DatastreamProducerConfigured", "configured"),
("DatastreamProducerConnected", "ready"),
("DatastreamSyntheticEventSent", "sent"),
("DatastreamSyntheticEventObserved", "observed"),
] {
orch_datastream.emit_bootstrap(
None,
config.run_id,
config.node_id,
phase,
status,
json!({
"producer":"mvp-orchestrator",
"producer_class":"rust-orchestrator",
"synthetic_id":orch_synthetic_id,
"datastream_endpoint":{
"role":"orchestrator-frame-archive",
"transport":"datastream-frame-log",
"configured":config.datastream_frame_log.is_some(),
"archive_path":config.datastream_frame_log.as_ref().map(|path| path.to_string_lossy().to_string()),
},
}),
);
}
drain_orch_stdio_capture(
orch_stdio_rx.as_ref(),
&mut orch_datastream,
@ -253,6 +300,22 @@ where
"ready",
json!({"endpoint":coordinator_endpoint.clone(),"has_relay":coordinator_endpoint.relay_urls().next().is_some(),"direct_addr_count":coordinator_endpoint.ip_addrs().count(),"relay_mode":format!("{:?}", config.relay.mode),"endpoint_addr_mask":config.endpoint_addr_mask.as_str()}),
);
orch_datastream.emit_bootstrap(
None,
config.run_id,
config.node_id,
"endpoint_config_snapshot",
"ready",
json!({
"producer":"mvp-orchestrator",
"coordinator_endpoint":coordinator_endpoint.clone(),
"has_relay":coordinator_endpoint.relay_urls().next().is_some(),
"direct_addr_count":coordinator_endpoint.ip_addrs().count(),
"relay_mode":format!("{:?}", config.relay.mode),
"endpoint_addr_mask":config.endpoint_addr_mask.as_str(),
"connectivity_preflight":"ready",
}),
);
let stack = DistributionRuntimeStack::new_with_codecs(
driver.node_id(),
DistributedNodeConfig::default(),
@ -4057,13 +4120,26 @@ impl OrchDatastream {
status: &str,
detail: Value,
) {
let benchmark = benchmark_observability::stamp("mvp-orchestrator");
let payload = serde_json::to_vec(&json!({
"schema_version": benchmark["schema_version"].clone(),
"type":"OrchBootstrap",
"event_type":"OrchBootstrap",
"event_name":phase,
"phase":phase,
"status":status,
"run_id":run_id,
"node_id":node_id,
"benchmark":benchmark_observability::stamp("mvp-orchestrator"),
"producer_component":benchmark["producer_component"].clone(),
"producer_instance_id":benchmark["producer_instance_id"].clone(),
"producer_process_id":benchmark["producer_process_id"].clone(),
"producer_sequence":benchmark["producer_sequence"].clone(),
"wall_clock_unix_ms":benchmark["wall_clock_unix_ms"].clone(),
"monotonic_ms":benchmark["monotonic_ms"].clone(),
"clock_source":benchmark["clock_source"].clone(),
"span_id":format!("mvp-orchestrator:{run_id}:{}:{phase}", benchmark["producer_sequence"]),
"parent_span_id":Value::Null,
"benchmark":benchmark,
"detail":detail,
}))
.expect("serialize orch bootstrap event");
@ -4080,14 +4156,27 @@ impl OrchDatastream {
status: &str,
detail: Value,
) {
let benchmark = benchmark_observability::stamp("mvp-orchestrator");
let payload = serde_json::to_vec(&json!({
"schema_version": benchmark["schema_version"].clone(),
"type":"OrchPromptEvent",
"event_type":"OrchPromptEvent",
"event_name":phase,
"phase":phase,
"status":status,
"run_id":run_id,
"node_id":node_id,
"request_id":request_id,
"benchmark":benchmark_observability::stamp("mvp-orchestrator"),
"producer_component":benchmark["producer_component"].clone(),
"producer_instance_id":benchmark["producer_instance_id"].clone(),
"producer_process_id":benchmark["producer_process_id"].clone(),
"producer_sequence":benchmark["producer_sequence"].clone(),
"wall_clock_unix_ms":benchmark["wall_clock_unix_ms"].clone(),
"monotonic_ms":benchmark["monotonic_ms"].clone(),
"clock_source":benchmark["clock_source"].clone(),
"span_id":format!("mvp-orchestrator:{run_id}:{request_id}:{}:{phase}", benchmark["producer_sequence"]),
"parent_span_id":format!("request:{request_id}"),
"benchmark":benchmark,
"detail":detail,
}))
.expect("serialize orch prompt event");

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

@ -115,6 +115,7 @@ impl VastAiBootstrapLauncher for FakeBootstrap {
endpoint: VastAiSshEndpoint,
_sink: PluginSink,
_producer: Option<datastream::DatastreamProducer>,
_lifecycle: LifecyclePolicy,
) -> Result<Self::Handle, String> {
if let Some(reason) = self.fail.clone() {
return Err(reason);

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

@ -1,5 +1,6 @@
use parking_lot::Mutex;
use std::collections::{BTreeMap, BTreeSet};
use std::collections::{BTreeMap, BTreeSet, HashSet};
use std::io::{BufRead, BufReader, Read};
use std::path::{Path, PathBuf};
use std::process::{Child, Command, Stdio};
use std::sync::{
@ -7,14 +8,16 @@ use std::sync::{
atomic::{AtomicBool, Ordering},
mpsc,
};
use std::time::Duration;
use std::thread::JoinHandle;
use std::time::{Duration, Instant};
use datastream::DatastreamProducer;
use serde::{Deserialize, Serialize};
use swactor::actor::{ActorAddress, ActorInterface};
use swactor::runtime::{Ctx, Runtime};
use swactor_vastai::{
LifecyclePolicy, ProvisionRequest, ProvisionedInstance, SelectionPolicy, classify_vastai_error,
CreateInstanceRequest, LifecyclePolicy, Offer, ProvisionRequest, ProvisionedInstance,
SelectionPolicy, classify_vastai_error, create_instance,
};
use crate::bootstrap_datastream::{BootstrapDatastreamBridge, node_stream_id};
@ -60,6 +63,33 @@ pub struct VastAiSshEndpoint {
pub user: String,
}
pub struct VastAiProviderMonitor {
stopping: Arc<AtomicBool>,
join: Option<JoinHandle<()>>,
}
impl VastAiProviderMonitor {
fn new(stopping: Arc<AtomicBool>, join: JoinHandle<()>) -> Self {
Self {
stopping,
join: Some(join),
}
}
fn stop(&mut self) {
self.stopping.store(true, Ordering::SeqCst);
if let Some(join) = self.join.take() {
let _ = join.join();
}
}
}
impl Drop for VastAiProviderMonitor {
fn drop(&mut self) {
self.stop();
}
}
pub trait VastAiLeaseClient: Send {
fn provision_one(&mut self, request: ProvisionRequest) -> Result<ProvisionedInstance, String>;
fn plan_first_wave_offers(
@ -76,6 +106,16 @@ pub trait VastAiLeaseClient: Send {
lifecycle: &LifecyclePolicy,
ssh_user: &str,
) -> Result<VastAiSshEndpoint, String>;
fn spawn_provider_monitor(
&mut self,
_contract_id: u64,
_label: String,
_lifecycle: LifecyclePolicy,
_spec: NodeProvisionSpec,
_sink: PluginSink,
) -> Option<VastAiProviderMonitor> {
None
}
fn destroy_contract(&mut self, contract_id: u64) -> Result<(), String>;
}
@ -83,6 +123,8 @@ pub trait VastAiLeaseClient: Send {
pub struct ToolsVastAiLeaseClient {
client: swactor_vastai::VastClient,
runtime: tokio::runtime::Runtime,
planned_offer_pool: Arc<Mutex<Vec<Offer>>>,
planned_offer_ids: Arc<Mutex<HashSet<u64>>>,
}
impl ToolsVastAiLeaseClient {
@ -91,7 +133,12 @@ impl ToolsVastAiLeaseClient {
.enable_all()
.build()
.map_err(|e| format!("vastai tokio runtime: {e}"))?;
Ok(Self { client, runtime })
Ok(Self {
client,
runtime,
planned_offer_pool: Arc::new(Mutex::new(Vec::new())),
planned_offer_ids: Arc::new(Mutex::new(HashSet::new())),
})
}
pub fn from_api_key(api_key: impl Into<String>) -> Result<Self, String> {
@ -101,25 +148,269 @@ impl ToolsVastAiLeaseClient {
pub fn client(&self) -> &swactor_vastai::VastClient {
&self.client
}
fn create_request_for_offer(
request: &ProvisionRequest,
offer_id: u64,
) -> CreateInstanceRequest {
let mut env = request.env.clone();
if let Some(overlay) = request.per_instance_env.first() {
env.extend(
overlay
.iter()
.map(|(key, value)| (key.clone(), value.clone())),
);
}
CreateInstanceRequest {
offer_id,
image: request.image.clone(),
disk_gb: request.disk_gb,
label: request.label.clone(),
env,
onstart: request.onstart.clone(),
}
}
fn candidate_pool(&mut self, request: &ProvisionRequest) -> Result<Vec<Offer>, String> {
let cached = self.planned_offer_pool.lock().clone();
if request
.preferred_offer_id
.is_some_and(|offer_id| cached.iter().any(|offer| offer.id == offer_id))
{
return Ok(cached);
}
self.runtime
.block_on(self.client.search_offers(&request.selection, 1))
}
fn create_from_offer(
&mut self,
request: &ProvisionRequest,
offer: &Offer,
) -> Result<ProvisionedInstance, String> {
let create = Self::create_request_for_offer(request, offer.id);
let info = self.runtime.block_on(create_instance(
self.client.http(),
self.client.base_url(),
self.client.api_key(),
&create,
))?;
Ok(ProvisionedInstance {
index: 0,
contract_id: info.contract_id,
offer_id: offer.id,
host_id: offer.host_id,
gpu_name: offer.gpu_name.clone(),
gpu_ram: offer.gpu_ram,
dph_total: offer.dph_total,
})
}
fn monitor_provider_status(
&mut self,
contract_id: u64,
label: String,
lifecycle: LifecyclePolicy,
spec: NodeProvisionSpec,
sink: PluginSink,
stopping: Arc<AtomicBool>,
) {
let mut last_state: Option<String> = None;
let mut state_since = Instant::now();
let mut poll = 0_u64;
while !stopping.load(Ordering::SeqCst) {
poll = poll.saturating_add(1);
let status = match self
.runtime
.block_on(self.client.instance_status(contract_id))
{
Ok(status) => status,
Err(error)
if error.contains("not found while fetching provider status")
|| error.contains("parse failed") =>
{
let reason = classified_start_error(format!(
"vastai provider monitor node {} contract {contract_id}: {error}",
spec.node_id
));
sink.observe(PluginObservation::ProviderLine {
run_id: spec.run_id,
node_id: spec.node_id,
line: serde_json::json!({
"type": "VastAiProviderStatusFailure",
"run_id": spec.run_id,
"node_id": spec.node_id,
"label": &label,
"contract_id": contract_id,
"poll": poll,
"reason": &reason,
})
.to_string(),
});
sink.observe(PluginObservation::Failed {
run_id: spec.run_id,
node_id: spec.node_id,
reason,
});
return;
}
Err(error) => {
sink.observe(PluginObservation::ProviderLine {
run_id: spec.run_id,
node_id: spec.node_id,
line: serde_json::json!({
"type": "VastAiProviderStatusPollRetry",
"run_id": spec.run_id,
"node_id": spec.node_id,
"label": &label,
"contract_id": contract_id,
"poll": poll,
"reason": error,
})
.to_string(),
});
if !sleep_provider_monitor(lifecycle.poll_interval, &stopping) {
return;
}
continue;
}
};
let actual = status.actual_status.as_str();
if last_state.as_deref() != Some(actual) {
state_since = Instant::now();
last_state = Some(actual.to_owned());
}
let in_state_ms = state_since.elapsed().as_millis();
sink.observe(PluginObservation::ProviderLine {
run_id: spec.run_id,
node_id: spec.node_id,
line: serde_json::json!({
"type": "VastAiProviderStatusObserved",
"run_id": spec.run_id,
"node_id": spec.node_id,
"label": &label,
"contract_id": contract_id,
"poll": poll,
"actual_status": &status.actual_status,
"intended_status": &status.intended_status,
"status_msg": &status.status_msg,
"disk_usage": status.disk_usage,
"in_state_ms": in_state_ms,
})
.to_string(),
});
if let Some(error) = provider_terminal_start_error(
contract_id,
&status.actual_status,
&status.intended_status,
status.status_msg.as_deref(),
) {
let reason = classified_start_error(format!(
"vastai provider monitor node {}: {error}",
spec.node_id
));
sink.observe(PluginObservation::ProviderLine {
run_id: spec.run_id,
node_id: spec.node_id,
line: serde_json::json!({
"type": "VastAiProviderTerminalBeforeRuntimeReady",
"run_id": spec.run_id,
"node_id": spec.node_id,
"label": &label,
"contract_id": contract_id,
"reason": &reason,
})
.to_string(),
});
sink.observe(PluginObservation::Failed {
run_id: spec.run_id,
node_id: spec.node_id,
reason,
});
return;
}
if !sleep_provider_monitor(lifecycle.poll_interval, &stopping) {
return;
}
}
}
}
impl Clone for ToolsVastAiLeaseClient {
fn clone(&self) -> Self {
Self::new(self.client.clone()).expect("clone VastAI lease client runtime")
Self {
client: self.client.clone(),
runtime: tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.expect("clone VastAI lease client runtime"),
planned_offer_pool: Arc::clone(&self.planned_offer_pool),
planned_offer_ids: Arc::clone(&self.planned_offer_ids),
}
}
}
impl VastAiLeaseClient for ToolsVastAiLeaseClient {
fn provision_one(&mut self, request: ProvisionRequest) -> Result<ProvisionedInstance, String> {
let fleet = self.runtime.block_on(self.client.provision(request))?;
let mut instances = fleet.instances;
if instances.len() != 1 {
if request.count != 1 {
return Err(format!(
"vastai provision expected one instance, got {}",
instances.len()
"vastai provision_one expected count=1, got {}",
request.count
));
}
Ok(instances.remove(0))
let pool = self.candidate_pool(&request)?;
let planned_offer_ids = self.planned_offer_ids.lock().clone();
let blocked_hosts = request
.selection
.blacklist_hosts
.iter()
.copied()
.collect::<HashSet<_>>();
let mut failed_hosts = HashSet::new();
let mut tried_offer_ids = HashSet::new();
let mut ordered = Vec::with_capacity(pool.len());
if let Some(preferred_offer_id) = request.preferred_offer_id
&& let Some(offer) = pool.iter().find(|offer| offer.id == preferred_offer_id)
{
ordered.push(offer.clone());
}
ordered.extend(pool.into_iter());
let mut last_error = None;
for offer in ordered {
if !tried_offer_ids.insert(offer.id) {
continue;
}
if request.preferred_offer_id != Some(offer.id) && planned_offer_ids.contains(&offer.id)
{
continue;
}
if offer.host_id.is_some_and(|host_id| {
blocked_hosts.contains(&host_id) || failed_hosts.contains(&host_id)
}) {
continue;
}
match self.create_from_offer(&request, &offer) {
Ok(instance) => return Ok(instance),
Err(error) => {
if let Some(host_id) = offer.host_id {
failed_hosts.insert(host_id);
}
last_error = Some(format!("offer {}: {error}", offer.id));
}
}
}
Err(format!(
"vastai provision node exhausted eligible offers{}",
last_error
.map(|error| format!(" after create failure ({error})"))
.unwrap_or_default()
))
}
fn plan_first_wave_offers(
@ -127,6 +418,8 @@ impl VastAiLeaseClient for ToolsVastAiLeaseClient {
requests: &[ProvisionRequest],
) -> Result<Vec<Option<u64>>, String> {
let Some(first) = requests.first() else {
self.planned_offer_pool.lock().clear();
self.planned_offer_ids.lock().clear();
return Ok(Vec::new());
};
let pool = self.runtime.block_on(
@ -139,6 +432,9 @@ impl VastAiLeaseClient for ToolsVastAiLeaseClient {
&first.selection.blacklist_hosts,
&[],
);
let planned_ids = planned.iter().map(|offer| offer.id).collect::<HashSet<_>>();
*self.planned_offer_pool.lock() = pool;
*self.planned_offer_ids.lock() = planned_ids;
let mut out = planned
.into_iter()
.map(|offer| Some(offer.id))
@ -154,23 +450,36 @@ impl VastAiLeaseClient for ToolsVastAiLeaseClient {
lifecycle: &LifecyclePolicy,
ssh_user: &str,
) -> Result<VastAiSshEndpoint, String> {
self.runtime.block_on(async {
let instances = self.client.list_by_label(label).await?;
if let Some(instance) = instances
.into_iter()
.find(|instance| instance.contract_id == contract_id)
{
let host = if instance.ssh_host.is_empty() {
instance.public_ipaddr
} else {
instance.ssh_host
};
return endpoint_from_parts(contract_id, host, instance.ssh_port, ssh_user);
let endpoint = self.runtime.block_on(self.client.wait_for_ssh_endpoint(
contract_id,
label,
lifecycle,
))?;
endpoint_from_parts(contract_id, endpoint.ip, endpoint.port, ssh_user)
}
let running = self.client.wait_for_running(contract_id, lifecycle).await?;
endpoint_from_parts(contract_id, running.ip, running.port, ssh_user)
})
fn spawn_provider_monitor(
&mut self,
contract_id: u64,
label: String,
lifecycle: LifecyclePolicy,
spec: NodeProvisionSpec,
sink: PluginSink,
) -> Option<VastAiProviderMonitor> {
let stopping = Arc::new(AtomicBool::new(false));
let thread_stopping = Arc::clone(&stopping);
let mut client = self.clone();
let join = std::thread::spawn(move || {
client.monitor_provider_status(
contract_id,
label,
lifecycle,
spec,
sink,
thread_stopping,
);
});
Some(VastAiProviderMonitor::new(stopping, join))
}
fn destroy_contract(&mut self, contract_id: u64) -> Result<(), String> {
@ -203,6 +512,44 @@ fn endpoint_from_parts(
})
}
fn provider_terminal_start_error(
contract_id: u64,
actual: &str,
intended: &str,
msg: Option<&str>,
) -> Option<String> {
if let Some(message) = msg {
let lower = message.to_ascii_lowercase();
if lower.contains("error") || lower.contains("failed") {
return Some(format!("instance {contract_id} error: {message}"));
}
}
if intended == "stopped" && actual != "running" {
return Some(format!(
"instance {contract_id} stopped: {}",
msg.unwrap_or_default()
));
}
match actual {
"exited" | "error" | "stopped" => Some(format!(
"instance {contract_id} reached terminal status: {actual}"
)),
_ => None,
}
}
fn sleep_provider_monitor(duration: Duration, stopping: &AtomicBool) -> bool {
let deadline = Instant::now() + duration;
while Instant::now() < deadline {
if stopping.load(Ordering::SeqCst) {
return false;
}
let remaining = deadline.saturating_duration_since(Instant::now());
std::thread::sleep(std::cmp::min(remaining, Duration::from_millis(100)));
}
!stopping.load(Ordering::SeqCst)
}
#[derive(Clone, Debug)]
pub struct VastAiProviderPlugin<C>
where
@ -421,6 +768,7 @@ pub trait VastAiBootstrapLauncher: Send {
endpoint: VastAiSshEndpoint,
sink: PluginSink,
producer: Option<DatastreamProducer>,
lifecycle: LifecyclePolicy,
) -> Result<Self::Handle, String>;
fn stop_bootstrap(&mut self, handle: &mut Self::Handle, reason: BootstrapStopReason);
@ -492,6 +840,7 @@ impl VastAiBootstrapLauncher for SshCommandBootstrapLauncher {
endpoint: VastAiSshEndpoint,
sink: PluginSink,
producer: Option<DatastreamProducer>,
lifecycle: LifecyclePolicy,
) -> Result<Self::Handle, String> {
if spec.args.is_empty() {
return Err(format!(
@ -514,6 +863,7 @@ impl VastAiBootstrapLauncher for SshCommandBootstrapLauncher {
self.ssh_identity.clone(),
child,
stopping,
lifecycle.state_timeout,
);
Ok(SshCommandBootstrapHandle {
@ -528,6 +878,70 @@ impl VastAiBootstrapLauncher for SshCommandBootstrapLauncher {
}
}
const POST_GRACE_BOOTSTRAP_FAILURE_LIMIT: u32 = 2;
fn classify_ssh_observation(line: &str) -> Option<&'static str> {
let lower = line.to_ascii_lowercase();
if lower.contains("permission denied (publickey")
|| lower.contains("publickey denied")
|| lower.contains("public key denied")
|| lower.contains("no supported authentication methods")
{
return Some("auth_denied");
}
if lower.contains("connection refused")
|| lower.contains("connect to host") && lower.contains("refused")
{
return Some("refused");
}
if lower.contains("operation timed out")
|| lower.contains("connection timed out")
|| lower.contains("connect timed out")
{
return Some("timeout");
}
None
}
fn post_grace_terminal_bootstrap_class(class: &str) -> bool {
matches!(class, "auth_denied" | "refused" | "timeout")
}
fn spawn_classifying_stderr_reader<R>(
stderr: R,
bridge: BootstrapDatastreamBridge,
observed_class: Arc<Mutex<Option<&'static str>>>,
) -> JoinHandle<()>
where
R: Read + Send + 'static,
{
std::thread::spawn(move || {
let reader = BufReader::new(stderr);
for next in reader.lines() {
match next {
Ok(line) => {
if let Some(class) = classify_ssh_observation(&line) {
*observed_class.lock() = Some(class);
bridge.observe_provider_line(
serde_json::json!({
"type": "VastAiBootstrapObservationClass",
"run_id": bridge.spec().run_id,
"node_id": bridge.spec().node_id,
"class": class,
})
.to_string(),
);
}
bridge.observe_stderr_line(line);
}
Err(error) => {
bridge.observe_provider_line(format!("read VastAI SSH stderr: {error}"));
break;
}
}
}
})
}
fn spawn_retrying_ssh_bootstrap(
spec: NodeProvisionSpec,
endpoint: VastAiSshEndpoint,
@ -536,12 +950,15 @@ fn spawn_retrying_ssh_bootstrap(
ssh_identity: Option<PathBuf>,
child_slot: Arc<Mutex<Option<Child>>>,
stopping: Arc<AtomicBool>,
post_grace_failure_after: Duration,
) {
std::thread::spawn(move || {
let run_id = spec.run_id;
let node_id = spec.node_id;
let mut attempt = 1u64;
let mut backoff = Duration::from_secs(1);
let bootstrap_started = Instant::now();
let mut post_grace_failures = 0_u32;
while !stopping.load(Ordering::SeqCst) {
sink.observe(PluginObservation::ProviderLine {
@ -562,7 +979,12 @@ fn spawn_retrying_ssh_bootstrap(
producer.clone(),
);
bridge.spawn_stdout_reader(stdout);
bridge.spawn_stderr_reader(stderr);
let observed_class = Arc::new(Mutex::new(None));
let mut stderr_reader = Some(spawn_classifying_stderr_reader(
stderr,
bridge.clone(),
Arc::clone(&observed_class),
));
loop {
if stopping.load(Ordering::SeqCst) {
@ -597,17 +1019,50 @@ fn spawn_retrying_ssh_bootstrap(
} else {
"not ready before runtime ready"
};
let line = format!(
"VastAI SSH bootstrap {readiness} (attempt {attempt}, {status}); retrying"
);
if let Some(reader) = stderr_reader.take() {
let _ = reader.join();
}
let observation_class =
(*observed_class.lock()).unwrap_or("process_exit");
sink.observe(PluginObservation::ProviderLine {
run_id,
node_id,
line,
line: serde_json::json!({
"type": "VastAiBootstrapAttemptCompleted",
"run_id": run_id,
"node_id": node_id,
"attempt": attempt,
"status": status.to_string(),
"class": observation_class,
"classification": readiness,
})
.to_string(),
});
if !status.success()
&& !post_grace_failure_after.is_zero()
&& bootstrap_started.elapsed() >= post_grace_failure_after
&& post_grace_terminal_bootstrap_class(observation_class)
{
post_grace_failures = post_grace_failures.saturating_add(1);
if post_grace_failures >= POST_GRACE_BOOTSTRAP_FAILURE_LIMIT {
sink.observe(PluginObservation::Failed {
run_id,
node_id,
reason: format!(
"VastAI SSH bootstrap repeated post-grace {observation_class} failure before runtime ready"
),
});
return;
}
} else {
post_grace_failures = 0;
}
break;
}
Some(Err(error)) => {
if let Some(reader) = stderr_reader.take() {
let _ = reader.join();
}
sink.observe(PluginObservation::ProviderLine {
run_id,
node_id,
@ -740,7 +1195,12 @@ where
struct VastAiNode<H> {
contract_id: u64,
bootstrap: Option<H>,
provider_monitor: Option<VastAiProviderMonitor>,
host_id: Option<u64>,
run_id: u64,
node_id: u64,
label: String,
sink: PluginSink,
}
impl<C, B> VastAiProvisioningPlugin<C, B>
@ -908,6 +1368,18 @@ where
})
.to_string(),
});
sink.observe(PluginObservation::ProviderLine {
run_id: spec.run_id,
node_id: spec.node_id,
line: serde_json::json!({
"type": "VastAiSshEndpointDiscoveryStarted",
"run_id": spec.run_id,
"node_id": spec.node_id,
"contract_id": instance.contract_id,
"label": &label,
})
.to_string(),
});
let endpoint = match self.client.ssh_endpoint(
instance.contract_id,
@ -944,11 +1416,27 @@ where
.to_string(),
});
sink.observe(PluginObservation::ProviderLine {
run_id: spec.run_id,
node_id: spec.node_id,
line: serde_json::json!({
"type": "VastAiBootstrapObservationStarted",
"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(),
endpoint,
sink,
sink.clone(),
self.bootstrap_producer.clone(),
self.config.lifecycle.clone(),
) {
Ok(handle) => handle,
Err(error) => {
@ -980,7 +1468,18 @@ where
VastAiNode {
contract_id: instance.contract_id,
bootstrap: Some(bootstrap),
provider_monitor: self.client.spawn_provider_monitor(
instance.contract_id,
label.clone(),
self.config.lifecycle.clone(),
spec.clone(),
sink.clone(),
),
host_id,
run_id: spec.run_id,
node_id: spec.node_id,
label,
sink,
},
);
Ok(handle)
@ -1059,13 +1558,56 @@ where
})
.to_string(),
});
} else {
sink.observe(PluginObservation::ProviderLine {
run_id: spec.run_id,
node_id: spec.node_id,
line: serde_json::json!({
"type": "VastAiFirstWaveOfferPlanUnavailable",
"run_id": spec.run_id,
"node_id": spec.node_id,
"label": &label,
})
.to_string(),
});
}
let mut client = self.client.clone();
let config = self.config.clone();
let worker_tx = completion_tx.clone();
let worker_sink = sink.clone();
std::thread::spawn(move || {
let started = match client.provision_one(request) {
Ok(instance) => {
worker_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(),
});
worker_sink.observe(PluginObservation::ProviderLine {
run_id: spec.run_id,
node_id: spec.node_id,
line: serde_json::json!({
"type": "VastAiSshEndpointDiscoveryStarted",
"run_id": spec.run_id,
"node_id": spec.node_id,
"contract_id": instance.contract_id,
"label": &label,
})
.to_string(),
});
match client.ssh_endpoint(
instance.contract_id,
&label,
@ -1112,29 +1654,17 @@ where
for (index, spec, started) in completion_rx {
match started {
Ok(started) => {
sink.observe(PluginObservation::ProviderLine {
run_id: spec.run_id,
node_id: spec.node_id,
line: format!(
"vastai contract {} ready for SSH lookup",
started.instance.contract_id
),
});
sink.observe(PluginObservation::ProviderLine {
run_id: spec.run_id,
node_id: spec.node_id,
line: serde_json::json!({
"type": "VastAiLeaseReady",
"type": "VastAiSshEndpointReady",
"run_id": spec.run_id,
"node_id": spec.node_id,
"label": &started.label,
"image": &spec.image,
"contract_id": started.instance.contract_id,
"offer_id": started.instance.offer_id,
"host_id": started.instance.host_id,
"gpu_name": &started.instance.gpu_name,
"gpu_ram": started.instance.gpu_ram,
"dph_total": started.instance.dph_total,
"host": &started.endpoint.host,
"port": started.endpoint.port,
"user": &started.endpoint.user,
})
.to_string(),
});
@ -1142,7 +1672,7 @@ where
run_id: spec.run_id,
node_id: spec.node_id,
line: serde_json::json!({
"type": "VastAiSshEndpointReady",
"type": "VastAiBootstrapObservationStarted",
"run_id": spec.run_id,
"node_id": spec.node_id,
"contract_id": started.instance.contract_id,
@ -1158,6 +1688,7 @@ where
started.endpoint,
sink.clone(),
self.bootstrap_producer.clone(),
self.config.lifecycle.clone(),
) {
Ok(handle) => handle,
Err(error) => {
@ -1192,7 +1723,18 @@ where
VastAiNode {
contract_id: started.instance.contract_id,
bootstrap: Some(bootstrap),
provider_monitor: self.client.spawn_provider_monitor(
started.instance.contract_id,
started.label.clone(),
self.config.lifecycle.clone(),
spec.clone(),
sink.clone(),
),
host_id,
run_id: spec.run_id,
node_id: spec.node_id,
label: started.label,
sink: sink.clone(),
},
);
results[index] = Some((spec, Ok(handle)));
@ -1232,6 +1774,23 @@ where
let Some(node) = self.nodes.get_mut(&handle.id) else {
return Ok(());
};
if let Some(monitor) = node.provider_monitor.as_mut() {
monitor.stop();
}
node.provider_monitor = None;
node.sink.observe(PluginObservation::ProviderLine {
run_id: node.run_id,
node_id: node.node_id,
line: serde_json::json!({
"type": "VastAiRuntimeReadyAccepted",
"run_id": node.run_id,
"node_id": node.node_id,
"label": &node.label,
"contract_id": node.contract_id,
"classification": "runtime_ready_over_provider_staleness",
})
.to_string(),
});
if let Some(mut bootstrap) = node.bootstrap.take() {
self.bootstrap
.stop_bootstrap(&mut bootstrap, BootstrapStopReason::RuntimeReady);
@ -1243,6 +1802,9 @@ where
let Some(mut node) = self.nodes.remove(&handle.id) else {
return Ok(());
};
if let Some(mut monitor) = node.provider_monitor.take() {
monitor.stop();
}
if let Some(host_id) = node.host_id {
self.leased_host_ids.remove(&host_id);
}
@ -1250,7 +1812,22 @@ where
self.bootstrap
.stop_bootstrap(&mut bootstrap, BootstrapStopReason::NodeStop);
}
self.client.destroy_contract(node.contract_id)
let result = self.client.destroy_contract(node.contract_id);
node.sink.observe(PluginObservation::ProviderLine {
run_id: node.run_id,
node_id: node.node_id,
line: serde_json::json!({
"type": "VastAiContractCleanup",
"run_id": node.run_id,
"node_id": node.node_id,
"label": &node.label,
"contract_id": node.contract_id,
"result": if result.is_ok() { "ok" } else { "failed" },
"error": result.as_ref().err(),
})
.to_string(),
});
result
}
}
@ -1294,6 +1871,7 @@ mod tests {
_endpoint: VastAiSshEndpoint,
_sink: PluginSink,
_producer: Option<DatastreamProducer>,
_lifecycle: LifecyclePolicy,
) -> Result<Self::Handle, String> {
panic!("build_request tests must not start SSH bootstrap")
}
@ -1391,6 +1969,7 @@ mod tests {
_endpoint: VastAiSshEndpoint,
_sink: PluginSink,
_producer: Option<DatastreamProducer>,
_lifecycle: LifecyclePolicy,
) -> Result<Self::Handle, String> {
Ok(spec.node_id)
}
@ -1471,6 +2050,21 @@ mod tests {
}
}
#[test]
fn ssh_bootstrap_observation_classifies_auth_and_transport_failures() {
for (line, expected) in [
("Permission denied (publickey).", Some("auth_denied")),
(
"ssh: connect to host ssh5.vast.ai port 22017: Connection refused",
Some("refused"),
),
("ssh: connect timed out", Some("timeout")),
("debug1: permanently_set_uid", None),
] {
assert_eq!(classify_ssh_observation(line), expected, "{line}");
}
}
#[test]
fn ssh_bootstrap_args_include_verbose_flag_and_identity_when_configured() {
let endpoint = VastAiSshEndpoint {

32
tools/vastai/src/client.rs
View file

@ -1,7 +1,8 @@
use std::time::Duration;
use crate::types::{
LabeledInstance, LifecyclePolicy, Offer, ProvisionRequest, ProvisionedFleet, RunningInstance,
LabeledInstance, LifecyclePolicy, Offer, ProviderInstanceStatus, ProvisionRequest,
ProvisionedFleet, RunningInstance,
};
/// Small convenience wrapper around a reqwest client + vast.ai endpoint.
@ -60,6 +61,18 @@ impl VastClient {
crate::lease::provision_fleet(&self.http, &self.base_url, &self.api_key, req).await
}
pub async fn instance_status(
&self,
contract_id: u64,
) -> Result<ProviderInstanceStatus, String> {
crate::monitor::fetch_instance_status(
&self.http,
&self.base_url,
&self.api_key,
contract_id,
)
.await
}
pub async fn wait_for_running(
&self,
contract_id: u64,
@ -75,6 +88,23 @@ impl VastClient {
.await
}
pub async fn wait_for_ssh_endpoint(
&self,
contract_id: u64,
label: &str,
policy: &LifecyclePolicy,
) -> Result<RunningInstance, String> {
crate::monitor::wait_for_ssh_endpoint_with_policy(
&self.http,
&self.base_url,
&self.api_key,
contract_id,
label,
policy,
)
.await
}
pub async fn list_by_label(&self, label: &str) -> Result<Vec<LabeledInstance>, String> {
crate::teardown::list_instances_by_label(&self.http, &self.base_url, &self.api_key, label)
.await

4
tools/vastai/src/lease.rs
View file

@ -398,9 +398,9 @@ mod tests {
.and(path("/api/v0/instances/101/"))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
"instances": {
"actual_status": "loading",
"actual_status": "error",
"intended_status": "running",
"status_msg": "still pulling"
"status_msg": "container failed before runtime readiness"
}
})))
.mount(&server)

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

@ -20,7 +20,10 @@ pub mod types;
pub use client::VastClient;
pub use lease::{confirm_lease, provision_fleet};
pub use logs::{fetch_logs, request_logs};
pub use monitor::{wait_for_running, wait_for_running_with_policy};
pub use monitor::{
fetch_instance_status, wait_for_running, wait_for_running_with_policy,
wait_for_ssh_endpoint_with_policy,
};
pub use pricing::CostModel;
pub use provision::create_instance;
pub use search::{plan_distinct_host_first_wave, select_offer_pool, select_offer_pool_with_policy};
@ -29,6 +32,6 @@ pub use teardown::{
};
pub use types::{
ContractRef, CreateInstanceRequest, FleetState, InstanceInfo, LabeledInstance, LifecyclePolicy,
Offer, ProvisionRequest, ProvisionedFleet, ProvisionedInstance, RunningInstance,
SelectionPolicy, VastAiFailureClass, classify_vastai_error,
Offer, ProviderInstanceStatus, ProvisionRequest, ProvisionedFleet, ProvisionedInstance,
RunningInstance, SelectionPolicy, VastAiFailureClass, classify_vastai_error,
};

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

@ -1,4 +1,92 @@
use crate::types::{InstanceResponse, LifecyclePolicy, RunningInstance};
use crate::types::{
InstanceResponse, LabeledInstance, LifecyclePolicy, ProviderInstanceStatus, RunningInstance,
};
pub async fn fetch_instance_status(
client: &reqwest::Client,
base_url: &str,
api_key: &str,
contract_id: u64,
) -> Result<ProviderInstanceStatus, String> {
let url = format!("{base_url}/api/v0/instances/{contract_id}/");
let resp = client
.get(&url)
.header("Authorization", format!("Bearer {api_key}"))
.send()
.await
.map_err(|e| format!("fetch_instance_status request failed: {e}"))?;
if resp.status() == reqwest::StatusCode::NOT_FOUND {
let body = resp.text().await.unwrap_or_default();
return Err(format!(
"instance {contract_id} not found while fetching provider status: {}",
body.chars().take(80).collect::<String>(),
));
}
if !resp.status().is_success() {
let status = resp.status();
let body = resp.text().await.unwrap_or_default();
return Err(format!(
"fetch_instance_status HTTP {status}: {}",
body.chars().take(80).collect::<String>(),
));
}
let wrapper: InstanceResponse = resp
.json()
.await
.map_err(|e| format!("fetch_instance_status parse failed: {e}"))?;
Ok(wrapper.instances.into())
}
fn provider_terminal_error(
contract_id: u64,
actual: &str,
intended: &str,
msg: Option<&str>,
) -> Option<String> {
if let Some(m) = msg {
let lower = m.to_ascii_lowercase();
if lower.contains("error") || lower.contains("failed") {
return Some(format!("instance {contract_id} error: {m}"));
}
}
if intended == "stopped" && actual != "running" {
return Some(format!(
"instance {contract_id} stopped: {}",
msg.unwrap_or_default()
));
}
match actual {
"exited" | "error" | "stopped" => Some(format!(
"instance {contract_id} reached terminal status: {actual}"
)),
_ => None,
}
}
fn labeled_endpoint(instance: &LabeledInstance) -> Option<RunningInstance> {
let host = if instance.ssh_host.trim().is_empty() {
instance.public_ipaddr.trim()
} else {
instance.ssh_host.trim()
};
if host.is_empty() || instance.ssh_port == 0 {
return None;
}
Some(RunningInstance {
ip: host.to_owned(),
port: instance.ssh_port,
})
}
fn missing_instance(error: &str) -> bool {
error.contains("not found while fetching provider status")
}
fn parse_failed(error: &str) -> bool {
error.contains("parse failed")
}
/// Historical env-backed polling wrapper.
pub async fn wait_for_running(
@ -20,67 +108,41 @@ pub async fn wait_for_running_with_policy(
contract_id: u64,
policy: &LifecyclePolicy,
) -> Result<RunningInstance, String> {
let url = format!("{base_url}/api/v0/instances/{contract_id}/");
let mut state_since = std::time::Instant::now();
let mut last_state: Option<String> = None;
let mut running_without_endpoint_since = None;
let mut poll = 0_u64;
loop {
poll += 1;
let resp = match client
.get(&url)
.header("Authorization", format!("Bearer {api_key}"))
.send()
.await
{
Ok(r) => r,
Err(e) => {
eprintln!(" contract {contract_id} poll {poll}: request error: {e} (retrying)");
let status = match fetch_instance_status(client, base_url, api_key, contract_id).await {
Ok(status) => status,
Err(error) if missing_instance(&error) => {
return Err(error.replace(
"while fetching provider status",
"while waiting for running",
));
}
Err(error) if parse_failed(&error) => return Err(error),
Err(error) => {
eprintln!(" contract {contract_id} poll {poll}: {error} (retrying)");
tokio::time::sleep(policy.poll_interval).await;
continue;
}
};
if resp.status() == reqwest::StatusCode::NOT_FOUND {
let body = resp.text().await.unwrap_or_default();
return Err(format!(
"instance {contract_id} not found while waiting for running: {}",
body.chars().take(80).collect::<String>(),
));
}
if !resp.status().is_success() {
let status = resp.status();
let body = resp.text().await.unwrap_or_default();
eprintln!(
" contract {contract_id} poll {poll}: HTTP {status} (retrying): {}",
body.chars().take(80).collect::<String>(),
);
tokio::time::sleep(policy.poll_interval).await;
continue;
}
let wrapper: InstanceResponse = resp
.json()
.await
.map_err(|e| format!("wait_for_running parse failed: {e}"))?;
let status = wrapper.instances;
let actual = status.actual_status.as_deref().unwrap_or("unknown");
let intended = status.intended_status.as_deref().unwrap_or("unknown");
let msg = status.status_msg.clone();
let disk = status.disk_usage;
let actual = status.actual_status.as_str();
let intended = status.intended_status.as_str();
if last_state.as_deref() != Some(actual) {
state_since = std::time::Instant::now();
}
last_state = Some(actual.to_string());
last_state = Some(actual.to_owned());
let in_state = state_since.elapsed().as_secs();
let msg_disp = match msg.as_deref() {
let msg_disp = match status.status_msg.as_deref() {
Some(m) if !m.is_empty() => format!(" msg=\"{m}\""),
_ => String::new(),
};
let disk_disp = match disk {
let disk_disp = match status.disk_usage {
Some(d) if d >= 0.0 => format!(" disk={d:.2}GB"),
_ => String::new(),
};
@ -88,41 +150,118 @@ pub async fn wait_for_running_with_policy(
" contract {contract_id} poll {poll}: status={actual} in-state={in_state}s{msg_disp}{disk_disp}",
);
if let Some(m) = &msg {
if m.contains("Error") || m.contains("failed") {
return Err(format!("instance {contract_id} error: {m}"));
if let Some(error) =
provider_terminal_error(contract_id, actual, intended, status.status_msg.as_deref())
{
return Err(error);
}
if actual == "running" {
if let Some(endpoint) = status.ssh_endpoint() {
return Ok(endpoint);
}
if intended == "stopped" && actual != "running" {
let since = running_without_endpoint_since
.get_or_insert_with(std::time::Instant::now)
.elapsed();
if !policy.state_timeout.is_zero() && since >= policy.state_timeout {
return Err(format!(
"instance {contract_id} stopped: {}",
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",
"instance {contract_id} running without usable SSH endpoint for {}s",
policy.state_timeout.as_secs()
));
}
} else {
running_without_endpoint_since = None;
}
match actual {
"running" => {
let ip = status
.public_ipaddr
.unwrap_or_else(|| "unknown".to_string());
let port = status.ssh_port.unwrap_or(0);
return Ok(RunningInstance { ip, port });
}
"exited" | "error" => {
return Err(format!(
"instance {contract_id} reached terminal status: {actual}"
));
}
_ => {
tokio::time::sleep(policy.poll_interval).await;
}
}
/// Poll provider data until a usable SSH endpoint exists, without requiring
/// provider `running` status first.
pub async fn wait_for_ssh_endpoint_with_policy(
client: &reqwest::Client,
base_url: &str,
api_key: &str,
contract_id: u64,
label: &str,
policy: &LifecyclePolicy,
) -> Result<RunningInstance, String> {
let mut running_without_endpoint_since = None;
let mut poll = 0_u64;
loop {
poll += 1;
match crate::teardown::list_instances_by_label(client, base_url, api_key, label).await {
Ok(instances) => {
if let Some(instance) = instances
.iter()
.find(|instance| instance.contract_id == contract_id)
{
if let Some(endpoint) = labeled_endpoint(instance) {
eprintln!(
" contract {contract_id} endpoint poll {poll}: endpoint discovered from label status={}",
instance.actual_status
);
return Ok(endpoint);
}
eprintln!(
" contract {contract_id} endpoint poll {poll}: label status={} endpoint missing",
instance.actual_status
);
}
}
Err(error) => {
eprintln!(
" contract {contract_id} endpoint poll {poll}: list-by-label error: {error} (retrying)"
);
}
}
let status = match fetch_instance_status(client, base_url, api_key, contract_id).await {
Ok(status) => status,
Err(error) if missing_instance(&error) => return Err(error),
Err(error) if parse_failed(&error) => return Err(error),
Err(error) => {
eprintln!(" contract {contract_id} endpoint poll {poll}: {error} (retrying)");
tokio::time::sleep(policy.poll_interval).await;
continue;
}
};
let actual = status.actual_status.as_str();
let intended = status.intended_status.as_str();
let msg = status.status_msg.as_deref();
if let Some(endpoint) = status.ssh_endpoint() {
eprintln!(
" contract {contract_id} endpoint poll {poll}: endpoint discovered from provider status={actual}",
);
return Ok(endpoint);
}
if let Some(error) = provider_terminal_error(contract_id, actual, intended, msg) {
return Err(error);
}
if actual == "running" {
let since = running_without_endpoint_since
.get_or_insert_with(std::time::Instant::now)
.elapsed();
if !policy.state_timeout.is_zero() && since >= policy.state_timeout {
return Err(format!(
"instance {contract_id} running without usable SSH endpoint for {}s",
policy.state_timeout.as_secs()
));
}
} else {
running_without_endpoint_since = None;
}
let msg_disp = match msg {
Some(m) if !m.is_empty() => format!(" msg=\"{m}\""),
_ => String::new(),
};
eprintln!(
" contract {contract_id} endpoint poll {poll}: status={actual} endpoint missing{msg_disp}; waiting"
);
tokio::time::sleep(policy.poll_interval).await;
}
}
@ -137,7 +276,7 @@ mod tests {
use super::*;
#[tokio::test]
async fn stuck_loading_state_returns_error_instead_of_polling_forever() {
async fn loading_state_remains_slow_progress_before_terminal_evidence() {
let server = MockServer::start().await;
Mock::given(method("GET"))
.and(path("/api/v0/instances/123/"))
@ -145,7 +284,7 @@ mod tests {
"instances": {
"actual_status": "loading",
"intended_status": "running",
"status_msg": "afad30e59d72: Already exists"
"status_msg": "pulling image layers"
}
})))
.mount(&server)
@ -153,26 +292,129 @@ mod tests {
let policy = LifecyclePolicy {
poll_interval: Duration::from_millis(1),
state_timeout: Duration::from_millis(5),
state_timeout: Duration::from_millis(1),
..LifecyclePolicy::default()
};
let error = wait_for_running_with_policy(
let still_pending = tokio::time::timeout(
Duration::from_millis(10),
wait_for_running_with_policy(
&reqwest::Client::new(),
&server.uri(),
"secret",
123,
&policy,
),
)
.await
.expect_err("stuck loading should be replaceable");
.await;
assert!(
error.contains("stuck in status loading"),
"error should name stuck provider state: {error}"
still_pending.is_err(),
"loading by itself should remain slow progress instead of replaceable failure"
);
}
#[tokio::test]
async fn endpoint_discovery_uses_label_endpoint_before_running_status() {
let server = MockServer::start().await;
Mock::given(method("GET"))
.and(path("/api/v0/instances/"))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
"instances": [{
"id": 789,
"label": "node-789",
"actual_status": "loading",
"ssh_host": "ssh5.vast.ai",
"ssh_port": 22017,
"public_ipaddr": ""
}]
})))
.mount(&server)
.await;
let policy = LifecyclePolicy {
poll_interval: Duration::from_secs(60),
state_timeout: Duration::from_secs(300),
..LifecyclePolicy::default()
};
let endpoint = wait_for_ssh_endpoint_with_policy(
&reqwest::Client::new(),
&server.uri(),
"secret",
789,
"node-789",
&policy,
)
.await
.expect("known endpoint should start bootstrap observation before running status");
assert_eq!(
endpoint,
RunningInstance {
ip: "ssh5.vast.ai".to_owned(),
port: 22017,
}
);
let requests = server.received_requests().await.expect("requests recorded");
assert!(
requests
.iter()
.all(|request| request.url.path() != "/api/v0/instances/789/"),
"endpoint discovery should not wait for provider-running status once label data has a usable endpoint"
);
}
#[tokio::test]
async fn endpoint_missing_after_running_grace_fails() {
let server = MockServer::start().await;
Mock::given(method("GET"))
.and(path("/api/v0/instances/"))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
"instances": [{
"id": 900,
"label": "node-900",
"actual_status": "running",
"ssh_host": "",
"ssh_port": 0,
"public_ipaddr": ""
}]
})))
.mount(&server)
.await;
Mock::given(method("GET"))
.and(path("/api/v0/instances/900/"))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
"instances": {
"actual_status": "running",
"intended_status": "running"
}
})))
.mount(&server)
.await;
let policy = LifecyclePolicy {
poll_interval: Duration::from_millis(1),
state_timeout: Duration::from_millis(1),
..LifecyclePolicy::default()
};
let error = wait_for_ssh_endpoint_with_policy(
&reqwest::Client::new(),
&server.uri(),
"secret",
900,
"node-900",
&policy,
)
.await
.expect_err("running without endpoint past grace is terminal");
assert!(
error.contains("running without usable SSH endpoint"),
"error should identify endpoint classification: {error}"
);
}
#[tokio::test]
async fn missing_instance_returns_error_instead_of_polling_forever() {
let server = MockServer::start().await;

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

@ -45,6 +45,31 @@ pub struct RunningInstance {
pub port: u16,
}
/// Provider status snapshot for one Vast.ai contract.
#[derive(Debug, Clone, PartialEq)]
pub struct ProviderInstanceStatus {
pub actual_status: String,
pub intended_status: String,
pub status_msg: Option<String>,
pub public_ipaddr: Option<String>,
pub ssh_port: Option<u16>,
pub disk_usage: Option<f64>,
}
impl ProviderInstanceStatus {
pub fn ssh_endpoint(&self) -> Option<RunningInstance> {
let ip = self
.public_ipaddr
.as_deref()
.filter(|ip| !ip.trim().is_empty())?;
let port = self.ssh_port.filter(|port| *port > 0)?;
Some(RunningInstance {
ip: ip.to_owned(),
port,
})
}
}
/// SSH endpoint + identity of a held instance, discovered by label.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct LabeledInstance {
@ -300,6 +325,23 @@ pub(crate) struct InstanceStatus {
pub disk_usage: Option<f64>,
}
impl From<InstanceStatus> for ProviderInstanceStatus {
fn from(status: InstanceStatus) -> Self {
Self {
actual_status: status
.actual_status
.unwrap_or_else(|| "unknown".to_string()),
intended_status: status
.intended_status
.unwrap_or_else(|| "unknown".to_string()),
status_msg: status.status_msg,
public_ipaddr: status.public_ipaddr,
ssh_port: status.ssh_port,
disk_usage: status.disk_usage,
}
}
}
#[derive(Debug, Deserialize)]
pub(crate) struct InstanceListResponse {
pub instances: Vec<InstanceListEntry>,

1422
xtask/src/main.rs
View file

File diff suppressed because it is too large Load diff