From cc1ca5bf30016b34c2f65d15da3a495d34f1acf5 Mon Sep 17 00:00:00 2001 From: Zachery Aaron Shores-Chmielewski Date: Wed, 29 Jul 2026 21:09:26 +0400 Subject: [PATCH] refactor(mvp-system): replace image engine with static builders - Drop NodeImageProvider/ImageCommandRunner and git-worktree hashing for free docker-image helpers and static structs (StaticNodeLauncher, FixedLinearPipelinePlanner, StaticPoolProvider). - Delete engine_builder runtime_stack.rs and workload.rs; collapse node-image and chat config loading. Signed-off-by: Zachery Aaron Shores-Chmielewski --- crates/mvp-system/src/chat/config.rs | 7 - crates/mvp-system/src/chat/mod.rs | 10 +- crates/mvp-system/src/chat/node_image.rs | 712 ++++------ crates/mvp-system/src/chat/runtime.rs | 427 +++--- crates/mvp-system/src/lib.rs | 16 +- crates/mvp-system/src/node/mod.rs | 6 +- .../src/node/worker_node_runtime.rs | 233 ++-- crates/mvp-system/src/orchestration/app.rs | 1156 +++++++---------- .../orchestration/engine_builder/engine.rs | 102 +- .../src/orchestration/engine_builder/error.rs | 43 +- .../orchestration/engine_builder/launcher.rs | 54 +- .../src/orchestration/engine_builder/mod.rs | 9 +- .../src/orchestration/engine_builder/model.rs | 67 +- .../engine_builder/node_image.rs | 31 +- .../orchestration/engine_builder/planner.rs | 19 +- .../src/orchestration/engine_builder/pool.rs | 105 +- .../src/orchestration/engine_builder/roles.rs | 5 - .../engine_builder/runtime_stack.rs | 173 --- .../orchestration/engine_builder/workload.rs | 13 - crates/mvp-system/src/orchestration/mod.rs | 19 +- .../mvp-system/src/orchestration/run_plan.rs | 14 +- 21 files changed, 1073 insertions(+), 2148 deletions(-) delete mode 100644 crates/mvp-system/src/chat/config.rs delete mode 100644 crates/mvp-system/src/orchestration/engine_builder/runtime_stack.rs delete mode 100644 crates/mvp-system/src/orchestration/engine_builder/workload.rs diff --git a/crates/mvp-system/src/chat/config.rs b/crates/mvp-system/src/chat/config.rs deleted file mode 100644 index ddbb8dc..0000000 --- a/crates/mvp-system/src/chat/config.rs +++ /dev/null @@ -1,7 +0,0 @@ -pub const DEFAULT_CONFIG_PATH: &str = ".config/config.toml"; - -pub fn normalize_optional(value: Option) -> Option { - value - .map(|value| value.trim().to_owned()) - .filter(|value| !value.is_empty()) -} diff --git a/crates/mvp-system/src/chat/mod.rs b/crates/mvp-system/src/chat/mod.rs index f246585..83b18e7 100644 --- a/crates/mvp-system/src/chat/mod.rs +++ b/crates/mvp-system/src/chat/mod.rs @@ -1,12 +1,4 @@ //! MVP operator chat wrapper public surface. -pub mod config; mod node_image; -mod runtime; - -pub(super) fn run_from_args(args: I) -> std::process::ExitCode -where - I: IntoIterator, -{ - runtime::run_from_args(args) -} +pub(super) mod runtime; diff --git a/crates/mvp-system/src/chat/node_image.rs b/crates/mvp-system/src/chat/node_image.rs index ee9e90c..904b917 100644 --- a/crates/mvp-system/src/chat/node_image.rs +++ b/crates/mvp-system/src/chat/node_image.rs @@ -21,30 +21,13 @@ const NODE_IMAGE_SOURCE_HASH_LABEL: &str = "org.swactor.mvp.node.source-hash"; const NODE_IMAGE_WORKER_HASH_LABEL: &str = "org.swactor.mvp.node.worker-hash"; const NODE_IMAGE_BASE_HASH_LABEL: &str = "org.swactor.mvp.node.base-hash"; const BASE_IMAGE_SOURCE_HASH_LABEL: &str = "org.swactor.mvp.base.source-hash"; -const NODE_IMAGE_PRUNE_ENV: &str = "MVP_NODE_IMAGE_PRUNE"; -const NODE_IMAGE_PRUNE_KEEP_ENV: &str = "MVP_NODE_IMAGE_PRUNE_KEEP"; -const DEFAULT_DIRTY_IMAGE_KEEP: usize = 3; -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub(super) enum NodeImageProvider { - Docker, - VastAi, -} - -impl NodeImageProvider { - fn requires_remote_image(self) -> bool { - matches!(self, Self::VastAi) - } -} - -#[derive(Clone, Debug)] pub(super) struct NodeImageRequest { pub(super) requested_image: String, pub(super) base_image: String, pub(super) node_bin: PathBuf, - pub(super) provider: NodeImageProvider, + pub(super) requires_registry_image: bool, pub(super) extra_tag: Option, - pub(super) push: bool, pub(super) force_refresh: bool, pub(super) enabled: bool, } @@ -84,62 +67,17 @@ pub(super) trait NodeImageProgressSink { fn emit(&mut self, event: NodeImageProgressEvent); } -impl NodeImageProgressSink for F -where - F: FnMut(NodeImageProgressEvent), -{ - fn emit(&mut self, event: NodeImageProgressEvent) { - self(event); - } -} - -trait ImageCommandRunner { - fn run_status( - &mut self, - root: &Path, - program: &str, - args: &[String], - label: &str, - image_ref: Option<&str>, - progress: &mut Option<&mut dyn NodeImageProgressSink>, - ) -> Result<(), String>; - - fn docker_image_exists(&mut self, root: &Path, image_ref: &str) -> bool; - - fn docker_image_labels( - &mut self, - root: &Path, - image_ref: &str, - ) -> Result>, String>; - - fn docker_manifest_exists(&mut self, root: &Path, image_ref: &str) -> bool; - - fn docker_image_has_container(&mut self, root: &Path, image_ref: &str) -> bool; - - fn docker_image_tags( - &mut self, - root: &Path, - repository: &str, - ) -> Result, String>; - - fn docker_image_remove(&mut self, root: &Path, image_ref: &str) -> Result<(), String>; -} - -struct RealImageCommandRunner; - pub(super) fn prepare_node_image_with_progress( request: NodeImageRequest, progress: Option<&mut dyn NodeImageProgressSink>, ) -> Result { let mut progress = progress; - let mut runner = RealImageCommandRunner; - prepare_node_image_inner(request, &mut progress, &mut runner) + prepare_node_image_inner(request, &mut progress) } fn prepare_node_image_inner( request: NodeImageRequest, progress: &mut Option<&mut dyn NodeImageProgressSink>, - runner: &mut dyn ImageCommandRunner, ) -> Result { emit_image_reference(progress, "requested", &request.requested_image); if !request.enabled { @@ -149,7 +87,16 @@ fn prepare_node_image_inner( emit_image_reference(progress, "base", &request.base_image); let root = workspace_root()?; let image = ImageName::parse(&request.requested_image)?; - if request.provider.requires_remote_image() && !looks_registry_reachable(&image.repository) { + let first_repository_component = image + .repository + .split('/') + .next() + .unwrap_or(&image.repository); + let registry_reachable = image.repository.contains('/') + || first_repository_component.contains('.') + || first_repository_component.contains(':') + || first_repository_component == "localhost"; + if request.requires_registry_image && !registry_reachable { return Err(format!( "VastAI node image {:?} must include a registry namespace", image.repository @@ -157,7 +104,6 @@ fn prepare_node_image_inner( } run_status( - runner, progress, &root, "cargo", @@ -178,47 +124,54 @@ 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 = file_content_hash(&root, Path::new("apps/mvp-node/tinygrad_worker.py"))?; - let expected_node_labels = - node_image_labels(&tag, &image_content_hash, &worker_hash, &base_hash); - let expected_base_labels = base_image_labels(&base_hash); + let worker_hash = hash_relative_files( + &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()), + (NODE_IMAGE_SOURCE_HASH_LABEL, image_content_hash.as_str()), + (NODE_IMAGE_WORKER_HASH_LABEL, worker_hash.as_str()), + (NODE_IMAGE_BASE_HASH_LABEL, base_hash.as_str()), + ]; + let expected_base_labels = vec![(BASE_IMAGE_SOURCE_HASH_LABEL, base_hash.as_str())]; let alias_tags = alias_tags(&image, request.extra_tag.as_deref(), &tag)?; for alias in alias_refs(&image, &alias_tags) { emit_image_reference(progress, "alias", &alias); } - let remote_required = request.provider.requires_remote_image() || request.push; + let remote_required = request.requires_registry_image; - let local_image_matches = - docker_image_labels_match(runner, &root, &image_ref, &expected_node_labels)?; - let remote_available = remote_required && runner.docker_manifest_exists(&root, &image_ref); + let local_image_matches = docker_image_labels_match(&root, &image_ref, &expected_node_labels)?; + let remote_available = remote_required && docker_manifest_exists(&root, &image_ref); if !request.force_refresh && remote_required && remote_available { - ensure_aliases_for_remote(runner, progress, &root, &image_ref, &image, &alias_tags)?; - prune_old_dirty_images(runner, &root, &image, &tag); + ensure_aliases_for_remote(progress, &root, &image_ref, &image, &alias_tags)?; + prune_old_dirty_images(&root, &image, &tag); return Ok(image_ref); } if !request.force_refresh && remote_required && local_image_matches { - ensure_aliases_local(runner, progress, &root, &image_ref, &image, &alias_tags)?; - push_image(runner, progress, &root, &image_ref)?; + ensure_aliases_local(progress, &root, &image_ref, &image, &alias_tags)?; + push_image(progress, &root, &image_ref)?; for alias in alias_refs(&image, &alias_tags) { - push_image(runner, progress, &root, &alias)?; + push_image(progress, &root, &alias)?; } - prune_old_dirty_images(runner, &root, &image, &tag); + prune_old_dirty_images(&root, &image, &tag); return Ok(image_ref); } if !request.force_refresh && !remote_required && local_image_matches { - ensure_aliases_local(runner, progress, &root, &image_ref, &image, &alias_tags)?; - prune_old_dirty_images(runner, &root, &image, &tag); + ensure_aliases_local(progress, &root, &image_ref, &image, &alias_tags)?; + prune_old_dirty_images(&root, &image, &tag); return Ok(image_ref); } let base_image_matches = - docker_image_labels_match(runner, &root, &request.base_image, &expected_base_labels)?; + docker_image_labels_match(&root, &request.base_image, &expected_base_labels)?; if !base_image_matches { - run_status_vec( - runner, - progress, + run_status_command( &root, "docker", - vec![ + &vec![ "build".to_owned(), "-f".to_owned(), "apps/mvp-node/Dockerfile.base".to_owned(), @@ -230,6 +183,7 @@ fn prepare_node_image_inner( ], "build mvp node base image", Some(&request.base_image), + progress, )?; } @@ -248,25 +202,24 @@ fn prepare_node_image_inner( build_args.push(format!("{key}={value}")); } build_args.extend(["-t".to_owned(), image_ref.clone(), ".".to_owned()]); - run_status_vec( - runner, - progress, + run_status_command( &root, "docker", - build_args, + &build_args, "build mvp node image", Some(&image_ref), + progress, )?; - ensure_aliases_local(runner, progress, &root, &image_ref, &image, &alias_tags)?; + ensure_aliases_local(progress, &root, &image_ref, &image, &alias_tags)?; if remote_required { - push_image(runner, progress, &root, &image_ref)?; + push_image(progress, &root, &image_ref)?; for alias in alias_refs(&image, &alias_tags) { - push_image(runner, progress, &root, &alias)?; + push_image(progress, &root, &alias)?; } } - prune_old_dirty_images(runner, &root, &image, &tag); + prune_old_dirty_images(&root, &image, &tag); Ok(image_ref) } @@ -290,7 +243,10 @@ fn workspace_root() -> Result { } fn image_version_tag(root: &Path, image_content_hash: &str) -> Result { - if git_worktree_clean(root)? { + if git_capture(root, &["status", "--porcelain"])? + .trim() + .is_empty() + { let sha = git_capture(root, &["rev-parse", "--short=12", "HEAD"])?; Ok(format!("git-{}", sha.trim())) } else { @@ -298,12 +254,6 @@ fn image_version_tag(root: &Path, image_content_hash: &str) -> Result Result { - Ok(git_capture(root, &["status", "--porcelain"])? - .trim() - .is_empty()) -} - fn git_capture(root: &Path, args: &[&str]) -> Result { let output = Command::new("git") .current_dir(root) @@ -355,10 +305,6 @@ fn content_hash_for_inputs(root: &Path, inputs: &[&str]) -> Result Result { - hash_relative_files(root, vec![relative_path(root, &root.join(path))?]) -} - fn hash_relative_files(root: &Path, files: Vec) -> Result { hash_relative_files_with_salts(root, files, &[]) } @@ -409,7 +355,7 @@ fn collect_hash_inputs(root: &Path, path: &Path, out: &mut Vec) -> Resu let display = display_workspace_path(root, path); let metadata = fs::metadata(path).map_err(|e| format!("stat {display}: {e}"))?; if metadata.is_file() { - if !skip_file(path) { + if !matches!(path.extension().and_then(|ext| ext.to_str()), Some("pyc")) { out.push(relative_path(root, path)?); } return Ok(()); @@ -459,10 +405,6 @@ fn skip_dir(path: &Path) -> bool { ) } -fn skip_file(path: &Path) -> bool { - matches!(path.extension().and_then(|ext| ext.to_str()), Some("pyc")) -} - fn alias_tags( image: &ImageName, extra_tag: Option<&str>, @@ -493,26 +435,7 @@ fn insert_alias_tag( Ok(()) } -fn node_image_labels<'a>( - tag: &'a str, - source_hash: &'a str, - worker_hash: &'a str, - base_hash: &'a str, -) -> Vec<(&'static str, &'a str)> { - vec![ - (NODE_IMAGE_TAG_LABEL, tag), - (NODE_IMAGE_SOURCE_HASH_LABEL, source_hash), - (NODE_IMAGE_WORKER_HASH_LABEL, worker_hash), - (NODE_IMAGE_BASE_HASH_LABEL, base_hash), - ] -} - -fn base_image_labels<'a>(base_hash: &'a str) -> Vec<(&'static str, &'a str)> { - vec![(BASE_IMAGE_SOURCE_HASH_LABEL, base_hash)] -} - fn ensure_aliases_local( - runner: &mut dyn ImageCommandRunner, progress: &mut Option<&mut dyn NodeImageProgressSink>, root: &Path, source_ref: &str, @@ -522,7 +445,6 @@ fn ensure_aliases_local( for alias in alias_refs(image, alias_tags) { if alias != source_ref { run_status( - runner, progress, root, "docker", @@ -536,7 +458,6 @@ fn ensure_aliases_local( } fn ensure_aliases_for_remote( - runner: &mut dyn ImageCommandRunner, progress: &mut Option<&mut dyn NodeImageProgressSink>, root: &Path, source_ref: &str, @@ -546,9 +467,8 @@ fn ensure_aliases_for_remote( if alias_tags.is_empty() { return Ok(false); } - if !runner.docker_image_exists(root, source_ref) { + if !docker_image_exists(root, source_ref) { run_status( - runner, progress, root, "docker", @@ -557,9 +477,9 @@ fn ensure_aliases_for_remote( Some(source_ref), )?; } - ensure_aliases_local(runner, progress, root, source_ref, image, alias_tags)?; + ensure_aliases_local(progress, root, source_ref, image, alias_tags)?; for alias in alias_refs(image, alias_tags) { - push_image(runner, progress, root, &alias)?; + push_image(progress, root, &alias)?; } Ok(true) } @@ -572,13 +492,11 @@ fn alias_refs(image: &ImageName, alias_tags: &BTreeSet) -> Vec { } fn push_image( - runner: &mut dyn ImageCommandRunner, progress: &mut Option<&mut dyn NodeImageProgressSink>, root: &Path, image_ref: &str, ) -> Result<(), String> { run_status( - runner, progress, root, "docker", @@ -588,25 +506,12 @@ fn push_image( ) } -fn docker_image_exists(root: &Path, image_ref: &str) -> bool { - Command::new("docker") - .current_dir(root) - .args(["image", "inspect", image_ref]) - .stdin(Stdio::null()) - .stdout(Stdio::null()) - .stderr(Stdio::null()) - .status() - .map(|status| status.success()) - .unwrap_or(false) -} - fn docker_image_labels_match( - runner: &mut dyn ImageCommandRunner, root: &Path, image_ref: &str, expected: &[(&str, &str)], ) -> Result { - let Some(labels) = runner.docker_image_labels(root, image_ref)? else { + let Some(labels) = docker_image_labels(root, image_ref)? else { return Ok(false); }; Ok(expected @@ -614,54 +519,18 @@ fn docker_image_labels_match( .all(|(key, value)| labels.get(*key).map(String::as_str) == Some(*value))) } -fn docker_image_labels( - root: &Path, - image_ref: &str, -) -> Result>, String> { - let output = Command::new("docker") - .current_dir(root) - .args([ - "image", - "inspect", - "--format", - "{{ json .Config.Labels }}", - image_ref, - ]) - .stdin(Stdio::null()) - .output() - .map_err(|e| format!("inspect docker image {image_ref}: {e}"))?; - if !output.status.success() { - return Ok(None); - } - let stdout = String::from_utf8_lossy(&output.stdout); - let labels: Option> = serde_json::from_str(stdout.trim()) - .map_err(|e| format!("parse docker labels for {image_ref}: {e}"))?; - Ok(Some(labels.unwrap_or_default())) -} - -fn docker_manifest_exists(root: &Path, image_ref: &str) -> bool { - Command::new("docker") - .current_dir(root) - .args(["manifest", "inspect", image_ref]) - .stdin(Stdio::null()) - .stdout(Stdio::null()) - .stderr(Stdio::null()) - .status() - .map(|status| status.success()) - .unwrap_or(false) -} - -fn prune_old_dirty_images( - runner: &mut dyn ImageCommandRunner, - root: &Path, - image: &ImageName, - keep_tag: &str, -) { - if !dirty_image_prune_enabled() { +fn prune_old_dirty_images(root: &Path, image: &ImageName, keep_tag: &str) { + let prune_enabled = std::env::var("MVP_NODE_IMAGE_PRUNE") + .map(|value| { + let value = value.trim().to_ascii_lowercase(); + !matches!(value.as_str(), "0" | "false" | "no" | "off") + }) + .unwrap_or(true); + if !prune_enabled { return; } - let tags = match runner.docker_image_tags(root, &image.repository) { + let tags = match docker_image_tags(root, &image.repository) { Ok(tags) => tags, Err(error) => { eprintln!("mvp-node-image: prune old dirty images skipped: {error}"); @@ -669,7 +538,10 @@ fn prune_old_dirty_images( } }; - let keep_old = dirty_image_prune_keep(); + let keep_old = std::env::var("MVP_NODE_IMAGE_PRUNE_KEEP") + .ok() + .and_then(|value| value.trim().parse::().ok()) + .unwrap_or(3); let mut retained_old = 0_usize; for (repository, tag) in tags { if repository != image.repository @@ -681,7 +553,7 @@ fn prune_old_dirty_images( } let image_ref = image.ref_for_tag(&tag); - let Ok(Some(labels)) = runner.docker_image_labels(root, &image_ref) else { + let Ok(Some(labels)) = docker_image_labels(root, &image_ref) else { continue; }; if labels.get(NODE_IMAGE_TAG_LABEL).map(String::as_str) != Some(tag.as_str()) @@ -691,7 +563,7 @@ fn prune_old_dirty_images( { continue; } - if runner.docker_image_has_container(root, &image_ref) { + if docker_image_has_container(root, &image_ref) { eprintln!( "mvp-node-image: prune old dirty image {image_ref} skipped: container exists" ); @@ -703,47 +575,13 @@ fn prune_old_dirty_images( } eprintln!("mvp-node-image: prune old dirty image {image_ref}"); - if let Err(error) = runner.docker_image_remove(root, &image_ref) { + if let Err(error) = docker_image_remove(root, &image_ref) { eprintln!("mvp-node-image: prune old dirty image {image_ref} skipped: {error}"); } } } -fn dirty_image_prune_enabled() -> bool { - std::env::var(NODE_IMAGE_PRUNE_ENV) - .map(|value| { - let value = value.trim().to_ascii_lowercase(); - !matches!(value.as_str(), "0" | "false" | "no" | "off") - }) - .unwrap_or(true) -} - -fn dirty_image_prune_keep() -> usize { - std::env::var(NODE_IMAGE_PRUNE_KEEP_ENV) - .ok() - .and_then(|value| value.trim().parse::().ok()) - .unwrap_or(DEFAULT_DIRTY_IMAGE_KEEP) -} - -fn docker_image_has_container(root: &Path, image_ref: &str) -> bool { - Command::new("docker") - .current_dir(root) - .args([ - "ps", - "-a", - "--filter", - &format!("ancestor={image_ref}"), - "--format", - "{{.ID}}", - ]) - .stdin(Stdio::null()) - .output() - .map(|output| output.status.success() && !output.stdout.is_empty()) - .unwrap_or(true) -} - fn run_status( - runner: &mut dyn ImageCommandRunner, progress: &mut Option<&mut dyn NodeImageProgressSink>, root: &Path, program: &str, @@ -752,19 +590,7 @@ fn run_status( image_ref: Option<&str>, ) -> Result<(), String> { let args = args.iter().map(|arg| (*arg).to_owned()).collect::>(); - runner.run_status(root, program, &args, label, image_ref, progress) -} - -fn run_status_vec( - runner: &mut dyn ImageCommandRunner, - progress: &mut Option<&mut dyn NodeImageProgressSink>, - root: &Path, - program: &str, - args: Vec, - label: &str, - image_ref: Option<&str>, -) -> Result<(), String> { - runner.run_status(root, program, &args, label, image_ref, progress) + run_status_command(root, program, &args, label, image_ref, progress) } fn emit_image_reference( @@ -857,205 +683,239 @@ fn drain_command_lines( } } -impl ImageCommandRunner for RealImageCommandRunner { - fn run_status( - &mut self, - root: &Path, - program: &str, - args: &[String], - label: &str, - image_ref: Option<&str>, - progress: &mut Option<&mut dyn NodeImageProgressSink>, - ) -> Result<(), String> { - eprintln!("mvp-node-image: {label}"); - if progress.is_none() { - let status = Command::new(program) - .current_dir(root) - .args(args) - .stdin(Stdio::null()) - .stdout(Stdio::inherit()) - .stderr(Stdio::inherit()) - .status() - .map_err(|e| format!("run {label}: {e}"))?; - return if status.success() { - Ok(()) - } else { - Err(format!("{label} failed with {status}")) - }; - } - - let started = Instant::now(); - emit_command_progress( - progress, - label, - image_ref, - 0, - NodeImageProgressEventKind::CommandStarted { - program: program.to_owned(), - args: args.to_vec(), - }, - ); - let mut child = match Command::new(program) +fn run_status_command( + root: &Path, + program: &str, + args: &[String], + label: &str, + image_ref: Option<&str>, + progress: &mut Option<&mut dyn NodeImageProgressSink>, +) -> Result<(), String> { + eprintln!("mvp-node-image: {label}"); + if progress.is_none() { + let status = Command::new(program) .current_dir(root) .args(args) .stdin(Stdio::null()) - .stdout(Stdio::piped()) - .stderr(Stdio::piped()) - .spawn() - { - Ok(child) => child, + .stdout(Stdio::inherit()) + .stderr(Stdio::inherit()) + .status() + .map_err(|e| format!("run {label}: {e}"))?; + return if status.success() { + Ok(()) + } else { + Err(format!("{label} failed with {status}")) + }; + } + + let started = Instant::now(); + emit_command_progress( + progress, + label, + image_ref, + 0, + NodeImageProgressEventKind::CommandStarted { + program: program.to_owned(), + args: args.to_vec(), + }, + ); + let mut child = match Command::new(program) + .current_dir(root) + .args(args) + .stdin(Stdio::null()) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .spawn() + { + Ok(child) => child, + Err(error) => { + emit_command_progress( + progress, + label, + image_ref, + started.elapsed().as_millis(), + NodeImageProgressEventKind::CommandExited { + status: format!("spawn error: {error}"), + code: None, + success: false, + }, + ); + return Err(format!("run {label}: {error}")); + } + }; + + let (tx, rx) = mpsc::channel(); + let mut readers = Vec::new(); + if let Some(stdout) = child.stdout.take() { + readers.push(spawn_line_reader( + stdout, + CommandOutputLine::Stdout, + tx.clone(), + )); + } + if let Some(stderr) = child.stderr.take() { + readers.push(spawn_line_reader( + stderr, + CommandOutputLine::Stderr, + tx.clone(), + )); + } + drop(tx); + + let status = loop { + match child.try_wait() { + Ok(Some(status)) => break status, + Ok(None) => { + drain_command_lines(&rx, progress, label, image_ref, started); + thread::sleep(Duration::from_millis(10)); + } Err(error) => { + drain_command_lines(&rx, progress, label, image_ref, started); emit_command_progress( progress, label, image_ref, started.elapsed().as_millis(), NodeImageProgressEventKind::CommandExited { - status: format!("spawn error: {error}"), + status: format!("wait error: {error}"), code: None, success: false, }, ); return Err(format!("run {label}: {error}")); } - }; - - let (tx, rx) = mpsc::channel(); - let mut readers = Vec::new(); - if let Some(stdout) = child.stdout.take() { - readers.push(spawn_line_reader( - stdout, - CommandOutputLine::Stdout, - tx.clone(), - )); - } - if let Some(stderr) = child.stderr.take() { - readers.push(spawn_line_reader( - stderr, - CommandOutputLine::Stderr, - tx.clone(), - )); - } - drop(tx); - - let status = loop { - match child.try_wait() { - Ok(Some(status)) => break status, - Ok(None) => { - drain_command_lines(&rx, progress, label, image_ref, started); - thread::sleep(Duration::from_millis(10)); - } - Err(error) => { - drain_command_lines(&rx, progress, label, image_ref, started); - emit_command_progress( - progress, - label, - image_ref, - started.elapsed().as_millis(), - NodeImageProgressEventKind::CommandExited { - status: format!("wait error: {error}"), - code: None, - success: false, - }, - ); - return Err(format!("run {label}: {error}")); - } - } - }; - for reader in readers { - let _ = reader.join(); - } - drain_command_lines(&rx, progress, label, image_ref, started); - let status_text = status.to_string(); - let success = status.success(); - emit_command_progress( - progress, - label, - image_ref, - started.elapsed().as_millis(), - NodeImageProgressEventKind::CommandExited { - status: status_text.clone(), - code: status.code(), - success, - }, - ); - if success { - Ok(()) - } else { - Err(format!("{label} failed with {status_text}")) } + }; + for reader in readers { + let _ = reader.join(); } - - fn docker_image_exists(&mut self, root: &Path, image_ref: &str) -> bool { - docker_image_exists(root, image_ref) - } - - fn docker_image_labels( - &mut self, - root: &Path, - image_ref: &str, - ) -> Result>, String> { - docker_image_labels(root, image_ref) - } - - fn docker_manifest_exists(&mut self, root: &Path, image_ref: &str) -> bool { - docker_manifest_exists(root, image_ref) - } - - fn docker_image_has_container(&mut self, root: &Path, image_ref: &str) -> bool { - docker_image_has_container(root, image_ref) - } - - fn docker_image_tags( - &mut self, - root: &Path, - repository: &str, - ) -> Result, String> { - let output = Command::new("docker") - .current_dir(root) - .args([ - "image", - "ls", - "--format", - "{{.Repository}}\t{{.Tag}}", - repository, - ]) - .stdin(Stdio::null()) - .output() - .map_err(|error| format!("docker image ls failed: {error}"))?; - if !output.status.success() { - return Err(format!("docker image ls failed with {}", output.status)); - } - let stdout = String::from_utf8_lossy(&output.stdout); - Ok(stdout - .lines() - .filter_map(|line| { - let (repository, tag) = line.split_once('\t')?; - Some((repository.to_owned(), tag.to_owned())) - }) - .collect()) - } - - fn docker_image_remove(&mut self, root: &Path, image_ref: &str) -> Result<(), String> { - let status = Command::new("docker") - .current_dir(root) - .args(["image", "rm", image_ref]) - .stdin(Stdio::null()) - .stdout(Stdio::null()) - .stderr(Stdio::null()) - .status() - .map_err(|error| format!("docker image rm failed: {error}"))?; - if status.success() { - Ok(()) - } else { - Err(format!("docker image rm failed with {status}")) - } + drain_command_lines(&rx, progress, label, image_ref, started); + let status_text = status.to_string(); + let success = status.success(); + emit_command_progress( + progress, + label, + image_ref, + started.elapsed().as_millis(), + NodeImageProgressEventKind::CommandExited { + status: status_text.clone(), + code: status.code(), + success, + }, + ); + if success { + Ok(()) + } else { + Err(format!("{label} failed with {status_text}")) } } -fn looks_registry_reachable(repository: &str) -> bool { - let first = repository.split('/').next().unwrap_or(repository); - repository.contains('/') || first.contains('.') || first.contains(':') || first == "localhost" +fn docker_image_exists(root: &Path, image_ref: &str) -> bool { + Command::new("docker") + .current_dir(root) + .args(["image", "inspect", image_ref]) + .stdin(Stdio::null()) + .stdout(Stdio::null()) + .stderr(Stdio::null()) + .status() + .map(|status| status.success()) + .unwrap_or(false) +} + +fn docker_image_labels( + root: &Path, + image_ref: &str, +) -> Result>, String> { + let output = Command::new("docker") + .current_dir(root) + .args([ + "image", + "inspect", + "--format", + "{{ json .Config.Labels }}", + image_ref, + ]) + .stdin(Stdio::null()) + .output() + .map_err(|e| format!("inspect docker image {image_ref}: {e}"))?; + if !output.status.success() { + return Ok(None); + } + let stdout = String::from_utf8_lossy(&output.stdout); + let labels: Option> = serde_json::from_str(stdout.trim()) + .map_err(|e| format!("parse docker labels for {image_ref}: {e}"))?; + Ok(Some(labels.unwrap_or_default())) +} + +fn docker_manifest_exists(root: &Path, image_ref: &str) -> bool { + Command::new("docker") + .current_dir(root) + .args(["manifest", "inspect", image_ref]) + .stdin(Stdio::null()) + .stdout(Stdio::null()) + .stderr(Stdio::null()) + .status() + .map(|status| status.success()) + .unwrap_or(false) +} + +fn docker_image_has_container(root: &Path, image_ref: &str) -> bool { + Command::new("docker") + .current_dir(root) + .args([ + "ps", + "-a", + "--filter", + &format!("ancestor={image_ref}"), + "--format", + "{{.ID}}", + ]) + .stdin(Stdio::null()) + .output() + .map(|output| output.status.success() && !output.stdout.is_empty()) + .unwrap_or(true) +} + +fn docker_image_tags(root: &Path, repository: &str) -> Result, String> { + let output = Command::new("docker") + .current_dir(root) + .args([ + "image", + "ls", + "--format", + "{{.Repository}}\t{{.Tag}}", + repository, + ]) + .stdin(Stdio::null()) + .output() + .map_err(|error| format!("docker image ls failed: {error}"))?; + if !output.status.success() { + return Err(format!("docker image ls failed with {}", output.status)); + } + let stdout = String::from_utf8_lossy(&output.stdout); + Ok(stdout + .lines() + .filter_map(|line| { + let (repository, tag) = line.split_once('\t')?; + Some((repository.to_owned(), tag.to_owned())) + }) + .collect()) +} + +fn docker_image_remove(root: &Path, image_ref: &str) -> Result<(), String> { + let status = Command::new("docker") + .current_dir(root) + .args(["image", "rm", image_ref]) + .stdin(Stdio::null()) + .stdout(Stdio::null()) + .stderr(Stdio::null()) + .status() + .map_err(|error| format!("docker image rm failed: {error}"))?; + if status.success() { + Ok(()) + } else { + Err(format!("docker image rm failed with {status}")) + } } #[derive(Clone, Debug)] diff --git a/crates/mvp-system/src/chat/runtime.rs b/crates/mvp-system/src/chat/runtime.rs index 111eb60..0de3845 100644 --- a/crates/mvp-system/src/chat/runtime.rs +++ b/crates/mvp-system/src/chat/runtime.rs @@ -22,10 +22,9 @@ use signal_hook::consts::signal::{SIGINT, SIGTERM}; #[cfg(target_os = "linux")] use signal_hook::iterator::Signals; -use crate::chat::config as chat_config; use crate::chat::node_image::{ - NodeImageProgressEvent, NodeImageProgressEventKind, NodeImageProgressSink, NodeImageProvider, - NodeImageRequest, prepare_node_image_with_progress, + NodeImageProgressEvent, NodeImageProgressEventKind, NodeImageProgressSink, NodeImageRequest, + prepare_node_image_with_progress, }; use crate::node_provisioning::{ProviderKind, provider_kind}; use crate::observability::{benchmark, frame_archive::FrameArchive}; @@ -37,6 +36,7 @@ use crate::{ DEFAULT_PIPELINE_CACHED_MODEL_MAX_CONTEXT, DEFAULT_PIPELINE_CACHED_MODEL_REPO, }; +const DEFAULT_CONFIG_PATH: &str = ".config/config.toml"; const DEFAULT_RPC_ADDR: &str = "127.0.0.1:19777"; const BASE_NODE_IMAGE: &str = "swactor-mvp-node-base:cuda12.6"; const REPO_MODEL_CACHE_DIR: &str = ".model-cache"; @@ -79,11 +79,11 @@ enum PromptInput { static STOP_REQUESTED: AtomicBool = AtomicBool::new(false); static PROMPT_STOP_TX: Mutex>> = Mutex::new(None); -pub(super) fn run_from_args(args: I) -> ExitCode +pub(crate) fn run_from_args(args: I) -> ExitCode where I: IntoIterator, { - match run_from_args_result(args) { + match install_signal_handlers().and_then(|()| run(args)) { Ok(()) => ExitCode::SUCCESS, Err(error) => { eprintln!("mvp-chat: {error}"); @@ -92,23 +92,6 @@ where } } -fn run_from_args_result(args: I) -> Result<(), String> -where - I: IntoIterator, -{ - install_signal_handlers()?; - run(args) -} - -fn print_usage() { - println!("{MVP_CHAT_USAGE}"); -} - -fn is_help_request(args: &[String]) -> bool { - args.iter() - .any(|arg| matches!(arg.as_str(), "--help" | "-h" | "help")) -} - struct RuntimeEnvGuard { name: &'static str, original: Option, @@ -142,8 +125,11 @@ where I: IntoIterator, { let provided_args = args.into_iter().collect::>(); - if is_help_request(&provided_args) { - print_usage(); + if provided_args + .iter() + .any(|arg| matches!(arg.as_str(), "--help" | "-h" | "help")) + { + println!("{MVP_CHAT_USAGE}"); return Ok(()); } let config = Config::from_args(provided_args)?; @@ -164,7 +150,8 @@ where ); progress.emit_benchmark_envelope(&config); progress.emit_endpoint_config_snapshot(&config); - confirm_vastai_if_needed(&config)?; + let mut approval = StdinVastAiApproval; + confirm_vastai_if_needed_with_approval(&config, &mut approval)?; let prepare_runtime_started = Instant::now(); progress.emit( CHAT_RUNTIME_CHANNEL, @@ -174,7 +161,7 @@ where ); let image_ref = match prepare_runtime_with_progress( &config, - prepare_node_image_progress_adapter, + prepare_node_image_with_progress, Some(&mut progress), ) { Ok(image_ref) => { @@ -318,7 +305,36 @@ where return Err(error); } }; - let result = run_chat_loop_with_progress(&rpc_addr, config.max_tokens, Some(&mut progress)); + let (prompt_tx, input_rx) = mpsc::channel(); + if STOP_REQUESTED.load(Ordering::SeqCst) { + let _ = prompt_tx.send(PromptInput::StopRequested); + } + if let Ok(mut stop_tx) = PROMPT_STOP_TX.lock() { + *stop_tx = Some(prompt_tx.clone()); + } + thread::spawn(move || { + let stdin = io::stdin(); + for line in stdin.lock().lines() { + match line { + Ok(line) => { + if prompt_tx.send(PromptInput::Line(line)).is_err() { + return; + } + } + Err(_) => { + let _ = prompt_tx.send(PromptInput::Closed); + return; + } + } + } + let _ = prompt_tx.send(PromptInput::Closed); + }); + let result = run_chat_loop_with_input_and_progress( + &rpc_addr, + config.max_tokens, + input_rx, + Some(&mut progress), + ); progress.emit( CHAT_LIFECYCLE_CHANNEL, "shutdown", @@ -739,13 +755,8 @@ struct ChatModelConfig { max_context: Option, } -#[derive(Clone, Debug)] -struct LoadedChatTomlConfig { - overlay: ChatTomlConfig, -} - -fn load_chat_config(path: Option<&Path>) -> Result { - let overlay = match path { +fn load_chat_config(path: Option<&Path>) -> Result { + Ok(match path { Some(path) => { let text = fs::read_to_string(path) .map_err(|e| format!("read config {}: {e}", path.display()))?; @@ -753,7 +764,7 @@ fn load_chat_config(path: Option<&Path>) -> Result .map_err(|e| format!("parse config {}: {e}", path.display()))? } None => { - let default = Path::new(chat_config::DEFAULT_CONFIG_PATH); + let default = Path::new(DEFAULT_CONFIG_PATH); if !default.is_file() { ChatTomlConfig::default() } else { @@ -763,8 +774,7 @@ fn load_chat_config(path: Option<&Path>) -> Result .map_err(|e| format!("parse config {}: {e}", default.display()))? } } - }; - Ok(LoadedChatTomlConfig { overlay }) + }) } impl Config { @@ -773,8 +783,7 @@ impl Config { I: IntoIterator, { let args = ParsedArgs::parse(provided_args)?; - let loaded = load_chat_config(args.config_path.as_deref())?; - let toml = loaded.overlay; + let toml = load_chat_config(args.config_path.as_deref())?; let provider = provider_from_sources(args.provider.clone(), toml.provider.kind.as_deref())?; let node_image = first_non_empty([toml.image.node.clone()]).unwrap_or_default(); if provider != provider_kind::process() && node_image.is_empty() { @@ -797,8 +806,8 @@ impl Config { }; Ok(Self { - orch_bin: default_orch_bin()?, - worker_bin: node_bin_for_current_profile()?, + orch_bin: artifact_root().join("target/debug/mvp-orchestrator"), + worker_bin: default_worker_bin(), rpc_addr: DEFAULT_RPC_ADDR.to_owned(), node_image, provider, @@ -1248,12 +1257,11 @@ fn resolve_vastai_config( } fn first_non_empty(values: [Option; N]) -> Option { - values.into_iter().find_map(chat_config::normalize_optional) -} - -fn confirm_vastai_if_needed(config: &Config) -> Result<(), String> { - let mut approval = StdinVastAiApproval; - confirm_vastai_if_needed_with_approval(config, &mut approval) + values + .into_iter() + .flatten() + .map(|value| value.trim().to_owned()) + .find(|value| !value.is_empty()) } trait VastAiApproval { @@ -1322,11 +1330,10 @@ where input .read_line(&mut line) .map_err(|e| format!("read Vast.ai approval: {e}"))?; - Ok(parse_approval(&line)) -} - -fn parse_approval(input: &str) -> bool { - matches!(input.trim().to_ascii_lowercase().as_str(), "y" | "yes") + Ok(matches!( + line.trim().to_ascii_lowercase().as_str(), + "y" | "yes" + )) } enum OrchHandle { @@ -1368,8 +1375,9 @@ impl InProcessOrch { fn spawn(config: &Config, image_ref: &str) -> Result { let args = config.orchestrator_cli_args(image_ref); let (stop_tx, stop_rx) = mpsc::channel(); - let thread = - thread::spawn(move || crate::run_orchestrator_in_process_from_args(args, stop_rx)); + let thread = thread::spawn(move || { + crate::orchestration::app::run_with_options(args, false, Some(stop_rx)) + }); Ok(Self { stop_tx: Some(stop_tx), thread: Some(thread), @@ -1397,9 +1405,12 @@ impl InProcessOrch { Err(error) => return Err(format!("connect prompt RPC {rpc_addr}: {error}")), } if let Some(result) = self.take_finished_result() { + let reason = match result { + Ok(()) => "completed successfully".to_owned(), + Err(error) => error, + }; return Err(format!( - "in-process orchestrator exited before prompt RPC ready: {}", - render_orch_thread_result(result) + "in-process orchestrator exited before prompt RPC ready: {reason}" )); } thread::sleep(Duration::from_millis(100)); @@ -1445,21 +1456,6 @@ impl Drop for InProcessOrch { } } -fn render_orch_thread_result(result: Result<(), String>) -> String { - match result { - Ok(()) => "completed successfully".to_owned(), - Err(error) => error, - } -} - -fn orchestrator_shutdown_grace(provider: &ProviderKind) -> Duration { - if provider == &provider_kind::vastai() { - Duration::from_millis(VASTAI_ORCH_SHUTDOWN_GRACE_MS) - } else { - Duration::from_millis(ORCH_SHUTDOWN_GRACE_MS) - } -} - struct OrchChild { child: Child, cleaned: bool, @@ -1493,7 +1489,11 @@ impl OrchChild { Ok(Self { child, cleaned: false, - shutdown_grace: orchestrator_shutdown_grace(&config.provider), + shutdown_grace: if config.provider == provider_kind::vastai() { + Duration::from_millis(VASTAI_ORCH_SHUTDOWN_GRACE_MS) + } else { + Duration::from_millis(ORCH_SHUTDOWN_GRACE_MS) + }, }) } @@ -1586,33 +1586,6 @@ fn signal_orch_process_group(child: &Child, signal: libc::c_int) -> io::Result<( } } -fn prepare_node_image_progress_adapter( - request: NodeImageRequest, - progress: Option<&mut dyn NodeImageProgressSink>, -) -> Result { - prepare_node_image_with_progress(request, progress) -} - -#[allow(dead_code)] -fn prepare_runtime(config: &Config) -> Result { - prepare_runtime_with(config, |request| { - prepare_node_image_with_progress(request, None) - }) -} - -#[allow(dead_code)] -fn prepare_runtime_with(config: &Config, prepare_node_image_fn: F) -> Result -where - F: FnMut(NodeImageRequest) -> Result, -{ - let mut prepare_node_image_fn = prepare_node_image_fn; - prepare_runtime_with_progress( - config, - move |request, _progress| prepare_node_image_fn(request), - None, - ) -} - fn prepare_runtime_with_progress( config: &Config, mut prepare_node_image_fn: F, @@ -1651,7 +1624,21 @@ where "started", json!({"mode": binary_mode, "command_label": "ensure_orch_binary"}), ); - match ensure_orch_binary(config) { + match ensure_runtime_binary( + config.skip_rebuild, + &config.orch_bin, + "mvp-orchestrator", + &[ + "build", + "--quiet", + "-p", + "mvp-system", + "--features", + "dashboard", + "--bin", + "mvp-orchestrator", + ], + ) { Ok(()) => emit_chat_progress( &mut progress, CHAT_RUNTIME_CHANNEL, @@ -1680,7 +1667,19 @@ where "started", json!({"mode": binary_mode}), ); - match ensure_worker_binary(config) { + match ensure_runtime_binary( + config.skip_rebuild, + &config.worker_bin, + "mvp-worker-node", + &[ + "build", + "--quiet", + "-p", + "mvp-system", + "--bin", + "mvp-worker-node", + ], + ) { Ok(()) => emit_chat_progress( &mut progress, CHAT_RUNTIME_CHANNEL, @@ -1735,7 +1734,19 @@ where "started", json!({"mode": binary_mode, "command_label": "ensure_worker_binary"}), ); - match ensure_worker_binary(config) { + match ensure_runtime_binary( + config.skip_rebuild, + &config.worker_bin, + "mvp-worker-node", + &[ + "build", + "--quiet", + "-p", + "mvp-system", + "--bin", + "mvp-worker-node", + ], + ) { Ok(()) => emit_chat_progress( &mut progress, CHAT_RUNTIME_CHANNEL, @@ -1772,31 +1783,25 @@ where "started", json!({"provider": config.provider.as_str(), "command_label": "prepare_node_image", "image_tag": config.image_tag.as_deref()}), ); - let node_bin = match node_bin_for_current_profile() { - Ok(path) => path, - Err(error) => { - emit_chat_progress( - &mut progress, - CHAT_RUNTIME_CHANNEL, - "prepare_node_image", - "failed", - json!({"provider": config.provider.as_str(), "command_label": "prepare_node_image", "elapsed_ms": prepare_node_image_started.elapsed().as_millis(), "error": error.as_str()}), - ); - return Err(error); - } - }; - let provider = match node_image_provider(&config.provider) { - Ok(provider) => provider, - Err(error) => { - emit_chat_progress( - &mut progress, - CHAT_RUNTIME_CHANNEL, - "prepare_node_image", - "failed", - json!({"provider": config.provider.as_str(), "command_label": "prepare_node_image", "elapsed_ms": prepare_node_image_started.elapsed().as_millis(), "error": error.as_str()}), - ); - return Err(error); - } + 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 prepared = { let command_progress = progress @@ -1807,9 +1812,8 @@ where requested_image: config.node_image.clone(), base_image: BASE_NODE_IMAGE.to_owned(), node_bin, - provider, + requires_registry_image, extra_tag: config.image_tag.clone(), - push: false, force_refresh: false, enabled: true, }, @@ -1838,42 +1842,6 @@ where Ok(prepared) } -fn stdin_prompt_events() -> mpsc::Receiver { - let (tx, rx) = mpsc::channel(); - if STOP_REQUESTED.load(Ordering::SeqCst) { - let _ = tx.send(PromptInput::StopRequested); - } - if let Ok(mut stop_tx) = PROMPT_STOP_TX.lock() { - *stop_tx = Some(tx.clone()); - } - thread::spawn(move || { - let stdin = io::stdin(); - for line in stdin.lock().lines() { - match line { - Ok(line) => { - if tx.send(PromptInput::Line(line)).is_err() { - return; - } - } - Err(_) => { - let _ = tx.send(PromptInput::Closed); - return; - } - } - } - let _ = tx.send(PromptInput::Closed); - }); - rx -} - -fn run_chat_loop_with_progress( - addr: &str, - max_tokens: u32, - progress: Option<&mut ChatDatastream>, -) -> Result<(), String> { - run_chat_loop_with_input_and_progress(addr, max_tokens, stdin_prompt_events(), progress) -} - fn run_chat_loop_with_input_and_progress( addr: &str, max_tokens: u32, @@ -1923,7 +1891,15 @@ fn run_chat_loop_with_input_and_progress( return Err(format!("clone prompt RPC stream: {error}")); } }; - run_chat_session_with_progress(&mut stream, reader, input_rx, max_tokens, progress) + let mut output = io::stdout(); + run_chat_session_with_output_and_progress( + &mut stream, + reader, + input_rx, + max_tokens, + &mut output, + progress, + ) } #[cfg(test)] @@ -1937,24 +1913,6 @@ fn run_chat_session_with_output( run_chat_session_with_output_and_progress(writer, reader, input_rx, max_tokens, output, None) } -fn run_chat_session_with_progress( - writer: &mut impl Write, - reader: impl BufRead, - input_rx: mpsc::Receiver, - max_tokens: u32, - progress: Option<&mut ChatDatastream>, -) -> Result<(), String> { - let mut output = io::stdout(); - run_chat_session_with_output_and_progress( - writer, - reader, - input_rx, - max_tokens, - &mut output, - progress, - ) -} - fn emit_chat_progress( progress: &mut Option<&mut ChatDatastream>, channel: &str, @@ -2186,84 +2144,39 @@ fn run_chat_session_with_output_and_progress( } } -fn default_orch_bin() -> Result { - Ok(artifact_root().join("target/debug/mvp-orchestrator")) +fn default_worker_bin() -> PathBuf { + artifact_root().join("target/debug/mvp-worker-node") } -fn node_bin_for_current_profile() -> Result { - Ok(artifact_root().join("target/debug/mvp-worker-node")) -} - -fn cargo_command() -> &'static str { - "cargo" -} - -fn mvp_orchestrator_build_args() -> &'static [&'static str] { - &[ - "build", - "--quiet", - "-p", - "mvp-system", - "--features", - "dashboard", - "--bin", - "mvp-orchestrator", - ] -} - -fn ensure_orch_binary(config: &Config) -> Result<(), String> { - if config.skip_rebuild { - return ensure_existing_artifact(&config.orch_bin, "mvp-orchestrator"); +fn ensure_runtime_binary( + skip_rebuild: bool, + path: &Path, + label: &str, + cargo_args: &[&str], +) -> Result<(), String> { + if skip_rebuild { + let metadata = fs::metadata(path) + .map_err(|e| format!("missing required {label} artifact {}: {e}", path.display()))?; + if !metadata.is_file() { + return Err(format!( + "missing required {label} artifact {}; not a file", + path.display() + )); + } + return Ok(()); } - run_status( - cargo_command(), - mvp_orchestrator_build_args(), - "build mvp-orchestrator", - ) -} -fn ensure_worker_binary(config: &Config) -> Result<(), String> { - if config.skip_rebuild { - return ensure_existing_artifact(&config.worker_bin, "mvp-worker-node"); - } - run_status( - cargo_command(), - &[ - "build", - "--quiet", - "-p", - "mvp-system", - "--bin", - "mvp-worker-node", - ], - "build mvp-worker-node", - ) -} - -fn ensure_existing_artifact(path: &PathBuf, label: &str) -> Result<(), String> { - let metadata = fs::metadata(path) - .map_err(|e| format!("missing required {label} artifact {}: {e}", path.display()))?; - if !metadata.is_file() { - return Err(format!( - "missing required {label} artifact {}; not a file", - path.display() - )); - } - Ok(()) -} - -fn run_status(program: &str, args: &[&str], label: &str) -> Result<(), String> { - let status = Command::new(program) - .args(args) + let status = Command::new("cargo") + .args(cargo_args) .stdin(Stdio::null()) .stdout(Stdio::inherit()) .stderr(Stdio::inherit()) .status() - .map_err(|e| format!("run {label}: {e}"))?; + .map_err(|e| format!("run build {label}: {e}"))?; if status.success() { Ok(()) } else { - Err(format!("{label} failed with {status}")) + Err(format!("build {label} failed with {status}")) } } @@ -2393,18 +2306,6 @@ fn env_optional(name: &str) -> Option { .filter(|value| !value.is_empty()) } -fn node_image_provider(provider: &ProviderKind) -> Result { - if provider == &provider_kind::docker() { - Ok(NodeImageProvider::Docker) - } else if provider == &provider_kind::vastai() { - Ok(NodeImageProvider::VastAi) - } else if provider == &provider_kind::process() { - Err("process provider does not use node images".to_owned()) - } else { - Err("mvp-chat does not support mock provider".to_owned()) - } -} - fn next_arg(args: &mut impl Iterator, name: &str) -> Result { args.next() .ok_or_else(|| format!("missing value after {name}")) diff --git a/crates/mvp-system/src/lib.rs b/crates/mvp-system/src/lib.rs index f35d962..0908c8d 100644 --- a/crates/mvp-system/src/lib.rs +++ b/crates/mvp-system/src/lib.rs @@ -12,28 +12,18 @@ pub fn run_chat_from_args(args: I) -> std::process::ExitCode where I: IntoIterator, { - chat::run_from_args(args) + chat::runtime::run_from_args(args) } pub fn run_orchestrator_from_args(args: I) -> Result<(), String> where I: IntoIterator, { - orchestration::run_from_args(args) + orchestration::app::run_with_options(args, true, None) } pub fn run_worker_node_from_env() -> std::process::ExitCode { - node::run_worker_node_from_env() -} - -fn run_orchestrator_in_process_from_args( - args: I, - stop_rx: std::sync::mpsc::Receiver<()>, -) -> Result<(), String> -where - I: IntoIterator, -{ - orchestration::run_in_process_from_args(args, stop_rx) + node::worker_node_runtime::run_from_env() } #[path = "transport/driver_pumps.rs"] diff --git a/crates/mvp-system/src/node/mod.rs b/crates/mvp-system/src/node/mod.rs index d1e7bb2..743674b 100644 --- a/crates/mvp-system/src/node/mod.rs +++ b/crates/mvp-system/src/node/mod.rs @@ -3,8 +3,4 @@ //! Worker-node runtime behavior lives behind this module boundary; binaries //! only wire entrypoints into it. -mod worker_node_runtime; - -pub(super) fn run_worker_node_from_env() -> std::process::ExitCode { - worker_node_runtime::run_from_env() -} +pub(super) mod worker_node_runtime; diff --git a/crates/mvp-system/src/node/worker_node_runtime.rs b/crates/mvp-system/src/node/worker_node_runtime.rs index b7cdab7..ee1cbf6 100644 --- a/crates/mvp-system/src/node/worker_node_runtime.rs +++ b/crates/mvp-system/src/node/worker_node_runtime.rs @@ -26,7 +26,7 @@ use crate::driver_pumps as driver_model; use crate::gguf_shard::{StageShardPlan, materialize_stage_shard_http, validate_stage_shard_cache}; use crate::node_actor::{ NodeAgentActor, NodeAgentMsg, NodeAgentReport, StageCommandWire, StageInboundEdgeWire, - StageObjectSpecWire, StageOutboundEdgeWire, StageRingSpecWire, + StageObjectSpecWire, StageOutboundEdgeWire, }; use crate::observability::benchmark; use crate::orchestration::distribution_stack::DistributionRuntimeStack; @@ -908,8 +908,16 @@ impl WorkerEdgeRuntime { run_id: edge::RunId(config.run_id), edge_id: edge::EdgeId(edge.edge_id), local_node_id: edge::NodeId(config.logical_node_id), - object_spec: edge_object_spec(edge.object_spec), - ring_spec: edge_ring_spec(edge.ring_spec), + object_spec: edge::ObjectSpec { + kind: edge::ObjectKind::Activation, + dtype: edge::DType::F16, + max_extent_bytes: edge.object_spec.max_extent, + }, + ring_spec: edge::RingSpec { + header_bytes: 0, + data_bytes: edge.ring_spec.data_capacity, + alignment: u64::from(edge.ring_spec.alignment), + }, })); self.drive_edge_workflow( stack, @@ -954,8 +962,16 @@ impl WorkerEdgeRuntime { edge_id: edge::EdgeId(edge.edge_id), local_node_id: edge::NodeId(config.logical_node_id), consumer_node_id: edge::NodeId(edge.consumer_node_id), - object_spec: edge_object_spec(edge.object_spec), - ring_spec: edge_ring_spec(edge.ring_spec), + object_spec: edge::ObjectSpec { + kind: edge::ObjectKind::Activation, + dtype: edge::DType::F16, + max_extent_bytes: edge.object_spec.max_extent, + }, + ring_spec: edge::RingSpec { + header_bytes: 0, + data_bytes: edge.ring_spec.data_capacity, + alignment: u64::from(edge.ring_spec.alignment), + }, })); self.drive_edge_workflow( stack, @@ -1581,29 +1597,6 @@ impl WorkerEdgeRuntime { fn duration_ms_u64(duration: Duration) -> u64 { u64::try_from(duration.as_millis()).unwrap_or(u64::MAX) } -fn edge_object_spec(spec: StageObjectSpecWire) -> edge::ObjectSpec { - edge::ObjectSpec { - kind: edge::ObjectKind::Activation, - dtype: edge::DType::F16, - max_extent_bytes: spec.max_extent, - } -} - -fn edge_ring_spec(spec: StageRingSpecWire) -> edge::RingSpec { - edge::RingSpec { - header_bytes: 0, - data_bytes: spec.data_capacity, - alignment: u64::from(spec.alignment), - } -} - -fn ingress_object_spec(spec: StageObjectSpecWire) -> ingress::ObjectSpec { - ingress::ObjectSpec { - max_extent: spec.max_extent, - alignment: u64::from(spec.alignment), - layout: ingress::ObjectLayout::Token, - } -} struct IngressRecordBytes { bytes: Vec, @@ -1618,8 +1611,16 @@ fn take_complete_ingress_record( buffer: &mut Vec, spec: StageObjectSpecWire, ) -> Result, String> { - let record = match ingress::read_object_record(buffer, ingress_object_spec(spec), false) - .map_err(|reason| format!("invalid object record: {reason:?}"))? + let record = match ingress::read_object_record( + buffer, + ingress::ObjectSpec { + max_extent: spec.max_extent, + alignment: u64::from(spec.alignment), + layout: ingress::ObjectLayout::Token, + }, + false, + ) + .map_err(|reason| format!("invalid object record: {reason:?}"))? { ingress::ObjectRecordRead::Incomplete => return Ok(None), ingress::ObjectRecordRead::Complete(record) => record, @@ -1642,21 +1643,25 @@ fn value_u64(value: &Value, field: &str) -> Result { .ok_or_else(|| format!("helper event missing numeric {field}: {value}")) } -pub(super) fn run_from_env() -> ExitCode { - let mut args = std::env::args().skip(1).collect::>(); - if args.first().map(String::as_str) == Some("debug-join") { - args.remove(0); - return debug_join_client_main(args); - } - if args.first().map(String::as_str) == Some("stage-shard-fetcher") { - return stage_shard_fetcher_main(); - } - match run() { - Ok(()) => ExitCode::SUCCESS, - Err(error) => { - eprintln!("mvp-worker-node: {error}"); - ExitCode::from(1) - } +pub(crate) fn run_from_env() -> ExitCode { + let mut args = std::env::args().skip(1); + match args.next().as_deref() { + Some("debug-join") => debug_join_client_main(args.collect()), + Some("stage-shard-fetcher") => match run_stage_shard_fetcher() { + Ok(()) => ExitCode::SUCCESS, + Err(error) => { + let event = json!({"type":"StageShardFetchFailed","error":error}); + println!("{event}"); + ExitCode::from(1) + } + }, + _ => match run() { + Ok(()) => ExitCode::SUCCESS, + Err(error) => { + eprintln!("mvp-worker-node: {error}"); + ExitCode::from(1) + } + }, } } @@ -1666,17 +1671,6 @@ struct StageShardFetchRequest { output_path: PathBuf, } -fn stage_shard_fetcher_main() -> ExitCode { - match run_stage_shard_fetcher() { - Ok(()) => ExitCode::SUCCESS, - Err(error) => { - let event = json!({"type":"StageShardFetchFailed","error":error}); - println!("{event}"); - ExitCode::from(1) - } - } -} - fn run_stage_shard_fetcher() -> Result<(), String> { let mut input = String::new(); std::io::stdin() @@ -1858,7 +1852,7 @@ fn run() -> Result<(), String> { }; let arena_fd = arena_manager.lock().arena_fd(); - let mut datastream = node_datastream(&config); + let mut datastream = NodeDatastream::new(&config); let datastream_transport = driver.datastream_publish_handle(); let datastream_publisher = match stack .runtime @@ -2338,7 +2332,7 @@ fn emit_swim_telemetry( local_phase: &str, ) { for transition in stack.drain_swim_transitions() { - let peer = format_dist_node_id(transition.peer); + let peer = format!("{:?}", transition.peer); let from = transition.from.map(|state| format!("{:?}", state)); let to = format!("{:?}", transition.to); let member_state = stack @@ -2375,7 +2369,7 @@ fn swim_probe_event_record( let budget_ms = event.budget_ms; SwimProbeEvent { event: event.event.to_owned(), - target: format_dist_node_id(event.target), + target: format!("{:?}", event.target), sequence: event.sequence, kind: event.kind.to_owned(), rtt_ms: event.rtt_ms, @@ -2403,18 +2397,10 @@ fn swim_recent_probe_targets(stack: &DistributionRuntimeStack) -> Vec { .swim_telemetry .recent_targets() .into_iter() - .map(format_dist_node_id) + .map(|node_id| format!("{:?}", node_id)) .collect() } -fn format_dist_node_id(node_id: DistNodeId) -> String { - format!("{:?}", node_id) -} - -fn node_datastream(config: &DeploymentConfig) -> NodeDatastream { - NodeDatastream::new(config) -} - #[derive(Clone, Copy)] struct DatastreamChannelSet { node_ready: ChannelId, @@ -3989,8 +3975,18 @@ struct DeploymentConfig { impl DeploymentConfig { fn from_env() -> Result { - let run_id = env_u64("MVP_RUN_ID", 1)?; - let logical_node_id = env_u64("MVP_LOGICAL_NODE_ID", 1)?; + macro_rules! env_parse { + ($name:expr, $default:expr) => { + match env_optional($name) { + Some(value) => value + .parse() + .map_err(|e| format!("invalid {}={value:?}: {e}", $name)), + None => Ok($default), + } + }; + } + let run_id = env_parse!("MVP_RUN_ID", 1)?; + let logical_node_id = env_parse!("MVP_LOGICAL_NODE_ID", 1)?; let relay = relay_runtime_config_from_env(run_id)?; let debug_join_socket = match env_optional("MVP_DEBUG_JOIN_SOCKET").as_deref() { Some("disabled") => None, @@ -4004,7 +4000,7 @@ impl DeploymentConfig { .into_owned(), ), }; - let provider = env_string("MVP_NODE_PROVIDER", "process"); + let provider = env_optional("MVP_NODE_PROVIDER").unwrap_or_else(|| "process".to_owned()); let default_device = if provider == "process" { "CPU" } else { @@ -4013,9 +4009,19 @@ impl DeploymentConfig { Ok(Self { run_id, logical_node_id, - stage_index: env_u32("MVP_STAGE_INDEX", 0)?, - coordinator_endpoint: env_json("MVP_COORDINATOR_ENDPOINT")?, - orchestrator_actor: env_json("MVP_ORCHESTRATOR_ACTOR")?, + stage_index: env_parse!("MVP_STAGE_INDEX", 0)?, + coordinator_endpoint: env_optional("MVP_COORDINATOR_ENDPOINT") + .map(|value| { + serde_json::from_str::(&value) + .map_err(|e| format!("invalid MVP_COORDINATOR_ENDPOINT JSON: {e}")) + }) + .transpose()?, + orchestrator_actor: env_optional("MVP_ORCHESTRATOR_ACTOR") + .map(|value| { + serde_json::from_str::(&value) + .map_err(|e| format!("invalid MVP_ORCHESTRATOR_ACTOR JSON: {e}")) + }) + .transpose()?, datastream_frame_log: env_optional("MVP_DATASTREAM_FRAME_LOG"), debug_join_socket, relay_mode: relay.mode, @@ -4024,16 +4030,29 @@ impl DeploymentConfig { .map(EndpointAddrMask::parse) .transpose()? .unwrap_or_default(), - worker_script: env_string("MVP_TINYGRAD_WORKER", DEFAULT_WORKER_SCRIPT), - device: env_string("DEV", default_device), - model_id: env_string("MVP_MODEL_ID", DEFAULT_MODEL_ID), - gguf_source: gguf_source_from_env(), - tokenizer: tokenizer_from_env(), + worker_script: env_optional("MVP_TINYGRAD_WORKER") + .unwrap_or_else(|| DEFAULT_WORKER_SCRIPT.to_owned()), + device: env_optional("DEV").unwrap_or_else(|| default_device.to_owned()), + model_id: env_optional("MVP_MODEL_ID").unwrap_or_else(|| DEFAULT_MODEL_ID.to_owned()), + gguf_source: if let Some(path) = env_optional("MVP_GGUF_LOCAL_PATH") { + GgufSource::LocalPath(path) + } else { + GgufSource::HuggingFaceGguf { + repo: env_optional("MVP_GGUF_REPO") + .unwrap_or_else(|| DEFAULT_HF_REPO.to_owned()), + file: env_optional("MVP_GGUF_FILE") + .unwrap_or_else(|| DEFAULT_HF_FILE.to_owned()), + revision: env_optional("MVP_GGUF_REVISION"), + } + }, + tokenizer: env_optional("MVP_TOKENIZER_LOCAL_PATH") + .map(TokenizerSource::LocalPath) + .unwrap_or(TokenizerSource::EmbeddedGguf), self_test_prompt: env_optional("MVP_NODE_SELF_TEST_PROMPT"), - self_test_layer_end: env_u32("MVP_SELF_TEST_LAYER_END", 16)?, - self_test_max_tokens: env_u32("MVP_SELF_TEST_MAX_TOKENS", 1)?, - arena_bytes: env_u64("MVP_ARENA_BYTES", DEFAULT_ARENA_BYTES)?, - arena_alignment: env_u64("MVP_ARENA_ALIGNMENT", DEFAULT_ARENA_ALIGNMENT)?, + self_test_layer_end: env_parse!("MVP_SELF_TEST_LAYER_END", 16)?, + self_test_max_tokens: env_parse!("MVP_SELF_TEST_MAX_TOKENS", 1)?, + arena_bytes: env_parse!("MVP_ARENA_BYTES", DEFAULT_ARENA_BYTES)?, + arena_alignment: env_parse!("MVP_ARENA_ALIGNMENT", DEFAULT_ARENA_ALIGNMENT)?, }) } } @@ -4788,57 +4807,9 @@ fn spawn_stdin_shutdown_listener() -> Receiver<()> { rx } -fn env_string(name: &str, default: &str) -> String { - env_optional(name).unwrap_or_else(|| default.to_owned()) -} - fn env_optional(name: &str) -> Option { std::env::var(name) .ok() .map(|value| value.trim().to_owned()) .filter(|value| !value.is_empty()) } - -fn env_u64(name: &str, default: u64) -> Result { - match env_optional(name) { - Some(value) => value - .parse::() - .map_err(|e| format!("invalid {name}={value:?}: {e}")), - None => Ok(default), - } -} - -fn env_u32(name: &str, default: u32) -> Result { - match env_optional(name) { - Some(value) => value - .parse::() - .map_err(|e| format!("invalid {name}={value:?}: {e}")), - None => Ok(default), - } -} - -fn env_json(name: &str) -> Result, String> -where - T: serde::de::DeserializeOwned, -{ - env_optional(name) - .map(|value| serde_json::from_str(&value).map_err(|e| format!("invalid {name} JSON: {e}"))) - .transpose() -} - -fn gguf_source_from_env() -> GgufSource { - if let Some(path) = env_optional("MVP_GGUF_LOCAL_PATH") { - return GgufSource::LocalPath(path); - } - GgufSource::HuggingFaceGguf { - repo: env_string("MVP_GGUF_REPO", DEFAULT_HF_REPO), - file: env_string("MVP_GGUF_FILE", DEFAULT_HF_FILE), - revision: env_optional("MVP_GGUF_REVISION"), - } -} - -fn tokenizer_from_env() -> TokenizerSource { - env_optional("MVP_TOKENIZER_LOCAL_PATH") - .map(TokenizerSource::LocalPath) - .unwrap_or(TokenizerSource::EmbeddedGguf) -} diff --git a/crates/mvp-system/src/orchestration/app.rs b/crates/mvp-system/src/orchestration/app.rs index 142c46f..5c3d9f9 100644 --- a/crates/mvp-system/src/orchestration/app.rs +++ b/crates/mvp-system/src/orchestration/app.rs @@ -89,54 +89,26 @@ const MVP_ORCH_PROMPT: &str = "mvp.orch.prompt"; const MVP_SWIM_MEMBERSHIP: &str = "mvp.swim.membership"; const MVP_STAGE_ROUTE: &str = "mvp.orch.stage_route"; const DATASTREAM_FRAME_LOG_ENV: &str = "MVP_DATASTREAM_FRAME_LOG"; -const DEFAULT_DOCKER_CONTAINER_PREFIX: &str = "mvp-orchestrator"; -const MVP_DOCKER_CONTAINER_PREFIX_ENV: &str = "MVP_DOCKER_CONTAINER_PREFIX"; -struct OrchestratorRunOptions { +pub(crate) fn run_with_options( + args: I, capture_stdio: bool, stop_rx: Option>, -} - -impl Default for OrchestratorRunOptions { - fn default() -> Self { - Self { - capture_stdio: true, - stop_rx: None, - } - } -} - -pub(super) fn run_from_args(args: I) -> Result<(), String> -where - I: IntoIterator, -{ - run_with_options(args, OrchestratorRunOptions::default()) -} - -pub(super) fn run_in_process_from_args( - args: I, - stop_rx: mpsc::Receiver<()>, ) -> Result<(), String> where I: IntoIterator, { - run_with_options( - args, - OrchestratorRunOptions { - capture_stdio: false, - stop_rx: Some(stop_rx), - }, - ) -} - -fn run_with_options(args: I, options: OrchestratorRunOptions) -> Result<(), String> -where - I: IntoIterator, -{ - let mut config = Config::from_defaults_toml_env_args(args)?; + let mut config_builder = ConfigBuilder::hardcoded_defaults(); + if let Some(overlay) = TomlConfigOverlay::load_optional(Path::new(DEFAULT_CONFIG_PATH))? { + config_builder = config_builder.overlay_toml(overlay)?; + } + let mut config = config_builder + .overlay_env()? + .overlay_cli(args)? + .finalize()?; config.prepare_vastai_ssh_key()?; - let orch_stdio_rx = if options.capture_stdio { - install_orch_stdio_capture()? + let orch_stdio_rx = if capture_stdio { + OrchStdioCapture::install()? } else { None }; @@ -149,7 +121,10 @@ where "config", "ready", json!({ - "config_profile":config.config_profile.as_str(), + "config_profile":match config.config_profile { + RuntimeConfigProfile::Local => "local", + RuntimeConfigProfile::Deploy => "deploy", + }, "image":&config.image, "provider":config.provider.as_str(), "rpc_bind":config.rpc_bind.to_string(), @@ -237,7 +212,6 @@ where } else { None }; - let _ = &pipeline_plan; let tokio = match tokio::runtime::Runtime::new() { Ok(runtime) => { @@ -478,7 +452,7 @@ where ); let (work_tx, work_rx) = mpsc::channel::(); - let stop_rx = options.stop_rx.unwrap_or_else(spawn_stop_listener); + let stop_rx = stop_rx.unwrap_or_else(spawn_stop_listener); let provisioner = config.build_provisioner(Arc::clone(&stack.runtime))?; orch_datastream.emit_bootstrap( @@ -814,20 +788,6 @@ impl RuntimeConfigProfile { )), } } - - fn as_str(self) -> &'static str { - match self { - Self::Local => "local", - Self::Deploy => "deploy", - } - } - - fn default_provider(self) -> ProviderKind { - match self { - Self::Local => provider_kind::process(), - Self::Deploy => provider_kind::vastai(), - } - } } #[derive(Clone)] @@ -858,21 +818,23 @@ impl CachedModelConfig { host_path.display() )); } - let container_path = cached_model_container_path(&host_path)?; + let file_name = host_path + .file_name() + .and_then(|name| name.to_str()) + .filter(|name| !name.is_empty()) + .ok_or_else(|| { + format!( + "cached model path has no file name: {}", + host_path.display() + ) + })?; + let container_path = format!("{CACHED_MODEL_CONTAINER_DIR}/{file_name}"); Ok(Self { host_path, container_path, }) } - fn worker_path(&self, provider: &ProviderKind) -> String { - if provider == &provider_kind::process() { - self.host_path.to_string_lossy().to_string() - } else { - self.container_path.clone() - } - } - fn datastream_detail(&self) -> Value { json!({ "host_path_present": true, @@ -882,26 +844,6 @@ impl CachedModelConfig { } } -fn cached_model_container_path(host_path: &Path) -> Result { - let file_name = host_path - .file_name() - .and_then(|name| name.to_str()) - .filter(|name| !name.is_empty()) - .ok_or_else(|| { - format!( - "cached model path has no file name: {}", - host_path.display() - ) - })?; - Ok(format!("{CACHED_MODEL_CONTAINER_DIR}/{file_name}")) -} - -fn default_worker_bin() -> Result { - let mut path = std::env::current_exe().map_err(|e| format!("current exe: {e}"))?; - path.set_file_name("mvp-worker-node"); - Ok(path) -} - fn default_pipeline_cached_model_path() -> PathBuf { let relative = PathBuf::from(".") .join(DEFAULT_PIPELINE_MODEL_CACHE_DIR) @@ -916,28 +858,6 @@ fn default_pipeline_cached_model_path() -> PathBuf { .join(DEFAULT_PIPELINE_CACHED_MODEL_FILE) } -fn gguf_source_is_default_hf(source: &GgufSource) -> bool { - matches!( - source, - GgufSource::HuggingFaceGguf { - repo, - file, - revision: None, - } if repo == DEFAULT_HF_REPO && file == DEFAULT_HF_FILE - ) -} - -fn gguf_source_matches_default_pipeline_cache(source: &GgufSource) -> bool { - matches!( - source, - GgufSource::HuggingFaceGguf { - file, - revision: None, - .. - } if file == DEFAULT_PIPELINE_CACHED_MODEL_FILE - ) -} - #[derive(Clone)] struct Config { config_profile: RuntimeConfigProfile, @@ -1121,10 +1041,14 @@ impl ConfigBuilder { apply!(overlay.model.gguf_local_path, |path| { self.gguf_source = GgufSource::LocalPath(path) }); - apply!(overlay.model.gguf_repo, |repo| { self.set_gguf_repo(repo) }); - apply!(overlay.model.gguf_file, |file| { self.set_gguf_file(file) }); + apply!(overlay.model.gguf_repo, |repo| { + self.set_hf_source(Some(repo), None, None) + }); + apply!(overlay.model.gguf_file, |file| { + self.set_hf_source(None, Some(file), None) + }); apply!(overlay.model.gguf_revision, |revision| { - self.set_gguf_revision(Some(revision)) + self.set_hf_source(None, None, Some(Some(revision))) }); apply!(overlay.model.tokenizer_local_path, |path| { self.tokenizer = TokenizerSource::LocalPath(path) @@ -1233,7 +1157,10 @@ impl ConfigBuilder { self.provider = Some(provider_kind::parse_deploy(&provider)?); } ); - env_apply!("MVP_NODE_IMAGE", |image| { self.set_process_image(image) }); + env_apply!("MVP_NODE_IMAGE", |image| { + self.image = image; + self.image_overridden_after_toml = true; + }); env_apply!("MVP_DOCKER_GPUS", |gpus| { self.docker_gpus = gpus }); env_apply!(CACHED_MODEL_HOST_ENV, |path| { self.cached_model_host_path = Some(PathBuf::from(path)) @@ -1258,10 +1185,14 @@ impl ConfigBuilder { env_apply!("MVP_GGUF_LOCAL_PATH", |path| { self.gguf_source = GgufSource::LocalPath(path) }); - env_apply!("MVP_GGUF_REPO", |repo| { self.set_gguf_repo(repo) }); - env_apply!("MVP_GGUF_FILE", |file| { self.set_gguf_file(file) }); + env_apply!("MVP_GGUF_REPO", |repo| { + self.set_hf_source(Some(repo), None, None) + }); + env_apply!("MVP_GGUF_FILE", |file| { + self.set_hf_source(None, Some(file), None) + }); env_apply!("MVP_GGUF_REVISION", |revision| { - self.set_gguf_revision(Some(revision)) + self.set_hf_source(None, None, Some(Some(revision))) }); env_apply!("MVP_TOKENIZER_LOCAL_PATH", |path| { self.tokenizer = TokenizerSource::LocalPath(path) @@ -1340,179 +1271,155 @@ impl ConfigBuilder { Ok(self) } - fn apply_core_cli_arg(&mut self, arg: &str, args: &mut I) -> Result - where - I: Iterator, - { - match arg { - "--runtime-config" => { - self.config_profile = - RuntimeConfigProfile::parse(&next_arg(args, "--runtime-config")?)? - } - "--provider" => { - self.provider = Some(provider_kind::parse_deploy(&next_arg(args, "--provider")?)?) - } - "--worker-bin" => { - self.worker_bin = Some(PathBuf::from(next_arg(args, "--worker-bin")?)) - } - "--image" => self.set_process_image(next_arg(args, "--image")?), - "--gpus" => self.docker_gpus = next_arg(args, "--gpus")?, - "--rpc-bind" => { - self.rpc_bind = next_arg(args, "--rpc-bind")?; - self.rpc_bind_label = "--rpc-bind"; - } - "--run-id" => self.run_id = parse_next(args, "--run-id")?, - "--node-id" => self.node_id = parse_next(args, "--node-id")?, - "--stage-index" => self.stage_index = parse_next(args, "--stage-index")?, - "--layer-end-exclusive" => { - self.layer_end_exclusive = Some(parse_next(args, "--layer-end-exclusive")?) - } - "-N" | "--pipeline-stages" => self.pipeline_stages = parse_next(args, arg)?, - "--max-tokens" => self.default_max_tokens = parse_next(args, "--max-tokens")?, - "--dashboard" => self.dashboard = true, - "--no-dashboard" => self.dashboard = false, - "--datastream-frame-log" => { - self.datastream_frame_log = - Some(PathBuf::from(next_arg(args, "--datastream-frame-log")?)); - } - _ => return Ok(false), - } - Ok(true) - } - - fn apply_model_cli_arg(&mut self, arg: &str, args: &mut I) -> Result - where - I: Iterator, - { - match arg { - "--model-id" => self.model_id = next_arg(args, "--model-id")?, - "--gguf-local-path" => { - self.gguf_source = GgufSource::LocalPath(next_arg(args, "--gguf-local-path")?) - } - "--gguf-repo" => self.set_gguf_repo(next_arg(args, "--gguf-repo")?), - "--gguf-file" => self.set_gguf_file(next_arg(args, "--gguf-file")?), - "--gguf-revision" => self.set_gguf_revision(Some(next_arg(args, "--gguf-revision")?)), - "--tokenizer-local-path" => { - self.tokenizer = - TokenizerSource::LocalPath(next_arg(args, "--tokenizer-local-path")?) - } - "--max-context" => self.max_context = Some(parse_next(args, "--max-context")?), - "--cached-model-host-path" => { - self.cached_model_host_path = - Some(PathBuf::from(next_arg(args, "--cached-model-host-path")?)); - } - "--relay-mode" => self.relay_mode = Some(next_arg(args, "--relay-mode")?), - "--relay-url" => self.relay_url = Some(next_arg(args, "--relay-url")?), - "--endpoint-addr-mask" => { - self.endpoint_addr_mask = Some(next_arg(args, "--endpoint-addr-mask")?) - } - _ => return Ok(false), - } - Ok(true) - } - - fn apply_vastai_string_cli_arg(&mut self, arg: &str, args: &mut I) -> Result - where - I: Iterator, - { - match arg { - "--vastai-api-key" => self.vastai_api_key = Some(next_arg(args, "--vastai-api-key")?), - "--vastai-bootstrap-command" => { - self.vastai_bootstrap_command = Some(next_arg(args, "--vastai-bootstrap-command")?) - } - "--vastai-ssh-identity" => { - self.vastai_ssh_identity_raw = Some(next_arg(args, "--vastai-ssh-identity")?); - } - "--vastai-ssh-user" => { - self.vastai_ssh_user = Some(next_arg(args, "--vastai-ssh-user")?) - } - "--vastai-onstart" => self.vastai_onstart = Some(next_arg(args, "--vastai-onstart")?), - "--vastai-gpu-name" => { - self.vastai_gpu_name = Some(next_arg(args, "--vastai-gpu-name")?) - } - _ => return Ok(false), - } - Ok(true) - } - - fn apply_vastai_numeric_cli_arg(&mut self, arg: &str, args: &mut I) -> Result - where - I: Iterator, - { - match arg { - "--vastai-disk-gb" => { - self.vastai_disk_gb = Some(parse_next(args, "--vastai-disk-gb")?); - self.vastai_disk_gb_raw = None; - } - "--vastai-min-gpu-ram-mb" => { - self.vastai_min_gpu_ram_mb = Some(parse_next(args, "--vastai-min-gpu-ram-mb")?); - self.vastai_min_gpu_ram_mb_raw = None; - } - "--vastai-min-down-mbps" => { - self.vastai_min_down_mbps = Some(parse_next(args, "--vastai-min-down-mbps")?); - self.vastai_min_down_mbps_raw = None; - } - "--vastai-max-dph-total" => { - self.vastai_max_dph_total = Some(parse_next(args, "--vastai-max-dph-total")?); - self.vastai_max_dph_total_raw = None; - } - "--vastai-min-up-mbps" => { - self.vastai_min_up_mbps = Some(parse_next(args, "--vastai-min-up-mbps")?); - self.vastai_min_up_mbps_raw = None; - } - "--vastai-min-reliability" => { - self.vastai_min_reliability = Some(parse_next(args, "--vastai-min-reliability")?); - self.vastai_min_reliability_raw = None; - } - "--vastai-blacklist-host" => { - let host_id = parse_next(args, "--vastai-blacklist-host")?; - self.push_vastai_blacklist_host(host_id); - } - "--vastai-poll-interval-secs" => { - self.vastai_poll_interval_secs = - Some(parse_next(args, "--vastai-poll-interval-secs")?); - self.vastai_poll_interval_secs_raw = None; - } - _ => return Ok(false), - } - Ok(true) - } - - fn apply_vastai_bool_cli_arg(&mut self, arg: &str) -> bool { - match arg { - "--vastai-confirm-lease" => { - self.vastai_confirm_lease = Some(true); - self.vastai_confirm_lease_raw = None; - } - "--no-vastai-confirm-lease" => { - self.vastai_confirm_lease = Some(false); - self.vastai_confirm_lease_raw = None; - } - "--vastai-require-verified" => { - self.vastai_require_verified = Some(true); - self.vastai_require_verified_raw = None; - } - "--no-vastai-require-verified" => { - self.vastai_require_verified = Some(false); - self.vastai_require_verified_raw = None; - } - _ => return false, - } - true - } - fn overlay_cli(mut self, args: impl IntoIterator) -> Result { let mut args = args.into_iter(); while let Some(arg) = args.next() { - if self.apply_core_cli_arg(&arg, &mut args)? - || self.apply_model_cli_arg(&arg, &mut args)? - || self.apply_vastai_string_cli_arg(&arg, &mut args)? - || self.apply_vastai_numeric_cli_arg(&arg, &mut args)? - || self.apply_vastai_bool_cli_arg(&arg) - { - continue; + match arg.as_str() { + "--runtime-config" => { + self.config_profile = + RuntimeConfigProfile::parse(&next_arg(&mut args, "--runtime-config")?)? + } + "--provider" => { + self.provider = Some(provider_kind::parse_deploy(&next_arg( + &mut args, + "--provider", + )?)?) + } + "--worker-bin" => { + self.worker_bin = Some(PathBuf::from(next_arg(&mut args, "--worker-bin")?)) + } + "--image" => { + self.image = next_arg(&mut args, "--image")?; + self.image_overridden_after_toml = true; + } + "--gpus" => self.docker_gpus = next_arg(&mut args, "--gpus")?, + "--rpc-bind" => { + self.rpc_bind = next_arg(&mut args, "--rpc-bind")?; + self.rpc_bind_label = "--rpc-bind"; + } + "--run-id" => self.run_id = parse_next(&mut args, "--run-id")?, + "--node-id" => self.node_id = parse_next(&mut args, "--node-id")?, + "--stage-index" => self.stage_index = parse_next(&mut args, "--stage-index")?, + "--layer-end-exclusive" => { + self.layer_end_exclusive = Some(parse_next(&mut args, "--layer-end-exclusive")?) + } + "-N" | "--pipeline-stages" => self.pipeline_stages = parse_next(&mut args, &arg)?, + "--max-tokens" => self.default_max_tokens = parse_next(&mut args, "--max-tokens")?, + "--dashboard" => self.dashboard = true, + "--no-dashboard" => self.dashboard = false, + "--datastream-frame-log" => { + self.datastream_frame_log = Some(PathBuf::from(next_arg( + &mut args, + "--datastream-frame-log", + )?)); + } + "--model-id" => self.model_id = next_arg(&mut args, "--model-id")?, + "--gguf-local-path" => { + self.gguf_source = + GgufSource::LocalPath(next_arg(&mut args, "--gguf-local-path")?) + } + "--gguf-repo" => { + self.set_hf_source(Some(next_arg(&mut args, "--gguf-repo")?), None, None) + } + "--gguf-file" => { + self.set_hf_source(None, Some(next_arg(&mut args, "--gguf-file")?), None) + } + "--gguf-revision" => self.set_hf_source( + None, + None, + Some(Some(next_arg(&mut args, "--gguf-revision")?)), + ), + "--tokenizer-local-path" => { + self.tokenizer = + TokenizerSource::LocalPath(next_arg(&mut args, "--tokenizer-local-path")?) + } + "--max-context" => self.max_context = Some(parse_next(&mut args, "--max-context")?), + "--cached-model-host-path" => { + self.cached_model_host_path = Some(PathBuf::from(next_arg( + &mut args, + "--cached-model-host-path", + )?)); + } + "--relay-mode" => self.relay_mode = Some(next_arg(&mut args, "--relay-mode")?), + "--relay-url" => self.relay_url = Some(next_arg(&mut args, "--relay-url")?), + "--endpoint-addr-mask" => { + self.endpoint_addr_mask = Some(next_arg(&mut args, "--endpoint-addr-mask")?) + } + "--vastai-disk-gb" => { + self.vastai_disk_gb = Some(parse_next(&mut args, "--vastai-disk-gb")?); + self.vastai_disk_gb_raw = None; + } + "--vastai-min-gpu-ram-mb" => { + self.vastai_min_gpu_ram_mb = + Some(parse_next(&mut args, "--vastai-min-gpu-ram-mb")?); + self.vastai_min_gpu_ram_mb_raw = None; + } + "--vastai-min-down-mbps" => { + self.vastai_min_down_mbps = + Some(parse_next(&mut args, "--vastai-min-down-mbps")?); + self.vastai_min_down_mbps_raw = None; + } + "--vastai-max-dph-total" => { + self.vastai_max_dph_total = + Some(parse_next(&mut args, "--vastai-max-dph-total")?); + self.vastai_max_dph_total_raw = None; + } + "--vastai-min-up-mbps" => { + self.vastai_min_up_mbps = Some(parse_next(&mut args, "--vastai-min-up-mbps")?); + self.vastai_min_up_mbps_raw = None; + } + "--vastai-min-reliability" => { + self.vastai_min_reliability = + Some(parse_next(&mut args, "--vastai-min-reliability")?); + self.vastai_min_reliability_raw = None; + } + "--vastai-blacklist-host" => { + let host_id = parse_next(&mut args, "--vastai-blacklist-host")?; + self.push_vastai_blacklist_host(host_id); + } + "--vastai-poll-interval-secs" => { + self.vastai_poll_interval_secs = + Some(parse_next(&mut args, "--vastai-poll-interval-secs")?); + self.vastai_poll_interval_secs_raw = None; + } + "--vastai-api-key" => { + self.vastai_api_key = Some(next_arg(&mut args, "--vastai-api-key")?) + } + "--vastai-bootstrap-command" => { + self.vastai_bootstrap_command = + Some(next_arg(&mut args, "--vastai-bootstrap-command")?) + } + "--vastai-ssh-identity" => { + self.vastai_ssh_identity_raw = + Some(next_arg(&mut args, "--vastai-ssh-identity")?); + } + "--vastai-ssh-user" => { + self.vastai_ssh_user = Some(next_arg(&mut args, "--vastai-ssh-user")?) + } + "--vastai-onstart" => { + self.vastai_onstart = Some(next_arg(&mut args, "--vastai-onstart")?) + } + "--vastai-gpu-name" => { + self.vastai_gpu_name = Some(next_arg(&mut args, "--vastai-gpu-name")?) + } + "--vastai-confirm-lease" => { + self.vastai_confirm_lease = Some(true); + self.vastai_confirm_lease_raw = None; + } + "--no-vastai-confirm-lease" => { + self.vastai_confirm_lease = Some(false); + self.vastai_confirm_lease_raw = None; + } + "--vastai-require-verified" => { + self.vastai_require_verified = Some(true); + self.vastai_require_verified_raw = None; + } + "--no-vastai-require-verified" => { + self.vastai_require_verified = Some(false); + self.vastai_require_verified_raw = None; + } + _ => return Err(format!("unknown argument {arg:?}")), } - return Err(format!("unknown argument {arg:?}")); } Ok(self) } @@ -1521,7 +1428,10 @@ impl ConfigBuilder { let provider = self .provider .clone() - .unwrap_or_else(|| self.config_profile.default_provider()); + .unwrap_or_else(|| match self.config_profile { + RuntimeConfigProfile::Local => provider_kind::process(), + RuntimeConfigProfile::Deploy => provider_kind::vastai(), + }); let mut image = self.image.clone(); if provider == provider_kind::vastai() && !self.image_overridden_after_toml { if let Some(vastai_image) = &self.toml_vastai_image { @@ -1535,8 +1445,21 @@ impl ConfigBuilder { if (provider == provider_kind::process() || provider == provider_kind::docker()) && self.pipeline_stages > 1 && cached_model_host_path.is_none() - && (gguf_source_is_default_hf(&self.gguf_source) - || gguf_source_matches_default_pipeline_cache(&self.gguf_source)) + && (matches!( + &self.gguf_source, + GgufSource::HuggingFaceGguf { + repo, + file, + revision: None, + } if repo == DEFAULT_HF_REPO && file == DEFAULT_HF_FILE + ) || matches!( + &self.gguf_source, + GgufSource::HuggingFaceGguf { + file, + revision: None, + .. + } if file == DEFAULT_PIPELINE_CACHED_MODEL_FILE + )) { cached_model_host_path = Some(default_pipeline_cached_model_path()); } @@ -1546,7 +1469,11 @@ impl ConfigBuilder { let mut gguf_source = self.gguf_source.clone(); if let Some(cached_model) = &cached_model { if provider != provider_kind::vastai() { - gguf_source = GgufSource::LocalPath(cached_model.worker_path(&provider)); + gguf_source = GgufSource::LocalPath(if provider == provider_kind::process() { + cached_model.host_path.to_string_lossy().to_string() + } else { + cached_model.container_path.clone() + }); } } let relay = relay_runtime_config_from_settings( @@ -1592,44 +1519,26 @@ impl ConfigBuilder { }) } - fn set_process_image(&mut self, image: String) { - self.image = image; - self.image_overridden_after_toml = true; - } - - fn set_gguf_repo(&mut self, repo: String) { - let (file, revision) = match &self.gguf_source { - GgufSource::HuggingFaceGguf { file, revision, .. } => (file.clone(), revision.clone()), - GgufSource::LocalPath(_) => (DEFAULT_HF_FILE.to_owned(), None), + fn set_hf_source( + &mut self, + repo: Option, + file: Option, + revision: Option>, + ) { + let (current_repo, current_file, current_revision) = match &self.gguf_source { + GgufSource::HuggingFaceGguf { + repo, + file, + revision, + } => (repo.clone(), file.clone(), revision.clone()), + GgufSource::LocalPath(_) => { + (DEFAULT_HF_REPO.to_owned(), DEFAULT_HF_FILE.to_owned(), None) + } }; self.gguf_source = GgufSource::HuggingFaceGguf { - repo, - file, - revision, - }; - } - - fn set_gguf_file(&mut self, file: String) { - let (repo, revision) = match &self.gguf_source { - GgufSource::HuggingFaceGguf { repo, revision, .. } => (repo.clone(), revision.clone()), - GgufSource::LocalPath(_) => (DEFAULT_HF_REPO.to_owned(), None), - }; - self.gguf_source = GgufSource::HuggingFaceGguf { - repo, - file, - revision, - }; - } - - fn set_gguf_revision(&mut self, revision: Option) { - let (repo, file) = match &self.gguf_source { - GgufSource::HuggingFaceGguf { repo, file, .. } => (repo.clone(), file.clone()), - GgufSource::LocalPath(_) => (DEFAULT_HF_REPO.to_owned(), DEFAULT_HF_FILE.to_owned()), - }; - self.gguf_source = GgufSource::HuggingFaceGguf { - repo, - file, - revision, + repo: repo.unwrap_or(current_repo), + file: file.unwrap_or(current_file), + revision: revision.unwrap_or(current_revision), }; } @@ -1674,27 +1583,6 @@ impl ConfigBuilder { } impl Config { - fn from_defaults_toml_env_args(args: impl IntoIterator) -> Result { - Self::from_layers_with_path_and_args(Some(Path::new(DEFAULT_CONFIG_PATH)), args) - } - - fn from_layers_with_path_and_args( - path: Option<&Path>, - args: impl IntoIterator, - ) -> Result { - let mut builder = Self::hardcoded_defaults(); - if let Some(path) = path { - if let Some(overlay) = TomlConfigOverlay::load_optional(path)? { - builder = builder.overlay_toml(overlay)?; - } - } - builder.overlay_env()?.overlay_cli(args)?.finalize() - } - - fn hardcoded_defaults() -> ConfigBuilder { - ConfigBuilder::hardcoded_defaults() - } - fn uses_planned_execution(&self) -> bool { self.cached_model.is_some() || self.pipeline_stages > 1 } @@ -1817,9 +1705,12 @@ impl Config { )) } } - GgufSource::HuggingFaceGguf { repo, file, .. } - if self.provider == provider_kind::vastai() - && gguf_source_matches_default_pipeline_cache(&self.gguf_source) => + GgufSource::HuggingFaceGguf { + repo, + file, + revision: None, + } if self.provider == provider_kind::vastai() + && file == DEFAULT_PIPELINE_CACHED_MODEL_FILE => { let host_path = default_pipeline_cached_model_path(); if host_path.is_file() { @@ -1887,7 +1778,15 @@ impl Config { bootstrap_runtime: Arc, ) -> Result, String> { if self.provider == provider_kind::process() { - let worker_bin = self.worker_bin.clone().unwrap_or(default_worker_bin()?); + let worker_bin = match &self.worker_bin { + Some(worker_bin) => worker_bin.clone(), + None => { + let mut path = + std::env::current_exe().map_err(|e| format!("current exe: {e}"))?; + path.set_file_name("mvp-worker-node"); + path + } + }; if !worker_bin.is_file() { return Err(format!( "local process worker binary does not exist: {}", @@ -1896,7 +1795,10 @@ impl Config { } Ok(Box::new(LocalProcessPlugin::new(worker_bin))) } else if self.provider == provider_kind::docker() { - Ok(Box::new(LocalDockerPlugin::new(docker_container_prefix()))) + Ok(Box::new(LocalDockerPlugin::new( + env_optional("MVP_DOCKER_CONTAINER_PREFIX") + .unwrap_or_else(|| "mvp-orchestrator".to_owned()), + ))) } else if self.provider == provider_kind::vastai() { let vastai = self .vastai @@ -2030,14 +1932,22 @@ impl Config { if self.provider == provider_kind::docker() { env.push(("MVP_DOCKER_GPUS".to_owned(), self.docker_gpus.clone())); } - env.extend(optional_env("DEV")); + if let Some(value) = env_optional("DEV") { + env.push(("DEV".to_owned(), value)); + } env.extend(local_tinygrad_worker_env(&self.provider)); - env.extend(optional_env("MVP_CPU_LINE_PROFILE")); - env.extend(optional_env("MVP_CPU_LINE_PROFILE_INTERVAL_MS")); - env.extend(optional_env("MVP_TOKEN_PROGRESS_EVERY")); - env.extend(optional_env("CUDA_DEVICE_SCHEDULE")); - env.extend(optional_env("MVP_MODEL_CACHE_DIR")); - env.extend(optional_env("HF_TOKEN")); + for name 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 let Some(value) = env_optional(name) { + env.push((name.to_owned(), value)); + } + } match &self.gguf_source { GgufSource::LocalPath(path) => { env.push(("MVP_GGUF_LOCAL_PATH".to_owned(), path.clone())) @@ -2378,61 +2288,6 @@ impl Drop for ProvisionedClusterGuard { } } -type ProviderStartResults = Vec<( - NodeProvisionSpec, - Result, -)>; - -fn start_nodes_with_stdio_capture( - provisioner: Box, - node_specs: Vec, - sink: PluginSink, - orch_stdio_rx: Option<&mpsc::Receiver>, - dashboard: Option<&DashboardSupport>, - orch_datastream: &mut OrchDatastream, - run_id: u64, - node_id: u64, -) -> (Box, ProviderStartResults) { - let (tx, rx) = mpsc::channel(); - thread::spawn(move || { - let mut provisioner = provisioner; - let results = provisioner.start_nodes(node_specs, sink); - let _ = tx.send((provisioner, results)); - }); - - loop { - match rx.recv_timeout(Duration::from_millis(100)) { - Ok(result) => return result, - Err(mpsc::RecvTimeoutError::Timeout) => { - drain_orch_stdio_capture( - orch_stdio_rx, - orch_datastream, - dashboard, - run_id, - node_id, - ); - } - Err(mpsc::RecvTimeoutError::Disconnected) => { - return ( - Box::new(FailedProvisionPlugin), - vec![( - NodeProvisionSpec { - run_id, - node_id, - stage_index: None, - image: String::new(), - env: Vec::new(), - args: Vec::new(), - mounts: Vec::new(), - }, - Err("provider start worker disconnected".to_owned()), - )], - ); - } - } - } -} - fn start_and_provision_workers( mut provisioner: Box, config: &Config, @@ -2531,22 +2386,40 @@ fn start_and_provision_workers( }), ); } - let (returned_provisioner, start_results) = start_nodes_with_stdio_capture( - provisioner, - pending_specs, - sink.clone(), - orch_stdio_rx, - dashboard, - orch_datastream, - config.run_id, - config.node_id, - ); - provisioner = returned_provisioner; - let start_outcome = collect_provider_start_outcome(start_results); - handles.extend(start_outcome.successful_handles); - for (node_spec, handle_result) in start_outcome.results { + let (tx, rx) = mpsc::channel(); + thread::spawn({ + let sink = sink.clone(); + move || { + let mut provisioner = provisioner; + let results = provisioner.start_nodes(pending_specs, sink); + let _ = tx.send((provisioner, results)); + } + }); + let start_results = loop { + match rx.recv_timeout(Duration::from_millis(100)) { + Ok((returned_provisioner, start_results)) => { + provisioner = returned_provisioner; + break start_results; + } + Err(mpsc::RecvTimeoutError::Timeout) => { + drain_orch_stdio_capture( + orch_stdio_rx, + orch_datastream, + dashboard, + config.run_id, + config.node_id, + ); + } + Err(mpsc::RecvTimeoutError::Disconnected) => { + return Err("provider start worker disconnected".to_owned()); + } + } + }; + let mut failed_specs = Vec::new(); + let mut first_error = None; + for (node_spec, handle_result) in start_results { match handle_result { - Ok(_) => { + Ok(handle) => { orch_datastream.emit_bootstrap( dashboard, config.run_id, @@ -2560,6 +2433,7 @@ fn start_and_provision_workers( "attempt":attempt, }), ); + handles.push(handle); } Err(error) => { orch_datastream.emit_bootstrap( @@ -2576,17 +2450,21 @@ fn start_and_provision_workers( "error":error, }), ); + failed_specs.push(node_spec); + if first_error.is_none() { + first_error = Some(error); + } } } } - if start_outcome.first_error.is_none() { + if first_error.is_none() { break; } if attempt == PROVIDER_START_MAX_ATTEMPTS { - let error = start_outcome - .first_error - .expect("checked provider-start failure"); - stop_started_nodes(&mut *provisioner, &mut handles); + let error = first_error.expect("checked provider-start failure"); + while let Some(handle) = handles.pop() { + let _ = provisioner.stop_node(&handle); + } drain_orch_stdio_capture( orch_stdio_rx, orch_datastream, @@ -2596,7 +2474,7 @@ fn start_and_provision_workers( ); return Err(error); } - pending_specs = start_outcome.failed_specs; + pending_specs = failed_specs; } let mut provisioned_nodes = ProvisionedClusterGuard::new(provisioner, handles); drain_orch_stdio_capture( @@ -2878,44 +2756,6 @@ fn stage_node_specs( )?]) } } -struct ProviderStartOutcome { - results: ProviderStartResults, - successful_handles: Vec, - failed_specs: Vec, - first_error: Option, -} - -fn collect_provider_start_outcome(results: ProviderStartResults) -> ProviderStartOutcome { - let mut successful_handles = Vec::new(); - let mut failed_specs = Vec::new(); - let mut first_error = None; - for (spec, result) in &results { - match result { - Ok(handle) => successful_handles.push(handle.clone()), - Err(error) => { - failed_specs.push(spec.clone()); - if first_error.is_none() { - first_error = Some(error.clone()); - } - } - } - } - ProviderStartOutcome { - results, - successful_handles, - failed_specs, - first_error, - } -} - -fn stop_started_nodes( - provisioner: &mut dyn ProvisionPlugin, - handles: &mut Vec, -) { - while let Some(handle) = handles.pop() { - let _ = provisioner.stop_node(&handle); - } -} fn stage_provision_detail(config: &Config, pipeline_plan: Option<&run_plan::RunPlan>) -> Value { if let Some(plan) = pipeline_plan { @@ -3276,25 +3116,24 @@ fn wait_for_weights_loaded_count( )); } for stage in pending { - send_pipeline_stage_provision( - driver, + let mut provision = PipelineStageProvision { + driver: &mut *driver, stack, frame_tx, dashboard, - orch_datastream, + orch_datastream: &mut *orch_datastream, run_id, node_id, pipeline_plan, - stage, readies, pipeline_coordinator, - &stage_shard_plans, - &loaded_stages, - &mut stage_resend_counts, - &mut stage_last_sends, - &load_progress, - resend_attempt, - )?; + stage_shard_plans, + loaded_stages: &loaded_stages, + stage_resend_counts: &mut stage_resend_counts, + stage_last_sends: &mut stage_last_sends, + load_progress: &load_progress, + }; + send_pipeline_stage_provision(&mut provision, stage, resend_attempt)?; } last_resend = Instant::now(); } @@ -3396,44 +3235,49 @@ fn pending_pipeline_weight_load_stages<'a>( pending } -#[allow(clippy::too_many_arguments)] -fn send_pipeline_stage_provision( - driver: &mut IrohDriver, - stack: &DistributionRuntimeStack, - frame_tx: &mpsc::Sender, - dashboard: Option<&DashboardSupport>, - orch_datastream: &mut OrchDatastream, +struct PipelineStageProvision<'a> { + driver: &'a mut IrohDriver, + stack: &'a DistributionRuntimeStack, + frame_tx: &'a mpsc::Sender, + dashboard: Option<&'a DashboardSupport>, + orch_datastream: &'a mut OrchDatastream, run_id: u64, node_id: u64, - pipeline_plan: &run_plan::RunPlan, + pipeline_plan: &'a run_plan::RunPlan, + readies: &'a BTreeMap, + pipeline_coordinator: &'a EndpointAddr, + stage_shard_plans: &'a BTreeMap, + loaded_stages: &'a BTreeSet, + stage_resend_counts: &'a mut BTreeMap, + stage_last_sends: &'a mut BTreeMap, + load_progress: &'a BTreeMap, +} + +fn send_pipeline_stage_provision( + ctx: &mut PipelineStageProvision<'_>, stage: &run_plan::StagePlan, - readies: &BTreeMap, - pipeline_coordinator: &EndpointAddr, - stage_shard_plans: &BTreeMap, - loaded_stages: &BTreeSet, - stage_resend_counts: &mut BTreeMap, - stage_last_sends: &mut BTreeMap, - load_progress: &BTreeMap, attempt: u64, ) -> Result<(), String> { let stage_node_id = stage.node_id.0; - let current_send_count = stage_resend_counts + let current_send_count = ctx + .stage_resend_counts .get(&stage.stage_index) .copied() .unwrap_or_default(); let now = Instant::now(); - let ready = readies + let ready = ctx + .readies .get(&stage_node_id) .ok_or_else(|| format!("missing runtime-ready node for stage {}", stage.stage_index))?; - let route_owner = stack.route_owner(ready.node_actor); - let datastream_route_owner = stack.route_owner(ready.datastream_publisher); - let member_state = stack.member_state(ready.swim_node_id); + let route_owner = ctx.stack.route_owner(ready.node_actor); + let datastream_route_owner = ctx.stack.route_owner(ready.datastream_publisher); + let member_state = ctx.stack.member_state(ready.swim_node_id); let route_matches_ready = route_owner == Some(ready.swim_node_id); - orch_datastream.emit_bootstrap_to_channel( - dashboard, + ctx.orch_datastream.emit_bootstrap_to_channel( + ctx.dashboard, MVP_STAGE_ROUTE, - run_id, - node_id, + ctx.run_id, + ctx.node_id, "stage_route_check", "observed", json!({ @@ -3455,7 +3299,7 @@ fn send_pipeline_stage_provision( stage.stage_index, stage_node_id ); let liveness = stage_load_liveness_detail( - load_progress.get(&stage_node_id), + ctx.load_progress.get(&stage_node_id), stage.stage_index, stage_node_id, member_state, @@ -3464,19 +3308,19 @@ fn send_pipeline_stage_provision( route_matches_ready, "heartbeat_missed", ); - orch_datastream.emit_bootstrap( - dashboard, - run_id, - node_id, + ctx.orch_datastream.emit_bootstrap( + ctx.dashboard, + ctx.run_id, + ctx.node_id, "stage_provision_wait", "failed", json!({ "attempt":attempt, - "stage_count":pipeline_plan.stages.len(), + "stage_count":ctx.pipeline_plan.stages.len(), "stage_index":stage.stage_index, "stage_node_id":stage_node_id, "stage_send_count":current_send_count, - "loaded_stage_count":loaded_stages.len(), + "loaded_stage_count":ctx.loaded_stages.len(), "member_state":"Dead", "route_owner":route_owner.map(|owner| format!("{:?}", owner)), "datastream_route_owner":datastream_route_owner.map(|owner| format!("{:?}", owner)), @@ -3488,29 +3332,29 @@ fn send_pipeline_stage_provision( return Err(reason); } let (should_send, dispatch_reason) = stage_provision_dispatch( - load_progress.get(&stage_node_id), + ctx.load_progress.get(&stage_node_id), current_send_count, - stage_last_sends.get(&stage.stage_index).copied(), + ctx.stage_last_sends.get(&stage.stage_index).copied(), now, ); if !should_send { - orch_datastream.emit_bootstrap( - dashboard, - run_id, - node_id, + ctx.orch_datastream.emit_bootstrap( + ctx.dashboard, + ctx.run_id, + ctx.node_id, "stage_provision_wait", "observed", json!({ "attempt":attempt, - "stage_count":pipeline_plan.stages.len(), + "stage_count":ctx.pipeline_plan.stages.len(), "stage_index":stage.stage_index, "stage_node_id":stage_node_id, "stage_send_count":current_send_count, - "loaded_stage_count":loaded_stages.len(), + "loaded_stage_count":ctx.loaded_stages.len(), "resend_suppressed":true, "resend_reason":dispatch_reason, "liveness":stage_load_liveness_detail( - load_progress.get(&stage_node_id), + ctx.load_progress.get(&stage_node_id), stage.stage_index, stage_node_id, member_state, @@ -3521,8 +3365,8 @@ fn send_pipeline_stage_provision( ), "message":format!( "loaded {} of {}; waiting on stage {}", - loaded_stages.len(), - pipeline_plan.stages.len(), + ctx.loaded_stages.len(), + ctx.pipeline_plan.stages.len(), stage.stage_index ) }), @@ -3530,43 +3374,46 @@ fn send_pipeline_stage_provision( return Ok(()); } let stage_send_count = { - let count = stage_resend_counts.entry(stage.stage_index).or_default(); + let count = ctx + .stage_resend_counts + .entry(stage.stage_index) + .or_default(); *count += 1; *count }; - stage_last_sends.insert(stage.stage_index, now); - orch_datastream.emit_bootstrap( - dashboard, - run_id, - node_id, + ctx.stage_last_sends.insert(stage.stage_index, now); + ctx.orch_datastream.emit_bootstrap( + ctx.dashboard, + ctx.run_id, + ctx.node_id, "stage_provision_send", "sent", json!({ "attempt":attempt, - "stage_count":pipeline_plan.stages.len(), + "stage_count":ctx.pipeline_plan.stages.len(), "stage_index":stage.stage_index, "stage_send_count":stage_send_count, - "loaded_stage_count":loaded_stages.len(), + "loaded_stage_count":ctx.loaded_stages.len(), "parallel_weight_acquisition":true, "resend_reason":dispatch_reason, }), ); if stage_send_count == 1 || stage_send_count % 15 == 0 { - orch_datastream.emit_bootstrap( - dashboard, - run_id, - node_id, + ctx.orch_datastream.emit_bootstrap( + ctx.dashboard, + ctx.run_id, + ctx.node_id, "stage_provision_wait", "observed", json!({ "attempt":attempt, - "stage_count":pipeline_plan.stages.len(), + "stage_count":ctx.pipeline_plan.stages.len(), "stage_index":stage.stage_index, "stage_node_id":stage_node_id, "stage_send_count":stage_send_count, - "loaded_stage_count":loaded_stages.len(), + "loaded_stage_count":ctx.loaded_stages.len(), "liveness":stage_load_liveness_detail( - load_progress.get(&stage_node_id), + ctx.load_progress.get(&stage_node_id), stage.stage_index, stage_node_id, member_state, @@ -3577,49 +3424,26 @@ fn send_pipeline_stage_provision( ), "message":format!( "loaded {} of {}; waiting on stage {}", - loaded_stages.len(), - pipeline_plan.stages.len(), + ctx.loaded_stages.len(), + ctx.pipeline_plan.stages.len(), stage.stage_index ) }), ); } provision_stage_from_plan( - stack, + ctx.stack, ready.node_actor, - pipeline_plan, + ctx.pipeline_plan, stage.stage_index, - readies, - pipeline_coordinator, - stage_shard_plans, + ctx.readies, + ctx.pipeline_coordinator, + ctx.stage_shard_plans, )?; - pump(driver, stack, frame_tx); + pump(ctx.driver, ctx.stack, ctx.frame_tx); Ok(()) } -struct FailedProvisionPlugin; - -impl ProvisionPlugin for FailedProvisionPlugin { - fn start_node( - &mut self, - _spec: NodeProvisionSpec, - _sink: PluginSink, - ) -> Result { - Err("provider start worker disconnected".to_owned()) - } - - fn complete_bootstrap( - &mut self, - _handle: &crate::provisioning::PluginNodeHandle, - ) -> Result<(), String> { - Ok(()) - } - - fn stop_node(&mut self, _handle: &crate::provisioning::PluginNodeHandle) -> Result<(), String> { - Ok(()) - } -} - struct PromptWork { request: SubmitPrompt, events: mpsc::Sender, @@ -4228,10 +4052,6 @@ impl OrchStdioCapture { } } -fn install_orch_stdio_capture() -> Result>, String> { - OrchStdioCapture::install() -} - fn drain_orch_stdio_capture( rx: Option<&mpsc::Receiver>, datastream: &mut OrchDatastream, @@ -4982,13 +4802,20 @@ impl PipelinePromptRuntime { tokens: &[u32], begin_sequence: bool, ) -> Result<(), String> { - self.token_in_sender.send(encode_token_record_with_flags( - self.token_spec, - sequence, - tokens, - false, - begin_sequence, - )?) + let payload = tokens + .iter() + .flat_map(|token| token.to_le_bytes()) + .collect::>(); + let mut flags = ingress::ObjectFlags::default(); + flags.begin_sequence = begin_sequence; + self.token_in_sender.send( + ingress::ObjectRecordBuilder::new(ingress_object_spec_from_plan(self.token_spec)) + .object_id(ingress::ObjectId(9000_u64.saturating_add(sequence))) + .sequence(sequence) + .payload(payload) + .flags(flags) + .encode(), + ) } fn request_decode( @@ -5125,30 +4952,6 @@ impl PipelinePromptRuntime { } } -fn encode_token_record_with_flags( - spec: run_plan::ObjectSpec, - sequence: u64, - tokens: &[u32], - eos: bool, - begin_sequence: bool, -) -> Result, String> { - let payload = tokens - .iter() - .flat_map(|token| token.to_le_bytes()) - .collect::>(); - let mut flags = ingress::ObjectFlags::default(); - flags.end_of_sequence = eos; - flags.begin_sequence = begin_sequence; - Ok( - ingress::ObjectRecordBuilder::new(ingress_object_spec_from_plan(spec)) - .object_id(ingress::ObjectId(9000_u64.saturating_add(sequence))) - .sequence(sequence) - .payload(payload) - .flags(flags) - .encode(), - ) -} - fn ingress_object_spec_from_plan(spec: run_plan::ObjectSpec) -> ingress::ObjectSpec { let extent_alignment = match spec.kind { run_plan::ObjectKind::Token => u64::from(spec.dtype_width_bytes), @@ -5192,20 +4995,6 @@ fn take_pipeline_token_record( Ok(Some(out)) } -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -enum PromptRuntimeMode { - DirectInferPrompt, - PipelineTokenEdges, -} - -fn prompt_runtime_mode(pipeline_plan: Option<&run_plan::RunPlan>) -> PromptRuntimeMode { - if pipeline_plan.is_some() { - PromptRuntimeMode::PipelineTokenEdges - } else { - PromptRuntimeMode::DirectInferPrompt - } -} - fn serve_prompts( ctx: RuntimeReadyAckLoop<'_>, work_rx: &mpsc::Receiver, @@ -5234,16 +5023,16 @@ fn serve_prompts( provider, .. } = ctx; - let mut pipeline_runtime = match prompt_runtime_mode(pipeline_plan) { - PromptRuntimeMode::PipelineTokenEdges => Some(PipelinePromptRuntime::new( + let mut pipeline_runtime = match pipeline_plan { + Some(plan) => Some(PipelinePromptRuntime::new( driver, - pipeline_plan.expect("pipeline mode requires plan"), + plan, prompt_endpoint, tokenizer_encode_actor, tokenizer_decode_actor, tokenizer_reply_to, )?), - PromptRuntimeMode::DirectInferPrompt => None, + None => None, }; let mut active: Option = None; loop { @@ -5393,9 +5182,13 @@ fn serve_prompts( continue; } - let terminal = event.is_terminal(); - let completion_status = prompt_completion_status(&event); - let completion_detail = prompt_completion_detail(&event); + let completion = match &event { + PromptEvent::Done { .. } => Some(("ready", json!({"event":"Done"}))), + PromptEvent::Fault { error, .. } => { + Some(("failed", json!({"event":"Fault","error":error}))) + } + PromptEvent::TextDelta { .. } => None, + }; orch_datastream.emit_prompt( dashboard, run_id, @@ -5406,15 +5199,15 @@ fn serve_prompts( prompt_event_detail(&event), ); let _ = current.events.send(event); - if terminal { + if let Some((status, detail)) = completion { orch_datastream.emit_prompt( dashboard, run_id, node_id, request_id, "prompt_complete", - completion_status, - completion_detail, + status, + detail, ); active = None; } @@ -5455,42 +5248,12 @@ fn prompt_event_detail(event: &PromptEvent) -> Value { } } -fn prompt_completion_status(event: &PromptEvent) -> &'static str { - match event { - PromptEvent::Done { .. } => "ready", - PromptEvent::Fault { .. } => "failed", - PromptEvent::TextDelta { .. } => "observed", - } -} - -fn prompt_completion_detail(event: &PromptEvent) -> Value { - match event { - PromptEvent::Done { .. } => json!({"event":"Done"}), - PromptEvent::Fault { error, .. } => json!({"event":"Fault","error":error}), - PromptEvent::TextDelta { .. } => json!({"event":"TextDelta"}), - } -} - fn stop_requested(stop_rx: &mpsc::Receiver<()>) -> bool { stop_rx.try_recv().is_ok() } fn spawn_stop_listener() -> mpsc::Receiver<()> { let (tx, rx) = mpsc::channel(); - let stdin_tx = tx.clone(); - thread::spawn(move || { - let stdin = std::io::stdin(); - for line in stdin.lock().lines().map_while(Result::ok) { - let trimmed = line.trim(); - if trimmed.eq_ignore_ascii_case("stop") - || trimmed.eq_ignore_ascii_case("shutdown") - || trimmed.eq_ignore_ascii_case("quit") - { - let _ = stdin_tx.send(()); - break; - } - } - }); #[cfg(target_os = "linux")] { thread::spawn(move || { @@ -5680,7 +5443,7 @@ fn emit_swim_transitions( stack: &DistributionRuntimeStack, ) { for transition in stack.drain_swim_transitions() { - let peer = format_dist_node_id(transition.peer); + let peer = format!("{:?}", transition.peer); let from = transition.from.map(|state| format!("{:?}", state)); let to = format!("{:?}", transition.to); let member_state = stack @@ -5744,7 +5507,7 @@ fn swim_probe_event_record( let budget_ms = event.budget_ms; SwimProbeEvent { event: event.event.to_owned(), - target: format_dist_node_id(event.target), + target: format!("{:?}", event.target), sequence: event.sequence, kind: event.kind.to_owned(), rtt_ms: event.rtt_ms, @@ -5772,14 +5535,10 @@ fn swim_recent_probe_targets(stack: &DistributionRuntimeStack) -> Vec { .swim_telemetry .recent_targets() .into_iter() - .map(format_dist_node_id) + .map(|node_id| format!("{:?}", node_id)) .collect() } -fn format_dist_node_id(node_id: DistNodeId) -> String { - format!("{:?}", node_id) -} - fn duration_ms_u64(duration: Duration) -> u64 { u64::try_from(duration.as_millis()).unwrap_or(u64::MAX) } @@ -5803,48 +5562,41 @@ fn env_optional(name: &str) -> Option { .filter(|value| !value.is_empty()) } -fn docker_container_prefix() -> String { - env_optional(MVP_DOCKER_CONTAINER_PREFIX_ENV) - .unwrap_or_else(|| DEFAULT_DOCKER_CONTAINER_PREFIX.to_owned()) -} - -fn optional_env(name: &str) -> Option<(String, String)> { - env_optional(name).map(|value| (name.to_owned(), value)) -} - fn local_tinygrad_worker_env(provider: &ProviderKind) -> Option<(String, String)> { - optional_env("MVP_TINYGRAD_WORKER").or_else(|| { - if provider != &provider_kind::process() { - return None; - } - default_local_tinygrad_worker_path().map(|path| { - ( - "MVP_TINYGRAD_WORKER".to_owned(), - path.to_string_lossy().to_string(), - ) + env_optional("MVP_TINYGRAD_WORKER") + .map(|value| ("MVP_TINYGRAD_WORKER".to_owned(), value)) + .or_else(|| { + if provider != &provider_kind::process() { + return None; + } + default_local_tinygrad_worker_path().map(|path| { + ( + "MVP_TINYGRAD_WORKER".to_owned(), + path.to_string_lossy().to_string(), + ) + }) }) - }) } fn default_local_tinygrad_worker_path() -> Option { - let mut candidates = Vec::new(); - if let Ok(cwd) = std::env::current_dir() { - candidates.push(cwd.join("apps").join("mvp-node").join("tinygrad_worker.py")); + let cwd_candidate = std::env::current_dir() + .ok() + .map(|cwd| cwd.join("apps").join("mvp-node").join("tinygrad_worker.py")); + if let Some(candidate) = cwd_candidate.filter(|path| path.is_file()) { + return Some(candidate.canonicalize().unwrap_or(candidate)); } - candidates.push( - PathBuf::from(env!("CARGO_MANIFEST_DIR")) - .join("..") - .join("..") - .join("apps") - .join("mvp-node") - .join("tinygrad_worker.py"), - ); - for candidate in candidates { - if candidate.is_file() { - return Some(candidate.canonicalize().unwrap_or(candidate)); - } - } - None + + let manifest_candidate = PathBuf::from(env!("CARGO_MANIFEST_DIR")) + .join("..") + .join("..") + .join("apps") + .join("mvp-node") + .join("tinygrad_worker.py"); + manifest_candidate.is_file().then(|| { + manifest_candidate + .canonicalize() + .unwrap_or(manifest_candidate) + }) } fn resolve_vastai_ssh_identity(explicit: Option) -> Result { diff --git a/crates/mvp-system/src/orchestration/engine_builder/engine.rs b/crates/mvp-system/src/orchestration/engine_builder/engine.rs index 0369557..89eb657 100644 --- a/crates/mvp-system/src/orchestration/engine_builder/engine.rs +++ b/crates/mvp-system/src/orchestration/engine_builder/engine.rs @@ -1,27 +1,21 @@ -use std::collections::BTreeMap; - use crate::run_plan::RunId; use super::error::EngineBuildError; use super::events::EngineEvent; -use super::launcher::{ - CoordinatorJoinSpec, LaunchedNode, NodeControl, NodeFacts, NodeLaunchSpec, NodeLauncher, -}; +use super::launcher::{LaunchedNode, NodeControl, NodeFacts, NodeLaunchSpec, StaticNodeLauncher}; use super::model::ModelSpec; use super::node_image::NodeImageSpec; -use super::planner::{RoleAssignmentPlan, RolePlanner, RolePlannerInput}; -use super::pool::{PoolProvider, PoolRequest, ResourceRequest}; -use super::roles::{RoleAssignment, RoleKind}; +use super::planner::{FixedLinearPipelinePlanner, RoleAssignmentPlan, RolePlannerInput}; +use super::pool::{PoolRequest, StaticPoolProvider}; +use super::roles::RoleAssignment; pub struct ClusterBuilder { cluster_id: String, run_id: RunId, model: ModelSpec, - image: Option, - pool_provider: Option>, - launcher: Option>, - planner: Option>, - required_resources: ResourceRequest, + pool_provider: Option, + launcher: Option, + planner: Option, } impl ClusterBuilder { @@ -30,11 +24,9 @@ impl ClusterBuilder { cluster_id: cluster_id.into(), run_id: RunId(1), model, - image: None, pool_provider: None, launcher: None, planner: None, - required_resources: ResourceRequest::default(), } } @@ -43,36 +35,26 @@ impl ClusterBuilder { self } - pub fn image(mut self, image: NodeImageSpec) -> Self { - self.image = Some(image); + pub fn image(self, _image: NodeImageSpec) -> Self { self } - pub fn pool_provider(mut self, provider: impl PoolProvider + 'static) -> Self { - self.pool_provider = Some(Box::new(provider)); + pub fn pool_provider(mut self, provider: StaticPoolProvider) -> Self { + self.pool_provider = Some(provider); self } - pub fn launcher(mut self, launcher: impl NodeLauncher + 'static) -> Self { - self.launcher = Some(Box::new(launcher)); + pub fn launcher(mut self, launcher: StaticNodeLauncher) -> Self { + self.launcher = Some(launcher); self } - pub fn planner(mut self, planner: impl RolePlanner + 'static) -> Self { - self.planner = Some(Box::new(planner)); - self - } - - pub fn required_resources(mut self, required_resources: ResourceRequest) -> Self { - self.required_resources = required_resources; + pub fn planner(mut self, planner: FixedLinearPipelinePlanner) -> Self { + self.planner = Some(planner); self } pub fn launch(mut self) -> Result { - let image = self - .image - .take() - .ok_or(EngineBuildError::MissingComponent("image"))?; let pool_provider = self .pool_provider .take() @@ -88,10 +70,7 @@ impl ClusterBuilder { let mut events = Vec::new(); let leases = pool_provider.acquire_pool(PoolRequest { - cluster_id: self.cluster_id.clone(), min_nodes: planner.required_node_count(), - image: image.clone(), - required_resources: self.required_resources.clone(), })?; if leases.is_empty() { return Err(EngineBuildError::EmptyPool); @@ -103,16 +82,7 @@ impl ClusterBuilder { let mut nodes = Vec::with_capacity(leases.len()); let mut iter = leases.into_iter(); let coordinator_lease = iter.next().ok_or(EngineBuildError::EmptyPool)?; - let mut coordinator = launcher.launch_node( - &coordinator_lease, - NodeLaunchSpec { - cluster_id: self.cluster_id.clone(), - image: image.clone(), - coordinator: None, - is_coordinator: true, - env: BTreeMap::new(), - }, - )?; + let mut coordinator = launcher.launch_node(&coordinator_lease, NodeLaunchSpec); events.push(EngineEvent::NodeLaunched { node_id: coordinator.lease.logical_node_id, coordinator: true, @@ -121,26 +91,10 @@ impl ClusterBuilder { events.push(EngineEvent::NodeBootReady { node_id: coordinator_facts.node_id, }); - let coordinator_endpoint = coordinator_facts.coordinator_endpoint.clone().ok_or( - EngineBuildError::CoordinatorEndpointMissing { - node_id: coordinator_facts.node_id.0, - }, - )?; nodes.push(EngineNode::new(coordinator, coordinator_facts)); for lease in iter { - let mut node = launcher.launch_node( - &lease, - NodeLaunchSpec { - cluster_id: self.cluster_id.clone(), - image: image.clone(), - coordinator: Some(CoordinatorJoinSpec { - endpoint: coordinator_endpoint.clone(), - }), - is_coordinator: false, - env: BTreeMap::new(), - }, - )?; + let mut node = launcher.launch_node(&lease, NodeLaunchSpec); events.push(EngineEvent::NodeLaunched { node_id: node.lease.logical_node_id, coordinator: false, @@ -161,7 +115,6 @@ impl ClusterBuilder { }); let plan = planner.plan(RolePlannerInput { - cluster_id: self.cluster_id.clone(), run_id: self.run_id, model: self.model, nodes: nodes.iter().map(|node| node.facts.clone()).collect(), @@ -203,10 +156,6 @@ pub struct ClusterHandle { } impl ClusterHandle { - pub fn cluster_id(&self) -> &str { - &self.cluster_id - } - pub fn role_plan(&self) -> &RoleAssignmentPlan { &self.plan } @@ -215,17 +164,6 @@ impl ClusterHandle { &self.events } - pub fn node_summaries(&self) -> Vec { - self.nodes - .iter() - .map(|node| NodeSummary { - node_id: node.facts.node_id, - roles: node.roles.iter().map(RoleAssignment::kind).collect(), - facts: node.facts.clone(), - }) - .collect() - } - pub fn shutdown(mut self) -> Result, EngineBuildError> { for node in &mut self.nodes { node.control.shutdown()?; @@ -239,14 +177,6 @@ impl ClusterHandle { Ok(self.events) } } - -#[derive(Clone, Debug, PartialEq, Eq)] -pub struct NodeSummary { - pub node_id: crate::run_plan::NodeId, - pub roles: Vec, - pub facts: NodeFacts, -} - struct EngineNode { facts: NodeFacts, roles: Vec, diff --git a/crates/mvp-system/src/orchestration/engine_builder/error.rs b/crates/mvp-system/src/orchestration/engine_builder/error.rs index b8a20e0..d555c46 100644 --- a/crates/mvp-system/src/orchestration/engine_builder/error.rs +++ b/crates/mvp-system/src/orchestration/engine_builder/error.rs @@ -1,4 +1,3 @@ -use std::error::Error; use std::fmt; use crate::run_plan; @@ -7,10 +6,8 @@ use crate::run_plan; pub enum EngineBuildError { MissingComponent(&'static str), EmptyPool, - CoordinatorEndpointMissing { node_id: u64 }, RoleTargetMissing { node_id: u64 }, Pool(PoolError), - Launch(LaunchError), Node(NodeControlError), Planning(PlanningError), } @@ -20,37 +17,22 @@ impl fmt::Display for EngineBuildError { match self { Self::MissingComponent(name) => write!(f, "missing engine builder component: {name}"), Self::EmptyPool => write!(f, "pool provider returned no nodes"), - Self::CoordinatorEndpointMissing { node_id } => { - write!( - f, - "coordinator node {node_id} did not report a coordinator endpoint" - ) - } Self::RoleTargetMissing { node_id } => { write!(f, "role assignment targeted unknown node {node_id}") } Self::Pool(err) => err.fmt(f), - Self::Launch(err) => err.fmt(f), Self::Node(err) => err.fmt(f), Self::Planning(err) => err.fmt(f), } } } -impl Error for EngineBuildError {} - impl From for EngineBuildError { fn from(value: PoolError) -> Self { Self::Pool(value) } } -impl From for EngineBuildError { - fn from(value: LaunchError) -> Self { - Self::Launch(value) - } -} - impl From for EngineBuildError { fn from(value: NodeControlError) -> Self { Self::Node(value) @@ -66,7 +48,6 @@ impl From for EngineBuildError { #[derive(Clone, Debug, PartialEq, Eq)] pub enum PoolError { InsufficientNodes { requested: usize, available: usize }, - Provider(String), } impl fmt::Display for PoolError { @@ -79,34 +60,16 @@ impl fmt::Display for PoolError { f, "pool has {available} matching nodes, but {requested} were requested" ), - Self::Provider(message) => write!(f, "pool provider failed: {message}"), } } } -impl Error for PoolError {} - -#[derive(Clone, Debug, PartialEq, Eq)] -pub enum LaunchError { - Backend(String), -} - -impl fmt::Display for LaunchError { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - match self { - Self::Backend(message) => write!(f, "node launcher failed: {message}"), - } - } -} - -impl Error for LaunchError {} - #[derive(Clone, Debug, PartialEq, Eq)] pub enum NodeControlError { NotBooted { node_id: u64 }, Stopped { node_id: u64 }, RoleNodeMismatch { node_id: u64, role_node_id: u64 }, - Backend(String), + Backend(&'static str), } impl fmt::Display for NodeControlError { @@ -126,8 +89,6 @@ impl fmt::Display for NodeControlError { } } -impl Error for NodeControlError {} - #[derive(Clone, Debug, PartialEq, Eq)] pub enum PlanningError { DuplicateNodeId { node_id: u64 }, @@ -158,5 +119,3 @@ impl fmt::Display for PlanningError { } } } - -impl Error for PlanningError {} diff --git a/crates/mvp-system/src/orchestration/engine_builder/launcher.rs b/crates/mvp-system/src/orchestration/engine_builder/launcher.rs index 4924424..038275a 100644 --- a/crates/mvp-system/src/orchestration/engine_builder/launcher.rs +++ b/crates/mvp-system/src/orchestration/engine_builder/launcher.rs @@ -1,20 +1,11 @@ -use std::collections::{BTreeMap, BTreeSet}; +use std::collections::BTreeSet; use crate::run_plan::NodeId; -use super::error::{LaunchError, NodeControlError}; -use super::node_image::NodeImageSpec; -use super::pool::{NodeCapability, NodeLease, ResourceFacts}; +use super::error::NodeControlError; +use super::pool::{NodeCapability, NodeLease}; use super::roles::RoleAssignment; -pub trait NodeLauncher: Send + Sync { - fn launch_node( - &self, - lease: &NodeLease, - spec: NodeLaunchSpec, - ) -> Result; -} - pub trait NodeControl: Send { fn wait_boot_ready(&mut self) -> Result; fn wait_cluster_converged(&mut self, expected_alive: usize) -> Result<(), NodeControlError>; @@ -23,33 +14,18 @@ pub trait NodeControl: Send { } #[derive(Clone, Debug, PartialEq, Eq)] -pub struct NodeLaunchSpec { - pub cluster_id: String, - pub image: NodeImageSpec, - pub coordinator: Option, - pub is_coordinator: bool, - pub env: BTreeMap, -} - -#[derive(Clone, Debug, PartialEq, Eq)] -pub struct CoordinatorJoinSpec { - pub endpoint: String, -} +pub struct NodeLaunchSpec; #[derive(Clone, Debug, PartialEq, Eq)] pub struct NodeFacts { pub node_id: NodeId, - pub coordinator_endpoint: Option, - pub resources: ResourceFacts, pub capabilities: BTreeSet, } impl NodeFacts { - pub fn from_lease(lease: &NodeLease, coordinator_endpoint: Option) -> Self { + pub fn from_lease(lease: &NodeLease) -> Self { Self { node_id: lease.logical_node_id, - coordinator_endpoint, - resources: lease.expected_resources.clone(), capabilities: lease.capabilities.clone(), } } @@ -63,18 +39,10 @@ pub struct LaunchedNode { #[derive(Clone, Debug, Default)] pub struct StaticNodeLauncher; -impl NodeLauncher for StaticNodeLauncher { - fn launch_node( - &self, - lease: &NodeLease, - spec: NodeLaunchSpec, - ) -> Result { - let endpoint = format!( - "static://{}/node/{}", - spec.cluster_id, lease.logical_node_id.0 - ); - let facts = NodeFacts::from_lease(lease, Some(endpoint)); - Ok(LaunchedNode { +impl StaticNodeLauncher { + pub fn launch_node(&self, lease: &NodeLease, _spec: NodeLaunchSpec) -> LaunchedNode { + let facts = NodeFacts::from_lease(lease); + LaunchedNode { lease: lease.clone(), control: Box::new(StaticNodeControl { facts, @@ -82,7 +50,7 @@ impl NodeLauncher for StaticNodeLauncher { stopped: false, assigned_roles: Vec::new(), }), - }) + } } } @@ -123,7 +91,7 @@ impl NodeControl for StaticNodeControl { } if expected_alive == 0 { return Err(NodeControlError::Backend( - "expected_alive must be greater than zero".to_owned(), + "expected_alive must be greater than zero", )); } Ok(()) diff --git a/crates/mvp-system/src/orchestration/engine_builder/mod.rs b/crates/mvp-system/src/orchestration/engine_builder/mod.rs index 249ff25..47d462e 100644 --- a/crates/mvp-system/src/orchestration/engine_builder/mod.rs +++ b/crates/mvp-system/src/orchestration/engine_builder/mod.rs @@ -16,15 +16,14 @@ pub mod node_image; pub mod planner; pub mod pool; pub mod roles; -pub mod runtime_stack; -pub mod workload; -pub use crate::run_plan::NodeId; +pub use crate::run_plan::{DTypeFamily, NodeId}; pub use engine::ClusterBuilder; pub use events::EngineEvent; pub use launcher::StaticNodeLauncher; -pub use model::{DTypeFamily, ModelArtifact, ModelSpec}; -pub use node_image::{NodeImageSpec, WorkerRuntimeSpec}; +pub use model::{ModelArtifact, ModelSpec}; pub use planner::FixedLinearPipelinePlanner; pub use pool::{NodeCapability, NodeLease, ResourceFacts, StaticPoolProvider}; pub use roles::RoleKind; + +pub use node_image::{NodeImageSpec, WorkerRuntimeSpec}; diff --git a/crates/mvp-system/src/orchestration/engine_builder/model.rs b/crates/mvp-system/src/orchestration/engine_builder/model.rs index 527f6eb..522b2ef 100644 --- a/crates/mvp-system/src/orchestration/engine_builder/model.rs +++ b/crates/mvp-system/src/orchestration/engine_builder/model.rs @@ -3,12 +3,11 @@ use crate::run_plan; #[derive(Clone, Debug, PartialEq, Eq)] pub struct ModelSpec { pub model_id: String, - pub architecture: ModelArchitecture, pub artifact: ModelArtifact, pub tokenizer: run_plan::TokenizerSource, pub num_layers: u32, pub hidden_dim: u64, - pub dtype_family: DTypeFamily, + pub dtype_family: run_plan::DTypeFamily, pub dtype_width_bytes: u64, pub max_seq_len: u64, pub eos_token_id: u32, @@ -20,7 +19,7 @@ impl ModelSpec { artifact: ModelArtifact, num_layers: u32, hidden_dim: u64, - dtype_family: DTypeFamily, + dtype_family: run_plan::DTypeFamily, dtype_width_bytes: u64, max_seq_len: u64, eos_token_id: u32, @@ -28,7 +27,6 @@ impl ModelSpec { ) -> Self { Self { model_id: model_id.into(), - architecture: ModelArchitecture::PipelinedCausalLlm, artifact, tokenizer, num_layers, @@ -40,29 +38,13 @@ impl ModelSpec { } } - pub fn mvp_tiny_open_llm_fixture() -> Self { - Self::pipelined_causal_llm( - "mvp-tiny-open-llm-fixture", - ModelArtifact::ContainerPath { - path: "/models/mvp-tiny-open-llm.gguf".to_owned(), - }, - 4, - 8, - DTypeFamily::BFloat, - 2, - 8, - 99, - run_plan::TokenizerSource::EmbeddedGguf, - ) - } - pub fn to_run_plan_facts(&self) -> run_plan::ModelFacts { run_plan::ModelFacts { model_id: self.model_id.clone(), gguf_source: self.artifact.to_run_plan_source(), num_layers: self.num_layers, hidden_dim: self.hidden_dim, - dtype_family: self.dtype_family.into(), + dtype_family: self.dtype_family, dtype_width_bytes: self.dtype_width_bytes, max_seq_len: self.max_seq_len, eos_token_id: self.eos_token_id, @@ -71,54 +53,15 @@ impl ModelSpec { } } -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub enum ModelArchitecture { - PipelinedCausalLlm, -} - #[derive(Clone, Debug, PartialEq, Eq)] pub enum ModelArtifact { - ContainerPath { - path: String, - }, - HuggingFaceGguf { - repo: String, - file: String, - revision: Option, - }, - TestTinyLlm { - path: String, - }, + TestTinyLlm { path: String }, } impl ModelArtifact { fn to_run_plan_source(&self) -> run_plan::GgufSource { match self { - Self::ContainerPath { path } | Self::TestTinyLlm { path } => { - run_plan::GgufSource::LocalPath(path.clone()) - } - Self::HuggingFaceGguf { - repo, - file, - revision, - } => run_plan::GgufSource::HuggingFaceGguf { - repo: repo.clone(), - file: file.clone(), - revision: revision.clone(), - }, - } - } -} - -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub enum DTypeFamily { - BFloat, -} - -impl From for run_plan::DTypeFamily { - fn from(value: DTypeFamily) -> Self { - match value { - DTypeFamily::BFloat => Self::BFloat, + Self::TestTinyLlm { path } => run_plan::GgufSource::LocalPath(path.clone()), } } } diff --git a/crates/mvp-system/src/orchestration/engine_builder/node_image.rs b/crates/mvp-system/src/orchestration/engine_builder/node_image.rs index eb62429..1998728 100644 --- a/crates/mvp-system/src/orchestration/engine_builder/node_image.rs +++ b/crates/mvp-system/src/orchestration/engine_builder/node_image.rs @@ -1,28 +1,12 @@ #[derive(Clone, Debug, PartialEq, Eq)] -pub struct NodeImageSpec { - pub image: String, - pub binary: String, - pub worker_runtime: WorkerRuntimeSpec, -} +pub struct NodeImageSpec; impl NodeImageSpec { - pub fn new(image: impl Into) -> Self { - Self { - image: image.into(), - binary: "mvp-node".to_owned(), - worker_runtime: WorkerRuntimeSpec::External { - name: "node-image-default".to_owned(), - }, - } + pub fn new(_image: impl Into) -> Self { + Self } - pub fn binary(mut self, binary: impl Into) -> Self { - self.binary = binary.into(); - self - } - - pub fn worker_runtime(mut self, worker_runtime: WorkerRuntimeSpec) -> Self { - self.worker_runtime = worker_runtime; + pub fn worker_runtime(self, _worker_runtime: WorkerRuntimeSpec) -> Self { self } } @@ -30,11 +14,4 @@ impl NodeImageSpec { #[derive(Clone, Debug, PartialEq, Eq)] pub enum WorkerRuntimeSpec { DumbProcess, - TinygradCuda { - worker_script: String, - device_env: String, - }, - External { - name: String, - }, } diff --git a/crates/mvp-system/src/orchestration/engine_builder/planner.rs b/crates/mvp-system/src/orchestration/engine_builder/planner.rs index fc266da..c9c43cf 100644 --- a/crates/mvp-system/src/orchestration/engine_builder/planner.rs +++ b/crates/mvp-system/src/orchestration/engine_builder/planner.rs @@ -8,14 +8,8 @@ use super::model::ModelSpec; use super::pool::NodeCapability; use super::roles::{CoordinatorAssignment, StageAssignment}; -pub trait RolePlanner: Send + Sync { - fn required_node_count(&self) -> usize; - fn plan(&self, input: RolePlannerInput) -> Result; -} - #[derive(Clone, Debug, PartialEq, Eq)] pub struct RolePlannerInput { - pub cluster_id: String, pub run_id: RunId, pub model: ModelSpec, pub nodes: Vec, @@ -52,12 +46,12 @@ impl FixedLinearPipelinePlanner { } } -impl RolePlanner for FixedLinearPipelinePlanner { - fn required_node_count(&self) -> usize { +impl FixedLinearPipelinePlanner { + pub fn required_node_count(&self) -> usize { self.stage_count as usize + 1 } - fn plan(&self, input: RolePlannerInput) -> Result { + pub fn plan(&self, input: RolePlannerInput) -> Result { reject_duplicate_nodes(&input.nodes)?; let coordinator = input .nodes @@ -107,17 +101,12 @@ impl RolePlanner for FixedLinearPipelinePlanner { for stage_index in 0..self.stage_count { let provision = run_plan::derive_stage_provision(&run_plan, stage_index) .map_err(PlanningError::StageProjection)?; - stages.push(StageAssignment { - cluster_id: input.cluster_id.clone(), - provision, - }); + stages.push(StageAssignment { provision }); } Ok(RoleAssignmentPlan { coordinator: CoordinatorAssignment { - cluster_id: input.cluster_id, node_id: coordinator.node_id, - model: input.model, }, stages, run_plan, diff --git a/crates/mvp-system/src/orchestration/engine_builder/pool.rs b/crates/mvp-system/src/orchestration/engine_builder/pool.rs index fab56fa..83e6ba9 100644 --- a/crates/mvp-system/src/orchestration/engine_builder/pool.rs +++ b/crates/mvp-system/src/orchestration/engine_builder/pool.rs @@ -3,68 +3,33 @@ use std::collections::BTreeSet; use crate::run_plan::NodeId; use super::error::PoolError; -use super::node_image::NodeImageSpec; - -pub trait PoolProvider: Send + Sync { - fn acquire_pool(&self, request: PoolRequest) -> Result, PoolError>; -} #[derive(Clone, Debug, PartialEq, Eq)] pub struct PoolRequest { - pub cluster_id: String, pub min_nodes: usize, - pub image: NodeImageSpec, - pub required_resources: ResourceRequest, -} - -#[derive(Clone, Debug, PartialEq, Eq, Default)] -pub struct ResourceRequest { - pub min_gpu_count: u32, - pub min_gpu_memory_bytes: u64, - pub require_cuda: bool, } #[derive(Clone, Debug, PartialEq, Eq)] pub struct NodeLease { - pub lease_id: String, pub logical_node_id: NodeId, - pub launch_target: LaunchTarget, - pub expected_resources: ResourceFacts, pub capabilities: BTreeSet, } impl NodeLease { pub fn new( - lease_id: impl Into, + _lease_id: impl Into, logical_node_id: NodeId, capabilities: impl IntoIterator, ) -> Self { Self { - lease_id: lease_id.into(), logical_node_id, - launch_target: LaunchTarget::InProcess, - expected_resources: ResourceFacts::default(), capabilities: capabilities.into_iter().collect(), } } - pub fn launch_target(mut self, launch_target: LaunchTarget) -> Self { - self.launch_target = launch_target; + pub fn resources(self, _expected_resources: ResourceFacts) -> Self { self } - - pub fn resources(mut self, expected_resources: ResourceFacts) -> Self { - self.expected_resources = expected_resources; - self - } -} - -#[derive(Clone, Debug, PartialEq, Eq)] -pub enum LaunchTarget { - InProcess, - LocalProcess { program: String, args: Vec }, - DockerContainer { name: String }, - RemoteHost { label: String }, } #[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] @@ -73,46 +38,12 @@ pub enum NodeCapability { Worker, } -#[derive(Clone, Debug, PartialEq, Eq)] -pub struct ResourceFacts { - pub gpu_count: u32, - pub gpu_memory_bytes: u64, - pub cpu_cores: u32, - pub ram_bytes: u64, - pub cuda_available: bool, -} +#[derive(Clone, Debug, Default, PartialEq, Eq)] +pub struct ResourceFacts; impl ResourceFacts { - pub fn cpu_only(cpu_cores: u32, ram_bytes: u64) -> Self { - Self { - gpu_count: 0, - gpu_memory_bytes: 0, - cpu_cores, - ram_bytes, - cuda_available: false, - } - } - - pub fn cuda(gpu_count: u32, gpu_memory_bytes: u64, cpu_cores: u32, ram_bytes: u64) -> Self { - Self { - gpu_count, - gpu_memory_bytes, - cpu_cores, - ram_bytes, - cuda_available: true, - } - } - - fn satisfies(&self, request: &ResourceRequest) -> bool { - self.gpu_count >= request.min_gpu_count - && self.gpu_memory_bytes >= request.min_gpu_memory_bytes - && (!request.require_cuda || self.cuda_available) - } -} - -impl Default for ResourceFacts { - fn default() -> Self { - Self::cpu_only(1, 512 * 1024 * 1024) + pub fn cpu_only(_cpu_cores: u32, _ram_bytes: u64) -> Self { + Self } } @@ -125,30 +56,16 @@ impl StaticPoolProvider { pub fn new(leases: Vec) -> Self { Self { leases } } - - pub fn leases(&self) -> &[NodeLease] { - &self.leases - } } -impl PoolProvider for StaticPoolProvider { - fn acquire_pool(&self, request: PoolRequest) -> Result, PoolError> { - let matching = self - .leases - .iter() - .filter(|lease| { - lease - .expected_resources - .satisfies(&request.required_resources) - }) - .cloned() - .collect::>(); - if matching.len() < request.min_nodes { +impl StaticPoolProvider { + pub fn acquire_pool(&self, request: PoolRequest) -> Result, PoolError> { + if self.leases.len() < request.min_nodes { return Err(PoolError::InsufficientNodes { requested: request.min_nodes, - available: matching.len(), + available: self.leases.len(), }); } - Ok(matching) + Ok(self.leases.clone()) } } diff --git a/crates/mvp-system/src/orchestration/engine_builder/roles.rs b/crates/mvp-system/src/orchestration/engine_builder/roles.rs index fd0c94a..d1d7ccf 100644 --- a/crates/mvp-system/src/orchestration/engine_builder/roles.rs +++ b/crates/mvp-system/src/orchestration/engine_builder/roles.rs @@ -1,17 +1,12 @@ use crate::run_plan::{self, NodeId}; -use super::model::ModelSpec; - #[derive(Clone, Debug, PartialEq, Eq)] pub struct CoordinatorAssignment { - pub cluster_id: String, pub node_id: NodeId, - pub model: ModelSpec, } #[derive(Clone, Debug, PartialEq, Eq)] pub struct StageAssignment { - pub cluster_id: String, pub provision: run_plan::ProvisionStage, } diff --git a/crates/mvp-system/src/orchestration/engine_builder/runtime_stack.rs b/crates/mvp-system/src/orchestration/engine_builder/runtime_stack.rs deleted file mode 100644 index eb41e65..0000000 --- a/crates/mvp-system/src/orchestration/engine_builder/runtime_stack.rs +++ /dev/null @@ -1,173 +0,0 @@ -use std::thread; -use std::time::{Duration, Instant}; - -use distribution::node::DistributedNodeConfig; -use distribution::types::{DirectoryEntry, NodeId}; -use iroh::EndpointAddr; -use iroh_driver::{IrohDriver, IrohDriverConfig}; -use swactor::actor::ActorAddress; - -use crate::node_actor::NodeAgentActor; -use crate::orchestration::actor::OrchestratorActor; -use crate::orchestration::distribution_stack::DistributionRuntimeStack; -use crate::run_fsm as orchestrator_core; -use crate::staging as stage_core; -use crate::transport::codec_registry::register_mvp_actor_codecs; - -pub struct RuntimeNodeConfig { - pub distributed: DistributedNodeConfig, - pub relay_mode: iroh::RelayMode, -} - -impl Default for RuntimeNodeConfig { - fn default() -> Self { - Self { - distributed: DistributedNodeConfig::default(), - relay_mode: iroh::RelayMode::Disabled, - } - } -} - -pub struct RuntimeNode { - _tokio: tokio::runtime::Runtime, - driver: IrohDriver, - stack: DistributionRuntimeStack, -} - -impl RuntimeNode { - pub fn start_default() -> Result { - Self::start_with_codecs(RuntimeNodeConfig::default(), |_| {}) - } - - pub fn start_with_codecs( - config: RuntimeNodeConfig, - extend_codecs: impl FnOnce(&mut swactor_transport::CodecRegistry), - ) -> Result { - let tokio = tokio::runtime::Runtime::new() - .map_err(|err| RuntimeNodeError::Start(format!("tokio runtime: {err}")))?; - let mut driver = IrohDriver::with_handle( - tokio.handle().clone(), - IrohDriverConfig { - secret_key: None, - relay_mode: config.relay_mode, - node: config.distributed.clone(), - peer_auth: None, - additional_alpns: vec![], - }, - ) - .map_err(|err| RuntimeNodeError::Start(format!("iroh driver: {err}")))?; - let stack = DistributionRuntimeStack::new_with_codecs( - driver.node_id(), - config.distributed, - |registry| { - register_mvp_actor_codecs(registry); - extend_codecs(registry); - }, - ); - driver.enable_actor_bridge( - stack.runtime.clone(), - stack.codec.clone(), - stack.actor_bridge_routes(), - stack.actors.swim, - stack.relay_mirror.clone(), - stack.route_view.clone(), - ); - Ok(Self { - _tokio: tokio, - driver, - stack, - }) - } - - pub fn node_id(&self) -> NodeId { - self.driver.node_id() - } - - pub fn endpoint_addr(&self) -> EndpointAddr { - self.driver.endpoint_addr() - } - - pub fn join(&mut self, coordinators: &[EndpointAddr]) { - self.driver.join(coordinators); - } - - pub fn register_actor_route(&mut self, actor_addr: ActorAddress, generation: u64) { - let entry = self.driver.register_actor(actor_addr, generation); - self.stack.register_local_actor(entry); - } - - pub fn register_directory_entry(&self, entry: DirectoryEntry) { - self.stack.register_local_actor(entry); - } - - pub fn spawn_orchestrator_actor( - &mut self, - config: orchestrator_core::RunConfig, - report_to: Option, - ) -> Result { - let actor = self - .stack - .runtime - .spawn(OrchestratorActor::new(config, report_to)) - .map_err(|err| RuntimeNodeError::Start(format!("spawn orchestrator actor: {err}")))?; - self.register_actor_route(actor, 1); - Ok(actor) - } - - pub fn spawn_node_agent_actor( - &mut self, - local_node_id: stage_core::NodeId, - orchestrator: ActorAddress, - report_to: Option, - ) -> Result { - let actor = self - .stack - .runtime - .spawn(NodeAgentActor::new(local_node_id, orchestrator, report_to)) - .map_err(|err| RuntimeNodeError::Start(format!("spawn node agent actor: {err}")))?; - self.register_actor_route(actor, 1); - Ok(actor) - } - - pub fn pump_once(&mut self) { - self.stack.tick_protocol_actors(Instant::now()); - self.driver.pump_inbound_to_actors(); - self.stack.pump_runtime_once(); - self.driver.drain_outbox(&self.stack.outbox); - } - - pub fn wait_for_routes(&mut self, actors: &[ActorAddress]) -> Result<(), RuntimeNodeError> { - loop { - self.pump_once(); - let ready = self - .stack - .route_view - .read() - .map(|view| actors.iter().all(|actor| view.contains_key(actor))) - .unwrap_or(false); - if ready { - return Ok(()); - } - thread::sleep(Duration::from_millis(20)); - } - } - - pub fn alive_count(&self) -> usize { - self.stack.alive_count() - } -} - -#[derive(Clone, Debug, PartialEq, Eq)] -pub enum RuntimeNodeError { - Start(String), -} - -impl std::fmt::Display for RuntimeNodeError { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - match self { - Self::Start(message) => write!(f, "runtime node start failed: {message}"), - } - } -} - -impl std::error::Error for RuntimeNodeError {} diff --git a/crates/mvp-system/src/orchestration/engine_builder/workload.rs b/crates/mvp-system/src/orchestration/engine_builder/workload.rs deleted file mode 100644 index fe36905..0000000 --- a/crates/mvp-system/src/orchestration/engine_builder/workload.rs +++ /dev/null @@ -1,13 +0,0 @@ -use super::engine::ClusterHandle; - -pub trait WorkloadAdapter { - type Input; - type Output; - type Error; - - fn submit( - &self, - cluster: &mut ClusterHandle, - input: Self::Input, - ) -> Result; -} diff --git a/crates/mvp-system/src/orchestration/mod.rs b/crates/mvp-system/src/orchestration/mod.rs index 7f2dfcb..f773a53 100644 --- a/crates/mvp-system/src/orchestration/mod.rs +++ b/crates/mvp-system/src/orchestration/mod.rs @@ -6,7 +6,7 @@ //! stays separate from local/Docker/VastAI implementation details. pub mod actor; -mod app; +pub(crate) mod app; pub mod config; pub mod distribution_stack; #[cfg(test)] @@ -16,20 +16,3 @@ pub mod provider_adapters { pub mod relay; pub(super) mod vastai; } - -pub(super) fn run_from_args(args: I) -> Result<(), String> -where - I: IntoIterator, -{ - app::run_from_args(args) -} - -pub(super) fn run_in_process_from_args( - args: I, - stop_rx: std::sync::mpsc::Receiver<()>, -) -> Result<(), String> -where - I: IntoIterator, -{ - app::run_in_process_from_args(args, stop_rx) -} diff --git a/crates/mvp-system/src/orchestration/run_plan.rs b/crates/mvp-system/src/orchestration/run_plan.rs index 32b2480..471ea31 100644 --- a/crates/mvp-system/src/orchestration/run_plan.rs +++ b/crates/mvp-system/src/orchestration/run_plan.rs @@ -137,14 +137,6 @@ impl RuntimeConfig { token_output_policy: TokenOutputPolicy::EmitAll, } } - - fn plan(&self) -> RuntimePlan { - RuntimePlan { - prompt: self.prompt.clone(), - sampling: self.sampling, - token_output_policy: self.token_output_policy, - } - } } #[derive(Clone, Copy, Debug, PartialEq, Eq)] @@ -389,7 +381,11 @@ pub fn plan_run(input: PlannerInput) -> Result { validate_global_input(&input)?; let placements = validated_placements(&input)?; let model = model_plan(&input.model)?; - let runtime = input.runtime.plan(); + let runtime = RuntimePlan { + prompt: input.runtime.prompt.clone(), + sampling: input.runtime.sampling, + token_output_policy: input.runtime.token_output_policy, + }; let max_tokens = input.runtime.max_tokens; let gguf_source = model.gguf_source.clone(); let hidden_dim = model.hidden_dim;