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, Recipe};
use crate::docker::{ServeOpts, shq};
use crate::native::flare_orbit::{controller_for, extra_flags_for};
// Re-exported: the driver is the FLARE entry point the rest of the launcher uses.
pub use crate::native::flare_orbit::{
    local_dir_for, orbit_cluster_for, orbit_cluster_paired, orbit_node_name,
};

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

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

/// A path for a shell command: a leading `$HOME` stays expandable (double
/// quotes) while everything else is single-quoted.
pub(crate) fn shell_path(p: &str) -> String {
    match p.strip_prefix("$HOME") {
        Some(rest) => format!("\"$HOME{}\"", rest.replace('"', "\\\"")),
        None => shq(p),
    }
}

/// Mirror a package into the controller's local dir and return the local
/// primary path.
///
/// The source is the package's **store root**, not `dirname(package)`: a
/// current (`format 2`) package lives in a blob store, where `blobs/<sha>`
/// holds the bytes once and `models/<version>/{model.safetensors,weights,…}`
/// are hardlinks into it. Two consequences, both measured on the 488 GB gb10
/// package:
///
/// * the copy must keep hardlinks (`rsync -aH`) — a hardlink-blind copy
///   expands every link and duplicates the whole package instead of mirroring
///   it;
/// * it must not `--delete` the destination root, which holds every other
///   prepared package in the same store.
///
/// A flat (legacy) package is a self-contained directory, so the destination
/// is dedicated to it and `--delete` still prunes stale companions.
pub fn mirror_cmd(local_dir: &str, remote_pkg: &str) -> String {
    use crate::native::flare_package::{Layout, layout_of, mirror_destination, mirror_source};
    let src = mirror_source(remote_pkg);
    let dst = shell_path(local_dir);
    let dest_pkg = shell_path(&mirror_destination(local_dir, remote_pkg));
    let (flags, prune) = match layout_of(remote_pkg) {
        Layout::BlobStore => ("-aH", ""),
        Layout::Flat => ("-aH", " --delete"),
    };
    format!(
        "set -e; src={}; dst={}; \
         if [ \"$src\" = \"$dst\" ]; then echo {dest_pkg}; exit 0; fi; \
         mkdir -p \"$dst\"; \
         rsync {flags}{prune} --exclude '.sm121-flare-package.json' \"$src\"/ \"$dst\"/; \
         echo {dest_pkg}",
        shq(&src),
        dst,
    )
}

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

/// 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 = controller_for(config, 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 prepared package from the Flare
            // store instead of failing after a successful download.
            if package.is_empty() {
                package = crate::native::flare_package::cached_package(config, recipe, controller)
                    .await?
                    .unwrap_or_default();
            }
            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?;
        // The path is printed by a real download; in a dry run there is none,
        // so show the shape (a store-relative `models/<version>/model.safetensors`,
        // which is what the mirror shape depends on).
        package = parse_package_path(&package)
            .unwrap_or_else(|| "<package-store>/models/<version>/model.safetensors".into());
        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<()> {
    // `orbit stop` must reach the orbit's controller, which the config may pin
    // to a node other than our rank 0.
    let controller = controller_for(config, 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;