Переглянути джерело

feat: section aliases for exec subcommand

Add a `[exec_aliases]` table to config.toml that maps short names to
full shell commands, usable via `cargo run exec --alias <name>`.

Example config:

    [exec_aliases]
    gpu-lock = "sudo nvidia-smi -lgc 300,2300"
    shutdown = "sudo shutdown now"
    reboot   = "sudo reboot"

Then:

    cargo run exec --cluster 4nodes --alias gpu-lock
    cargo run exec --cluster 4nodes --alias shutdown

CLI rules:
- `--alias <name>` is mutually exclusive with a trailing command.
- Unknown alias errors list the configured aliases.
- Either `--alias` or a trailing command must be given.

Implementation:
- New module `src/pre_serve.rs` houses the pre-serve related types
  (`OnFailure`, `PreServeCommand`, `BeforeServeTable`) so
  `config.rs` stays under the 500-line limit (AGENTS.md rule).
- New `Config::exec_aliases: HashMap<String, String>` for the
  alias table.
- New `exec::resolve_command` function picks between alias and
  trailing-command inputs, with informative errors.
- 6 new tests in exec.rs cover the alias-vs-args logic.
Alexander Korolev 4 місяців тому
батько
коміт
ce40c23b7a
7 змінених файлів з 212 додано та 55 видалено
  1. 22 0
      AGENTS.md
  2. 1 0
      src/before_serve.rs
  3. 6 1
      src/cli.rs
  4. 11 49
      src/config.rs
  5. 111 3
      src/exec.rs
  6. 5 2
      src/main.rs
  7. 56 0
      src/pre_serve.rs

+ 22 - 0
AGENTS.md

@@ -117,6 +117,28 @@ Commands are run sequentially on each node (so order matters) but nodes
 run in parallel.  Use `--skip-before-serve` to bypass for a one-off
 serve.
 
+### `exec` aliases (`[exec_aliases]`)
+
+Map short names to full shell commands for use with
+`cargo run exec --alias`.  Defined in `config.toml`:
+
+```toml
+[exec_aliases]
+gpu-lock = "sudo nvidia-smi -lgc 300,2300"
+shutdown = "sudo shutdown now"
+reboot = "sudo reboot"
+```
+
+Then:
+
+```bash
+cargo run exec --cluster 4nodes --alias gpu-lock
+cargo run exec --cluster 4nodes --alias shutdown
+```
+
+`--alias <name>` and a trailing command are mutually exclusive.  If the
+alias name is unknown, the error message lists the configured aliases.
+
 ## recipes/
 
 TOML format. Directory name is intentionally misspelled. Each recipe lives in its own subdirectory:

+ 1 - 0
src/before_serve.rs

@@ -178,6 +178,7 @@ mod tests {
             hf_cache_dir: None,
             init_port: 25000,
             before_serve,
+            exec_aliases: HashMap::new(),
         }
     }
 

+ 6 - 1
src/cli.rs

@@ -105,8 +105,13 @@ pub enum Commands {
         #[arg(long)]
         node: Option<String>,
 
+        /// Name of an alias from `[exec_aliases]` in config.toml.
+        /// Mutually exclusive with a trailing command.
+        #[arg(long)]
+        alias: Option<String>,
+
         /// Command to run (everything after `--`).
-        #[arg(trailing_var_arg = true, required = true)]
+        #[arg(trailing_var_arg = true)]
         command: Vec<String>,
     },
 }

+ 11 - 49
src/config.rs

@@ -3,6 +3,8 @@ use serde::Deserialize;
 use std::collections::HashMap;
 use tracing::info;
 
