launcher/src/tmux/serve.rs

use anyhow::{Context, Result};
use std::time::Duration;
use tokio::time::sleep;
use tracing::{info, warn};

use crate::config::{Config, Recipe};
use crate::tmux::TmuxManager;

/// Launches a cluster via tmux sessions.
/// For Ray backend: starts head node first, waits for readiness, then workers.
pub async fn tmux_serve(
    config: &Config,
    recipe_name: &str,
    cluster_name: &str,
    skip_before_serve: bool,
    dry_run: bool,
    auto_power_on: bool,
    force: bool,
) -> Result<()> {
    let recipe_path = format!("recipes/{}/recipe.toml", recipe_name);
    let recipe: Recipe = toml::from_str(&tokio::fs::read_to_string(&recipe_path).await?)?;

    // Feature gaps: tmux deployment doesn't support Docker-only features.
    if !recipe.mods.commands.is_empty() {
        anyhow::bail!(
            "{}",
            fluent_i18n::t!("tmux-mods-unsupported", { "name" => &recipe.name })
        );
    }
    if recipe.composed_model.is_some() {
        anyhow::bail!(
            "{}",
            fluent_i18n::t!("tmux-composed-model-unsupported", { "name" => &recipe.name })
        );
    }

    info!(
        "🚀 TMUX launch: {} ({}) [backend: {}] [engine: {}]",
        recipe.name,
        recipe_name,
        recipe
            .distributed_backend
            .map_or("none".into(), |b| b.to_string()),
        recipe.engine
    );

    let choice = if cluster_name != "auto" {
        cluster_name.to_string()
    } else if let Some(ref rc) = recipe.cluster {
        rc.clone()
    } else {
        "auto".to_string()
    };
    // Use prepare_cluster (not select_cluster) so --on (auto_power_on) and
    // --force (force_occupied) are supported for tmux deploy — they have
    // nothing to do with Docker.
    let recipe_for_select = recipe.clone();
    let (cluster, nccl) = if dry_run {
        crate::orchestrate::prepare_cluster_dry_run(config, &recipe_for_select, &choice, None)?
    } else {
        use crate::ends::{SshNodePinger, SshProber, ha_controller_from_config};
        use crate::orchestrate::{BootTiming, PrepareRequest, prepare_cluster};
        let power_ctrl = if auto_power_on {
            ha_controller_from_config(config)
        } else {
            None
        };
        prepare_cluster(
            config,
            &recipe_for_select,
            &choice,
            None,
            &PrepareRequest {
                prober: &SshProber,
                power: power_ctrl.as_ref(),
                pinger: &SshNodePinger,
                auto_power_on,
                boot_timing: BootTiming::default(),
                force_occupied: force,
            },
        )
        .await?
    };

    if let Some(required) = recipe.required_nodes {
        let actual = cluster.nodes.len() as u32;
        if actual != required {
            anyhow::bail!(
                "Recipe '{}' requires {} nodes, but cluster '{}' has {} nodes",
                recipe.name,
                required,
                cluster.name,
                actual
            );
        }
    }

    crate::before_serve::run(config, &cluster, &recipe, skip_before_serve, dry_run).await?;

    // Prepare the Python venv on each node if it doesn't exist yet.
    prepare_venv(config, &recipe, &cluster, recipe_name, dry_run).await?;

    let manager = TmuxManager;

    let is_ray = recipe.distributed_backend == Some(crate::config::DistributedBackend::Ray);

    if is_ray {
        info!("{}", fluent_i18n::t!("tmux-ray-backend"));

        let head_node = cluster
            .nodes
            .iter()
            .find(|n| n.rank == 0)
            .context("No head node (rank 0) in cluster")?;

        let head_cmds = manager.build_session_commands(head_node, &recipe, nccl.as_ref(), &cluster);
        manager
            .create_session(config, head_node, &head_cmds, dry_run)
            .await?;

        if !dry_run {
            info!("{}", fluent_i18n::t!("tmux-waiting-head"));
            let sess = manager.session_name(&recipe, 0);
            let head_ready = wait_for_head_ready(config, head_node, &sess, &recipe).await?;
            if head_ready {
                info!("{}", fluent_i18n::t!("tmux-head-ready"));
            } else {
                warn!("⚠️ Head node may not be fully ready — starting workers anyway.");
            }
        }

        for node in &cluster.nodes {
            if node.rank == 0 {
                continue;
            }
            let worker_cmds =
                manager.build_session_commands(node, &recipe, nccl.as_ref(), &cluster);
            manager
                .create_session(config, node, &worker_cmds, dry_run)
                .await?;
        }
    } else {
        info!("{}", fluent_i18n::t!("tmux-mp-backend"));

        let mut handles = Vec::new();
        for node in &cluster.nodes {
            let config = config.clone();
            let node = node.clone();
            let recipe = recipe.clone();
            let nccl = nccl.clone();
            let cluster = cluster.clone();
            let manager = manager.clone();
            let dry = dry_run;

            handles.push(crate::locale::spawn_with_locale(async move {
                let cmds = manager.build_session_commands(&node, &recipe, nccl.as_ref(), &cluster);
                manager
                    .create_session(&config, &node, &cmds, dry)
                    .await
                    .map(|_| node.name.clone())
            }));
        }

        let results = futures::future::join_all(handles).await;
        let mut failures = Vec::new();
        for result in results {
            match result {
                Ok(Ok(name)) => info!(
                    "{}",
                    fluent_i18n::t!("tmux-session-started", { "name" => &name })
                ),
                Ok(Err(e)) => failures.push(e),
                Err(e) => failures.push(anyhow::anyhow!("Task join: {}", e)),
            }
        }

        if !failures.is_empty() {
            for e in &failures {
                warn!("  ✗ {}", e);
            }
            anyhow::bail!("Failed to start sessions on {} node(s)", failures.len());
        }
    }

    if !dry_run {
        sleep(Duration::from_secs(5)).await;
        info!(
            "🎉 Cluster launched via tmux! API → http://{}:{}",
            cluster.effective_master_ip(),
            recipe.port
        );
        info!(
            "📌 To view sessions: `cargo run tmux status --recipe {}`",
            recipe_name
        );
        info!(
            "📌 To attach: `cargo run tmux attach --recipe {}`",
            recipe_name
        );
    }

    Ok(())
}

