Sfoglia il codice sorgente

feat(docker): authoritative build/sync source, single-node build clusters, chunked upload

- build: a build never requires NCCL config (it is a single-node docker build),
  so it can run on a cluster that has no NCCL (e.g. 1node-3); add
  resolve_build_cluster()
- build: handle_build returns whether the image was actually (re)built; serve
  treats that node as the authoritative image-sync source so a peer's older
  image (creation-timestamp heuristic) cannot overwrite a fresh build
- ssh upload_dir: stream the tar across SSH in chunks instead of one inline
  base64 argv, fixing 'Argument list too long' for large from-source build
  dirs (hash lockfiles ~250KB)
- docker command: do not double shell-quote recipe args that are already
  single-quoted
Alexander Korolev 2 mesi fa
parent
commit
e300668348
7 ha cambiato i file con 90 aggiunte e 17 eliminazioni
  1. 24 0
      src/config.rs
  2. 8 3
      src/docker/build.rs
  3. 7 1
      src/docker/command.rs
  4. 17 8
      src/docker/mod.rs
  5. 12 2
      src/docker/sync.rs
  6. 1 1
      src/main.rs
  7. 21 2
      src/ssh.rs

+ 24 - 0
src/config.rs

@@ -120,6 +120,30 @@ pub fn select_cluster_by_name(config: &Config, cluster_name: &str) -> Result<Clu
     Ok(resolve_cluster_name(&config.clusters, cluster_name)?.clone())
 }
 
+/// Resolve a cluster for a **build**.
+///
+/// A build runs the Docker image build on a single node, so NCCL config is
+/// never required here, even when the recipe's `required_nodes > 1` (the
+/// resulting image is later pushed/synced to the serving cluster). For
+/// `cluster_name != "auto"` the named cluster is used as-is; for `"auto"` the
+/// smallest cluster containing `build_node` is selected.
+pub fn resolve_build_cluster(
+    config: &Config,
+    cluster_name: &str,
+    build_node: &str,
+) -> Result<Cluster> {
+    if cluster_name != "auto" {
+        return select_cluster_by_name(config, cluster_name);
+    }
+    config
+        .clusters
+        .iter()
+        .filter(|c| c.nodes.iter().any(|n| n.name == build_node))
+        .min_by_key(|c| c.nodes.len())
+        .cloned()
+        .ok_or_else(|| anyhow::anyhow!("no cluster contains build node '{}'", build_node))
+}
+
 pub fn resolve_cluster<'c>(
     config: &'c Config,
     recipe: &Recipe,

+ 8 - 3
src/docker/build.rs

@@ -292,6 +292,11 @@ pub async fn build_image_on_node(
 }
 
 /// Decide whether to build and dispatch the build for the serve flow.
