Эх сурвалжийг харах

feat: add --rm to containers by default, --keep flag to persist

- Add --rm to docker run (auto-remove on stop) by default
- Add --keep flag on serve command to suppress --rm and persist
  containers after stop (useful for debugging)
- Default behavior: containers auto-clean on stop, no manual
  docker rm needed
- Update CLI, all Docker plumbing, and tests
- 26 tests pass, clippy clean, fmt clean
Alexander Korolev 4 сар өмнө
parent
commit
4d3c297c50
4 өөрчлөгдсөн 54 нэмэгдсэн , 4 устгасан
  1. 4 0
      src/cli.rs
  2. 11 2
      src/docker/mod.rs
  3. 33 0
      src/docker/tests.rs
  4. 6 2
      src/main.rs

+ 4 - 0
src/cli.rs

@@ -22,6 +22,10 @@ pub enum Commands {
 
         #[arg(long, default_value = "auto")]
         cluster: String,
+
+        /// Keep containers after stop (default: auto-remove on stop)
+        #[arg(long, default_value = "false")]
+        keep: bool,
     },
 
     /// Показать логи контейнеров

+ 11 - 2
src/docker/mod.rs

@@ -11,6 +11,7 @@ pub async fn serve(
     config: &Config,
     recipe_name: &str,
     cluster_name: &str,
+    keep: bool,
     dry_run: bool,
 ) -> Result<()> {
     let recipe = load_recipe(recipe_name).await?;
@@ -47,6 +48,7 @@ pub async fn serve(
         recipe_name,
         &recipe,
         cache_dir,
+        keep,
         dry_run,
     )
     .await?;
@@ -220,6 +222,7 @@ pub async fn health_check(config: &Config, cluster: &Cluster, dry_run: bool) ->
 }
 
 /// Starts containers on all nodes in parallel.
+#[allow(clippy::too_many_arguments)]
 async fn start_cluster(
     config: &Config,
     cluster: &Cluster,
@@ -227,6 +230,7 @@ async fn start_cluster(
     recipe_name: &str,
     recipe: &Recipe,
     cache_dir: String,
+    keep: bool,
     dry_run: bool,
 ) -> Result<()> {
     let nnodes = cluster.nodes.len() as u32;
@@ -258,7 +262,7 @@ async fn start_cluster(
                 let _ = run_ssh(&config, &node, cleanup, dry).await;
             }
             let cmd = build_docker_command(
-                &cluster, &nccl, &recipe, &rn, &node, is_head, pp_size, nnodes, &cd,
+                &cluster, &nccl, &recipe, &rn, &node, is_head, pp_size, nnodes, &cd, keep,
             )
             .join(" ");
             let result = run_ssh_with_retry(&config, &node, cmd, dry, 2).await;
@@ -306,11 +310,16 @@ fn build_docker_command(
     pp_size: u32,
     nnodes: u32,
     cache_dir: &str,
+    keep: bool,
 ) -> Vec<String> {
     let is_sglang = recipe.engine == "sglang";
 
+    let mut flags = vec!["--gpus all", "--ipc=host", "--privileged", "--network=host"];
+    if !keep {
+        flags.push("--rm");
+    }
     let mut cmd = vec![
-        "docker run -d --gpus all --ipc=host --privileged --network=host".to_string(),
+        format!("docker run -d {}", flags.join(" ")),
         format!("-e NCCL_IB_DISABLE={}", nccl.ib_disable),
         format!("-e NCCL_IB_HCA={}", nccl.ib_hca),
         format!("-e NCCL_SOCKET_IFNAME={}", nccl.socket_ifname),

+ 33 - 0
src/docker/tests.rs

@@ -62,6 +62,23 @@ fn run(recipe: &Recipe, node: &Node, head: bool) -> String {
         2,
         2,
         "/home/user/.cache/huggingface",
+        false,
+    )
+    .join(" ")
+}
+
+fn run_keep(recipe: &Recipe, node: &Node, head: bool) -> String {
+    build_docker_command(
+        &tcls(),
+        &tnccl(),
+        recipe,
+        "r",
+        node,
+        head,
+        2,
+        2,
+        "/home/user/.cache/huggingface",
+        true,
     )
     .join(" ")
 }
@@ -85,6 +102,7 @@ fn mp() {
     assert!(cmd.contains("--master-addr 1"));
     assert!(cmd.contains("--distributed-executor-backend mp"));
     assert!(!cmd.contains("RAY_ADDRESS"));
+    assert!(cmd.contains("--rm"), "default --rm should be present");
 }
 
 #[test]
@@ -158,6 +176,7 @@ fn srun(recipe: &Recipe, node: &Node, head: bool) -> String {
         4,
         4,
         "/home/user/.cache/huggingface",
+        false,
     )
     .join(" ")
 }
@@ -186,6 +205,10 @@ fn sglang_head() {
     assert!(cmd.contains("--port 8000"));
     assert!(cmd.contains("-e HF_HOME=/cache/huggingface"));
     assert!(cmd.contains("-v /home/user/.cache/huggingface:/cache/huggingface"));
+    assert!(
+        cmd.contains("--rm"),
+        "SGLang default --rm should be present"
+    );
 }
 
 #[test]
@@ -211,3 +234,13 @@ fn sglang_entrypoint_model_path_order() {
     let port = cmd.find("--port 8000").unwrap();
     assert!(tp < port, "--tp before --port");
 }
+
+#[test]
+fn keep_disables_rm() {
+    let cmd = run_keep(&trec(), &tcls().nodes[0], true);
+    assert!(!cmd.contains("--rm"), "--keep should suppress --rm");
+    assert!(cmd.contains("docker run -d"));
+    assert!(cmd.contains("--gpus all"));
+    assert!(cmd.contains("-e HF_HOME=/cache/huggingface"));
+    assert!(cmd.contains("-v /home/user/.cache/huggingface:/cache/huggingface"));
+}

+ 6 - 2
src/main.rs

@@ -23,8 +23,12 @@ async fn main() -> Result<()> {
     let dry_run = cli.dry_run;
 
     match cli.command {
-        Commands::Serve { recipe, cluster } => {
-            docker::serve(&config, &recipe, &cluster, dry_run).await?;
+        Commands::Serve {
+            recipe,
+            cluster,
+            keep,
+        } => {
+            docker::serve(&config, &recipe, &cluster, keep, dry_run).await?;
         }
         Commands::Logs { recipe, follow } => {
             let (cluster, recipe) = resolve_cluster_for_recipe(&config, &recipe).await?;