overbrainer 0.10.0

//! `overbrainer pod`: `ls` and `rm <run-id>` for the Runpod pods overbrainer
//! created, and the read-only catalog listings (`gpus`, `datacenters`,
//! `volumes`, `templates`).

use std::path::Path;
use std::sync::atomic::AtomicBool;

use anyhow::Context;

use super::progress::duration_words;
use super::{GpuArgs, PodCommand};
use crate::config::{DEFAULT_RUNPOD_BASE_URL, EnvSource, Settings};
use crate::events::EventBus;
use crate::runpod::{
    GpuFilter, PodCtx, PodError, RunpodClient, Timing, data_center_table, gpu_table, pod_rows,
    remove_run_pods, select_gpus, table, template_table, volume_table,
};
use crate::runs::{JobKind, Runs};

/// Runs an `overbrainer pod` subcommand.
///
/// # Errors
///
/// Returns an error when the configuration, the API key or the Runpod API fails,
/// or when `pod rm` refuses a running run.
pub async fn run(project_dir: &Path, command: &PodCommand) -> anyhow::Result<()> {
    let settings = crate::config::load(project_dir, EnvSource::Process)?;
    let client = client(&settings).await?;
    match command {
        PodCommand::Gpus(args) => gpus(&client, args).await,
        PodCommand::Datacenters => datacenters(&client).await,
        PodCommand::Volumes => volumes(&client).await,
        PodCommand::Templates => templates(&client).await,
        PodCommand::Ls => with_pods(project_dir, &client, Pods::Ls).await,
        PodCommand::Rm { run_id, force } => {
            let rm = Pods::Rm {
                run_id,
                force: *force,
            };
            with_pods(project_dir, &client, rm).await
        },
    }
}

/// The pod commands that need `runs/` and the progress renderer.
enum Pods<'a> {
    /// `pod ls`.
    Ls,
    /// `pod rm <run-id>`.
    Rm { run_id: &'a str, force: bool },
}

/// Runs `command` with the runs of `project_dir` and a progress renderer.
async fn with_pods(
    project_dir: &Path,
    client: &RunpodClient,
    command: Pods<'_>,
) -> anyhow::Result<()> {
    let runs = Runs::new(project_dir);
    // An export's pod is recorded with the export, inside its run.
    let runs = match command {
        Pods::Rm { run_id, .. } => runs.holding(run_id),
        Pods::Ls => runs,
    };
    let bus = EventBus::new();
    let renderer = tokio::spawn(super::progress::render(bus.subscribe()));
    let timing = Timing::standard();
    let interrupted = AtomicBool::new(false);
    let ctx = PodCtx {
        client,
        runs: &runs,
        bus: &bus,
        timing: &timing,
        interrupted: &interrupted,
    };
    let result = match command {
        Pods::Ls => ls(&ctx).await,
        Pods::Rm { run_id, force } => rm(&ctx, run_id, force).await,
    };
    drop(bus);
    renderer.await.ok();
    result
}

/// A client of the Runpod API with the account key, resolved only now, and
/// the job's secrets, which the pod's logs never show.
///
/// # Errors
///
/// Returns an error when the key is not set or cannot be resolved.
pub(crate) async fn client(settings: &Settings) -> anyhow::Result<RunpodClient> {
    let key = settings
        .runpod
        .api_key
        .as_ref()
        .context("no Runpod API key: set OVERBRAINER_RUNPOD__API_KEY")?;
    let key = super::resolver()
        .resolve(key)
        .await
        .context("cannot resolve runpod.api_key")?;
    let base_url = settings
        .runpod
        .base_url
        .as_deref()
        .unwrap_or(DEFAULT_RUNPOD_BASE_URL);
    Ok(RunpodClient::new(base_url, &key)?.with_secrets(job_secrets(settings).await))
}

/// The secrets a run's job may get, which its pod's logs must never show:
/// the Hugging Face token, when set and resolvable. Best effort: one that
/// cannot be resolved is only noted, since the job gets none then either.
async fn job_secrets(settings: &Settings) -> Vec<secrecy::SecretString> {
    let Some(token) = &settings.hf_token else {
        return Vec::new();
    };
    match super::resolver().resolve(token).await {
        Ok(token) => vec![token],
        Err(error) => {
            tracing::debug!("cannot resolve hf_token for the pod log redaction: {error}");
            Vec::new()
        },
    }
}

async fn ls(ctx: &PodCtx<'_>) -> anyhow::Result<()> {
    let rows = pod_rows(ctx).await?;
    if rows.is_empty() {
        println!("pod: no overbrainer pod on this account");
    }
    for line in table(&rows) {
        println!("{line}");
    }
    Ok(())
}

