launcher/src/native/flare.rs

//! FLARE deploy (`engine = "flare"`, `deploy = "native"`).
//!
//! FLARE (`local-inference-lab/flare`) serves a *signed package* instead of a
//! raw HF checkpoint: `flare download <repo> --revision <rev>` fetches it, and
//! `flare orbit run "<controller>:<path>" --cluster <orbit cluster> …` brings
//! up one rank per paired agent and serves the API on the controller. The
//! launcher therefore does what a human would: preflight the agents, download
//! the package once on rank 0 (stamped, so a repeat serve costs nothing), run
//! orbit detached, then probe `/health` + `/v1/models`.
//!
//! Recipe block:
//! ```toml
//! engine = "flare"
//! deploy = "native"
//! [flare]
//! package   = "local-inference-lab/DeepSeek-V4.1-Flash-1017-FLARE"
//! revision  = "gb10"
//! name      = "deepseek1017"
//! http_port = 8000
//! serve_flags = ["--max-concurrency", "16", "--mtp-draft-tokens", "5"]
//! ```
//! Credentials (`--dangerously-approve-key`, publisher) live in the private
//! config's `[flare]` block, never in the recipe repo.

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

use crate::config::{Config, Node, Recipe};
use crate::docker::{ServeOpts, shq};

/// FLARE API port on the controller (recipe `[flare].http_port`, else the
/// cluster default 30000 — the same port every other lane serves on).
pub fn http_port(recipe: &Recipe) -> u16 {
    recipe
        .flare
        .as_ref()
        .and_then(|f| f.http_port)
        .unwrap_or(30000)
}

/// Orbit cluster name for one of OUR clusters: the per-cluster mapping wins,
/// then the flat `[flare].orbit_cluster`, then the default. This is what makes
/// one launcher config work on `dgx58-switch` and `prod4` alike.
pub fn orbit_cluster_for(config: &Config, cluster: &str) -> String {
    let f = config.flare.as_ref();
    f.and_then(|f| f.clusters.get(cluster))
        .and_then(|c| c.orbit_cluster.clone())
        .or_else(|| f.and_then(|f| f.orbit_cluster.clone()))
        .unwrap_or_else(|| "sparks".to_string())
}

/// The orbit cluster a **specific** cluster is paired with, or `None` when the
/// config does not say (the flat `[flare].orbit_cluster` is then only a
/// fallback for `serve`, not evidence of a pairing).
///
/// This is the identification rule the dashboard and `health` use: a cluster
/// whose own nodes were never paired with an orbit must not be probed with the
/// flat default (`sparks`) — `6nodes-switch`, `1node-*`, `ds4` and `qwen38`
/// share nodes with the paired clusters, so asking one of them `--cluster
/// sparks` made the sparks lane look like a lane of *that* cluster.
pub fn orbit_cluster_paired(config: &Config, cluster: &str) -> Option<String> {
    let f = config.flare.as_ref()?;
    if let Some(own) = f.clusters.get(cluster) {
        // A per-cluster block without an orbit name inherits the flat one —
        // the operator already said "this cluster speaks FLARE".
        return own
            .orbit_cluster
            .clone()
            .or_else(|| f.orbit_cluster.clone());
    }
    // Once the map exists it is the source of truth: the flat block describes
    // the machine it was written for, and our clusters share nodes (dgx5 is
    // rank 0 of `dgx58-switch` AND rank 4 of `6nodes-switch`), so a flat
    // `sparks` would pair `6nodes-switch` too.
    if !f.clusters.is_empty() {
        return None;
    }
    let flat = f.orbit_cluster.clone()?;
    match f.controller_node.as_deref() {
        // Legacy flat block: it speaks for the cluster that holds its
        // controller node.
        Some(ctl) => config
            .clusters
            .iter()
            .find(|c| c.name == cluster)
            .is_some_and(|c| c.nodes.iter().any(|n| n.name == ctl))
            .then_some(flat),
        // No controller either: unambiguous only with a single cluster.
        None => (config.clusters.len() == 1 && config.clusters[0].name == cluster).then_some(flat),
    }
}

/// Extra `orbit run` flags for a cluster (per-cluster map first, then flat).
/// Fabric shape is a cluster property. The prod 4-node switch uses no extra
/// flag: FLARE finds the star links itself once the legacy CX7 ports are down;
/// `--roce-ring` was a symptom-level fix and is retired.
pub fn extra_flags_for(config: &Config, cluster: &str) -> Vec<String> {
    let Some(f) = config.flare.as_ref() else {
        return Vec::new();
    };
    let per = f
        .clusters
        .get(cluster)
        .map(|c| c.extra_flags.clone())
        .unwrap_or_default();
    if per.is_empty() {
        f.extra_flags.clone()
    } else {
        per
    }
}

