Parcourir la source

refactor: move prebuilt-image check to sync::ensure_prebuilt_image (mod.rs back under 500)

Alexander Korolev il y a 3 semaines
Parent
commit
a1472225ad
2 fichiers modifiés avec 29 ajouts et 25 suppressions
  1. 8 25
      src/docker/mod.rs
  2. 21 0
      src/docker/sync.rs

+ 8 - 25
src/docker/mod.rs

@@ -81,12 +81,8 @@ pub async fn serve(
         fluent_i18n::t!("recipe-starting", { "name" => rn, "engine" => re, "backend" => recipe.distributed_backend.map_or("none".to_string(), |b| b.to_string()) })
     );
     // `--fit`: override required_nodes + tensor_parallel_size so TP tracks
-    // the node count (the recipe's `required_nodes == Some(tp_size)` path in
-    // `start_cluster` makes tp_size follow nnodes). `pipeline_parallel_size`
-    // is left untouched. Bare `--fit` (whole cluster) is handled after
-    // cluster selection, when nnodes is known; for now just clear the cap so
-    // selection/trim keeps every cluster node.
-    //
+    // the node count. Bare `--fit` (whole cluster) is handled after cluster
+    // selection; for now clear the cap so selection keeps every node.
     // Mutually exclusive with `--solo` (solo = 1 node, no dist flags).
     let fit: Option<u32> = match opts.fit.as_deref() {
         None => None,
@@ -248,30 +244,17 @@ pub async fn serve(
         opts.sync_image && recipe.sync_image && !opts.dry_run && cluster.nodes.len() > 1;
     if should_sync {
         // After a fresh build the build node holds the authoritative image:
-        // force it as the sync source so a peer's older image (by creation
-        // timestamp, or from a no-op `docker load` round-trip) can't overwrite
-        // the new build. The "most recently created" heuristic only applies
-        // when nothing was built here and no explicit --sync-image-source was
-        // given — i.e. the nodes merely disagree and we pick the newest.
+        // force it as the sync source so a peer's older image can't overwrite
+        // the new build. The "newest creation timestamp" heuristic only
+        // applies when nothing was built here and no explicit
+        // --sync-image-source was given.
         let source = built_on.as_deref().or(opts.sync_image_source.as_deref());
         self::sync::sync_image_to_nodes(config, &cluster, &recipe.docker_image, source).await?;
     }
     // Prebuilt external images (no `docker_image_dir`) are not built by the
-    // launcher — verify the image exists on the launch nodes BEFORE any
-    // `docker run`: a missing image would make docker silently pull it from
-    // the registry over the internet, and the SSH timeout would kill the
-    // pull while the serve still reports "started" and waits out a 30-min
-    // readiness timeout for a container that never loads. Sync it node-to-
-    // node from a peer that has it; if no node has it, fail with a hint.
+    // launcher — ensure they exist on the launch nodes BEFORE any `docker run`.
     if !opts.dry_run && recipe.docker_image_dir.is_none() {
-        self::sync::sync_image_to_nodes(config, &cluster, &recipe.docker_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)",
-                    recipe.docker_image
-                )
-            })?;
+        self::sync::ensure_prebuilt_image(config, &cluster, &recipe.docker_image).await?;
     }
 
     // Drop page cache on all nodes right before starting containers.

+ 21 - 0
src/docker/sync.rs

@@ -29,6 +29,27 @@ async fn get_image_created(config: &Config, node: &Node, image: &str) -> Option<
         .map(|c| c.trim().to_string())
 }
 
+/// 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 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.