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

/// Local package mirror dir for a cluster (per-cluster map first).
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())
}

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)));
    }
    if let Some(key) = config.flare.as_ref().and_then(|c| c.accept_key.clone()) {
        cmd.push_str(&format!(" --dangerously-approve-key {}", shq(&key)));
    }
    Ok(cmd)
}

/// `flare orbit run "<node>:<package>" --cluster <c> --name <n> … --detach`.
pub fn run_cmd(
    config: &Config,
    recipe: &Recipe,
    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)));
    }
    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);
    format!(
        "set -e; src=$(dirname {pkg}); dst={dir}; base=$(basename {pkg}); \
         mkdir -p \"$dst\"; \
         cp -f {pkg} \"$dst/$base\"; \
         for f in \"$src\"/*; do b=$(basename \"$f\"); [ \"$b\" = \"$base\" ] && continue; \
           if [ ! -f \"$dst/$b\" ] || [ \"$(stat -c%s \"$f\")\" != \"$(stat -c%s \"$dst/$b\")\" ]; then \
             cp -f \"$f\" \"$dst/$b\" || exit 1; fi; done; \
         echo \"$dst/$base\""
    )
}

/// 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> {
    stdout
        .lines()
        .rev()
        .map(str::trim)
        .find(|l| l.ends_with(".safetensors") && l.contains('/'))
        .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 dl = download_cmd(config, recipe)?;
    let package = crate::ssh::run_ssh(config, controller, dl, dry_run).await?;
    let package = parse_package_path(&package).unwrap_or_default();
    if !dry_run && 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.
    let mut package = package;
    if let Some(dir) = local_dir_for(config, &cluster.name) {
        let dir = dir.replace('~', "$HOME");
        let cmd = mirror_cmd(&dir, &package);
        let out = crate::ssh::run_ssh(config, controller, cmd, dry_run).await?;
        if !dry_run {
            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;
        }
    }
    let orbit_name = orbit_node_name(config, &cluster.name, controller);
    let run = run_cmd(config, recipe, &orbit, &orbit_name, &package)?;
    crate::ssh::run_ssh(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)]
mod tests {
    use super::*;
    use crate::config::FlareRecipe;

    fn cfg() -> Config {
        let c: Config = toml::from_str(
            "ssh_user = \"u\"\nssh_key_path = \"/k\"\n\n[flare]\naccept_key = \"key-123\"\naccept_publisher = \"local-development\"\norbit_cluster = \"sparks\"\n",
        )
        .expect("test config");
        c
    }

    fn recipe() -> Recipe {
        let mut r = crate::engine::tests::base_recipe();
        r.engine = "flare".into();
        r.model = "deepseek-ai/DeepSeek-V4.1-Flash-1017".into();
        r.flare = Some(FlareRecipe {
            package: "local-inference-lab/DeepSeek-V4.1-Flash-1017-FLARE".into(),
            revision: Some("gb10".into()),
            name: Some("deepseek1017".into()),
            http_port: Some(8000),
            serve_flags: vec![
                "--max-concurrency".into(),
                "16".into(),
                "--mtp-adaptive-verification".into(),
            ],
        });
        r
    }

    #[test]
    fn download_carries_revision_and_key() {
        let cmd = download_cmd(&cfg(), &recipe()).unwrap();
        assert!(cmd.starts_with("flare download "), "{cmd}");
        // shq() shell-quotes every value, so compare the quoted forms
        assert!(cmd.contains("--revision 'gb10'"), "{cmd}");
        assert!(cmd.contains("--dangerously-approve-key 'key-123'"), "{cmd}");
    }

