소스 검색

feat: record cluster label on containers for auto-detection

Add --label sm121.cluster=<name> to all launched containers so that
stop and logs can look up which cluster a recipe was started on,
rather than guessing via auto-selection.

- build_docker_command: add sm121.cluster label
- resolve_cluster_by_container_label: new fn that finds the cluster
  by querying running containers' labels via SSH
- resolve_cluster_for_op: new helper for stop/logs that prefers
  label lookup, falls back to config-based auto for dry-run
Alexander Korolev 4 달 전
부모
커밋
6ff1e710aa
2개의 변경된 파일과 49개의 추가작업 그리고 4개의 파일을 삭제
  1. 24 0
      src/docker/mod.rs
  2. 25 4
      src/main.rs

+ 24 - 0
src/docker/mod.rs

@@ -224,6 +224,29 @@ pub async fn stop_cluster(
     Ok(())
 }
 
+/// Find the cluster that has running containers for a recipe by querying
+/// the `sm121.cluster` label.  Used by stop/logs when `--cluster auto`.
+pub async fn resolve_cluster_by_container_label<'c>(
+    config: &'c Config,
+    recipe_name: &str,
+) -> Result<&'c Cluster> {
+    for cluster in &config.clusters {
+        for node in &cluster.nodes {
+            let cmd = format!(
+                "docker ps -q --filter \"label=sm121.recipe={}\" \
+                 --format '{{{{.Names}}}}' | head -1",
+                recipe_name
+            );
+            if let Ok(output) = run_ssh(config, node, cmd, false).await
+                && !output.trim().is_empty()
+            {
+                return Ok(cluster);
+            }
+        }
+    }
+    anyhow::bail!("No running containers found for recipe '{}'", recipe_name);
+}
+
 #[allow(clippy::too_many_arguments)]
 async fn start_cluster(
     config: &Config,
@@ -391,6 +414,7 @@ fn build_docker_command(
     cmd.push(format!("--name sm121_{}_{}", recipe_name, node.rank));
     cmd.push(format!("--label sm121.recipe={}", recipe_name));
     cmd.push(format!("--label sm121.rank={}", node.rank));
+    cmd.push(format!("--label sm121.cluster={}", cluster.name));
     if recipe.distributed_backend == "ray" {
         env.push((
             "RAY_ADDRESS",

+ 25 - 4
src/main.rs

@@ -100,8 +100,7 @@ async fn main() -> Result<()> {
             cluster,
             follow,
         } => {
-            let recipe = load_recipe(&recipe_name).await?;
-            let (cluster, _) = resolve_cluster(&config, &recipe, &cluster)?;
+            let cluster = resolve_cluster_for_op(&config, &recipe_name, &cluster).await?;
             docker::show_logs(&config, &cluster, &recipe_name, follow, dry_run).await?;
         }
         Commands::Stop {
@@ -109,8 +108,7 @@ async fn main() -> Result<()> {
             cluster,
             node,
         } => {
-            let recipe = load_recipe(&recipe_name).await?;
-            let (cluster, _) = resolve_cluster(&config, &recipe, &cluster)?;
+            let cluster = resolve_cluster_for_op(&config, &recipe_name, &cluster).await?;
             docker::stop_cluster(&config, &cluster, &recipe_name, node.as_deref(), dry_run).await?;
         }
         Commands::Health {
@@ -198,3 +196,26 @@ async fn main() -> Result<()> {
 
     Ok(())
 }
+
+/// Resolve a cluster for container operations (stop/logs).
+/// When `auto`, tries live container label lookup first, then falls back
+/// to config-based auto-selection (for dry-run / offline scenarios).
+async fn resolve_cluster_for_op(
+    config: &Config,
+    recipe_name: &str,
+    cluster_name: &str,
+) -> Result<Cluster> {
+    if cluster_name != "auto" {
+        let recipe = load_recipe(recipe_name).await?;
+        let (cluster, _) = resolve_cluster(config, &recipe, cluster_name)?;
+        return Ok(cluster);
+    }
+    // Try live container label lookup first
+    if let Ok(c) = docker::resolve_cluster_by_container_label(config, recipe_name).await {
+        return Ok(c.clone());
+    }
+    // Fall back to config-based auto (dry-run / containers not running)
+    let recipe = load_recipe(recipe_name).await?;
+    let (cluster, _) = resolve_cluster(config, &recipe, "auto")?;
+    Ok(cluster)
+}