feat: per-stage GGUF weight sharding and deploy hardening
Distribute only each stage's GGUF layer slice over HTTP, add sampler and weight-load health telemetry, and harden node-image build, orchestrator provisioning, and the VastAI lease/search path. - gguf_shard (new): StageShardPlan and plan_stage_shard parse the GGUF directory and compute coalesced per-stage tensor byte ranges; materialize_stage_shard_http fetches only those ranges (plus the header) to build a stage-local GGUF, with planned_fetch_bytes accounting. - orchestrator_app: build a BTreeMap<u32, StageShardPlan> from the run plan for HuggingFace sources, thread stage_shard_plan through StageProvisionWire and weight-load, emit stage_shard_plan summaries, and add liveness phases (prefetching/fetching_stage_shard, cache_ready, stage_shard_ready). - worker_node: add a stage-shard-fetcher subcommand and materialize_stage_shard_with_process that spawns the fetcher, streams its stdout/stderr as stage_shard_fetch events (StageShardCacheReady/StageShardReady), caches under MVP_MODEL_CACHE_DIR, and feeds the local shard path into load_weights. - worker_node: add NODE_SAMPLER_CHANNEL and SamplerHealth telemetry (gpu/cpu/net samplers emit started/waiting/ready/failed) plus structured helper stdout/stderr streaming (wait_for_helper_event/drain_worker_stderr). - node_image: expand node-image build/push handling for the deploy path. - tools/vastai: extend lease, search, and types and drop unused pricing code. Signed-off-by: Zachery Aaron Shores-Chmielewski <zacheryasc@gmail.com>
This commit is contained in:
parent
4605700af6
commit
187c81498a
22 changed files with 5750 additions and 521 deletions
77
Cargo.lock
generated
77
Cargo.lock
generated
|
|
@ -2,6 +2,12 @@
|
|||
# It is not intended for manual editing.
|
||||
version = 4
|
||||
|
||||
[[package]]
|
||||
name = "adler2"
|
||||
version = "2.0.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "320119579fcad9c21884f5c4861d16174d0e06250625266f50fe6898340abefa"
|
||||
|
||||
[[package]]
|
||||
name = "aead"
|
||||
version = "0.5.2"
|
||||
|
|
@ -671,6 +677,15 @@ dependencies = [
|
|||
"libc",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "crc32fast"
|
||||
version = "1.5.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "9481c1c90cbf2ac953f07c8d4a58aa3945c425b7185c9154d67a65e4230da511"
|
||||
dependencies = [
|
||||
"cfg-if",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "criterion"
|
||||
version = "0.5.1"
|
||||
|
|
@ -1265,6 +1280,16 @@ version = "0.1.9"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "5baebc0774151f905a1a2cc41989300b1e6fbb29aff0ceffa1064fdd3088d582"
|
||||
|
||||
[[package]]
|
||||
name = "flate2"
|
||||
version = "1.1.9"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "843fba2746e448b37e26a819579957415c8cef339bf08564fe8b7ddbd959573c"
|
||||
dependencies = [
|
||||
"crc32fast",
|
||||
"miniz_oxide",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "fnv"
|
||||
version = "1.0.7"
|
||||
|
|
@ -1757,7 +1782,7 @@ dependencies = [
|
|||
"tokio",
|
||||
"tokio-rustls",
|
||||
"tower-service",
|
||||
"webpki-roots",
|
||||
"webpki-roots 1.0.8",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
|
@ -2062,7 +2087,7 @@ dependencies = [
|
|||
"tracing",
|
||||
"url",
|
||||
"wasm-bindgen-futures",
|
||||
"webpki-roots",
|
||||
"webpki-roots 1.0.8",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
|
@ -2207,7 +2232,7 @@ dependencies = [
|
|||
"tracing-subscriber",
|
||||
"url",
|
||||
"vergen-gitcl",
|
||||
"webpki-roots",
|
||||
"webpki-roots 1.0.8",
|
||||
"ws_stream_wasm",
|
||||
]
|
||||
|
||||
|
|
@ -2443,6 +2468,16 @@ version = "0.2.1"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a"
|
||||
|
||||
[[package]]
|
||||
name = "miniz_oxide"
|
||||
version = "0.8.9"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "1fa76a2c86f704bdb222d66965fb3d63269ce38518b83cb0575fca855ebb6316"
|
||||
dependencies = [
|
||||
"adler2",
|
||||
"simd-adler32",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "mio"
|
||||
version = "1.2.1"
|
||||
|
|
@ -2491,6 +2526,7 @@ dependencies = [
|
|||
"swactor-vastai",
|
||||
"tokio",
|
||||
"toml 0.8.23",
|
||||
"ureq",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
|
@ -3662,7 +3698,7 @@ dependencies = [
|
|||
"wasm-bindgen",
|
||||
"wasm-bindgen-futures",
|
||||
"web-sys",
|
||||
"webpki-roots",
|
||||
"webpki-roots 1.0.8",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
|
@ -4184,6 +4220,12 @@ version = "3.0.0"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "28d567dcbaf0049cb8ac2608a76cd95ff9e4412e1899d389ee400918ca7537f5"
|
||||
|
||||
[[package]]
|
||||
name = "simd-adler32"
|
||||
version = "0.3.10"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "3a219298ac11a56ea9a6d2120044824d6f01aeb034955e7af7bc16858527deea"
|
||||
|
||||
[[package]]
|
||||
name = "simd_cesu8"
|
||||
version = "1.1.1"
|
||||
|
|
@ -4639,7 +4681,7 @@ dependencies = [
|
|||
"time",
|
||||
"tokio",
|
||||
"tokio-rustls",
|
||||
"webpki-roots",
|
||||
"webpki-roots 1.0.8",
|
||||
"x509-parser",
|
||||
]
|
||||
|
||||
|
|
@ -4956,6 +4998,22 @@ version = "0.9.0"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "8ecb6da28b8a351d773b68d5825ac39017e680750f980f3a1a85cd8dd28a47c1"
|
||||
|
||||
[[package]]
|
||||
name = "ureq"
|
||||
version = "2.12.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "02d1a66277ed75f640d608235660df48c8e3c19f3b4edb6a263315626cc3c01d"
|
||||
dependencies = [
|
||||
"base64",
|
||||
"flate2",
|
||||
"log",
|
||||
"once_cell",
|
||||
"rustls",
|
||||
"rustls-pki-types",
|
||||
"url",
|
||||
"webpki-roots 0.26.11",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "url"
|
||||
version = "2.5.8"
|
||||
|
|
@ -5223,6 +5281,15 @@ dependencies = [
|
|||
"rustls-pki-types",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "webpki-roots"
|
||||
version = "0.26.11"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "521bc38abb08001b01866da9f51eb7c5d647a19260e00054a8c7fd5f9e57f7a9"
|
||||
dependencies = [
|
||||
"webpki-roots 1.0.8",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "webpki-roots"
|
||||
version = "1.0.8"
|
||||
|
|
|
|||
|
|
@ -24,6 +24,7 @@ swactor-vastai = { path = "../../tools/vastai" }
|
|||
parking_lot = "0.12"
|
||||
blake3 = "1"
|
||||
toml = "0.8"
|
||||
ureq = "2"
|
||||
|
||||
[target.'cfg(target_os = "linux")'.dependencies]
|
||||
libc = "0.2"
|
||||
|
|
|
|||
|
|
@ -4,7 +4,7 @@ use swactor::actor::{ActorAddress, ActorInterface};
|
|||
use swactor::runtime::Ctx;
|
||||
use swactor_transport::{CodecRegistry, NetworkMessage};
|
||||
|
||||
use crate::{run_plan, stage_controller as stage};
|
||||
use crate::{gguf_shard::StageShardPlan, run_plan, stage_controller as stage};
|
||||
|
||||
use super::codec::JsonCodec;
|
||||
use super::orchestrator::OrchestratorMsg;
|
||||
|
|
@ -98,6 +98,7 @@ pub struct StageProvisionWire {
|
|||
pub model_id: String,
|
||||
pub gguf_source: run_plan::GgufSource,
|
||||
pub tokenizer: run_plan::TokenizerSource,
|
||||
pub stage_shard_plan: Option<StageShardPlan>,
|
||||
}
|
||||
|
||||
impl StageProvisionWire {
|
||||
|
|
@ -119,6 +120,7 @@ impl StageProvisionWire {
|
|||
self.gguf_source.clone(),
|
||||
self.tokenizer.clone(),
|
||||
),
|
||||
shard_plan: self.stage_shard_plan.clone(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -228,6 +230,7 @@ pub enum StageCommandWire {
|
|||
tokenizer: run_plan::TokenizerSource,
|
||||
layer_start: u32,
|
||||
layer_end_exclusive: u32,
|
||||
stage_shard_plan: Option<StageShardPlan>,
|
||||
},
|
||||
RewireEdge {
|
||||
edge_id: u64,
|
||||
|
|
@ -663,12 +666,17 @@ impl From<&stage::StageCommand> for StageCommandWire {
|
|||
layer_start: layer_range.start,
|
||||
layer_end_exclusive: layer_range.end_exclusive,
|
||||
},
|
||||
stage::StageCommand::LoadWeights { source, range } => Self::LoadWeights {
|
||||
stage::StageCommand::LoadWeights {
|
||||
source,
|
||||
range,
|
||||
shard_plan,
|
||||
} => Self::LoadWeights {
|
||||
model_id: source.model_id.clone(),
|
||||
gguf_source: source.gguf_source.clone(),
|
||||
tokenizer: source.tokenizer.clone(),
|
||||
layer_start: range.start,
|
||||
layer_end_exclusive: range.end_exclusive,
|
||||
stage_shard_plan: shard_plan.clone(),
|
||||
},
|
||||
stage::StageCommand::RewireEdge { edge_id } => Self::RewireEdge { edge_id: edge_id.0 },
|
||||
stage::StageCommand::ExecuteStep(step) => Self::ExecuteStep {
|
||||
|
|
|
|||
|
|
@ -27,7 +27,8 @@ use mvp_system::config as chat_config;
|
|||
use mvp_system::config::ResolvedVastAiConfig;
|
||||
use mvp_system::endpoint_advertisement::EndpointAddrMask;
|
||||
use mvp_system::node_image::{
|
||||
NodeImageProvider, NodeImageRequest, PreparedNodeImage, prepare_node_image,
|
||||
NodeImageProgressEvent, NodeImageProgressEventKind, NodeImageProgressSink, NodeImageProvider,
|
||||
NodeImageRequest, PreparedNodeImage, prepare_node_image_with_progress,
|
||||
};
|
||||
use mvp_system::node_provisioning::ProviderKind;
|
||||
use mvp_system::prompt_rpc::{PromptEvent, SubmitPrompt, write_json_line};
|
||||
|
|
@ -154,34 +155,38 @@ where
|
|||
);
|
||||
progress.emit_benchmark_envelope(&config);
|
||||
confirm_vastai_if_needed(&config)?;
|
||||
let prepare_runtime_started = Instant::now();
|
||||
progress.emit(
|
||||
CHAT_RUNTIME_CHANNEL,
|
||||
"prepare_runtime",
|
||||
"started",
|
||||
json!({"provider": config.provider.as_str()}),
|
||||
);
|
||||
let image_ref =
|
||||
match prepare_runtime_with_progress(&config, prepare_node_image, Some(&mut progress)) {
|
||||
Ok(image_ref) => {
|
||||
progress.emit(
|
||||
let image_ref = match prepare_runtime_with_progress(
|
||||
&config,
|
||||
prepare_node_image_progress_adapter,
|
||||
Some(&mut progress),
|
||||
) {
|
||||
Ok(image_ref) => {
|
||||
progress.emit(
|
||||
CHAT_RUNTIME_CHANNEL,
|
||||
"prepare_runtime",
|
||||
"ready",
|
||||
json!({"image_ref": image_ref}),
|
||||
json!({"image_ref": image_ref, "elapsed_ms": prepare_runtime_started.elapsed().as_millis()}),
|
||||
);
|
||||
image_ref
|
||||
}
|
||||
Err(error) => {
|
||||
progress.emit(
|
||||
image_ref
|
||||
}
|
||||
Err(error) => {
|
||||
progress.emit(
|
||||
CHAT_RUNTIME_CHANNEL,
|
||||
"prepare_runtime",
|
||||
"failed",
|
||||
json!({"error": error}),
|
||||
json!({"error": error, "elapsed_ms": prepare_runtime_started.elapsed().as_millis()}),
|
||||
);
|
||||
progress.archive_pending()?;
|
||||
return Err(error);
|
||||
}
|
||||
};
|
||||
progress.archive_pending()?;
|
||||
return Err(error);
|
||||
}
|
||||
};
|
||||
progress.emit(
|
||||
CHAT_COMPONENT_CHANNEL,
|
||||
"orchestrator_process_spawn",
|
||||
|
|
@ -510,6 +515,66 @@ impl ChatDatastream {
|
|||
}
|
||||
}
|
||||
|
||||
impl NodeImageProgressSink for ChatDatastream {
|
||||
fn emit(&mut self, event: NodeImageProgressEvent) {
|
||||
let mut detail = serde_json::Map::new();
|
||||
if let Some(command_label) = event.command_label {
|
||||
detail.insert("command_label".to_owned(), json!(command_label));
|
||||
}
|
||||
if let Some(image_ref) = event.image_ref {
|
||||
detail.insert("image_ref".to_owned(), json!(image_ref));
|
||||
}
|
||||
if let Some(elapsed_ms) = event.elapsed_ms {
|
||||
detail.insert("elapsed_ms".to_owned(), json!(elapsed_ms));
|
||||
}
|
||||
|
||||
let (phase, status) = match event.kind {
|
||||
NodeImageProgressEventKind::ImageReference { role, image_ref } => {
|
||||
detail.insert("event".to_owned(), json!("image_ref"));
|
||||
detail.insert("role".to_owned(), json!(role));
|
||||
detail.insert("image_ref".to_owned(), json!(image_ref));
|
||||
("prepare_node_image", "image_ref")
|
||||
}
|
||||
NodeImageProgressEventKind::CommandStarted { program, args } => {
|
||||
detail.insert("event".to_owned(), json!("command_start"));
|
||||
detail.insert("program".to_owned(), json!(program));
|
||||
detail.insert("args".to_owned(), json!(args));
|
||||
("node_image_command", "started")
|
||||
}
|
||||
NodeImageProgressEventKind::CommandStdout { line } => {
|
||||
detail.insert("event".to_owned(), json!("stdout"));
|
||||
detail.insert("stream".to_owned(), json!("stdout"));
|
||||
detail.insert("line".to_owned(), json!(line));
|
||||
("node_image_command", "stdout")
|
||||
}
|
||||
NodeImageProgressEventKind::CommandStderr { line } => {
|
||||
detail.insert("event".to_owned(), json!("stderr"));
|
||||
detail.insert("stream".to_owned(), json!("stderr"));
|
||||
detail.insert("line".to_owned(), json!(line));
|
||||
("node_image_command", "stderr")
|
||||
}
|
||||
NodeImageProgressEventKind::CommandExited {
|
||||
status: command_status,
|
||||
code,
|
||||
success,
|
||||
} => {
|
||||
detail.insert("event".to_owned(), json!("command_exit"));
|
||||
detail.insert("command_status".to_owned(), json!(command_status));
|
||||
detail.insert("exit_code".to_owned(), json!(code));
|
||||
detail.insert("success".to_owned(), json!(success));
|
||||
if let Some(elapsed_ms) = detail.get("elapsed_ms").cloned() {
|
||||
detail.insert("duration_ms".to_owned(), elapsed_ms);
|
||||
}
|
||||
(
|
||||
"node_image_command",
|
||||
if success { "exited" } else { "failed" },
|
||||
)
|
||||
}
|
||||
};
|
||||
self.emit(CHAT_RUNTIME_CHANNEL, phase, status, Value::Object(detail));
|
||||
}
|
||||
}
|
||||
|
||||
struct ChatFrameArchive {
|
||||
file: File,
|
||||
next_seq: u64,
|
||||
|
|
@ -1366,26 +1431,44 @@ fn signal_orch_process_group(child: &Child, signal: libc::c_int) -> io::Result<(
|
|||
}
|
||||
}
|
||||
|
||||
type PrepareNodeImageFn = fn(NodeImageRequest) -> Result<PreparedNodeImage, String>;
|
||||
fn prepare_node_image_progress_adapter(
|
||||
request: NodeImageRequest,
|
||||
progress: Option<&mut dyn NodeImageProgressSink>,
|
||||
) -> Result<PreparedNodeImage, String> {
|
||||
prepare_node_image_with_progress(request, progress)
|
||||
}
|
||||
|
||||
#[allow(dead_code)]
|
||||
fn prepare_runtime(config: &Config) -> Result<String, String> {
|
||||
prepare_runtime_with(config, prepare_node_image)
|
||||
prepare_runtime_with(config, |request| {
|
||||
prepare_node_image_with_progress(request, None)
|
||||
})
|
||||
}
|
||||
|
||||
#[allow(dead_code)]
|
||||
fn prepare_runtime_with(
|
||||
config: &Config,
|
||||
prepare_node_image_fn: PrepareNodeImageFn,
|
||||
) -> Result<String, String> {
|
||||
prepare_runtime_with_progress(config, prepare_node_image_fn, None)
|
||||
fn prepare_runtime_with<F>(config: &Config, prepare_node_image_fn: F) -> Result<String, String>
|
||||
where
|
||||
F: FnMut(NodeImageRequest) -> Result<PreparedNodeImage, String>,
|
||||
{
|
||||
let mut prepare_node_image_fn = prepare_node_image_fn;
|
||||
prepare_runtime_with_progress(
|
||||
config,
|
||||
move |request, _progress| prepare_node_image_fn(request),
|
||||
None,
|
||||
)
|
||||
}
|
||||
|
||||
fn prepare_runtime_with_progress(
|
||||
fn prepare_runtime_with_progress<F>(
|
||||
config: &Config,
|
||||
prepare_node_image_fn: PrepareNodeImageFn,
|
||||
mut prepare_node_image_fn: F,
|
||||
progress: Option<&mut ChatDatastream>,
|
||||
) -> Result<String, String> {
|
||||
) -> Result<String, String>
|
||||
where
|
||||
F: FnMut(
|
||||
NodeImageRequest,
|
||||
Option<&mut dyn NodeImageProgressSink>,
|
||||
) -> Result<PreparedNodeImage, String>,
|
||||
{
|
||||
let mut progress = progress;
|
||||
let binary_mode = if config.skip_rebuild {
|
||||
"existing_artifact"
|
||||
|
|
@ -1408,12 +1491,13 @@ fn prepare_runtime_with_progress(
|
|||
json!({"mode": config.orchestrator_launch_mode()}),
|
||||
);
|
||||
} else {
|
||||
let ensure_orch_started = Instant::now();
|
||||
emit_chat_progress(
|
||||
&mut progress,
|
||||
CHAT_RUNTIME_CHANNEL,
|
||||
"ensure_orch_binary",
|
||||
"started",
|
||||
json!({"mode": binary_mode}),
|
||||
json!({"mode": binary_mode, "command_label": "ensure_orch_binary"}),
|
||||
);
|
||||
match ensure_orch_binary(config) {
|
||||
Ok(()) => emit_chat_progress(
|
||||
|
|
@ -1421,7 +1505,7 @@ fn prepare_runtime_with_progress(
|
|||
CHAT_RUNTIME_CHANNEL,
|
||||
"ensure_orch_binary",
|
||||
"ready",
|
||||
json!({"mode": binary_mode}),
|
||||
json!({"mode": binary_mode, "command_label": "ensure_orch_binary", "elapsed_ms": ensure_orch_started.elapsed().as_millis()}),
|
||||
),
|
||||
Err(error) => {
|
||||
emit_chat_progress(
|
||||
|
|
@ -1429,7 +1513,7 @@ fn prepare_runtime_with_progress(
|
|||
CHAT_RUNTIME_CHANNEL,
|
||||
"ensure_orch_binary",
|
||||
"failed",
|
||||
json!({"mode": binary_mode, "error": error.as_str()}),
|
||||
json!({"mode": binary_mode, "command_label": "ensure_orch_binary", "elapsed_ms": ensure_orch_started.elapsed().as_millis(), "error": error.as_str()}),
|
||||
);
|
||||
return Err(error);
|
||||
}
|
||||
|
|
@ -1468,7 +1552,7 @@ fn prepare_runtime_with_progress(
|
|||
CHAT_RUNTIME_CHANNEL,
|
||||
"prepare_node_image",
|
||||
"skipped",
|
||||
json!({"provider": config.provider.as_str(), "reason": "process_provider"}),
|
||||
json!({"provider": config.provider.as_str(), "command_label": "prepare_node_image", "reason": "process_provider"}),
|
||||
);
|
||||
return Ok(config.node_image.clone());
|
||||
}
|
||||
|
|
@ -1480,23 +1564,24 @@ fn prepare_runtime_with_progress(
|
|||
CHAT_RUNTIME_CHANNEL,
|
||||
"ensure_worker_binary",
|
||||
"skipped",
|
||||
json!({"mode": binary_mode, "reason": "vastai_remote_image"}),
|
||||
json!({"mode": binary_mode, "command_label": "ensure_worker_binary", "reason": "vastai_remote_image"}),
|
||||
);
|
||||
emit_chat_progress(
|
||||
&mut progress,
|
||||
CHAT_RUNTIME_CHANNEL,
|
||||
"prepare_node_image",
|
||||
"skipped",
|
||||
json!({"provider": config.provider.as_str(), "reason": "skip_rebuild"}),
|
||||
json!({"provider": config.provider.as_str(), "command_label": "prepare_node_image", "reason": "skip_rebuild"}),
|
||||
);
|
||||
return Ok(config.node_image.clone());
|
||||
}
|
||||
let ensure_worker_started = Instant::now();
|
||||
emit_chat_progress(
|
||||
&mut progress,
|
||||
CHAT_RUNTIME_CHANNEL,
|
||||
"ensure_worker_binary",
|
||||
"started",
|
||||
json!({"mode": binary_mode}),
|
||||
json!({"mode": binary_mode, "command_label": "ensure_worker_binary"}),
|
||||
);
|
||||
match ensure_worker_binary(config) {
|
||||
Ok(()) => emit_chat_progress(
|
||||
|
|
@ -1504,7 +1589,7 @@ fn prepare_runtime_with_progress(
|
|||
CHAT_RUNTIME_CHANNEL,
|
||||
"ensure_worker_binary",
|
||||
"ready",
|
||||
json!({"mode": binary_mode}),
|
||||
json!({"mode": binary_mode, "command_label": "ensure_worker_binary", "elapsed_ms": ensure_worker_started.elapsed().as_millis()}),
|
||||
),
|
||||
Err(error) => {
|
||||
emit_chat_progress(
|
||||
|
|
@ -1512,7 +1597,7 @@ fn prepare_runtime_with_progress(
|
|||
CHAT_RUNTIME_CHANNEL,
|
||||
"ensure_worker_binary",
|
||||
"failed",
|
||||
json!({"mode": binary_mode, "error": error.as_str()}),
|
||||
json!({"mode": binary_mode, "command_label": "ensure_worker_binary", "elapsed_ms": ensure_worker_started.elapsed().as_millis(), "error": error.as_str()}),
|
||||
);
|
||||
return Err(error);
|
||||
}
|
||||
|
|
@ -1522,17 +1607,18 @@ fn prepare_runtime_with_progress(
|
|||
CHAT_RUNTIME_CHANNEL,
|
||||
"prepare_node_image",
|
||||
"skipped",
|
||||
json!({"provider": config.provider.as_str(), "reason": "skip_rebuild"}),
|
||||
json!({"provider": config.provider.as_str(), "command_label": "prepare_node_image", "reason": "skip_rebuild"}),
|
||||
);
|
||||
return Ok(config.node_image.clone());
|
||||
}
|
||||
|
||||
let prepare_node_image_started = Instant::now();
|
||||
emit_chat_progress(
|
||||
&mut progress,
|
||||
CHAT_RUNTIME_CHANNEL,
|
||||
"prepare_node_image",
|
||||
"started",
|
||||
json!({"provider": config.provider.as_str()}),
|
||||
json!({"provider": config.provider.as_str(), "command_label": "prepare_node_image", "image_tag": config.image_tag.as_deref()}),
|
||||
);
|
||||
let node_bin = match node_bin_for_current_profile() {
|
||||
Ok(path) => path,
|
||||
|
|
@ -1542,7 +1628,7 @@ fn prepare_runtime_with_progress(
|
|||
CHAT_RUNTIME_CHANNEL,
|
||||
"prepare_node_image",
|
||||
"failed",
|
||||
json!({"provider": config.provider.as_str(), "error": error.as_str()}),
|
||||
json!({"provider": config.provider.as_str(), "command_label": "prepare_node_image", "elapsed_ms": prepare_node_image_started.elapsed().as_millis(), "error": error.as_str()}),
|
||||
);
|
||||
return Err(error);
|
||||
}
|
||||
|
|
@ -1555,31 +1641,39 @@ fn prepare_runtime_with_progress(
|
|||
CHAT_RUNTIME_CHANNEL,
|
||||
"prepare_node_image",
|
||||
"failed",
|
||||
json!({"provider": config.provider.as_str(), "error": error.as_str()}),
|
||||
json!({"provider": config.provider.as_str(), "command_label": "prepare_node_image", "elapsed_ms": prepare_node_image_started.elapsed().as_millis(), "error": error.as_str()}),
|
||||
);
|
||||
return Err(error);
|
||||
}
|
||||
};
|
||||
let prepared = match prepare_node_image_fn(NodeImageRequest {
|
||||
requested_image: config.node_image.clone(),
|
||||
base_image: BASE_NODE_IMAGE.to_owned(),
|
||||
node_bin,
|
||||
provider,
|
||||
extra_tag: config.image_tag.clone(),
|
||||
push: false,
|
||||
force_refresh: false,
|
||||
enabled: true,
|
||||
}) {
|
||||
Ok(prepared) => prepared,
|
||||
Err(error) => {
|
||||
emit_chat_progress(
|
||||
&mut progress,
|
||||
CHAT_RUNTIME_CHANNEL,
|
||||
"prepare_node_image",
|
||||
"failed",
|
||||
json!({"provider": config.provider.as_str(), "error": error.as_str()}),
|
||||
);
|
||||
return Err(error);
|
||||
let prepared = {
|
||||
let command_progress = progress
|
||||
.as_deref_mut()
|
||||
.map(|sink| sink as &mut dyn NodeImageProgressSink);
|
||||
match prepare_node_image_fn(
|
||||
NodeImageRequest {
|
||||
requested_image: config.node_image.clone(),
|
||||
base_image: BASE_NODE_IMAGE.to_owned(),
|
||||
node_bin,
|
||||
provider,
|
||||
extra_tag: config.image_tag.clone(),
|
||||
push: false,
|
||||
force_refresh: false,
|
||||
enabled: true,
|
||||
},
|
||||
command_progress,
|
||||
) {
|
||||
Ok(prepared) => prepared,
|
||||
Err(error) => {
|
||||
emit_chat_progress(
|
||||
&mut progress,
|
||||
CHAT_RUNTIME_CHANNEL,
|
||||
"prepare_node_image",
|
||||
"failed",
|
||||
json!({"provider": config.provider.as_str(), "command_label": "prepare_node_image", "elapsed_ms": prepare_node_image_started.elapsed().as_millis(), "error": error.as_str()}),
|
||||
);
|
||||
return Err(error);
|
||||
}
|
||||
}
|
||||
};
|
||||
emit_chat_progress(
|
||||
|
|
@ -1587,7 +1681,7 @@ fn prepare_runtime_with_progress(
|
|||
CHAT_RUNTIME_CHANNEL,
|
||||
"prepare_node_image",
|
||||
"ready",
|
||||
json!({"provider": config.provider.as_str(), "image_ref": prepared.image_ref}),
|
||||
json!({"provider": config.provider.as_str(), "command_label": "prepare_node_image", "image_ref": prepared.image_ref, "elapsed_ms": prepare_node_image_started.elapsed().as_millis()}),
|
||||
);
|
||||
Ok(prepared.image_ref)
|
||||
}
|
||||
|
|
@ -3088,6 +3182,112 @@ relay_url = "https://relay.example"
|
|||
panic!("image preparer must not be called when --skip-rebuild is set")
|
||||
}
|
||||
|
||||
fn panic_prepare_node_image_with_progress(
|
||||
_request: NodeImageRequest,
|
||||
_progress: Option<&mut dyn NodeImageProgressSink>,
|
||||
) -> Result<PreparedNodeImage, String> {
|
||||
panic!("image preparer must not be called when --skip-rebuild is set")
|
||||
}
|
||||
|
||||
fn runtime_events(path: &Path) -> Vec<serde_json::Value> {
|
||||
fs::read_to_string(path)
|
||||
.expect("read progress archive")
|
||||
.lines()
|
||||
.filter_map(|line| {
|
||||
let outer: serde_json::Value = serde_json::from_str(line).ok()?;
|
||||
if outer.get("channel").and_then(serde_json::Value::as_str)
|
||||
!= Some(CHAT_RUNTIME_CHANNEL)
|
||||
{
|
||||
return None;
|
||||
}
|
||||
outer
|
||||
.get("payload")?
|
||||
.get("value")?
|
||||
.as_str()
|
||||
.and_then(|text| serde_json::from_str::<serde_json::Value>(text).ok())
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn emit_fake_node_image_progress(
|
||||
progress: Option<&mut dyn NodeImageProgressSink>,
|
||||
success: bool,
|
||||
) {
|
||||
let Some(sink) = progress else {
|
||||
return;
|
||||
};
|
||||
sink.emit(NodeImageProgressEvent {
|
||||
command_label: None,
|
||||
image_ref: Some("docker.io/acme/node:prepared".to_owned()),
|
||||
elapsed_ms: None,
|
||||
kind: NodeImageProgressEventKind::ImageReference {
|
||||
role: "resolved".to_owned(),
|
||||
image_ref: "docker.io/acme/node:prepared".to_owned(),
|
||||
},
|
||||
});
|
||||
sink.emit(NodeImageProgressEvent {
|
||||
command_label: Some("build mvp node image".to_owned()),
|
||||
image_ref: Some("docker.io/acme/node:prepared".to_owned()),
|
||||
elapsed_ms: Some(0),
|
||||
kind: NodeImageProgressEventKind::CommandStarted {
|
||||
program: "fake-docker".to_owned(),
|
||||
args: vec!["build".to_owned()],
|
||||
},
|
||||
});
|
||||
sink.emit(NodeImageProgressEvent {
|
||||
command_label: Some("build mvp node image".to_owned()),
|
||||
image_ref: Some("docker.io/acme/node:prepared".to_owned()),
|
||||
elapsed_ms: Some(1),
|
||||
kind: NodeImageProgressEventKind::CommandStdout {
|
||||
line: "building layer".to_owned(),
|
||||
},
|
||||
});
|
||||
sink.emit(NodeImageProgressEvent {
|
||||
command_label: Some("build mvp node image".to_owned()),
|
||||
image_ref: Some("docker.io/acme/node:prepared".to_owned()),
|
||||
elapsed_ms: Some(2),
|
||||
kind: NodeImageProgressEventKind::CommandStderr {
|
||||
line: "pushing metadata".to_owned(),
|
||||
},
|
||||
});
|
||||
sink.emit(NodeImageProgressEvent {
|
||||
command_label: Some("build mvp node image".to_owned()),
|
||||
image_ref: Some("docker.io/acme/node:prepared".to_owned()),
|
||||
elapsed_ms: Some(3),
|
||||
kind: NodeImageProgressEventKind::CommandExited {
|
||||
status: if success {
|
||||
"exit status: 0".to_owned()
|
||||
} else {
|
||||
"exit status: 42".to_owned()
|
||||
},
|
||||
code: Some(if success { 0 } else { 42 }),
|
||||
success,
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
fn fake_prepare_node_image_with_progress(
|
||||
_request: NodeImageRequest,
|
||||
progress: Option<&mut dyn NodeImageProgressSink>,
|
||||
) -> Result<PreparedNodeImage, String> {
|
||||
emit_fake_node_image_progress(progress, true);
|
||||
Ok(PreparedNodeImage {
|
||||
image_ref: "docker.io/acme/node:prepared".to_owned(),
|
||||
tag: "prepared".to_owned(),
|
||||
already_available: false,
|
||||
built: true,
|
||||
pushed: false,
|
||||
})
|
||||
}
|
||||
|
||||
fn failing_prepare_node_image_with_progress(
|
||||
_request: NodeImageRequest,
|
||||
progress: Option<&mut dyn NodeImageProgressSink>,
|
||||
) -> Result<PreparedNodeImage, String> {
|
||||
emit_fake_node_image_progress(progress, false);
|
||||
Err("build mvp node image failed with exit status: 42".to_owned())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn skip_rebuild_requires_existing_artifacts_and_skips_image_preparation() {
|
||||
let temp = TempDir::new("skip-rebuild");
|
||||
|
|
@ -3110,6 +3310,188 @@ relay_url = "https://relay.example"
|
|||
assert_eq!(image_ref, "docker.io/acme/node:latest");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prepare_runtime_progress_records_local_prep_details() {
|
||||
let temp = TempDir::new("prep-progress");
|
||||
let archive_path = temp.path().join("frames.ndjson");
|
||||
let orch_bin = temp.path().join("mvp-orchestrator");
|
||||
let worker_bin = temp.path().join("mvp-worker-node");
|
||||
fs::write(&orch_bin, b"orch").expect("write orchestrator artifact");
|
||||
fs::write(&worker_bin, b"worker").expect("write worker artifact");
|
||||
let mut config = base_config(ProviderKind::Docker);
|
||||
config.skip_rebuild = true;
|
||||
config.orch_bin = orch_bin;
|
||||
config.worker_bin = worker_bin;
|
||||
config.node_image = "docker.io/acme/node:latest".to_owned();
|
||||
let mut progress =
|
||||
ChatDatastream::new(91, Some(archive_path.clone())).expect("datastream constructs");
|
||||
|
||||
prepare_runtime_with_progress(
|
||||
&config,
|
||||
panic_prepare_node_image_with_progress,
|
||||
Some(&mut progress),
|
||||
)
|
||||
.expect("skip rebuild uses existing artifacts");
|
||||
progress.archive_pending().expect("archive prep frames");
|
||||
|
||||
let events = fs::read_to_string(&archive_path).expect("read archive");
|
||||
let inner_events = events
|
||||
.lines()
|
||||
.filter_map(|line| {
|
||||
let outer: serde_json::Value = serde_json::from_str(line).ok()?;
|
||||
if outer.get("channel").and_then(serde_json::Value::as_str)
|
||||
!= Some(CHAT_RUNTIME_CHANNEL)
|
||||
{
|
||||
return None;
|
||||
}
|
||||
let inner = outer
|
||||
.get("payload")?
|
||||
.get("value")?
|
||||
.as_str()
|
||||
.and_then(|text| serde_json::from_str::<serde_json::Value>(text).ok())?;
|
||||
Some(inner)
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
let ensure_ready = inner_events
|
||||
.iter()
|
||||
.find(|event| {
|
||||
event.get("phase").and_then(serde_json::Value::as_str) == Some("ensure_orch_binary")
|
||||
&& event.get("status").and_then(serde_json::Value::as_str) == Some("ready")
|
||||
})
|
||||
.expect("ensure_orch_binary ready event");
|
||||
assert_eq!(
|
||||
ensure_ready
|
||||
.pointer("/detail/command_label")
|
||||
.and_then(serde_json::Value::as_str),
|
||||
Some("ensure_orch_binary")
|
||||
);
|
||||
assert!(
|
||||
ensure_ready
|
||||
.pointer("/detail/elapsed_ms")
|
||||
.and_then(serde_json::Value::as_u64)
|
||||
.is_some()
|
||||
);
|
||||
let image_skip = inner_events
|
||||
.iter()
|
||||
.find(|event| {
|
||||
event.get("phase").and_then(serde_json::Value::as_str) == Some("prepare_node_image")
|
||||
&& event.get("status").and_then(serde_json::Value::as_str) == Some("skipped")
|
||||
})
|
||||
.expect("prepare_node_image skipped event");
|
||||
assert_eq!(
|
||||
image_skip
|
||||
.pointer("/detail/reason")
|
||||
.and_then(serde_json::Value::as_str),
|
||||
Some("skip_rebuild")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prepare_runtime_streams_node_image_command_progress() {
|
||||
let temp = TempDir::new("node-image-command-progress");
|
||||
let archive_path = temp.path().join("frames.ndjson");
|
||||
let mut config = base_config(ProviderKind::Docker);
|
||||
config.skip_rebuild = false;
|
||||
config.gpu_run = true;
|
||||
config.node_image = "docker.io/acme/node:latest".to_owned();
|
||||
let mut progress =
|
||||
ChatDatastream::new(92, Some(archive_path.clone())).expect("datastream constructs");
|
||||
|
||||
let image_ref = prepare_runtime_with_progress(
|
||||
&config,
|
||||
fake_prepare_node_image_with_progress,
|
||||
Some(&mut progress),
|
||||
)
|
||||
.expect("fake image preparation succeeds");
|
||||
progress.archive_pending().expect("archive prep frames");
|
||||
assert_eq!(image_ref, "docker.io/acme/node:prepared");
|
||||
|
||||
let events = runtime_events(&archive_path);
|
||||
assert!(events.iter().any(|event| {
|
||||
event.get("phase").and_then(serde_json::Value::as_str) == Some("prepare_node_image")
|
||||
&& event.get("status").and_then(serde_json::Value::as_str) == Some("image_ref")
|
||||
&& event
|
||||
.pointer("/detail/image_ref")
|
||||
.and_then(serde_json::Value::as_str)
|
||||
== Some("docker.io/acme/node:prepared")
|
||||
}));
|
||||
assert!(events.iter().any(|event| {
|
||||
event.get("phase").and_then(serde_json::Value::as_str) == Some("node_image_command")
|
||||
&& event.get("status").and_then(serde_json::Value::as_str) == Some("started")
|
||||
&& event
|
||||
.pointer("/detail/command_label")
|
||||
.and_then(serde_json::Value::as_str)
|
||||
== Some("build mvp node image")
|
||||
&& event
|
||||
.pointer("/detail/program")
|
||||
.and_then(serde_json::Value::as_str)
|
||||
== Some("fake-docker")
|
||||
}));
|
||||
assert!(events.iter().any(|event| {
|
||||
event.get("status").and_then(serde_json::Value::as_str) == Some("stdout")
|
||||
&& event
|
||||
.pointer("/detail/line")
|
||||
.and_then(serde_json::Value::as_str)
|
||||
== Some("building layer")
|
||||
}));
|
||||
assert!(events.iter().any(|event| {
|
||||
event.get("status").and_then(serde_json::Value::as_str) == Some("stderr")
|
||||
&& event
|
||||
.pointer("/detail/line")
|
||||
.and_then(serde_json::Value::as_str)
|
||||
== Some("pushing metadata")
|
||||
}));
|
||||
assert!(events.iter().any(|event| {
|
||||
event.get("status").and_then(serde_json::Value::as_str) == Some("exited")
|
||||
&& event
|
||||
.pointer("/detail/command_status")
|
||||
.and_then(serde_json::Value::as_str)
|
||||
== Some("exit status: 0")
|
||||
&& event.pointer("/detail/duration_ms").is_some()
|
||||
}));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prepare_runtime_command_failure_preserves_label_and_status() {
|
||||
let temp = TempDir::new("node-image-command-failure");
|
||||
let archive_path = temp.path().join("frames.ndjson");
|
||||
let mut config = base_config(ProviderKind::Docker);
|
||||
config.skip_rebuild = false;
|
||||
config.gpu_run = true;
|
||||
let mut progress =
|
||||
ChatDatastream::new(93, Some(archive_path.clone())).expect("datastream constructs");
|
||||
|
||||
let error = prepare_runtime_with_progress(
|
||||
&config,
|
||||
failing_prepare_node_image_with_progress,
|
||||
Some(&mut progress),
|
||||
)
|
||||
.expect_err("fake image preparation failure propagates");
|
||||
progress.archive_pending().expect("archive prep frames");
|
||||
assert!(error.contains("build mvp node image"), "{error}");
|
||||
|
||||
let events = runtime_events(&archive_path);
|
||||
let failure = events
|
||||
.iter()
|
||||
.find(|event| {
|
||||
event.get("phase").and_then(serde_json::Value::as_str) == Some("node_image_command")
|
||||
&& event.get("status").and_then(serde_json::Value::as_str) == Some("failed")
|
||||
})
|
||||
.expect("failed command progress event");
|
||||
assert_eq!(
|
||||
failure
|
||||
.pointer("/detail/command_label")
|
||||
.and_then(serde_json::Value::as_str),
|
||||
Some("build mvp node image")
|
||||
);
|
||||
assert_eq!(
|
||||
failure
|
||||
.pointer("/detail/command_status")
|
||||
.and_then(serde_json::Value::as_str),
|
||||
Some("exit status: 42")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn vastai_skip_rebuild_uses_remote_image_without_worker_artifact() {
|
||||
let temp = TempDir::new("vastai-skip-rebuild");
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load diff
1320
crates/mvp-system/src/gguf_shard.rs
Normal file
1320
crates/mvp-system/src/gguf_shard.rs
Normal file
File diff suppressed because it is too large
Load diff
|
|
@ -15,6 +15,7 @@ pub mod edge_establisher;
|
|||
pub mod endpoint_advertisement;
|
||||
pub mod engine_builder;
|
||||
pub mod gguf_metadata;
|
||||
pub mod gguf_shard;
|
||||
pub mod gpu_worker_ctl;
|
||||
pub mod gpu_worker_egress_producer;
|
||||
pub mod gpu_worker_ingress_parser;
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load diff
File diff suppressed because it is too large
Load diff
|
|
@ -138,6 +138,20 @@ pub trait ProvisionPlugin: Send {
|
|||
sink: PluginSink,
|
||||
) -> Result<PluginNodeHandle, String>;
|
||||
|
||||
fn start_nodes(
|
||||
&mut self,
|
||||
specs: Vec<NodeProvisionSpec>,
|
||||
sink: PluginSink,
|
||||
) -> Vec<(NodeProvisionSpec, Result<PluginNodeHandle, String>)> {
|
||||
specs
|
||||
.into_iter()
|
||||
.map(|spec| {
|
||||
let result = self.start_node(spec.clone(), sink.clone());
|
||||
(spec, result)
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn complete_bootstrap(&mut self, handle: &PluginNodeHandle) -> Result<(), String>;
|
||||
|
||||
fn stop_node(&mut self, handle: &PluginNodeHandle) -> Result<(), String>;
|
||||
|
|
|
|||
|
|
@ -1,3 +1,4 @@
|
|||
use crate::gguf_shard::StageShardPlan;
|
||||
use crate::run_plan::{GgufSource, TokenizerSource};
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
|
||||
|
|
@ -97,6 +98,7 @@ pub struct ProvisionStage {
|
|||
pub inbound: EdgeProvision,
|
||||
pub outbound: EdgeProvision,
|
||||
pub weight_source: WeightSource,
|
||||
pub shard_plan: Option<StageShardPlan>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
|
|
@ -223,6 +225,7 @@ pub enum StageCommand {
|
|||
LoadWeights {
|
||||
source: WeightSource,
|
||||
range: LayerRange,
|
||||
shard_plan: Option<StageShardPlan>,
|
||||
},
|
||||
RewireEdge {
|
||||
edge_id: EdgeId,
|
||||
|
|
@ -366,6 +369,7 @@ impl StageController {
|
|||
self.commands.push(StageCommand::LoadWeights {
|
||||
source: provision.weight_source.clone(),
|
||||
range: provision.layer_range,
|
||||
shard_plan: provision.shard_plan.clone(),
|
||||
});
|
||||
self.provision = Some(provision);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -802,6 +802,7 @@ impl LocalMockCluster {
|
|||
provision.gguf_source,
|
||||
provision.tokenizer,
|
||||
),
|
||||
shard_plan: None,
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -27,6 +27,7 @@ fn valid_provision() -> stage::ProvisionStage {
|
|||
inbound: stage::EdgeProvision::inbound(stage::EdgeId(7001)),
|
||||
outbound: stage::EdgeProvision::outbound(stage::EdgeId(7002)),
|
||||
weight_source: stage::WeightSource::embedded_gguf("model", "model.gguf"),
|
||||
shard_plan: None,
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
use std::collections::VecDeque;
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use mvp_system::node_provisioning as provision;
|
||||
use mvp_system::node_provisioning::ProviderPlugin;
|
||||
|
|
@ -10,7 +11,7 @@ use mvp_system::vastai_provisioning::{
|
|||
BootstrapStopReason, VastAiBootstrapLauncher, VastAiLeaseClient, VastAiProviderPlugin,
|
||||
VastAiProvisioningConfig, VastAiProvisioningPlugin, VastAiSshEndpoint,
|
||||
};
|
||||
use parking_lot::Mutex;
|
||||
use parking_lot::{Condvar, Mutex};
|
||||
use swactor_vastai::{LifecyclePolicy, ProvisionRequest, ProvisionedInstance, SelectionPolicy};
|
||||
|
||||
#[derive(Default)]
|
||||
|
|
@ -28,7 +29,7 @@ fn sink() -> PluginSink {
|
|||
PluginSink::new(Arc::new(RecordingSink::default()))
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
#[derive(Clone, Default)]
|
||||
struct FakeLeaseClient {
|
||||
requests: Vec<ProvisionRequest>,
|
||||
endpoint_lookups: Vec<(u64, String, String)>,
|
||||
|
|
@ -37,6 +38,7 @@ struct FakeLeaseClient {
|
|||
destroy_result: Option<Result<(), String>>,
|
||||
next_contract_id: u64,
|
||||
host_ids: VecDeque<Option<u64>>,
|
||||
first_wave_plan: Vec<Option<u64>>,
|
||||
}
|
||||
|
||||
impl FakeLeaseClient {
|
||||
|
|
@ -63,6 +65,15 @@ impl VastAiLeaseClient for FakeLeaseClient {
|
|||
})
|
||||
}
|
||||
|
||||
fn plan_first_wave_offers(
|
||||
&mut self,
|
||||
requests: &[ProvisionRequest],
|
||||
) -> Result<Vec<Option<u64>>, String> {
|
||||
let mut plan = self.first_wave_plan.clone();
|
||||
plan.resize(requests.len(), None);
|
||||
Ok(plan)
|
||||
}
|
||||
|
||||
fn ssh_endpoint(
|
||||
&mut self,
|
||||
contract_id: u64,
|
||||
|
|
@ -118,6 +129,209 @@ impl VastAiBootstrapLauncher for FakeBootstrap {
|
|||
}
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
struct ParallelLeaseClient {
|
||||
state: Arc<Mutex<ParallelLeaseState>>,
|
||||
gate: Arc<ParallelLeaseGate>,
|
||||
}
|
||||
|
||||
struct ParallelLeaseGate {
|
||||
target: usize,
|
||||
started: Mutex<usize>,
|
||||
all_started: Condvar,
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct ParallelLeaseState {
|
||||
requests: Vec<ProvisionRequest>,
|
||||
endpoint_lookups: Vec<u64>,
|
||||
destroyed: Vec<u64>,
|
||||
first_endpoint_request_count: Option<usize>,
|
||||
next_contract_id: u64,
|
||||
host_ids: VecDeque<Option<u64>>,
|
||||
first_wave_plan_requests: usize,
|
||||
first_wave_plan: Vec<Option<u64>>,
|
||||
}
|
||||
|
||||
impl ParallelLeaseClient {
|
||||
fn new(target: usize) -> Self {
|
||||
Self {
|
||||
state: Arc::new(Mutex::new(ParallelLeaseState {
|
||||
next_contract_id: 100,
|
||||
host_ids: (0..target)
|
||||
.map(|index| Some(10_000 + u64::try_from(index).unwrap()))
|
||||
.collect(),
|
||||
first_wave_plan: (0..target)
|
||||
.map(|index| Some(9_000 + u64::try_from(index).unwrap()))
|
||||
.collect(),
|
||||
..ParallelLeaseState::default()
|
||||
})),
|
||||
gate: Arc::new(ParallelLeaseGate {
|
||||
target,
|
||||
started: Mutex::new(0),
|
||||
all_started: Condvar::new(),
|
||||
}),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl VastAiLeaseClient for ParallelLeaseClient {
|
||||
fn plan_first_wave_offers(
|
||||
&mut self,
|
||||
requests: &[ProvisionRequest],
|
||||
) -> Result<Vec<Option<u64>>, String> {
|
||||
let mut state = self.state.lock();
|
||||
state.first_wave_plan_requests = requests.len();
|
||||
let mut plan = state.first_wave_plan.clone();
|
||||
plan.resize(requests.len(), None);
|
||||
Ok(plan)
|
||||
}
|
||||
|
||||
fn provision_one(&mut self, request: ProvisionRequest) -> Result<ProvisionedInstance, String> {
|
||||
{
|
||||
self.state.lock().requests.push(request);
|
||||
}
|
||||
let mut started = self.gate.started.lock();
|
||||
*started += 1;
|
||||
if *started < self.gate.target {
|
||||
let wait = self
|
||||
.gate
|
||||
.all_started
|
||||
.wait_for(&mut started, Duration::from_secs(2));
|
||||
assert!(
|
||||
!wait.timed_out(),
|
||||
"all concurrent Vast.ai lease requests should start before any waits for SSH"
|
||||
);
|
||||
} else {
|
||||
self.gate.all_started.notify_all();
|
||||
}
|
||||
drop(started);
|
||||
|
||||
let mut state = self.state.lock();
|
||||
let contract_id = state.next_contract_id;
|
||||
state.next_contract_id = state.next_contract_id.wrapping_add(1).max(1);
|
||||
let host_id = state.host_ids.pop_front().unwrap_or(Some(77));
|
||||
Ok(ProvisionedInstance {
|
||||
index: 0,
|
||||
contract_id,
|
||||
offer_id: 55,
|
||||
host_id,
|
||||
gpu_name: "RTX 4090".to_owned(),
|
||||
gpu_ram: Some(24_000.0),
|
||||
dph_total: 0.42,
|
||||
})
|
||||
}
|
||||
|
||||
fn ssh_endpoint(
|
||||
&mut self,
|
||||
contract_id: u64,
|
||||
_label: &str,
|
||||
_lifecycle: &LifecyclePolicy,
|
||||
ssh_user: &str,
|
||||
) -> Result<VastAiSshEndpoint, String> {
|
||||
let mut state = self.state.lock();
|
||||
let request_count = state.requests.len();
|
||||
state
|
||||
.first_endpoint_request_count
|
||||
.get_or_insert(request_count);
|
||||
state.endpoint_lookups.push(contract_id);
|
||||
Ok(VastAiSshEndpoint {
|
||||
host: "ssh5.vast.ai".to_owned(),
|
||||
port: 22017,
|
||||
user: ssh_user.to_owned(),
|
||||
})
|
||||
}
|
||||
|
||||
fn destroy_contract(&mut self, contract_id: u64) -> Result<(), String> {
|
||||
self.state.lock().destroyed.push(contract_id);
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
struct OutOfOrderLeaseClient {
|
||||
state: Arc<Mutex<OutOfOrderLeaseState>>,
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct OutOfOrderLeaseState {
|
||||
requests: Vec<u64>,
|
||||
endpoint_lookups: Vec<u64>,
|
||||
destroyed: Vec<u64>,
|
||||
slow_node_ids: Vec<u64>,
|
||||
endpoint_fail_node_ids: Vec<u64>,
|
||||
}
|
||||
|
||||
impl OutOfOrderLeaseClient {
|
||||
fn new(slow_node_ids: Vec<u64>, endpoint_fail_node_ids: Vec<u64>) -> Self {
|
||||
Self {
|
||||
state: Arc::new(Mutex::new(OutOfOrderLeaseState {
|
||||
slow_node_ids,
|
||||
endpoint_fail_node_ids,
|
||||
..OutOfOrderLeaseState::default()
|
||||
})),
|
||||
}
|
||||
}
|
||||
|
||||
fn node_id_from_label(label: Option<&str>) -> u64 {
|
||||
label
|
||||
.and_then(|label| label.rsplit('-').next())
|
||||
.and_then(|node| node.parse::<u64>().ok())
|
||||
.expect("test request labels include node id suffix")
|
||||
}
|
||||
}
|
||||
|
||||
impl VastAiLeaseClient for OutOfOrderLeaseClient {
|
||||
fn provision_one(&mut self, request: ProvisionRequest) -> Result<ProvisionedInstance, String> {
|
||||
let node_id = Self::node_id_from_label(request.label.as_deref());
|
||||
let should_sleep = {
|
||||
let mut state = self.state.lock();
|
||||
state.requests.push(node_id);
|
||||
state.slow_node_ids.contains(&node_id)
|
||||
};
|
||||
if should_sleep {
|
||||
std::thread::sleep(Duration::from_millis(150));
|
||||
}
|
||||
Ok(ProvisionedInstance {
|
||||
index: 0,
|
||||
contract_id: 1_000 + node_id,
|
||||
offer_id: 55 + node_id,
|
||||
host_id: Some(10_000 + node_id),
|
||||
gpu_name: "RTX 4090".to_owned(),
|
||||
gpu_ram: Some(24_000.0),
|
||||
dph_total: 0.42,
|
||||
})
|
||||
}
|
||||
|
||||
fn ssh_endpoint(
|
||||
&mut self,
|
||||
contract_id: u64,
|
||||
_label: &str,
|
||||
_lifecycle: &LifecyclePolicy,
|
||||
ssh_user: &str,
|
||||
) -> Result<VastAiSshEndpoint, String> {
|
||||
let node_id = contract_id - 1_000;
|
||||
let should_fail = {
|
||||
let mut state = self.state.lock();
|
||||
state.endpoint_lookups.push(node_id);
|
||||
state.endpoint_fail_node_ids.contains(&node_id)
|
||||
};
|
||||
if should_fail {
|
||||
return Err("connection refused".to_owned());
|
||||
}
|
||||
Ok(VastAiSshEndpoint {
|
||||
host: "ssh5.vast.ai".to_owned(),
|
||||
port: 22017,
|
||||
user: ssh_user.to_owned(),
|
||||
})
|
||||
}
|
||||
|
||||
fn destroy_contract(&mut self, contract_id: u64) -> Result<(), String> {
|
||||
self.state.lock().destroyed.push(contract_id);
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
fn spec() -> NodeProvisionSpec {
|
||||
NodeProvisionSpec {
|
||||
run_id: 9,
|
||||
|
|
@ -272,6 +486,171 @@ fn pipeline_starts_blacklist_hosts_already_leased_in_run() {
|
|||
plugin.stop_node(&first_handle).unwrap();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn failed_vastai_host_is_blacklisted_for_later_requests() {
|
||||
let mut client = FakeLeaseClient::default().with_contract(100);
|
||||
client.host_ids.extend([Some(77), Some(88)]);
|
||||
client
|
||||
.endpoint_results
|
||||
.push_back(Err("connection refused".to_owned()));
|
||||
let mut plugin = VastAiProvisioningPlugin::new(client, FakeBootstrap::default(), config());
|
||||
let first_error = plugin.start_node(spec(), sink()).unwrap_err();
|
||||
assert!(first_error.contains("connection refused"));
|
||||
|
||||
let mut second = spec();
|
||||
second.node_id = 12;
|
||||
second.stage_index = Some(3);
|
||||
let second_handle = plugin.start_node(second, sink()).unwrap();
|
||||
|
||||
assert!(
|
||||
plugin.client().requests[1]
|
||||
.selection
|
||||
.blacklist_hosts
|
||||
.contains(&77),
|
||||
"host that failed before runtime-ready must be excluded from later Vast.ai requests"
|
||||
);
|
||||
plugin.stop_node(&second_handle).unwrap();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn vastai_start_nodes_starts_lease_requests_concurrently() {
|
||||
let client = ParallelLeaseClient::new(4);
|
||||
let state = client.state.clone();
|
||||
let mut plugin = VastAiProvisioningPlugin::new(client, FakeBootstrap::default(), config());
|
||||
let specs = (0..4)
|
||||
.map(|index| {
|
||||
let mut spec = spec();
|
||||
spec.node_id = 11 + index;
|
||||
spec.stage_index = Some(u32::try_from(index).unwrap());
|
||||
spec
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
let results = plugin.start_nodes(specs, sink());
|
||||
|
||||
assert!(results.iter().all(|(_, result)| result.is_ok()));
|
||||
assert_eq!(plugin.active_contract_count(), 4);
|
||||
assert_eq!(plugin.bootstrap().starts.len(), 4);
|
||||
let state = state.lock();
|
||||
assert_eq!(state.requests.len(), 4);
|
||||
assert_eq!(state.endpoint_lookups.len(), 4);
|
||||
assert_eq!(
|
||||
state.first_endpoint_request_count,
|
||||
Some(4),
|
||||
"SSH lookup must not begin before every lease request has started"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn vastai_start_nodes_assigns_shared_first_wave_offer_plan() {
|
||||
let client = ParallelLeaseClient::new(3);
|
||||
let state = client.state.clone();
|
||||
let mut plugin = VastAiProvisioningPlugin::new(client, FakeBootstrap::default(), config());
|
||||
let specs = (0..3)
|
||||
.map(|index| {
|
||||
let mut spec = spec();
|
||||
spec.node_id = 11 + index;
|
||||
spec.stage_index = Some(u32::try_from(index).unwrap());
|
||||
spec
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
let results = plugin.start_nodes(specs, sink());
|
||||
|
||||
assert!(results.iter().all(|(_, result)| result.is_ok()));
|
||||
let state = state.lock();
|
||||
assert_eq!(state.first_wave_plan_requests, 3);
|
||||
let mut assigned = state
|
||||
.requests
|
||||
.iter()
|
||||
.map(|request| {
|
||||
(
|
||||
request.label.clone().expect("request label"),
|
||||
request.preferred_offer_id,
|
||||
)
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
assigned.sort_by(|left, right| left.0.cmp(&right.0));
|
||||
assert_eq!(
|
||||
assigned
|
||||
.into_iter()
|
||||
.map(|(_, preferred)| preferred)
|
||||
.collect::<Vec<_>>(),
|
||||
vec![Some(9_000), Some(9_001), Some(9_002)],
|
||||
"per-node requests should carry the coordinated first-wave offer plan"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn vastai_start_nodes_bootstraps_fast_completion_before_earlier_slow_node() {
|
||||
let client = OutOfOrderLeaseClient::new(vec![11], Vec::new());
|
||||
let state = client.state.clone();
|
||||
let mut plugin = VastAiProvisioningPlugin::new(client, FakeBootstrap::default(), config());
|
||||
let mut slow = spec();
|
||||
slow.node_id = 11;
|
||||
slow.stage_index = Some(0);
|
||||
let mut fast = spec();
|
||||
fast.node_id = 12;
|
||||
fast.stage_index = Some(1);
|
||||
|
||||
let results = plugin.start_nodes(vec![slow, fast], sink());
|
||||
|
||||
assert!(results.iter().all(|(_, result)| result.is_ok()));
|
||||
assert_eq!(
|
||||
plugin
|
||||
.bootstrap()
|
||||
.starts
|
||||
.iter()
|
||||
.map(|(spec, _)| spec.node_id)
|
||||
.collect::<Vec<_>>(),
|
||||
vec![12, 11],
|
||||
"later fast completion should bootstrap before earlier slow completion"
|
||||
);
|
||||
assert_eq!(state.lock().endpoint_lookups.len(), 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn vastai_start_nodes_one_failure_does_not_block_completed_node_bootstrap() {
|
||||
let client = OutOfOrderLeaseClient::new(vec![11], vec![11]);
|
||||
let state = client.state.clone();
|
||||
let mut plugin = VastAiProvisioningPlugin::new(client, FakeBootstrap::default(), config());
|
||||
let mut slow_failure = spec();
|
||||
slow_failure.node_id = 11;
|
||||
slow_failure.stage_index = Some(0);
|
||||
let mut fast_success = spec();
|
||||
fast_success.node_id = 12;
|
||||
fast_success.stage_index = Some(1);
|
||||
|
||||
let results = plugin.start_nodes(vec![slow_failure, fast_success], sink());
|
||||
|
||||
assert!(
|
||||
results[0]
|
||||
.1
|
||||
.as_ref()
|
||||
.unwrap_err()
|
||||
.contains("connection refused")
|
||||
);
|
||||
assert!(
|
||||
results[0]
|
||||
.1
|
||||
.as_ref()
|
||||
.unwrap_err()
|
||||
.contains("class=connection_refused")
|
||||
);
|
||||
assert!(results[1].1.is_ok());
|
||||
assert_eq!(
|
||||
plugin
|
||||
.bootstrap()
|
||||
.starts
|
||||
.iter()
|
||||
.map(|(spec, _)| spec.node_id)
|
||||
.collect::<Vec<_>>(),
|
||||
vec![12],
|
||||
"successful completed node should bootstrap even though another node fails"
|
||||
);
|
||||
assert_eq!(state.lock().destroyed.as_slice(), &[1_011]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn stop_destroys_known_vastai_contract_exactly_once() {
|
||||
let mut plugin = VastAiProvisioningPlugin::new(
|
||||
|
|
@ -293,7 +672,7 @@ fn stop_destroys_known_vastai_contract_exactly_once() {
|
|||
}
|
||||
|
||||
#[test]
|
||||
fn vastai_complete_bootstrap_keeps_log_tail_until_node_stop() {
|
||||
fn vastai_complete_bootstrap_stops_optional_log_tail_before_node_stop() {
|
||||
let mut plugin = VastAiProvisioningPlugin::new(
|
||||
FakeLeaseClient::default().with_contract(100),
|
||||
FakeBootstrap::default(),
|
||||
|
|
@ -303,9 +682,9 @@ fn vastai_complete_bootstrap_keeps_log_tail_until_node_stop() {
|
|||
|
||||
plugin.complete_bootstrap(&handle).unwrap();
|
||||
|
||||
assert!(
|
||||
plugin.bootstrap().stops.is_empty(),
|
||||
"runtime-ready completion should keep the SSH log tail alive"
|
||||
assert_eq!(
|
||||
plugin.bootstrap().stops,
|
||||
vec![(1, BootstrapStopReason::RuntimeReady)]
|
||||
);
|
||||
assert_eq!(plugin.client().destroyed, Vec::<u64>::new());
|
||||
assert_eq!(plugin.active_contract_count(), 1);
|
||||
|
|
@ -315,7 +694,7 @@ fn vastai_complete_bootstrap_keeps_log_tail_until_node_stop() {
|
|||
assert_eq!(plugin.client().destroyed, vec![100]);
|
||||
assert_eq!(
|
||||
plugin.bootstrap().stops,
|
||||
vec![(1, BootstrapStopReason::NodeStop)]
|
||||
vec![(1, BootstrapStopReason::RuntimeReady)]
|
||||
);
|
||||
assert_eq!(plugin.active_contract_count(), 0);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -58,6 +58,7 @@ fn provision_wire() -> StageProvisionWire {
|
|||
model_id: "smollm2-135m-q4".to_owned(),
|
||||
gguf_source: GgufSource::LocalPath("/models/smollm.gguf".to_owned()),
|
||||
tokenizer: TokenizerSource::EmbeddedGguf,
|
||||
stage_shard_plan: None,
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -139,6 +140,7 @@ fn provision_wire_from_plan(plan: &run_plan::RunPlan, stage_index: u32) -> Stage
|
|||
model_id: provision.model.model_id,
|
||||
gguf_source: provision.gguf_source,
|
||||
tokenizer: provision.tokenizer,
|
||||
stage_shard_plan: None,
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -5,6 +5,7 @@ use std::process::{Child, Command, Stdio};
|
|||
use std::sync::{
|
||||
Arc,
|
||||
atomic::{AtomicBool, Ordering},
|
||||
mpsc,
|
||||
};
|
||||
use std::time::Duration;
|
||||
|
||||
|
|
@ -12,7 +13,9 @@ use datastream::DatastreamProducer;
|
|||
use serde::{Deserialize, Serialize};
|
||||
use swactor::actor::{ActorAddress, ActorInterface};
|
||||
use swactor::runtime::{Ctx, Runtime};
|
||||
use swactor_vastai::{LifecyclePolicy, ProvisionRequest, ProvisionedInstance, SelectionPolicy};
|
||||
use swactor_vastai::{
|
||||
LifecyclePolicy, ProvisionRequest, ProvisionedInstance, SelectionPolicy, classify_vastai_error,
|
||||
};
|
||||
|
||||
use crate::bootstrap_datastream::{BootstrapDatastreamBridge, node_stream_id};
|
||||
use crate::node_provisioning::{
|
||||
|
|
@ -59,6 +62,12 @@ pub struct VastAiSshEndpoint {
|
|||
|
||||
pub trait VastAiLeaseClient: Send {
|
||||
fn provision_one(&mut self, request: ProvisionRequest) -> Result<ProvisionedInstance, String>;
|
||||
fn plan_first_wave_offers(
|
||||
&mut self,
|
||||
requests: &[ProvisionRequest],
|
||||
) -> Result<Vec<Option<u64>>, String> {
|
||||
Ok(vec![None; requests.len()])
|
||||
}
|
||||
|
||||
fn ssh_endpoint(
|
||||
&mut self,
|
||||
|
|
@ -94,6 +103,12 @@ impl ToolsVastAiLeaseClient {
|
|||
}
|
||||
}
|
||||
|
||||
impl Clone for ToolsVastAiLeaseClient {
|
||||
fn clone(&self) -> Self {
|
||||
Self::new(self.client.clone()).expect("clone VastAI lease client runtime")
|
||||
}
|
||||
}
|
||||
|
||||
impl VastAiLeaseClient for ToolsVastAiLeaseClient {
|
||||
fn provision_one(&mut self, request: ProvisionRequest) -> Result<ProvisionedInstance, String> {
|
||||
let fleet = self.runtime.block_on(self.client.provision(request))?;
|
||||
|
|
@ -107,6 +122,31 @@ impl VastAiLeaseClient for ToolsVastAiLeaseClient {
|
|||
Ok(instances.remove(0))
|
||||
}
|
||||
|
||||
fn plan_first_wave_offers(
|
||||
&mut self,
|
||||
requests: &[ProvisionRequest],
|
||||
) -> Result<Vec<Option<u64>>, String> {
|
||||
let Some(first) = requests.first() else {
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
let pool = self.runtime.block_on(
|
||||
self.client
|
||||
.search_offers(&first.selection, requests.len() as u32),
|
||||
)?;
|
||||
let planned = swactor_vastai::plan_distinct_host_first_wave(
|
||||
&pool,
|
||||
requests.len() as u32,
|
||||
&first.selection.blacklist_hosts,
|
||||
&[],
|
||||
);
|
||||
let mut out = planned
|
||||
.into_iter()
|
||||
.map(|offer| Some(offer.id))
|
||||
.collect::<Vec<_>>();
|
||||
out.resize(requests.len(), None);
|
||||
Ok(out)
|
||||
}
|
||||
|
||||
fn ssh_endpoint(
|
||||
&mut self,
|
||||
contract_id: u64,
|
||||
|
|
@ -240,6 +280,7 @@ where
|
|||
disk_gb: spec.shape.disk_gb,
|
||||
env,
|
||||
per_instance_env: vec![BTreeMap::new()],
|
||||
preferred_offer_id: None,
|
||||
onstart: self.config.onstart.clone(),
|
||||
selection: self.selection_for(spec),
|
||||
lifecycle: self.config.lifecycle.clone(),
|
||||
|
|
@ -602,13 +643,27 @@ fn spawn_retrying_ssh_bootstrap(
|
|||
backoff.as_secs()
|
||||
),
|
||||
});
|
||||
std::thread::sleep(backoff);
|
||||
if !sleep_ssh_backoff(backoff, &stopping) {
|
||||
return;
|
||||
}
|
||||
backoff = next_ssh_backoff(backoff);
|
||||
attempt += 1;
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
fn sleep_ssh_backoff(backoff: Duration, stopping: &AtomicBool) -> bool {
|
||||
let deadline = std::time::Instant::now() + backoff;
|
||||
while std::time::Instant::now() < deadline {
|
||||
if stopping.load(Ordering::SeqCst) {
|
||||
return false;
|
||||
}
|
||||
let remaining = deadline.saturating_duration_since(std::time::Instant::now());
|
||||
std::thread::sleep(std::cmp::min(remaining, Duration::from_millis(50)));
|
||||
}
|
||||
!stopping.load(Ordering::SeqCst)
|
||||
}
|
||||
|
||||
fn spawn_ssh_bootstrap_attempt(
|
||||
spec: &NodeProvisionSpec,
|
||||
endpoint: &VastAiSshEndpoint,
|
||||
|
|
@ -677,6 +732,7 @@ where
|
|||
config: VastAiProvisioningConfig,
|
||||
bootstrap_producer: Option<DatastreamProducer>,
|
||||
leased_host_ids: BTreeSet<u64>,
|
||||
failed_host_ids: BTreeSet<u64>,
|
||||
next_handle_id: u64,
|
||||
nodes: BTreeMap<u64, VastAiNode<B::Handle>>,
|
||||
}
|
||||
|
|
@ -700,6 +756,7 @@ where
|
|||
bootstrap_producer: None,
|
||||
next_handle_id: 1,
|
||||
leased_host_ids: BTreeSet::new(),
|
||||
failed_host_ids: BTreeSet::new(),
|
||||
nodes: BTreeMap::new(),
|
||||
}
|
||||
}
|
||||
|
|
@ -751,7 +808,11 @@ 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 {
|
||||
for host_id in self
|
||||
.leased_host_ids
|
||||
.iter()
|
||||
.chain(self.failed_host_ids.iter())
|
||||
{
|
||||
if !selection.blacklist_hosts.contains(host_id) {
|
||||
selection.blacklist_hosts.push(*host_id);
|
||||
}
|
||||
|
|
@ -764,6 +825,7 @@ where
|
|||
disk_gb: self.config.disk_gb,
|
||||
env,
|
||||
per_instance_env: vec![BTreeMap::new()],
|
||||
preferred_offer_id: None,
|
||||
onstart: self.config.onstart.clone(),
|
||||
selection,
|
||||
lifecycle: self.config.lifecycle.clone(),
|
||||
|
|
@ -779,9 +841,25 @@ where
|
|||
}
|
||||
}
|
||||
|
||||
fn classified_start_error(reason: String) -> String {
|
||||
let class = classify_vastai_error(&reason).as_str();
|
||||
format!("{reason} [class={class}]")
|
||||
}
|
||||
|
||||
struct VastAiStartedLease {
|
||||
label: String,
|
||||
instance: ProvisionedInstance,
|
||||
endpoint: VastAiSshEndpoint,
|
||||
}
|
||||
|
||||
struct VastAiBatchStartError {
|
||||
reason: String,
|
||||
failed_host_id: Option<u64>,
|
||||
}
|
||||
|
||||
impl<C, B> ProvisionPlugin for VastAiProvisioningPlugin<C, B>
|
||||
where
|
||||
C: VastAiLeaseClient + 'static,
|
||||
C: VastAiLeaseClient + Clone + 'static,
|
||||
B: VastAiBootstrapLauncher + 'static,
|
||||
{
|
||||
fn start_node(
|
||||
|
|
@ -801,10 +879,9 @@ where
|
|||
});
|
||||
|
||||
let request = self.build_request(&spec, label.clone());
|
||||
let instance = self
|
||||
.client
|
||||
.provision_one(request)
|
||||
.map_err(|e| format!("vastai provision node {}: {e}", spec.node_id))?;
|
||||
let instance = self.client.provision_one(request).map_err(|e| {
|
||||
classified_start_error(format!("vastai provision node {}: {e}", spec.node_id))
|
||||
})?;
|
||||
sink.observe(PluginObservation::ProviderLine {
|
||||
run_id: spec.run_id,
|
||||
node_id: spec.node_id,
|
||||
|
|
@ -840,9 +917,15 @@ where
|
|||
) {
|
||||
Ok(endpoint) => endpoint,
|
||||
Err(error) => {
|
||||
if let Some(host_id) = instance.host_id {
|
||||
self.failed_host_ids.insert(host_id);
|
||||
}
|
||||
return Err(self.cleanup_contract_after_start_error(
|
||||
instance.contract_id,
|
||||
format!("vastai SSH endpoint node {}: {error}", spec.node_id),
|
||||
classified_start_error(format!(
|
||||
"vastai SSH endpoint node {}: {error}",
|
||||
spec.node_id
|
||||
)),
|
||||
));
|
||||
}
|
||||
};
|
||||
|
|
@ -869,9 +952,15 @@ where
|
|||
) {
|
||||
Ok(handle) => handle,
|
||||
Err(error) => {
|
||||
if let Some(host_id) = instance.host_id {
|
||||
self.failed_host_ids.insert(host_id);
|
||||
}
|
||||
return Err(self.cleanup_contract_after_start_error(
|
||||
instance.contract_id,
|
||||
format!("vastai bootstrap node {}: {error}", spec.node_id),
|
||||
classified_start_error(format!(
|
||||
"vastai bootstrap node {}: {error}",
|
||||
spec.node_id
|
||||
)),
|
||||
));
|
||||
}
|
||||
};
|
||||
|
|
@ -897,7 +986,256 @@ where
|
|||
Ok(handle)
|
||||
}
|
||||
|
||||
fn complete_bootstrap(&mut self, _handle: &PluginNodeHandle) -> Result<(), String> {
|
||||
fn start_nodes(
|
||||
&mut self,
|
||||
specs: Vec<NodeProvisionSpec>,
|
||||
sink: PluginSink,
|
||||
) -> Vec<(NodeProvisionSpec, Result<PluginNodeHandle, String>)> {
|
||||
if specs.len() <= 1 {
|
||||
return specs
|
||||
.into_iter()
|
||||
.map(|spec| {
|
||||
let result = self.start_node(spec.clone(), sink.clone());
|
||||
(spec, result)
|
||||
})
|
||||
.collect();
|
||||
}
|
||||
|
||||
let mut results = (0..specs.len()).map(|_| None).collect::<Vec<_>>();
|
||||
let mut start_inputs = Vec::new();
|
||||
for (index, spec) in specs.into_iter().enumerate() {
|
||||
if !spec.mounts.is_empty() {
|
||||
results[index] = Some((
|
||||
spec,
|
||||
Err("vastai provider does not support host file mounts".to_owned()),
|
||||
));
|
||||
continue;
|
||||
}
|
||||
let stream_id = node_stream_id(spec.run_id, spec.node_id);
|
||||
let label = self.label_for(&spec);
|
||||
sink.observe(PluginObservation::ProviderLine {
|
||||
run_id: spec.run_id,
|
||||
node_id: spec.node_id,
|
||||
line: format!("vastai provisioning label={label} stream={stream_id}"),
|
||||
});
|
||||
let request = self.build_request(&spec, label.clone());
|
||||
start_inputs.push((index, spec, label, request));
|
||||
}
|
||||
|
||||
let request_plan = start_inputs
|
||||
.iter()
|
||||
.map(|(_, _, _, request)| request.clone())
|
||||
.collect::<Vec<_>>();
|
||||
let offer_plan = match self.client.plan_first_wave_offers(&request_plan) {
|
||||
Ok(plan) => plan,
|
||||
Err(error) => {
|
||||
for (_, spec, _, _) in &start_inputs {
|
||||
sink.observe(PluginObservation::ProviderLine {
|
||||
run_id: spec.run_id,
|
||||
node_id: spec.node_id,
|
||||
line: format!(
|
||||
"vastai first-wave offer planning failed; falling back to per-node selection: {error}"
|
||||
),
|
||||
});
|
||||
}
|
||||
vec![None; start_inputs.len()]
|
||||
}
|
||||
};
|
||||
|
||||
let (completion_tx, completion_rx) = mpsc::channel();
|
||||
for (plan_index, (index, spec, label, mut request)) in start_inputs.into_iter().enumerate()
|
||||
{
|
||||
request.preferred_offer_id = offer_plan.get(plan_index).copied().flatten();
|
||||
if let Some(offer_id) = request.preferred_offer_id {
|
||||
sink.observe(PluginObservation::ProviderLine {
|
||||
run_id: spec.run_id,
|
||||
node_id: spec.node_id,
|
||||
line: serde_json::json!({
|
||||
"type": "VastAiFirstWaveOfferPlanned",
|
||||
"run_id": spec.run_id,
|
||||
"node_id": spec.node_id,
|
||||
"label": &label,
|
||||
"offer_id": offer_id,
|
||||
})
|
||||
.to_string(),
|
||||
});
|
||||
}
|
||||
let mut client = self.client.clone();
|
||||
let config = self.config.clone();
|
||||
let worker_tx = completion_tx.clone();
|
||||
std::thread::spawn(move || {
|
||||
let started = match client.provision_one(request) {
|
||||
Ok(instance) => {
|
||||
match client.ssh_endpoint(
|
||||
instance.contract_id,
|
||||
&label,
|
||||
&config.lifecycle,
|
||||
&config.ssh_user,
|
||||
) {
|
||||
Ok(endpoint) => Ok(VastAiStartedLease {
|
||||
label,
|
||||
instance,
|
||||
endpoint,
|
||||
}),
|
||||
Err(error) => {
|
||||
let failed_host_id = instance.host_id;
|
||||
let reason = match client.destroy_contract(instance.contract_id) {
|
||||
Ok(()) => classified_start_error(format!(
|
||||
"vastai SSH endpoint node {}: {error}",
|
||||
spec.node_id
|
||||
)),
|
||||
Err(cleanup) => classified_start_error(format!(
|
||||
"vastai SSH endpoint node {}: {error}; cleanup destroy {} failed: {cleanup}",
|
||||
spec.node_id, instance.contract_id
|
||||
)),
|
||||
};
|
||||
Err(VastAiBatchStartError {
|
||||
reason,
|
||||
failed_host_id,
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(error) => Err(VastAiBatchStartError {
|
||||
reason: classified_start_error(format!(
|
||||
"vastai provision node {}: {error}",
|
||||
spec.node_id
|
||||
)),
|
||||
failed_host_id: None,
|
||||
}),
|
||||
};
|
||||
let _ = worker_tx.send((index, spec, started));
|
||||
});
|
||||
}
|
||||
drop(completion_tx);
|
||||
|
||||
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",
|
||||
"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,
|
||||
})
|
||||
.to_string(),
|
||||
});
|
||||
sink.observe(PluginObservation::ProviderLine {
|
||||
run_id: spec.run_id,
|
||||
node_id: spec.node_id,
|
||||
line: serde_json::json!({
|
||||
"type": "VastAiSshEndpointReady",
|
||||
"run_id": spec.run_id,
|
||||
"node_id": spec.node_id,
|
||||
"contract_id": started.instance.contract_id,
|
||||
"host": &started.endpoint.host,
|
||||
"port": started.endpoint.port,
|
||||
"user": &started.endpoint.user,
|
||||
})
|
||||
.to_string(),
|
||||
});
|
||||
|
||||
let bootstrap = match self.bootstrap.start_bootstrap(
|
||||
spec.clone(),
|
||||
started.endpoint,
|
||||
sink.clone(),
|
||||
self.bootstrap_producer.clone(),
|
||||
) {
|
||||
Ok(handle) => handle,
|
||||
Err(error) => {
|
||||
if let Some(host_id) = started.instance.host_id {
|
||||
self.failed_host_ids.insert(host_id);
|
||||
}
|
||||
let node_id = spec.node_id;
|
||||
results[index] = Some((
|
||||
spec,
|
||||
Err(self.cleanup_contract_after_start_error(
|
||||
started.instance.contract_id,
|
||||
classified_start_error(format!(
|
||||
"vastai bootstrap node {node_id}: {error}"
|
||||
)),
|
||||
)),
|
||||
));
|
||||
continue;
|
||||
}
|
||||
};
|
||||
|
||||
let host_id = started.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,
|
||||
};
|
||||
self.next_handle_id = self.next_handle_id.wrapping_add(1).max(1);
|
||||
self.nodes.insert(
|
||||
handle.id,
|
||||
VastAiNode {
|
||||
contract_id: started.instance.contract_id,
|
||||
bootstrap: Some(bootstrap),
|
||||
host_id,
|
||||
},
|
||||
);
|
||||
results[index] = Some((spec, Ok(handle)));
|
||||
}
|
||||
Err(error) => {
|
||||
if let Some(host_id) = error.failed_host_id {
|
||||
self.failed_host_ids.insert(host_id);
|
||||
}
|
||||
results[index] = Some((spec, Err(error.reason)));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
results
|
||||
.into_iter()
|
||||
.enumerate()
|
||||
.map(|(index, result)| {
|
||||
result.unwrap_or_else(|| {
|
||||
(
|
||||
NodeProvisionSpec {
|
||||
run_id: 0,
|
||||
node_id: u64::try_from(index).unwrap_or(u64::MAX),
|
||||
stage_index: None,
|
||||
image: String::new(),
|
||||
env: Vec::new(),
|
||||
args: Vec::new(),
|
||||
mounts: Vec::new(),
|
||||
},
|
||||
Err("vastai provision worker panicked".to_owned()),
|
||||
)
|
||||
})
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn complete_bootstrap(&mut self, handle: &PluginNodeHandle) -> Result<(), String> {
|
||||
let Some(node) = self.nodes.get_mut(&handle.id) else {
|
||||
return Ok(());
|
||||
};
|
||||
if let Some(mut bootstrap) = node.bootstrap.take() {
|
||||
self.bootstrap
|
||||
.stop_bootstrap(&mut bootstrap, BootstrapStopReason::RuntimeReady);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
|
|
@ -1083,7 +1421,7 @@ mod tests {
|
|||
}
|
||||
|
||||
#[test]
|
||||
fn vastai_complete_bootstrap_keeps_log_tail_until_node_stop() {
|
||||
fn vastai_complete_bootstrap_stops_optional_log_tail_before_node_stop() {
|
||||
let destroyed_contracts = Arc::new(Mutex::new(Vec::new()));
|
||||
let stop_reasons = Arc::new(Mutex::new(Vec::new()));
|
||||
let sink = PluginSink::new(Arc::new(ObservationSink::default()));
|
||||
|
|
@ -1104,13 +1442,16 @@ mod tests {
|
|||
plugin
|
||||
.complete_bootstrap(&handle)
|
||||
.expect("runtime-ready bootstrap completion succeeds");
|
||||
assert!(
|
||||
stop_reasons.lock().is_empty(),
|
||||
"VastAI bootstrap SSH tail must remain alive for post-ready worker logs"
|
||||
assert_eq!(
|
||||
*stop_reasons.lock(),
|
||||
vec![BootstrapStopReason::RuntimeReady]
|
||||
);
|
||||
|
||||
plugin.stop_node(&handle).expect("VastAI node stops");
|
||||
assert_eq!(*stop_reasons.lock(), vec![BootstrapStopReason::NodeStop]);
|
||||
assert_eq!(
|
||||
*stop_reasons.lock(),
|
||||
vec![BootstrapStopReason::RuntimeReady]
|
||||
);
|
||||
assert_eq!(*destroyed_contracts.lock(), vec![42]);
|
||||
}
|
||||
|
||||
|
|
@ -1189,6 +1530,18 @@ mod tests {
|
|||
assert!(!args.iter().any(|arg| arg == "IdentitiesOnly=yes"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ssh_bootstrap_backoff_sleep_observes_stop_without_waiting_full_backoff() {
|
||||
let stopping = AtomicBool::new(true);
|
||||
let started = std::time::Instant::now();
|
||||
|
||||
assert!(!sleep_ssh_backoff(Duration::from_secs(5), &stopping));
|
||||
assert!(
|
||||
started.elapsed() < Duration::from_millis(250),
|
||||
"stopped bootstrap backoff should not wait for the full retry delay"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ssh_bootstrap_actor_stop_kills_child() {
|
||||
let runtime = Arc::new(swactor::runtime::Runtime::new(
|
||||
|
|
|
|||
|
|
@ -4,9 +4,9 @@ use std::time::Duration;
|
|||
|
||||
use crate::config::{ENV_ASSUME_YES, truthy_env};
|
||||
use crate::monitor::wait_for_running_with_policy;
|
||||
use crate::pricing::{CostModel, plan_picks};
|
||||
use crate::pricing::CostModel;
|
||||
use crate::provision::create_instance;
|
||||
use crate::search::select_offer_pool_with_policy;
|
||||
use crate::search::{plan_distinct_host_first_wave, select_offer_pool_with_policy};
|
||||
use crate::teardown::{destroy_instance_with_retry, rollback};
|
||||
use crate::types::{
|
||||
CreateInstanceRequest, InstanceInfo, Offer, ProvisionRequest, ProvisionedFleet,
|
||||
|
|
@ -15,7 +15,7 @@ use crate::types::{
|
|||
|
||||
/// Print the planned lease + hourly cost and, on TTY, require y/N confirmation.
|
||||
pub fn confirm_lease(pool: &[Offer], num_instances: u32, cost: &CostModel) -> Result<(), String> {
|
||||
let picks = plan_picks(pool, num_instances);
|
||||
let picks = plan_distinct_host_first_wave(pool, num_instances, &[], &[]);
|
||||
let total_dph: f64 = picks.iter().map(|o| o.dph_total).sum();
|
||||
let total_eff: f64 = picks.iter().map(|o| cost.effective_price(o)).sum();
|
||||
|
||||
|
|
@ -80,9 +80,23 @@ fn next_eligible_offer<'a>(
|
|||
pool: &'a [Offer],
|
||||
tried_offer_ids: &[u64],
|
||||
used_host_ids: &HashSet<u64>,
|
||||
failed_host_ids: &HashSet<u64>,
|
||||
preferred_offer_id: Option<u64>,
|
||||
) -> Option<&'a Offer> {
|
||||
if let Some(offer_id) = preferred_offer_id {
|
||||
if let Some(offer) = pool.iter().find(|o| o.id == offer_id) {
|
||||
if !tried_offer_ids.contains(&offer.id)
|
||||
&& offer.host_id.is_none_or(|h| !failed_host_ids.contains(&h))
|
||||
{
|
||||
return Some(offer);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pool.iter().find(|o| {
|
||||
!tried_offer_ids.contains(&o.id) && o.host_id.map_or(true, |h| !used_host_ids.contains(&h))
|
||||
!tried_offer_ids.contains(&o.id)
|
||||
&& o.host_id
|
||||
.is_none_or(|h| !used_host_ids.contains(&h) && !failed_host_ids.contains(&h))
|
||||
})
|
||||
}
|
||||
|
||||
|
|
@ -103,14 +117,22 @@ async fn provision_one(
|
|||
index: u32,
|
||||
tried_offer_ids: &mut Vec<u64>,
|
||||
used_host_ids: &mut HashSet<u64>,
|
||||
failed_host_ids: &mut HashSet<u64>,
|
||||
preferred_offer_id: Option<u64>,
|
||||
) -> Result<ProvisionedInstance, String> {
|
||||
let mut attempt = 1_u64;
|
||||
loop {
|
||||
let offer = match next_eligible_offer(pool, tried_offer_ids, used_host_ids) {
|
||||
let offer = match next_eligible_offer(
|
||||
pool,
|
||||
tried_offer_ids,
|
||||
used_host_ids,
|
||||
failed_host_ids,
|
||||
preferred_offer_id.filter(|offer_id| !tried_offer_ids.contains(offer_id)),
|
||||
) {
|
||||
Some(o) => o.clone(),
|
||||
None => {
|
||||
return Err(format!(
|
||||
"pool exhausted for index {index} (no untried offer on an unused host)"
|
||||
"pool exhausted for index {index} (no untried offer outside failed hosts)"
|
||||
));
|
||||
}
|
||||
};
|
||||
|
|
@ -201,10 +223,12 @@ pub async fn provision_fleet(
|
|||
confirm_lease(&pool, req.count, &CostModel::from_policy(&req.selection))?;
|
||||
}
|
||||
|
||||
let first_wave =
|
||||
plan_distinct_host_first_wave(&pool, req.count, &req.selection.blacklist_hosts, &[]);
|
||||
let mut tried_offer_ids = Vec::new();
|
||||
let mut created: Vec<ProvisionedInstance> = Vec::with_capacity(req.count as usize);
|
||||
let mut used_host_ids = HashSet::new();
|
||||
|
||||
let mut failed_host_ids = HashSet::new();
|
||||
for index in 0..req.count {
|
||||
match provision_one(
|
||||
client,
|
||||
|
|
@ -215,6 +239,10 @@ pub async fn provision_fleet(
|
|||
index,
|
||||
&mut tried_offer_ids,
|
||||
&mut used_host_ids,
|
||||
&mut failed_host_ids,
|
||||
req.preferred_offer_id
|
||||
.filter(|_| req.count == 1)
|
||||
.or_else(|| first_wave.get(index as usize).map(|offer| offer.id)),
|
||||
)
|
||||
.await
|
||||
{
|
||||
|
|
@ -237,6 +265,9 @@ pub async fn provision_fleet(
|
|||
{
|
||||
Ok(_) => break,
|
||||
Err(e) => {
|
||||
if let Some(host_id) = created[idx].host_id {
|
||||
failed_host_ids.insert(host_id);
|
||||
}
|
||||
eprintln!(
|
||||
"lease_chain: index {index} contract {cid} did not reach running: {e}"
|
||||
);
|
||||
|
|
@ -257,6 +288,8 @@ pub async fn provision_fleet(
|
|||
index,
|
||||
&mut tried_offer_ids,
|
||||
&mut used_host_ids,
|
||||
&mut failed_host_ids,
|
||||
None,
|
||||
)
|
||||
.await
|
||||
{
|
||||
|
|
@ -286,3 +319,136 @@ pub async fn provision_fleet(
|
|||
instances: created,
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::collections::BTreeMap;
|
||||
use std::time::Duration;
|
||||
|
||||
use serde_json::json;
|
||||
use wiremock::matchers::{method, path};
|
||||
use wiremock::{Mock, MockServer, ResponseTemplate};
|
||||
|
||||
use super::*;
|
||||
use crate::types::{LifecyclePolicy, SelectionPolicy};
|
||||
|
||||
fn offer(id: u64, host_id: u64) -> serde_json::Value {
|
||||
json!({
|
||||
"id": id,
|
||||
"gpu_name": "RTX 4090",
|
||||
"dph_total": id as f64 / 100.0,
|
||||
"gpu_ram": 24_000.0,
|
||||
"compute_cap": 890,
|
||||
"geolocation": "US",
|
||||
"internet_down_cost_per_tb": 0.0,
|
||||
"internet_up_cost_per_tb": 0.0,
|
||||
"host_id": host_id,
|
||||
"verification": "verified"
|
||||
})
|
||||
}
|
||||
|
||||
fn request(count: u32) -> ProvisionRequest {
|
||||
ProvisionRequest {
|
||||
count,
|
||||
image: "registry.example/mvp-worker:latest".to_owned(),
|
||||
label: Some("lease-test".to_owned()),
|
||||
disk_gb: 80,
|
||||
env: BTreeMap::new(),
|
||||
per_instance_env: Vec::new(),
|
||||
preferred_offer_id: None,
|
||||
onstart: None,
|
||||
selection: SelectionPolicy {
|
||||
drop_cheap_frac: 0.0,
|
||||
..SelectionPolicy::default()
|
||||
},
|
||||
lifecycle: LifecyclePolicy {
|
||||
lease_pace: Duration::ZERO,
|
||||
poll_interval: Duration::from_millis(1),
|
||||
state_timeout: Duration::from_millis(5),
|
||||
},
|
||||
confirm_lease: false,
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn replacement_excludes_failed_host_from_shared_offer_pool() {
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(method("GET"))
|
||||
.and(path("/api/v0/bundles/"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
|
||||
"offers": [
|
||||
offer(1, 10),
|
||||
offer(2, 20),
|
||||
offer(3, 10),
|
||||
offer(4, 30)
|
||||
]
|
||||
})))
|
||||
.mount(&server)
|
||||
.await;
|
||||
for (offer_id, contract_id) in [(1, 101), (2, 102), (4, 104)] {
|
||||
Mock::given(method("PUT"))
|
||||
.and(path(format!("/api/v0/asks/{offer_id}/")))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
|
||||
"new_contract": contract_id
|
||||
})))
|
||||
.mount(&server)
|
||||
.await;
|
||||
}
|
||||
Mock::given(method("GET"))
|
||||
.and(path("/api/v0/instances/101/"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
|
||||
"instances": {
|
||||
"actual_status": "loading",
|
||||
"intended_status": "running",
|
||||
"status_msg": "still pulling"
|
||||
}
|
||||
})))
|
||||
.mount(&server)
|
||||
.await;
|
||||
for contract_id in [102, 104] {
|
||||
Mock::given(method("GET"))
|
||||
.and(path(format!("/api/v0/instances/{contract_id}/")))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
|
||||
"instances": {
|
||||
"actual_status": "running",
|
||||
"intended_status": "running",
|
||||
"public_ipaddr": "127.0.0.1",
|
||||
"ssh_port": 22
|
||||
}
|
||||
})))
|
||||
.mount(&server)
|
||||
.await;
|
||||
}
|
||||
Mock::given(method("DELETE"))
|
||||
.and(path("/api/v0/instances/101/"))
|
||||
.respond_with(ResponseTemplate::new(200))
|
||||
.mount(&server)
|
||||
.await;
|
||||
|
||||
let fleet = provision_fleet(&reqwest::Client::new(), &server.uri(), "secret", request(2))
|
||||
.await
|
||||
.expect("replacement should use non-failed host");
|
||||
|
||||
assert_eq!(
|
||||
fleet
|
||||
.instances
|
||||
.iter()
|
||||
.map(|instance| (instance.index, instance.offer_id, instance.host_id))
|
||||
.collect::<Vec<_>>(),
|
||||
vec![(0, 4, Some(30)), (1, 2, Some(20))]
|
||||
);
|
||||
let requests = server.received_requests().await.expect("recorded requests");
|
||||
assert!(
|
||||
requests
|
||||
.iter()
|
||||
.any(|request| request.url.path() == "/api/v0/asks/4/"),
|
||||
"replacement should rent an offer from a non-failed host"
|
||||
);
|
||||
assert!(
|
||||
!requests
|
||||
.iter()
|
||||
.any(|request| request.url.path() == "/api/v0/asks/3/"),
|
||||
"replacement must skip untried offers on the failed host"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -23,12 +23,12 @@ pub use logs::{fetch_logs, request_logs};
|
|||
pub use monitor::{wait_for_running, wait_for_running_with_policy};
|
||||
pub use pricing::CostModel;
|
||||
pub use provision::create_instance;
|
||||
pub use search::{select_offer_pool, select_offer_pool_with_policy};
|
||||
pub use search::{plan_distinct_host_first_wave, select_offer_pool, select_offer_pool_with_policy};
|
||||
pub use teardown::{
|
||||
destroy_all_instances, destroy_instance, destroy_instance_with_retry, list_instances_by_label,
|
||||
};
|
||||
pub use types::{
|
||||
ContractRef, CreateInstanceRequest, FleetState, InstanceInfo, LabeledInstance, LifecyclePolicy,
|
||||
Offer, ProvisionRequest, ProvisionedFleet, ProvisionedInstance, RunningInstance,
|
||||
SelectionPolicy,
|
||||
SelectionPolicy, VastAiFailureClass, classify_vastai_error,
|
||||
};
|
||||
|
|
|
|||
|
|
@ -48,20 +48,3 @@ pub(crate) fn rank_survivors(offers: Vec<Offer>, cost: &CostModel, drop_frac: f6
|
|||
survivors.sort_by(&by_price);
|
||||
survivors
|
||||
}
|
||||
|
||||
pub(crate) fn plan_picks(pool: &[Offer], num_instances: u32) -> Vec<&Offer> {
|
||||
let mut picks = Vec::with_capacity(num_instances as usize);
|
||||
let mut used = std::collections::HashSet::new();
|
||||
for o in pool {
|
||||
if picks.len() == num_instances as usize {
|
||||
break;
|
||||
}
|
||||
if let Some(h) = o.host_id {
|
||||
if !used.insert(h) {
|
||||
continue;
|
||||
}
|
||||
}
|
||||
picks.push(o);
|
||||
}
|
||||
picks
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,6 +1,60 @@
|
|||
use crate::filters::reachable_offers;
|
||||
use crate::pricing::{CostModel, rank_survivors};
|
||||
use crate::types::{Offer, SearchResponse, SelectionPolicy};
|
||||
use std::collections::HashSet;
|
||||
|
||||
/// Choose the first batch of offers from one ranked pool, preferring different
|
||||
/// hosts whenever the filtered pool can satisfy that.
|
||||
pub fn plan_distinct_host_first_wave(
|
||||
pool: &[Offer],
|
||||
target_count: u32,
|
||||
blacklisted_hosts: &[u64],
|
||||
failed_hosts: &[u64],
|
||||
) -> Vec<Offer> {
|
||||
let target = target_count as usize;
|
||||
let blocked = blacklisted_hosts
|
||||
.iter()
|
||||
.chain(failed_hosts.iter())
|
||||
.copied()
|
||||
.collect::<HashSet<_>>();
|
||||
let mut selected = Vec::with_capacity(target);
|
||||
let mut selected_ids = HashSet::new();
|
||||
let mut selected_hosts = HashSet::new();
|
||||
|
||||
for offer in pool.iter().filter(|offer| {
|
||||
offer
|
||||
.host_id
|
||||
.is_none_or(|host_id| !blocked.contains(&host_id))
|
||||
}) {
|
||||
if selected.len() == target {
|
||||
break;
|
||||
}
|
||||
if let Some(host_id) = offer.host_id {
|
||||
if !selected_hosts.insert(host_id) {
|
||||
continue;
|
||||
}
|
||||
}
|
||||
selected_ids.insert(offer.id);
|
||||
selected.push(offer.clone());
|
||||
}
|
||||
|
||||
if selected.len() < target {
|
||||
for offer in pool.iter().filter(|offer| {
|
||||
offer
|
||||
.host_id
|
||||
.is_none_or(|host_id| !blocked.contains(&host_id))
|
||||
}) {
|
||||
if selected.len() == target {
|
||||
break;
|
||||
}
|
||||
if selected_ids.insert(offer.id) {
|
||||
selected.push(offer.clone());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
selected
|
||||
}
|
||||
|
||||
/// Historical env-backed offer search wrapper.
|
||||
pub async fn select_offer_pool(
|
||||
|
|
@ -103,3 +157,70 @@ pub async fn select_offer_pool_with_policy(
|
|||
);
|
||||
Ok(pool)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn offer(id: u64, host_id: Option<u64>) -> Offer {
|
||||
Offer {
|
||||
id,
|
||||
gpu_name: "RTX 4090".to_owned(),
|
||||
dph_total: id as f64 / 100.0,
|
||||
gpu_ram: Some(24_000.0),
|
||||
compute_cap: 890,
|
||||
geolocation: Some("US".to_owned()),
|
||||
inet_down_cost_per_tb: 0.0,
|
||||
inet_up_cost_per_tb: 0.0,
|
||||
host_id,
|
||||
verification: Some("verified".to_owned()),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn first_wave_plan_prefers_distinct_hosts_and_preserves_blacklists() {
|
||||
let pool = vec![
|
||||
offer(1, Some(10)),
|
||||
offer(2, Some(10)),
|
||||
offer(3, Some(20)),
|
||||
offer(4, Some(30)),
|
||||
offer(5, Some(40)),
|
||||
];
|
||||
|
||||
let plan = plan_distinct_host_first_wave(&pool, 3, &[30], &[]);
|
||||
|
||||
assert_eq!(
|
||||
plan.iter().map(|offer| offer.id).collect::<Vec<_>>(),
|
||||
vec![1, 3, 5]
|
||||
);
|
||||
assert_eq!(
|
||||
plan.iter()
|
||||
.filter_map(|offer| offer.host_id)
|
||||
.collect::<std::collections::HashSet<_>>()
|
||||
.len(),
|
||||
3
|
||||
);
|
||||
assert!(
|
||||
plan.iter().all(|offer| offer.host_id != Some(30)),
|
||||
"operator blacklist must remain authoritative"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn first_wave_plan_excludes_failed_hosts_from_replacements() {
|
||||
let pool = vec![
|
||||
offer(1, Some(10)),
|
||||
offer(2, Some(20)),
|
||||
offer(3, Some(30)),
|
||||
offer(4, Some(20)),
|
||||
];
|
||||
|
||||
let plan = plan_distinct_host_first_wave(&pool, 2, &[], &[20]);
|
||||
|
||||
assert_eq!(
|
||||
plan.iter().map(|offer| offer.id).collect::<Vec<_>>(),
|
||||
vec![1, 3]
|
||||
);
|
||||
assert!(plan.iter().all(|offer| offer.host_id != Some(20)));
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -133,6 +133,9 @@ pub struct ProvisionRequest {
|
|||
pub env: BTreeMap<String, String>,
|
||||
/// Per-index env overlays, merged after `env`.
|
||||
pub per_instance_env: Vec<BTreeMap<String, String>>,
|
||||
/// Preferred offer for single-instance requests after an app-level shared
|
||||
/// first-wave planner has already coordinated distinct hosts.
|
||||
pub preferred_offer_id: Option<u64>,
|
||||
pub onstart: Option<String>,
|
||||
pub selection: SelectionPolicy,
|
||||
pub lifecycle: LifecyclePolicy,
|
||||
|
|
@ -158,6 +161,99 @@ pub struct ProvisionedFleet {
|
|||
pub instances: Vec<ProvisionedInstance>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub enum VastAiFailureClass {
|
||||
ConnectionRefused,
|
||||
PublicKeyDenied,
|
||||
EndpointMissing,
|
||||
ProviderLoadingTimeout,
|
||||
VanishedOffer,
|
||||
Other,
|
||||
}
|
||||
|
||||
impl VastAiFailureClass {
|
||||
pub fn as_str(self) -> &'static str {
|
||||
match self {
|
||||
Self::ConnectionRefused => "connection_refused",
|
||||
Self::PublicKeyDenied => "publickey_denied",
|
||||
Self::EndpointMissing => "endpoint_missing",
|
||||
Self::ProviderLoadingTimeout => "provider_loading_timeout",
|
||||
Self::VanishedOffer => "vanished_offer",
|
||||
Self::Other => "other",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn classify_vastai_error(raw: &str) -> VastAiFailureClass {
|
||||
let lower = raw.to_ascii_lowercase();
|
||||
if lower.contains("connection refused") || lower.contains("os error 111") {
|
||||
return VastAiFailureClass::ConnectionRefused;
|
||||
}
|
||||
if lower.contains("permission denied (publickey")
|
||||
|| lower.contains("publickey denied")
|
||||
|| lower.contains("public key denied")
|
||||
|| lower.contains("no supported authentication methods")
|
||||
{
|
||||
return VastAiFailureClass::PublicKeyDenied;
|
||||
}
|
||||
if lower.contains("no ssh host")
|
||||
|| lower.contains("no ssh port")
|
||||
|| lower.contains("no ssh endpoint")
|
||||
|| lower.contains("endpoint missing")
|
||||
|| lower.contains("ssh missing")
|
||||
{
|
||||
return VastAiFailureClass::EndpointMissing;
|
||||
}
|
||||
if lower.contains("stuck in status loading")
|
||||
|| lower.contains("provider loading timeout")
|
||||
|| (lower.contains("status loading") && lower.contains("timeout"))
|
||||
{
|
||||
return VastAiFailureClass::ProviderLoadingTimeout;
|
||||
}
|
||||
if lower.contains("no_such_ask")
|
||||
|| lower.contains("no such ask")
|
||||
|| lower.contains("vanished offer")
|
||||
|| (lower.contains("404") && lower.contains("/asks/"))
|
||||
{
|
||||
return VastAiFailureClass::VanishedOffer;
|
||||
}
|
||||
VastAiFailureClass::Other
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod failure_class_tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn representative_vastai_errors_classify_to_stable_failure_classes() {
|
||||
for (raw, class) in [
|
||||
(
|
||||
"ssh: connect to host ssh5.vast.ai port 22017: Connection refused",
|
||||
VastAiFailureClass::ConnectionRefused,
|
||||
),
|
||||
(
|
||||
"Permission denied (publickey).",
|
||||
VastAiFailureClass::PublicKeyDenied,
|
||||
),
|
||||
(
|
||||
"vastai contract 123 has no SSH host",
|
||||
VastAiFailureClass::EndpointMissing,
|
||||
),
|
||||
(
|
||||
"instance 123 stuck in status loading for 300s",
|
||||
VastAiFailureClass::ProviderLoadingTimeout,
|
||||
),
|
||||
(
|
||||
"create_instance HTTP 400: {\"error\":\"no_such_ask\"}",
|
||||
VastAiFailureClass::VanishedOffer,
|
||||
),
|
||||
] {
|
||||
assert_eq!(classify_vastai_error(raw), class, "{raw}");
|
||||
assert_ne!(class.as_str(), "other");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Generic held-fleet handle file.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
|
||||
pub struct FleetState {
|
||||
|
|
|
|||
|
|
@ -755,6 +755,8 @@ fn run_mvp_chat_check(args: Vec<String>) -> ExitCode {
|
|||
};
|
||||
let run_id = mvp_chat_check_run_id();
|
||||
println!("mvp-chat-check: scenario {}", invocation.name());
|
||||
println!("mvp-chat-check: artifacts {}", paths.root.display());
|
||||
println!("mvp-chat-check: datastream {}", paths.dump_log.display());
|
||||
|
||||
let output = match run_mvp_chat_check_process(&workspace, &paths, run_id, &invocation) {
|
||||
Ok(output) => output,
|
||||
|
|
|
|||
Loading…
Reference in a new issue