launcher/src/docker/stop.rs

//! Builder for the `sm121 stop` per-node remote command.
//!
//! `docker stop` only returns once containers reach the `Exited` state; the
//! `--rm` removal and the GPU/CUDA-context teardown can lag behind it.
//! Re-serving on the same node in that window makes the new container race
//! the dying one (immediate worker death). By default the command therefore
//! verifies — polls up to 60 s — that no labeled container is running before
//! reporting success; `--no-wait` restores the legacy fire-and-forget behavior.

use super::shq;
use crate::config::{Cluster, Config, Node};
use tracing::info;

/// Builds the per-node remote stop command.
///
/// `wait = true` (default): after `docker stop`, poll until `docker ps` shows
/// no container with the `sm121.recipe` / `sm121.cluster` labels in ANY state
/// (running or exited-but-not-yet-removed) for 30 s, then explicitly `rm` the
/// stragglers and poll running-only for another 30 s. Reporting success while
/// any labeled container still existed is what made an immediately following
/// `serve` see the node as `occupied` and bail out.
pub fn remote_command(recipe_name: &str, cluster_name: &str, rank: u32, wait: bool) -> String {
    let filter = format!(
        "--filter {} --filter {}",
        shq(&format!("label=sm121.recipe={}", recipe_name)),
        shq(&format!("label=sm121.cluster={}", cluster_name)),
    );
    let wait_block = if wait {
        "\
for _ in $(seq 1 30); do \
    n=$(docker ps -aq {filter}); \
    [ -z \"$n\" ] && break; \
    sleep 1; \
done; \
docker ps -aq {filter} | xargs -r docker rm >/dev/null 2>&1; \
for _ in $(seq 1 30); do \
    n=$(docker ps -q {filter}); \
    [ -z \"$n\" ] && break; \
    sleep 1; \
done; \
n=$(docker ps -aq {filter}); \
if [ -n \"$n\" ]; then \
    echo \"sm121 stop: containers still present after wait: $n\" >&2; \
    exit 1; \
fi; \
n=$(docker ps -q {filter}); \
if [ -n \"$n\" ]; then \
    echo \"sm121 stop: containers still running after wait: $n\" >&2; \
    exit 1; \
fi; "
            .to_string()
    } else {
        String::new()
    };
    format!(
        "docker ps -q {filter} | xargs -r docker stop; rm -rf {tmp} 2>/dev/null; {wait_block}true",
        tmp = shq(&format!("/tmp/sm121_{}_{}", recipe_name, rank)),
        wait_block = wait_block,
    )
}

/// Find the cluster containing all named nodes (for `--node` without
/// `--cluster`): the only cluster whose node set is a superset of the given
/// names. Ambiguous (nodes spread across several clusters) → error naming
/// the candidates.
fn cluster_for_node(config: &Config, nodes: &[String]) -> anyhow::Result<Cluster> {
    let mut matches: Vec<&Cluster> = config
        .clusters
        .iter()
        .filter(|c| nodes.iter().all(|n| c.nodes.iter().any(|cn| cn.name == *n)))
        .collect();
    match matches.len() {
        0 => anyhow::bail!("Node(s) {} not found in any cluster", nodes.join(", ")),
        1 => Ok(matches.remove(0).clone()),
        n => anyhow::bail!(
            "Nodes {} are in {} clusters ({}); add --cluster",
            nodes.join(", "),
            n,
            matches
                .iter()
                .map(|c| c.name.as_str())
                .collect::<Vec<_>>()
                .join(", ")
        ),
    }
}
/// Stop the recipe's containers across `nodes` (all cluster nodes when the
/// list is empty) and verify they are really gone on EVERY node: each node
/// runs the verifying remote command, all nodes are attempted even when a
/// sibling fails, and any failure names the straggler nodes.
pub async fn stop_cluster(
    config: &Config,
    cluster: &Cluster,
    recipe_name: &str,
    nodes: &[String],
    wait: bool,
    dry_run: bool,
) -> Result<(), anyhow::Error> {
    info!(
        "{}",
        fluent_i18n::t!("cluster-stopping", { "name" => &cluster.name })
    );
    let nodes: Vec<&Node> = if nodes.is_empty() {
        cluster.nodes.iter().collect()
    } else {
        let mut found: Vec<&Node> = Vec::new();
        let mut not_found: Vec<String> = Vec::new();
        for nn in nodes {
            if let Some(node) = cluster.nodes.iter().find(|n| n.name == *nn) {
                found.push(node);
            } else {
                not_found.push(nn.clone());
            }
        }
        if !not_found.is_empty() {
            let names = not_found.join(", ");
            anyhow::bail!(
                fluent_i18n::t!("exec-node-not-found", { "name" => &names, "cluster" => &cluster.name })
            );
        }
        found
    };
    let mut handles = Vec::new();
    for node in &nodes {
        let config = config.clone();
        let node = (*node).clone();
        let cmd = remote_command(recipe_name, &cluster.name, node.rank, wait);
        let dry = dry_run;
        let name = node.name.clone();
        handles.push(crate::locale::spawn_with_locale(async move {
            (name, super::run_ssh(&config, &node, cmd, dry).await)
        }));
    }
    let mut failed = Vec::new();
    for r in futures::future::join_all(handles).await {
        match r {
            Ok((_, Ok(_))) => {}
            Ok((name, Err(e))) => failed.push(format!("{name}: {e}")),
            // a panicked node task is an unverified node
            Err(e) => failed.push(format!("task: {e}")),
        }
    }
    if failed.is_empty() {
        info!("{}", fluent_i18n::t!("cluster-all-stopped"));
        Ok(())
    } else {
        let nodes = failed.join(", ");
        anyhow::bail!(fluent_i18n::t!("stop-not-verified", { "nodes" => &nodes }))
    }
}

