Browse Source

stop: wait until containers are actually stopped (60s verify); add --no-wait

Alexander Korolev 1 tháng trước cách đây
mục cha
commit
a33cfa87a2
7 tập tin đã thay đổi với 127 bổ sung và 39 xóa
  1. 4 0
      AGENTS.md
  2. 6 0
      src/cli.rs
  3. 23 0
      src/docker/deploy.rs
  4. 3 6
      src/docker/mod.rs
  5. 66 0
      src/docker/stop.rs
  6. 24 32
      src/main.rs
  7. 1 1
      src/web/api.rs

+ 4 - 0
AGENTS.md

@@ -20,6 +20,10 @@ See [`DEVELOPMENT.md`](DEVELOPMENT.md) for architecture, module layout, and dev
 serve/stop/logs/health/doctor/exec/power/tmux/init-host-keys/config/web
 ```
 
+`stop` waits until the containers are actually stopped and removed (polls up
+60 s per node) — an immediate `serve` on the same nodes would otherwise race
+the dying containers' GPU release. `--no-wait` opts out.
+
 `sm121 web` — local axum + SSE dashboard: LED board (cluster-level health,
 non-blocking `extra_entities` as ⚙ badges), themes, serve/stop/power jobs
 with type-the-name power gate, live log-less jobs window fed by docker

+ 6 - 0
src/cli.rs

@@ -136,6 +136,12 @@ pub enum Commands {
         /// Power off cluster after stopping containers (requires HA config).
         #[arg(long = "off")]
         power_off: bool,
+
+        /// Don't wait until the containers are actually stopped and removed
+        /// (default: wait and verify, up to 60 s per node, so an immediate
+        /// re-serve on the same nodes doesn't race the dying containers).
+        #[arg(long)]
+        no_wait: bool,
     },
 
     /// Check cluster health

+ 23 - 0
src/docker/deploy.rs

@@ -16,6 +16,29 @@ pub struct DeploymentVariant {
 
 /// Discover running containers for a recipe across all unique nodes.
 /// Groups into deployment variants by (cluster, node-set).
+/// Builds the "multiple deployments found" error text: one line per cluster
+/// with its `--node` names, so the operator can disambiguate with `--node`.
+pub fn multi_deployment_error(recipe_name: &str, variants: &[DeploymentVariant]) -> String {
+    let mut msg = format!(
+        "Multiple deployments of '{}' found — use --node to select:\n",
+        recipe_name
+    );
+    for (i, v) in variants.iter().enumerate() {
+        let ns = v
+            .nodes
+            .iter()
+            .map(|n| format!("--node {}", n))
+            .collect::<Vec<_>>()
+            .join(" ");
+        msg.push_str(&format!(
+            "  {}: cluster '{}' with {}\n",
+            i + 1,
+            v.cluster,
+            ns
+        ));
+    }
+    msg
+}
 pub async fn discover_recipe_deployments(
     config: &Config,
     recipe_name: &str,

+ 3 - 6
src/docker/mod.rs

@@ -11,6 +11,7 @@ pub(crate) mod path_ref;
 pub use logs::show_logs;
 pub mod build_dir;
 pub mod stale;
+pub mod stop;
 pub mod sync;
 pub mod tail;
 use crate::{
@@ -280,6 +281,7 @@ pub async fn stop_cluster(
     cluster: &Cluster,
     recipe_name: &str,
     nodes: &[String],
+    wait: bool,
     dry_run: bool,
 ) -> Result<()> {
     info!(
@@ -310,12 +312,7 @@ pub async fn stop_cluster(
         run_ssh(
             config,
             node,
-            format!(
-                "docker ps -q --filter {} --filter {} | xargs -r docker stop; rm -rf {} 2>/dev/null; true",
-                shq(&format!("label=sm121.recipe={}", recipe_name)),
-                shq(&format!("label=sm121.cluster={}", cluster.name)),
-                shq(&format!("/tmp/sm121_{}_{}", recipe_name, node.rank)),
-            ),
+            stop::remote_command(recipe_name, &cluster.name, node.rank, wait),
             dry_run,
         )
         .await?;

+ 66 - 0
src/docker/stop.rs

@@ -0,0 +1,66 @@
+//! 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,
+    )
+}
+
+#[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"));
+    }
+}

+ 24 - 32
src/main.rs

@@ -140,8 +140,18 @@ async fn main() -> Result<()> {
             cluster,
             node,
             power_off,
+            no_wait,
         } => {
-            handle_stop(&config, &recipe_name, &cluster, &node, power_off, dry_run).await?;
+            handle_stop(
+                &config,
+                &recipe_name,
+                &cluster,
+                &node,
+                !no_wait,
+                power_off,
+                dry_run,
+            )
+            .await?;
         }
         Commands::Health {
             cluster: cluster_name,
@@ -325,21 +335,19 @@ async fn handle_logs(
     .await?;
     Ok(())
 }
-
 /// Handle `stop` command: stop containers and optionally power off.
 pub(crate) async fn handle_stop(
     config: &Config,
     recipe_name: &str,
     cluster: &str,
     node: &[String],
+    wait: bool,
     power_off: bool,
     dry_run: bool,
 ) -> Result<()> {
-    let (stopped_cluster, stopped_nodes): (Cluster, Vec<String>) = if !node.is_empty() {
-        let cluster = resolve_cluster_for_op(config, recipe_name, cluster).await?;
-        let nodes = node.to_vec();
-        docker::stop_cluster(config, &cluster, recipe_name, &nodes, dry_run).await?;
-        (cluster, nodes)
+    let (cluster, nodes): (Cluster, Vec<String>) = if !node.is_empty() {
+        let c = resolve_cluster_for_op(config, recipe_name, cluster).await?;
+        (c, node.to_vec())
     } else {
         let variants = docker::deploy::discover_recipe_deployments(config, recipe_name).await?;
         if variants.is_empty() {
@@ -347,42 +355,26 @@ pub(crate) async fn handle_stop(
         }
         if variants.len() == 1 {
             let v = &variants[0];
-            let cluster = select_cluster_by_name(config, &v.cluster)?;
-            let nodes = v.nodes.clone();
-            docker::stop_cluster(config, &cluster, recipe_name, &nodes, dry_run).await?;
-            (cluster, nodes)
+            let c = select_cluster_by_name(config, &v.cluster)?;
+            (c, v.nodes.clone())
         } else {
-            let mut msg = format!(
-                "Multiple deployments of '{}' found — use --node to select:\n",
-                recipe_name
-            );
-            for (i, v) in variants.iter().enumerate() {
-                let nodes = v
-                    .nodes
-                    .iter()
-                    .map(|n| format!("--node {}", n))
-                    .collect::<Vec<_>>()
-                    .join(" ");
-                msg.push_str(&format!(
-                    "  {}: cluster '{}' with {}\n",
-                    i + 1,
-                    v.cluster,
-                    nodes
-                ));
-            }
-            anyhow::bail!(msg);
+            anyhow::bail!(docker::deploy::multi_deployment_error(
+                recipe_name,
+                &variants
+            ));
         }
     };
+    docker::stop_cluster(config, &cluster, recipe_name, &nodes, wait, dry_run).await?;
 
     // Run recipe `after_stop` hooks (e.g. restore_dgx_services) on the
     // nodes that were actually stopped. Best-effort against the loaded
     // recipe; if the recipe file is unavailable, skip silently.
     if let Ok(recipe) = crate::config::load_recipe(recipe_name, &config.recipe_sources).await {
-        crate::after_stop::run(config, &stopped_cluster, &recipe, dry_run).await?;
+        crate::after_stop::run(config, &cluster, &recipe, dry_run).await?;
     }
 
     if power_off {
-        crate::power::power_off(config, &stopped_cluster, &stopped_nodes, dry_run).await?;
+        crate::power::power_off(config, &cluster, &nodes, dry_run).await?;
     }
     Ok(())
 }

+ 1 - 1
src/web/api.rs

@@ -107,7 +107,7 @@ pub(crate) async fn stop(
     super::jobs::spawn(app2, id.clone(), async move {
         let nodes = vec![node.unwrap_or_default()];
         let nodes = if nodes[0].is_empty() { vec![] } else { nodes };
-        crate::handle_stop(&config, &recipe, &cluster, &nodes, false, false).await
+        crate::handle_stop(&config, &recipe, &cluster, &nodes, true, false, false).await
     });
     Ok(Json(serde_json::json!({ "job": id })))
 }