|
@@ -0,0 +1,285 @@
|
|
|
|
|
+//! `local-view` staging block executor.
|
|
|
|
|
+//!
|
|
|
|
|
+//! Some checkpoints stream a large side table at every token (the DSV4.1
|
|
|
|
|
+//! Engram tables read ~1 MB of rows per token). Serving those rows from the
|
|
|
|
|
+//! shared NAS costs an order of magnitude more than from local NVMe, while the
|
|
|
|
|
+//! bulk of the weights belongs on shared storage. A `local-view` block
|
|
|
|
|
+//! materializes, per launch node, a directory holding **real copies** of the
|
|
|
|
|
+//! files matching `local` and symlinks (into the shared cache's container
|
|
|
|
|
+//! path) for everything else; `serve` mounts it and points `model` at it.
|
|
|
|
|
+//!
|
|
|
|
|
+//! The node-side work is `scripts/local_view.py`; this module only builds the
|
|
|
|
|
+//! command lines, decides the container mount and computes the bytes a node
|
|
|
|
|
+//! must have free (used by the pre-start free-space check).
|
|
|
|
|
+
|
|
|
|
|
+use crate::config::{Config, Node, Staging, expand_tilde};
|
|
|
|
|
+use crate::ssh::{run_ssh_long_running, upload_script};
|
|
|
|
|
+use anyhow::{Result, bail};
|
|
|
|
|
+use std::path::Path;
|
|
|
|
|
+use tracing::info;
|
|
|
|
|
+
|
|
|
|
|
+pub const SCRIPT_BODY: &str = include_str!("scripts/local_view.py");
|
|
|
|
|
+
|
|
|
|
|
+/// Container-side path the view is mounted at (the model references it).
|
|
|
|
|
+pub const CONTAINER_ROOT: &str = "/root/.cache/hf-view";
|
|
|
|
|
+
|
|
|
|
|
+fn script_path(id: &str) -> String {
|
|
|
|
|
+ format!("/tmp/sm121-local-view-{id}.py")
|
|
|
|
|
+}
|
|
|
|
|
+
|
|
|
|
|
+/// The repo a block views: explicit `repo`, else the recipe's model.
|
|
|
|
|
+pub fn repo_for(block: &Staging, model: &str) -> String {
|
|
|
|
|
+ block.repo.clone().unwrap_or_else(|| model.to_string())
|
|
|
|
|
+}
|
|
|
|
|
+
|
|
|
|
|
+/// Node-side view directory (tilde-expanded).
|
|
|
|
|
+pub fn dst_for(block: &Staging, id: &str, ssh_user: &str) -> String {
|
|
|
|
|
+ block
|
|
|
|
|
+ .dst
|
|
|
|
|
+ .as_deref()
|
|
|
|
|
+ .map(|d| expand_tilde(d, ssh_user))
|
|
|
|
|
+ .unwrap_or_else(|| format!("~/.cache/hf-view/{id}"))
|
|
|
|
|
+}
|
|
|
|
|
+
|
|
|
|
|
+/// `measure` (bytes a node needs) or `apply` (build the view).
|
|
|
|
|
+pub struct ViewSpec<'a> {
|
|
|
|
|
+ pub repo: &'a str,
|
|
|
|
|
+ pub local: &'a [String],
|
|
|
|
|
+ pub verify: &'a str,
|
|
|
|
|
+ pub min_free: Option<&'a str>,
|
|
|
|
|
+ pub force: bool,
|
|
|
|
|
+}
|
|
|
|
|
+
|
|
|
|
|
+/// Command line for `python3 <script> apply --src … --repo … --dst … --local …`.
|
|
|
|
|
+pub fn command_for(id: &str, mode: &str, src: &str, dst: &str, spec: &ViewSpec<'_>) -> String {
|
|
|
|
|
+ let mut cmd = format!(
|
|
|
|
|
+ "python3 {} {} --src {} --repo {} --dst {} --container-root {} --verify {}",
|
|
|
|
|
+ crate::docker::shq(&script_path(id)),
|
|
|
|
|
+ mode,
|
|
|
|
|
+ crate::docker::shq(src),
|
|
|
|
|
+ crate::docker::shq(spec.repo),
|
|
|
|
|
+ crate::docker::shq(dst),
|
|
|
|
|
+ crate::docker::shq(CONTAINER_ROOT),
|
|
|
|
|
+ spec.verify,
|
|
|
|
|
+ );
|
|
|
|
|
+ for g in spec.local {
|
|
|
|
|
+ cmd.push_str(&format!(" --local {}", crate::docker::shq(g)));
|
|
|
|
|
+ }
|
|
|
|
|
+ if let Some(m) = spec.min_free {
|
|
|
|
|
+ cmd.push_str(&format!(" --min-free {}", crate::docker::shq(m)));
|
|
|
|
|
+ }
|
|
|
|
|
+ if spec.force {
|
|
|
|
|
+ cmd.push_str(" --force");
|
|
|
|
|
+ }
|
|
|
|
|
+ cmd
|
|
|
|
|
+}
|
|
|
|
|
+
|
|
|
|
|
+/// What one node needs free for a view, plus what it has (bytes).
|
|
|
|
|
+#[derive(Debug, Clone, PartialEq, Eq)]
|
|
|
|
|
+pub struct Need {
|
|
|
|
|
+ pub node: String,
|
|
|
|
|
+ pub repo: String,
|
|
|
|
|
+ pub need: u64,
|
|
|
|
|
+ pub free: u64,
|
|
|
|
|
+ pub dst: String,
|
|
|
|
|
+}
|
|
|
|
|
+
|
|
|
|
|
+impl Need {
|
|
|
|
|
+ pub fn short(&self) -> bool {
|
|
|
|
|
+ self.free < self.need
|
|
|
|
|
+ }
|
|
|
|
|
+
|
|
|
|
|
+ pub fn line(&self) -> String {
|
|
|
|
|
+ format!(
|
|
|
|
|
+ "{}: {} → needs {:.1} GiB, has {:.1} GiB{}",
|
|
|
|
|
+ self.node,
|
|
|
|
|
+ self.repo,
|
|
|
|
|
+ self.need as f64 / (1u64 << 30) as f64,
|
|
|
|
|
+ self.free as f64 / (1u64 << 30) as f64,
|
|
|
|
|
+ if self.short() { " ⚠ SHORT" } else { "" }
|
|
|
|
|
+ )
|
|
|
|
|
+ }
|
|
|
|
|
+}
|
|
|
|
|
+
|
|
|
|
|
+/// Parse the `measure` output (`{"need": …, "free": …}`) of the node script.
|
|
|
|
|
+pub fn parse_measure(node: &str, repo: &str, dst: &str, stdout: &str) -> Result<Need> {
|
|
|
|
|
+ let v: serde_json::Value = serde_json::from_str(stdout.trim()).map_err(|e| {
|
|
|
|
|
+ anyhow::anyhow!("local-view measure on {node}: unparsable output {stdout:?}: {e}")
|
|
|
|
|
+ })?;
|
|
|
|
|
+ let num = |k: &str| v.get(k).and_then(serde_json::Value::as_u64).unwrap_or(0);
|
|
|
|
|
+ Ok(Need {
|
|
|
|
|
+ node: node.to_string(),
|
|
|
|
|
+ repo: repo.to_string(),
|
|
|
|
|
+ need: num("need"),
|
|
|
|
|
+ free: num("free"),
|
|
|
|
|
+ dst: dst.to_string(),
|
|
|
|
|
+ })
|
|
|
|
|
+}
|
|
|
|
|
+
|
|
|
|
|
+/// The aggregated, human-facing shortfall message (`None` when every node fits).
|
|
|
|
|
+pub fn shortfall(needs: &[Need]) -> Option<String> {
|
|
|
|
|
+ let short: Vec<&Need> = needs.iter().filter(|n| n.short()).collect();
|
|
|
|
|
+ if short.is_empty() {
|
|
|
|
|
+ return None;
|
|
|
|
|
+ }
|
|
|
|
|
+ let mut msg =
|
|
|
|
|
+ String::from("not enough free space on the stage target(s) — nothing was started:\n");
|
|
|
|
|
+ for n in &short {
|
|
|
|
|
+ msg.push_str(&format!(
|
|
|
|
|
+ " {}: {} needs {:.1} GiB in {}, has {:.1} GiB\n",
|
|
|
|
|
+ n.node,
|
|
|
|
|
+ n.repo,
|
|
|
|
|
+ n.need as f64 / (1u64 << 30) as f64,
|
|
|
|
|
+ n.dst,
|
|
|
|
|
+ n.free as f64 / (1u64 << 30) as f64
|
|
|
|
|
+ ));
|
|
|
|
|
+ }
|
|
|
|
|
+ msg.push_str("free space (docker rmi / stale staging dirs) or point dst at a bigger volume.");
|
|
|
|
|
+ Some(msg)
|
|
|
|
|
+}
|
|
|
|
|
+
|
|
|
|
|
+/// Ask every launch node how much space a block's `local` copies need.
|
|
|
|
|
+pub async fn measure(
|
|
|
|
|
+ config: &Config,
|
|
|
|
|
+ nodes: &[Node],
|
|
|
|
|
+ block: &Staging,
|
|
|
|
|
+ index: usize,
|
|
|
|
|
+ model: &str,
|
|
|
|
|
+ shared_source: &Path,
|
|
|
|
|
+) -> Result<Vec<Need>> {
|
|
|
|
|
+ let id = block.ledger_id(index);
|
|
|
|
|
+ let repo = repo_for(block, model);
|
|
|
|
|
+ let spec = ViewSpec {
|
|
|
|
|
+ repo: &repo,
|
|
|
|
|
+ local: &block.local,
|
|
|
|
|
+ verify: block.verify.as_str(),
|
|
|
|
|
+ min_free: block.min_free.as_deref(),
|
|
|
|
|
+ force: false,
|
|
|
|
|
+ };
|
|
|
|
|
+ let mut out = Vec::new();
|
|
|
|
|
+ for node in nodes {
|
|
|
|
|
+ let dst = dst_for(block, &id, &config.ssh_user);
|
|
|
|
|
+ upload_script(config, node, SCRIPT_BODY, &script_path(&id), false).await?;
|
|
|
|
|
+ let cmd = command_for(
|
|
|
|
|
+ &id,
|
|
|
|
|
+ "measure",
|
|
|
|
|
+ &shared_source.display().to_string(),
|
|
|
|
|
+ &dst,
|
|
|
|
|
+ &spec,
|
|
|
|
|
+ );
|
|
|
|
|
+ let stdout = crate::ssh::run_ssh(config, node, cmd, false).await?;
|
|
|
|
|
+ out.push(parse_measure(&node.name, &repo, &dst, &stdout)?);
|
|
|
|
|
+ }
|
|
|
|
|
+ Ok(out)
|
|
|
|
|
+}
|
|
|
|
|
+
|
|
|
|
|
+/// Everything one node's `local-view` pass needs.
|
|
|
|
|
+pub struct StageCtx<'a> {
|
|
|
|
|
+ pub config: &'a Config,
|
|
|
|
|
+ pub recipe: &'a crate::config::Recipe,
|
|
|
|
|
+ pub shared_source: &'a Path,
|
|
|
|
|
+ pub force: bool,
|
|
|
|
|
+ pub dry_run: bool,
|
|
|
|
|
+}
|
|
|
|
|
+
|
|
|
|
|
+/// Stage one `local-view` block on one node.
|
|
|
|
|
+pub async fn stage_node(
|
|
|
|
|
+ ctx: &StageCtx<'_>,
|
|
|
|
|
+ node: &Node,
|
|
|
|
|
+ block: &Staging,
|
|
|
|
|
+ index: usize,
|
|
|
|
|
+) -> Result<()> {
|
|
|
|
|
+ let (config, model, shared_source, force, dry_run) = (
|
|
|
|
|
+ ctx.config,
|
|
|
|
|
+ ctx.recipe.model.as_str(),
|
|
|
|
|
+ ctx.shared_source,
|
|
|
|
|
+ ctx.force,
|
|
|
|
|
+ ctx.dry_run,
|
|
|
|
|
+ );
|
|
|
|
|
+ if block.local.is_empty() {
|
|
|
|
|
+ bail!(
|
|
|
|
|
+ "[[staging]] local-view block '{}' needs at least one `local` glob (the files that must be real node-local copies)",
|
|
|
|
|
+ block.ledger_id(index)
|
|
|
|
|
+ );
|
|
|
|
|
+ }
|
|
|
|
|
+ let id = block.ledger_id(index);
|
|
|
|
|
+ let repo = repo_for(block, model);
|
|
|
|
|
+ let dst = dst_for(block, &id, &config.ssh_user);
|
|
|
|
|
+ let spec = ViewSpec {
|
|
|
|
|
+ repo: &repo,
|
|
|
|
|
+ local: &block.local,
|
|
|
|
|
+ verify: block.verify.as_str(),
|
|
|
|
|
+ min_free: block.min_free.as_deref(),
|
|
|
|
|
+ force,
|
|
|
|
|
+ };
|
|
|
|
|
+ let src = shared_source.display().to_string();
|
|
|
|
|
+ if dry_run {
|
|
|
|
|
+ info!(
|
|
|
|
|
+ "dry-run staging[{}] on {}: local-view {} → {} (local: {:?}, min_free {:?})",
|
|
|
|
|
+ id, node.name, repo, dst, block.local, block.min_free
|
|
|
|
|
+ );
|
|
|
|
|
+ return Ok(());
|
|
|
|
|
+ }
|
|
|
|
|
+ upload_script(config, node, SCRIPT_BODY, &script_path(&id), false).await?;
|
|
|
|
|
+ let cmd = command_for(&id, "apply", &src, &dst, &spec);
|
|
|
|
|
+ run_ssh_long_running(config, node, cmd, false).await?;
|
|
|
|
|
+ Ok(())
|
|
|
|
|
+}
|
|
|
|
|
+
|
|
|
|
|
+#[cfg(test)]
|
|
|
|
|
+mod tests {
|
|
|
|
|
+ use super::*;
|
|
|
|
|
+
|
|
|
|
|
+ fn need(node: &str, need: u64, free: u64) -> Need {
|
|
|
|
|
+ Need {
|
|
|
|
|
+ node: node.into(),
|
|
|
|
|
+ repo: "org/model".into(),
|
|
|
|
|
+ need,
|
|
|
|
|
+ free,
|
|
|
|
|
+ dst: "/v".into(),
|
|
|
|
|
+ }
|
|
|
|
|
+ }
|
|
|
|
|
+
|
|
|
|
|
+ #[test]
|
|
|
|
|
+ fn measure_output_parses() {
|
|
|
|
|
+ let n = parse_measure(
|
|
|
|
|
+ "dgx5",
|
|
|
|
|
+ "org/m",
|
|
|
|
|
+ "/v",
|
|
|
|
|
+ r#"{"repo":"org/m","need":203,"free":400}"#,
|
|
|
|
|
+ )
|
|
|
|
|
+ .unwrap();
|
|
|
|
|
+ assert_eq!(n.need, 203);
|
|
|
|
|
+ assert_eq!(n.free, 400);
|
|
|
|
|
+ assert!(!n.short());
|
|
|
|
|
+ }
|
|
|
|
|
+
|
|
|
|
|
+ #[test]
|
|
|
|
|
+ fn shortfall_lists_only_the_short_nodes() {
|
|
|
|
|
+ let needs = vec![need("a", 100, 200), need("b", 300, 250)];
|
|
|
|
|
+ let msg = shortfall(&needs).expect("b is short");
|
|
|
|
|
+ assert!(msg.contains("b:"), "{msg}");
|
|
|
|
|
+ assert!(!msg.contains("a:"), "{msg}");
|
|
|
|
|
+ assert!(msg.contains("nothing was started"), "{msg}");
|
|
|
|
|
+ assert!(shortfall(&[need("a", 1, 2)]).is_none());
|
|
|
|
|
+ }
|
|
|
|
|
+
|
|
|
|
|
+ #[test]
|
|
|
|
|
+ fn command_carries_globs_and_container_root() {
|
|
|
|
|
+ let globs = vec!["model-00047-of-00048.safetensors".to_string()];
|
|
|
|
|
+ let spec = ViewSpec {
|
|
|
|
|
+ repo: "org/model",
|
|
|
|
|
+ local: &globs,
|
|
|
|
|
+ verify: "sample-hash",
|
|
|
|
|
+ min_free: Some("10G"),
|
|
|
|
|
+ force: true,
|
|
|
|
|
+ };
|
|
|
|
|
+ let cmd = command_for("s0", "apply", "/src", "/dst", &spec);
|
|
|
|
|
+ // shq() shell-quotes arguments, so the glob may come back quoted
|
|
|
|
|
+ assert!(cmd.contains("model-00047-of-00048.safetensors"), "{cmd}");
|
|
|
|
|
+ assert!(cmd.contains("--local"), "{cmd}");
|
|
|
|
|
+ assert!(cmd.contains(CONTAINER_ROOT), "{cmd}");
|
|
|
|
|
+ assert!(cmd.contains("--min-free '10G'"), "{cmd}");
|
|
|
|
|
+ assert!(cmd.ends_with("--force"), "{cmd}");
|
|
|
|
|
+ }
|
|
|
|
|
+}
|