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;
use anyhow::Result;
use std::time::{Duration, Instant};

/// `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_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.
pub(crate) async fn wait_engine(cluster: &Cluster, port: u16, timeout_secs: u64) -> Result<()> {
    wait_ready(&head_url(cluster, port), timeout_secs).await
}

/// 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.
pub(crate) async fn watch(cluster: &Cluster, port: u16) -> Result<()> {
    let mut misses = 0;
    loop {
        tokio::time::sleep(Duration::from_secs(WATCH_POLL_SECS)).await;
        if alive(cluster, port).await {
            misses = 0;
        } else {
            misses += 1;
            if misses >= WATCH_MISSES {
                anyhow::bail!(
                    "engine stopped answering at {} — containers died mid-follow",
                    head_url(cluster, port)
                );
            }
        }
    }
}