+pub use crate::pre_serve::{BeforeServeTable, OnFailure, PreServeCommand};
+
 #[derive(Clone, Deserialize)]
 pub struct Node {
     pub name: String,
@@ -28,48 +30,6 @@ pub struct Cluster {
     pub nodes: Vec<Node>,
 }
 
-/// Behaviour on a non-zero exit code for a [`PreServeCommand`].
-#[derive(Clone, Copy, Debug, PartialEq, Eq, Default, Deserialize)]
-#[serde(rename_all = "lowercase")]
-pub enum OnFailure {
-    /// Abort the serve (default).
-    #[default]
-    Abort,
-    /// Log a warning but continue.
-    Warn,
-    /// Continue silently.
-    Continue,
-}
-
-/// One entry in a `before_serve` block.  Either a bare command string
-/// (defaults to `on_failure = abort`) or a table with explicit `on_failure`.
-#[derive(Clone, Debug, Deserialize)]
-#[serde(untagged)]
-pub enum PreServeCommand {
-    Bare(String),
-    Configured {
-        command: String,
-        #[serde(default)]
-        on_failure: OnFailure,
-    },
-}
-
-impl PreServeCommand {
-    pub fn command(&self) -> &str {
-        match self {
-            PreServeCommand::Bare(c) => c,
-            PreServeCommand::Configured { command, .. } => command,
-        }
-    }
-
-    pub fn on_failure(&self) -> OnFailure {
-        match self {
-            PreServeCommand::Bare(_) => OnFailure::default(),
-            PreServeCommand::Configured { on_failure, .. } => *on_failure,
-        }
-    }
-}
-
 /// Top-level config containing cluster definitions and per-cluster
 /// pre-serve commands.
 #[derive(Clone, Deserialize)]
@@ -88,14 +48,10 @@ pub struct Config {
     /// Per-cluster pre-serve commands.  Looked up by cluster name.
     #[serde(default)]
     pub before_serve: HashMap<String, BeforeServeTable>,
-}
-
-/// TOML sub-table under `[before_serve."<cluster>"]`.  Currently only
-/// carries `commands`; reserved for future per-cluster tuning.
-#[derive(Clone, Debug, Default, Deserialize)]
-pub struct BeforeServeTable {
+    /// Command aliases for `cargo run exec --alias`.  Keyed by alias
+    /// name, value is the shell command to run.
     #[serde(default)]
-    pub commands: Vec<PreServeCommand>,
+    pub exec_aliases: HashMap<String, String>,
 }
 
 fn default_init_port() -> u16 {
@@ -400,6 +356,7 @@ nodes = [
             hf_cache_dir: None,
             init_port: 25000,
             before_serve: HashMap::new(),
+            exec_aliases: HashMap::new(),
         };
         config.nccl.insert("test-cluster".to_string(), test_nccl());
 
@@ -418,6 +375,7 @@ nodes = [
             hf_cache_dir: None,
             init_port: 25000,
             before_serve: HashMap::new(),
+            exec_aliases: HashMap::new(),
         };
         config.nccl.insert("test-cluster".to_string(), test_nccl());
 
@@ -436,6 +394,7 @@ nodes = [
             hf_cache_dir: None,
             init_port: 25000,
             before_serve: HashMap::new(),
+            exec_aliases: HashMap::new(),
         };
         let result = select_cluster(&config, "auto");
         assert!(result.is_err());
@@ -451,6 +410,7 @@ nodes = [
             hf_cache_dir: None,
             init_port: 25000,
             before_serve: HashMap::new(),
+            exec_aliases: HashMap::new(),
         };
         let result = select_cluster(&config, "nonexistent");
         assert!(result.is_err());
@@ -468,6 +428,7 @@ nodes = [
             hf_cache_dir: None,
             init_port: 25000,
             before_serve: HashMap::new(),
+            exec_aliases: HashMap::new(),
         };
         let cluster = select_cluster_by_name(&config, "test-cluster").unwrap();
         assert_eq!(cluster.name, "test-cluster");
@@ -484,6 +445,7 @@ nodes = [
             hf_cache_dir: None,
             init_port: 25000,
             before_serve: HashMap::new(),
+            exec_aliases: HashMap::new(),
         };
         let result = select_cluster_by_name(&config, "no-such-cluster");
         assert!(result.is_err());

+ 111 - 3
src/exec.rs

@@ -1,4 +1,4 @@
-use anyhow::{Context, Result};
+use anyhow::{Context, Result, anyhow};
 use tracing::{info, warn};
 
 use crate::config::{Cluster, Config, Recipe, select_cluster};
@@ -98,8 +98,116 @@ pub async fn run(
     Ok(())
 }
 
+/// Resolve the command to execute for `cargo run exec`.
+///
+/// - `--alias <name>` looks up `name` in `config.exec_aliases` and returns
+///   the configured shell command.  Errors if not found.
+/// - Otherwise joins the trailing positional args into a single command
+///   line (must be non-empty).
+///
+/// Exactly one of `alias_name` or `command_args` must be provided.
+pub fn resolve_command(
+    config: &Config,
+    alias_name: Option<&str>,
+    command_args: &[String],
+) -> Result<String> {
+    let has_alias = alias_name.is_some();
+    let has_args = !command_args.is_empty();
+    match (has_alias, has_args) {
+        (true, false) => {
+            let name = alias_name.expect("checked by has_alias");
+            config.exec_aliases.get(name).cloned().ok_or_else(|| {
+                let mut known: Vec<String> = config.exec_aliases.keys().cloned().collect();
+                known.sort();
+                if known.is_empty() {
+                    anyhow!("Unknown alias @{} (no aliases configured)", name)
+                } else {
+                    anyhow!(
+                        "Unknown alias @{}; known aliases: {}",
+                        name,
+                        known.join(", ")
+                    )
+                }
+            })
+        }
+        (true, true) => {
+            anyhow::bail!("Cannot use --alias together with a trailing command")
+        }
+        (false, false) => {
+            anyhow::bail!("Must specify either --alias <name> or a command after `--`")
+        }
+        (false, true) => Ok(command_args.join(" ")),
+    }
+}
+
 #[cfg(test)]
 mod tests {
-    // Most logic here is thin wrappers over `run_ssh`; integration
-    // tests against a live cluster cover the bulk of the behaviour.
+    use super::*;
+    use std::collections::HashMap;
+
+    fn test_config(aliases: &[(&str, &str)]) -> Config {
+        let exec_aliases: HashMap<String, String> = aliases
+            .iter()
+            .map(|(k, v)| (k.to_string(), v.to_string()))
+            .collect();
+        Config {
+            ssh_user: "u".to_string(),
+            ssh_key_path: "/k".to_string(),
+            clusters: vec![],
+            nccl: HashMap::new(),
+            hf_cache_dir: None,
+            init_port: 25000,
+            before_serve: HashMap::new(),
+            exec_aliases,
+        }
+    }
+
+    #[test]
+    fn resolve_command_alias() {
+        let cfg = test_config(&[("gpu-lock", "sudo nvidia-smi -lgc 300,2300")]);
+        let cmd = resolve_command(&cfg, Some("gpu-lock"), &[]).unwrap();
+        assert_eq!(cmd, "sudo nvidia-smi -lgc 300,2300");
+    }
+
+    #[test]
+    fn resolve_command_unknown_alias_lists_known() {
+        let cfg = test_config(&[("gpu-lock", "a"), ("shutdown", "b")]);
+        let err = resolve_command(&cfg, Some("nope"), &[]).unwrap_err();
+        let msg = err.to_string();
+        assert!(msg.contains("Unknown alias @nope"), "got: {msg}");
+        assert!(msg.contains("gpu-lock"), "got: {msg}");
+        assert!(msg.contains("shutdown"), "got: {msg}");
+    }
+
+    #[test]
+    fn resolve_command_unknown_alias_no_aliases() {
+        let cfg = test_config(&[]);
+        let err = resolve_command(&cfg, Some("anything"), &[]).unwrap_err();
+        let msg = err.to_string();
+        assert!(msg.contains("no aliases configured"), "got: {msg}");
+    }
+
+    #[test]
+    fn resolve_command_no_alias_no_args_errors() {
+        let cfg = test_config(&[]);
+        let err = resolve_command(&cfg, None, &[]).unwrap_err();
+        let msg = err.to_string();
+        assert!(msg.contains("Must specify"), "got: {msg}");
+    }
+
+    #[test]
+    fn resolve_command_alias_and_args_errors() {
+        let cfg = test_config(&[("gpu", "a")]);
+        let err = resolve_command(&cfg, Some("gpu"), &["ls".into()]).unwrap_err();
+        let msg = err.to_string();
+        assert!(msg.contains("Cannot use --alias together"), "got: {msg}");
+    }
+
+    #[test]
+    fn resolve_command_literal_join() {
+        let cfg = test_config(&[]);
+        let cmd =
+            resolve_command(&cfg, None, &["echo".into(), "hello".into(), "world".into()]).unwrap();
+        assert_eq!(cmd, "echo hello world");
+    }
 }

+ 5 - 2
src/main.rs

@@ -9,6 +9,7 @@ mod doctor;
 mod exec;
 mod health;
 mod host_keys;
+mod pre_serve;
 mod ssh;
 mod tmux;
 
@@ -98,12 +99,14 @@ async fn main() -> Result<()> {
             cluster,
             recipe,
             node,
+            alias,
             command,
         } => {
             let (cluster, _recipe) =
                 exec::resolve_target(&config, cluster.as_deref(), recipe.as_deref()).await?;
-            // Join the trailing args into a single command line.
-            let cmd_str = command.join(" ");
+            // Resolve command: alias takes precedence, otherwise join
+            // the trailing positional args into a single command line.
+            let cmd_str = exec::resolve_command(&config, alias.as_deref(), &command)?;
             exec::run(&config, &cluster, node.as_deref(), &cmd_str, dry_run).await?;
         }
     }

+ 56 - 0
src/pre_serve.rs

@@ -0,0 +1,56 @@
+//! Types for the `before_serve` configuration and `exec_aliases` table.
+//!
+//! Lives outside `config.rs` so the latter stays under the 500-line limit
+//! the project enforces (see AGENTS.md).
+
+use serde::Deserialize;
+
+/// Behaviour on a non-zero exit code for a [`PreServeCommand`].
+#[derive(Clone, Copy, Debug, PartialEq, Eq, Default, Deserialize)]
+#[serde(rename_all = "lowercase")]
+pub enum OnFailure {
+    /// Abort the serve (default).
+    #[default]
+    Abort,
+    /// Log a warning but continue.
+    Warn,
+    /// Continue silently.
+    Continue,
+}
+
+/// One entry in a `before_serve` block.  Either a bare command string
+/// (defaults to `on_failure = abort`) or a table with explicit `on_failure`.
+#[derive(Clone, Debug, Deserialize)]
+#[serde(untagged)]
+pub enum PreServeCommand {
+    Bare(String),
+    Configured {
+        command: String,
+        #[serde(default)]
+        on_failure: OnFailure,
+    },
+}
+
+impl PreServeCommand {
+    pub fn command(&self) -> &str {
+        match self {
+            PreServeCommand::Bare(c) => c,
+            PreServeCommand::Configured { command, .. } => command,
+        }
+    }
+
+    pub fn on_failure(&self) -> OnFailure {
+        match self {
+            PreServeCommand::Bare(_) => OnFailure::default(),
+            PreServeCommand::Configured { on_failure, .. } => *on_failure,
+        }
+    }
+}
+
+/// TOML sub-table under `[before_serve."<cluster>"]`.  Currently only
+/// carries `commands`; reserved for future per-cluster tuning.
+#[derive(Clone, Debug, Default, Deserialize)]
+pub struct BeforeServeTable {
+    #[serde(default)]
+    pub commands: Vec<PreServeCommand>,
+}