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

refactor(mvp-system): drop node-image and stage-controller actors

Browse source 
- Remove docker build-context hashing, MvpLifecycleRecord, LocalShimRelayProvider, and
  StageControllerActor.
- Replace gguf scalar/array skip readers with shared helpers; simplify vastai adapter
  (~-360 lines).
- Add rpc-ready wait, binary-progress ensure, and host-gpu sampler helpers.


Signed-off-by: Zachery Aaron Shores-Chmielewski <zacheryasc@gmail.com>
This commit is contained in:
Zachery Aaron Shores-Chmielewski 2026-07-31 12:04:09 +04:00
parent 13728e3732
commit ee848d9cf9
20 changed files with 519 additions and 1038 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

90
crates/mvp-system/src/chat/node_image.rs
View file

@ -103,8 +103,7 @@ fn prepare_node_image_inner(
));
}
run_status(
progress,
run_status_command(
&root,
"cargo",
&[
@ -117,6 +116,7 @@ fn prepare_node_image_inner(
],
"build mvp-worker-node",
None,
progress,
)?;
let base_hash = content_hash_for_inputs(&root, BASE_IMAGE_SOURCE_INPUTS)?;
@ -124,12 +124,13 @@ fn prepare_node_image_inner(
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 worker_hash = hash_relative_files(
let worker_hash = hash_relative_files_with_salts(
&root,
vec![relative_path(
&root,
&root.join("apps/mvp-node/tinygrad_worker.py"),
)?],
&[],
)?;
let expected_node_labels = vec![
(NODE_IMAGE_TAG_LABEL, tag.as_str()),
@ -168,18 +169,19 @@ fn prepare_node_image_inner(
let base_image_matches =
docker_image_labels_match(&root, &request.base_image, &expected_base_labels)?;
if !base_image_matches {
let base_source_hash_label = format!("{BASE_IMAGE_SOURCE_HASH_LABEL}={base_hash}");
run_status_command(
&root,
"docker",
&vec![
"build".to_owned(),
"-f".to_owned(),
"apps/mvp-node/Dockerfile.base".to_owned(),
"--label".to_owned(),
format!("{BASE_IMAGE_SOURCE_HASH_LABEL}={base_hash}"),
"-t".to_owned(),
request.base_image.clone(),
".".to_owned(),
&[
"build",
"-f",
"apps/mvp-node/Dockerfile.base",
"--label",
base_source_hash_label.as_str(),
"-t",
request.base_image.as_str(),
".",
],
"build mvp node base image",
Some(&request.base_image),
@ -187,7 +189,14 @@ fn prepare_node_image_inner(
)?;
}
let node_bin = docker_build_context_path(&root, &request.node_bin)?;
let node_bin = {
let full = if request.node_bin.is_absolute() {
request.node_bin.to_path_buf()
} else {
root.join(&request.node_bin)
};
relative_path(&root, &full).map(|relative| relative.to_string_lossy().to_string())
}?;
let mut build_args = vec![
"build".to_owned(),
"-f".to_owned(),
@ -205,7 +214,7 @@ fn prepare_node_image_inner(
run_status_command(
&root,
"docker",
&build_args,
&build_args.iter().map(String::as_str).collect::<Vec<_>>(),
"build mvp node image",
Some(&image_ref),
progress,
@ -302,10 +311,6 @@ fn content_hash_for_inputs(root: &Path, inputs: &[&str]) -> Result<String, Strin
}
files.sort();
files.dedup();
hash_relative_files(root, files)
}
fn hash_relative_files(root: &Path, files: Vec<PathBuf>) -> Result<String, String> {
hash_relative_files_with_salts(root, files, &[])
}
@ -360,7 +365,7 @@ fn collect_hash_inputs(root: &Path, path: &Path, out: &mut Vec<PathBuf>) -> Resu
}
return Ok(());
}
if !metadata.is_dir() || skip_dir(path) {
if !metadata.is_dir() || matches!(path.file_name().and_then(|name| name.to_str()), Some(".git" | "target" | "__pycache__")) {
return Ok(());
}
let entries = fs::read_dir(path).map_err(|e| format!("read dir {display}: {e}"))?;
@ -381,15 +386,6 @@ fn relative_path(root: &Path, path: &Path) -> Result<PathBuf, String> {
})
}
fn docker_build_context_path(root: &Path, path: &Path) -> Result<String, String> {
let full = if path.is_absolute() {
path.to_path_buf()
} else {
root.join(path)
};
relative_path(root, &full).map(|relative| relative.to_string_lossy().to_string())
}
fn display_workspace_path(root: &Path, path: &Path) -> String {
match path.strip_prefix(root) {
Ok(relative) if relative.as_os_str().is_empty() => ".".to_owned(),
@ -398,13 +394,6 @@ fn display_workspace_path(root: &Path, path: &Path) -> String {
}
}
fn skip_dir(path: &Path) -> bool {
matches!(
path.file_name().and_then(|name| name.to_str()),
Some(".git" | "target" | "__pycache__")
)
}
fn alias_tags(
image: &ImageName,
extra_tag: Option<&str>,
@ -444,13 +433,13 @@ fn ensure_aliases_local(
) -> Result<(), String> {
for alias in alias_refs(image, alias_tags) {
if alias != source_ref {
run_status(
progress,
run_status_command(
root,
"docker",
&["tag", source_ref, &alias],
"tag mvp node image",
Some(&alias),
progress,
)?;
}
}
@ -468,13 +457,13 @@ fn ensure_aliases_for_remote(
return Ok(false);
}
if !docker_image_exists(root, source_ref) {
run_status(
progress,
run_status_command(
root,
"docker",
&["pull", source_ref],
"pull mvp node image",
Some(source_ref),
progress,
)?;
}
ensure_aliases_local(progress, root, source_ref, image, alias_tags)?;
@ -496,13 +485,13 @@ fn push_image(
root: &Path,
image_ref: &str,
) -> Result<(), String> {
run_status(
progress,
run_status_command(
root,
"docker",
&["push", image_ref],
"push mvp node image",
Some(image_ref),
progress,
)
}
@ -581,18 +570,6 @@ fn prune_old_dirty_images(root: &Path, image: &ImageName, keep_tag: &str) {
}
}
fn run_status(
progress: &mut Option<&mut dyn NodeImageProgressSink>,
root: &Path,
program: &str,
args: &[&str],
label: &str,
image_ref: Option<&str>,
) -> Result<(), String> {
let args = args.iter().map(|arg| (*arg).to_owned()).collect::<Vec<_>>();
run_status_command(root, program, &args, label, image_ref, progress)
}
fn emit_image_reference(
progress: &mut Option<&mut dyn NodeImageProgressSink>,
role: &str,
@ -686,16 +663,17 @@ fn drain_command_lines(
fn run_status_command(
root: &Path,
program: &str,
args: &[String],
args: &[&str],
label: &str,
image_ref: Option<&str>,
progress: &mut Option<&mut dyn NodeImageProgressSink>,
) -> Result<(), String> {
let args: Vec<String> = args.iter().map(|arg| (*arg).to_owned()).collect();
eprintln!("mvp-node-image: {label}");
if progress.is_none() {
let status = Command::new(program)
.current_dir(root)
.args(args)
.args(&args)
.stdin(Stdio::null())
.stdout(Stdio::inherit())
.stderr(Stdio::inherit())
@ -721,7 +699,7 @@ fn run_status_command(
);
let mut child = match Command::new(program)
.current_dir(root)
.args(args)
.args(&args)
.stdin(Stdio::null())
.stdout(Stdio::piped())
.stderr(Stdio::piped())

295
crates/mvp-system/src/chat/runtime.rs
View file

@ -1351,6 +1351,33 @@ struct InProcessOrch {
cleaned: bool,
}
fn wait_for_rpc_ready(
rpc_addr: &str,
mut check_dead: impl FnMut() -> Result<(), String>,
) -> Result<String, String> {
loop {
if STOP_REQUESTED.load(Ordering::SeqCst) {
return Err("interrupted before orchestrator became ready".to_owned());
}
match TcpStream::connect(rpc_addr) {
Ok(stream) => {
let _ = stream.shutdown(Shutdown::Both);
return Ok(rpc_addr.to_owned());
}
Err(error)
if matches!(
error.kind(),
std::io::ErrorKind::ConnectionRefused
| std::io::ErrorKind::TimedOut
| std::io::ErrorKind::AddrNotAvailable
) => {}
Err(error) => return Err(format!("connect prompt RPC {rpc_addr}: {error}")),
}
check_dead()?;
thread::sleep(Duration::from_millis(100));
}
}
impl InProcessOrch {
fn spawn(config: &Config, image_ref: &str) -> Result<Self, String> {
let args = config.orchestrator_cli_args(image_ref);
@ -1366,24 +1393,7 @@ impl InProcessOrch {
}
fn wait_ready(&mut self, rpc_addr: String) -> Result<String, String> {
loop {
if STOP_REQUESTED.load(Ordering::SeqCst) {
return Err("interrupted before orchestrator became ready".to_owned());
}
match TcpStream::connect(&rpc_addr) {
Ok(stream) => {
let _ = stream.shutdown(Shutdown::Both);
return Ok(rpc_addr);
}
Err(error)
if matches!(
error.kind(),
std::io::ErrorKind::ConnectionRefused
| std::io::ErrorKind::TimedOut
| std::io::ErrorKind::AddrNotAvailable
) => {}
Err(error) => return Err(format!("connect prompt RPC {rpc_addr}: {error}")),
}
wait_for_rpc_ready(&rpc_addr, || {
if let Some(result) = self.take_finished_result() {
let reason = match result {
Ok(()) => "completed successfully".to_owned(),
@ -1393,8 +1403,8 @@ impl InProcessOrch {
"in-process orchestrator exited before prompt RPC ready: {reason}"
));
}
thread::sleep(Duration::from_millis(100));
}
Ok(())
})
}
fn shutdown(&mut self) {
@ -1478,24 +1488,7 @@ impl OrchChild {
}
fn wait_ready(&mut self, rpc_addr: String) -> Result<String, String> {
loop {
if STOP_REQUESTED.load(Ordering::SeqCst) {
return Err("interrupted before orchestrator became ready".to_owned());
}
match TcpStream::connect(&rpc_addr) {
Ok(stream) => {
let _ = stream.shutdown(Shutdown::Both);
return Ok(rpc_addr);
}
Err(error)
if matches!(
error.kind(),
std::io::ErrorKind::ConnectionRefused
| std::io::ErrorKind::TimedOut
| std::io::ErrorKind::AddrNotAvailable
) => {}
Err(error) => return Err(format!("connect prompt RPC {rpc_addr}: {error}")),
}
wait_for_rpc_ready(&rpc_addr, || {
if let Some(status) = self
.child
.try_wait()
@ -1505,8 +1498,8 @@ impl OrchChild {
"orchestrator exited before prompt RPC ready: {status}"
));
}
thread::sleep(Duration::from_millis(100));
}
Ok(())
})
}
// The orchestrator shutdown spec is still pending. Replace this with the approved
@ -1566,6 +1559,60 @@ fn signal_orch_process_group(child: &Child, signal: libc::c_int) -> io::Result<(
}
}
fn ensure_binary_with_progress(
progress: &mut Option<&mut ChatDatastream>,
phase: &str,
mode: &str,
verbose: bool,
skip_rebuild: bool,
bin: &Path,
label: &str,
cargo_args: &[&str],
) -> Result<(), String> {
let started = Instant::now();
emit_chat_progress(
progress,
CHAT_RUNTIME_CHANNEL,
phase,
"started",
if verbose {
json!({"mode": mode, "command_label": phase})
} else {
json!({"mode": mode})
},
);
match ensure_runtime_binary(skip_rebuild, bin, label, cargo_args) {
Ok(()) => {
emit_chat_progress(
progress,
CHAT_RUNTIME_CHANNEL,
phase,
"ready",
if verbose {
json!({"mode": mode, "command_label": phase, "elapsed_ms": started.elapsed().as_millis()})
} else {
json!({"mode": mode})
},
);
Ok(())
}
Err(error) => {
emit_chat_progress(
progress,
CHAT_RUNTIME_CHANNEL,
phase,
"failed",
if verbose {
json!({"mode": mode, "command_label": phase, "elapsed_ms": started.elapsed().as_millis(), "error": error.as_str()})
} else {
json!({"mode": mode, "error": error.as_str()})
},
);
Err(error)
}
}
}
fn prepare_runtime_with_progress<F>(
config: &Config,
mut prepare_node_image_fn: F,
@ -1596,15 +1643,11 @@ where
json!({"mode": config.orchestrator_launch_mode()}),
);
} else {
let ensure_orch_started = Instant::now();
emit_chat_progress(
ensure_binary_with_progress(
&mut progress,
CHAT_RUNTIME_CHANNEL,
"ensure_orch_binary",
"started",
json!({"mode": binary_mode, "command_label": "ensure_orch_binary"}),
);
match ensure_runtime_binary(
binary_mode,
true,
config.skip_rebuild,
&config.orch_bin,
"mvp-orchestrator",
@ -1618,36 +1661,15 @@ where
"--bin",
"mvp-orchestrator",
],
) {
Ok(()) => emit_chat_progress(
&mut progress,
CHAT_RUNTIME_CHANNEL,
"ensure_orch_binary",
"ready",
json!({"mode": binary_mode, "command_label": "ensure_orch_binary", "elapsed_ms": ensure_orch_started.elapsed().as_millis()}),
),
Err(error) => {
emit_chat_progress(
&mut progress,
CHAT_RUNTIME_CHANNEL,
"ensure_orch_binary",
"failed",
json!({"mode": binary_mode, "command_label": "ensure_orch_binary", "elapsed_ms": ensure_orch_started.elapsed().as_millis(), "error": error.as_str()}),
);
return Err(error);
}
}
)?;
}
if config.provider == provider_kind::process() {
emit_chat_progress(
ensure_binary_with_progress(
&mut progress,
CHAT_RUNTIME_CHANNEL,
"ensure_worker_binary",
"started",
json!({"mode": binary_mode}),
);
match ensure_runtime_binary(
binary_mode,
false,
config.skip_rebuild,
&config.worker_bin,
"mvp-worker-node",
@ -1659,25 +1681,7 @@ where
"--bin",
"mvp-worker-node",
],
) {
Ok(()) => emit_chat_progress(
&mut progress,
CHAT_RUNTIME_CHANNEL,
"ensure_worker_binary",
"ready",
json!({"mode": binary_mode}),
),
Err(error) => {
emit_chat_progress(
&mut progress,
CHAT_RUNTIME_CHANNEL,
"ensure_worker_binary",
"failed",
json!({"mode": binary_mode, "error": error.as_str()}),
);
return Err(error);
}
}
)?;
emit_chat_progress(
&mut progress,
CHAT_RUNTIME_CHANNEL,
@ -1706,15 +1710,11 @@ where
);
return Ok(config.node_image.clone());
}
let ensure_worker_started = Instant::now();
emit_chat_progress(
ensure_binary_with_progress(
&mut progress,
CHAT_RUNTIME_CHANNEL,
"ensure_worker_binary",
"started",
json!({"mode": binary_mode, "command_label": "ensure_worker_binary"}),
);
match ensure_runtime_binary(
binary_mode,
true,
config.skip_rebuild,
&config.worker_bin,
"mvp-worker-node",
@ -1726,25 +1726,7 @@ where
"--bin",
"mvp-worker-node",
],
) {
Ok(()) => emit_chat_progress(
&mut progress,
CHAT_RUNTIME_CHANNEL,
"ensure_worker_binary",
"ready",
json!({"mode": binary_mode, "command_label": "ensure_worker_binary", "elapsed_ms": ensure_worker_started.elapsed().as_millis()}),
),
Err(error) => {
emit_chat_progress(
&mut progress,
CHAT_RUNTIME_CHANNEL,
"ensure_worker_binary",
"failed",
json!({"mode": binary_mode, "command_label": "ensure_worker_binary", "elapsed_ms": ensure_worker_started.elapsed().as_millis(), "error": error.as_str()}),
);
return Err(error);
}
}
)?;
emit_chat_progress(
&mut progress,
CHAT_RUNTIME_CHANNEL,
@ -1764,25 +1746,7 @@ where
json!({"provider": config.provider.as_str(), "command_label": "prepare_node_image", "image_tag": config.image_tag.as_deref()}),
);
let node_bin = default_worker_bin();
let requires_registry_image = if config.provider == provider_kind::docker() {
false
} else if config.provider == provider_kind::vastai() {
true
} else {
let error = if config.provider == provider_kind::process() {
"process provider does not use node images"
} else {
"mvp-chat does not support mock provider"
};
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}),
);
return Err(error.to_owned());
};
let requires_registry_image = config.provider == provider_kind::vastai();
let prepared = {
let command_progress = progress
.as_deref_mut()
@ -1905,10 +1869,6 @@ fn emit_chat_progress(
}
}
fn prompt_hash_hex(prompt: &str) -> String {
blake3::hash(prompt.as_bytes()).to_hex().to_string()
}
fn run_chat_session_with_output_and_progress(
writer: &mut impl Write,
mut reader: impl BufRead,
@ -1920,16 +1880,11 @@ fn run_chat_session_with_output_and_progress(
let mut progress = progress;
let mut next_request_id = 1_u64;
let mut next_prompt_index = 1_u64;
let prompt_exited = |progress: &mut Option<&mut ChatDatastream>, reason: &str| emit_chat_progress(progress, CHAT_PROMPT_CHANNEL, "prompt_loop", "exited", json!({"reason": reason}));
loop {
if STOP_REQUESTED.load(Ordering::SeqCst) {
emit_chat_progress(
&mut progress,
CHAT_PROMPT_CHANNEL,
"prompt_loop",
"exited",
json!({"reason": "stop_requested"}),
);
prompt_exited(&mut progress, "stop_requested");
return Ok(());
}
emit_chat_progress(
@ -1944,23 +1899,11 @@ fn run_chat_session_with_output_and_progress(
let prompt = match input_rx.recv() {
Ok(PromptInput::Line(line)) => line.trim_end().to_owned(),
Ok(PromptInput::Closed) | Err(_) => {
emit_chat_progress(
&mut progress,
CHAT_PROMPT_CHANNEL,
"prompt_loop",
"exited",
json!({"reason": "input_closed"}),
);
prompt_exited(&mut progress, "input_closed");
return Ok(());
}
Ok(PromptInput::StopRequested) => {
emit_chat_progress(
&mut progress,
CHAT_PROMPT_CHANNEL,
"prompt_loop",
"exited",
json!({"reason": "stop_requested"}),
);
prompt_exited(&mut progress, "stop_requested");
return Ok(());
}
};
@ -1972,7 +1915,7 @@ fn run_chat_session_with_output_and_progress(
next_request_id = next_request_id.wrapping_add(1).max(1);
let prompt_index = next_prompt_index;
next_prompt_index = next_prompt_index.wrapping_add(1).max(1);
let prompt_hash = prompt_hash_hex(&prompt);
let prompt_hash = blake3::hash(prompt.as_bytes()).to_hex().to_string();
emit_chat_progress(
&mut progress,
CHAT_PROMPT_CHANNEL,
@ -2000,13 +1943,7 @@ fn run_chat_session_with_output_and_progress(
loop {
if STOP_REQUESTED.load(Ordering::SeqCst) {
emit_chat_progress(
&mut progress,
CHAT_PROMPT_CHANNEL,
"prompt_loop",
"exited",
json!({"reason": "stop_requested"}),
);
prompt_exited(&mut progress, "stop_requested");
return Ok(());
}
let mut line = String::new();
@ -2269,16 +2206,6 @@ fn provider_from_sources(
Ok(provider_kind::process())
}
fn env_flag(name: &str, default: bool) -> bool {
match env_optional(name) {
Some(value) => !matches!(
value.to_ascii_lowercase().as_str(),
"0" | "false" | "no" | "off"
),
None => default,
}
}
fn env_optional(name: &str) -> Option<String> {
std::env::var(name)
.ok()
@ -2302,6 +2229,16 @@ where
.map_err(|e| format!("invalid {name}={value:?}: {e}"))
}
fn env_flag(name: &str, default: bool) -> bool {
match env_optional(name) {
Some(value) => !matches!(
value.to_ascii_lowercase().as_str(),
"0" | "false" | "no" | "off"
),
None => default,
}
}
fn parse_pipeline_stages_value(
args: &mut impl Iterator<Item = String>,
name: &str,

236
crates/mvp-system/src/node/worker_node_runtime.rs
View file

@ -557,26 +557,75 @@ fn submit_sampler_sample_health(
seq: u64,
error: Option<&str>,
) {
match error {
Some(error) => submit_sampler_health(
producer,
health_channel,
context,
sampler,
sample_channel,
let (status, detail) = match error {
Some(error) => (
"failed",
json!({"state":"error","sample_seq":seq,"error":error}),
),
None => submit_sampler_health(
producer,
health_channel,
context,
sampler,
sample_channel,
None => (
"ready",
json!({"state":"sample_observed","sample_seq":seq}),
),
}
};
submit_sampler_health(
producer,
health_channel,
context,
sampler,
sample_channel,
status,
detail,
);
}
fn spawn_blocking_sampler<S: Record + Send + 'static>(
handle: tokio::runtime::Handle,
producer: DatastreamProducer,
channel: ChannelId,
health_channel: ChannelId,
health_context: SamplerHealthContext,
sampler: &'static str,
sample_channel: &'static str,
interval: Duration,
error_label: &'static str,
sample_fn: fn(u64) -> S,
error_fn: fn(u64, String) -> S,
error_of: fn(&S) -> Option<&str>,
) {
handle.spawn(async move {
submit_sampler_started(
&producer,
health_channel,
health_context,
sampler,
sample_channel,
interval,
);
let mut seq = 0_u64;
let mut interval = tokio::time::interval(interval);
loop {
interval.tick().await;
let sample_seq = seq;
let sample = match tokio::task::spawn_blocking(move || sample_fn(sample_seq)).await {
Ok(sample) => sample,
Err(error) => error_fn(sample_seq, format!("{error_label}: {error}")),
};
submit_sampler_sample_health(
&producer,
health_channel,
health_context,
sampler,
sample_channel,
sample_seq,
error_of(&sample),
);
seq = seq.saturating_add(1);
producer.submit_record(channel, &sample);
}
});
}
fn spawn_host_gpu_sampler(
@ -586,48 +635,20 @@ fn spawn_host_gpu_sampler(
health_channel: ChannelId,
health_context: SamplerHealthContext,
) {
handle.spawn(async move {
let sample_channel = datastream::hardware::gpu::HOST_GPU_CHANNEL;
submit_sampler_started(
&producer,
health_channel,
health_context,
"gpu",
sample_channel,
datastream::hardware::gpu::GPU_SAMPLE_INTERVAL,
);
let mut seq = 0_u64;
let mut interval = tokio::time::interval(datastream::hardware::gpu::GPU_SAMPLE_INTERVAL);
loop {
interval.tick().await;
let sample_seq = seq;
let sample = match tokio::task::spawn_blocking(move || {
datastream::hardware::gpu::sample(sample_seq)
})
.await
{
Ok(sample) => sample,
Err(error) => datastream::hardware::gpu::HostGpuSample::error(
sample_seq,
format!("gpu sampler task failed: {error}"),
),
};
submit_sampler_sample_health(
&producer,
health_channel,
health_context,
"gpu",
sample_channel,
sample_seq,
sample.error.as_deref(),
);
seq = seq.saturating_add(1);
producer.submit_record(channel, &sample);
}
});
spawn_blocking_sampler(
handle,
producer,
channel,
health_channel,
health_context,
"gpu",
datastream::hardware::gpu::HOST_GPU_CHANNEL,
datastream::hardware::gpu::GPU_SAMPLE_INTERVAL,
"gpu sampler task failed",
datastream::hardware::gpu::sample,
datastream::hardware::gpu::HostGpuSample::error,
|s| s.error.as_deref(),
);
}
fn spawn_host_cpu_sampler(
@ -677,49 +698,20 @@ fn spawn_host_net_sampler(
health_channel: ChannelId,
health_context: SamplerHealthContext,
) {
handle.spawn(async move {
let sample_channel = datastream::hardware::net::HOST_NET_CHANNEL;
submit_sampler_started(
&producer,
health_channel,
health_context,
"net",
sample_channel,
datastream::hardware::net::HOST_NET_SAMPLE_INTERVAL,
);
let mut seq = 0_u64;
let mut interval =
tokio::time::interval(datastream::hardware::net::HOST_NET_SAMPLE_INTERVAL);
loop {
interval.tick().await;
let sample_seq = seq;
let sample = match tokio::task::spawn_blocking(move || {
datastream::hardware::net::sample(sample_seq)
})
.await
{
Ok(sample) => sample,
Err(error) => datastream::hardware::net::HostNetSample::error(
sample_seq,
format!("network sampler task failed: {error}"),
),
};
submit_sampler_sample_health(
&producer,
health_channel,
health_context,
"net",
sample_channel,
sample_seq,
sample.error.as_deref(),
);
seq = seq.saturating_add(1);
producer.submit_record(channel, &sample);
}
});
spawn_blocking_sampler(
handle,
producer,
channel,
health_channel,
health_context,
"net",
datastream::hardware::net::HOST_NET_CHANNEL,
datastream::hardware::net::HOST_NET_SAMPLE_INTERVAL,
"network sampler task failed",
datastream::hardware::net::sample,
datastream::hardware::net::HostNetSample::error,
|s| s.error.as_deref(),
);
}
fn spawn_arena_sampler(
@ -857,7 +849,6 @@ impl WorkerEdgeRuntime {
driver,
)?;
}
EdgeTransportEvent::StreamEnded { .. } => {}
EdgeTransportEvent::StreamFault {
edge_id: Some(edge_id),
..
@ -873,7 +864,8 @@ impl WorkerEdgeRuntime {
driver,
)?;
}
EdgeTransportEvent::StreamFault { edge_id: None, .. } => {}
EdgeTransportEvent::StreamEnded { .. }
| EdgeTransportEvent::StreamFault { edge_id: None, .. } => {}
}
}
Ok(())
@ -1952,7 +1944,7 @@ fn run() -> Result<(), String> {
datastream.channels.host_cpu,
sampler_health_channel,
sampler_health_context,
vec![std::process::id(), worker.pid()],
vec![std::process::id(), worker.child.id()],
);
worker_evt(
"worker_initialize",
@ -2032,7 +2024,7 @@ fn run() -> Result<(), String> {
emit_swim_telemetry(&mut datastream, &stack, "main_loop");
drain_debug_join_commands(&mut debug_join_rx, &mut driver, &config, &mut datastream);
datastream.tick();
worker.drain_stderr(&config, &mut datastream);
drain_worker_stderr(&worker.stderr_rx, &config, &mut datastream);
edge_runtime.poll_iroh(
&mut driver,
&stack,
@ -3171,15 +3163,6 @@ fn stop_stage_shard_child(child: &mut Option<Child>) {
let _ = child.wait();
}
fn stage_shard_cache_path(plan: &StageShardPlan) -> PathBuf {
let root = std::env::var("MVP_MODEL_CACHE_DIR")
.ok()
.filter(|value| !value.trim().is_empty())
.map(PathBuf::from)
.unwrap_or_else(|| PathBuf::from("/var/cache/mvp-models"));
root.join(plan.cache_file_name())
}
fn spawn_stage_shard_reader<R: Read + Send + 'static>(
stream: StageShardProcessStream,
reader: R,
@ -3251,7 +3234,12 @@ fn materialize_stage_shard_with_process(
driver: &mut IrohDriver,
stack: &DistributionRuntimeStack,
) -> Result<PathBuf, String> {
let output_path = stage_shard_cache_path(plan);
let output_path = std::env::var("MVP_MODEL_CACHE_DIR")
.ok()
.filter(|value| !value.trim().is_empty())
.map(PathBuf::from)
.unwrap_or_else(|| PathBuf::from("/var/cache/mvp-models"))
.join(plan.cache_file_name());
if output_path.is_file() {
match validate_stage_shard_cache(&output_path, plan) {
Ok(()) => {
@ -3678,6 +3666,13 @@ fn run_self_test(
Ok(())
}
fn env_optional(name: &str) -> Option<String> {
std::env::var(name)
.ok()
.map(|value| value.trim().to_owned())
.filter(|value| !value.is_empty())
}
#[derive(Clone)]
struct DeploymentConfig {
run_id: u64,
@ -4071,10 +4066,6 @@ impl TinygradWorker {
})
}
fn pid(&self) -> u32 {
self.child.id()
}
fn initialize(
&mut self,
device: &str,
@ -4402,10 +4393,6 @@ impl TinygradWorker {
.map_err(|e| format!("poll tinygrad helper: {e}"))
}
fn drain_stderr(&mut self, config: &DeploymentConfig, datastream: &mut NodeDatastream) {
drain_worker_stderr(&self.stderr_rx, config, datastream);
}
fn command(
&mut self,
command: Value,
@ -4512,10 +4499,3 @@ fn spawn_stdin_shutdown_listener() -> Receiver<()> {
});
rx
}
fn env_optional(name: &str) -> Option<String> {
std::env::var(name)
.ok()
.map(|value| value.trim().to_owned())
.filter(|value| !value.is_empty())
}

4
crates/mvp-system/src/observability/frame_archive.rs
View file

@ -19,10 +19,6 @@ pub(crate) struct FrameArchive {
}
impl FrameArchive {
pub(crate) fn open(path: &Path) -> Result<Self, String> {
Self::open_with_label(path, "datastream frame log")
}
pub(crate) fn open_with_label(path: &Path, label: &'static str) -> Result<Self, String> {
if let Some(parent) = path.parent()
&& !parent.as_os_str().is_empty()

4
crates/mvp-system/src/observability/lifecycle.rs
View file

@ -1,4 +1,3 @@
#![allow(dead_code)]
use serde::{Deserialize, Serialize};
@ -179,11 +178,13 @@ impl Event {
}
}
#[cfg(test)]
pub(crate) struct TraceBuilder {
run_id: RunId,
events: Vec<Event>,
}
#[cfg(test)]
impl TraceBuilder {
pub(crate) fn new(run_id: RunId) -> Self {
Self {
@ -473,6 +474,7 @@ impl TraceBuilder {
}
}
#[cfg(test)]
pub(crate) fn requires_log_scraping(_events: &[Event]) -> bool {
false
}

14
crates/mvp-system/src/observability/provisioning_logs.rs
View file

@ -1,4 +1,3 @@
#![allow(dead_code)]
use std::io::{BufRead, BufReader, Read};
use std::thread::{self, JoinHandle};
@ -14,12 +13,8 @@ use crate::provisioning::{
NodeProvisionSpec, PluginObservation, PluginSink, ProvisionLogLine, ProvisionLogStream,
};
pub(crate) fn node_datastream_id(node_id: u64) -> String {
node_id.to_string()
}
pub(crate) fn node_stream_id(run_id: u64, node_id: u64) -> StreamId {
StreamId::new(NodeId::new(&node_datastream_id(node_id)), Lifetime(run_id))
StreamId::new(NodeId::new(&node_id.to_string()), Lifetime(run_id))
}
#[derive(Clone)]
@ -46,10 +41,6 @@ impl BootstrapDatastreamBridge {
&self.spec
}
pub(crate) fn stream_id(&self) -> StreamId {
node_stream_id(self.spec.run_id, self.spec.node_id)
}
pub(crate) fn observe_stdout_line(&self, line: impl Into<String>) {
let line = line.into();
if let Some(frame) = parse_stdio_datastream_frame(&self.spec, &line) {
@ -185,6 +176,3 @@ pub(crate) fn parse_stdio_datastream_frame(
})
}
pub(crate) fn bootstrap_log_channel(node_id: u64, stream: ProvisionLogStream) -> String {
mvp_provision_log_channel(node_id, stream)
}

31
crates/mvp-system/src/observability/telemetry.rs
View file

@ -1,14 +1,11 @@
#![allow(dead_code)]
//! MVP-system-owned datastream channel records.
use datastream::hardware::net::HostNetSample;
use datastream::{ChannelRegistry, Record};
use datastream::Record;
use serde::{Deserialize, Serialize};
use crate::observability::lifecycle as obs;
use crate::provisioning::{self, ProvisionLogStream};
use data_plane::arena::ArenaSample;
/// Structured MVP lifecycle facts: run, node, stage, edge, ring, object, step, and worker events.
pub(crate) const MVP_LIFECYCLE: &str = "mvp.lifecycle";
@ -24,22 +21,6 @@ pub(crate) struct MvpLifecycleRecord {
pub event: obs::Event,
}
impl MvpLifecycleRecord {
pub(crate) fn new(event: obs::Event) -> Self {
Self { event }
}
pub(crate) fn kind(&self) -> obs::EventKind {
self.event.kind()
}
}
impl From<obs::Event> for MvpLifecycleRecord {
fn from(event: obs::Event) -> Self {
Self::new(event)
}
}
impl Record for MvpLifecycleRecord {
const CHANNEL: &'static str = MVP_LIFECYCLE;
}
@ -84,13 +65,3 @@ impl Record for MvpProvisionLogRecord {
const CHANNEL: &'static str = MVP_PROVISIONING_LOGS;
}
/// Registry fragment for consumers that want typed MVP datastream decoding.
pub(crate) fn channel_registry() -> ChannelRegistry {
let registry = ChannelRegistry::new()
.with_record::<MvpLifecycleRecord>()
.with_record::<MvpProvisionEventRecord>()
.with_record::<MvpProvisionLogRecord>()
.with_record::<HostNetSample>()
.with_record::<ArenaSample>();
registry
}

5
crates/mvp-system/src/orchestration/actor.rs
View file

@ -228,9 +228,7 @@ impl OrchestratorActor {
run_id: core::RunId(run_id),
stage_index,
}),
OrchestratorMsg::ObserveNodeRuntimeReady { .. } => {}
OrchestratorMsg::ObserveNodeRuntimeReadyAck { .. } => {}
OrchestratorMsg::ObserveWeightsReady { .. } => {}
OrchestratorMsg::ObserveNodeRuntimeReady { .. } | OrchestratorMsg::ObserveNodeRuntimeReadyAck { .. } | OrchestratorMsg::ObserveWeightsReady { .. } | OrchestratorMsg::Snapshot { .. } => {}
OrchestratorMsg::ObserveTokenInEndpointReady => {
self.core.observe(core::RunEvent::TokenInEndpointReady)
}
@ -275,7 +273,6 @@ impl OrchestratorActor {
self.core.observe(core::RunEvent::TokenEndpointsStopped)
}
OrchestratorMsg::AdvanceTimeMs(delta) => self.core.advance_time_ms(delta),
OrchestratorMsg::Snapshot { .. } => {}
}
}

176
crates/mvp-system/src/orchestration/app.rs
View file

@ -1792,8 +1792,8 @@ impl Config {
}
}
fn node_spec_env_keys(&self) -> Vec<&'static str> {
let mut keys = vec![
fn node_spec_env_keys(&self) -> Vec<String> {
let mut keys: Vec<String> = vec![
"MVP_RUN_ID",
"MVP_LOGICAL_NODE_ID",
"MVP_NODE_PROVIDER",
@ -1804,90 +1804,18 @@ impl Config {
"MVP_IROH_RELAY_MODE",
MVP_IROH_ENDPOINT_ADDR_MASK_ENV,
"MVP_PIPELINE_STAGES",
];
if self.relay.url.is_some() {
keys.push(MVP_IROH_RELAY_URL_ENV);
}
if self.provider.as_str() == "docker" {
keys.push("MVP_DOCKER_GPUS");
}
if std::env::var_os("DEV").is_some() {
keys.push("DEV");
}
if local_tinygrad_worker_env(self.provider.as_str()).is_some() {
keys.push("MVP_TINYGRAD_WORKER");
}
for key in [
"MVP_CPU_LINE_PROFILE",
"MVP_CPU_LINE_PROFILE_INTERVAL_MS",
"MVP_TOKEN_PROGRESS_EVERY",
"CUDA_DEVICE_SCHEDULE",
"MVP_MODEL_CACHE_DIR",
"HF_TOKEN",
] {
if std::env::var_os(key).is_some() {
keys.push(key);
}
}
match &self.gguf_source {
GgufSource::LocalPath(_) => keys.push("MVP_GGUF_LOCAL_PATH"),
GgufSource::HuggingFaceGguf { revision, .. } => {
keys.push("MVP_GGUF_REPO");
keys.push("MVP_GGUF_FILE");
if revision.is_some() {
keys.push("MVP_GGUF_REVISION");
}
}
}
if matches!(self.tokenizer, TokenizerSource::LocalPath(_)) {
keys.push("MVP_TOKENIZER_LOCAL_PATH");
}
if self.max_context.is_some() {
keys.push("MVP_MAX_CONTEXT");
}
]
.into_iter()
.map(str::to_owned)
.collect();
keys.extend(self.extra_worker_env().into_iter().map(|(k, _)| k));
keys
}
fn node_spec_for_stage(
&self,
coordinator: EndpointAddr,
orchestrator_actor: ActorAddress,
logical_node_id: u64,
stage_index: u32,
) -> Result<NodeProvisionSpec, String> {
/// Conditional worker env pairs shared by `node_spec_env_keys` and `node_spec_for_stage`.
fn extra_worker_env(&self) -> Vec<(String, String)> {
let provider_name = self.provider.as_str();
let mut env = vec![
("MVP_RUN_ID".to_owned(), self.run_id.to_string()),
(
"MVP_LOGICAL_NODE_ID".to_owned(),
logical_node_id.to_string(),
),
("MVP_STAGE_INDEX".to_owned(), stage_index.to_string()),
(
"MVP_PIPELINE_STAGES".to_owned(),
self.pipeline_stages.to_string(),
),
(
MVP_IROH_ENDPOINT_ADDR_MASK_ENV.to_owned(),
self.endpoint_addr_mask.as_str().to_owned(),
),
("MVP_NODE_PROVIDER".to_owned(), provider_name.to_owned()),
(
"MVP_COORDINATOR_ENDPOINT".to_owned(),
serde_json::to_string(&coordinator)
.map_err(|e| format!("serialize coordinator endpoint: {e}"))?,
),
(
"MVP_ORCHESTRATOR_ACTOR".to_owned(),
serde_json::to_string(&orchestrator_actor)
.map_err(|e| format!("serialize orchestrator actor: {e}"))?,
),
("MVP_MODEL_ID".to_owned(), self.model_id.clone()),
(
"MVP_IROH_RELAY_MODE".to_owned(),
relay_mode_env_value(&self.relay.mode).to_owned(),
),
];
let mut env = Vec::new();
if let Some(url) = &self.relay.url {
env.push((MVP_IROH_RELAY_URL_ENV.to_owned(), url.clone()));
}
@ -1932,6 +1860,50 @@ impl Config {
if let Some(max_context) = self.max_context {
env.push(("MVP_MAX_CONTEXT".to_owned(), max_context.to_string()));
}
env
}
fn node_spec_for_stage(
&self,
coordinator: EndpointAddr,
orchestrator_actor: ActorAddress,
logical_node_id: u64,
stage_index: u32,
) -> Result<NodeProvisionSpec, String> {
let provider_name = self.provider.as_str();
let mut env = vec![
("MVP_RUN_ID".to_owned(), self.run_id.to_string()),
(
"MVP_LOGICAL_NODE_ID".to_owned(),
logical_node_id.to_string(),
),
("MVP_STAGE_INDEX".to_owned(), stage_index.to_string()),
(
"MVP_PIPELINE_STAGES".to_owned(),
self.pipeline_stages.to_string(),
),
(
MVP_IROH_ENDPOINT_ADDR_MASK_ENV.to_owned(),
self.endpoint_addr_mask.as_str().to_owned(),
),
("MVP_NODE_PROVIDER".to_owned(), provider_name.to_owned()),
(
"MVP_COORDINATOR_ENDPOINT".to_owned(),
serde_json::to_string(&coordinator)
.map_err(|e| format!("serialize coordinator endpoint: {e}"))?,
),
(
"MVP_ORCHESTRATOR_ACTOR".to_owned(),
serde_json::to_string(&orchestrator_actor)
.map_err(|e| format!("serialize orchestrator actor: {e}"))?,
),
("MVP_MODEL_ID".to_owned(), self.model_id.clone()),
(
"MVP_IROH_RELAY_MODE".to_owned(),
relay_mode_env_value(&self.relay.mode).to_owned(),
),
];
env.extend(self.extra_worker_env());
let args = match provider_name {
"vastai" => self
.vastai
@ -2956,8 +2928,8 @@ fn wait_for_runtime_readies(
while let Ok(observation) = obs_rx.try_recv() {
emit_plugin_observation(orch_datastream, dashboard, provider, &observation);
match observation {
PluginObservation::DatastreamFrame { .. } => {}
PluginObservation::ProviderLine { .. }
PluginObservation::DatastreamFrame { .. }
| PluginObservation::ProviderLine { .. }
| PluginObservation::StdoutLine { .. }
| PluginObservation::StderrLine { .. } => {}
PluginObservation::Failed { reason, .. } => return Err(reason),
@ -3052,7 +3024,12 @@ fn wait_for_weights_loaded_count(
}
if last_resend.elapsed() >= Duration::from_secs(15) {
resend_attempt += 1;
let pending = pending_pipeline_weight_load_stages(pipeline_plan, &loaded_stages);
let mut pending: Vec<&run_plan::StagePlan> = pipeline_plan
.stages
.iter()
.filter(|stage| !loaded_stages.contains(&stage.stage_index))
.collect();
pending.sort_by_key(|stage| stage.stage_index);
if pending.is_empty() {
return Err(format!(
"missing unloaded pipeline weight stage; loaded {} of {expected_count}",
@ -3127,8 +3104,8 @@ fn wait_for_weights_loaded_count(
);
return Err(reason);
}
PluginObservation::DatastreamFrame { .. } => {}
PluginObservation::ProviderLine { .. }
PluginObservation::DatastreamFrame { .. }
| PluginObservation::ProviderLine { .. }
| PluginObservation::StdoutLine { .. }
| PluginObservation::StderrLine { .. } => {}
}
@ -3166,19 +3143,6 @@ fn wait_for_weights_loaded_count(
}
}
fn pending_pipeline_weight_load_stages<'a>(
pipeline_plan: &'a run_plan::RunPlan,
loaded_stages: &BTreeSet<u32>,
) -> Vec<&'a run_plan::StagePlan> {
let mut pending = pipeline_plan
.stages
.iter()
.filter(|stage| !loaded_stages.contains(&stage.stage_index))
.collect::<Vec<_>>();
pending.sort_by_key(|stage| stage.stage_index);
pending
}
struct PipelineStageProvision<'a> {
driver: &'a mut IrohDriver,
stack: &'a DistributionRuntimeStack,
@ -3732,7 +3696,7 @@ impl OrchDatastream {
producer,
channels: BTreeMap::new(),
channel_names: BTreeMap::new(),
archive: frame_log.map(FrameArchive::open).transpose()?,
archive: frame_log.map(|p| FrameArchive::open_with_label(p, "datastream frame log")).transpose()?
};
for name in [
MVP_PROVISIONING_EVENTS,
@ -4435,10 +4399,6 @@ impl PipelinePromptRuntime {
})
}
fn is_active(&self) -> bool {
self.active.is_some()
}
fn note_progress(&mut self) {
let now = Instant::now();
self.last_progress_at = Some(now);
@ -5026,7 +4986,7 @@ fn serve_prompts(
if active.is_none()
&& pipeline_runtime
.as_ref()
.is_none_or(|pipeline| !pipeline.is_active())
.is_none_or(|pipeline| pipeline.active.is_none())
&& let Ok(work) = work_rx.try_recv()
{
let request = work.request;
@ -5225,8 +5185,8 @@ fn drain_observations_with_exit(
PluginObservation::Exited {
node_id, status, ..
} => return Err(exit_message(node_id, status)),
PluginObservation::DatastreamFrame { .. } => {}
PluginObservation::ProviderLine { .. }
PluginObservation::DatastreamFrame { .. }
| PluginObservation::ProviderLine { .. }
| PluginObservation::StdoutLine { .. }
| PluginObservation::StderrLine { .. } => {}
}

15
crates/mvp-system/src/orchestration/distribution_stack.rs
View file

@ -1,4 +1,3 @@
#![allow(dead_code)]
//! MVP-system swactor distribution runtime wiring.
//!
@ -58,10 +57,6 @@ pub(crate) struct DistributionRuntimeStack {
}
impl DistributionRuntimeStack {
pub(crate) fn new(node_id: NodeId, config: DistributedNodeConfig) -> Self {
Self::new_with_codecs(node_id, config, |_| {})
}
pub(crate) fn new_with_codecs(
node_id: NodeId,
config: DistributedNodeConfig,
@ -210,16 +205,6 @@ impl DistributionRuntimeStack {
.send_to(self.actors.directory, DirectoryIn::Register(entry));
}
pub(crate) fn alive_count(&self) -> usize {
self.membership_mirror
.lock()
.expect("membership mirror poisoned")
.all_members()
.iter()
.filter(|entry| entry.state == MemberState::Alive)
.count()
}
pub(crate) fn member_state(&self, node_id: NodeId) -> Option<MemberState> {
self.membership_mirror
.lock()

1
crates/mvp-system/src/orchestration/engine_builder/mod.rs
View file

@ -1,4 +1,3 @@
#![allow(dead_code)]
//! Pool-based engine/node builder primitives.
//!

26
crates/mvp-system/src/orchestration/provider_adapters/relay.rs
View file

@ -1,4 +1,3 @@
#![allow(dead_code)]
//! Relay provisioning shims for MVP runtimes.
//!
@ -56,27 +55,8 @@ pub(crate) trait RelayProvider: Send {
fn provision_relay(&mut self, request: RelayProvisionRequest) -> Result<RelayLease, String>;
fn relay_mode(&self, lease: &RelayLease) -> Result<RelayMode, String>;
fn release_relay(&mut self, _lease: RelayLease) -> Result<(), String> {
Ok(())
}
}
#[derive(Clone, Copy, Debug, Default)]
pub(crate) struct LocalShimRelayProvider;
impl RelayProvider for LocalShimRelayProvider {
fn provision_relay(&mut self, request: RelayProvisionRequest) -> Result<RelayLease, String> {
Ok(RelayLease {
id: RelayLeaseId(format!("local-shim:{}", request.run_id)),
endpoints: Vec::new(),
})
}
fn relay_mode(&self, _lease: &RelayLease) -> Result<RelayMode, String> {
Ok(RelayMode::Disabled)
}
}
#[derive(Clone, Debug)]
pub(crate) struct StaticRelayProvider {
@ -128,7 +108,7 @@ impl RelayProvider for StaticRelayProvider {
}
pub(crate) fn relay_runtime_config_from_env(run_id: u64) -> Result<RelayRuntimeConfig, String> {
let mode = relay_mode_setting_from_env();
let mode = env_optional(MVP_IROH_RELAY_MODE_ENV).map(|value| value.to_ascii_lowercase());
let url = selected_relay_url_from_env();
relay_runtime_config_from_settings(run_id, mode.as_deref(), url.as_deref())
}
@ -183,10 +163,6 @@ fn relay_runtime_config_from_optional_static_provider(
})
}
fn relay_mode_setting_from_env() -> Option<String> {
env_optional(MVP_IROH_RELAY_MODE_ENV).map(|value| value.to_ascii_lowercase())
}
fn env_optional(name: &str) -> Option<String> {
std::env::var(name)
.ok()

364
crates/mvp-system/src/orchestration/provider_adapters/vastai/mod.rs
View file

@ -8,7 +8,6 @@ use std::thread::{self, JoinHandle};
use std::time::{Duration, Instant};
use datastream::DatastreamProducer;
use serde::{Deserialize, Serialize};
use swactor::actor::{ActorAddress, ActorInterface};
use swactor::runtime::{Ctx, ExternalSender, Runtime, RuntimeConfig, RuntimeHandle};
use swactor_vastai::{
@ -517,7 +516,19 @@ impl VastAiLeaseClient for ToolsVastAiLeaseClient {
label,
lifecycle,
))?;
endpoint_from_parts(contract_id, endpoint.ip, endpoint.port, ssh_user)
let host = endpoint.ip;
let port = endpoint.port;
if host.is_empty() || host == "unknown" {
return Err(format!("vastai contract {contract_id} has no SSH host"));
}
if port == 0 {
return Err(format!("vastai contract {contract_id} has no SSH port"));
}
Ok(VastAiSshEndpoint {
host,
port,
user: ssh_user.to_owned(),
})
}
fn spawn_provider_monitor(
@ -556,25 +567,6 @@ impl VastAiLeaseClient for ToolsVastAiLeaseClient {
}
}
fn endpoint_from_parts(
contract_id: u64,
host: String,
port: u16,
ssh_user: &str,
) -> Result<VastAiSshEndpoint, String> {
if host.is_empty() || host == "unknown" {
return Err(format!("vastai contract {contract_id} has no SSH host"));
}
if port == 0 {
return Err(format!("vastai contract {contract_id} has no SSH port"));
}
Ok(VastAiSshEndpoint {
host,
port,
user: ssh_user.to_owned(),
})
}
fn provider_terminal_start_error(
contract_id: u64,
actual: &str,
@ -611,12 +603,6 @@ fn provider_status_message_has_terminal_failure(message: &str) -> bool {
})
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub(crate) enum BootstrapStopReason {
RuntimeReady,
NodeStop,
}
pub(crate) trait VastAiBootstrapLauncher: Send {
type Handle: Send;
@ -629,7 +615,7 @@ pub(crate) trait VastAiBootstrapLauncher: Send {
lifecycle: LifecyclePolicy,
) -> Result<Self::Handle, String>;
fn stop_bootstrap(&mut self, handle: &mut Self::Handle, reason: BootstrapStopReason);
fn stop_bootstrap(&mut self, handle: &mut Self::Handle);
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
@ -672,7 +658,6 @@ struct SshBootstrapActor {
backoff: Duration,
observation_class: Option<&'static str>,
stopped: bool,
start_on_boot: bool,
}
impl SshBootstrapActor {
@ -698,7 +683,6 @@ impl SshBootstrapActor {
backoff: Duration::from_secs(1),
observation_class: None,
stopped: false,
start_on_boot: true,
}
}
@ -873,7 +857,7 @@ impl SshBootstrapActor {
delay.as_secs(),
self.attempt
));
self.backoff = next_ssh_backoff(self.backoff);
self.backoff = std::cmp::min(self.backoff.saturating_mul(2), Duration::from_secs(30));
self.attempt = self.attempt.saturating_add(1);
schedule_ssh_message(
self.sender.clone(),
@ -905,9 +889,7 @@ impl ActorInterface for SshBootstrapActor {
type Response = ();
fn on_start(&mut self, ctx: &Ctx) {
if self.start_on_boot {
let _ = ctx.send(ctx.self_addr(), SshBootstrapMsg::StartAttempt);
}
let _ = ctx.send(ctx.self_addr(), SshBootstrapMsg::StartAttempt);
}
fn handle(&mut self, ctx: &Ctx, msg: Self::Incoming) {
@ -996,7 +978,7 @@ impl VastAiBootstrapLauncher for SshCommandBootstrapLauncher {
})
}
fn stop_bootstrap(&mut self, handle: &mut Self::Handle, _reason: BootstrapStopReason) {
fn stop_bootstrap(&mut self, handle: &mut Self::Handle) {
let _ = handle.runtime.send_to(handle.actor, SshBootstrapMsg::Stop);
handle.runtime.tick();
}
@ -1123,8 +1105,12 @@ fn ssh_bootstrap_args(
args
}
fn next_ssh_backoff(current: Duration) -> Duration {
std::cmp::min(current.saturating_mul(2), Duration::from_secs(30))
fn emit_node_line(sink: &PluginSink, run_id: u64, node_id: u64, line: impl Into<String>) {
sink.observe(PluginObservation::ProviderLine {
run_id,
node_id,
line: line.into(),
});
}
pub(crate) struct VastAiProvisioningPlugin<C, B>
@ -1253,54 +1239,33 @@ where
}
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}"),
});
emit_node_line(&sink, spec.run_id, spec.node_id, format!("vastai provisioning label={label} stream={stream_id}"));
let request = self.build_request(&spec, label.clone());
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,
line: format!(
"vastai contract {} ready for SSH lookup",
instance.contract_id
),
});
sink.observe(PluginObservation::ProviderLine {
run_id: spec.run_id,
node_id: spec.node_id,
line: serde_json::json!({
"type": "VastAiLeaseReady",
"run_id": spec.run_id,
"node_id": spec.node_id,
"label": &label,
"image": &spec.image,
"contract_id": instance.contract_id,
"offer_id": instance.offer_id,
"host_id": instance.host_id,
"gpu_name": &instance.gpu_name,
"gpu_ram": instance.gpu_ram,
"dph_total": instance.dph_total,
})
.to_string(),
});
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(),
});
emit_node_line(&sink, spec.run_id, spec.node_id, format!("vastai contract {} ready for SSH lookup", instance.contract_id));
emit_node_line(&sink, spec.run_id, spec.node_id, 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());
emit_node_line(&sink, spec.run_id, spec.node_id, 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,
@ -1322,35 +1287,25 @@ where
));
}
};
sink.observe(PluginObservation::ProviderLine {
run_id: spec.run_id,
node_id: spec.node_id,
line: serde_json::json!({
"type": "VastAiSshEndpointReady",
"run_id": spec.run_id,
"node_id": spec.node_id,
"contract_id": instance.contract_id,
"host": &endpoint.host,
"port": endpoint.port,
"user": &endpoint.user,
})
.to_string(),
});
emit_node_line(&sink, spec.run_id, spec.node_id, serde_json::json!({
"type": "VastAiSshEndpointReady",
"run_id": spec.run_id,
"node_id": spec.node_id,
"contract_id": instance.contract_id,
"host": &endpoint.host,
"port": endpoint.port,
"user": &endpoint.user,
}).to_string());
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(),
});
emit_node_line(&sink, spec.run_id, spec.node_id, 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(),
@ -1433,11 +1388,7 @@ where
}
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}"),
});
emit_node_line(&sink, spec.run_id, spec.node_id, format!("vastai provisioning label={label} stream={stream_id}"));
let request = self.build_request(&spec, label.clone());
start_inputs.push((index, spec, label, request));
}
@ -1450,13 +1401,7 @@ where
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}"
),
});
emit_node_line(&sink, spec.run_id, spec.node_id, format!("vastai first-wave offer planning failed; falling back to per-node selection: {error}"));
}
vec![None; start_inputs.len()]
}
@ -1467,30 +1412,20 @@ where
{
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(),
});
emit_node_line(&sink, spec.run_id, spec.node_id, serde_json::json!({
"type": "VastAiFirstWaveOfferPlanned",
"run_id": spec.run_id,
"node_id": spec.node_id,
"label": &label,
"offer_id": offer_id,
}).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(),
});
emit_node_line(&sink, spec.run_id, spec.node_id, 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();
@ -1499,36 +1434,26 @@ where
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(),
});
emit_node_line(&worker_sink, spec.run_id, spec.node_id, 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());
emit_node_line(&worker_sink, spec.run_id, spec.node_id, 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,
@ -1575,34 +1500,24 @@ 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: 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(),
});
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": started.instance.contract_id,
"host": &started.endpoint.host,
"port": started.endpoint.port,
"user": &started.endpoint.user,
})
.to_string(),
});
emit_node_line(&sink, spec.run_id, spec.node_id, 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());
emit_node_line(&sink, spec.run_id, spec.node_id, serde_json::json!({
"type": "VastAiBootstrapObservationStarted",
"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(),
@ -1695,19 +1610,14 @@ where
let Some(node) = self.nodes.get_mut(&handle.id) else {
return Ok(());
};
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(),
});
emit_node_line(&node.sink, node.run_id, node.node_id, 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());
Ok(())
}
@ -1722,24 +1632,18 @@ where
self.leased_host_ids.remove(&host_id);
}
if let Some(mut bootstrap) = node.bootstrap.take() {
self.bootstrap
.stop_bootstrap(&mut bootstrap, BootstrapStopReason::NodeStop);
self.bootstrap.stop_bootstrap(&mut bootstrap);
}
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(),
});
emit_node_line(&node.sink, node.run_id, node.node_id, 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
}
}

82
crates/mvp-system/src/orchestration/run_fsm.rs
View file

@ -1,4 +1,3 @@
#![allow(dead_code)]
#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub(crate) struct RunId(pub(crate) u64);
@ -279,7 +278,13 @@ impl OrchestratorRun {
}
fn stage_ready(&mut self, run_id: RunId, stage_index: u32) {
if run_id != self.config.run_id || !self.plan_has_stage(stage_index) {
if run_id != self.config.run_id
|| !self.plan.as_ref().is_some_and(|plan| {
plan.stages
.iter()
.any(|stage| stage.stage_index == stage_index)
})
{
self.fault(RunFaultReason::UnknownStageReady { stage_index });
return;
}
@ -295,7 +300,13 @@ impl OrchestratorRun {
if eos {
self.complete();
} else if (self.injected_sequences.len() as u64) < self.config.max_tokens {
self.inject_decode(sequence + 1, token_id, sequence);
self.inject(TokenObjectInjection {
sequence: sequence + 1,
payload: TokenObjectPayload::Decode {
token_id,
sampling: SamplingData { source_sequence: sequence },
},
});
} else {
self.complete();
}
@ -344,30 +355,23 @@ impl OrchestratorRun {
if self.terminal || !self.provisioned || !self.injected_sequences.is_empty() {
return;
}
if self.token_in_ready && self.token_out_ready && self.all_stages_ready() {
self.inject_prompt(0);
if self.token_in_ready
&& self.token_out_ready
&& self.plan.as_ref().is_some_and(|plan| {
plan.stages
.iter()
.all(|stage| self.ready_stages.contains(&stage.stage_index))
})
{
self.inject(TokenObjectInjection {
sequence: 0,
payload: TokenObjectPayload::Prompt {
tokens: self.config.prompt.clone(),
},
});
}
}
fn inject_prompt(&mut self, sequence: u64) {
self.inject(TokenObjectInjection {
sequence,
payload: TokenObjectPayload::Prompt {
tokens: self.config.prompt.clone(),
},
});
}
fn inject_decode(&mut self, sequence: u64, token_id: u32, source_sequence: u64) {
self.inject(TokenObjectInjection {
sequence,
payload: TokenObjectPayload::Decode {
token_id,
sampling: SamplingData { source_sequence },
},
});
}
fn inject(&mut self, object: TokenObjectInjection) {
if self.terminal {
return;
@ -434,7 +438,11 @@ impl OrchestratorRun {
if !self.teardown_started || !self.token_endpoints_stopped {
return;
}
if !self.all_stages_stopped() {
if !self.plan.as_ref().is_some_and(|plan| {
plan.stages
.iter()
.all(|stage| self.stopped_stages.contains(&stage.stage_index))
}) {
return;
}
self.mark_torn_down();
@ -451,28 +459,4 @@ impl OrchestratorRun {
});
}
}
fn plan_has_stage(&self, stage_index: u32) -> bool {
self.plan.as_ref().is_some_and(|plan| {
plan.stages
.iter()
.any(|stage| stage.stage_index == stage_index)
})
}
fn all_stages_ready(&self) -> bool {
self.plan.as_ref().is_some_and(|plan| {
plan.stages
.iter()
.all(|stage| self.ready_stages.contains(&stage.stage_index))
})
}
fn all_stages_stopped(&self) -> bool {
self.plan.as_ref().is_some_and(|plan| {
plan.stages
.iter()
.all(|stage| self.stopped_stages.contains(&stage.stage_index))
})
}
}

28
crates/mvp-system/src/orchestration/run_plan.rs
View file

@ -1,4 +1,3 @@
#![allow(dead_code)]
pub(crate) const MO01_HEADER_BYTES: u64 = 40;
const TOKEN_ID_WIDTH_BYTES: u32 = 4;
@ -126,13 +125,11 @@ pub(crate) enum RingDirection {
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) enum HostPinning {
Pageable,
PinnedRequired,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) enum WakeCoalescing {
PendingBit,
ReadySet,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
@ -169,15 +166,12 @@ pub(crate) enum ObjectKind {
Token,
Activation,
Weight,
ModelShard,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) enum ShapeRule {
TokenIds,
ActivationRows { max_seq_len: u32, hidden_dim: u32 },
WeightTensor,
ModelShardBytes,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
@ -298,7 +292,6 @@ pub(crate) enum PlanRejectionKind {
DuplicateStageAssignment,
MissingStage,
InvalidStageCount,
EdgeEndpointMismatch,
ModelStageLayoutMismatch,
InvalidObjectSpec,
UnsupportedShapeOrLayout,
@ -439,7 +432,8 @@ pub(crate) fn plan_run(input: PlannerInput) -> Result<RunPlan, PlanRejection> {
let mut stages = Vec::with_capacity(input.stage_count as usize);
for placement in &placements {
let stage_index = placement.stage_index;
let (start, end) = layer_range(input.model.num_layers, input.stage_count, stage_index);
let start = (u64::from(input.model.num_layers) * u64::from(stage_index) / u64::from(input.stage_count)) as u32;
let end = (u64::from(input.model.num_layers) * u64::from(stage_index + 1) / u64::from(input.stage_count)) as u32;
let inbound_edge = if stage_index == 0 {
token_in_edge
} else {
@ -510,7 +504,11 @@ pub(crate) fn derive_stage_provision(
outbound: OutboundEdgeProvision {
edge_id: outbound.edge_id,
kind: outbound.kind,
consumer_node_id: endpoint_node_id(&outbound.consumer),
consumer_node_id: match &outbound.consumer {
EdgeEndpoint::Orchestrator { node_id } | EdgeEndpoint::Stage { node_id, .. } => {
*node_id
}
},
object_spec: outbound.object_spec,
ring_spec: ring_spec_for_direction(outbound.ring_spec, RingDirection::Egress),
},
@ -607,18 +605,6 @@ fn ring_spec_for_direction(mut spec: RingSpec, direction: RingDirection) -> Ring
spec
}
fn layer_range(num_layers: u32, stage_count: u32, stage_index: u32) -> (u32, u32) {
let start = (u64::from(num_layers) * u64::from(stage_index) / u64::from(stage_count)) as u32;
let end = (u64::from(num_layers) * u64::from(stage_index + 1) / u64::from(stage_count)) as u32;
(start, end)
}
fn endpoint_node_id(endpoint: &EdgeEndpoint) -> NodeId {
match endpoint {
EdgeEndpoint::Orchestrator { node_id } | EdgeEndpoint::Stage { node_id, .. } => *node_id,
}
}
fn valid_ring(spec: RingSpec) -> bool {
spec.data_capacity > 0 && spec.alignment > 0 && spec.alignment.is_power_of_two()
}

79
crates/mvp-system/src/staging/actor.rs
View file

@ -1,79 +0,0 @@
use swactor::actor::{ActorAddress, ActorInterface};
use swactor::runtime::Ctx;
use crate::staging::control as core;
#[derive(Clone, Debug, PartialEq, Eq)]
pub(crate) enum StageControllerMsg {
Observe(core::StageEvent),
Snapshot { reply_to: ActorAddress },
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub(crate) enum StageControllerReport {
Command(core::StageCommand),
Lifecycle(core::StageLifecycleEvent),
Snapshot {
commands: Vec<core::StageCommand>,
events: Vec<core::StageLifecycleEvent>,
},
}
pub(crate) struct StageControllerActor {
core: core::StageController,
report_to: Option<ActorAddress>,
command_cursor: usize,
event_cursor: usize,
}
impl StageControllerActor {
pub(crate) fn new(local_node_id: core::NodeId, report_to: Option<ActorAddress>) -> Self {
Self {
core: core::StageController::new(local_node_id),
report_to,
command_cursor: 0,
event_cursor: 0,
}
}
fn drain_outputs(&mut self, ctx: &Ctx) {
let Some(report_to) = self.report_to else {
self.command_cursor = self.core.commands().len();
self.event_cursor = self.core.events().len();
return;
};
for command in &self.core.commands()[self.command_cursor..] {
let _ = ctx.send(report_to, StageControllerReport::Command(command.clone()));
}
self.command_cursor = self.core.commands().len();
for event in &self.core.events()[self.event_cursor..] {
let _ = ctx.send(report_to, StageControllerReport::Lifecycle(event.clone()));
}
self.event_cursor = self.core.events().len();
}
}
impl ActorInterface for StageControllerActor {
type Incoming = StageControllerMsg;
type Response = ();
fn handle(&mut self, ctx: &Ctx, msg: Self::Incoming) {
match msg {
StageControllerMsg::Observe(event) => {
self.core.observe(event);
self.drain_outputs(ctx);
}
StageControllerMsg::Snapshot { reply_to } => {
let _ = ctx.send(
reply_to,
StageControllerReport::Snapshot {
commands: self.core.commands().to_vec(),
events: self.core.events().to_vec(),
},
);
}
}
}
}

1
crates/mvp-system/src/staging/control.rs
View file

@ -1,4 +1,3 @@
#![allow(dead_code)]
use crate::gguf_shard::StageShardPlan;
use crate::run_plan::{GgufSource, TokenizerSource};

14
crates/mvp-system/src/staging/gguf_metadata.rs
View file

@ -6,8 +6,8 @@ use std::path::Path;
use crate::gguf_common::{GgufValueType, read_integer_value, read_u32, read_u64};
use crate::run_plan::{self, DTypeFamily, GgufSource, TokenizerSource};
const GGUF_MAGIC: &[u8; 4] = b"GGUF";
const SUPPORTED_GGUF_VERSION: u32 = 3;
pub(crate) const GGUF_MAGIC: &[u8; 4] = b"GGUF";
pub(crate) const SUPPORTED_GGUF_VERSION: u32 = 3;
const DEFAULT_EFFECTIVE_CONTEXT: u64 = 512;
const MAX_METADATA_STRING_BYTES: u64 = 16 * 1024 * 1024;
const MAX_METADATA_KEY_BYTES: u64 = 1024 * 1024;
@ -161,7 +161,7 @@ fn required_u32(map: &BTreeMap<String, u64>, key: &str, label: &str) -> Result<u
u32::try_from(value).map_err(|_| format!("GGUF metadata {label} key {key} exceeds u32"))
}
fn skip_scalar<R: Read + Seek>(reader: &mut R, value_type: GgufValueType) -> Result<(), String> {
pub(crate) fn skip_scalar<R: Read + Seek>(reader: &mut R, value_type: GgufValueType) -> Result<(), String> {
match value_type {
GgufValueType::String => skip_gguf_string(reader),
GgufValueType::Array => skip_array(reader),
@ -169,7 +169,7 @@ fn skip_scalar<R: Read + Seek>(reader: &mut R, value_type: GgufValueType) -> Res
}
}
fn skip_array<R: Read + Seek>(reader: &mut R) -> Result<(), String> {
pub(crate) fn skip_array<R: Read + Seek>(reader: &mut R) -> Result<(), String> {
let element_type = GgufValueType::read(reader, "GGUF metadata value type")?;
let len = read_u64(reader)?;
match element_type {
@ -195,7 +195,7 @@ fn skip_array<R: Read + Seek>(reader: &mut R) -> Result<(), String> {
}
}
fn read_gguf_string<R: Read + Seek>(reader: &mut R, max_len: u64) -> Result<String, String> {
pub(crate) fn read_gguf_string<R: Read + Seek>(reader: &mut R, max_len: u64) -> Result<String, String> {
let len = read_u64(reader)?;
if len > max_len {
return Err(format!(
@ -211,12 +211,12 @@ fn read_gguf_string<R: Read + Seek>(reader: &mut R, max_len: u64) -> Result<Stri
String::from_utf8(bytes).map_err(|e| format!("GGUF metadata string is not UTF-8: {e}"))
}
fn skip_gguf_string<R: Read + Seek>(reader: &mut R) -> Result<(), String> {
pub(crate) fn skip_gguf_string<R: Read + Seek>(reader: &mut R) -> Result<(), String> {
let len = read_u64(reader)?;
skip_bytes(reader, len)
}
fn skip_bytes<R: Seek>(reader: &mut R, mut bytes: u64) -> Result<(), String> {
pub(crate) fn skip_bytes<R: Seek>(reader: &mut R, mut bytes: u64) -> Result<(), String> {
while bytes > 0 {
let chunk = bytes.min(i64::MAX as u64);
reader

89
crates/mvp-system/src/staging/gguf_shard.rs
View file

@ -1,14 +1,13 @@
use std::fs::File;
use std::io::{Read, Seek, SeekFrom, Write};
use std::io::{Read, Seek, Write};
use std::path::{Path, PathBuf};
use crate::gguf_common::{GgufValueType, read_integer_value, read_u32, read_u64};
use serde::{Deserialize, Serialize};
use crate::run_plan::GgufSource;
use crate::staging::gguf_metadata::{skip_scalar as skip_value, read_gguf_string, GGUF_MAGIC, SUPPORTED_GGUF_VERSION};
const GGUF_MAGIC: &[u8; 4] = b"GGUF";
const SUPPORTED_GGUF_VERSION: u32 = 3;
const DEFAULT_ALIGNMENT: u64 = 32;
const MAX_STRING_BYTES: u64 = 64 * 1024 * 1024;
const STAGE_SHARD_CACHE_FORMAT_VERSION: &str = "stage-shard-cache-v2";
@ -65,20 +64,9 @@ impl StageShardPlan {
format!("{}.stage-{:05}.gguf", self.cache_key, self.stage_index)
}
pub(crate) fn source_url(&self) -> Result<String, String> {
source_url(&self.source)
}
pub(crate) fn planned_tensor_fetch_bytes(&self) -> u64 {
self.merged_tensor_ranges
.iter()
.map(|range| range.len)
.sum()
}
pub(crate) fn planned_fetch_bytes(&self) -> u64 {
self.metadata_end
.saturating_add(self.planned_tensor_fetch_bytes())
.saturating_add(self.merged_tensor_ranges.iter().map(|r| r.len).sum())
}
pub(crate) fn planned_range_count(&self) -> usize {
@ -216,7 +204,7 @@ pub(crate) fn source_url(source: &GgufSource) -> Result<String, String> {
} => Ok(format!(
"https://huggingface.co/{repo}/resolve/{}/{}",
revision.as_deref().unwrap_or("main"),
encode_hf_path(file)
file.split('/').map(percent_encode_path_segment).collect::<Vec<_>>().join("/")
)),
GgufSource::LocalPath(path) => Err(format!(
"stage shard range fetching requires a remote Hugging Face source; got local path {path:?}"
@ -224,13 +212,6 @@ pub(crate) fn source_url(source: &GgufSource) -> Result<String, String> {
}
}
fn encode_hf_path(path: &str) -> String {
path.split('/')
.map(percent_encode_path_segment)
.collect::<Vec<_>>()
.join("/")
}
fn percent_encode_path_segment(segment: &str) -> String {
let mut out = String::new();
for byte in segment.bytes() {
@ -477,7 +458,7 @@ pub(crate) fn materialize_stage_shard_http<F>(
where
F: FnMut(serde_json::Value),
{
let url = plan.source_url()?;
let url = source_url(&plan.source)?;
materialize_stage_shard_from_url(plan, &url, output_path, emit)
}
@ -799,66 +780,6 @@ fn pad_writer_to_alignment<W: Write + Seek>(writer: &mut W, alignment: u64) -> R
Ok(())
}
fn skip_value<R: Read + Seek>(reader: &mut R, value_type: GgufValueType) -> Result<(), String> {
match value_type {
GgufValueType::String => skip_gguf_string(reader),
GgufValueType::Array => skip_array(reader),
scalar => skip_bytes(reader, scalar.fixed_width().expect("scalar width")),
}
}
fn skip_array<R: Read + Seek>(reader: &mut R) -> Result<(), String> {
let element_type = GgufValueType::read(reader, "GGUF value type")?;
let len = read_u64(reader)?;
match element_type {
GgufValueType::String => {
for _ in 0..len {
skip_gguf_string(reader)?;
}
Ok(())
}
GgufValueType::Array => {
for _ in 0..len {
skip_array(reader)?;
}
Ok(())
}
scalar => {
let width = scalar.fixed_width().expect("scalar array width");
let bytes = width
.checked_mul(len)
.ok_or_else(|| "GGUF array byte count overflow".to_owned())?;
skip_bytes(reader, bytes)
}
}
}
fn read_gguf_string<R: Read>(reader: &mut R, max_len: u64) -> Result<String, String> {
let len = read_u64(reader)?;
if len > max_len {
return Err(format!("GGUF string length {len} exceeds {max_len}"));
}
let len = usize::try_from(len).map_err(|_| "GGUF string length exceeds usize".to_owned())?;
let mut bytes = vec![0_u8; len];
reader
.read_exact(&mut bytes)
.map_err(|e| format!("read GGUF string: {e}"))?;
String::from_utf8(bytes).map_err(|e| format!("GGUF string is not UTF-8: {e}"))
}
fn skip_gguf_string<R: Read + Seek>(reader: &mut R) -> Result<(), String> {
let len = read_u64(reader)?;
skip_bytes(reader, len)
}
fn skip_bytes<R: Seek>(reader: &mut R, bytes: u64) -> Result<(), String> {
let offset = i64::try_from(bytes).map_err(|_| format!("cannot seek over {bytes} bytes"))?;
reader
.seek(SeekFrom::Current(offset))
.map_err(|e| format!("skip bytes: {e}"))?;
Ok(())
}
fn align_to(value: u64, alignment: u64) -> Result<u64, String> {
if alignment == 0 {
return Err("GGUF alignment must be non-zero".to_owned());

3
crates/mvp-system/src/staging/mod.rs
View file

@ -1,9 +1,6 @@
#![allow(dead_code)]
//! MVP stage control, shard planning, and weight lifecycle public surface.
#[cfg(test)]
pub(crate) mod actor;
pub(crate) mod control;
pub(crate) mod gguf_metadata;
#[cfg(test)]