/// Polls the head node until the API endpoint responds or timeout.
async fn wait_for_head_ready(
    config: &Config,
    head_node: &crate::config::Node,
    session_name: &str,
    recipe: &Recipe,
) -> Result<bool> {
    let max_attempts = 30;
    for i in 0..max_attempts {
        let pid_cmd = format!(
            "tmux list-panes -t {}:0 -F '#{{pane_pid}}' | head -1",
            session_name
        );
        match crate::ssh::run_ssh(config, head_node, pid_cmd, false).await {
            Ok(pid_str) => {
                let pid = pid_str.trim();
                if pid.is_empty() || pid == "0" {
                    sleep(Duration::from_secs(5)).await;
                    continue;
                }
            }
            Err(_) => {
                sleep(Duration::from_secs(5)).await;
                continue;
            }
        }

        let api_check = format!(
            "curl -s -o /dev/null -w '%{{http_code}}' http://127.0.0.1:{}/v1/models 2>/dev/null || echo 'not_ready'",
            recipe.port
        );
        if let Ok(out) = crate::ssh::run_ssh(config, head_node, api_check, false).await {
            let code = out.trim();
            if code == "200" {
                info!(
                    "{}",
                    fluent_i18n::t!("tmux-head-api-ready", { "attempt" => i + 1 })
                );
                return Ok(true);
            }
        }

        if (i + 1) % 5 == 0 {
            info!(
                "  Still waiting for head node... ({}/{})",
                i + 1,
                max_attempts
            );
        }
        sleep(Duration::from_secs(5)).await;
    }

    warn!("  Head node did not become ready within timeout");
    Ok(false)
}

/// Prepares the Python venv on each node if it doesn't exist.
///
/// Checks `test -f <venv>/bin/activate` on each node; if absent, runs
/// the setup commands in order, then re-checks. The setup commands come
/// from either `recipe.venv_setup_script` (a script file read from the
/// recipe directory) or `recipe.venv_setup` (inline list or multi-line
/// string). A failure on any command aborts the serve. Nodes are
/// prepared in parallel.
async fn prepare_venv(
    config: &Config,
    recipe: &Recipe,
    cluster: &crate::config::Cluster,
    recipe_name: &str,
    dry_run: bool,
) -> Result<()> {
    let venv = match &recipe.venv {
        Some(v) => v,
        None => {
            info!("prepare_venv: no venv field in recipe — skipping");
            return Ok(());
        }
    };

    info!(
        "prepare_venv: venv={}, has_venv_setup_script={}, has_venv_setup={}",
        venv,
        recipe.venv_setup_script.is_some(),
        !recipe.venv_setup.to_vec().is_empty()
    );

    // Resolve setup commands: script file takes precedence over inline.
    let setup: Vec<String> = if let Some(script_path) = &recipe.venv_setup_script {
        let full_path = format!("recipes/{}/{}", recipe_name, script_path);
        if dry_run {
            info!("  read venv_setup_script: {} (dry-run)", full_path);
        } else {
            info!("  read venv_setup_script: {}", full_path);
        }
        let body = tokio::fs::read_to_string(&full_path)
            .await
            .map_err(|e| anyhow::anyhow!("Cannot read venv_setup_script '{}': {}", full_path, e))?;
        body.lines()
            .map(|l| l.trim().to_string())
            .filter(|l| !l.is_empty() && !l.starts_with('#'))
            .collect()
    } else {
        recipe.venv_setup.to_vec()
    };

    if setup.is_empty() {
        info!("venv setup is empty for {} — skipping", venv);
        return Ok(());
    }

    info!(
        "{} ({} commands)",
        fluent_i18n::t!("venv-preparing", { "venv" => venv }),
        setup.len()
    );

    // Prepare venv on each node sequentially.
    for node in &cluster.nodes {
        let check_cmd = format!("test -f {}/bin/activate", venv);

        if !dry_run {
            // Skip if venv already exists.
            if crate::ssh::run_ssh(config, node, check_cmd.clone(), false)
                .await
                .is_ok()
            {
                info!(
                    "{}",
                    fluent_i18n::t!("venv-exists", { "name" => &node.name })
                );
                continue;
            }
            info!(
                "{}",
                fluent_i18n::t!("venv-creating", { "name" => &node.name })
            );

            // Create venv, then run setup with it activated so that
            // uv pip install finds VIRTUAL_ENV. All commands are joined
            // with && so the script halts on the first failure.
            let joined = setup.join(" && ");
            let full_script = format!(
                "uv venv {} -p 3.12 && source {}/bin/activate && {}",
                venv, venv, joined
            );
            let script = format!("bash -c '{}'", full_script.replace('\'', "'\\''"));
            crate::ssh::run_ssh(config, node, script, false).await?;

            // Re-check.
            if crate::ssh::run_ssh(config, node, check_cmd, false)
                .await
                .is_ok()
            {
                info!(
                    "{}",
                    fluent_i18n::t!("venv-created", { "name" => &node.name })
                );
            } else {
                anyhow::bail!(
                    "{}",
                    fluent_i18n::t!("venv-not-created", { "name" => &node.name })
                );
            }
        }
    }

    Ok(())
}