|
|
@@ -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="));
|
|
|
+ }
|
|
|
+}
|