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;

/// Builds the per-node remote stop command.
///
/// `wait = true` (default): after `docker stop`, poll every 2 s (30 times,
/// 60 s max) until `docker ps` shows no container with the `sm121.recipe` /
/// `sm121.cluster` labels in the running state; if any remain, exit 1 —
/// `run_ssh` surfaces that as an error.
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 -q {filter}); \
            [ -z \"$n\" ] && break; \
            sleep 2; \
        done; \
        n=$(docker ps -q {filter}); \
        if [ -n \"$n\" ]; then \
            echo \"sm121 stop: containers still running after 60 s: $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,
    )
}

/// 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 !nodes.is_empty() {
        let cluster = crate::resolve_cluster_for_op(config, recipe_name, cluster_name).await?;
        (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
            ));
        }
    };
    crate::docker::stop_cluster(config, &cluster, recipe_name, &nodes, wait, dry_run).await?;
    // 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("seq 1 30"));
        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("exit 1"));
        assert!(cmd.contains("/tmp/sm121_r_1"));
    }
}