launcher/src/native/mod.rs

//! Native (head-coordinated) deploy: an engine's own cluster manager runs on
//! rank 0 and brings up the peers itself. Currently only DGPP
//! (`HawkBearPig/dgpp`, `engine = "dgpp"`, `deploy = "native"`).
//!
//! The launcher generates DGPP's cluster JSON (model + world size + engine
//! knobs + http port) and a sibling `.env` (node list + ssh user + ports) from
//! the recipe + selected cluster, uploads both to rank 0's DGPP checkout, and
//! runs `scripts/dgpp-cluster up --config <path>` there. Stop runs
//! `scripts/dgpp-cluster down`. Readiness is the uniform `/v1/models` probe on
//! the DGPP HTTP port (default 18080).

use anyhow::Result;
use tracing::info;

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

pub mod flare;
pub mod flare_package;
mod provision;

/// DGPP HTTP port for a recipe (its `[dgpp].http_port`, else 18080).
pub fn http_port(recipe: &Recipe) -> u16 {
    recipe
        .dgpp
        .as_ref()
        .and_then(|d| d.http_port)
        .unwrap_or(18080)
}

/// DGPP repo checkout path on the nodes (its `[dgpp].dir`, else `~/dgpp`).
fn dgpp_dir(d: &DgppConfig) -> String {
    d.dir.clone().unwrap_or_else(|| "~/dgpp".to_string())
}

/// Rewrite a leading `~` to `$HOME` so the remote shell expands it (a `~`
/// inside single quotes would be a literal directory named `~`). Callers wrap
/// the result in double quotes.
fn expand_home(path: &str) -> String {
    if path == "~" {
        "$HOME".to_string()
    } else if let Some(rest) = path.strip_prefix("~/") {
        format!("$HOME/{rest}")
    } else {
        path.to_string()
    }
}

/// Build the DGPP cluster JSON from the recipe. `model` and `world_size` come
/// from the recipe + launch node count; `http` from `[dgpp].http_port`; the
/// `engine` block is the recipe's `[dgpp.engine]` map serialized verbatim.
pub fn generate_cluster_json(recipe: &Recipe, world_size: usize) -> Result<String> {
    let d = recipe
        .dgpp
        .as_ref()
        .ok_or_else(|| anyhow::anyhow!("recipe has no [dgpp] block"))?;
    let mut root = serde_json::Map::new();
    root.insert(
        "model".into(),
        serde_json::Value::String(recipe.model.clone()),
    );
    root.insert(
        "world_size".into(),
        serde_json::Value::from(world_size as i64),
    );
    if let Some(r) = &d.release {
        root.insert("release".into(), serde_json::Value::String(r.clone()));
    }
    if !d.engine.is_empty() {
        let mut eng = serde_json::Map::new();
        for (k, v) in &d.engine {
            eng.insert(k.clone(), v.clone());
        }
        root.insert("engine".into(), serde_json::Value::Object(eng));
    }
    let mut http = serde_json::Map::new();
    http.insert(
        "bind_host".into(),
        serde_json::Value::String("0.0.0.0".into()),
    );
    http.insert(
        "port".into(),
        serde_json::Value::from(http_port(recipe) as i64),
    );
    root.insert("http".into(), serde_json::Value::Object(http));
    // `paths` always present (DGPP examples carry it); resident_cache empty =
    // DGPP default.
    let mut paths = serde_json::Map::new();
    paths.insert(
        "resident_cache".into(),
        serde_json::Value::String(d.resident_cache.clone().unwrap_or_default()),
    );
    root.insert("paths".into(), serde_json::Value::Object(paths));
    Ok(serde_json::to_string_pretty(&serde_json::Value::Object(
        root,
    ))?)
}

/// Build the DGPP site `.env` (allowlisted KEY=value lines DGPP's site_env.py
/// reads): the launch nodes (ssh_ip, rank order), the ssh user, and the ports.
pub fn generate_site_env(ssh_user: &str, cluster: &Cluster, recipe: &Recipe) -> String {
    let nodes = cluster
        .nodes
        .iter()
        .map(|n| n.ssh_ip.clone())
        .collect::<Vec<_>>()
        .join(" ");
    let port = http_port(recipe);
    format!(
        "DGPP_NODES={nodes}\nDGPP_SSH_USER={user}\nDGPP_HTTP_PORT={port}\nDGPP_FABRIC_PORT=29970\nDGPP_JOURNAL_PORT=29971\n",
        user = ssh_user,
    )
}

