Bläddra i källkod

sync: one docker save + parallel fabric fan-out (3x on a 4-node set), inner hop key = node-side ~/.ssh/<basename> (mac path broke dgx5-8), known_hosts lives until task end (10G fallback resurrected); sm121 sync-image command

Alexander Korolev 3 veckor sedan
förälder
incheckning
6b16590b1f
3 ändrade filer med 325 tillägg och 130 borttagningar
  1. 17 0
      src/cli.rs
  2. 303 130
      src/docker/sync.rs
  3. 5 0
      src/main.rs

+ 17 - 0
src/cli.rs

@@ -24,6 +24,10 @@ pub struct Cli {
 
 #[derive(Subcommand)]
 pub enum Commands {
+    /// Distribute a recipe's image across the cluster nodes (node-to-node,
+    /// one save + parallel fan-out; no serve started).
+    SyncImage(SyncImageArgs),
+
     /// Launch a cluster
     Serve {
         #[arg(long, required = true)]
@@ -471,3 +475,16 @@ pub enum TmuxCommands {
 #[cfg(test)]
 #[path = "cli_tests.rs"]
 mod tests;
+
+#[derive(Args)]
+pub struct SyncImageArgs {
+    #[arg(long, required = true)]
+    pub recipe: String,
+
+    #[arg(long, default_value = "auto")]
+    pub cluster: String,
+
+    /// Explicit source node (default: newest image present in the set).
+    #[arg(long)]
+    pub source: Option<String>,
+}

+ 303 - 130
src/docker/sync.rs

@@ -26,16 +26,55 @@ async fn get_image_digest(config: &Config, node: &Node, image: &str) -> Result<O
     }
 }
 
-/// Get the Docker image creation timestamp.
-async fn get_image_created(config: &Config, node: &Node, image: &str) -> Option<String> {
-    let cmd = format!(
-        "docker image inspect --format '{{{{.Created}}}}' {} 2>/dev/null",
-        shq(image)
-    );
-    run_ssh(config, node, cmd, false)
+/// Probe one node for the image: digest + creation timestamp (the latter
+/// only when present — two serial round-trips inside, but nodes are
+/// probed in parallel by [`probe_all`]).
+async fn probe_node(
+    config: &Config,
+    node: &Node,
+    image: &str,
+) -> Result<(Node, Option<String>, Option<String>)> {
+    let node = node.clone();
+    let digest = get_image_digest(config, &node, image).await?;
+    let created = if digest.is_some() {
+        run_ssh(
+            config,
+            &node,
+            format!(
+                "docker image inspect --format '{{{{.Created}}}}' {} 2>/dev/null",
+                shq(image)
+            ),
+            false,
+        )
         .await
         .ok()
         .map(|c| c.trim().to_string())
+    } else {
+        None
+    };
+    Ok((node, digest, created))
+}
+
+/// Probe every cluster node for the image digest/creation time in parallel.
+async fn probe_all(
+    config: &Config,
+    cluster: &Cluster,
+    image: &str,
+) -> Result<Vec<(Node, Option<String>, Option<String>)>> {
+    let mut handles = Vec::new();
+    for node in &cluster.nodes {
+        let config = config.clone();
+        let node = (*node).clone();
+        let image = image.to_string();
+        handles.push(crate::locale::spawn_with_locale(async move {
+            probe_node(&config, &node, &image).await
+        }));
+    }
+    let mut probed = Vec::new();
+    for h in futures::future::join_all(handles).await {
+        probed.push(h??); // join error first, then the probe's own ssh error
+    }
+    Ok(probed)
 }
 
 /// Ensure a prebuilt (externally built) image is present on every node of
@@ -55,67 +94,68 @@ pub async fn ensure_prebuilt_image(config: &Config, cluster: &Cluster, image: &s
     })
 }
 
-/// Sync Docker image across cluster nodes.
-/// Finds the node with the *most recently built* image (by creation timestamp),
-/// then pipes `docker save | ssh docker load` to nodes missing or with a different digest.
+/// Sync a Docker image across cluster nodes — one `docker save` on the
+/// source, then a PARALLEL fan-out of `ssh <target> docker load < tar` to
+/// every target that is missing or holds a different digest.
+///
+/// Why not one save per target (the old shape): with N peers the source
+/// re-serialized the whole image N times and the transfers ran strictly
+/// serially — N × (save + load) wall time for one image. The fabric
+/// carries multiple streams at full rate, so one save + N parallel loads
+/// costs ≈ save + max(load), roughly N× faster for a 4-node set.
 ///
-/// The inner `ssh docker load` hop is **host-key verified**: the target's full
-/// public key is fetched over a russh connection that already enforces the
-/// pinned fingerprint, written to a per-target `known_hosts` on the source
-/// node, and passed to the inner `ssh` with `StrictHostKeyChecking=yes`. So the
-/// operator key is only ever presented to a host that proves it holds the
-/// pinned key — an impersonator at `target.ssh_ip` is rejected before auth and
-/// cannot harvest the key or inject a malicious `docker load`. The temporary
-/// `known_hosts` is removed after the pipe completes.
+/// Data path: node-to-node, never through the Mac (mac→node is ~32 MB/s vs
+/// ~547 MB/s per stream on the 200G fabric). The inner hop prefers the
+/// target's fabric IP (ib_ip); if that fails it falls back to the 10G mgmt
+/// IP with a warning (the per-target known_hosts covers both hops and lives
+/// until the target task ends — the old per-pipe `rm` deleted it before the
+/// fallback attempt).
+///
+/// Security of the inner hop: the target's full public key is fetched over
+/// a russh connection that enforces the pinned fingerprint, written to a
+/// per-target `known_hosts` on the source node, and passed to the inner
+/// `ssh` with `StrictHostKeyChecking=yes` — an impersonator at the target
+/// IP is rejected before auth. The file is removed after the target task.
 pub async fn sync_image_to_nodes(
     config: &Config,
     cluster: &Cluster,
     image: &str,
     explicit_source: Option<&str>,
 ) -> Result<()> {
-    // Collect digests and creation timestamps per node.
-    let mut digests: Vec<(&Node, Option<String>)> = Vec::new();
-    let mut created: Vec<Option<String>> = Vec::new();
-    for node in &cluster.nodes {
-        let d = get_image_digest(config, node, image).await?;
-        let c = if d.is_some() {
-            get_image_created(config, node, image).await
-        } else {
-            None
-        };
-        digests.push((node, d));
-        created.push(c);
-    }
+    let probed = probe_all(config, cluster, image).await?;
 
     let source_idx = if let Some(src_name) = explicit_source {
-        digests
+        probed
             .iter()
-            .position(|(n, _)| n.name == src_name)
+            .position(|(n, _, _)| n.name == src_name)
             .ok_or_else(|| {
                 anyhow::anyhow!("sync-image-source '{}' not found in cluster", src_name)
             })?
     } else {
-        let mut indexed: Vec<usize> = digests
+        let mut indexed: Vec<usize> = probed
             .iter()
             .enumerate()
-            .filter(|(_, (_, d))| d.is_some())
+            .filter(|(_, (_, d, _))| d.is_some())
             .map(|(i, _)| i)
             .collect();
         indexed.sort_by(|&a, &b| {
-            created[b]
+            probed[b]
+                .2
                 .as_deref()
                 .unwrap_or("")
-                .cmp(created[a].as_deref().unwrap_or(""))
+                .cmp(probed[a].2.as_deref().unwrap_or(""))
         });
         *indexed
             .first()
             .ok_or_else(|| anyhow::anyhow!("Image '{}' not found on any node", image))?
     };
 
-    let (src_node, src_digest_opt) = &digests[source_idx];
-    let source_digest = src_digest_opt.as_ref().unwrap();
-    // Explain WHY this node won the source selection: an explicit source
-    // (a just-built image or --sync-image-source) is authoritative, whereas
+    let (src_node, src_digest_opt, _) = &probed[source_idx];
+    let source_digest = src_digest_opt
+        .as_ref()
+        .ok_or_else(|| anyhow::anyhow!("source node '{}' lost the image", src_node.name))?;
+    // WHY this node won the source selection: an explicit source (a
+    // just-built image or --sync-image-source) is authoritative, whereas
     // the heuristic source is simply the node with the newest creation
     // timestamp — a peer could differ and this is the tie-break that was
     // overwriting fresh builds.
@@ -125,113 +165,246 @@ pub async fn sync_image_to_nodes(
         "newest image (creation timestamp tie-break)"
     };
     info!(
-        "Syncing image '{}' FROM '{}' ({}) to peers...",
+        "Syncing image '{}' FROM '{}' ({}) ...",
         image, src_node.name, reason
     );
-    for (target, target_digest) in &digests {
-        if target.name == src_node.name {
+
+    // Targets: peers with no image or a different digest.
+    let mut targets: Vec<Node> = Vec::new();
+    for (node, digest, _) in &probed {
+        if node.name == src_node.name {
             continue;
         }
-        if target_digest.as_deref() == Some(source_digest) {
+        if digest.as_deref() == Some(source_digest.as_str()) {
             info!(
                 "{}",
-                fluent_i18n::t!("sync-digest-match", { "name" => &target.name })
+                fluent_i18n::t!("sync-digest-match", { "name" => &node.name })
             );
             continue;
         }
-        info!(
-            "{}",
-            fluent_i18n::t!("sync-syncing", { "name" => &target.name, "ip" => &target.ssh_ip })
-        );
-        // Data path: node-to-node, never through the Mac (mac→node is ~32
-        // MB/s vs ~547 MB/s per stream on the 200G fabric). The inner hop
-        // prefers the target's fabric IP (ib_ip); if that fails (broken
-        // node-to-node ssh over 200.x) it falls back to the 10G mgmt IP with
-        // a warning, so a stale authorized_keys never blocks a serve.
-        let attempts: Vec<String> = match &target.ib_ip {
-            Some(ib) => vec![ib.clone(), target.ssh_ip.clone()],
-            None => vec![target.ssh_ip.clone()],
-        };
-        // Fetch the target's full public key over a verified russh connection
-        // (which enforces the pinned fingerprint), then write a per-target
-        // known_hosts on the SOURCE node (one line per IP we may hop to) and
-        // run the inner ssh with it, StrictHostKeyChecking=yes. A host that
-        // doesn't hold the pinned key is rejected before any auth — no key
-        // leak, no forgery.
-        let pubkey = crate::ssh::fetch_host_pubkey(config, target).await?;
-        let known_hosts = format!("/tmp/sm121_known_hosts_{}", target.name);
-        let mut kh_ips = vec![target.ssh_ip.clone()];
-        if let Some(ib) = &target.ib_ip {
-            kh_ips.push(ib.clone());
-        }
-        let entry = kh_ips
-            .iter()
-            .map(|ip| format!("{} {}", ip, pubkey))
-            .collect::<Vec<_>>()
-            .join("\n")
-            + "\n";
-        let write_known = format!(
-            "printf '%s' {} > {} && chmod 600 {}",
-            shq(&entry),
-            shq(&known_hosts),
-            shq(&known_hosts),
-        );
-        run_ssh(config, src_node, write_known, false).await?;
+        targets.push(node.clone());
+    }
+    if targets.is_empty() {
+        return Ok(());
+    }
 
-        let pipe_for = |ip: &str| {
-            format!(
-                "docker save {} | ssh -o StrictHostKeyChecking=yes -o HashKnownHosts=no -o UserKnownHostsFile={} -i {} {}@{} docker load; ec=$?; rm -f {}; exit $ec",
-                shq(image),
+    // ONE save on the source: /var/tmp on the root NVMe (docker save writes
+    // owned by the invoking user), filename-safe image name.
+    let tar = format!("/var/tmp/sm121_sync_{}.tar", safe_name(image));
+    info!(
+        "Saving '{}' once on '{}' ({} peer(s) in parallel) ...",
+        image,
+        src_node.name,
+        targets.len()
+    );
+    let save_cmd = format!(
+        "rm -f {tar} && docker save -o {tar} {image} || {{ rm -f {tar}; exit 1; }}",
+        tar = shq(&tar),
+        image = shq(image),
+    );
+    if let Err(e) = run_ssh_long_running(config, src_node, save_cmd, false).await {
+        run_ssh(config, src_node, format!("rm -f {}", shq(&tar)), false)
+            .await
+            .ok();
+        anyhow::bail!("docker save on '{}' failed: {e}", src_node.name);
+    }
+
+    // Parallel fan-out: one task per target (fabric IP first, mgmt IP
+    // fallback, digest verification, known_hosts cleanup inside the task).
+    let mut handles = Vec::new();
+    for target in &targets {
+        let config = config.clone();
+        let src_node = src_node.clone();
+        let target = target.clone();
+        let image = image.to_string();
+        let tar = tar.clone();
+        let expected = source_digest.clone();
+        let key = node_key_path(&config);
+        let user = config.ssh_user.clone();
+        handles.push(crate::locale::spawn_with_locale(async move {
+            info!(
+                "{}",
+                fluent_i18n::t!("sync-syncing", { "name" => &target.name, "ip" => &target.ssh_ip })
+            );
+            let name = target.name.clone();
+            let pubkey = match crate::ssh::fetch_host_pubkey(&config, &target).await {
+                Ok(k) => k,
+                Err(e) => return (name, Some(format!("host key fetch: {e}"))),
+            };
+            let known_hosts = format!("/tmp/sm121_known_hosts_{}", target.name);
+            let mut kh_ips = vec![target.ssh_ip.clone()];
+            if let Some(ib) = &target.ib_ip {
+                kh_ips.push(ib.clone());
+            }
+            let entry = kh_ips
+                .iter()
+                .map(|ip| format!("{} {}", ip, pubkey))
+                .collect::<Vec<_>>()
+                .join("\n")
+                + "\n";
+            let write_known = format!(
+                "printf '%s' {} > {} && chmod 600 {}",
+                shq(&entry),
                 shq(&known_hosts),
-                shq(&config.ssh_key_path.display().to_string()),
-                shq(&config.ssh_user),
-                shq(ip),
                 shq(&known_hosts),
-            )
-        };
-        let mut synced = false;
-        let mut last_err: Option<String> = None;
-        for (idx, ip) in attempts.iter().enumerate() {
-            match run_ssh_long_running(config, src_node, pipe_for(ip), false).await {
-                Ok(_) => {
-                    let verify = get_image_digest(config, target, image).await?;
-                    if verify.as_deref() == Some(source_digest) {
-                        info!(
-                            "{}",
-                            fluent_i18n::t!("sync-synced", { "name" => &target.name })
-                        );
-                        synced = true;
-                    } else {
-                        // a digest mismatch is not a connectivity problem —
-                        // the other IP won't help, fail outright
-                        last_err = Some(format!(
-                            "digest mismatch (source={}, target={:?})",
-                            source_digest, verify
-                        ));
-                        break;
+            );
+            if let Err(e) = run_ssh(&config, &src_node, write_known, false).await {
+                return (name, Some(format!("known_hosts write: {e}")));
+            }
+            let attempts: Vec<String> = match &target.ib_ip {
+                Some(ib) => vec![ib.clone(), target.ssh_ip.clone()],
+                None => vec![target.ssh_ip.clone()],
+            };
+            let mut last_err: Option<String> = None;
+            let mut synced = false;
+            for (idx, ip) in attempts.iter().enumerate() {
+                let pipe = format!(
+                    "ssh -o StrictHostKeyChecking=yes -o HashKnownHosts=no -o UserKnownHostsFile={kh} -i {key} {user}@{ip} docker load < {tar}; exit $?",
+                    kh = shq(&known_hosts),
+                    key = key,
+                    user = shq(&user),
+                    ip = shq(ip),
+                    tar = shq(&tar),
+                );
+                match run_ssh_long_running(&config, &src_node, pipe, false).await {
+                    Ok(_) => {
+                        let verify = get_image_digest(&config, &target, &image).await;
+                        match verify {
+                            Ok(v) if v.as_deref() == Some(expected.as_str()) => {
+                                synced = true;
+                                break;
+                            }
+                            Ok(v) => {
+                                // a digest mismatch is not a connectivity
+                                // problem — the other IP won't help, bail out
+                                last_err = Some(format!(
+                                    "digest mismatch (source={expected}, target={v:?})"
+                                ));
+                                break;
+                            }
+                            Err(e) => last_err = Some(format!("digest verify: {e}")),
+                        }
                     }
-                }
-                Err(e) => {
-                    if idx == 0 {
-                        tracing::warn!(
-                            "{}",
-                            fluent_i18n::t!("sync-fabric-fallback", {
-                                "name" => &target.name,
-                                "err" => e.to_string(),
-                            })
-                        );
+                    Err(e) => {
+                        if idx == 0 {
+                            tracing::warn!(
+                                "{}",
+                                fluent_i18n::t!("sync-fabric-fallback", {
+                                    "name" => &target.name,
+                                    "err" => e.to_string(),
+                                })
+                            );
+                        }
+                        last_err = Some(e.to_string());
                     }
-                    last_err = Some(e.to_string());
                 }
             }
+            // known_hosts matters only for this target — keep it until the
+            // last attempt has run, then drop it on the source node.
+            run_ssh(
+                &config,
+                &src_node,
+                format!("rm -f {}", shq(&known_hosts)),
+                false,
+            )
+            .await
+            .ok();
+            if synced {
+                info!(
+                    "{}",
+                    fluent_i18n::t!("sync-synced", { "name" => &target.name })
+                );
+                (name, None)
+            } else {
+                (name, Some(last_err.unwrap_or_default()))
+            }
+        }));
+    }
+
+    let mut failed: Vec<String> = Vec::new();
+    for h in futures::future::join_all(handles).await {
+        match h {
+            Ok((_, None)) => {}
+            Ok((name, Some(e))) => failed.push(format!("{name}: {e}")),
+            Err(e) => failed.push(format!("task: {e}")),
         }
-        if !synced {
-            anyhow::bail!(
-                "Image sync to '{}' failed: {}",
-                target.name,
-                last_err.unwrap_or_default()
+    }
+    run_ssh(config, src_node, format!("rm -f {}", shq(&tar)), false)
+        .await
+        .ok();
+    if failed.is_empty() {
+        Ok(())
+    } else {
+        anyhow::bail!("Image sync failed on: {}", failed.join("; "))
+    }
+}
+
+/// Identity path for an INNER ssh hop that runs ON a node. The configured
+/// key path belongs to the Mac (`/Users/kilork/...`), which nodes do not
+/// have — nodes hold the shared key at `~/.ssh/<basename>` (dgx5-8 lack the
+/// `~/.ssh/config` IdentityFile fallback dgx1-4 have, so there the mac path
+/// failed with exit 255 / "Permission denied", 2026-09-29). The basename is
+/// what the inner hop must reference; the shell expands `~` itself.
+fn node_key_path(config: &Config) -> String {
+    config
+        .ssh_key_path
+        .file_name()
+        .map(|f| format!("~/.ssh/{}", f.to_string_lossy()))
+        .filter(|f| f.len() > "~/.ssh/".len())
+        .unwrap_or_else(|| config.ssh_key_path.display().to_string())
+}
+
+/// Filename-safe image name for the one-shot save tarball on the source.
+fn safe_name(image: &str) -> String {
+    image
+        .chars()
+        .map(|c| if c.is_ascii_alphanumeric() { c } else { '_' })
+        .collect()
+}
+
+/// `sm121 sync-image` handler: distribute a recipe's image across the
+/// cluster without starting a serve (pre-warm before bench/serve).
+pub async fn handle_sync_image(
+    config: &Config,
+    recipe_name: &str,
+    cluster_name: &str,
+    source: Option<&str>,
+    dry_run: bool,
+) -> Result<()> {
+    let recipe = crate::config::load_recipe(recipe_name, &config.recipe_sources).await?;
+    let cluster = crate::config::select_cluster_by_name(config, cluster_name)?;
+    let image = recipe.docker_image.clone();
+    if dry_run {
+        // No SSH in dry-run: print the plan only.
+        println!(
+            "sync-image: {} over {} node(s) of '{}'",
+            image,
+            cluster.nodes.len(),
+            cluster.name
+        );
+        println!(
+            "sync-image: ssh <source>: rm -f /var/tmp/sm121_sync_<image>.tar && docker save -o /var/tmp/sm121_sync_<image>.tar {image}"
+        );
+        for n in &cluster.nodes {
+            let hop = n.ib_ip.as_deref().unwrap_or(&n.ssh_ip);
+            println!(
+                "sync-image:   {} <- ssh {}@{} docker load < /var/tmp/sm121_sync_<image>.tar (parallel)",
+                n.name, config.ssh_user, hop
             );
         }
+        return Ok(());
+    }
+    sync_image_to_nodes(config, &cluster, &image, source).await
+}
+
+#[cfg(test)]
+mod tests {
+    use super::safe_name;
+
+    #[test]
+    fn safe_image_name_is_filename_safe() {
+        let safe = safe_name("eugr/spark-vllm-b12x:latest@sha256:abc123");
+        assert!(!safe.contains('/') && !safe.contains(':') && !safe.contains('@'));
+        assert!(!safe.is_empty());
     }
-    Ok(())
 }

+ 5 - 0
src/main.rs

@@ -1,5 +1,6 @@
 #![deny(clippy::wildcard_imports)]
 
+use crate::docker::sync;
 use anyhow::Result;
 use clap::Parser;
 
@@ -128,6 +129,10 @@ async fn main() -> Result<()> {
             )
             .await?;
         }
+        Commands::SyncImage(a) => {
+            sync::handle_sync_image(&config, &a.recipe, &a.cluster, a.source.as_deref(), dry_run)
+                .await?;
+        }
         Commands::Build {
             recipe,
             build_node,