/// `pod rm`: deletes the pods of run or job `run_id` and prints each one; a
/// running export or compare is refused with how to stop it instead.
async fn rm(ctx: &PodCtx<'_>, run_id: &str, force: bool) -> anyhow::Result<()> {
    let removal = remove_run_pods(ctx, run_id, force).await;
    if removal.removed.is_empty() && removal.result.is_ok() {
        println!("pod: no pod for run {run_id}");
    }
    // Every deleted pod is printed, even when the command fails overall.
    for pod in &removal.removed {
        let after = pod.uptime.map_or_else(String::new, |uptime| {
            format!(" after {}", duration_words(uptime))
        });
        let spend = pod
            .estimated_spend
            .map_or_else(String::new, |spend| format!(", about ${spend:.2}"));
        println!("pod: {} deleted{after}{spend}", pod.pod_id);
    }
    match (removal.result, ctx.runs.job_of()) {
        (Err(error), Some((run, kind))) => Err(job_refusal(error, run_id, &run, kind)),
        (result, _) => Ok(result?),
    }
}

/// The error of `pod rm` of the job `id` (an export or a compare) of run
/// `run`: such a job is not stopped with `train cancel`, but by interrupting
/// the command following it, which cancels it.
fn job_refusal(error: PodError, id: &str, run: &str, kind: JobKind) -> anyhow::Error {
    let noun = kind.noun();
    let command = match kind {
        JobKind::Export => format!("overbrainer export {run}"),
        JobKind::Compare => format!("overbrainer compare --run {run}"),
    };
    match error {
        PodError::RunStillRunning { kept, .. } => anyhow::anyhow!(
            "{noun} {id} of run {run} is still running{kept}; stop the `{command}` following it \
             with Ctrl-C, which cancels it, or pass --force"
        ),
        PodError::StillStarting(_) => anyhow::anyhow!(
            "{noun} {id} of run {run} is still starting its pod; wait for it, or stop the \
             `{command}` starting it with Ctrl-C, or use `pod rm {id} --force`"
        ),
        other => other.into(),
    }
}

/// Runs `overbrainer pod gpus`: the Secure Cloud GPU types, filtered by
/// `args` and sorted cheapest first.
async fn gpus(client: &RunpodClient, args: &GpuArgs) -> anyhow::Result<()> {
    // `gpu_count` is not asked on the command line: the listing is for one GPU.
    let gpus = client.list_gpu_types(1).await?;
    let filter = GpuFilter {
        min_vram_gb: args.min_vram,
        max_price: args.max_price,
        data_center: args.data_center.clone(),
        in_stock: args.in_stock,
        gpu_count: None,
    };
    let selected = select_gpus(&gpus, &filter);
    if selected.is_empty() {
        println!("pod: no GPU type matches");
    }
    for line in gpu_table(&selected, args.data_center.as_deref()) {
        println!("{line}");
    }
    Ok(())
}

/// Runs `overbrainer pod datacenters`: every data center of the catalog.
async fn datacenters(client: &RunpodClient) -> anyhow::Result<()> {
    let data_centers = client.list_data_centers().await?;
    if data_centers.is_empty() {
        println!("pod: no data center found");
    }
    for line in data_center_table(&data_centers) {
        println!("{line}");
    }
    Ok(())
}

/// Runs `overbrainer pod volumes`: the account's network volumes.
async fn volumes(client: &RunpodClient) -> anyhow::Result<()> {
    let volumes = client.list_network_volumes().await?;
    if volumes.is_empty() {
        println!("pod: no network volume on this account");
    }
    for line in volume_table(&volumes) {
        println!("{line}");
    }
    Ok(())
}

/// Runs `overbrainer pod templates`: the account's pod templates.
async fn templates(client: &RunpodClient) -> anyhow::Result<()> {
    let templates = client.list_templates().await?;
    if templates.is_empty() {
        println!("pod: no pod template on this account");
    }
    for line in template_table(&templates) {
        println!("{line}");
    }
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;

    /// `pod rm` of a running export or compare says to stop the command
    /// following it, not `train cancel`; the export text is unchanged.
    #[test]
    fn pod_rm_of_a_running_export_says_how_to_stop_it() {
        let running = |id: &str| PodError::RunStillRunning {
            run_id: id.into(),
            kept: ": kept pod p1".into(),
        };
        assert_eq!(
            job_refusal(running("e1"), "e1", "r1", JobKind::Export).to_string(),
            "export e1 of run r1 is still running: kept pod p1; stop the `overbrainer export r1` \
             following it with Ctrl-C, which cancels it, or pass --force"
        );
        let starting = job_refusal(
            PodError::StillStarting("e1".into()),
            "e1",
            "r1",
            JobKind::Export,
        );
        assert!(!starting.to_string().contains("train cancel"), "{starting}");
        assert_eq!(
            job_refusal(running("c1"), "c1", "r1", JobKind::Compare).to_string(),
            "compare c1 of run r1 is still running: kept pod p1; stop the \
             `overbrainer compare --run r1` following it with Ctrl-C, which cancels it, or pass \
             --force"
        );
    }
}