launcher/src/tmux/serve.rs

use anyhow::{Context, Result};
use std::collections::HashSet;
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,
) -> Result<()> {
    let recipe_path = format!("recipes/{}/recipe.toml", recipe_name);
    let recipe: Recipe = toml::from_str(&tokio::fs::read_to_string(&recipe_path).await?)?;

    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()
    };
    let (cluster, nccl) =
        crate::config::select_cluster(config, &choice, recipe.required_nodes, &HashSet::new())?;

    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?;

    let manager = TmuxManager {
        session_prefix: recipe_name.to_string(),
    };

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

    if is_ray {
        info!("🐍 Ray backend detected — starting head node first...");

        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, &cluster);
        manager
            .create_session(config, head_node, &head_cmds, dry_run)
            .await?;

        if !dry_run {
            info!("⏳ Waiting for head node to be ready...");
            let sess = manager.session_name(&recipe, 0);
            let head_ready = wait_for_head_ready(config, head_node, &sess, &recipe).await?;
            if head_ready {
                info!("✅ Head node ready. Starting worker nodes...");
            } 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, &cluster);
            manager
                .create_session(config, node, &worker_cmds, dry_run)
                .await?;
        }
    } else {
        info!("🔗 MP backend detected — starting all nodes in parallel...");

        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(tokio::spawn(async move {
                let cmds = manager.build_session_commands(&node, &recipe, &nccl, &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!("  ✅ Session started on {}", 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.master_ib_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!("  Head node 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)
}