refactor: more grind on mvp-system
This commit is contained in:
parent
0902de6394
commit
79e1a48d28
20 changed files with 519 additions and 1038 deletions
|
|
@ -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())
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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())
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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 { .. } => {}
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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 { .. } => {}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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,4 +1,3 @@
|
|||
#![allow(dead_code)]
|
||||
|
||||
//! Pool-based engine/node builder primitives.
|
||||
//!
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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,4 +1,3 @@
|
|||
#![allow(dead_code)]
|
||||
|
||||
use crate::gguf_shard::StageShardPlan;
|
||||
use crate::run_plan::{GgufSource, TokenizerSource};
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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());
|
||||
|
|
|
|||
|
|
@ -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)]
|
||||
|
|
|
|||
Loading…
Reference in a new issue