launcher/src/docker/sync.rs

use crate::config::{Cluster, Config, Node};
use crate::ssh::{run_ssh, run_ssh_long_running};
use anyhow::Result;
use tracing::info;

use super::shq;

/// Get the Docker image digest (short ID) from a node.
async fn get_image_digest(config: &Config, node: &Node, image: &str) -> Result<Option<String>> {
    // `docker image ls` does not accept digest refs (`repo@sha256:...`) —
    // use `docker image inspect` for those; it also returns the image ID.
    // `|| true`: inspect exits 1 when the ref is absent (a peer that simply
    // doesn't have it yet) — that is "no digest", not an ssh failure.
    let cmd = if image.contains('@') {
        format!(
            "docker image inspect --format '{{{{.Id}}}}' {} 2>/dev/null || true",
            shq(image)
        )
    } else {
        format!("docker image ls --format '{{{{.ID}}}}' {}", shq(image))
    };
    let out = run_ssh(config, node, cmd, false).await?;
    let digest = out.trim().to_string();
    if digest.is_empty() {
        Ok(None)
    } else {
        Ok(Some(digest))
    }
}

/// 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
/// the launch set BEFORE `docker run`: sync it node-to-node from a peer that
/// already has it, or fail loudly when no node has it.
///
/// Without this check a missing image makes `docker run` silently pull it
/// from the registry over the internet; the SSH timeout then kills the pull
/// while the serve still reports "started" and waits out the readiness
/// timeout for a container that never loads.
pub async fn ensure_prebuilt_image(config: &Config, cluster: &Cluster, image: &str) -> Result<()> {
    sync_image_to_nodes(config, cluster, image, None).await.map_err(|e| {
        anyhow::anyhow!(
            "image '{}' not available on the launch nodes: {e} \u{2014} pull it on one node first (or build it with --build-node)",
            image
        )
    })
}

/// 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.
///
/// 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<()> {
    let probed = probe_all(config, cluster, image).await?;

    let source_idx = if let Some(src_name) = explicit_source {
        probed
            .iter()
            .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> = probed
            .iter()
            .enumerate()
            .filter(|(_, (_, d, _))| d.is_some())
            .map(|(i, _)| i)
            .collect();
        indexed.sort_by(|&a, &b| {
            probed[b]
                .2
                .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, _) = &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.
    let reason = if explicit_source.is_some() {
        "authoritative image"
    } else {
        "newest image (creation timestamp tie-break)"
    };
    info!(
        "Syncing image '{}' FROM '{}' ({}) ...",
        image, src_node.name, reason
    );

    // 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 digest.as_deref() == Some(source_digest.as_str()) {
            info!(
                "{}",
                fluent_i18n::t!("sync-digest-match", { "name" => &node.name })
            );
            continue;
        }
        targets.push(node.clone());
    }
    if targets.is_empty() {
        return Ok(());
    }

    // 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(&known_hosts),
            );
            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(_) => {
                        // A digest-only ref (`repo@sha256:...`) loses its
                        // RepoDigest association through save/load — the tar
                        // carries no tags, so the loaded image is untagged
                        // and `docker run <ref>` fails with "No such image"
                        // on the peer. Re-attach it with a pull-by-digest:
                        // all layers are local, this fetches nothing but the
                        // manifest (and pins the association).
                        if image.contains('@') {
                            let pull = format!("docker pull {}", shq(&image));
                            if let Err(e) = run_ssh(&config, &target, pull, false).await {
                                last_err = Some(format!("pull-by-digest re-attach: {e}"));
                                continue;
                            }
                        }
                        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(),
                                })
                            );
                        }
                        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}")),
        }
    }
    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());
    }
}