/// Handle the `stop` command: stop containers and optionally power off.
///
/// Waits for the containers to actually stop (see [`remote_command`]), then
/// runs the recipe's `after_stop` hooks and the optional power off.
pub async fn handle(
    config: &crate::config::Config,
    recipe_name: &str,
    cluster_name: &str,
    nodes: &[String],
    wait: bool,
    power_off: bool,
    dry_run: bool,
) -> Result<(), anyhow::Error> {
    let (cluster, nodes): (crate::config::Cluster, Vec<String>) =
        if !cluster_name.is_empty() && cluster_name != "auto" {
            // Explicit cluster selection (exact or prefix): stop that cluster's
            // deployment directly — no discovery, no multi-deployment ambiguity.
            let cluster = crate::config::select_cluster_by_name(config, cluster_name)?;
            (cluster, nodes.to_vec())
        } else if !nodes.is_empty() {
            // --node without --cluster: find the cluster that contains the node.
            let cluster = cluster_for_node(config, nodes)?;
            (cluster, nodes.to_vec())
        } else {
            let variants =
                crate::docker::deploy::discover_recipe_deployments(config, recipe_name).await?;
            if variants.is_empty() {
                anyhow::bail!("No running containers found for recipe '{}'", recipe_name);
            }
            if variants.len() == 1 {
                let v = &variants[0];
                let cluster = crate::config::select_cluster_by_name(config, &v.cluster)?;
                (cluster, v.nodes.clone())
            } else {
                anyhow::bail!(crate::docker::deploy::multi_deployment_error(
                    recipe_name,
                    &variants
                ));
            }
        };
    stop_cluster(config, &cluster, recipe_name, &nodes, wait, dry_run).await?;
    // the stop is verified (every node ran the verifying command) — close
    // the recipe's open run in the ledger
    if !dry_run {
        crate::runlock::mark_stop(recipe_name, &config.recipe_sources)?;
    }
    // Run recipe `after_stop` hooks (e.g. restore_dgx_services) on the stopped nodes; best-effort.
    if let Ok(recipe) = crate::config::load_recipe(recipe_name, &config.recipe_sources).await {
        crate::after_stop::run(config, &cluster, &recipe, dry_run).await?;
    }
    if power_off {
        crate::power::power_off(config, &cluster, &nodes, dry_run).await?;
    }
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::remote_command;

    #[test]
    fn default_includes_verification_loop() {
        let cmd = remote_command("r", "c", 0, true);
        assert!(cmd.contains("xargs -r docker stop"));
        assert!(cmd.contains("docker ps -aq"));
        assert!(cmd.contains("xargs -r docker rm"));
        assert!(cmd.contains("exit 1"));
        assert!(cmd.ends_with("true"));
    }

    #[test]
    fn no_wait_omits_verification_loop() {
        let cmd = remote_command("r", "c", 1, false);
        assert!(!cmd.contains("seq 1 30"));
        assert!(!cmd.contains("docker rm"));
        assert!(!cmd.contains("exit 1"));
        assert!(cmd.contains("/tmp/sm121_r_1"));
    }
}