stash: tested working yoke grind runs
This commit is contained in:
parent
599e8e870b
commit
e8821c85be
21 changed files with 1073 additions and 2148 deletions
|
|
@ -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())
|
||||
}
|
||||
|
|
@ -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;
|
||||
|
|
|
|||
|
|
@ -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)]
|
||||
|
|
|
|||
|
|
@ -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}"))
|
||||
|
|
|
|||
|
|
@ -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"]
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
|
|||
|
|
@ -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
|
|
@ -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>,
|
||||
|
|
|
|||
|
|
@ -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 {}
|
||||
|
|
|
|||
|
|
@ -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(())
|
||||
|
|
|
|||
|
|
@ -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};
|
||||
|
|
|
|||
|
|
@ -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()),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
},
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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())
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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 {}
|
||||
|
|
@ -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>;
|
||||
}
|
||||
|
|
@ -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)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
|
|||
Loading…
Reference in a new issue