Sfoglia il codice sorgente

feat(cli): sm121 publish — recipe -> standalone reproducible bundle (and gist)

Generalizes this morning's manual gist: bundles recipe.toml + flat image
overlay (Dockerfile with patches/ paths rewritten to root, patches,
versions.env) + generated run.sh (dist flags, recipe args/env, cluster
fabric env + rank-0 IB ip when --cluster given) + gist-aware install.sh +
README with a pinned-coordinates table.

Reproducibility baked in: HF model revision resolved live from the API
(run.sh downloads + serves the pinned sha; README warns when unpinnable),
BASE_IMAGE shown as written (digest-pinned = ok, bare tag = warn).

--gist shells out to authenticated gh (gist create --public), then bakes
the returned id back into install.sh and updates the gist (gh gist edit)
so the curl-pipe one-liner self-bootstraps.

4 unit tests (flatten, versions parsing, vllm/sglang/solo render).
Alexander Korolev 1 mese fa
parent
commit
79d1a76c74
4 ha cambiato i file con 514 aggiunte e 0 eliminazioni
  1. 1 0
      .gitignore
  2. 18 0
      src/cli.rs
  3. 9 0
      src/main.rs
  4. 486 0
      src/publish.rs

+ 1 - 0
.gitignore

@@ -7,3 +7,4 @@ __pycache__/
 
 # scratch from docker-images/*/gen-locks.sh
 .gen/
+publish/

+ 18 - 0
src/cli.rs

@@ -155,6 +155,24 @@ pub enum Commands {
         repair: bool,
     },
 
+    /// Publish a recipe (+ image overlay) as a standalone reproducible bundle
+    Publish {
+        #[arg(long, required = true)]
+        recipe: String,
+
+        /// Cluster whose fabric env + node count fill the generated run script
+        #[arg(long)]
+        cluster: Option<String>,
+
+        /// Output dir (default: ./publish/<recipe>)
+        #[arg(long)]
+        out: Option<std::path::PathBuf>,
+
+        /// Also create a public GitHub gist (requires `gh` CLI, authenticated)
+        #[arg(long)]
+        gist: bool,
+    },
+
     /// Manage tmux sessions (Docker alternative)
     #[command(subcommand)]
     Tmux(TmuxCommands),

+ 9 - 0
src/main.rs

@@ -25,6 +25,7 @@ mod locale;
 mod orchestrate;
 mod power;
 mod pre_serve;
+mod publish;
 mod ssh;
 mod tmux;
 mod web;
@@ -150,6 +151,14 @@ async fn main() -> Result<()> {
         } => {
             doctor::doctor(&config, &recipe, &cluster, repair, dry_run).await?;
         }
+        Commands::Publish {
+            recipe,
+            cluster,
+            out,
+            gist,
+        } => {
+            publish::publish(&config, &recipe, cluster.as_deref(), out, gist).await?;
+        }
         Commands::Tmux(tmux_cmd) => {
             crate::tmux::handle_tmux(&config, tmux_cmd, dry_run).await?;
         }

+ 486 - 0
src/publish.rs

