stash: tested working yoke grind runs

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

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())
}

View file

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

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_BASE_HASH_LABEL: &str = "org.swactor.mvp.node.base-hash";
const BASE_IMAGE_SOURCE_HASH_LABEL: &str = "org.swactor.mvp.base.source-hash";
const NODE_IMAGE_PRUNE_ENV: &str = "MVP_NODE_IMAGE_PRUNE";
const NODE_IMAGE_PRUNE_KEEP_ENV: &str = "MVP_NODE_IMAGE_PRUNE_KEEP";
const DEFAULT_DIRTY_IMAGE_KEEP: usize = 3;
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(super) enum NodeImageProvider {
Docker,
VastAi,
}
impl NodeImageProvider {
fn requires_remote_image(self) -> bool {
matches!(self, Self::VastAi)
}
}
#[derive(Clone, Debug)]
pub(super) struct NodeImageRequest {
pub(super) requested_image: String,
pub(super) base_image: String,
pub(super) node_bin: PathBuf,
pub(super) provider: NodeImageProvider,
pub(super) requires_registry_image: bool,
pub(super) extra_tag: Option<String>,
pub(super) push: bool,
pub(super) force_refresh: bool,
pub(super) enabled: bool,
}
@ -84,62 +67,17 @@ pub(super) trait NodeImageProgressSink {
fn emit(&mut self, event: NodeImageProgressEvent);
}
impl<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(
request: NodeImageRequest,
progress: Option<&mut dyn NodeImageProgressSink>,
) -> Result<String, String> {
let mut progress = progress;
let mut runner = RealImageCommandRunner;
prepare_node_image_inner(request, &mut progress, &mut runner)
prepare_node_image_inner(request, &mut progress)
}
fn prepare_node_image_inner(
request: NodeImageRequest,
progress: &mut Option<&mut dyn NodeImageProgressSink>,
runner: &mut dyn ImageCommandRunner,
) -> Result<String, String> {
emit_image_reference(progress, "requested", &request.requested_image);
if !request.enabled {
@ -149,7 +87,16 @@ fn prepare_node_image_inner(
emit_image_reference(progress, "base", &request.base_image);
let root = workspace_root()?;
let image = ImageName::parse(&request.requested_image)?;
if request.provider.requires_remote_image() && !looks_registry_reachable(&image.repository) {
let first_repository_component = image
.repository
.split('/')
.next()
.unwrap_or(&image.repository);
let registry_reachable = image.repository.contains('/')
|| first_repository_component.contains('.')
|| first_repository_component.contains(':')
|| first_repository_component == "localhost";
if request.requires_registry_image && !registry_reachable {
return Err(format!(
"VastAI node image {:?} must include a registry namespace",
image.repository
@ -157,7 +104,6 @@ fn prepare_node_image_inner(
}
run_status(
runner,
progress,
&root,
"cargo",
@ -178,47 +124,54 @@ fn prepare_node_image_inner(
let tag = image_version_tag(&root, &image_content_hash)?;
let image_ref = image.ref_for_tag(&tag);
emit_image_reference(progress, "resolved", &image_ref);
let worker_hash = file_content_hash(&root, Path::new("apps/mvp-node/tinygrad_worker.py"))?;
let expected_node_labels =
node_image_labels(&tag, &image_content_hash, &worker_hash, &base_hash);
let expected_base_labels = base_image_labels(&base_hash);
let worker_hash = hash_relative_files(
&root,
vec![relative_path(
&root,
&root.join("apps/mvp-node/tinygrad_worker.py"),
)?],
)?;
let expected_node_labels = vec![
(NODE_IMAGE_TAG_LABEL, tag.as_str()),
(NODE_IMAGE_SOURCE_HASH_LABEL, image_content_hash.as_str()),
(NODE_IMAGE_WORKER_HASH_LABEL, worker_hash.as_str()),
(NODE_IMAGE_BASE_HASH_LABEL, base_hash.as_str()),
];
let expected_base_labels = vec![(BASE_IMAGE_SOURCE_HASH_LABEL, base_hash.as_str())];
let alias_tags = alias_tags(&image, request.extra_tag.as_deref(), &tag)?;
for alias in alias_refs(&image, &alias_tags) {
emit_image_reference(progress, "alias", &alias);
}
let remote_required = request.provider.requires_remote_image() || request.push;
let remote_required = request.requires_registry_image;
let local_image_matches =
docker_image_labels_match(runner, &root, &image_ref, &expected_node_labels)?;
let remote_available = remote_required && runner.docker_manifest_exists(&root, &image_ref);
let local_image_matches = docker_image_labels_match(&root, &image_ref, &expected_node_labels)?;
let remote_available = remote_required && docker_manifest_exists(&root, &image_ref);
if !request.force_refresh && remote_required && remote_available {
ensure_aliases_for_remote(runner, progress, &root, &image_ref, &image, &alias_tags)?;
prune_old_dirty_images(runner, &root, &image, &tag);
ensure_aliases_for_remote(progress, &root, &image_ref, &image, &alias_tags)?;
prune_old_dirty_images(&root, &image, &tag);
return Ok(image_ref);
}
if !request.force_refresh && remote_required && local_image_matches {
ensure_aliases_local(runner, progress, &root, &image_ref, &image, &alias_tags)?;
push_image(runner, progress, &root, &image_ref)?;
ensure_aliases_local(progress, &root, &image_ref, &image, &alias_tags)?;
push_image(progress, &root, &image_ref)?;
for alias in alias_refs(&image, &alias_tags) {
push_image(runner, progress, &root, &alias)?;
push_image(progress, &root, &alias)?;
}
prune_old_dirty_images(runner, &root, &image, &tag);
prune_old_dirty_images(&root, &image, &tag);
return Ok(image_ref);
}
if !request.force_refresh && !remote_required && local_image_matches {
ensure_aliases_local(runner, progress, &root, &image_ref, &image, &alias_tags)?;
prune_old_dirty_images(runner, &root, &image, &tag);
ensure_aliases_local(progress, &root, &image_ref, &image, &alias_tags)?;
prune_old_dirty_images(&root, &image, &tag);
return Ok(image_ref);
}
let base_image_matches =
docker_image_labels_match(runner, &root, &request.base_image, &expected_base_labels)?;
docker_image_labels_match(&root, &request.base_image, &expected_base_labels)?;
if !base_image_matches {
run_status_vec(
runner,
progress,
run_status_command(
&root,
"docker",
vec![
&vec![
"build".to_owned(),
"-f".to_owned(),
"apps/mvp-node/Dockerfile.base".to_owned(),
@ -230,6 +183,7 @@ fn prepare_node_image_inner(
],
"build mvp node base image",
Some(&request.base_image),
progress,
)?;
}
@ -248,25 +202,24 @@ fn prepare_node_image_inner(
build_args.push(format!("{key}={value}"));
}
build_args.extend(["-t".to_owned(), image_ref.clone(), ".".to_owned()]);
run_status_vec(
runner,
progress,
run_status_command(
&root,
"docker",
build_args,
&build_args,
"build mvp node image",
Some(&image_ref),
progress,
)?;
ensure_aliases_local(runner, progress, &root, &image_ref, &image, &alias_tags)?;
ensure_aliases_local(progress, &root, &image_ref, &image, &alias_tags)?;
if remote_required {
push_image(runner, progress, &root, &image_ref)?;
push_image(progress, &root, &image_ref)?;
for alias in alias_refs(&image, &alias_tags) {
push_image(runner, progress, &root, &alias)?;
push_image(progress, &root, &alias)?;
}
}
prune_old_dirty_images(runner, &root, &image, &tag);
prune_old_dirty_images(&root, &image, &tag);
Ok(image_ref)
}
@ -290,7 +243,10 @@ fn workspace_root() -> Result<PathBuf, 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"])?;
Ok(format!("git-{}", sha.trim()))
} 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> {
let output = Command::new("git")
.current_dir(root)
@ -355,10 +305,6 @@ fn content_hash_for_inputs(root: &Path, inputs: &[&str]) -> Result<String, Strin
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> {
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 metadata = fs::metadata(path).map_err(|e| format!("stat {display}: {e}"))?;
if metadata.is_file() {
if !skip_file(path) {
if !matches!(path.extension().and_then(|ext| ext.to_str()), Some("pyc")) {
out.push(relative_path(root, path)?);
}
return Ok(());
@ -459,10 +405,6 @@ fn skip_dir(path: &Path) -> bool {
)
}
fn skip_file(path: &Path) -> bool {
matches!(path.extension().and_then(|ext| ext.to_str()), Some("pyc"))
}
fn alias_tags(
image: &ImageName,
extra_tag: Option<&str>,
@ -493,26 +435,7 @@ fn insert_alias_tag(
Ok(())
}
fn node_image_labels<'a>(
tag: &'a str,
source_hash: &'a str,
worker_hash: &'a str,
base_hash: &'a str,
) -> Vec<(&'static str, &'a str)> {
vec![
(NODE_IMAGE_TAG_LABEL, tag),
(NODE_IMAGE_SOURCE_HASH_LABEL, source_hash),
(NODE_IMAGE_WORKER_HASH_LABEL, worker_hash),
(NODE_IMAGE_BASE_HASH_LABEL, base_hash),
]
}
fn base_image_labels<'a>(base_hash: &'a str) -> Vec<(&'static str, &'a str)> {
vec![(BASE_IMAGE_SOURCE_HASH_LABEL, base_hash)]
}
fn ensure_aliases_local(
runner: &mut dyn ImageCommandRunner,
progress: &mut Option<&mut dyn NodeImageProgressSink>,
root: &Path,
source_ref: &str,
@ -522,7 +445,6 @@ fn ensure_aliases_local(
for alias in alias_refs(image, alias_tags) {
if alias != source_ref {
run_status(
runner,
progress,
root,
"docker",
@ -536,7 +458,6 @@ fn ensure_aliases_local(
}
fn ensure_aliases_for_remote(
runner: &mut dyn ImageCommandRunner,
progress: &mut Option<&mut dyn NodeImageProgressSink>,
root: &Path,
source_ref: &str,
@ -546,9 +467,8 @@ fn ensure_aliases_for_remote(
if alias_tags.is_empty() {
return Ok(false);
}
if !runner.docker_image_exists(root, source_ref) {
if !docker_image_exists(root, source_ref) {
run_status(
runner,
progress,
root,
"docker",
@ -557,9 +477,9 @@ fn ensure_aliases_for_remote(
Some(source_ref),
)?;
}
ensure_aliases_local(runner, progress, root, source_ref, image, alias_tags)?;
ensure_aliases_local(progress, root, source_ref, image, alias_tags)?;
for alias in alias_refs(image, alias_tags) {
push_image(runner, progress, root, &alias)?;
push_image(progress, root, &alias)?;
}
Ok(true)
}
@ -572,13 +492,11 @@ fn alias_refs(image: &ImageName, alias_tags: &BTreeSet<String>) -> Vec<String> {
}
fn push_image(
runner: &mut dyn ImageCommandRunner,
progress: &mut Option<&mut dyn NodeImageProgressSink>,
root: &Path,
image_ref: &str,
) -> Result<(), String> {
run_status(
runner,
progress,
root,
"docker",
@ -588,25 +506,12 @@ fn push_image(
)
}
fn docker_image_exists(root: &Path, image_ref: &str) -> bool {
Command::new("docker")
.current_dir(root)
.args(["image", "inspect", image_ref])
.stdin(Stdio::null())
.stdout(Stdio::null())
.stderr(Stdio::null())
.status()
.map(|status| status.success())
.unwrap_or(false)
}
fn docker_image_labels_match(
runner: &mut dyn ImageCommandRunner,
root: &Path,
image_ref: &str,
expected: &[(&str, &str)],
) -> Result<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);
};
Ok(expected
@ -614,54 +519,18 @@ fn docker_image_labels_match(
.all(|(key, value)| labels.get(*key).map(String::as_str) == Some(*value)))
}
fn docker_image_labels(
root: &Path,
image_ref: &str,
) -> Result<Option<BTreeMap<String, String>>, String> {
let output = Command::new("docker")
.current_dir(root)
.args([
"image",
"inspect",
"--format",
"{{ json .Config.Labels }}",
image_ref,
])
.stdin(Stdio::null())
.output()
.map_err(|e| format!("inspect docker image {image_ref}: {e}"))?;
if !output.status.success() {
return Ok(None);
}
let stdout = String::from_utf8_lossy(&output.stdout);
let labels: Option<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() {
fn prune_old_dirty_images(root: &Path, image: &ImageName, keep_tag: &str) {
let prune_enabled = std::env::var("MVP_NODE_IMAGE_PRUNE")
.map(|value| {
let value = value.trim().to_ascii_lowercase();
!matches!(value.as_str(), "0" | "false" | "no" | "off")
})
.unwrap_or(true);
if !prune_enabled {
return;
}
let tags = match runner.docker_image_tags(root, &image.repository) {
let tags = match docker_image_tags(root, &image.repository) {
Ok(tags) => tags,
Err(error) => {
eprintln!("mvp-node-image: prune old dirty images skipped: {error}");
@ -669,7 +538,10 @@ fn prune_old_dirty_images(
}
};
let keep_old = dirty_image_prune_keep();
let keep_old = std::env::var("MVP_NODE_IMAGE_PRUNE_KEEP")
.ok()
.and_then(|value| value.trim().parse::<usize>().ok())
.unwrap_or(3);
let mut retained_old = 0_usize;
for (repository, tag) in tags {
if repository != image.repository
@ -681,7 +553,7 @@ fn prune_old_dirty_images(
}
let image_ref = image.ref_for_tag(&tag);
let Ok(Some(labels)) = runner.docker_image_labels(root, &image_ref) else {
let Ok(Some(labels)) = docker_image_labels(root, &image_ref) else {
continue;
};
if labels.get(NODE_IMAGE_TAG_LABEL).map(String::as_str) != Some(tag.as_str())
@ -691,7 +563,7 @@ fn prune_old_dirty_images(
{
continue;
}
if runner.docker_image_has_container(root, &image_ref) {
if docker_image_has_container(root, &image_ref) {
eprintln!(
"mvp-node-image: prune old dirty image {image_ref} skipped: container exists"
);
@ -703,47 +575,13 @@ fn prune_old_dirty_images(
}
eprintln!("mvp-node-image: prune old dirty image {image_ref}");
if let Err(error) = runner.docker_image_remove(root, &image_ref) {
if let Err(error) = docker_image_remove(root, &image_ref) {
eprintln!("mvp-node-image: prune old dirty image {image_ref} skipped: {error}");
}
}
}
fn dirty_image_prune_enabled() -> bool {
std::env::var(NODE_IMAGE_PRUNE_ENV)
.map(|value| {
let value = value.trim().to_ascii_lowercase();
!matches!(value.as_str(), "0" | "false" | "no" | "off")
})
.unwrap_or(true)
}
fn dirty_image_prune_keep() -> usize {
std::env::var(NODE_IMAGE_PRUNE_KEEP_ENV)
.ok()
.and_then(|value| value.trim().parse::<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(
runner: &mut dyn ImageCommandRunner,
progress: &mut Option<&mut dyn NodeImageProgressSink>,
root: &Path,
program: &str,
@ -752,19 +590,7 @@ fn run_status(
image_ref: Option<&str>,
) -> Result<(), String> {
let args = args.iter().map(|arg| (*arg).to_owned()).collect::<Vec<_>>();
runner.run_status(root, program, &args, label, image_ref, progress)
}
fn run_status_vec(
runner: &mut dyn ImageCommandRunner,
progress: &mut Option<&mut dyn NodeImageProgressSink>,
root: &Path,
program: &str,
args: Vec<String>,
label: &str,
image_ref: Option<&str>,
) -> Result<(), String> {
runner.run_status(root, program, &args, label, image_ref, progress)
run_status_command(root, program, &args, label, image_ref, progress)
}
fn emit_image_reference(
@ -857,205 +683,239 @@ fn drain_command_lines(
}
}
impl ImageCommandRunner for RealImageCommandRunner {
fn run_status(
&mut self,
root: &Path,
program: &str,
args: &[String],
label: &str,
image_ref: Option<&str>,
progress: &mut Option<&mut dyn NodeImageProgressSink>,
) -> Result<(), String> {
eprintln!("mvp-node-image: {label}");
if progress.is_none() {
let status = Command::new(program)
.current_dir(root)
.args(args)
.stdin(Stdio::null())
.stdout(Stdio::inherit())
.stderr(Stdio::inherit())
.status()
.map_err(|e| format!("run {label}: {e}"))?;
return if status.success() {
Ok(())
} else {
Err(format!("{label} failed with {status}"))
};
}
let started = Instant::now();
emit_command_progress(
progress,
label,
image_ref,
0,
NodeImageProgressEventKind::CommandStarted {
program: program.to_owned(),
args: args.to_vec(),
},
);
let mut child = match Command::new(program)
fn run_status_command(
root: &Path,
program: &str,
args: &[String],
label: &str,
image_ref: Option<&str>,
progress: &mut Option<&mut dyn NodeImageProgressSink>,
) -> Result<(), String> {
eprintln!("mvp-node-image: {label}");
if progress.is_none() {
let status = Command::new(program)
.current_dir(root)
.args(args)
.stdin(Stdio::null())
.stdout(Stdio::piped())
.stderr(Stdio::piped())
.spawn()
{
Ok(child) => child,
.stdout(Stdio::inherit())
.stderr(Stdio::inherit())
.status()
.map_err(|e| format!("run {label}: {e}"))?;
return if status.success() {
Ok(())
} else {
Err(format!("{label} failed with {status}"))
};
}
let started = Instant::now();
emit_command_progress(
progress,
label,
image_ref,
0,
NodeImageProgressEventKind::CommandStarted {
program: program.to_owned(),
args: args.to_vec(),
},
);
let mut child = match Command::new(program)
.current_dir(root)
.args(args)
.stdin(Stdio::null())
.stdout(Stdio::piped())
.stderr(Stdio::piped())
.spawn()
{
Ok(child) => child,
Err(error) => {
emit_command_progress(
progress,
label,
image_ref,
started.elapsed().as_millis(),
NodeImageProgressEventKind::CommandExited {
status: format!("spawn error: {error}"),
code: None,
success: false,
},
);
return Err(format!("run {label}: {error}"));
}
};
let (tx, rx) = mpsc::channel();
let mut readers = Vec::new();
if let Some(stdout) = child.stdout.take() {
readers.push(spawn_line_reader(
stdout,
CommandOutputLine::Stdout,
tx.clone(),
));
}
if let Some(stderr) = child.stderr.take() {
readers.push(spawn_line_reader(
stderr,
CommandOutputLine::Stderr,
tx.clone(),
));
}
drop(tx);
let status = loop {
match child.try_wait() {
Ok(Some(status)) => break status,
Ok(None) => {
drain_command_lines(&rx, progress, label, image_ref, started);
thread::sleep(Duration::from_millis(10));
}
Err(error) => {
drain_command_lines(&rx, progress, label, image_ref, started);
emit_command_progress(
progress,
label,
image_ref,
started.elapsed().as_millis(),
NodeImageProgressEventKind::CommandExited {
status: format!("spawn error: {error}"),
status: format!("wait error: {error}"),
code: None,
success: false,
},
);
return Err(format!("run {label}: {error}"));
}
};
let (tx, rx) = mpsc::channel();
let mut readers = Vec::new();
if let Some(stdout) = child.stdout.take() {
readers.push(spawn_line_reader(
stdout,
CommandOutputLine::Stdout,
tx.clone(),
));
}
if let Some(stderr) = child.stderr.take() {
readers.push(spawn_line_reader(
stderr,
CommandOutputLine::Stderr,
tx.clone(),
));
}
drop(tx);
let status = loop {
match child.try_wait() {
Ok(Some(status)) => break status,
Ok(None) => {
drain_command_lines(&rx, progress, label, image_ref, started);
thread::sleep(Duration::from_millis(10));
}
Err(error) => {
drain_command_lines(&rx, progress, label, image_ref, started);
emit_command_progress(
progress,
label,
image_ref,
started.elapsed().as_millis(),
NodeImageProgressEventKind::CommandExited {
status: format!("wait error: {error}"),
code: None,
success: false,
},
);
return Err(format!("run {label}: {error}"));
}
}
};
for reader in readers {
let _ = reader.join();
}
drain_command_lines(&rx, progress, label, image_ref, started);
let status_text = status.to_string();
let success = status.success();
emit_command_progress(
progress,
label,
image_ref,
started.elapsed().as_millis(),
NodeImageProgressEventKind::CommandExited {
status: status_text.clone(),
code: status.code(),
success,
},
);
if success {
Ok(())
} else {
Err(format!("{label} failed with {status_text}"))
}
};
for reader in readers {
let _ = reader.join();
}
fn docker_image_exists(&mut self, root: &Path, image_ref: &str) -> bool {
docker_image_exists(root, image_ref)
}
fn docker_image_labels(
&mut self,
root: &Path,
image_ref: &str,
) -> Result<Option<BTreeMap<String, String>>, String> {
docker_image_labels(root, image_ref)
}
fn docker_manifest_exists(&mut self, root: &Path, image_ref: &str) -> bool {
docker_manifest_exists(root, image_ref)
}
fn docker_image_has_container(&mut self, root: &Path, image_ref: &str) -> bool {
docker_image_has_container(root, image_ref)
}
fn docker_image_tags(
&mut self,
root: &Path,
repository: &str,
) -> Result<Vec<(String, String)>, String> {
let output = Command::new("docker")
.current_dir(root)
.args([
"image",
"ls",
"--format",
"{{.Repository}}\t{{.Tag}}",
repository,
])
.stdin(Stdio::null())
.output()
.map_err(|error| format!("docker image ls failed: {error}"))?;
if !output.status.success() {
return Err(format!("docker image ls failed with {}", output.status));
}
let stdout = String::from_utf8_lossy(&output.stdout);
Ok(stdout
.lines()
.filter_map(|line| {
let (repository, tag) = line.split_once('\t')?;
Some((repository.to_owned(), tag.to_owned()))
})
.collect())
}
fn docker_image_remove(&mut self, root: &Path, image_ref: &str) -> Result<(), String> {
let status = Command::new("docker")
.current_dir(root)
.args(["image", "rm", image_ref])
.stdin(Stdio::null())
.stdout(Stdio::null())
.stderr(Stdio::null())
.status()
.map_err(|error| format!("docker image rm failed: {error}"))?;
if status.success() {
Ok(())
} else {
Err(format!("docker image rm failed with {status}"))
}
drain_command_lines(&rx, progress, label, image_ref, started);
let status_text = status.to_string();
let success = status.success();
emit_command_progress(
progress,
label,
image_ref,
started.elapsed().as_millis(),
NodeImageProgressEventKind::CommandExited {
status: status_text.clone(),
code: status.code(),
success,
},
);
if success {
Ok(())
} else {
Err(format!("{label} failed with {status_text}"))
}
}
fn looks_registry_reachable(repository: &str) -> bool {
let first = repository.split('/').next().unwrap_or(repository);
repository.contains('/') || first.contains('.') || first.contains(':') || first == "localhost"
fn docker_image_exists(root: &Path, image_ref: &str) -> bool {
Command::new("docker")
.current_dir(root)
.args(["image", "inspect", image_ref])
.stdin(Stdio::null())
.stdout(Stdio::null())
.stderr(Stdio::null())
.status()
.map(|status| status.success())
.unwrap_or(false)
}
fn docker_image_labels(
root: &Path,
image_ref: &str,
) -> Result<Option<BTreeMap<String, String>>, String> {
let output = Command::new("docker")
.current_dir(root)
.args([
"image",
"inspect",
"--format",
"{{ json .Config.Labels }}",
image_ref,
])
.stdin(Stdio::null())
.output()
.map_err(|e| format!("inspect docker image {image_ref}: {e}"))?;
if !output.status.success() {
return Ok(None);
}
let stdout = String::from_utf8_lossy(&output.stdout);
let labels: Option<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 docker_image_has_container(root: &Path, image_ref: &str) -> bool {
Command::new("docker")
.current_dir(root)
.args([
"ps",
"-a",
"--filter",
&format!("ancestor={image_ref}"),
"--format",
"{{.ID}}",
])
.stdin(Stdio::null())
.output()
.map(|output| output.status.success() && !output.stdout.is_empty())
.unwrap_or(true)
}
fn docker_image_tags(root: &Path, repository: &str) -> Result<Vec<(String, String)>, String> {
let output = Command::new("docker")
.current_dir(root)
.args([
"image",
"ls",
"--format",
"{{.Repository}}\t{{.Tag}}",
repository,
])
.stdin(Stdio::null())
.output()
.map_err(|error| format!("docker image ls failed: {error}"))?;
if !output.status.success() {
return Err(format!("docker image ls failed with {}", output.status));
}
let stdout = String::from_utf8_lossy(&output.stdout);
Ok(stdout
.lines()
.filter_map(|line| {
let (repository, tag) = line.split_once('\t')?;
Some((repository.to_owned(), tag.to_owned()))
})
.collect())
}
fn docker_image_remove(root: &Path, image_ref: &str) -> Result<(), String> {
let status = Command::new("docker")
.current_dir(root)
.args(["image", "rm", image_ref])
.stdin(Stdio::null())
.stdout(Stdio::null())
.stderr(Stdio::null())
.status()
.map_err(|error| format!("docker image rm failed: {error}"))?;
if status.success() {
Ok(())
} else {
Err(format!("docker image rm failed with {status}"))
}
}
#[derive(Clone, Debug)]

View file

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

View file

@ -12,28 +12,18 @@ pub fn run_chat_from_args<I>(args: I) -> std::process::ExitCode
where
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>
where
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 {
node::run_worker_node_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)
node::worker_node_runtime::run_from_env()
}
#[path = "transport/driver_pumps.rs"]

View file

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

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

File diff suppressed because it is too large Load diff

View file

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

View file

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

View file

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

View file

@ -16,15 +16,14 @@ pub mod node_image;
pub mod planner;
pub mod pool;
pub mod roles;
pub mod runtime_stack;
pub mod workload;
pub use crate::run_plan::NodeId;
pub use crate::run_plan::{DTypeFamily, NodeId};
pub use engine::ClusterBuilder;
pub use events::EngineEvent;
pub use launcher::StaticNodeLauncher;
pub use model::{DTypeFamily, ModelArtifact, ModelSpec};
pub use node_image::{NodeImageSpec, WorkerRuntimeSpec};
pub use model::{ModelArtifact, ModelSpec};
pub use planner::FixedLinearPipelinePlanner;
pub use pool::{NodeCapability, NodeLease, ResourceFacts, StaticPoolProvider};
pub use roles::RoleKind;
pub use node_image::{NodeImageSpec, WorkerRuntimeSpec};

View file

@ -3,12 +3,11 @@ use crate::run_plan;
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct ModelSpec {
pub model_id: String,
pub architecture: ModelArchitecture,
pub artifact: ModelArtifact,
pub tokenizer: run_plan::TokenizerSource,
pub num_layers: u32,
pub hidden_dim: u64,
pub dtype_family: DTypeFamily,
pub dtype_family: run_plan::DTypeFamily,
pub dtype_width_bytes: u64,
pub max_seq_len: u64,
pub eos_token_id: u32,
@ -20,7 +19,7 @@ impl ModelSpec {
artifact: ModelArtifact,
num_layers: u32,
hidden_dim: u64,
dtype_family: DTypeFamily,
dtype_family: run_plan::DTypeFamily,
dtype_width_bytes: u64,
max_seq_len: u64,
eos_token_id: u32,
@ -28,7 +27,6 @@ impl ModelSpec {
) -> Self {
Self {
model_id: model_id.into(),
architecture: ModelArchitecture::PipelinedCausalLlm,
artifact,
tokenizer,
num_layers,
@ -40,29 +38,13 @@ impl ModelSpec {
}
}
pub fn mvp_tiny_open_llm_fixture() -> Self {
Self::pipelined_causal_llm(
"mvp-tiny-open-llm-fixture",
ModelArtifact::ContainerPath {
path: "/models/mvp-tiny-open-llm.gguf".to_owned(),
},
4,
8,
DTypeFamily::BFloat,
2,
8,
99,
run_plan::TokenizerSource::EmbeddedGguf,
)
}
pub fn to_run_plan_facts(&self) -> run_plan::ModelFacts {
run_plan::ModelFacts {
model_id: self.model_id.clone(),
gguf_source: self.artifact.to_run_plan_source(),
num_layers: self.num_layers,
hidden_dim: self.hidden_dim,
dtype_family: self.dtype_family.into(),
dtype_family: self.dtype_family,
dtype_width_bytes: self.dtype_width_bytes,
max_seq_len: self.max_seq_len,
eos_token_id: self.eos_token_id,
@ -71,54 +53,15 @@ impl ModelSpec {
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum ModelArchitecture {
PipelinedCausalLlm,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum ModelArtifact {
ContainerPath {
path: String,
},
HuggingFaceGguf {
repo: String,
file: String,
revision: Option<String>,
},
TestTinyLlm {
path: String,
},
TestTinyLlm { path: String },
}
impl ModelArtifact {
fn to_run_plan_source(&self) -> run_plan::GgufSource {
match self {
Self::ContainerPath { path } | Self::TestTinyLlm { path } => {
run_plan::GgufSource::LocalPath(path.clone())
}
Self::HuggingFaceGguf {
repo,
file,
revision,
} => run_plan::GgufSource::HuggingFaceGguf {
repo: repo.clone(),
file: file.clone(),
revision: revision.clone(),
},
}
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum DTypeFamily {
BFloat,
}
impl From<DTypeFamily> for run_plan::DTypeFamily {
fn from(value: DTypeFamily) -> Self {
match value {
DTypeFamily::BFloat => Self::BFloat,
Self::TestTinyLlm { path } => run_plan::GgufSource::LocalPath(path.clone()),
}
}
}

View file

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

View file

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

View file

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

View file

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

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 {}

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>;
}

View file

@ -6,7 +6,7 @@
//! stays separate from local/Docker/VastAI implementation details.
pub mod actor;
mod app;
pub(crate) mod app;
pub mod config;
pub mod distribution_stack;
#[cfg(test)]
@ -16,20 +16,3 @@ pub mod provider_adapters {
pub mod relay;
pub(super) mod vastai;
}
pub(super) fn run_from_args<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)
}

View file

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