/// Local package mirror dir for a cluster (per-cluster map first).
pub(crate) fn local_dir_for(config: &Config, cluster: &str) -> Option<String> {
    let f = config.flare.as_ref()?;
    f.clusters
        .get(cluster)
        .and_then(|c| c.local_dir.clone())
        .or_else(|| f.local_dir.clone())
}

/// The orbit-side name for one of our nodes (rank order → `orbit_nodes`,
/// falling back to the cluster's own name).
fn orbit_node_name(config: &Config, cluster: &str, node: &Node) -> String {
    let f = config.flare.as_ref();
    let per_cluster = f
        .and_then(|f| f.clusters.get(cluster))
        .map(|c| c.orbit_nodes.as_slice())
        .filter(|v| !v.is_empty());
    per_cluster
        .or_else(|| f.map(|f| f.orbit_nodes.as_slice()))
        .and_then(|v| v.get(node.rank as usize).cloned())
        .unwrap_or_else(|| node.name.clone())
}

pub(crate) fn binary(config: &Config) -> String {
    config
        .flare
        .as_ref()
        .and_then(|f| f.binary.clone())
        .unwrap_or_else(|| "flare".to_string())
}

/// `flare download <package> --revision <rev>` — prints the package path.
pub fn download_cmd(config: &Config, recipe: &Recipe) -> Result<String> {
    let f = recipe
        .flare
        .as_ref()
        .ok_or_else(|| anyhow::anyhow!("engine=flare requires a [flare] block"))?;
    let mut cmd = format!("{} download {}", binary(config), shq(&f.package));
    if let Some(rev) = &f.revision {
        cmd.push_str(&format!(" --revision {}", shq(rev)));
    }
    // A prepared package is built for one TP size; Spark has one visible GPU,
    // so `flare download` defaults to TP1 and rejects a TP4 package
    // ("package exports no deployment with tensor-parallel size 1").
    // Use the launch-node count as the package's TP (a recipe with
    // `required_nodes` set is a TP-N package; base recipes fall back to
    // tensor_parallel_size).
    let tp = recipe
        .required_nodes
        .unwrap_or(recipe.tensor_parallel_size)
        .max(1);
    cmd.push_str(&format!(" --tensor-parallel {}", tp));
    // NOTE: `flare download` does NOT accept the approval flags (it fails with
    // "unexpected argument") — the signature gate lives on `orbit run`, which
    // is where `run_cmd` passes them. A download has always been accepted
    // unattended, so approving a publisher here would break the serve.
    Ok(cmd)
}

/// `flare orbit run "<node>:<package>" --cluster <c> --name <n> … --detach`.
pub fn run_cmd(
    config: &Config,
    recipe: &Recipe,
    our_cluster: &str,
    orbit_cluster: &str,
    controller: &str,
    package: &str,
) -> Result<String> {
    let f = recipe
        .flare
        .as_ref()
        .ok_or_else(|| anyhow::anyhow!("engine=flare requires a [flare] block"))?;
    let model_name = f.name.clone().unwrap_or_else(|| recipe.model.clone());
    let mut cmd = format!(
        "{} orbit run {} --cluster {} --name {} --host 0.0.0.0 --port {} --detach",
        binary(config),
        shq(&format!("{controller}:{package}")),
        shq(orbit_cluster),
        shq(&model_name),
        http_port(recipe),
    );
    if let Some(k) = config.flare.as_ref().and_then(|c| c.accept_key.clone()) {
        cmd.push_str(&format!(" --dangerously-approve-key {}", shq(&k)));
    }
    if let Some(p) = config
        .flare
        .as_ref()
        .and_then(|c| c.accept_publisher.clone())
    {
        cmd.push_str(&format!(" --dangerously-accept-publisher {}", shq(&p)));
    }
    for flag in &f.serve_flags {
        cmd.push_str(&format!(" {}", shq(flag)));
    }
    // fabric shape is a cluster property (a switch-emulated ring, a DAC chain)
    for flag in extra_flags_for(config, our_cluster) {
        cmd.push_str(&format!(" {}", shq(&flag)));
    }
    Ok(cmd)
}

/// `flare orbit stop <name> --cluster <c>`.
pub fn stop_cmd(config: &Config, recipe: &Recipe, orbit_cluster: &str) -> Result<String> {
    let f = recipe
        .flare
        .as_ref()
        .ok_or_else(|| anyhow::anyhow!("engine=flare requires a [flare] block"))?;
    let model_name = f.name.clone().unwrap_or_else(|| recipe.model.clone());
    Ok(format!(
        "{} orbit stop {} --cluster {}",
        binary(config),
        shq(&model_name),
        shq(orbit_cluster),
    ))
}