@@ -0,0 +1,486 @@
+//! `sm121 publish` — turn a recipe + its `docker-images/<dir>` overlay into a
+//! standalone, reproducible recipe-card bundle (README, flat Dockerfile +
+//! patches, run/install scripts, the recipe itself). `--gist` additionally
+//! publishes the bundle as a public GitHub gist via `gh`.
+//!
+//! The bundle is intentionally flat (gists have no directories): `patches/`
+//! paths in the Dockerfile are rewritten to the bundle root.
+
+use std::collections::BTreeMap;
+use std::path::{Path, PathBuf};
+use std::process::Command;
+
+use anyhow::{Context, Result, bail};
+use tracing::info;
+
+use crate::config::{Cluster, Config, NcclConfig, Recipe, recipe_dirs};
+use crate::ssh::shq;
+
+/// Publish `recipe_name` into `out` (default `./publish/<name>`), optionally
+/// creating a gist. `cluster` (optional) supplies fabric env + node count in
+/// the generated run script.
+pub async fn publish(
+    config: &Config,
+    recipe_name: &str,
+    cluster: Option<&str>,
+    out: Option<PathBuf>,
+    gist: bool,
+) -> Result<()> {
+    let recipe = crate::config::load_recipe(recipe_name, &config.recipe_sources).await?;
+    if recipe.deploy != crate::config::DeployTarget::Docker {
+        bail!("publish supports docker recipes only (deploy = tmux has no image bundle)");
+    }
+    let bundle = out.unwrap_or_else(|| Path::new("publish").join(recipe_name));
+    if bundle.exists() {
+        std::fs::remove_dir_all(&bundle)
+            .with_context(|| format!("clearing {}", bundle.display()))?;
+    }
+    std::fs::create_dir_all(&bundle)?;
+
+    let recipe_src = recipe_dirs(config)
+        .into_iter()
+        .map(|d| d.join(recipe_name).join("recipe.toml"))
+        .find(|p| p.exists())
+        .with_context(|| format!("recipe.toml for '{recipe_name}' not found in recipe dirs"))?;
+    std::fs::copy(&recipe_src, bundle.join("recipe.toml"))?;
+
+    let mut versions: BTreeMap<String, String> = BTreeMap::new();
+    let image_dir =
+        Path::new("docker-images").join(recipe.docker_image_dir.as_deref().unwrap_or(recipe_name));
+    if image_dir.join("Dockerfile").exists() {
+        let df = std::fs::read_to_string(image_dir.join("Dockerfile"))?;
+        std::fs::write(bundle.join("Dockerfile"), flatten_dockerfile(&df))?;
+        copy_flat(&image_dir, &bundle)?;
+        if let Ok(v) = std::fs::read_to_string(image_dir.join("versions.env")) {
+            std::fs::write(bundle.join("versions.env"), &v)?;
+            versions = versions_build_args(&v);
+        }
+        info!("📦 bundled image overlay {}", image_dir.display());
+    } else {
+        info!(
+            "ℹ️  no docker-images/{} overlay — bundling recipe + scripts only",
+            image_dir.display()
+        );
+    }
+
+    let cluster = match cluster {
+        Some(want) => Some(
+            config
+                .clusters
+                .iter()
+                .find(|c| c.name == want)
+                .with_context(|| format!("cluster '{want}' not in config"))?,
+        ),
+        None => None,
+    };
+    let nccl = cluster.map(|c| config.nccl.get(&c.name));
+    let revision = hf_revision(&recipe.model).await;
+    let hf_cache = config
+        .mount_aliases
+        .get("hf_cache")
+        .cloned()
+        .or_else(|| config.hf_cache_dir.clone())
+        .unwrap_or_else(|| "~/.cache/huggingface".into());
+
+    let run = render_run_script(
+        &recipe,
+        cluster,
+        nccl.flatten(),
+        &hf_cache,
+        config.init_port,
+        revision.as_deref(),
+    );
+    std::fs::write(bundle.join("run.sh"), run)?;
+    let readme = render_readme(
+        recipe_name,
+        &recipe,
+        &versions,
+        revision.as_deref(),
+        image_dir.exists(),
+    );
+    std::fs::write(bundle.join("README.md"), readme)?;
+    info!("📝 bundle written to {}", bundle.display());
+
+    if gist {
+        let id = publish_gist(&bundle, &recipe.name)?;
+        // The gist knows its own id now — bake it into install.sh and push the
+        // update, so `curl …/raw/install.sh | bash` self-bootstraps from anywhere.
+        let raw = std::fs::read_to_string(bundle.join("install.sh"))?.replacen(
+            "GIST_ID='          '",
+            &format!("GIST_ID={}", shq(&id)),
+            1,
+        );
+        std::fs::write(bundle.join("install.sh"), &raw)?;
+        match Command::new("gh")
+            .args(["gist", "edit", &id])
+            .arg(bundle.join("install.sh"))
+            .output()
+        {
+            Ok(o) if o.status.success() => info!("🔗 install.sh updated with gist id"),
+            _ => info!("ℹ️  could not auto-update install.sh — set GIST_ID={id} manually"),
+        }
+    } else {
+        std::fs::write(
+            bundle.join("install.sh"),
+            render_install_sh(recipe_name, &recipe.docker_image, &versions, None),
+        )?;
+    }
+    Ok(())
+}
+
+/// Rewrite `patches/…` references to the bundle root on COPY/ADD lines
+/// (gists cannot carry directories). Comments keep their original text.
+pub(crate) fn flatten_dockerfile(text: &str) -> String {
+    let mut in_copy = false;
+    text.lines()
+        .map(|line| {
+            let t = line.trim_start();
+            let is_copy = in_copy || t.starts_with("COPY ") || t.starts_with("ADD ");
+            in_copy = is_copy && t.ends_with('\\');
+            if is_copy {
+                line.replace("patches/", "")
+            } else {
+                line.to_string()
+            }
+        })
+        .collect::<Vec<_>>()
+        .join("\n")
+        + if text.ends_with('\n') { "\n" } else { "" }
+}
+
+/// Parse `versions.env` (`KEY=VAL` lines, `#` comments, optional quotes) into
+/// `--build-arg` pairs.
+pub(crate) fn versions_build_args(text: &str) -> BTreeMap<String, String> {
+    let mut out = BTreeMap::new();
+    for line in text.lines() {
+        let line = line.split('#').next().unwrap_or("").trim();
+        if let Some((k, v)) = line.split_once('=') {
+            let v = v.trim().trim_matches('"').trim_matches('\'');
+            if !k.trim().is_empty() && !v.is_empty() {
+                out.insert(k.trim().to_string(), v.to_string());
+            }
+        }
+    }
+    out
+}
+
+/// Copy every file under `dir` (including `patches/*`) into `dst` root,
+/// flattening paths. Sub-directories other than `patches/` are rejected —
+/// a flat bundle can only reference root-level files.
+fn copy_flat(src_dir: &Path, dst: &Path) -> Result<()> {
+    for entry in std::fs::read_dir(src_dir)? {
+        let entry = entry?;
+        let ft = entry.file_type()?;
+        let name = entry.file_name().to_string_lossy().to_string();
+        if ft.is_dir() {
+            if name != "patches" {
+                bail!(
+                    "publish: unsupported sub-dir '{}' in image overlay (flat bundles only)",
+                    name
+                );
+            }
+            for f in std::fs::read_dir(entry.path())? {
+                let f = f?;
+                if f.file_type()?.is_file() {
+                    std::fs::copy(f.path(), dst.join(f.file_name()))
+                        .with_context(|| "copying patch")?;
+                }
+            }
+        } else if ft.is_file() && name != "build.sh" && name != "Dockerfile" {
+            std::fs::copy(entry.path(), dst.join(&name))?;
+        }
+    }
+    Ok(())
+}
+
+fn render_run_script(
+    recipe: &Recipe,
+    cluster: Option<&Cluster>,
+    nccl: Option<&NcclConfig>,
+    hf_cache: &str,
+    init_port: u16,
+    revision: Option<&str>,
+) -> String {
+    let nodes = cluster.map_or_else(
+        || recipe.required_nodes.unwrap_or(2) as usize,
+        |c| c.nodes.len(),
+    );
+    let solo = recipe.solo == Some(true) || nodes <= 1;
+    let master = cluster
+        .and_then(|c| c.master_ib_ip.clone())
+        .or_else(|| cluster.and_then(|c| c.nodes.iter().find(|n| n.rank == 0)?.ib_ip.clone()))
+        .unwrap_or_else(|| "<rank0-IB-ip>".into());
+    let mut s = String::new();
+    s.push_str(&format!(
+        "#!/usr/bin/env bash\n\
+         # {} — generated by `sm121 publish`. Run on rank 0; on each worker set\n\
+         # NODE_RANK=<rank> and HEADLESS=\"--headless\" (drop --host/--port).\n\
+         set -euo pipefail\n\n\
+         MODEL={}\n\
+         IMAGE={}\n",
+        recipe.name,
+        shq(&recipe.model),
+        shq(&recipe.docker_image),
+    ));
+    if let Some(r) = revision {
+        s.push_str(&format!(
+            "MODEL_REVISION={}   # pin: download + serve THIS revision (repo branches move)\n",
+            shq(r)
+        ));
+    } else if recipe.model.contains('/') {
+        s.push_str(
+            "# MODEL_REVISION=   # consider pinning the HF revision (the repo branch moves)\n",
+        );
+    }
+    if !solo {
+        s.push_str(&format!(
+            "MASTER_ADDR={}   # rank-0 IB ip\nNODE_RANK=0\nHEADLESS=\"\"\n",
+            shq(&master)
+        ));
+    }
+    s.push_str("\n# --- fabric env (adjust per cluster) ------------------------------------\n");
+    match nccl {
+        Some(n) => {
+            s.push_str(&format!(
+                "FABRIC_ENV=(\n  -e NCCL_IB_HCA={}\n  -e NCCL_SOCKET_IFNAME={}\n  -e GLOO_SOCKET_IFNAME={}\n  -e NCCL_IB_DISABLE={}\n",
+                shq(&n.ib_hca), shq(&n.socket_ifname),
+                shq(n.gloo_socket_ifname.as_deref().unwrap_or(&n.socket_ifname)), n.ib_disable));
+            match n.ib_gid_index {
+                Some(g) => s.push_str(&format!("  -e NCCL_IB_GID_INDEX={g}\n)\n")),
+                None => s.push_str("  # NCCL_IB_GID_INDEX renumbers across reboots on DGX Spark —\n  # detect the RoCE-v2 IPv4 GID per node, don't copy a stale number:\n  # for i in $(seq 0 15); do grep -l RoCE* /sys/class/infiniband/$HCA/ports/1/gid_attrs/types/$i 2>/dev/null && echo $i; done\n)\n"),
+            }
+        }
+        None => s.push_str("FABRIC_ENV=(\n  # -e NCCL_IB_HCA=... -e NCCL_SOCKET_IFNAME=... -e NCCL_IB_GID_INDEX=...\n)\n"),
+    }
+    s.push_str(&format!(
+        "\nexec docker run -d --name pub-{} --gpus all --ipc=host --network host \\\n  --shm-size {} \\\n  -v {}:/root/.cache/huggingface \\\n  \"${{FABRIC_ENV[@]}}\" \\\n",
+        recipe.model.replace('/', "-"),
+        if recipe.shm_size.is_empty() { "32g".into() } else { recipe.shm_size.clone() },
+        hf_cache, // unquoted: `~` must expand in bash
+    ));
+    for (k, v) in recipe.env.iter().collect::<BTreeMap<_, _>>() {
+        s.push_str(&format!("  -e {} \\\n", shq(&format!("{k}={v}"))));
+    }
+    s.push_str("  \"$IMAGE\" \\\n");
+    s.push_str(&format!(
+        "  --model \"$MODEL\"{}",
+        if revision.is_some() {
+            " --revision \"$MODEL_REVISION\""
+        } else {
+            ""
+        }
+    ));
+    if !solo {
+        if recipe.engine == "sglang" {
+            s.push_str(&format!(
+                " \\\n  --tp {} --nnodes {nodes} --node-rank \"$NODE_RANK\" --dist-init-addr \"$MASTER_ADDR:{init_port}\"",
+                recipe.tensor_parallel_size
+            ));
+        } else {
+            s.push_str(&format!(
+                " \\\n  --tensor-parallel-size {} --nnodes {nodes} --node-rank \"$NODE_RANK\" --master-addr \"$MASTER_ADDR\" --master-port {init_port}",
+                recipe.tensor_parallel_size
+            ));
+            if let Some(pp) = recipe.pipeline_parallel_size.filter(|p| *p > 1) {
+                s.push_str(&format!(" --pipeline-parallel-size {pp}"));
+            }
+        }
+    }
+    for a in recipe.args.to_vec() {
+        s.push_str(&format!(" \\\n  {}", shq(&a)));
+    }
+    if !solo {
+        s.push_str(" \\\n  ${HEADLESS:+$HEADLESS}");
+    }
+    s.push_str(&format!(" \\\n  --host 0.0.0.0 --port {}\n", recipe.port));
+    s
+}
+
+fn render_install_sh(
+    recipe_name: &str,
+    image: &str,
+    versions: &BTreeMap<String, String>,
+    gist_id: Option<&str>,
+) -> String {
+    let args = versions
+        .iter()
+        .map(|(k, v)| format!("  --build-arg {}={}\n", shq(k), shq(v)))
+        .collect::<String>();
+    let gid = gist_id.unwrap_or("'          '");
+    format!(
+        "#!/usr/bin/env bash\n\
+         # Build {image}. Run from this folder, or set GIST_ID and run from\n\
+         # anywhere to clone the published gist into /tmp first.\n\
+         # Usage: curl -fsSL <raw-url>/install.sh | bash   (after publishing as a gist)\n\
+         set -euo pipefail\n\n\
+         GIST_ID={}  # fill in after `sm121 publish --gist`\n\
+         DIR=/tmp/{recipe_name}-build\n\
+         IMAGE={}\n\n\
+         HERE=\"$(cd \"$(dirname \"$0\")\" && pwd)\"\n\
+         if [ -f \"$HERE/Dockerfile\" ]; then\n  SRC=\"$HERE\"\n\
+         else\n  rm -rf \"$DIR\"; git clone -q \"https://gist.github.com/${{GIST_ID}}.git\" \"$DIR\"; SRC=\"$DIR\"\n\
+         fi\n\n\
+         docker build -t \"$IMAGE\" \\\n{}\"$SRC\"\n\
+         echo \"built: $IMAGE — see run.sh\"\n",
+        shq(image),
+        gid,
+        if args.is_empty() {
+            String::new()
+        } else {
+            format!(" \\\n{args}")
+        },
+    )
+}
+
+fn render_readme(
+    recipe_name: &str,
+    recipe: &Recipe,
+    versions: &BTreeMap<String, String>,
+    revision: Option<&str>,
+    has_image: bool,
+) -> String {
+    let mut s = format!("# {} — reproducible recipe card\n\n", recipe.name);
+    s.push_str(&format!(
+        "Generated by `sm121 publish` from the [sm121-launcher](https://git.kilork.org/kilork/sm121-launcher) \
+         recipe `{recipe_name}` (engine: {}, model: `{}`).\n\n",
+        recipe.engine, recipe.model
+    ));
+    s.push_str(
+        "**`recipe.toml` is the source of truth — its header comments document every \
+                non-obvious flag and the measured boot markers. Read it first.**\n\n",
+    );
+    s.push_str("## Files\n\n| file | what |\n|---|---|\n| `recipe.toml` | the launcher recipe (rationale in comments) |\n");
+    if has_image {
+        s.push_str(
+            "| `Dockerfile`, patches, `versions.env` | image overlay (flat — build here) |\n",
+        );
+    }
+    s.push_str("| `install.sh` | one-shot build (gist-aware) |\n| `run.sh` | docker run, rank 0 (worker deltas in comments) |\n\n");
+    if has_image {
+        s.push_str("## Build\n\n```bash\nbash install.sh\n```\n\n");
+        s.push_str("## Pinned coordinates\n\n| thing | ref |\n|---|---|\n");
+        if let Some(b) = versions.get("BASE_IMAGE") {
+            let note = if b.contains('@') {
+                "digest-pinned"
+            } else {
+                "⚠ mutable tag — re-pin if reproducibility matters"
+            };
+            s.push_str(&format!("| base image | `{b}` ({note}) |\n"));
+        }
+        match revision {
+            Some(r) => s.push_str(&format!("| model | `{}` @ `{r}` |\n", recipe.model)),
+            None if recipe.model.contains('/') => s.push_str(&format!(
+                "| model | `{}` @ main — ⚠ unpinned branch |\n",
+                recipe.model
+            )),
+            _ => {}
+        }
+        s.push('\n');
+    }
+    s.push_str(
+        "## Run\n\nCache the model on every node (see `recipe.toml` PREREQUISITES), edit \
+                `MASTER_ADDR`/fabric env in `run.sh`, run on rank 0, then on each worker with \
+                `NODE_RANK`/`HEADLESS` set. `recipe.toml` `[[extra_mounts]]` and built-in \
+                `before_serve` host scripts are launcher-side — translate manually if needed.\n",
+    );
+    s
+}
+
+async fn hf_revision(model: &str) -> Option<String> {
+    if !model.contains('/') || model.starts_with('.') || model.starts_with('/') {
+        return None;
+    }
+    let url = format!("https://huggingface.co/api/models/{model}");
+    let resp = reqwest::Client::builder()
+        .timeout(std::time::Duration::from_secs(8))
+        .build()
+        .ok()?
+        .get(url)
+        .send()
+        .await
+        .ok()?
+        .json::<serde_json::Value>()
+        .await
+        .ok()?;
+    resp.get("sha")?.as_str().map(str::to_string)
+}
+
+fn publish_gist(bundle: &Path, desc: &str) -> Result<String> {
+    let mut cmd = Command::new("gh");
+    cmd.arg("gist")
+        .arg("create")
+        .arg("--public")
+        .arg("--desc")
+        .arg(desc);
+    let mut files = Vec::new();
+    for e in std::fs::read_dir(bundle)? {
+        let p = e?.path();
+        if p.is_file() {
+            files.push(p);
+        }
+    }
+    files.sort();
+    cmd.args(&files);
+    let out = cmd
+        .output()
+        .context("`gh` not found — install GitHub CLI or copy the bundle manually")?;
+    if !out.status.success() {
+        bail!(
+            "gh gist create failed: {}",
+            String::from_utf8_lossy(&out.stderr)
+        );
+    }
+    let url = String::from_utf8_lossy(&out.stdout).trim().to_string();
+    info!("🌐 gist: {url}");
+    // gh prints the full html url; the gist *id* is the last path segment.
+    Ok(url.rsplit('/').next().unwrap_or(&url).to_string())
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn flatten_rewrites_copy_lines_only() {
+        let df = "FROM x\n# see patches/UPSTREAM.md\nCOPY patches/a.py /tmp/a.py\nCOPY b /tmp/ \\\n  patches/c.patch\nRUN grep patches/ /tmp/x\n";
+        let f = flatten_dockerfile(df);
+        assert!(f.contains("# see patches/UPSTREAM.md"));
+        assert!(f.contains("COPY a.py /tmp/a.py"));
+        assert!(f.contains("  c.patch\nRUN grep patches/"));
+    }
+
+    #[test]
+    fn versions_env_parses_quoted_and_commented() {
+        let v = versions_build_args("# c\nBASE_IMAGE=\"img@sha256:1\"\nX = 'y' \nNOVAL=\n");
+        assert_eq!(v.get("BASE_IMAGE").unwrap(), "img@sha256:1");
+        assert_eq!(v.get("X").unwrap(), "y");
+        assert!(!v.contains_key("NOVAL"));
+    }
+
+    fn rec(engine: &str) -> Recipe {
+        toml::from_str(&format!(
+            "engine = \"{engine}\"\nname = \"T\"\ndocker_image = \"i:t\"\nmodel = \"o/m\"\n\
+             port = 30000\nrequired_nodes = 2\ntensor_parallel_size = 2\n\
+             args = [\"--speculative-config\", \"'{{\\\"method\\\":\\\"mtp\\\"}}'\"]\n"
+        ))
+        .unwrap()
+    }
+
+    #[test]
+    fn run_script_vllm_has_dist_flags_and_json_arg() {
+        let s = render_run_script(&rec("vllm"), None, None, "/hf", 29500, Some("abc"));
+        assert!(s.contains("--tensor-parallel-size 2 --nnodes 2 --node-rank \"$NODE_RANK\""));
+        assert!(s.contains("--revision \"$MODEL_REVISION\""));
+        assert!(s.contains("'--speculative-config'"));
+        assert!(s.contains(&shq("'{\"method\":\"mtp\"}'"))); // json arg survives shq quoting intact
+    }
+
+    #[test]
+    fn run_script_sglang_solo_omits_dist() {
+        let mut r = rec("sglang");
+        r.solo = Some(true);
+        let s = render_run_script(&r, None, None, "/hf", 29500, None);
+        assert!(!s.contains("--nnodes") && !s.contains("MASTER_ADDR"));
+        assert!(s.contains("# MODEL_REVISION="));
+    }
+}