Selaa lähdekoodia

feat: add --node parameter to logs command

- Add --node parameter to Logs command to select specific node's logs
- In follow mode, logs from head node (rank 0) or specified node with rank 0
- In non-follow mode, logs from specified node or all nodes
- Extract log functions to src/docker/logs.rs to keep mod.rs under 500 lines
Alexander Korolev 4 kuukautta sitten
vanhempi
sitoutus
0ae454c259
5 muutettua tiedostoa jossa 124 lisäystä ja 76 poistoa
  1. 4 0
      src/cli.rs
  2. 87 0
      src/docker/logs.rs
  3. 17 70
      src/docker/mod.rs
  4. 4 4
      src/docker/tests.rs
  5. 12 2
      src/main.rs

+ 4 - 0
src/cli.rs

@@ -50,6 +50,10 @@ pub enum Commands {
 
         #[arg(long, default_value = "auto")]
         cluster: String,
+
+        /// Target a specific node within the cluster (default: head node if following, all nodes if not).
+        #[arg(long)]
+        node: Option<String>,
     },
 
     /// Stop a cluster

+ 87 - 0
src/docker/logs.rs

@@ -0,0 +1,87 @@
+//! Container log operations — show logs and follow logs for a recipe.
+
+use anyhow::{Context, Result};
+use tracing::info;
+
+use crate::config::{Cluster, Config, Node};
+use crate::ssh::{run_ssh_follow, run_ssh_with_retry};
+
+pub fn logs_tail_cmd(recipe_name: &str, rank: u32, tail: usize) -> String {
+    format!(
+        "docker ps -q --filter \"label=sm121.recipe={}\" --filter \"label=sm121.rank={}\" | xargs -r docker logs --timestamps --tail={} 2>&1",
+        recipe_name, rank, tail
+    )
+}
+
+pub fn logs_follow_cmd(recipe_name: &str) -> String {
+    format!(
+        "docker ps -q --filter \"label=sm121.recipe={}\" --filter \"label=sm121.rank=0\" | xargs -r docker logs -f --timestamps 2>&1",
+        recipe_name
+    )
+}
+
+pub async fn show_logs(
+    config: &Config,
+    cluster: &Cluster,
+    recipe_name: &str,
+    follow: bool,
+    node_name: Option<&str>,
+    dry_run: bool,
+) -> Result<()> {
+    // Resolve target nodes
+    let targets: Vec<&Node> = match node_name {
+        Some(name) => cluster.nodes.iter().filter(|n| n.name == *name).collect(),
+        None => cluster.nodes.iter().collect(),
+    };
+
+    if targets.is_empty() {
+        if let Some(n) = node_name {
+            let name = n.to_string();
+            anyhow::bail!(
+                fluent_i18n::t!("exec-node-not-found", { "name" => &name, "cluster" => &cluster.name })
+            );
+        }
+        anyhow::bail!(fluent_i18n::t!("exec-cluster-empty", { "name" => &cluster.name }));
+    }
+
+    if follow {
+        // Follow mode: only head node (rank 0) or specified node with rank 0
+        let target_node = targets
+            .iter()
+            .find(|n| n.rank == 0)
+            .or(targets.first())
+            .context("No nodes in cluster")?;
+        info!(
+            "{}",
+            fluent_i18n::t!("logs-follow-head", { "name" => &target_node.name })
+        );
+        let cmd = logs_follow_cmd(recipe_name);
+        run_ssh_follow(config, target_node, cmd, dry_run).await?;
+    } else {
+        let cn = cluster.name.clone();
+        info!(
+            "{}",
+            fluent_i18n::t!("logs-header", { "name" => cn, "count" => targets.len() })
+        );
+        let mut handles = Vec::new();
+        for node in targets {
+            let cmd = logs_tail_cmd(recipe_name, node.rank, 100);
+            let config = config.clone();
+            let node = (*node).clone();
+            let dry = dry_run;
+            let name = node.name.clone();
+            handles.push(tokio::spawn(async move {
+                (name, run_ssh_with_retry(&config, &node, cmd, dry, 1).await)
+            }));
+        }
+        for result in futures::future::join_all(handles).await {
+            if let Ok((name, Ok(output))) = result
+                && !output.trim().is_empty()
+            {
+                println!("{}", fluent_i18n::t!("logs-section", { "name" => &name }));
+                println!("{}", output);
+            }
+        }
+    }
+    Ok(())
+}

+ 17 - 70
src/docker/mod.rs

@@ -1,11 +1,13 @@
 pub mod deploy;
 pub mod health;
+pub mod logs;
 pub mod mods;
+pub use logs::show_logs;
 pub mod tail;
 use crate::{
     config::{Cluster, Config, NcclConfig, Node, Recipe, load_recipe, resolve_cluster},
     doctor::try_repaired_image,
-    ssh::{run_ssh, run_ssh_follow, run_ssh_with_retry},
+    ssh::{run_ssh, run_ssh_with_retry},
 };
 use anyhow::{Context, Result};
 use std::borrow::Cow;
@@ -13,6 +15,20 @@ use std::collections::HashSet;
 use tokio::time::{Duration, sleep};
 use tracing::info;
 