/// `flare orbit attach <name> --cluster <c>` (log tail; the CLI's `logs`).
pub fn attach_cmd(config: &Config, recipe: &Recipe, orbit_cluster: &str) -> Result<String> {
    let f = recipe
        .flare
        .as_ref()
        .ok_or_else(|| anyhow::anyhow!("engine=flare requires a [flare] block"))?;
    let model_name = f.name.clone().unwrap_or_else(|| recipe.model.clone());
    Ok(format!(
        "{} orbit attach {} --cluster {}",
        binary(config),
        shq(&model_name),
        shq(orbit_cluster),
    ))
}

/// Preflight on every launch node: the agent installed and the orbit cluster
/// visible. Runs before any download or launch — a peer that never paired
/// otherwise fails minutes into a 500 GB fetch.
///
/// Only the controller has the orbit file with `[cluster.<name>]` — a peer
/// just needs the binary and a live agent, so the cluster check is added
/// there and nowhere else.
pub fn preflight_cmd(config: &Config, orbit_cluster: &str, controller: bool) -> String {
    let base = format!(
        "{} --version >/dev/null && systemctl --user is-active flare-agent >/dev/null",
        binary(config)
    );
    if controller {
        format!(
            "{base} && {} orbit status --cluster {} >/dev/null",
            binary(config),
            shq(orbit_cluster),
        )
    } else {
        base
    }
}

/// Mirror a package's files into the controller's local dir and return the
/// local primary path.
///
/// FLARE names every companion by its sha256, so "present and the same size"
/// is a sound consistency check — a truncated or stale file re-copies, and an
/// unchanged package costs a `stat` per file. The primary (metadata) is tiny
/// and always refreshed.
pub fn mirror_cmd(local_dir: &str, remote_pkg: &str) -> String {
    let dir = shq(local_dir);
    let pkg = shq(remote_pkg);
    // New packages keep companion data in nested `weights/` and
    // `deployments/` directories, so a flat file loop is wrong: copy the
    // package directory recursively, and skip the copy when the cache root
    // is already the package directory itself.
    format!(
        "set -e; src=$(dirname {pkg}); dst={dir}; \
         if [ \"$src\" = \"$dst\" ]; then echo {pkg}; exit 0; fi; \
         mkdir -p \"$dst\"; \
         rsync -a --delete --exclude '.sm121-flare-package.json' \"$src\"/ \"$dst\"/; \
         echo \"$dst/$(basename {pkg})\""
    )
}

/// Pick the package path out of `flare download` output: it prints progress
/// lines ("Downloading model file: …") and finally the resolved package path
/// ending in `.safetensors`. Taking the last line blindly handed `orbit run`
/// a progress line and the group died silently — so scan backwards for a path.
pub fn parse_package_path(stdout: &str) -> Option<String> {
    // Cached downloads print only a "Resolved package: … / model.safetensors"
    // line, so require an absolute path and ignore the metadata lines.
    stdout
        .lines()
        .rev()
        .map(str::trim)
        .find(|l| l.starts_with('/') && l.ends_with(".safetensors"))
        .map(str::to_string)
}

fn head(cluster: &crate::config::Cluster) -> Result<&Node> {
    cluster
        .nodes
        .iter()
        .find(|n| n.rank == 0)
        .ok_or_else(|| anyhow::anyhow!("no head node (rank 0) in cluster"))
}