/// Write a file on `node` via a single-quoted heredoc (content has no line
/// equal to the delimiter — JSON/env never do). `path` is already `$HOME`-
/// expanded and is wrapped in double quotes so it expands on the remote shell.
async fn write_remote(
    config: &Config,
    node: &Node,
    path: &str,
    content: &str,
    dry_run: bool,
) -> Result<()> {
    let cmd = format!("cat > \"{path}\" <<'SM121DGPP'\n{content}\nSM121DGPP\n");
    crate::ssh::run_ssh(config, node, cmd, dry_run).await?;
    Ok(())
}

/// Serve a DGPP recipe: select the cluster, run before_serve hooks, upload the
/// generated cluster JSON + `.env` to rank 0, and launch `dgpp-cluster up`
/// there (head-coordinated — the manager brings up the peers).
pub async fn native_serve(
    config: &Config,
    recipe_name: &str,
    cluster_name: &str,
    node_name: Option<&str>,
    opts: &ServeOpts,
) -> Result<Cluster> {
    let dry_run = opts.dry_run;
    let recipe_obj = crate::config::load_recipe(recipe_name, &config.recipe_sources).await?;
    if recipe_obj.engine == "flare" {
        return flare_serve(
            config,
            &recipe_obj,
            recipe_name,
            cluster_name,
            node_name,
            opts,
        )
        .await;
    }
    let dgpp = recipe_obj
        .dgpp
        .clone()
        .ok_or_else(|| anyhow::anyhow!("engine=dgpp requires a [dgpp] block"))?;

    let (mut cluster, _nccl) = if dry_run {
        crate::orchestrate::prepare_cluster_dry_run(config, &recipe_obj, cluster_name, node_name)?
    } else {
        use crate::ends::{SshNodePinger, SshProber, ha_controller_from_config};
        use crate::orchestrate::{BootTiming, PrepareRequest, prepare_cluster};
        let power = opts
            .auto_power_on
            .then(|| ha_controller_from_config(config))
            .flatten();
        prepare_cluster(
            config,
            &recipe_obj,
            cluster_name,
            node_name,
            &PrepareRequest {
                prober: &SshProber,
                power: power.as_ref(),
                pinger: &SshNodePinger,
                auto_power_on: opts.auto_power_on,
                boot_timing: BootTiming::default(),
                force_occupied: false,
            },
        )
        .await?
    };
    // Trim to the launch set (world size = required_nodes).
    if let Some(n) = recipe_obj.required_nodes
        && (n as usize) < cluster.nodes.len()
    {
        cluster.nodes.truncate(n as usize);
    }
    let head = cluster
        .nodes
        .iter()
        .find(|n| n.rank == 0)
        .cloned()
        .ok_or_else(|| anyhow::anyhow!("no head node (rank 0) in cluster"))?;

    crate::before_serve::run(
        config,
        &cluster,
        &recipe_obj,
        recipe_name,
        opts.skip_before_serve,
        dry_run,
    )
    .await?;

    // Provision the DGPP checkout at the pinned rev on every launch node (the
    // native analog of building the docker image on demand): clone-at-pin +
    // release.sh + install, with a rev-marker skip. No-op when the recipe
    // declares no [dgpp].repo/rev (assume pre-installed). Nodes in parallel.
    let dir = expand_home(&dgpp_dir(&dgpp));
    let cfg_name = format!("cluster_sm121_{recipe_name}.json");
    let cfg_path = format!("{dir}/{cfg_name}");
    let env_path = format!("{dir}/.env");
    let json = generate_cluster_json(&recipe_obj, cluster.nodes.len())?;
    let env = generate_site_env(&config.ssh_user, &cluster, &recipe_obj);

    // Stage the generated cluster JSON on EVERY launch node — DGPP's
    // `dgpp-cluster install` needs a deployment config, and each node runs a
    // rank — plus the site `.env` on every node too (install validates
    // DGPP_NODES per node; the rank-0 manager reads it for `up`).
    for n in &cluster.nodes {
        crate::ssh::run_ssh(
            config,
            n,
            format!("mkdir -p \"{dir}\" && mkdir -p \"{dir}/log\""),
            dry_run,
        )
        .await?;
        write_remote(config, n, &cfg_path, &json, dry_run).await?;
        write_remote(config, n, &env_path, &env, dry_run).await?;
    }

    // Provision (clone-at-pin + release.sh + install --config <staged json>) on
    // RANK 0 ONLY — `dgpp-cluster install` is head-coordinated: it installs the
    // release on rank 0 AND SSHes to the peers (from DGPP_NODES) to install
    // there too, like the docker build-once + sync model. Running it on every
    // node would double-install the peers ("release already exists").
    let prov = provision::provision_node(config, &head, recipe_name, &dgpp, &cfg_path, dry_run);
    if let Err(e) = prov.await {
        anyhow::bail!("DGPP provisioning failed on head '{}': {e}", head.name);
    }

    // Launch on rank 0 only; the DGPP manager SSHes to the peers. Detached so
    // it survives the SSH session (nohup + log file).
    let up = format!(
        "cd \"{dir}\" && set -a && . ./.env && set +a && nohup scripts/dgpp-cluster up --config '{cfg_name}' > 'log/sm121-{recipe_name}.log' 2>&1 &"
    );
    info!("🚀 {} (dgpp, head {})", recipe_obj.name, head.name);
    crate::ssh::run_ssh(config, &head, up, dry_run).await?;

    if !dry_run {
        crate::runlock::mark_start(recipe_name, &config.recipe_sources, Some(&cluster.name))?;
        if !opts.no_wait_ready
            && let Err(e) = crate::readiness::wait_engine(
                config,
                &cluster,
                http_port(&recipe_obj),
                1800,
                recipe_name,
            )
            .await
        {
            let _ = crate::runlock::mark_stop(recipe_name, &config.recipe_sources);
            return Err(e);
        }
    }
    Ok(cluster)
}

