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`.
//!
//! `patches/` files are flattened into the bundle root (gists cannot carry
//! directories); other sub-directories are preserved and `--gist` is refused
//! for such overlays.

use anyhow::{Context, Result, bail};
use std::collections::BTreeMap;
use std::path::{Path, PathBuf};
use tracing::info;

use crate::config::{Config, Recipe, recipe_dirs};
use crate::docker::build_dir::{build_checksum, flatten_overlay};
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 struct PushOpts {
    pub push: bool,
    pub node: Option<String>,
    pub force: bool,
}

/// The (separate) recipe-bundle registry: where to upload the bundle and with
/// which token. Not to be confused with the image  push target.
pub struct RecipeRegistryOpts {
    pub url: Option<String>,
    pub token: Option<String>,
}

pub async fn publish(
    config: &Config,
    recipe_name: &str,
    cluster: Option<&str>,
    out: Option<PathBuf>,
    gist: bool,
    recipe_registry: &RecipeRegistryOpts,
    push_opts: PushOpts,
) -> 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 profile_used = recipe.build_profile.clone().is_some_and(|p| !p.is_empty())
        || !recipe.build.versions.is_empty();
    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"))?;
        let (df_flat, nested) = flatten_overlay(&image_dir, &bundle, &df)?;
        std::fs::write(bundle.join("Dockerfile"), &df_flat)?;
        if gist && nested {
            bail!(
                "publish --gist needs a flat overlay ({image_dir:?} nests non-flattenable dirs) — use --out/--registry instead"
            );
        }
        // Merge the profile + recipe [build.versions] pins the same way the
        // build path does, so the bundle builds the image that was served.
        let (merged_env, _) = build_checksum(
            &image_dir,
            recipe.build_profile.as_deref(),
            &recipe.build.versions,
        )?;
        if let Some(merged) = merged_env {
            std::fs::write(bundle.join("versions.env"), &merged)?;
            versions = versions_build_args(&merged);
            info!(
                "📦 bundled image overlay {} with merged pin set (profile: {})",
                image_dir.display(),
                recipe
                    .build_profile
                    .as_deref()
                    .unwrap_or("— recipe [build.versions] only")
            );
        } else 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());
        }
        // Bundle the workspace's generic build driver (forwards every pin as
        // --build-arg, prints the `Done: <tag>` marker the launcher labels
        // with) when one exists; otherwise a minimal driver.
        let ws_build = image_dir.join("build.sh");
        if let Ok(bs) = std::fs::read_to_string(&ws_build) {
            std::fs::write(bundle.join("build.sh"), bs)?;
        } else {
            std::fs::write(
                bundle.join("build.sh"),
                "#!/usr/bin/env bash\nset -euo pipefail; cd \"$(dirname \"$0\")\"\n\
                 TAG=\"${TAG:?set TAG=<recipe docker_image>}\"\nDOCKER_BUILDKIT=1 docker build -t \"$TAG\" -f Dockerfile . 2>&1\n\
                 echo \"=== Done: ${TAG} ===\"\n",
            )?;
        }
    } else {
        info!(
            "ℹ️  no docker-images/{} overlay — bundling recipe + scripts only",
            image_dir.display()
        );
        std::fs::write(
            bundle.join("build.sh"),
            "#!/usr/bin/env bash\nset -euo pipefail; cd \"$(dirname \"$0\")\"\n\
             TAG=\"${TAG:?set TAG=<recipe docker_image>}\"\nDOCKER_BUILDKIT=1 docker build -t \"$TAG\" -f Dockerfile . 2>&1\n\
             echo \"=== Done: ${TAG} ===\"\n",
        )?;
    }

    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 = crate::publish_render::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 = crate::publish_render::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 = crate::publish_render::render_readme(
        recipe_name,
        &recipe,
        &versions,
        revision.as_deref(),
        image_dir.exists(),
        profile_used,
    );
    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"),
        crate::publish_render::render_install_sh(
            recipe_name,
            &recipe.docker_image,
            &versions,
            None,
        ),
    )?;

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

    if gist {
        let id = crate::publish_render::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)?;
        // Push the updated install.sh back to the gist (gh gist edit is
        // interactive-only — use the API):
        let body = serde_json::json!({ "files": { "install.sh": { "content": raw } } }).to_string();
        let pushed = std::process::Command::new("gh")
            .args(["api", "-X", "POST", &format!("gists/{id}"), "--input", "-"])
            .stdin(std::process::Stdio::piped())
            .stdout(std::process::Stdio::null())
            .stderr(std::process::Stdio::null())
            .spawn()
            .and_then(|mut child| {
                if let Some(mut stdin) = child.stdin.take() {
                    use std::io::Write;
                    stdin.write_all(body.as_bytes())?;
                }
                child.wait()
            });
        match pushed {
            Ok(s) if s.success() => info!("🔗 install.sh updated with gist id"),
            _ => info!("ℹ️  could not auto-update install.sh — set GIST_ID={id} manually"),
        }
    }

    if push_opts.push {
        push_recipe_image(config, &recipe, &push_opts).await?;
    }
    Ok(())
}

/// `publish --push`: push the recipe's image to the `[registry]` registry
/// under its version tag (the `sm121.checksum` label, or the image ID when
/// unlabeled) + `latest`. The image is found on `--node` or by probing all
/// cluster nodes; the push runs from that node.
async fn push_recipe_image(config: &Config, recipe: &Recipe, opts: &PushOpts) -> Result<()> {
    if config.registry.is_none() {
        anyhow::bail!(
            "publish --push requires a [registry] section in config.toml (base = ghcr.io/<you>)"
        );
    }
    let image = recipe.docker_image.trim();
    if image.is_empty() {
        anyhow::bail!(
            "recipe '{}' has no docker_image — nothing to push",
            recipe.name
        );
    }
    // 1. Find a node that has the image.
    let image = image.to_string();
    let found: Option<crate::config::Node> = if let Some(want) = &opts.node {
        let node = config
            .clusters
            .iter()
            .flat_map(|c| c.nodes.iter())
            .find(|n| n.name.as_str() == want.as_str())
            .with_context(|| format!("node '{want}' not in any cluster"))?
            .clone();
        let id = crate::ssh::run_ssh(
            config,
            &node,
            format!(
                "docker image ls --format '{{{{.ID}}}}' {img} | head -1",
                img = shq(&image)
            ),
            false,
        )
        .await
        .unwrap_or_default();
        if id.trim().is_empty() {
            None
        } else {
            Some(node)
        }
    } else {
        let mut seen = std::collections::HashSet::new();
        let all: Vec<crate::config::Node> = config
            .clusters
            .iter()
            .flat_map(|c| c.nodes.iter().cloned())
            .filter(|n| seen.insert(n.ssh_ip.clone()))
            .collect();
        let handles = all.iter().map(|node| {
            let config = config.clone();
            let node = node.clone();
            let image = image.clone();
            crate::locale::spawn_with_locale(async move {
                let id = crate::ssh::run_ssh(
                    &config,
                    &node,
                    format!(
                        "docker image ls --format '{{{{.ID}}}}' {img} | head -1",
                        img = shq(&image)
                    ),
                    false,
                )
                .await
                .unwrap_or_default();
                (node.name.clone(), id.trim().to_string())
            })
        });
        let results = futures::future::join_all(handles).await;
        all.into_iter()
            .zip(results)
            .find_map(|(node, h)| h.ok().filter(|(_, id)| !id.is_empty()).map(|_| node))
    };
    let Some(node) = found else {
        anyhow::bail!(
            "image '{image}' not found on {} — build it first (sm121 build --recipe {} --build-node <n>)",
            opts.node
                .as_deref()
                .map(|n| format!("'{n}'"))
                .unwrap_or_else(|| "any node".to_string()),
            recipe.name
        );
    };
    // 2. Version = the sm121.checksum label (build checksum), else the image ID.
    let label = crate::ssh::run_ssh(
        config,
        &node,
        format!(
            "docker inspect --format '{{{{index .Config.Labels \"sm121.checksum\"}}}}' {img} 2>/dev/null || true",
            img = shq(&image)
        ),
        false,
    )
    .await
    .unwrap_or_default();
    let version = if label.trim().len() >= 12 {
        label.trim()[..12].to_string()
    } else {
        let id = crate::ssh::run_ssh(
            config,
            &node,
            format!(
                "docker image inspect --format '{{{{.Id}}}}' {img}",
                img = shq(&image)
            ),
            false,
        )
        .await
        .unwrap_or_default();
        let id = id.trim().strip_prefix("sha256:").unwrap_or(id.trim());
        id.get(..12).unwrap_or(id).to_string()
    };
    info!(
        "publish --push: pushing '{image}' v{version} from '{}'",
        node.name
    );
    crate::docker::push::push_image(config, &node, &image, &version, opts.force, "publish").await
}
/// 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
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::config::Recipe;

    #[test]
    fn flatten_overlay_rewrites_copy_lines_and_collision_renames() {
        let tmp = std::env::temp_dir().join("sm121_flatten_overlay_test");
        let _ = std::fs::remove_dir_all(&tmp);
        let src = tmp.join("src");
        std::fs::create_dir_all(src.join("locks")).unwrap();
        std::fs::create_dir_all(src.join("patches")).unwrap();
        std::fs::create_dir_all(src.join("profile")).unwrap();
        std::fs::write(src.join("apt-sources.list"), "root").unwrap();
        std::fs::write(src.join("locks/apt-sources.list"), "lock").unwrap(); // collision
        std::fs::write(src.join("locks/apt-packages.txt"), "pkg").unwrap();
        std::fs::write(src.join("locks/python-bootstrap.txt"), "boot").unwrap(); // unreferenced -> not shipped
        std::fs::write(src.join("patches/p.patch"), "z").unwrap();
        std::fs::write(src.join("profile/x.env"), "y").unwrap();
        std::fs::create_dir_all(src.join("extra")).unwrap();
        std::fs::write(src.join("extra/y.env"), "z2").unwrap(); // unreferenced dir -> not shipped
        std::fs::write(src.join("Dockerfile"), "f").unwrap();
        let df = "FROM x\n# see patches/UPSTREAM.md\nCOPY apt-sources.list /tmp/root-a.list\nCOPY locks/apt-sources.list /tmp/a.list\nCOPY locks/apt-packages.txt /tmp/b.txt \\\n  patches/p.patch\nCOPY profile/x.env /tmp/x\nRUN grep patches/ /tmp/x\n";
        let dst = tmp.join("dst");
        std::fs::create_dir_all(&dst).unwrap();
        let (flat, nested) = crate::docker::build_dir::flatten_overlay(&src, &dst, df).unwrap();
        assert!(
            !nested,
            "nothing nests anymore — unreferenced files are dropped"
        );
        // collision rename for the referenced locks file
        assert!(
            dst.join("locks-apt-sources.list").exists(),
            "collision rename"
        );
        assert!(dst.join("apt-sources.list").exists(), "root file shipped");
        assert_eq!(
            std::fs::read_to_string(dst.join("apt-packages.txt")).unwrap(),
            "pkg"
        );
        assert_eq!(std::fs::read_to_string(dst.join("p.patch")).unwrap(), "z");
        assert!(dst.join("x.env").exists(), "profile dir flattened");
        // unreferenced files / dirs are not shipped
        assert!(!dst.join("python-bootstrap.txt").exists());
        assert!(!dst.join("extra").exists());
        // Dockerfile rewritten: renamed file, prefix drops (root + continuation),
        // comments + non-COPY lines untouched
        assert!(flat.contains("COPY locks-apt-sources.list /tmp/a.list"));
        assert!(flat.contains("COPY apt-packages.txt /tmp/b.txt \\\n  p.patch"));
        assert!(flat.contains("COPY x.env /tmp/x"));
        assert!(flat.contains("COPY apt-sources.list /tmp/root-a.list"));
        assert!(flat.contains("# see patches/UPSTREAM.md"));
        assert!(flat.contains("RUN grep patches/ /tmp/x"));
        assert!(!dst.join("Dockerfile").exists());
    }

    #[test]
    fn flatten_overlay_glob_collision_keeps_dir_nested() {
        let tmp = std::env::temp_dir().join("sm121_flatten_glob_test");
        let _ = std::fs::remove_dir_all(&tmp);
        let src = tmp.join("src");
        std::fs::create_dir_all(src.join("patches")).unwrap();
        std::fs::write(src.join("a.py"), "root").unwrap(); // referenced + collides with patches/a.py
        std::fs::write(src.join("patches/a.py"), "patched").unwrap();
        let df = "FROM x\nCOPY a.py /tmp/a0.py\nCOPY patches/*.py /tmp/\n";
        let dst = tmp.join("dst");
        std::fs::create_dir_all(&dst).unwrap();
        let (flat, nested) = crate::docker::build_dir::flatten_overlay(&src, &dst, df).unwrap();
        assert!(nested, "glob collision must not flatten");
        assert!(dst.join("patches/a.py").exists());
        assert!(
            flat.contains("COPY patches/*.py /tmp/"),
            "line left untouched"
        );
    }

    #[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 = crate::publish_render::render_run_script(
            &rec("vllm"),
            None,
            None,
            "/hf",
            29500,
            Some("abc"),
        )
        .unwrap();
        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 =
            crate::publish_render::render_run_script(&r, None, None, "/hf", 29500, None).unwrap();
        assert!(!s.contains("--nnodes") && !s.contains("MASTER_ADDR"));
        assert!(s.contains("# MODEL_REVISION="));
    }
}