/// Launch a FLARE recipe: preflight, package download on the controller, then
/// `orbit run` detached. Readiness (`/health`, `/v1/models`) is polled by the
/// caller, like every other engine.
pub async fn serve(
    config: &Config,
    recipe: &Recipe,
    cluster: &crate::config::Cluster,
    dry_run: bool,
) -> Result<()> {
    let controller = head(cluster)?;
    let orbit = orbit_cluster_for(config, &cluster.name);
    if dry_run {
        println!(
            "  [dry-run] preflight on {} nodes (cluster check on the controller only): {}",
            cluster.nodes.len(),
            preflight_cmd(config, &orbit, true)
        );
    } else {
        for node in &cluster.nodes {
            let pre = preflight_cmd(config, &orbit, node.rank == 0);
            if let Err(e) = crate::ssh::run_ssh(config, node, pre, false).await {
                bail!(
                    "FLARE preflight failed on {}: {e} — install `flare`, then `flare agent install` and pair the node with the controller",
                    node.name
                );
            }
        }
        info!("FLARE agents ok on {} node(s)", cluster.nodes.len());
    }
    let id = crate::native::flare_package::package_id(recipe)?;
    // Identification: a complete local mirror of this exact package (repo,
    // revision, TP, and the CURRENT prepared version) means the download
    // verification and the 500 GB recopy are both unnecessary.
    let mut package: String;
    if !dry_run {
        if let Some(local) =
            crate::native::flare_package::local_package(config, recipe, &cluster.name, controller)
                .await?
        {
            info!(
                "FLARE package already local on {}; skipping download and mirror",
                controller.name
            );
            package = local;
        } else {
            let dl = download_cmd(config, recipe)?;
            let out = crate::ssh::run_ssh_long_running(config, controller, dl, dry_run).await?;
            package = parse_package_path(&out).unwrap_or_default();
            // A cached download prints only "Resolved package: … / model.safetensors"
            // and no final path; resolve the default prepared model from Flare's
            // model cache instead of failing after a successful download.
            if package.is_empty() {
                let cache_cmd = format!(
                    "{} models list --json 2>/dev/null | python3 -c 'import json,sys; d=json.load(sys.stdin); m=d.get(\"models\") or []; vers=(m[0].get(\"versions\") or []) if m else []; v=vers[0] if vers else dict(); print(d.get(\"cache\",\"\") + \"/models/\" + m[0].get(\"model_key\",\"\") + \"/\" + v.get(\"version\",\"\"))'",
                    binary(config)
                );
                let dir = crate::ssh::run_ssh(config, controller, cache_cmd, false).await?;
                let dir = dir.trim();
                if !dir.is_empty() {
                    let find = format!(
                        "find {} -maxdepth 1 -name '*.safetensors' -print -quit",
                        shq(dir)
                    );
                    let primary = crate::ssh::run_ssh(config, controller, find, false).await?;
                    package = primary.trim().to_string();
                }
            }
            if package.is_empty() {
                bail!(
                    "flare download printed no package path on {}",
                    controller.name
                );
            }

            // Local mirror: launch from the controller's own disk when configured —
            // engram rows are read at request time, so a shared mount is the wrong
            // place for the package even with a RAM row cache.
            if let Some(dir) = local_dir_for(config, &cluster.name) {
                let dir = dir.replace('~', "$HOME");
                // Version from the NAS package directory (the local mirror root
                // has no version-named subdir).
                let version = crate::native::flare_package::version_of(&package);
                let cmd = mirror_cmd(&dir, &package);
                let out =
                    crate::ssh::run_ssh_long_running(config, controller, cmd, dry_run).await?;
                let local = out.lines().last().unwrap_or_default().trim().to_string();
                if local.is_empty() {
                    bail!("flare local mirror produced no path on {}", controller.name);
                }
                package = local;
                crate::native::flare_package::write_marker(
                    config, controller, &dir, &package, &id, &version,
                )
                .await?;
            }
        }
    } else {
        // Dry-run: print the download and mirror commands, no SSH.
        let dl = download_cmd(config, recipe)?;
        package = crate::ssh::run_ssh_long_running(config, controller, dl, true).await?;
        package = parse_package_path(&package).unwrap_or_default();
        if let Some(dir) = local_dir_for(config, &cluster.name) {
            let dir = dir.replace('~', "$HOME");
            let cmd = mirror_cmd(&dir, &package);
            crate::ssh::run_ssh_long_running(config, controller, cmd, true).await?;
        }
    }
    let orbit_name = orbit_node_name(config, &cluster.name, controller);
    let run = run_cmd(config, recipe, &cluster.name, &orbit, &orbit_name, &package)?;
    crate::ssh::run_ssh_long_running(config, controller, run.clone(), dry_run).await?;
    if dry_run {
        println!("  [dry-run] {}: {run}", controller.name);
    } else {
        info!("FLARE group launched on {}", controller.name);
    }
    Ok(())
}

/// Post-launch: wait for the API (the package load takes many minutes) and
/// point the operator at the log command.
pub async fn after_serve(
    config: &Config,
    recipe: &Recipe,
    cluster: &crate::config::Cluster,
    recipe_name: &str,
    opts: &ServeOpts,
    follow: bool,
) -> Result<()> {
    let port = http_port(recipe);
    if !opts.no_wait_ready {
        crate::readiness::wait_engine(config, cluster, port, 30 * 60, recipe_name).await?;
    }
    if follow {
        info!(
            "FLARE serving on {}:{} — logs: sm121 logs --recipe {recipe_name}",
            cluster
                .nodes
                .first()
                .map(|n| n.ssh_ip.as_str())
                .unwrap_or("?"),
            port
        );
    }
    Ok(())
}

/// Stop the FLARE group on the controller.
pub async fn stop(
    config: &Config,
    recipe: &Recipe,
    cluster: &crate::config::Cluster,
    dry_run: bool,
) -> Result<()> {
    let controller = head(cluster)?;
    let cmd = stop_cmd(config, recipe, &orbit_cluster_for(config, &cluster.name))?;
    crate::ssh::run_ssh(config, controller, cmd, dry_run).await?;
    Ok(())
}

#[cfg(test)]
#[path = "flare_tests.rs"]
mod tests;