/// Stop a DGPP run: `dgpp-cluster down` on rank 0 (the manager stops the
/// peers), then the caller closes the ledger.
/// FLARE path: prepare the cluster (same SSH/HA prepare as dgpp), then launch
/// the signed package through `flare orbit`.
async fn flare_serve(
    config: &Config,
    recipe: &Recipe,
    recipe_name: &str,
    cluster_name: &str,
    node_name: Option<&str>,
    opts: &ServeOpts,
) -> Result<Cluster> {
    let (cluster, _nccl) = if opts.dry_run {
        crate::orchestrate::prepare_cluster_dry_run(config, recipe, cluster_name, node_name)?
    } else {
        use crate::ends::{SshNodePinger, SshProber, ha_controller_from_config};
        use crate::orchestrate::{PrepareRequest, prepare_cluster};
        let power = opts
            .auto_power_on
            .then(|| ha_controller_from_config(config))
            .flatten();
        prepare_cluster(
            config,
            recipe,
            cluster_name,
            node_name,
            &PrepareRequest {
                prober: &SshProber,
                power: power.as_ref(),
                pinger: &SshNodePinger,
                auto_power_on: opts.auto_power_on,
                boot_timing: crate::orchestrate::BootTiming::default(),
                force_occupied: false,
            },
        )
        .await?
    };
    if recipe.flare.is_none() {
        anyhow::bail!("engine=flare requires a [flare] block");
    }
    // Cluster identification: the orbit name must come from this cluster's own
    // `[flare.clusters.<name>]` entry. Falling back to the flat `sparks`
    // default silently launches onto another machine's orbit (our clusters
    // share nodes, and `orbit run` creates the group wherever it is told).
    let flat = flare::orbit_cluster_for(config, &cluster.name);
    if let Some(orbit) = flare::orbit_cluster_paired(config, &cluster.name) {
        info!("FLARE: cluster '{}' → orbit '{}'", cluster.name, orbit);
    } else {
        anyhow::bail!(
            "FLARE cluster '{}' is not paired with an orbit cluster — add a \
             [flare.clusters.\"{}\"] block (orbit_cluster = \"…\") to the private \
             config. Without it the launch would silently target the flat \
             default '{}'.",
            cluster.name,
            cluster.name,
            flat
        );
    }
    // before_serve hooks (the global GPU clock lock among them): a Spark that
    // runs unclocked drifts and stalls a collective, so native launches get
    // the same prepare step docker/tmux lanes do.
    crate::before_serve::run(
        config,
        &cluster,
        recipe,
        cluster_name,
        opts.skip_before_serve,
        opts.dry_run,
    )
    .await?;
    flare::serve(config, recipe, &cluster, opts.dry_run).await?;
    // The run ledger is the single-launcher serialization point: a FLARE group
    // is as real as a container, so it must be recorded (otherwise a second
    // launcher sees "no open run" and starts a second lane on the same nodes).
    if !opts.dry_run {
        crate::runlock::mark_start(recipe_name, &config.recipe_sources, Some(&cluster.name))?;
    }
    Ok(cluster)
}