    #[test]
    fn run_passes_flags_verbatim_and_the_controller_package() {
        let cmd = run_cmd(
            &cfg(),
            &recipe(),
            "sparks",
            "dgx5",
            "/pkgs/dsv41.safetensors",
        )
        .unwrap();
        assert!(
            cmd.contains("orbit run 'dgx5:/pkgs/dsv41.safetensors'"),
            "{cmd}"
        );
        assert!(cmd.contains("--cluster 'sparks'"), "{cmd}");
        assert!(cmd.contains("--name 'deepseek1017'"), "{cmd}");
        assert!(cmd.contains("--port 8000"), "{cmd}"); // recipe value wins
        assert!(cmd.contains("--detach"), "{cmd}");
        assert!(cmd.contains("'--mtp-adaptive-verification'"), "{cmd}");
        assert!(cmd.contains("--port 8000"), "{cmd}"); // recipe value wins
        assert!(
            cmd.contains("--dangerously-accept-publisher 'local-development'"),
            "{cmd}"
        );
    }

    #[test]
    fn orbit_node_names_follow_rank_order() {
        let mut c = cfg();
        c.flare.as_mut().unwrap().orbit_nodes = vec!["spark1".into()];
        let node = Node {
            name: "dgx5.lan".into(),
            ssh_ip: "10.0.0.1".into(),
            ib_ip: None,
            rank: 0,
            host_key: None,
            entity: None,
        };
        assert_eq!(orbit_node_name(&c, "dgx58-switch", &node), "spark1");
        let node1 = Node { rank: 1, ..node };
        // no entry for rank 1 → the cluster's own name is used
        assert_eq!(orbit_node_name(&c, "dgx58-switch", &node1), "dgx5.lan");
    }

    #[test]
    fn per_cluster_mapping_wins_over_the_flat_block() {
        // The same launcher config must serve FLARE on another cluster: the
        // per-cluster map supplies that cluster's orbit name and node list.
        let mut c = cfg();
        c.flare.as_mut().unwrap().clusters.insert(
            "prod4".into(),
            crate::config::flare_cluster::FlareCluster {
                orbit_cluster: Some("prod4".into()),
                orbit_nodes: vec!["gx10-dgx1".into()],
                controller_node: None,
                local_dir: Some("/data/flare".into()),
            },
        );
        assert_eq!(orbit_cluster_for(&c, "prod4"), "prod4");
        assert_eq!(orbit_cluster_for(&c, "dgx58-switch"), "sparks");
        assert_eq!(local_dir_for(&c, "prod4").as_deref(), Some("/data/flare"));
        let node = Node {
            name: "dgx1.lan".into(),
            ssh_ip: "10.0.0.1".into(),
            ib_ip: None,
            rank: 0,
            host_key: None,
            entity: None,
        };
        assert_eq!(orbit_node_name(&c, "prod4", &node), "gx10-dgx1");
        let cmd = stop_cmd(&c, &recipe(), &orbit_cluster_for(&c, "prod4")).unwrap();
        assert!(cmd.contains("--cluster 'prod4'"), "{cmd}");
    }

    #[test]
    fn stop_names_the_group_and_cluster() {
        let cmd = stop_cmd(&cfg(), &recipe(), "sparks").unwrap();
        assert_eq!(cmd, "flare orbit stop 'deepseek1017' --cluster 'sparks'");
    }

    #[test]
    fn preflight_checks_agent_and_pairing() {
        let cmd = preflight_cmd(&cfg(), "sparks", true);
        assert!(cmd.contains("is-active flare-agent"), "{cmd}");
        assert!(cmd.contains("orbit status --cluster 'sparks'"), "{cmd}");
        // a peer has no orbit file: binary + agent only
        assert!(!preflight_cmd(&cfg(), "sparks", false).contains("orbit status"));
    }

    #[test]
    fn package_path_is_taken_from_the_download_output() {
        let out = "Resolved package: org/repo@abc / pkg.safetensors\n\
                   Downloading model file: aa.safetensors...\n\
                   Downloading model file: bb.safetensors...\n\
                   /mnt/nas/hub/snapshots/abc/pkg.safetensors\n";
        assert_eq!(
            parse_package_path(out).as_deref(),
            Some("/mnt/nas/hub/snapshots/abc/pkg.safetensors")
        );
        assert_eq!(
            parse_package_path("Downloading model file: x.safetensors..."),
            None
        );
    }

    #[test]
    fn missing_block_is_an_error() {
        let mut r = recipe();
        r.flare = None;
        assert!(download_cmd(&cfg(), &r).is_err());
    }
}