launcher/src/publish.rs

//! `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,
    registry: Option<&str>,
    token: Option<&str>,
) -> 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());

    // install.sh is part of the bundle even for the gist flow — the gist
    // branch only patches the returned id into it afterwards.
    std::fs::write(
        bundle.join("install.sh"),
        render_install_sh(recipe_name, &recipe.docker_image, &versions, None),
    )?;

    if let Some(url) = registry {
        crate::registry::upload(&bundle, recipe_name, url, token).await?;
    }

    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"),
        }
    }
    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/sm121/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="));
    }
}