pub async fn native_stop(
    config: &Config,
    recipe: &Recipe,
    cluster: &Cluster,
    dry_run: bool,
) -> Result<()> {
    if recipe.engine == "flare" {
        return flare::stop(config, recipe, cluster, dry_run).await;
    }
    let Some(d) = recipe.dgpp.as_ref() else {
        anyhow::bail!("engine=dgpp requires a [dgpp] block to stop");
    };
    let dir = expand_home(&dgpp_dir(d));
    let head = cluster
        .nodes
        .iter()
        .find(|n| n.rank == 0)
        .ok_or_else(|| anyhow::anyhow!("no head node (rank 0) in cluster"))?;
    let down = format!("cd \"{dir}\" && scripts/dgpp-cluster down");
    crate::ssh::run_ssh(config, head, down, dry_run).await?;
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::config::{DgppConfig, Node};
    use serde_json::json;

    fn dgpp_recipe() -> Recipe {
        let mut r = crate::engine::tests::base_recipe();
        r.engine = "dgpp".into();
        r.name = "dgpp-qwen".into();
        r.model = "unsloth/Qwen3.8-27B-NVFP4".into();
        r.port = 18080;
        r.required_nodes = Some(1);
        r.deploy = crate::config::DeployTarget::Native;
        r.dgpp = Some(DgppConfig {
            http_port: Some(18080),
            engine: [
                ("kv_capacity".to_string(), json!(262144)),
                ("mtp".to_string(), json!(true)),
                (
                    "dflash_model".to_string(),
                    json!("z-lab/Qwen3.8-27B-DFlash2"),
                ),
            ]
            .into_iter()
            .collect(),
            ..Default::default()
        });
        r
    }

    #[test]
    fn cluster_json_carries_model_world_size_engine_and_http() {
        let r = dgpp_recipe();
        let json: serde_json::Value =
            serde_json::from_str(&generate_cluster_json(&r, 1).unwrap()).unwrap();
        assert_eq!(json["model"], "unsloth/Qwen3.8-27B-NVFP4");
        assert_eq!(json["world_size"], 1);
        assert_eq!(json["engine"]["kv_capacity"], 262144);
        assert_eq!(json["engine"]["mtp"], true);
        assert_eq!(json["engine"]["dflash_model"], "z-lab/Qwen3.8-27B-DFlash2");
        assert_eq!(json["http"]["port"], 18080);
        assert_eq!(json["http"]["bind_host"], "0.0.0.0");
        assert_eq!(json["paths"]["resident_cache"], "");
    }

    #[test]
    fn expand_home_rewrites_leading_tilde() {
        assert_eq!(expand_home("~/dgpp"), "$HOME/dgpp");
        assert_eq!(expand_home("~"), "$HOME");
        assert_eq!(expand_home("/srv/dgpp"), "/srv/dgpp");
    }

    #[test]
    fn site_env_lists_nodes_ssh_user_and_ports() {
        let r = dgpp_recipe();
        let cluster = Cluster {
            name: "dgx1".into(),
            master_ib_ip: None,
            nodes: vec![
                Node {
                    name: "dgx1".into(),
                    ssh_ip: "192.168.88.10".into(),
                    ib_ip: None,
                    rank: 0,
                    host_key: None,
                    entity: None,
                },
                Node {
                    name: "dgx2".into(),
                    ssh_ip: "192.168.88.14".into(),
                    ib_ip: None,
                    rank: 1,
                    host_key: None,
                    entity: None,
                },
            ],
            extra_entities: vec![],
        };
        let env = generate_site_env("kilork", &cluster, &r);
        assert!(env.contains("DGPP_NODES=192.168.88.10 192.168.88.14"));
        assert!(env.contains("DGPP_SSH_USER=kilork"));
        assert!(env.contains("DGPP_HTTP_PORT=18080"));
    }
}