+///
+/// Returns `true` when an image was actually (re)built on the node, `false` when
+/// the build was skipped (checksum match) or it was a dry-run. Callers use the
+/// return value to treat the freshly built image as the authoritative sync
+/// source instead of re-running the "most recently created" heuristic.
 pub async fn handle_build(
     config: &Config,
     cluster: &Cluster,
@@ -299,7 +304,7 @@ pub async fn handle_build(
     recipe: &Recipe,
     opts: &ServeOpts,
     build_node_name: &str,
-) -> Result<()> {
+) -> Result<bool> {
     if opts.dry_run {
         let dir = recipe.docker_image_dir.as_deref().unwrap_or(recipe_name);
         let overrides = if recipe.build.versions.is_empty() {
@@ -314,7 +319,7 @@ pub async fn handle_build(
             "[dry-run] would build '{}' on '{}' from docker-images/{}/build.sh{}",
             recipe.docker_image, build_node_name, dir, overrides
         );
-        return Ok(());
+        return Ok(false);
     }
     let image_dir = recipe
         .docker_image_dir
@@ -390,7 +395,7 @@ pub async fn handle_build(
         )
         .await?;
     }
-    Ok(())
+    Ok(should_build)
 }
 
 /// Format a `Duration` as a compact human string: `45s`, `1m 30s`, `12m 04s`.

+ 7 - 1
src/docker/command.rs

@@ -315,7 +315,13 @@ pub(crate) fn build_docker_command(
         cmd.push(img);
         // Shell-quote every tail word (entrypoint, model, args, flags) so a
         // crafted recipe arg/env cannot break out of the remote `docker run`.
-        cmd.extend(tail.iter().map(|t| shq(t)));
+        cmd.extend(tail.iter().map(|t| {
+            if t.starts_with('\'') {
+                t.clone()
+            } else {
+                shq(t)
+            }
+        }));
     } else {
         let script = mods::build_mod_script(&ctx.recipe.mods.commands, &tail);
         // `--entrypoint sh` and `-c` stay literal; `img` and `script` are

+ 17 - 8
src/docker/mod.rs

@@ -128,12 +128,17 @@ pub async fn serve(
     use self::sync::sync_image_to_nodes;
 
     // Build image on specified node if needed.
+    // `built_on` records the node that actually (re)built the image. After a
+    // fresh build that image is authoritative — the sync step below must pull
+    // FROM it, never let a peer's older image (by creation timestamp) overwrite
+    // the new build.
+    let mut built_on: Option<String> = None;
     if let Some(ref build_node_name) = opts.build_node {
         if !opts.dry_run && !cluster.nodes.iter().any(|n| n.name == *build_node_name) {
             anyhow::bail!("build-node '{}' not in cluster", build_node_name);
         }
         if !opts.dry_run {
-            crate::docker::build::handle_build(
+            let did_build = crate::docker::build::handle_build(
                 config,
                 &cluster,
                 recipe_name,
@@ -142,6 +147,9 @@ pub async fn serve(
                 build_node_name,
             )
             .await?;
+            if did_build {
+                built_on = Some(build_node_name.clone());
+            }
         }
     }
 
@@ -168,13 +176,14 @@ pub async fn serve(
     let should_sync =
         opts.sync_image && recipe.sync_image && !opts.dry_run && cluster.nodes.len() > 1;
     if should_sync {
-        sync_image_to_nodes(
-            config,
-            &cluster,
-            &recipe.docker_image,
-            opts.sync_image_source.as_deref(),
-        )
-        .await?;
+        // 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.
+        let sync_source = built_on.as_deref().or(opts.sync_image_source.as_deref());
+        sync_image_to_nodes(config, &cluster, &recipe.docker_image, sync_source).await?;
     }
 
     // Drop page cache on all nodes right before starting containers.

+ 12 - 2
src/docker/sync.rs

@@ -88,9 +88,19 @@ pub async fn sync_image_to_nodes(
 
     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
+    // 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 '{}' (most recent build) to peers...",
-        image, src_node.name
+        "Syncing image '{}' FROM '{}' ({}) to peers...",
+        image, src_node.name, reason
     );
     for (target, target_digest) in &digests {
         if target.name == src_node.name {

+ 1 - 1
src/main.rs

@@ -278,7 +278,7 @@ async fn handle_build_cmd(
     dry_run: bool,
 ) -> Result<()> {
     let recipe_obj = crate::config::load_recipe(recipe).await?;
-    let cluster = resolve_cluster_for_op(config, recipe, cluster).await?;
+    let cluster = crate::config::resolve_build_cluster(config, cluster, build_node)?;
     let opts = docker::ServeOpts {
         keep: false,
         skip_before_serve: true,

+ 21 - 2
src/ssh.rs

@@ -385,12 +385,31 @@ pub async fn upload_dir(
 
     // Base64 is safe inside single quotes (charset: A-Za-z0-9+/=).
     // `printf '%s'` is portable and doesn't interpret escape sequences.
+    //
+    // The full payload can exceed the platform ARG_MAX (e.g. from-source build
+    // dirs with large hash-lock files -> ~250KB base64 > the ~256KB limit the
+    // remote `zsh -c` shell enforces), so we never inline it all in one argv.
+    // Instead stream it in fixed-size chunks via `>>` appends to a temp file on
+    // the node, then base64-decode + extract in a final command.
+    let tmp = format!("/tmp/sm121-up-{}.b64", std::process::id());
+    const CHUNK: usize = 60_000;
+    let chunks: Vec<&str> = b64
+        .as_bytes()
+        .chunks(CHUNK)
+        .map(|c| std::str::from_utf8(c).expect("base64 is ASCII"))
+        .collect();
+    for (i, piece) in chunks.iter().enumerate() {
+        let op = if i == 0 { ">" } else { ">>" };
+        let cmd = format!("printf '%s' '{}' {} {}", piece, op, shq(&tmp));
+        run_ssh(config, node, cmd, false).await?;
+    }
+
     // `remote_dir` is shell-quoted — a path with spaces or metachars would
     // otherwise break `mkdir -p` / `cd`.
     let cmd = format!(
-        "mkdir -p {dir} && cd {dir} && printf '%s' '{b64}' | base64 -d | tar xzf -",
+        "mkdir -p {dir} && cd {dir} && base64 -d < {tmp} | tar xzf - && rm -f {tmp}",
         dir = shq(remote_dir),
-        b64 = b64
+        tmp = shq(&tmp)
     );
 
     run_ssh(config, node, cmd, false).await?;