launcher/src/publish_render.rs

//! Bundle rendering for `sm121 publish`: run script, install script, README,
//! HF revision lookup, and the `gh gist` flow.

use std::collections::BTreeMap;
use std::path::Path;
use std::process::Command;

use anyhow::{Context, Result, bail};
use tracing::info;

use crate::config::{Cluster, NcclConfig, Recipe};
use crate::ssh::shq;

pub(crate) fn render_run_script(
    recipe: &Recipe,
    cluster: Option<&Cluster>,
    nccl: Option<&NcclConfig>,
    hf_cache: &str,
    init_port: u16,
    revision: Option<&str>,
) -> Result<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
    ));
    let env = recipe.resolved_env()?;
    for (k, v) in 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));
    Ok(s)
}

pub(crate) 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
        .map(shq)
        .unwrap_or_else(|| "'          '".to_string());
    let image_q = shq(image);
    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={gid}  # filled in by `sm121 publish --gist`\n\
         DIR=/tmp/{recipe_name}-build\n\
         IMAGE={image_q}\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{args}\"$SRC\"\n\
         echo \"built: $IMAGE — see run.sh\"\n",
    )
}

pub(crate) fn render_readme(
    recipe_name: &str,
    recipe: &Recipe,
    versions: &BTreeMap<String, String>,
    revision: Option<&str>,
    has_image: bool,
    profile_used: 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 profile_used {
            s.push_str(&format!(
                "| image pins | `versions.env` with {} merged in (precedence: image base < profile < recipe [build.versions]) |\n",
                recipe
                    .build_profile
                    .as_deref()
                    .map(|p| format!("profile `{p}`"))
                    .unwrap_or_else(|| "`[build.versions]` overrides".into())
            ));
        }
        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
}

pub(crate) 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)
}

pub(crate) 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())
}