feat: auto-sync Docker image from head node to workers before serve

src/docker/mod.rs

@@ -90,6 +90,10 @@ pub async fn serve(
     if image != recipe.docker_image {
         recipe.docker_image = image.clone();
     }
+    // Sync Docker image from head node to workers if needed.
+    if !opts.dry_run && cluster.nodes.len() > 1 {
+        sync_image_to_nodes(config, &cluster, &recipe.docker_image).await?;
+    }
     let cache_dir = config.resolved_hf_cache_dir();
     start_cluster(
         config,
@@ -406,5 +410,77 @@ fn build_docker_command(ctx: &DockerCommandCtx, node: &Node, is_head: bool) -> V
     }
     cmd
 }
+
+/// Get the Docker image digest (short ID) from a node.
+async fn get_image_digest(
+    config: &Config,
+    node: &Node,
+    image: &str,
+) -> Result<Option<String>> {
+    let cmd = format!(
+        "docker image ls --format '{{{{.ID}}}}' {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))
+    }
+}
+
+/// Sync Docker image from the head node (rank 0) to all worker nodes.
+/// Skips workers that already have the same image digest.
+pub async fn sync_image_to_nodes(
+    config: &Config,
+    cluster: &Cluster,
+    image: &str,
+) -> Result<()> {
+    let head = cluster.nodes.iter().find(|n| n.rank == 0);
+    let head = match head {
+        Some(h) => h,
+        None => return Ok(()),  // single-node, nothing to sync
+    };
+    let head_digest = get_image_digest(config, head, image).await?;
+    if head_digest.is_none() {
+        anyhow::bail!(
+            "Image '{}' not found on head node '{}'",
+            image,
+            head.name
+        );
+    }
+    let head_digest = head_digest.unwrap();
+    info!("Syncing image '{}' from '{}' to workers...", image, head.name);
+    for worker in cluster.nodes.iter().filter(|n| n.rank > 0) {
+        let worker_digest = get_image_digest(config, worker, image).await?;
+        if worker_digest.as_deref() == Some(&head_digest) {
+            info!("  '{}' OK (digest matches)", worker.name);
+            continue;
+        }
+        info!("  {} -> {}...", worker.name, worker.ssh_ip);
+        // Run docker save | ssh docker load from the head node to the worker.
+        let pipe_cmd = format!(
+            "docker save '{}' | ssh -o StrictHostKeyChecking=no -i {} {}@{} docker load",
+            image,
+            config.ssh_key_path.display(),
+            config.ssh_user,
+            worker.ssh_ip
+        );
+        run_ssh(config, head, pipe_cmd, false).await?;
+        // Verify the image landed with the right digest.
+        let verify = get_image_digest(config, worker, image).await?;
+        if verify.as_deref() != Some(&head_digest) {
+            anyhow::bail!(
+                "Image sync to '{}' failed: digest mismatch (head={}, worker={:?})",
+                worker.name,
+                head_digest,
+                verify
+            );
+        }
+        info!("  {} OK", worker.name);
+    }
+    Ok(())
+}
+
 #[cfg(test)]
 mod tests;