+pub(crate) fn shq(s: &str) -> String {
+    let mut out = String::with_capacity(s.len() + 2);
+    out.push('\'');
+    for c in s.chars() {
+        if c == '\'' {
+            out.push_str(r"'\''");
+        } else {
+            out.push(c);
+        }
+    }
+    out.push('\'');
+    out
+}
+
 pub async fn serve(
     config: &Config,
     recipe_name: &str,
@@ -109,75 +125,6 @@ pub async fn serve(
     .await?;
     Ok(cluster)
 }
-pub(crate) fn logs_tail_cmd(recipe_name: &str, rank: u32, tail: usize) -> String {
-    format!(
-        "docker ps -q --filter \"label=sm121.recipe={}\" --filter \"label=sm121.rank={}\" | xargs -r docker logs --timestamps --tail={} 2>&1",
-        recipe_name, rank, tail
-    )
-}
-pub(crate) fn logs_follow_cmd(recipe_name: &str) -> String {
-    format!(
-        "docker ps -q --filter \"label=sm121.recipe={}\" --filter \"label=sm121.rank=0\" | xargs -r docker logs -f --timestamps 2>&1",
-        recipe_name
-    )
-}
-pub(crate) fn shq(s: &str) -> String {
-    let mut out = String::with_capacity(s.len() + 2);
-    out.push('\'');
-    for c in s.chars() {
-        if c == '\'' {
-            out.push_str(r"'\''");
-        } else {
-            out.push(c);
-        }
-    }
-    out.push('\'');
-    out
-}
-pub async fn show_logs(
-    config: &Config,
-    cluster: &Cluster,
-    recipe_name: &str,
-    follow: bool,
-    dry_run: bool,
-) -> Result<()> {
-    if follow {
-        let head_node = cluster.nodes.first().context("No nodes in cluster")?;
-        info!(
-            "{}",
-            fluent_i18n::t!("logs-follow-head", { "name" => &head_node.name })
-        );
-        let cmd = logs_follow_cmd(recipe_name);
-        run_ssh_follow(config, head_node, cmd, dry_run).await?;
-    } else {
-        let cn = cluster.name.clone();
-        info!(
-            "{}",
-            fluent_i18n::t!("logs-header", { "name" => cn, "count" => cluster.nodes.len() })
-        );
-        let mut handles = Vec::new();
-        for node in &cluster.nodes {
-            let cmd = logs_tail_cmd(recipe_name, node.rank, 100);
-            let config = config.clone();
-            let node = node.clone();
-            let dry = dry_run;
-            let name = node.name.clone();
-            handles.push(tokio::spawn(async move {
-                (name, run_ssh_with_retry(&config, &node, cmd, dry, 1).await)
-            }));
-        }
-        for result in futures::future::join_all(handles).await {
-            if let Ok((name, Ok(output))) = result
-                && !output.trim().is_empty()
-            {
-                println!("{}", fluent_i18n::t!("logs-section", { "name" => &name }));
-                println!("{}", output);
-            }
-        }
-    }
-    Ok(())
-}
-/// Stops containers for a recipe on specific nodes of a cluster.
 pub async fn stop_cluster(
     config: &Config,
     cluster: &Cluster,

+ 4 - 4
src/docker/tests.rs

@@ -386,7 +386,7 @@ fn custom_init_port() {
 
 #[test]
 fn logs_tail_cmd_format() {
-    let cmd = super::logs_tail_cmd("test-sleep", 0, 100);
+    let cmd = super::logs::logs_tail_cmd("test-sleep", 0, 100);
     assert!(cmd.contains("docker ps -q"));
     assert!(cmd.contains("--filter \"label=sm121.recipe=test-sleep\""));
     assert!(cmd.contains("--filter \"label=sm121.rank=0\""));
@@ -398,8 +398,8 @@ fn logs_tail_cmd_format() {
 
 #[test]
 fn logs_tail_cmd_per_rank() {
-    let cmd0 = super::logs_tail_cmd("glm51", 0, 100);
-    let cmd2 = super::logs_tail_cmd("glm51", 2, 100);
+    let cmd0 = super::logs::logs_tail_cmd("glm51", 0, 100);
+    let cmd2 = super::logs::logs_tail_cmd("glm51", 2, 100);
     assert!(cmd0.contains("label=sm121.rank=0"));
     assert!(cmd2.contains("label=sm121.rank=2"));
     assert!(!cmd0.contains("label=sm121.rank=2"));
@@ -407,7 +407,7 @@ fn logs_tail_cmd_per_rank() {
 
 #[test]
 fn logs_follow_cmd_format() {
-    let cmd = super::logs_follow_cmd("test-sleep");
+    let cmd = super::logs::logs_follow_cmd("test-sleep");
     assert!(cmd.contains("xargs -r docker logs -f"));
     assert!(cmd.contains("--timestamps"));
     // Follow targets the head node only (rank=0).

+ 12 - 2
src/main.rs

@@ -97,16 +97,26 @@ async fn main() -> Result<()> {
             .await?;
             if follow && !dry_run {
                 tracing::info!("{}", fluent_i18n::t!("container-log-header"));
-                docker::show_logs(&config, &cluster, &recipe, true, dry_run).await?;
+                docker::show_logs(&config, &cluster, &recipe, true, node.as_deref(), dry_run)
+                    .await?;
             }
         }
         Commands::Logs {
             recipe: recipe_name,
             cluster,
             follow,
+            node,
         } => {
             let cluster = resolve_cluster_for_op(&config, &recipe_name, &cluster).await?;
-            docker::show_logs(&config, &cluster, &recipe_name, follow, dry_run).await?;
+            docker::show_logs(
+                &config,
+                &cluster,
+                &recipe_name,
+                follow,
+                node.as_deref(),
+                dry_run,
+            )
+            .await?;
         }
         Commands::Stop {
             recipe: recipe_name,