launcher/src/tmux/mod.rs

pub mod commands;
pub mod ops;
pub mod serve;
pub mod session;

use anyhow::Result;

use crate::cli::TmuxCommands;
use crate::config::{Config, NcclConfig, Node, Recipe};

/// Manages tmux sessions for model inference across a DGX Spark cluster.
#[derive(Clone)]
pub struct TmuxManager {
    #[allow(dead_code)]
    pub session_prefix: String,
}

impl TmuxManager {
    /// Generates a namespaced session name for a recipe and node rank.
    /// Format: `{recipe_name}-{engine}-r{rank}`
    pub fn session_name(&self, recipe: &Recipe, rank: u32) -> String {
        let name = recipe.name.replace(' ', "_");
        format!("{}-{}-r{}", name, recipe.engine, rank)
    }

    /// Builds the list of `export` commands for NCCL environment variables.
    pub fn build_nccl_setup(&self, nccl: &NcclConfig) -> Vec<String> {
        vec![
            format!("export NCCL_IB_HCA=\"{}\"", nccl.ib_hca),
            format!("export NCCL_SOCKET_IFNAME=\"{}\"", nccl.socket_ifname),
            format!("export NCCL_IB_GID_INDEX={}", nccl.ib_gid_index),
            format!("export NCCL_IB_DISABLE={}", nccl.ib_disable),
        ]
    }

    /// Builds engine-specific host IP environment variable.
    pub fn build_host_ip_export(&self, engine: &str, node: &Node) -> Vec<String> {
        match engine {
            "vllm" => vec![format!("export VLLM_HOST_IP={}", node.ib_ip)],
            "sglang" => vec![format!("export SGLANG_WORKER_IP={}", node.ib_ip)],
            _ => vec![
                format!("export VLLM_HOST_IP={}", node.ib_ip),
                format!("export SGLANG_WORKER_IP={}", node.ib_ip),
            ],
        }
    }
}

/// Detailed health status of a tmux session.
#[derive(Debug, Clone)]
pub struct SessionHealth {
    pub session_exists: bool,
    pub process_alive: bool,
    pub gpu_ok: bool,
    pub api_ok: bool,
}

impl SessionHealth {
    pub fn is_healthy(&self) -> bool {
        self.session_exists && self.process_alive && self.gpu_ok && self.api_ok
    }

    #[allow(dead_code)]
    pub fn is_degraded(&self) -> bool {
        self.session_exists && !self.process_alive
    }
}

/// Top-level handler for all `cargo run tmux <subcommand>` commands.
pub async fn handle_tmux(config: &Config, cmd: TmuxCommands, dry_run: bool) -> Result<()> {
    match cmd {
        TmuxCommands::Serve {
            recipe,
            cluster,
            skip_before_serve,
        } => {
            serve::tmux_serve(config, &recipe, &cluster, skip_before_serve, dry_run).await?;
        }
        TmuxCommands::Stop { recipe, cluster } => {
            ops::tmux_stop(config, &recipe, &cluster, dry_run).await?;
        }
        TmuxCommands::Status { recipe, cluster } => {
            ops::tmux_status(config, &recipe, &cluster, dry_run).await?;
        }
        TmuxCommands::Logs {
            recipe,
            follow,
            cluster,
        } => {
            ops::tmux_logs(config, &recipe, &cluster, follow, dry_run).await?;
        }
        TmuxCommands::Attach { recipe, cluster } => {
            ops::tmux_attach(config, &recipe, &cluster, dry_run).await?;
        }
    }
    Ok(())
}