launcher/src/readiness.rs

//! Engine readiness probes: poll `GET /v1/models` on the head node until
//! the inference API answers — after `sm121 serve` launches containers
//! (the model load can take many minutes) and before `sm121 bench` runs —
//! plus a liveness watch for follow mode (the serve must fail loudly when
//! the containers die mid-stream instead of tailing a dead engine).

use crate::config::{Cluster, Config};
use crate::ssh::run_ssh;
use anyhow::Result;
use std::time::{Duration, Instant};
use tokio::net::TcpStream;

/// `sm121 serve`: big checkpoints load for many minutes.
pub(crate) const SERVE_WAIT_SECS: u64 = 1800;
/// `sm121 bench`: the serve is expected to already be up.
pub(crate) const BENCH_WAIT_SECS: u64 = 900;
/// Liveness watch: probe interval / consecutive misses that mean "died".
const WATCH_POLL_SECS: u64 = 30;
const WATCH_MISSES: u32 = 2;

fn head_host_port(cluster: &Cluster, port: u16) -> (String, u16) {
    let head = cluster
        .nodes
        .iter()
        .find(|n| n.rank == 0)
        .or_else(|| cluster.nodes.first());
    (head.map_or("", |n| n.ssh_ip.as_str()).to_string(), port)
}

fn head_url(cluster: &Cluster, port: u16) -> String {
    let head = cluster
        .nodes
        .iter()
        .find(|n| n.rank == 0)
        .or_else(|| cluster.nodes.first());
    format!(
        "http://{}:{}/v1",
        head.map_or("", |n| n.ssh_ip.as_str()),
        port
    )
}

/// Poll `GET <base>/models` until it answers (engine up) or the timeout
/// hits; progress note every 30 s.
pub(crate) async fn wait_ready(base_url: &str, timeout_secs: u64) -> Result<()> {
    let client = reqwest::Client::builder()
        .timeout(Duration::from_secs(10))
        .build()?;
    let deadline = Instant::now() + Duration::from_secs(timeout_secs);
    let mut since_note = Instant::now();
    loop {
        if let Ok(r) = client.get(format!("{base_url}/models")).send().await
            && r.status().is_success()
        {
            return Ok(());
        }
        if since_note.elapsed() >= Duration::from_secs(30) {
            since_note = Instant::now();
            eprintln!("waiting for API at {base_url} (up to {timeout_secs} s)...");
        }
        if Instant::now() > deadline {
            anyhow::bail!("API at {base_url} not ready after {timeout_secs} s");
        }
        tokio::time::sleep(Duration::from_secs(5)).await;
    }
}

/// Serve-side readiness: the head node's engine API must answer, AND every
/// node's recipe container must stay alive. A container that exits during the
/// (multi-minute) model load is detected within ~30 s and fails the serve fast
/// — instead of polling a dead API until the full timeout.
pub(crate) async fn wait_engine(
    config: &Config,
    cluster: &Cluster,
    port: u16,
    timeout_secs: u64,
    recipe_name: &str,
) -> Result<()> {
    let client = reqwest::Client::builder()
        .timeout(Duration::from_secs(10))
        .build()?;
    let url = head_url(cluster, port);
    let deadline = Instant::now() + Duration::from_secs(timeout_secs);
    let mut since_note = Instant::now();
    let mut since_alive = Instant::now();
    loop {
        if let Ok(r) = client.get(format!("{url}/models")).send().await
            && r.status().is_success()
        {
            return Ok(());
        }
        // Fail fast if a recipe container has exited (the API is dead, so the
        // pure poll above would otherwise run to the full timeout).
        if since_alive.elapsed() >= Duration::from_secs(30) {
            since_alive = Instant::now();
            containers_alive(config, cluster, recipe_name).await?;
        }
        if since_note.elapsed() >= Duration::from_secs(30) {
            since_note = Instant::now();
            eprintln!("waiting for API at {url} (up to {timeout_secs} s)...");
        }
        if Instant::now() > deadline {
            anyhow::bail!("API at {url} not ready after {timeout_secs} s");
        }
        tokio::time::sleep(Duration::from_secs(5)).await;
    }
}

/// Every node's recipe container must be running. Returns an error naming the
/// dead node/container + a log tail, so a crashed launch fails fast instead of
/// polling a dead API. A node that is itself unreachable is skipped (the API
/// probe covers that case); only a container that has *exited* is a death.
pub(crate) async fn containers_alive(
    config: &Config,
    cluster: &Cluster,
    recipe_name: &str,
) -> Result<()> {
    let mut dead: Vec<String> = Vec::new();
    for node in &cluster.nodes {
        let cmd = format!(
            "docker ps -a --filter label=sm121.recipe={recipe_name} --format '{{{{.Names}}}}|{{{{.Status}}}}'"
        );
        let out = match run_ssh(config, node, cmd, false).await {
            Ok(o) => o,
            Err(_) => continue, // node unreachable — not a container death
        };
        for line in out.lines() {
            let (name, status) = match line.trim().split_once('|') {
                Some((n, s)) => (n, s.trim()),
                None => continue,
            };
            if !status.starts_with("Up") {
                let tail = run_ssh(
                    config,
                    node,
                    format!("docker logs --tail 12 {name} 2>&1 | tail -12"),
                    false,
                )
                .await
                .unwrap_or_default();
                let tail = tail
                    .lines()
                    .map(|l| format!("      {l}"))
                    .collect::<Vec<_>>()
                    .join("\n");
                dead.push(format!("  {} ({}): {}\n{tail}", node.name, name, status));
            }
        }
    }
    if !dead.is_empty() {
        anyhow::bail!(
            "serve crashed before readiness — container(s) exited:\n{}",
            dead.join("\n")
        );
    }
    Ok(())
}

/// One-shot liveness: is the engine answering right now?
pub(crate) async fn alive(cluster: &Cluster, port: u16) -> bool {
    let client = reqwest::Client::builder()
        .timeout(Duration::from_secs(10))
        .build()
        .ok();
    let Some(client) = client else {
        return false;
    };
    client
        .get(format!("{}/models", head_url(cluster, port)))
        .send()
        .await
        .is_ok_and(|r| r.status().is_success())
}

/// Follow-mode liveness watch: resolves `Err` once the engine misses
/// `WATCH_MISSES` consecutive probes — the containers died mid-follow.
/// The API port opens only when the engine is ready (weight load can take
/// 15+ min), so a connection-level failure before first readiness is "still
/// loading", not "died"; misses count only after the engine answered once
/// (or the port is listening but /v1/models misbehaves).
pub(crate) async fn watch(cluster: &Cluster, port: u16) -> Result<()> {
    let (host, p) = head_host_port(cluster, port);
    let mut misses = 0;
    let mut ever_ready = false;
    loop {
        tokio::time::sleep(Duration::from_secs(WATCH_POLL_SECS)).await;
        if alive(cluster, port).await {
            ever_ready = true;
            misses = 0;
        } else {
            let listening = TcpStream::connect((host.as_str(), p)).await.is_ok();
            if !ever_ready && !listening {
                continue; // API not open yet — normal while loading
            }
            misses += 1;
            if misses >= WATCH_MISSES {
                anyhow::bail!(
                    "engine stopped answering at {} — containers died mid-follow",
                    head_url(cluster, port)
                );
            }
        }
    }
}