launcher/src/cli.rs

use clap::{Parser, Subcommand};

#[derive(Parser)]
#[command(author, version, about = "DGX Spark Launcher for vLLM and SGLang")]
pub struct Cli {
    #[arg(long, short, default_value = "config.toml")]
    pub config: String,

    #[arg(long, short, help = "Show commands without executing")]
    pub dry_run: bool,

    #[command(subcommand)]
    pub command: Commands,
}

#[derive(Subcommand)]
pub enum Commands {
    /// Launch a cluster
    Serve {
        #[arg(long, required = true)]
        recipe: String,

        #[arg(long, default_value = "auto")]
        cluster: String,

        /// Target a specific node within the cluster (default: all).
        #[arg(long)]
        node: Option<String>,

        /// Keep containers after stop (default: auto-remove on stop)
        #[arg(long)]
        keep: bool,

        /// Skip pre-serve hooks (cluster-level + recipe-level).
        #[arg(long)]
        skip_before_serve: bool,

        /// Follow logs after serving (use --no-follow to skip).
        #[arg(long, action = clap::ArgAction::SetFalse, default_value_t = true)]
        follow: bool,

        /// Power on cluster if switch is off (requires HA config).
        #[arg(short = 'O', long = "on")]
        auto_power_on: bool,

        /// Sync Docker image across nodes before starting.
        #[arg(long, default_value_t = true)]
        sync_image: bool,

        /// Explicit source node for image sync (default: auto-detect).
        #[arg(long)]
        sync_image_source: Option<String>,

        /// Build image on this node if not found (requires docker-images/<recipe>/).
        #[arg(long)]
        build_node: Option<String>,

        /// Force rebuild even if checksum matches.
        #[arg(long)]
        force_build: bool,

        /// Run even if nodes are occupied (another container running).
        /// Skips the occupancy check during cluster selection — use to run
        /// a lightweight test recipe alongside a live serve, or to force
        /// onto a busy node. Dangerous for GPU recipes (OOM, port
        /// conflicts); the operator is responsible for coexistence.
        #[arg(long)]
        force: bool,

        /// Serve single-node on the first free node found (overrides recipe
        /// `solo`). Use `--no-solo` to force multi-node even if the recipe
        /// sets `solo = true`.
        #[arg(long)]
        solo: Option<bool>,
    },

    /// Show container logs
    Logs {
        #[arg(long, required = true)]
        recipe: String,

        #[arg(long, action = clap::ArgAction::SetTrue)]
        follow: bool,

        #[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
    Stop {
        #[arg(long, required = true)]
        recipe: String,

        /// Target specific nodes (repeatable, e.g. --node dgx1 --node dgx2).
        #[arg(long)]
        node: Vec<String>,

        #[arg(long, default_value = "auto")]
        cluster: String,

        /// Power off cluster after stopping containers (requires HA config).
        #[arg(long = "off")]
        power_off: bool,
    },

    /// Check cluster health
    Health {
        #[arg(
            long,
            help = "Cluster name (e.g. 4nodes-switch). If omitted — checks all."
        )]
        cluster: Option<String>,
        #[arg(long, help = "Show node host keys (for adding to config.toml)")]
        fingerprint: bool,
    },

    /// Check and repair a recipe (doctor mode)
    Doctor {
        #[arg(long, required = true, help = "Recipe name to check")]
        recipe: String,
        #[arg(
            long,
            default_value = "auto",
            help = "Cluster (if not auto — uses recipe.cluster)"
        )]
        cluster: String,
        #[arg(long, help = "Auto-fix issues (create Dockerfile + build image)")]
        repair: bool,
    },

    /// Manage tmux sessions (Docker alternative)
    #[command(subcommand)]
    Tmux(TmuxCommands),

    /// Auto-discover host keys and write them into config.toml
    InitHostKeys {
        #[arg(long, help = "Cluster name (e.g. 4nodes-switch). If omitted — all.")]
        cluster: Option<String>,
    },

    /// Run an arbitrary shell command on a cluster or single node
    Exec {
        /// Cluster name (mutually exclusive with --recipe).
        #[arg(long, conflicts_with = "recipe")]
        cluster: Option<String>,

        /// Recipe name; resolves to the recipe's pinned cluster.
        #[arg(long, conflicts_with = "cluster")]
        recipe: Option<String>,

        /// Target a specific node within the cluster (default: all).
        #[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)]
        command: Vec<String>,
    },

    /// Build Docker image for a recipe on a remote node
    Build {
        #[arg(long, required = true, help = "Recipe name")]
        recipe: String,

        #[arg(long, required = true, help = "Node to build on (e.g. dgx1.lan)")]
        build_node: String,

        #[arg(long, default_value = "auto", help = "Cluster name")]
        cluster: String,

        /// Force rebuild even if checksum matches.
        #[arg(long)]
        force_build: bool,
    },

    /// Power on/off nodes via Home Assistant (smart plugs / Wake-on-LAN).
    Power {
        #[command(subcommand)]
        action: PowerAction,
    },
}

#[derive(Subcommand)]
pub enum PowerAction {
    /// Power off: SSH `sudo shutdown now` + wait + HA `turn_off`.
    Off {
        #[arg(long)]
        cluster: String,

        /// Target a specific node within the cluster (default: all).
        #[arg(long)]
        node: Option<String>,
    },
    /// Power on: HA `turn_on` only (no SSH; the smart plug wakes the machine).
    On {
        #[arg(long)]
        cluster: String,

        /// Target a specific node within the cluster (default: all).
        #[arg(long)]
        node: Option<String>,
    },
    /// Switch clusters: disable nodes from source, enable nodes in target.
    Switch {
        /// Current cluster (nodes + entities to shut down).
        #[arg(long)]
        from: String,
        /// Target cluster (nodes + entities to wake up).
        #[arg(long)]
        to: String,
    },
}

#[derive(Subcommand)]
pub enum TmuxCommands {
    /// Launch cluster via tmux
    Serve {
        #[arg(long, required = true)]
        recipe: String,
        #[arg(long, default_value = "auto")]
        cluster: String,
        /// Skip pre-serve hooks (cluster-level + recipe-level).
        #[arg(long)]
        skip_before_serve: bool,
    },
    /// Stop tmux sessions
    Stop {
        #[arg(long, required = true)]
        recipe: String,
        #[arg(long, default_value = "auto")]
        cluster: String,
    },
    /// Detailed tmux session status
    Status {
        #[arg(long, required = true)]
        recipe: String,
        #[arg(long, default_value = "auto")]
        cluster: String,
    },
    /// Show tmux session logs
    Logs {
        #[arg(long, required = true)]
        recipe: String,
        #[arg(long, action = clap::ArgAction::SetTrue)]
        follow: bool,
        #[arg(long, default_value = "auto")]
        cluster: String,
    },
    /// Attach to head node tmux session
    Attach {
        #[arg(long, required = true)]
        recipe: String,
        #[arg(long, default_value = "auto")]
        cluster: String,
    },
}