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

stash: tested working yoke grind runs

Browse source
This commit is contained in:
Zachery Aaron Shores-Chmielewski 2026-07-29 21:09:26 +04:00
parent 599e8e870b
commit e8821c85be
21 changed files with 1073 additions and 2148 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

7
crates/mvp-system/src/chat/config.rs
View file

@ -1,7 +0,0 @@
pub const DEFAULT_CONFIG_PATH: &str = ".config/config.toml";
pub fn normalize_optional(value: Option<String>) -> Option<String> {
value
.map(|value| value.trim().to_owned())
.filter(|value| !value.is_empty())
}

10
crates/mvp-system/src/chat/mod.rs
View file

@ -1,12 +1,4 @@
//! MVP operator chat wrapper public surface. //! MVP operator chat wrapper public surface.
pub mod config;
mod node_image; mod node_image;
mod runtime; pub(super) mod runtime;
pub(super) fn run_from_args<I>(args: I) -> std::process::ExitCode
where
I: IntoIterator<Item = String>,
{
runtime::run_from_args(args)
}

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

@ -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_WORKER_HASH_LABEL: &str = "org.swactor.mvp.node.worker-hash";
const NODE_IMAGE_BASE_HASH_LABEL: &str = "org.swactor.mvp.node.base-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 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) struct NodeImageRequest {
pub(super) requested_image: String, pub(super) requested_image: String,
pub(super) base_image: String, pub(super) base_image: String,
pub(super) node_bin: PathBuf, pub(super) node_bin: PathBuf,
pub(super) provider: NodeImageProvider, pub(super) requires_registry_image: bool,
pub(super) extra_tag: Option<String>, pub(super) extra_tag: Option<String>,
pub(super) push: bool,
pub(super) force_refresh: bool, pub(super) force_refresh: bool,
pub(super) enabled: bool, pub(super) enabled: bool,
} }
@ -84,62 +67,17 @@ pub(super) trait NodeImageProgressSink {
fn emit(&mut self, event: NodeImageProgressEvent); fn emit(&mut self, event: NodeImageProgressEvent);
} }
impl<F> 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<Option<BTreeMap<String, String>>, 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<Vec<(String, String)>, String>;
fn docker_image_remove(&mut self, root: &Path, image_ref: &str) -> Result<(), String>;
}
struct RealImageCommandRunner;
pub(super) fn prepare_node_image_with_progress( pub(super) fn prepare_node_image_with_progress(
request: NodeImageRequest, request: NodeImageRequest,
progress: Option<&mut dyn NodeImageProgressSink>, progress: Option<&mut dyn NodeImageProgressSink>,
) -> Result<String, String> { ) -> Result<String, String> {
let mut progress = progress; let mut progress = progress;
let mut runner = RealImageCommandRunner; prepare_node_image_inner(request, &mut progress)
prepare_node_image_inner(request, &mut progress, &mut runner)
} }
fn prepare_node_image_inner( fn prepare_node_image_inner(
request: NodeImageRequest, request: NodeImageRequest,
progress: &mut Option<&mut dyn NodeImageProgressSink>, progress: &mut Option<&mut dyn NodeImageProgressSink>,
runner: &mut dyn ImageCommandRunner,
) -> Result<String, String> { ) -> Result<String, String> {
emit_image_reference(progress, "requested", &request.requested_image); emit_image_reference(progress, "requested", &request.requested_image);
if !request.enabled { if !request.enabled {
@ -149,7 +87,16 @@ fn prepare_node_image_inner(
emit_image_reference(progress, "base", &request.base_image); emit_image_reference(progress, "base", &request.base_image);
let root = workspace_root()?; let root = workspace_root()?;
let image = ImageName::parse(&request.requested_image)?; 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!( return Err(format!(
"VastAI node image {:?} must include a registry namespace", "VastAI node image {:?} must include a registry namespace",
image.repository image.repository
@ -157,7 +104,6 @@ fn prepare_node_image_inner(
} }
run_status( run_status(
runner,
progress, progress,
&root, &root,
"cargo", "cargo",
@ -178,47 +124,54 @@ fn prepare_node_image_inner(
let tag = image_version_tag(&root, &image_content_hash)?; let tag = image_version_tag(&root, &image_content_hash)?;
let image_ref = image.ref_for_tag(&tag); let image_ref = image.ref_for_tag(&tag);
emit_image_reference(progress, "resolved", &image_ref); emit_image_reference(progress, "resolved", &image_ref);
let worker_hash = file_content_hash(&root, Path::new("apps/mvp-node/tinygrad_worker.py"))?; let worker_hash = hash_relative_files(
let expected_node_labels = &root,
node_image_labels(&tag, &image_content_hash, &worker_hash, &base_hash); vec![relative_path(
let expected_base_labels = base_image_labels(&base_hash); &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)?; let alias_tags = alias_tags(&image, request.extra_tag.as_deref(), &tag)?;
for alias in alias_refs(&image, &alias_tags) { for alias in alias_refs(&image, &alias_tags) {
emit_image_reference(progress, "alias", &alias); 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 = let local_image_matches = docker_image_labels_match(&root, &image_ref, &expected_node_labels)?;
docker_image_labels_match(runner, &root, &image_ref, &expected_node_labels)?; let remote_available = remote_required && docker_manifest_exists(&root, &image_ref);
let remote_available = remote_required && runner.docker_manifest_exists(&root, &image_ref);
if !request.force_refresh && remote_required && remote_available { if !request.force_refresh && remote_required && remote_available {
ensure_aliases_for_remote(runner, progress, &root, &image_ref, &image, &alias_tags)?; ensure_aliases_for_remote(progress, &root, &image_ref, &image, &alias_tags)?;
prune_old_dirty_images(runner, &root, &image, &tag); prune_old_dirty_images(&root, &image, &tag);
return Ok(image_ref); return Ok(image_ref);
} }
if !request.force_refresh && remote_required && local_image_matches { if !request.force_refresh && remote_required && local_image_matches {
ensure_aliases_local(runner, progress, &root, &image_ref, &image, &alias_tags)?; ensure_aliases_local(progress, &root, &image_ref, &image, &alias_tags)?;
push_image(runner, progress, &root, &image_ref)?; push_image(progress, &root, &image_ref)?;
for alias in alias_refs(&image, &alias_tags) { 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); return Ok(image_ref);
} }
if !request.force_refresh && !remote_required && local_image_matches { if !request.force_refresh && !remote_required && local_image_matches {
ensure_aliases_local(runner, progress, &root, &image_ref, &image, &alias_tags)?; ensure_aliases_local(progress, &root, &image_ref, &image, &alias_tags)?;
prune_old_dirty_images(runner, &root, &image, &tag); prune_old_dirty_images(&root, &image, &tag);
return Ok(image_ref); return Ok(image_ref);
} }
let base_image_matches = 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 { if !base_image_matches {
run_status_vec( run_status_command(
runner,
progress,
&root, &root,
"docker", "docker",
vec![ &vec![
"build".to_owned(), "build".to_owned(),
"-f".to_owned(), "-f".to_owned(),
"apps/mvp-node/Dockerfile.base".to_owned(), "apps/mvp-node/Dockerfile.base".to_owned(),
@ -230,6 +183,7 @@ fn prepare_node_image_inner(
], ],
"build mvp node base image", "build mvp node base image",
Some(&request.base_image), Some(&request.base_image),
progress,
)?; )?;
} }
@ -248,25 +202,24 @@ fn prepare_node_image_inner(
build_args.push(format!("{key}={value}")); build_args.push(format!("{key}={value}"));
} }
build_args.extend(["-t".to_owned(), image_ref.clone(), ".".to_owned()]); build_args.extend(["-t".to_owned(), image_ref.clone(), ".".to_owned()]);
run_status_vec( run_status_command(
runner,
progress,
&root, &root,
"docker", "docker",
build_args, &build_args,
"build mvp node image", "build mvp node image",
Some(&image_ref), 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 { if remote_required {
push_image(runner, progress, &root, &image_ref)?; push_image(progress, &root, &image_ref)?;
for alias in alias_refs(&image, &alias_tags) { 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) Ok(image_ref)
} }
@ -290,7 +243,10 @@ fn workspace_root() -> Result<PathBuf, String> {
} }
fn image_version_tag(root: &Path, image_content_hash: &str) -> Result<String, String> { fn image_version_tag(root: &Path, image_content_hash: &str) -> Result<String, String> {
if git_worktree_clean(root)? { if git_capture(root, &["status", "--porcelain"])?
.trim()
.is_empty()
{
let sha = git_capture(root, &["rev-parse", "--short=12", "HEAD"])?; let sha = git_capture(root, &["rev-parse", "--short=12", "HEAD"])?;
Ok(format!("git-{}", sha.trim())) Ok(format!("git-{}", sha.trim()))
} else { } else {
@ -298,12 +254,6 @@ fn image_version_tag(root: &Path, image_content_hash: &str) -> Result<String, St
} }
} }
fn git_worktree_clean(root: &Path) -> Result<bool, String> {
Ok(git_capture(root, &["status", "--porcelain"])?
.trim()
.is_empty())
}
fn git_capture(root: &Path, args: &[&str]) -> Result<String, String> { fn git_capture(root: &Path, args: &[&str]) -> Result<String, String> {
let output = Command::new("git") let output = Command::new("git")
.current_dir(root) .current_dir(root)
@ -355,10 +305,6 @@ fn content_hash_for_inputs(root: &Path, inputs: &[&str]) -> Result<String, Strin
hash_relative_files(root, files) hash_relative_files(root, files)
} }
fn file_content_hash(root: &Path, path: &Path) -> Result<String, String> {
hash_relative_files(root, vec![relative_path(root, &root.join(path))?])
}
fn hash_relative_files(root: &Path, files: Vec<PathBuf>) -> Result<String, String> { fn hash_relative_files(root: &Path, files: Vec<PathBuf>) -> Result<String, String> {
hash_relative_files_with_salts(root, files, &[]) hash_relative_files_with_salts(root, files, &[])
} }
@ -409,7 +355,7 @@ fn collect_hash_inputs(root: &Path, path: &Path, out: &mut Vec<PathBuf>) -> Resu
let display = display_workspace_path(root, path); let display = display_workspace_path(root, path);
let metadata = fs::metadata(path).map_err(|e| format!("stat {display}: {e}"))?; let metadata = fs::metadata(path).map_err(|e| format!("stat {display}: {e}"))?;
if metadata.is_file() { 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)?); out.push(relative_path(root, path)?);
} }
return Ok(()); 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( fn alias_tags(
image: &ImageName, image: &ImageName,
extra_tag: Option<&str>, extra_tag: Option<&str>,
@ -493,26 +435,7 @@ fn insert_alias_tag(
Ok(()) 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( fn ensure_aliases_local(
runner: &mut dyn ImageCommandRunner,
progress: &mut Option<&mut dyn NodeImageProgressSink>, progress: &mut Option<&mut dyn NodeImageProgressSink>,
root: &Path, root: &Path,
source_ref: &str, source_ref: &str,
@ -522,7 +445,6 @@ fn ensure_aliases_local(
for alias in alias_refs(image, alias_tags) { for alias in alias_refs(image, alias_tags) {
if alias != source_ref { if alias != source_ref {
run_status( run_status(
runner,
progress, progress,
root, root,
"docker", "docker",
@ -536,7 +458,6 @@ fn ensure_aliases_local(
} }
fn ensure_aliases_for_remote( fn ensure_aliases_for_remote(
runner: &mut dyn ImageCommandRunner,
progress: &mut Option<&mut dyn NodeImageProgressSink>, progress: &mut Option<&mut dyn NodeImageProgressSink>,
root: &Path, root: &Path,
source_ref: &str, source_ref: &str,
@ -546,9 +467,8 @@ fn ensure_aliases_for_remote(
if alias_tags.is_empty() { if alias_tags.is_empty() {
return Ok(false); return Ok(false);
} }
if !runner.docker_image_exists(root, source_ref) { if !docker_image_exists(root, source_ref) {
run_status( run_status(
runner,
progress, progress,
root, root,
"docker", "docker",
@ -557,9 +477,9 @@ fn ensure_aliases_for_remote(
Some(source_ref), 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) { for alias in alias_refs(image, alias_tags) {
push_image(runner, progress, root, &alias)?; push_image(progress, root, &alias)?;
} }
Ok(true) Ok(true)
} }
@ -572,13 +492,11 @@ fn alias_refs(image: &ImageName, alias_tags: &BTreeSet<String>) -> Vec<String> {
} }
fn push_image( fn push_image(
runner: &mut dyn ImageCommandRunner,
progress: &mut Option<&mut dyn NodeImageProgressSink>, progress: &mut Option<&mut dyn NodeImageProgressSink>,
root: &Path, root: &Path,
image_ref: &str, image_ref: &str,
) -> Result<(), String> { ) -> Result<(), String> {
run_status( run_status(
runner,
progress, progress,
root, root,
"docker", "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( fn docker_image_labels_match(
runner: &mut dyn ImageCommandRunner,
root: &Path, root: &Path,
image_ref: &str, image_ref: &str,
expected: &[(&str, &str)], expected: &[(&str, &str)],
) -> Result<bool, String> { ) -> Result<bool, String> {
let Some(labels) = runner.docker_image_labels(root, image_ref)? else { let Some(labels) = docker_image_labels(root, image_ref)? else {
return Ok(false); return Ok(false);
}; };
Ok(expected Ok(expected
@ -614,54 +519,18 @@ fn docker_image_labels_match(
.all(|(key, value)| labels.get(*key).map(String::as_str) == Some(*value))) .all(|(key, value)| labels.get(*key).map(String::as_str) == Some(*value)))
} }
fn docker_image_labels( fn prune_old_dirty_images(root: &Path, image: &ImageName, keep_tag: &str) {
root: &Path, let prune_enabled = std::env::var("MVP_NODE_IMAGE_PRUNE")
image_ref: &str, .map(|value| {
) -> Result<Option<BTreeMap<String, String>>, String> { let value = value.trim().to_ascii_lowercase();
let output = Command::new("docker") !matches!(value.as_str(), "0" | "false" | "no" | "off")
.current_dir(root) })
.args([ .unwrap_or(true);
"image", if !prune_enabled {
"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<BTreeMap<String, String>> = 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() {
return; return;
} }
let tags = match runner.docker_image_tags(root, &image.repository) { let tags = match docker_image_tags(root, &image.repository) {
Ok(tags) => tags, Ok(tags) => tags,
Err(error) => { Err(error) => {
eprintln!("mvp-node-image: prune old dirty images skipped: {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::<usize>().ok())
.unwrap_or(3);
let mut retained_old = 0_usize; let mut retained_old = 0_usize;
for (repository, tag) in tags { for (repository, tag) in tags {
if repository != image.repository if repository != image.repository
@ -681,7 +553,7 @@ fn prune_old_dirty_images(
} }
let image_ref = image.ref_for_tag(&tag); 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; continue;
}; };
if labels.get(NODE_IMAGE_TAG_LABEL).map(String::as_str) != Some(tag.as_str()) 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; continue;
} }
if runner.docker_image_has_container(root, &image_ref) { if docker_image_has_container(root, &image_ref) {
eprintln!( eprintln!(
"mvp-node-image: prune old dirty image {image_ref} skipped: container exists" "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}"); 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}"); 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::<usize>().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( fn run_status(
runner: &mut dyn ImageCommandRunner,
progress: &mut Option<&mut dyn NodeImageProgressSink>, progress: &mut Option<&mut dyn NodeImageProgressSink>,
root: &Path, root: &Path,
program: &str, program: &str,
@ -752,19 +590,7 @@ fn run_status(
image_ref: Option<&str>, image_ref: Option<&str>,
) -> Result<(), String> { ) -> Result<(), String> {
let args = args.iter().map(|arg| (*arg).to_owned()).collect::<Vec<_>>(); let args = args.iter().map(|arg| (*arg).to_owned()).collect::<Vec<_>>();
runner.run_status(root, program, &args, label, image_ref, progress) run_status_command(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<String>,
label: &str,
image_ref: Option<&str>,
) -> Result<(), String> {
runner.run_status(root, program, &args, label, image_ref, progress)
} }
fn emit_image_reference( fn emit_image_reference(
@ -857,16 +683,14 @@ fn drain_command_lines(
} }
} }
impl ImageCommandRunner for RealImageCommandRunner { fn run_status_command(
fn run_status(
&mut self,
root: &Path, root: &Path,
program: &str, program: &str,
args: &[String], args: &[String],
label: &str, label: &str,
image_ref: Option<&str>, image_ref: Option<&str>,
progress: &mut Option<&mut dyn NodeImageProgressSink>, progress: &mut Option<&mut dyn NodeImageProgressSink>,
) -> Result<(), String> { ) -> Result<(), String> {
eprintln!("mvp-node-image: {label}"); eprintln!("mvp-node-image: {label}");
if progress.is_none() { if progress.is_none() {
let status = Command::new(program) let status = Command::new(program)
@ -984,33 +808,75 @@ impl ImageCommandRunner for RealImageCommandRunner {
} else { } else {
Err(format!("{label} failed with {status_text}")) Err(format!("{label} failed with {status_text}"))
} }
} }
fn docker_image_exists(&mut self, root: &Path, image_ref: &str) -> bool { fn docker_image_exists(root: &Path, image_ref: &str) -> bool {
docker_image_exists(root, image_ref) 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( fn docker_image_labels(
&mut self,
root: &Path, root: &Path,
image_ref: &str, image_ref: &str,
) -> Result<Option<BTreeMap<String, String>>, String> { ) -> Result<Option<BTreeMap<String, String>>, String> {
docker_image_labels(root, image_ref) 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<BTreeMap<String, String>> = 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(&mut self, root: &Path, image_ref: &str) -> bool { fn docker_manifest_exists(root: &Path, image_ref: &str) -> bool {
docker_manifest_exists(root, image_ref) 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(&mut self, root: &Path, image_ref: &str) -> bool { fn docker_image_has_container(root: &Path, image_ref: &str) -> bool {
docker_image_has_container(root, image_ref) 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( fn docker_image_tags(root: &Path, repository: &str) -> Result<Vec<(String, String)>, String> {
&mut self,
root: &Path,
repository: &str,
) -> Result<Vec<(String, String)>, String> {
let output = Command::new("docker") let output = Command::new("docker")
.current_dir(root) .current_dir(root)
.args([ .args([
@ -1034,9 +900,9 @@ impl ImageCommandRunner for RealImageCommandRunner {
Some((repository.to_owned(), tag.to_owned())) Some((repository.to_owned(), tag.to_owned()))
}) })
.collect()) .collect())
} }
fn docker_image_remove(&mut self, root: &Path, image_ref: &str) -> Result<(), String> { fn docker_image_remove(root: &Path, image_ref: &str) -> Result<(), String> {
let status = Command::new("docker") let status = Command::new("docker")
.current_dir(root) .current_dir(root)
.args(["image", "rm", image_ref]) .args(["image", "rm", image_ref])
@ -1050,12 +916,6 @@ impl ImageCommandRunner for RealImageCommandRunner {
} else { } else {
Err(format!("docker image rm failed with {status}")) Err(format!("docker image rm failed with {status}"))
} }
}
}
fn looks_registry_reachable(repository: &str) -> bool {
let first = repository.split('/').next().unwrap_or(repository);
repository.contains('/') || first.contains('.') || first.contains(':') || first == "localhost"
} }
#[derive(Clone, Debug)] #[derive(Clone, Debug)]

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

@ -22,10 +22,9 @@ use signal_hook::consts::signal::{SIGINT, SIGTERM};
#[cfg(target_os = "linux")] #[cfg(target_os = "linux")]
use signal_hook::iterator::Signals; use signal_hook::iterator::Signals;
use crate::chat::config as chat_config;
use crate::chat::node_image::{ use crate::chat::node_image::{
NodeImageProgressEvent, NodeImageProgressEventKind, NodeImageProgressSink, NodeImageProvider, NodeImageProgressEvent, NodeImageProgressEventKind, NodeImageProgressSink, NodeImageRequest,
NodeImageRequest, prepare_node_image_with_progress, prepare_node_image_with_progress,
}; };
use crate::node_provisioning::{ProviderKind, provider_kind}; use crate::node_provisioning::{ProviderKind, provider_kind};
use crate::observability::{benchmark, frame_archive::FrameArchive}; use crate::observability::{benchmark, frame_archive::FrameArchive};
@ -37,6 +36,7 @@ use crate::{
DEFAULT_PIPELINE_CACHED_MODEL_MAX_CONTEXT, DEFAULT_PIPELINE_CACHED_MODEL_REPO, 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 DEFAULT_RPC_ADDR: &str = "127.0.0.1:19777";
const BASE_NODE_IMAGE: &str = "swactor-mvp-node-base:cuda12.6"; const BASE_NODE_IMAGE: &str = "swactor-mvp-node-base:cuda12.6";
const REPO_MODEL_CACHE_DIR: &str = ".model-cache"; const REPO_MODEL_CACHE_DIR: &str = ".model-cache";
@ -79,11 +79,11 @@ enum PromptInput {
static STOP_REQUESTED: AtomicBool = AtomicBool::new(false); static STOP_REQUESTED: AtomicBool = AtomicBool::new(false);
static PROMPT_STOP_TX: Mutex<Option<mpsc::Sender<PromptInput>>> = Mutex::new(None); static PROMPT_STOP_TX: Mutex<Option<mpsc::Sender<PromptInput>>> = Mutex::new(None);
pub(super) fn run_from_args<I>(args: I) -> ExitCode pub(crate) fn run_from_args<I>(args: I) -> ExitCode
where where
I: IntoIterator<Item = String>, I: IntoIterator<Item = String>,
{ {
match run_from_args_result(args) { match install_signal_handlers().and_then(|()| run(args)) {
Ok(()) => ExitCode::SUCCESS, Ok(()) => ExitCode::SUCCESS,
Err(error) => { Err(error) => {
eprintln!("mvp-chat: {error}"); eprintln!("mvp-chat: {error}");
@ -92,23 +92,6 @@ where
} }
} }
fn run_from_args_result<I>(args: I) -> Result<(), String>
where
I: IntoIterator<Item = String>,
{
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 { struct RuntimeEnvGuard {
name: &'static str, name: &'static str,
original: Option<std::ffi::OsString>, original: Option<std::ffi::OsString>,
@ -142,8 +125,11 @@ where
I: IntoIterator<Item = String>, I: IntoIterator<Item = String>,
{ {
let provided_args = args.into_iter().collect::<Vec<_>>(); let provided_args = args.into_iter().collect::<Vec<_>>();
if is_help_request(&provided_args) { if provided_args
print_usage(); .iter()
.any(|arg| matches!(arg.as_str(), "--help" | "-h" | "help"))
{
println!("{MVP_CHAT_USAGE}");
return Ok(()); return Ok(());
} }
let config = Config::from_args(provided_args)?; let config = Config::from_args(provided_args)?;
@ -164,7 +150,8 @@ where
); );
progress.emit_benchmark_envelope(&config); progress.emit_benchmark_envelope(&config);
progress.emit_endpoint_config_snapshot(&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(); let prepare_runtime_started = Instant::now();
progress.emit( progress.emit(
CHAT_RUNTIME_CHANNEL, CHAT_RUNTIME_CHANNEL,
@ -174,7 +161,7 @@ where
); );
let image_ref = match prepare_runtime_with_progress( let image_ref = match prepare_runtime_with_progress(
&config, &config,
prepare_node_image_progress_adapter, prepare_node_image_with_progress,
Some(&mut progress), Some(&mut progress),
) { ) {
Ok(image_ref) => { Ok(image_ref) => {
@ -318,7 +305,36 @@ where
return Err(error); 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( progress.emit(
CHAT_LIFECYCLE_CHANNEL, CHAT_LIFECYCLE_CHANNEL,
"shutdown", "shutdown",
@ -739,13 +755,8 @@ struct ChatModelConfig {
max_context: Option<u32>, max_context: Option<u32>,
} }
#[derive(Clone, Debug)] fn load_chat_config(path: Option<&Path>) -> Result<ChatTomlConfig, String> {
struct LoadedChatTomlConfig { Ok(match path {
overlay: ChatTomlConfig,
}
fn load_chat_config(path: Option<&Path>) -> Result<LoadedChatTomlConfig, String> {
let overlay = match path {
Some(path) => { Some(path) => {
let text = fs::read_to_string(path) let text = fs::read_to_string(path)
.map_err(|e| format!("read config {}: {e}", path.display()))?; .map_err(|e| format!("read config {}: {e}", path.display()))?;
@ -753,7 +764,7 @@ fn load_chat_config(path: Option<&Path>) -> Result<LoadedChatTomlConfig, String>
.map_err(|e| format!("parse config {}: {e}", path.display()))? .map_err(|e| format!("parse config {}: {e}", path.display()))?
} }
None => { None => {
let default = Path::new(chat_config::DEFAULT_CONFIG_PATH); let default = Path::new(DEFAULT_CONFIG_PATH);
if !default.is_file() { if !default.is_file() {
ChatTomlConfig::default() ChatTomlConfig::default()
} else { } else {
@ -763,8 +774,7 @@ fn load_chat_config(path: Option<&Path>) -> Result<LoadedChatTomlConfig, String>
.map_err(|e| format!("parse config {}: {e}", default.display()))? .map_err(|e| format!("parse config {}: {e}", default.display()))?
} }
} }
}; })
Ok(LoadedChatTomlConfig { overlay })
} }
impl Config { impl Config {
@ -773,8 +783,7 @@ impl Config {
I: IntoIterator<Item = String>, I: IntoIterator<Item = String>,
{ {
let args = ParsedArgs::parse(provided_args)?; let args = ParsedArgs::parse(provided_args)?;
let loaded = load_chat_config(args.config_path.as_deref())?; let toml = load_chat_config(args.config_path.as_deref())?;
let toml = loaded.overlay;
let provider = provider_from_sources(args.provider.clone(), toml.provider.kind.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(); let node_image = first_non_empty([toml.image.node.clone()]).unwrap_or_default();
if provider != provider_kind::process() && node_image.is_empty() { if provider != provider_kind::process() && node_image.is_empty() {
@ -797,8 +806,8 @@ impl Config {
}; };
Ok(Self { Ok(Self {
orch_bin: default_orch_bin()?, orch_bin: artifact_root().join("target/debug/mvp-orchestrator"),
worker_bin: node_bin_for_current_profile()?, worker_bin: default_worker_bin(),
rpc_addr: DEFAULT_RPC_ADDR.to_owned(), rpc_addr: DEFAULT_RPC_ADDR.to_owned(),
node_image, node_image,
provider, provider,
@ -1248,12 +1257,11 @@ fn resolve_vastai_config(
} }
fn first_non_empty<const N: usize>(values: [Option<String>; N]) -> Option<String> { fn first_non_empty<const N: usize>(values: [Option<String>; N]) -> Option<String> {
values.into_iter().find_map(chat_config::normalize_optional) values
} .into_iter()
.flatten()
fn confirm_vastai_if_needed(config: &Config) -> Result<(), String> { .map(|value| value.trim().to_owned())
let mut approval = StdinVastAiApproval; .find(|value| !value.is_empty())
confirm_vastai_if_needed_with_approval(config, &mut approval)
} }
trait VastAiApproval { trait VastAiApproval {
@ -1322,11 +1330,10 @@ where
input input
.read_line(&mut line) .read_line(&mut line)
.map_err(|e| format!("read Vast.ai approval: {e}"))?; .map_err(|e| format!("read Vast.ai approval: {e}"))?;
Ok(parse_approval(&line)) Ok(matches!(
} line.trim().to_ascii_lowercase().as_str(),
"y" | "yes"
fn parse_approval(input: &str) -> bool { ))
matches!(input.trim().to_ascii_lowercase().as_str(), "y" | "yes")
} }
enum OrchHandle { enum OrchHandle {
@ -1368,8 +1375,9 @@ impl InProcessOrch {
fn spawn(config: &Config, image_ref: &str) -> Result<Self, String> { fn spawn(config: &Config, image_ref: &str) -> Result<Self, String> {
let args = config.orchestrator_cli_args(image_ref); let args = config.orchestrator_cli_args(image_ref);
let (stop_tx, stop_rx) = mpsc::channel(); let (stop_tx, stop_rx) = mpsc::channel();
let thread = let thread = thread::spawn(move || {
thread::spawn(move || crate::run_orchestrator_in_process_from_args(args, stop_rx)); crate::orchestration::app::run_with_options(args, false, Some(stop_rx))
});
Ok(Self { Ok(Self {
stop_tx: Some(stop_tx), stop_tx: Some(stop_tx),
thread: Some(thread), thread: Some(thread),
@ -1397,9 +1405,12 @@ impl InProcessOrch {
Err(error) => return Err(format!("connect prompt RPC {rpc_addr}: {error}")), Err(error) => return Err(format!("connect prompt RPC {rpc_addr}: {error}")),
} }
if let Some(result) = self.take_finished_result() { if let Some(result) = self.take_finished_result() {
let reason = match result {
Ok(()) => "completed successfully".to_owned(),
Err(error) => error,
};
return Err(format!( return Err(format!(
"in-process orchestrator exited before prompt RPC ready: {}", "in-process orchestrator exited before prompt RPC ready: {reason}"
render_orch_thread_result(result)
)); ));
} }
thread::sleep(Duration::from_millis(100)); 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 { struct OrchChild {
child: Child, child: Child,
cleaned: bool, cleaned: bool,
@ -1493,7 +1489,11 @@ impl OrchChild {
Ok(Self { Ok(Self {
child, child,
cleaned: false, 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<String, String> {
prepare_node_image_with_progress(request, progress)
}
#[allow(dead_code)]
fn prepare_runtime(config: &Config) -> Result<String, String> {
prepare_runtime_with(config, |request| {
prepare_node_image_with_progress(request, None)
})
}
#[allow(dead_code)]
fn prepare_runtime_with<F>(config: &Config, prepare_node_image_fn: F) -> Result<String, String>
where
F: FnMut(NodeImageRequest) -> Result<String, String>,
{
let mut prepare_node_image_fn = prepare_node_image_fn;
prepare_runtime_with_progress(
config,
move |request, _progress| prepare_node_image_fn(request),
None,
)
}
fn prepare_runtime_with_progress<F>( fn prepare_runtime_with_progress<F>(
config: &Config, config: &Config,
mut prepare_node_image_fn: F, mut prepare_node_image_fn: F,
@ -1651,7 +1624,21 @@ where
"started", "started",
json!({"mode": binary_mode, "command_label": "ensure_orch_binary"}), 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( Ok(()) => emit_chat_progress(
&mut progress, &mut progress,
CHAT_RUNTIME_CHANNEL, CHAT_RUNTIME_CHANNEL,
@ -1680,7 +1667,19 @@ where
"started", "started",
json!({"mode": binary_mode}), 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( Ok(()) => emit_chat_progress(
&mut progress, &mut progress,
CHAT_RUNTIME_CHANNEL, CHAT_RUNTIME_CHANNEL,
@ -1735,7 +1734,19 @@ where
"started", "started",
json!({"mode": binary_mode, "command_label": "ensure_worker_binary"}), 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( Ok(()) => emit_chat_progress(
&mut progress, &mut progress,
CHAT_RUNTIME_CHANNEL, CHAT_RUNTIME_CHANNEL,
@ -1772,31 +1783,25 @@ where
"started", "started",
json!({"provider": config.provider.as_str(), "command_label": "prepare_node_image", "image_tag": config.image_tag.as_deref()}), 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() { let node_bin = default_worker_bin();
Ok(path) => path, let requires_registry_image = if config.provider == provider_kind::docker() {
Err(error) => { false
emit_chat_progress( } else if config.provider == provider_kind::vastai() {
&mut progress, true
CHAT_RUNTIME_CHANNEL, } else {
"prepare_node_image", let error = if config.provider == provider_kind::process() {
"failed", "process provider does not use node images"
json!({"provider": config.provider.as_str(), "command_label": "prepare_node_image", "elapsed_ms": prepare_node_image_started.elapsed().as_millis(), "error": error.as_str()}), } else {
); "mvp-chat does not support mock provider"
return Err(error);
}
}; };
let provider = match node_image_provider(&config.provider) {
Ok(provider) => provider,
Err(error) => {
emit_chat_progress( emit_chat_progress(
&mut progress, &mut progress,
CHAT_RUNTIME_CHANNEL, CHAT_RUNTIME_CHANNEL,
"prepare_node_image", "prepare_node_image",
"failed", "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()}), 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); return Err(error.to_owned());
}
}; };
let prepared = { let prepared = {
let command_progress = progress let command_progress = progress
@ -1807,9 +1812,8 @@ where
requested_image: config.node_image.clone(), requested_image: config.node_image.clone(),
base_image: BASE_NODE_IMAGE.to_owned(), base_image: BASE_NODE_IMAGE.to_owned(),
node_bin, node_bin,
provider, requires_registry_image,
extra_tag: config.image_tag.clone(), extra_tag: config.image_tag.clone(),
push: false,
force_refresh: false, force_refresh: false,
enabled: true, enabled: true,
}, },
@ -1838,42 +1842,6 @@ where
Ok(prepared) Ok(prepared)
} }
fn stdin_prompt_events() -> mpsc::Receiver<PromptInput> {
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( fn run_chat_loop_with_input_and_progress(
addr: &str, addr: &str,
max_tokens: u32, max_tokens: u32,
@ -1923,7 +1891,15 @@ fn run_chat_loop_with_input_and_progress(
return Err(format!("clone prompt RPC stream: {error}")); 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)] #[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) 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<PromptInput>,
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( fn emit_chat_progress(
progress: &mut Option<&mut ChatDatastream>, progress: &mut Option<&mut ChatDatastream>,
channel: &str, channel: &str,
@ -2186,61 +2144,17 @@ fn run_chat_session_with_output_and_progress(
} }
} }
fn default_orch_bin() -> Result<PathBuf, String> { fn default_worker_bin() -> PathBuf {
Ok(artifact_root().join("target/debug/mvp-orchestrator")) artifact_root().join("target/debug/mvp-worker-node")
} }
fn node_bin_for_current_profile() -> Result<PathBuf, String> { fn ensure_runtime_binary(
Ok(artifact_root().join("target/debug/mvp-worker-node")) skip_rebuild: bool,
} path: &Path,
label: &str,
fn cargo_command() -> &'static str { cargo_args: &[&str],
"cargo" ) -> Result<(), String> {
} if skip_rebuild {
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");
}
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) let metadata = fs::metadata(path)
.map_err(|e| format!("missing required {label} artifact {}: {e}", path.display()))?; .map_err(|e| format!("missing required {label} artifact {}: {e}", path.display()))?;
if !metadata.is_file() { if !metadata.is_file() {
@ -2249,21 +2163,20 @@ fn ensure_existing_artifact(path: &PathBuf, label: &str) -> Result<(), String> {
path.display() path.display()
)); ));
} }
Ok(()) return Ok(());
} }
fn run_status(program: &str, args: &[&str], label: &str) -> Result<(), String> { let status = Command::new("cargo")
let status = Command::new(program) .args(cargo_args)
.args(args)
.stdin(Stdio::null()) .stdin(Stdio::null())
.stdout(Stdio::inherit()) .stdout(Stdio::inherit())
.stderr(Stdio::inherit()) .stderr(Stdio::inherit())
.status() .status()
.map_err(|e| format!("run {label}: {e}"))?; .map_err(|e| format!("run build {label}: {e}"))?;
if status.success() { if status.success() {
Ok(()) Ok(())
} else { } else {
Err(format!("{label} failed with {status}")) Err(format!("build {label} failed with {status}"))
} }
} }
@ -2393,18 +2306,6 @@ fn env_optional(name: &str) -> Option<String> {
.filter(|value| !value.is_empty()) .filter(|value| !value.is_empty())
} }
fn node_image_provider(provider: &ProviderKind) -> Result<NodeImageProvider, String> {
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<Item = String>, name: &str) -> Result<String, String> { fn next_arg(args: &mut impl Iterator<Item = String>, name: &str) -> Result<String, String> {
args.next() args.next()
.ok_or_else(|| format!("missing value after {name}")) .ok_or_else(|| format!("missing value after {name}"))

16
crates/mvp-system/src/lib.rs
View file

@ -12,28 +12,18 @@ pub fn run_chat_from_args<I>(args: I) -> std::process::ExitCode
where where
I: IntoIterator<Item = String>, I: IntoIterator<Item = String>,
{ {
chat::run_from_args(args) chat::runtime::run_from_args(args)
} }
pub fn run_orchestrator_from_args<I>(args: I) -> Result<(), String> pub fn run_orchestrator_from_args<I>(args: I) -> Result<(), String>
where where
I: IntoIterator<Item = String>, I: IntoIterator<Item = String>,
{ {
orchestration::run_from_args(args) orchestration::app::run_with_options(args, true, None)
} }
pub fn run_worker_node_from_env() -> std::process::ExitCode { pub fn run_worker_node_from_env() -> std::process::ExitCode {
node::run_worker_node_from_env() node::worker_node_runtime::run_from_env()
}
fn run_orchestrator_in_process_from_args<I>(
args: I,
stop_rx: std::sync::mpsc::Receiver<()>,
) -> Result<(), String>
where
I: IntoIterator<Item = String>,
{
orchestration::run_in_process_from_args(args, stop_rx)
} }
#[path = "transport/driver_pumps.rs"] #[path = "transport/driver_pumps.rs"]

6
crates/mvp-system/src/node/mod.rs
View file

@ -3,8 +3,4 @@
//! Worker-node runtime behavior lives behind this module boundary; binaries //! Worker-node runtime behavior lives behind this module boundary; binaries
//! only wire entrypoints into it. //! only wire entrypoints into it.
mod worker_node_runtime; pub(super) mod worker_node_runtime;
pub(super) fn run_worker_node_from_env() -> std::process::ExitCode {
worker_node_runtime::run_from_env()
}

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

@ -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::gguf_shard::{StageShardPlan, materialize_stage_shard_http, validate_stage_shard_cache};
use crate::node_actor::{ use crate::node_actor::{
NodeAgentActor, NodeAgentMsg, NodeAgentReport, StageCommandWire, StageInboundEdgeWire, NodeAgentActor, NodeAgentMsg, NodeAgentReport, StageCommandWire, StageInboundEdgeWire,
StageObjectSpecWire, StageOutboundEdgeWire, StageRingSpecWire, StageObjectSpecWire, StageOutboundEdgeWire,
}; };
use crate::observability::benchmark; use crate::observability::benchmark;
use crate::orchestration::distribution_stack::DistributionRuntimeStack; use crate::orchestration::distribution_stack::DistributionRuntimeStack;
@ -908,8 +908,16 @@ impl WorkerEdgeRuntime {
run_id: edge::RunId(config.run_id), run_id: edge::RunId(config.run_id),
edge_id: edge::EdgeId(edge.edge_id), edge_id: edge::EdgeId(edge.edge_id),
local_node_id: edge::NodeId(config.logical_node_id), local_node_id: edge::NodeId(config.logical_node_id),
object_spec: edge_object_spec(edge.object_spec), object_spec: edge::ObjectSpec {
ring_spec: edge_ring_spec(edge.ring_spec), 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( self.drive_edge_workflow(
stack, stack,
@ -954,8 +962,16 @@ impl WorkerEdgeRuntime {
edge_id: edge::EdgeId(edge.edge_id), edge_id: edge::EdgeId(edge.edge_id),
local_node_id: edge::NodeId(config.logical_node_id), local_node_id: edge::NodeId(config.logical_node_id),
consumer_node_id: edge::NodeId(edge.consumer_node_id), consumer_node_id: edge::NodeId(edge.consumer_node_id),
object_spec: edge_object_spec(edge.object_spec), object_spec: edge::ObjectSpec {
ring_spec: edge_ring_spec(edge.ring_spec), 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( self.drive_edge_workflow(
stack, stack,
@ -1581,29 +1597,6 @@ impl WorkerEdgeRuntime {
fn duration_ms_u64(duration: Duration) -> u64 { fn duration_ms_u64(duration: Duration) -> u64 {
u64::try_from(duration.as_millis()).unwrap_or(u64::MAX) 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 { struct IngressRecordBytes {
bytes: Vec<u8>, bytes: Vec<u8>,
@ -1618,7 +1611,15 @@ fn take_complete_ingress_record(
buffer: &mut Vec<u8>, buffer: &mut Vec<u8>,
spec: StageObjectSpecWire, spec: StageObjectSpecWire,
) -> Result<Option<IngressRecordBytes>, String> { ) -> Result<Option<IngressRecordBytes>, String> {
let record = match ingress::read_object_record(buffer, ingress_object_spec(spec), false) 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:?}"))? .map_err(|reason| format!("invalid object record: {reason:?}"))?
{ {
ingress::ObjectRecordRead::Incomplete => return Ok(None), ingress::ObjectRecordRead::Incomplete => return Ok(None),
@ -1642,21 +1643,25 @@ fn value_u64(value: &Value, field: &str) -> Result<u64, String> {
.ok_or_else(|| format!("helper event missing numeric {field}: {value}")) .ok_or_else(|| format!("helper event missing numeric {field}: {value}"))
} }
pub(super) fn run_from_env() -> ExitCode { pub(crate) fn run_from_env() -> ExitCode {
let mut args = std::env::args().skip(1).collect::<Vec<_>>(); let mut args = std::env::args().skip(1);
if args.first().map(String::as_str) == Some("debug-join") { match args.next().as_deref() {
args.remove(0); Some("debug-join") => debug_join_client_main(args.collect()),
return debug_join_client_main(args); 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)
} }
if args.first().map(String::as_str) == Some("stage-shard-fetcher") { },
return stage_shard_fetcher_main(); _ => match run() {
}
match run() {
Ok(()) => ExitCode::SUCCESS, Ok(()) => ExitCode::SUCCESS,
Err(error) => { Err(error) => {
eprintln!("mvp-worker-node: {error}"); eprintln!("mvp-worker-node: {error}");
ExitCode::from(1) ExitCode::from(1)
} }
},
} }
} }
@ -1666,17 +1671,6 @@ struct StageShardFetchRequest {
output_path: PathBuf, 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> { fn run_stage_shard_fetcher() -> Result<(), String> {
let mut input = String::new(); let mut input = String::new();
std::io::stdin() std::io::stdin()
@ -1858,7 +1852,7 @@ fn run() -> Result<(), String> {
}; };
let arena_fd = arena_manager.lock().arena_fd(); 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_transport = driver.datastream_publish_handle();
let datastream_publisher = match stack let datastream_publisher = match stack
.runtime .runtime
@ -2338,7 +2332,7 @@ fn emit_swim_telemetry(
local_phase: &str, local_phase: &str,
) { ) {
for transition in stack.drain_swim_transitions() { 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 from = transition.from.map(|state| format!("{:?}", state));
let to = format!("{:?}", transition.to); let to = format!("{:?}", transition.to);
let member_state = stack let member_state = stack
@ -2375,7 +2369,7 @@ fn swim_probe_event_record(
let budget_ms = event.budget_ms; let budget_ms = event.budget_ms;
SwimProbeEvent { SwimProbeEvent {
event: event.event.to_owned(), event: event.event.to_owned(),
target: format_dist_node_id(event.target), target: format!("{:?}", event.target),
sequence: event.sequence, sequence: event.sequence,
kind: event.kind.to_owned(), kind: event.kind.to_owned(),
rtt_ms: event.rtt_ms, rtt_ms: event.rtt_ms,
@ -2403,18 +2397,10 @@ fn swim_recent_probe_targets(stack: &DistributionRuntimeStack) -> Vec<String> {
.swim_telemetry .swim_telemetry
.recent_targets() .recent_targets()
.into_iter() .into_iter()
.map(format_dist_node_id) .map(|node_id| format!("{:?}", node_id))
.collect() .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)] #[derive(Clone, Copy)]
struct DatastreamChannelSet { struct DatastreamChannelSet {
node_ready: ChannelId, node_ready: ChannelId,
@ -3989,8 +3975,18 @@ struct DeploymentConfig {
impl DeploymentConfig { impl DeploymentConfig {
fn from_env() -> Result<Self, String> { fn from_env() -> Result<Self, String> {
let run_id = env_u64("MVP_RUN_ID", 1)?; macro_rules! env_parse {
let logical_node_id = env_u64("MVP_LOGICAL_NODE_ID", 1)?; ($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 relay = relay_runtime_config_from_env(run_id)?;
let debug_join_socket = match env_optional("MVP_DEBUG_JOIN_SOCKET").as_deref() { let debug_join_socket = match env_optional("MVP_DEBUG_JOIN_SOCKET").as_deref() {
Some("disabled") => None, Some("disabled") => None,
@ -4004,7 +4000,7 @@ impl DeploymentConfig {
.into_owned(), .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" { let default_device = if provider == "process" {
"CPU" "CPU"
} else { } else {
@ -4013,9 +4009,19 @@ impl DeploymentConfig {
Ok(Self { Ok(Self {
run_id, run_id,
logical_node_id, logical_node_id,
stage_index: env_u32("MVP_STAGE_INDEX", 0)?, stage_index: env_parse!("MVP_STAGE_INDEX", 0)?,
coordinator_endpoint: env_json("MVP_COORDINATOR_ENDPOINT")?, coordinator_endpoint: env_optional("MVP_COORDINATOR_ENDPOINT")
orchestrator_actor: env_json("MVP_ORCHESTRATOR_ACTOR")?, .map(|value| {
serde_json::from_str::<EndpointAddr>(&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::<ActorAddress>(&value)
.map_err(|e| format!("invalid MVP_ORCHESTRATOR_ACTOR JSON: {e}"))
})
.transpose()?,
datastream_frame_log: env_optional("MVP_DATASTREAM_FRAME_LOG"), datastream_frame_log: env_optional("MVP_DATASTREAM_FRAME_LOG"),
debug_join_socket, debug_join_socket,
relay_mode: relay.mode, relay_mode: relay.mode,
@ -4024,16 +4030,29 @@ impl DeploymentConfig {
.map(EndpointAddrMask::parse) .map(EndpointAddrMask::parse)
.transpose()? .transpose()?
.unwrap_or_default(), .unwrap_or_default(),
worker_script: env_string("MVP_TINYGRAD_WORKER", DEFAULT_WORKER_SCRIPT), worker_script: env_optional("MVP_TINYGRAD_WORKER")
device: env_string("DEV", default_device), .unwrap_or_else(|| DEFAULT_WORKER_SCRIPT.to_owned()),
model_id: env_string("MVP_MODEL_ID", DEFAULT_MODEL_ID), device: env_optional("DEV").unwrap_or_else(|| default_device.to_owned()),
gguf_source: gguf_source_from_env(), model_id: env_optional("MVP_MODEL_ID").unwrap_or_else(|| DEFAULT_MODEL_ID.to_owned()),
tokenizer: tokenizer_from_env(), 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_prompt: env_optional("MVP_NODE_SELF_TEST_PROMPT"),
self_test_layer_end: env_u32("MVP_SELF_TEST_LAYER_END", 16)?, self_test_layer_end: env_parse!("MVP_SELF_TEST_LAYER_END", 16)?,
self_test_max_tokens: env_u32("MVP_SELF_TEST_MAX_TOKENS", 1)?, self_test_max_tokens: env_parse!("MVP_SELF_TEST_MAX_TOKENS", 1)?,
arena_bytes: env_u64("MVP_ARENA_BYTES", DEFAULT_ARENA_BYTES)?, arena_bytes: env_parse!("MVP_ARENA_BYTES", DEFAULT_ARENA_BYTES)?,
arena_alignment: env_u64("MVP_ARENA_ALIGNMENT", DEFAULT_ARENA_ALIGNMENT)?, arena_alignment: env_parse!("MVP_ARENA_ALIGNMENT", DEFAULT_ARENA_ALIGNMENT)?,
}) })
} }
} }
@ -4788,57 +4807,9 @@ fn spawn_stdin_shutdown_listener() -> Receiver<()> {
rx rx
} }
fn env_string(name: &str, default: &str) -> String {
env_optional(name).unwrap_or_else(|| default.to_owned())
}
fn env_optional(name: &str) -> Option<String> { fn env_optional(name: &str) -> Option<String> {
std::env::var(name) std::env::var(name)
.ok() .ok()
.map(|value| value.trim().to_owned()) .map(|value| value.trim().to_owned())
.filter(|value| !value.is_empty()) .filter(|value| !value.is_empty())
} }
fn env_u64(name: &str, default: u64) -> Result<u64, String> {
match env_optional(name) {
Some(value) => value
.parse::<u64>()
.map_err(|e| format!("invalid {name}={value:?}: {e}")),
None => Ok(default),
}
}
fn env_u32(name: &str, default: u32) -> Result<u32, String> {
match env_optional(name) {
Some(value) => value
.parse::<u32>()
.map_err(|e| format!("invalid {name}={value:?}: {e}")),
None => Ok(default),
}
}
fn env_json<T>(name: &str) -> Result<Option<T>, 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)
}

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

File diff suppressed because it is too large Load diff

102
crates/mvp-system/src/orchestration/engine_builder/engine.rs
View file

@ -1,27 +1,21 @@
use std::collections::BTreeMap;
use crate::run_plan::RunId; use crate::run_plan::RunId;
use super::error::EngineBuildError; use super::error::EngineBuildError;
use super::events::EngineEvent; use super::events::EngineEvent;
use super::launcher::{ use super::launcher::{LaunchedNode, NodeControl, NodeFacts, NodeLaunchSpec, StaticNodeLauncher};
CoordinatorJoinSpec, LaunchedNode, NodeControl, NodeFacts, NodeLaunchSpec, NodeLauncher,
};
use super::model::ModelSpec; use super::model::ModelSpec;
use super::node_image::NodeImageSpec; use super::node_image::NodeImageSpec;
use super::planner::{RoleAssignmentPlan, RolePlanner, RolePlannerInput}; use super::planner::{FixedLinearPipelinePlanner, RoleAssignmentPlan, RolePlannerInput};
use super::pool::{PoolProvider, PoolRequest, ResourceRequest}; use super::pool::{PoolRequest, StaticPoolProvider};
use super::roles::{RoleAssignment, RoleKind}; use super::roles::RoleAssignment;
pub struct ClusterBuilder { pub struct ClusterBuilder {
cluster_id: String, cluster_id: String,
run_id: RunId, run_id: RunId,
model: ModelSpec, model: ModelSpec,
image: Option<NodeImageSpec>, pool_provider: Option<StaticPoolProvider>,
pool_provider: Option<Box<dyn PoolProvider>>, launcher: Option<StaticNodeLauncher>,
launcher: Option<Box<dyn NodeLauncher>>, planner: Option<FixedLinearPipelinePlanner>,
planner: Option<Box<dyn RolePlanner>>,
required_resources: ResourceRequest,
} }
impl ClusterBuilder { impl ClusterBuilder {
@ -30,11 +24,9 @@ impl ClusterBuilder {
cluster_id: cluster_id.into(), cluster_id: cluster_id.into(),
run_id: RunId(1), run_id: RunId(1),
model, model,
image: None,
pool_provider: None, pool_provider: None,
launcher: None, launcher: None,
planner: None, planner: None,
required_resources: ResourceRequest::default(),
} }
} }
@ -43,36 +35,26 @@ impl ClusterBuilder {
self self
} }
pub fn image(mut self, image: NodeImageSpec) -> Self { pub fn image(self, _image: NodeImageSpec) -> Self {
self.image = Some(image);
self self
} }
pub fn pool_provider(mut self, provider: impl PoolProvider + 'static) -> Self { pub fn pool_provider(mut self, provider: StaticPoolProvider) -> Self {
self.pool_provider = Some(Box::new(provider)); self.pool_provider = Some(provider);
self self
} }
pub fn launcher(mut self, launcher: impl NodeLauncher + 'static) -> Self { pub fn launcher(mut self, launcher: StaticNodeLauncher) -> Self {
self.launcher = Some(Box::new(launcher)); self.launcher = Some(launcher);
self self
} }
pub fn planner(mut self, planner: impl RolePlanner + 'static) -> Self { pub fn planner(mut self, planner: FixedLinearPipelinePlanner) -> Self {
self.planner = Some(Box::new(planner)); self.planner = Some(planner);
self
}
pub fn required_resources(mut self, required_resources: ResourceRequest) -> Self {
self.required_resources = required_resources;
self self
} }
pub fn launch(mut self) -> Result<ClusterHandle, EngineBuildError> { pub fn launch(mut self) -> Result<ClusterHandle, EngineBuildError> {
let image = self
.image
.take()
.ok_or(EngineBuildError::MissingComponent("image"))?;
let pool_provider = self let pool_provider = self
.pool_provider .pool_provider
.take() .take()
@ -88,10 +70,7 @@ impl ClusterBuilder {
let mut events = Vec::new(); let mut events = Vec::new();
let leases = pool_provider.acquire_pool(PoolRequest { let leases = pool_provider.acquire_pool(PoolRequest {
cluster_id: self.cluster_id.clone(),
min_nodes: planner.required_node_count(), min_nodes: planner.required_node_count(),
image: image.clone(),
required_resources: self.required_resources.clone(),
})?; })?;
if leases.is_empty() { if leases.is_empty() {
return Err(EngineBuildError::EmptyPool); return Err(EngineBuildError::EmptyPool);
@ -103,16 +82,7 @@ impl ClusterBuilder {
let mut nodes = Vec::with_capacity(leases.len()); let mut nodes = Vec::with_capacity(leases.len());
let mut iter = leases.into_iter(); let mut iter = leases.into_iter();
let coordinator_lease = iter.next().ok_or(EngineBuildError::EmptyPool)?; let coordinator_lease = iter.next().ok_or(EngineBuildError::EmptyPool)?;
let mut coordinator = launcher.launch_node( let mut coordinator = launcher.launch_node(&coordinator_lease, NodeLaunchSpec);
&coordinator_lease,
NodeLaunchSpec {
cluster_id: self.cluster_id.clone(),
image: image.clone(),
coordinator: None,
is_coordinator: true,
env: BTreeMap::new(),
},
)?;
events.push(EngineEvent::NodeLaunched { events.push(EngineEvent::NodeLaunched {
node_id: coordinator.lease.logical_node_id, node_id: coordinator.lease.logical_node_id,
coordinator: true, coordinator: true,
@ -121,26 +91,10 @@ impl ClusterBuilder {
events.push(EngineEvent::NodeBootReady { events.push(EngineEvent::NodeBootReady {
node_id: coordinator_facts.node_id, 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)); nodes.push(EngineNode::new(coordinator, coordinator_facts));
for lease in iter { for lease in iter {
let mut node = launcher.launch_node( let mut node = launcher.launch_node(&lease, NodeLaunchSpec);
&lease,
NodeLaunchSpec {
cluster_id: self.cluster_id.clone(),
image: image.clone(),
coordinator: Some(CoordinatorJoinSpec {
endpoint: coordinator_endpoint.clone(),
}),
is_coordinator: false,
env: BTreeMap::new(),
},
)?;
events.push(EngineEvent::NodeLaunched { events.push(EngineEvent::NodeLaunched {
node_id: node.lease.logical_node_id, node_id: node.lease.logical_node_id,
coordinator: false, coordinator: false,
@ -161,7 +115,6 @@ impl ClusterBuilder {
}); });
let plan = planner.plan(RolePlannerInput { let plan = planner.plan(RolePlannerInput {
cluster_id: self.cluster_id.clone(),
run_id: self.run_id, run_id: self.run_id,
model: self.model, model: self.model,
nodes: nodes.iter().map(|node| node.facts.clone()).collect(), nodes: nodes.iter().map(|node| node.facts.clone()).collect(),
@ -203,10 +156,6 @@ pub struct ClusterHandle {
} }
impl ClusterHandle { impl ClusterHandle {
pub fn cluster_id(&self) -> &str {
&self.cluster_id
}
pub fn role_plan(&self) -> &RoleAssignmentPlan { pub fn role_plan(&self) -> &RoleAssignmentPlan {
&self.plan &self.plan
} }
@ -215,17 +164,6 @@ impl ClusterHandle {
&self.events &self.events
} }
pub fn node_summaries(&self) -> Vec<NodeSummary> {
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<Vec<EngineEvent>, EngineBuildError> { pub fn shutdown(mut self) -> Result<Vec<EngineEvent>, EngineBuildError> {
for node in &mut self.nodes { for node in &mut self.nodes {
node.control.shutdown()?; node.control.shutdown()?;
@ -239,14 +177,6 @@ impl ClusterHandle {
Ok(self.events) Ok(self.events)
} }
} }
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct NodeSummary {
pub node_id: crate::run_plan::NodeId,
pub roles: Vec<RoleKind>,
pub facts: NodeFacts,
}
struct EngineNode { struct EngineNode {
facts: NodeFacts, facts: NodeFacts,
roles: Vec<RoleAssignment>, roles: Vec<RoleAssignment>,

43
crates/mvp-system/src/orchestration/engine_builder/error.rs
View file

@ -1,4 +1,3 @@
use std::error::Error;
use std::fmt; use std::fmt;
use crate::run_plan; use crate::run_plan;
@ -7,10 +6,8 @@ use crate::run_plan;
pub enum EngineBuildError { pub enum EngineBuildError {
MissingComponent(&'static str), MissingComponent(&'static str),
EmptyPool, EmptyPool,
CoordinatorEndpointMissing { node_id: u64 },
RoleTargetMissing { node_id: u64 }, RoleTargetMissing { node_id: u64 },
Pool(PoolError), Pool(PoolError),
Launch(LaunchError),
Node(NodeControlError), Node(NodeControlError),
Planning(PlanningError), Planning(PlanningError),
} }
@ -20,37 +17,22 @@ impl fmt::Display for EngineBuildError {
match self { match self {
Self::MissingComponent(name) => write!(f, "missing engine builder component: {name}"), Self::MissingComponent(name) => write!(f, "missing engine builder component: {name}"),
Self::EmptyPool => write!(f, "pool provider returned no nodes"), 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 } => { Self::RoleTargetMissing { node_id } => {
write!(f, "role assignment targeted unknown node {node_id}") write!(f, "role assignment targeted unknown node {node_id}")
} }
Self::Pool(err) => err.fmt(f), Self::Pool(err) => err.fmt(f),
Self::Launch(err) => err.fmt(f),
Self::Node(err) => err.fmt(f), Self::Node(err) => err.fmt(f),
Self::Planning(err) => err.fmt(f), Self::Planning(err) => err.fmt(f),
} }
} }
} }
impl Error for EngineBuildError {}
impl From<PoolError> for EngineBuildError { impl From<PoolError> for EngineBuildError {
fn from(value: PoolError) -> Self { fn from(value: PoolError) -> Self {
Self::Pool(value) Self::Pool(value)
} }
} }
impl From<LaunchError> for EngineBuildError {
fn from(value: LaunchError) -> Self {
Self::Launch(value)
}
}
impl From<NodeControlError> for EngineBuildError { impl From<NodeControlError> for EngineBuildError {
fn from(value: NodeControlError) -> Self { fn from(value: NodeControlError) -> Self {
Self::Node(value) Self::Node(value)
@ -66,7 +48,6 @@ impl From<PlanningError> for EngineBuildError {
#[derive(Clone, Debug, PartialEq, Eq)] #[derive(Clone, Debug, PartialEq, Eq)]
pub enum PoolError { pub enum PoolError {
InsufficientNodes { requested: usize, available: usize }, InsufficientNodes { requested: usize, available: usize },
Provider(String),
} }
impl fmt::Display for PoolError { impl fmt::Display for PoolError {
@ -79,34 +60,16 @@ impl fmt::Display for PoolError {
f, f,
"pool has {available} matching nodes, but {requested} were requested" "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)] #[derive(Clone, Debug, PartialEq, Eq)]
pub enum NodeControlError { pub enum NodeControlError {
NotBooted { node_id: u64 }, NotBooted { node_id: u64 },
Stopped { node_id: u64 }, Stopped { node_id: u64 },
RoleNodeMismatch { node_id: u64, role_node_id: u64 }, RoleNodeMismatch { node_id: u64, role_node_id: u64 },
Backend(String), Backend(&'static str),
} }
impl fmt::Display for NodeControlError { impl fmt::Display for NodeControlError {
@ -126,8 +89,6 @@ impl fmt::Display for NodeControlError {
} }
} }
impl Error for NodeControlError {}
#[derive(Clone, Debug, PartialEq, Eq)] #[derive(Clone, Debug, PartialEq, Eq)]
pub enum PlanningError { pub enum PlanningError {
DuplicateNodeId { node_id: u64 }, DuplicateNodeId { node_id: u64 },
@ -158,5 +119,3 @@ impl fmt::Display for PlanningError {
} }
} }
} }
impl Error for PlanningError {}

54
crates/mvp-system/src/orchestration/engine_builder/launcher.rs
View file

@ -1,20 +1,11 @@
use std::collections::{BTreeMap, BTreeSet}; use std::collections::BTreeSet;
use crate::run_plan::NodeId; use crate::run_plan::NodeId;
use super::error::{LaunchError, NodeControlError}; use super::error::NodeControlError;
use super::node_image::NodeImageSpec; use super::pool::{NodeCapability, NodeLease};
use super::pool::{NodeCapability, NodeLease, ResourceFacts};
use super::roles::RoleAssignment; use super::roles::RoleAssignment;
pub trait NodeLauncher: Send + Sync {
fn launch_node(
&self,
lease: &NodeLease,
spec: NodeLaunchSpec,
) -> Result<LaunchedNode, LaunchError>;
}
pub trait NodeControl: Send { pub trait NodeControl: Send {
fn wait_boot_ready(&mut self) -> Result<NodeFacts, NodeControlError>; fn wait_boot_ready(&mut self) -> Result<NodeFacts, NodeControlError>;
fn wait_cluster_converged(&mut self, expected_alive: usize) -> Result<(), NodeControlError>; fn wait_cluster_converged(&mut self, expected_alive: usize) -> Result<(), NodeControlError>;
@ -23,33 +14,18 @@ pub trait NodeControl: Send {
} }
#[derive(Clone, Debug, PartialEq, Eq)] #[derive(Clone, Debug, PartialEq, Eq)]
pub struct NodeLaunchSpec { pub struct NodeLaunchSpec;
pub cluster_id: String,
pub image: NodeImageSpec,
pub coordinator: Option<CoordinatorJoinSpec>,
pub is_coordinator: bool,
pub env: BTreeMap<String, String>,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct CoordinatorJoinSpec {
pub endpoint: String,
}
#[derive(Clone, Debug, PartialEq, Eq)] #[derive(Clone, Debug, PartialEq, Eq)]
pub struct NodeFacts { pub struct NodeFacts {
pub node_id: NodeId, pub node_id: NodeId,
pub coordinator_endpoint: Option<String>,
pub resources: ResourceFacts,
pub capabilities: BTreeSet<NodeCapability>, pub capabilities: BTreeSet<NodeCapability>,
} }
impl NodeFacts { impl NodeFacts {
pub fn from_lease(lease: &NodeLease, coordinator_endpoint: Option<String>) -> Self { pub fn from_lease(lease: &NodeLease) -> Self {
Self { Self {
node_id: lease.logical_node_id, node_id: lease.logical_node_id,
coordinator_endpoint,
resources: lease.expected_resources.clone(),
capabilities: lease.capabilities.clone(), capabilities: lease.capabilities.clone(),
} }
} }
@ -63,18 +39,10 @@ pub struct LaunchedNode {
#[derive(Clone, Debug, Default)] #[derive(Clone, Debug, Default)]
pub struct StaticNodeLauncher; pub struct StaticNodeLauncher;
impl NodeLauncher for StaticNodeLauncher { impl StaticNodeLauncher {
fn launch_node( pub fn launch_node(&self, lease: &NodeLease, _spec: NodeLaunchSpec) -> LaunchedNode {
&self, let facts = NodeFacts::from_lease(lease);
lease: &NodeLease, LaunchedNode {
spec: NodeLaunchSpec,
) -> Result<LaunchedNode, LaunchError> {
let endpoint = format!(
"static://{}/node/{}",
spec.cluster_id, lease.logical_node_id.0
);
let facts = NodeFacts::from_lease(lease, Some(endpoint));
Ok(LaunchedNode {
lease: lease.clone(), lease: lease.clone(),
control: Box::new(StaticNodeControl { control: Box::new(StaticNodeControl {
facts, facts,
@ -82,7 +50,7 @@ impl NodeLauncher for StaticNodeLauncher {
stopped: false, stopped: false,
assigned_roles: Vec::new(), assigned_roles: Vec::new(),
}), }),
}) }
} }
} }
@ -123,7 +91,7 @@ impl NodeControl for StaticNodeControl {
} }
if expected_alive == 0 { if expected_alive == 0 {
return Err(NodeControlError::Backend( return Err(NodeControlError::Backend(
"expected_alive must be greater than zero".to_owned(), "expected_alive must be greater than zero",
)); ));
} }
Ok(()) Ok(())

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

@ -16,15 +16,14 @@ pub mod node_image;
pub mod planner; pub mod planner;
pub mod pool; pub mod pool;
pub mod roles; 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 engine::ClusterBuilder;
pub use events::EngineEvent; pub use events::EngineEvent;
pub use launcher::StaticNodeLauncher; pub use launcher::StaticNodeLauncher;
pub use model::{DTypeFamily, ModelArtifact, ModelSpec}; pub use model::{ModelArtifact, ModelSpec};
pub use node_image::{NodeImageSpec, WorkerRuntimeSpec};
pub use planner::FixedLinearPipelinePlanner; pub use planner::FixedLinearPipelinePlanner;
pub use pool::{NodeCapability, NodeLease, ResourceFacts, StaticPoolProvider}; pub use pool::{NodeCapability, NodeLease, ResourceFacts, StaticPoolProvider};
pub use roles::RoleKind; pub use roles::RoleKind;
pub use node_image::{NodeImageSpec, WorkerRuntimeSpec};

67
crates/mvp-system/src/orchestration/engine_builder/model.rs
View file

@ -3,12 +3,11 @@ use crate::run_plan;
#[derive(Clone, Debug, PartialEq, Eq)] #[derive(Clone, Debug, PartialEq, Eq)]
pub struct ModelSpec { pub struct ModelSpec {
pub model_id: String, pub model_id: String,
pub architecture: ModelArchitecture,
pub artifact: ModelArtifact, pub artifact: ModelArtifact,
pub tokenizer: run_plan::TokenizerSource, pub tokenizer: run_plan::TokenizerSource,
pub num_layers: u32, pub num_layers: u32,
pub hidden_dim: u64, pub hidden_dim: u64,
pub dtype_family: DTypeFamily, pub dtype_family: run_plan::DTypeFamily,
pub dtype_width_bytes: u64, pub dtype_width_bytes: u64,
pub max_seq_len: u64, pub max_seq_len: u64,
pub eos_token_id: u32, pub eos_token_id: u32,
@ -20,7 +19,7 @@ impl ModelSpec {
artifact: ModelArtifact, artifact: ModelArtifact,
num_layers: u32, num_layers: u32,
hidden_dim: u64, hidden_dim: u64,
dtype_family: DTypeFamily, dtype_family: run_plan::DTypeFamily,
dtype_width_bytes: u64, dtype_width_bytes: u64,
max_seq_len: u64, max_seq_len: u64,
eos_token_id: u32, eos_token_id: u32,
@ -28,7 +27,6 @@ impl ModelSpec {
) -> Self { ) -> Self {
Self { Self {
model_id: model_id.into(), model_id: model_id.into(),
architecture: ModelArchitecture::PipelinedCausalLlm,
artifact, artifact,
tokenizer, tokenizer,
num_layers, 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 { pub fn to_run_plan_facts(&self) -> run_plan::ModelFacts {
run_plan::ModelFacts { run_plan::ModelFacts {
model_id: self.model_id.clone(), model_id: self.model_id.clone(),
gguf_source: self.artifact.to_run_plan_source(), gguf_source: self.artifact.to_run_plan_source(),
num_layers: self.num_layers, num_layers: self.num_layers,
hidden_dim: self.hidden_dim, hidden_dim: self.hidden_dim,
dtype_family: self.dtype_family.into(), dtype_family: self.dtype_family,
dtype_width_bytes: self.dtype_width_bytes, dtype_width_bytes: self.dtype_width_bytes,
max_seq_len: self.max_seq_len, max_seq_len: self.max_seq_len,
eos_token_id: self.eos_token_id, 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)] #[derive(Clone, Debug, PartialEq, Eq)]
pub enum ModelArtifact { pub enum ModelArtifact {
ContainerPath { TestTinyLlm { path: String },
path: String,
},
HuggingFaceGguf {
repo: String,
file: String,
revision: Option<String>,
},
TestTinyLlm {
path: String,
},
} }
impl ModelArtifact { impl ModelArtifact {
fn to_run_plan_source(&self) -> run_plan::GgufSource { fn to_run_plan_source(&self) -> run_plan::GgufSource {
match self { match self {
Self::ContainerPath { path } | Self::TestTinyLlm { path } => { Self::TestTinyLlm { path } => run_plan::GgufSource::LocalPath(path.clone()),
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<DTypeFamily> for run_plan::DTypeFamily {
fn from(value: DTypeFamily) -> Self {
match value {
DTypeFamily::BFloat => Self::BFloat,
} }
} }
} }

31
crates/mvp-system/src/orchestration/engine_builder/node_image.rs
View file

@ -1,28 +1,12 @@
#[derive(Clone, Debug, PartialEq, Eq)] #[derive(Clone, Debug, PartialEq, Eq)]
pub struct NodeImageSpec { pub struct NodeImageSpec;
pub image: String,
pub binary: String,
pub worker_runtime: WorkerRuntimeSpec,
}
impl NodeImageSpec { impl NodeImageSpec {
pub fn new(image: impl Into<String>) -> Self { pub fn new(_image: impl Into<String>) -> Self {
Self { Self
image: image.into(),
binary: "mvp-node".to_owned(),
worker_runtime: WorkerRuntimeSpec::External {
name: "node-image-default".to_owned(),
},
}
} }
pub fn binary(mut self, binary: impl Into<String>) -> Self { pub fn worker_runtime(self, _worker_runtime: WorkerRuntimeSpec) -> Self {
self.binary = binary.into();
self
}
pub fn worker_runtime(mut self, worker_runtime: WorkerRuntimeSpec) -> Self {
self.worker_runtime = worker_runtime;
self self
} }
} }
@ -30,11 +14,4 @@ impl NodeImageSpec {
#[derive(Clone, Debug, PartialEq, Eq)] #[derive(Clone, Debug, PartialEq, Eq)]
pub enum WorkerRuntimeSpec { pub enum WorkerRuntimeSpec {
DumbProcess, DumbProcess,
TinygradCuda {
worker_script: String,
device_env: String,
},
External {
name: String,
},
} }

19
crates/mvp-system/src/orchestration/engine_builder/planner.rs
View file

@ -8,14 +8,8 @@ use super::model::ModelSpec;
use super::pool::NodeCapability; use super::pool::NodeCapability;
use super::roles::{CoordinatorAssignment, StageAssignment}; use super::roles::{CoordinatorAssignment, StageAssignment};
pub trait RolePlanner: Send + Sync {
fn required_node_count(&self) -> usize;
fn plan(&self, input: RolePlannerInput) -> Result<RoleAssignmentPlan, PlanningError>;
}
#[derive(Clone, Debug, PartialEq, Eq)] #[derive(Clone, Debug, PartialEq, Eq)]
pub struct RolePlannerInput { pub struct RolePlannerInput {
pub cluster_id: String,
pub run_id: RunId, pub run_id: RunId,
pub model: ModelSpec, pub model: ModelSpec,
pub nodes: Vec<NodeFacts>, pub nodes: Vec<NodeFacts>,
@ -52,12 +46,12 @@ impl FixedLinearPipelinePlanner {
} }
} }
impl RolePlanner for FixedLinearPipelinePlanner { impl FixedLinearPipelinePlanner {
fn required_node_count(&self) -> usize { pub fn required_node_count(&self) -> usize {
self.stage_count as usize + 1 self.stage_count as usize + 1
} }
fn plan(&self, input: RolePlannerInput) -> Result<RoleAssignmentPlan, PlanningError> { pub fn plan(&self, input: RolePlannerInput) -> Result<RoleAssignmentPlan, PlanningError> {
reject_duplicate_nodes(&input.nodes)?; reject_duplicate_nodes(&input.nodes)?;
let coordinator = input let coordinator = input
.nodes .nodes
@ -107,17 +101,12 @@ impl RolePlanner for FixedLinearPipelinePlanner {
for stage_index in 0..self.stage_count { for stage_index in 0..self.stage_count {
let provision = run_plan::derive_stage_provision(&run_plan, stage_index) let provision = run_plan::derive_stage_provision(&run_plan, stage_index)
.map_err(PlanningError::StageProjection)?; .map_err(PlanningError::StageProjection)?;
stages.push(StageAssignment { stages.push(StageAssignment { provision });
cluster_id: input.cluster_id.clone(),
provision,
});
} }
Ok(RoleAssignmentPlan { Ok(RoleAssignmentPlan {
coordinator: CoordinatorAssignment { coordinator: CoordinatorAssignment {
cluster_id: input.cluster_id,
node_id: coordinator.node_id, node_id: coordinator.node_id,
model: input.model,
}, },
stages, stages,
run_plan, run_plan,

105
crates/mvp-system/src/orchestration/engine_builder/pool.rs
View file

@ -3,68 +3,33 @@ use std::collections::BTreeSet;
use crate::run_plan::NodeId; use crate::run_plan::NodeId;
use super::error::PoolError; use super::error::PoolError;
use super::node_image::NodeImageSpec;
pub trait PoolProvider: Send + Sync {
fn acquire_pool(&self, request: PoolRequest) -> Result<Vec<NodeLease>, PoolError>;
}
#[derive(Clone, Debug, PartialEq, Eq)] #[derive(Clone, Debug, PartialEq, Eq)]
pub struct PoolRequest { pub struct PoolRequest {
pub cluster_id: String,
pub min_nodes: usize, 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)] #[derive(Clone, Debug, PartialEq, Eq)]
pub struct NodeLease { pub struct NodeLease {
pub lease_id: String,
pub logical_node_id: NodeId, pub logical_node_id: NodeId,
pub launch_target: LaunchTarget,
pub expected_resources: ResourceFacts,
pub capabilities: BTreeSet<NodeCapability>, pub capabilities: BTreeSet<NodeCapability>,
} }
impl NodeLease { impl NodeLease {
pub fn new( pub fn new(
lease_id: impl Into<String>, _lease_id: impl Into<String>,
logical_node_id: NodeId, logical_node_id: NodeId,
capabilities: impl IntoIterator<Item = NodeCapability>, capabilities: impl IntoIterator<Item = NodeCapability>,
) -> Self { ) -> Self {
Self { Self {
lease_id: lease_id.into(),
logical_node_id, logical_node_id,
launch_target: LaunchTarget::InProcess,
expected_resources: ResourceFacts::default(),
capabilities: capabilities.into_iter().collect(), capabilities: capabilities.into_iter().collect(),
} }
} }
pub fn launch_target(mut self, launch_target: LaunchTarget) -> Self { pub fn resources(self, _expected_resources: ResourceFacts) -> Self {
self.launch_target = launch_target;
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<String> },
DockerContainer { name: String },
RemoteHost { label: String },
} }
#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] #[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
@ -73,46 +38,12 @@ pub enum NodeCapability {
Worker, Worker,
} }
#[derive(Clone, Debug, PartialEq, Eq)] #[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct ResourceFacts { pub struct ResourceFacts;
pub gpu_count: u32,
pub gpu_memory_bytes: u64,
pub cpu_cores: u32,
pub ram_bytes: u64,
pub cuda_available: bool,
}
impl ResourceFacts { impl ResourceFacts {
pub fn cpu_only(cpu_cores: u32, ram_bytes: u64) -> Self { pub fn cpu_only(_cpu_cores: u32, _ram_bytes: u64) -> Self {
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)
} }
} }
@ -125,30 +56,16 @@ impl StaticPoolProvider {
pub fn new(leases: Vec<NodeLease>) -> Self { pub fn new(leases: Vec<NodeLease>) -> Self {
Self { leases } Self { leases }
} }
pub fn leases(&self) -> &[NodeLease] {
&self.leases
}
} }
impl PoolProvider for StaticPoolProvider { impl StaticPoolProvider {
fn acquire_pool(&self, request: PoolRequest) -> Result<Vec<NodeLease>, PoolError> { pub fn acquire_pool(&self, request: PoolRequest) -> Result<Vec<NodeLease>, PoolError> {
let matching = self if self.leases.len() < request.min_nodes {
.leases
.iter()
.filter(|lease| {
lease
.expected_resources
.satisfies(&request.required_resources)
})
.cloned()
.collect::<Vec<_>>();
if matching.len() < request.min_nodes {
return Err(PoolError::InsufficientNodes { return Err(PoolError::InsufficientNodes {
requested: request.min_nodes, requested: request.min_nodes,
available: matching.len(), available: self.leases.len(),
}); });
} }
Ok(matching) Ok(self.leases.clone())
} }
} }

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

@ -1,17 +1,12 @@
use crate::run_plan::{self, NodeId}; use crate::run_plan::{self, NodeId};
use super::model::ModelSpec;
#[derive(Clone, Debug, PartialEq, Eq)] #[derive(Clone, Debug, PartialEq, Eq)]
pub struct CoordinatorAssignment { pub struct CoordinatorAssignment {
pub cluster_id: String,
pub node_id: NodeId, pub node_id: NodeId,
pub model: ModelSpec,
} }
#[derive(Clone, Debug, PartialEq, Eq)] #[derive(Clone, Debug, PartialEq, Eq)]
pub struct StageAssignment { pub struct StageAssignment {
pub cluster_id: String,
pub provision: run_plan::ProvisionStage, pub provision: run_plan::ProvisionStage,
} }

173
crates/mvp-system/src/orchestration/engine_builder/runtime_stack.rs
View file

@ -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, RuntimeNodeError> {
Self::start_with_codecs(RuntimeNodeConfig::default(), |_| {})
}
pub fn start_with_codecs(
config: RuntimeNodeConfig,
extend_codecs: impl FnOnce(&mut swactor_transport::CodecRegistry),
) -> Result<Self, RuntimeNodeError> {
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<ActorAddress>,
) -> Result<ActorAddress, RuntimeNodeError> {
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<ActorAddress>,
) -> Result<ActorAddress, RuntimeNodeError> {
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 {}

13
crates/mvp-system/src/orchestration/engine_builder/workload.rs
View file

@ -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<Self::Output, Self::Error>;
}

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

@ -6,7 +6,7 @@
//! stays separate from local/Docker/VastAI implementation details. //! stays separate from local/Docker/VastAI implementation details.
pub mod actor; pub mod actor;
mod app; pub(crate) mod app;
pub mod config; pub mod config;
pub mod distribution_stack; pub mod distribution_stack;
#[cfg(test)] #[cfg(test)]
@ -16,20 +16,3 @@ pub mod provider_adapters {
pub mod relay; pub mod relay;
pub(super) mod vastai; pub(super) mod vastai;
} }
pub(super) fn run_from_args<I>(args: I) -> Result<(), String>
where
I: IntoIterator<Item = String>,
{
app::run_from_args(args)
}
pub(super) fn run_in_process_from_args<I>(
args: I,
stop_rx: std::sync::mpsc::Receiver<()>,
) -> Result<(), String>
where
I: IntoIterator<Item = String>,
{
app::run_in_process_from_args(args, stop_rx)
}

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

@ -137,14 +137,6 @@ impl RuntimeConfig {
token_output_policy: TokenOutputPolicy::EmitAll, 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)] #[derive(Clone, Copy, Debug, PartialEq, Eq)]
@ -389,7 +381,11 @@ pub fn plan_run(input: PlannerInput) -> Result<RunPlan, PlanRejection> {
validate_global_input(&input)?; validate_global_input(&input)?;
let placements = validated_placements(&input)?; let placements = validated_placements(&input)?;
let model = model_plan(&input.model)?; 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 max_tokens = input.runtime.max_tokens;
let gguf_source = model.gguf_source.clone(); let gguf_source = model.gguf_source.clone();
let hidden_dim = model.hidden_dim; let hidden_dim = model.hidden_dim;