overbrainer 0.8.2

//! The steps of a Runpod run around the generic run flows: provisioning its pod,
//! following its job under a client-side deadline, and ending the pod once the
//! results are retrieved. The CLI drives them and decides what Ctrl-C does.

use std::fs;
use std::io;
use std::path::Path;
use std::time::{Duration, SystemTime, UNIX_EPOCH};

use secrecy::SecretString;
use tokio::sync::broadcast::Receiver;
use tokio::sync::broadcast::error::RecvError;

use crate::config::{SshClient, process_client};
use crate::events::Event;
use crate::exec::{ExecError, Executor, PodEndpoint, SshDestination, SshExecutor};
use crate::runs::{Outcome, RunCtx, RunError, RunRecord, RunState, Runs, SnapshotReason, watch};
use crate::train::{Pace, SNAPSHOT_REQUEST, Trainer};

use super::disk::{DiskWatch, VolumeDisk};
use super::provision::note_strays;
use super::{
    BOOTSTRAP_LOG, CLIENT_KEY, CostCap, DeleteReason, DeletedBy, Pod, PodCtx, PodError, PodId,
    PodKeys, PodPlan, PodRecord, PodState, PodStatus, Provisioned, RemoteStatus, RunpodTarget,
    SSH_DIR, SshEndpoint, VOLUME_MOUNT, VOLUME_WORKDIR, alias, forget_keys, keep_file,
    pod_endpoint, provision, remove, short_cap_warning, sweep, with_pod_logs,
};
use crate::secrets::Redactor;

/// How long after the watchdog's deadline the client deletes the pod itself, when
/// it is still there: the watchdog should have done it first.
pub const DEADLINE_MARGIN: Duration = Duration::from_mins(5);

/// Where the watchdog looks for the client's "retrieved" marker, in a run
/// directory on the pod.
pub const RETRIEVED_MARKER: &str = ".pod/retrieved";

/// Why a run fails when its pod was deleted on `max_hours`, by the client's own
/// guard or by the pod's watchdog.
pub const MAX_HOURS_REACHED: &str = "max_hours reached: the pod was deleted before the job ended";

/// Why a run fails when its pod was deleted on `max_cost_usd` by its watchdog.
pub const MAX_COST_REACHED: &str = "max_cost_usd reached: the pod was deleted before the job ended";

/// Where the client writes, on the pod, when the job is stopped with a snapshot
/// for the cost cap (Unix seconds), in the run directory: read by the watchdog.
pub const SNAPSHOT_AT_FILE: &str = ".pod/snapshot_at";

/// Where the client writes, on the pod, when the watchdog deletes the pod for
/// the cost cap (Unix seconds), in the run directory.
pub const COST_CAP_FILE: &str = ".pod/cost_cap_at";

/// The lease a client following a job keeps on its pod, in the run directory:
/// the watchdog skips its `max_hours` deadline while the file was touched less
/// than [`LEASE_TTL`] ago.
pub const LEASE_FILE: &str = ".pod/lease";

/// How long a renewed lease holds the deadline off; the watchdog's
/// `OVERBRAINER_LEASE_TTL` default.
pub const LEASE_TTL: Duration = Duration::from_mins(15);

/// How often the following client renews the lease.
const LEASE_RENEW: Duration = Duration::from_mins(5);

/// A job with no new metric for this long counts as stalled: its lease is no
/// longer renewed, so `max_hours` applies again.
const STALL: Duration = Duration::from_mins(30);

/// How often the lease holder looks at the time.
const LEASE_TICK: Duration = Duration::from_secs(60);

/// Looks in a row that must answer Runpod's 404 before a pod is declared gone.
const GONE_LOOKS: u32 = 3;

/// Creates the pod of the new run `run` (from `runs::create`) and waits until it is
/// ready: keys in `runs/<id>/ssh/`, then `pod.json`, then provisioning. When that
/// fails, the run is saved `Failed` with the reason (`interrupted before the job
/// started` after Ctrl-C), whatever pod was created is deleted, and once every
/// pod of the run is confirmed deleted, the run's private client key is removed
/// and its host key secret deleted.
/// `vram_floor_gb` is the least VRAM `auto` GPU types need when the target sets
/// no `min_vram_gb` (see [`PodPlan::vram_floor_gb`]).
///
/// # Errors
///
/// Returns the [`PodError`] of the step that failed.
pub async fn start_pod(
    ctx: &PodCtx<'_>,
    target: &RunpodTarget,
    mut run: RunRecord,
    keep: bool,
    vram_floor_gb: Option<u32>,
) -> Result<(PodRecord, Provisioned), PodError> {
    let result = provision_run(ctx, target, &run, keep, vram_floor_gb).await;
    if let Err(error) = &result {
        run.state = RunState::Failed;
        run.message = Some(run_message(error));
        if let Err(save_error) = ctx.runs.save(&run) {
            tracing::warn!("cannot record run {} as failed: {save_error}", run.id);
        }
        forget_keys(ctx, &run.id).await;
    }
    result
}

/// The reason a failed provisioning gives in `run.json`: the error's message,
/// except that an answer of the Runpod API is reduced to its status. A
/// non-create call's error carries Runpod's own text (cleaned of the account
/// key), which the log shows but `run.json` never keeps.
fn run_message(error: &PodError) -> String {
    match error {
        PodError::Api(api) => match api.status() {
            Some(status) => format!("Runpod answered {status}"),
            None => error.to_string(),
        },
        _ => error.to_string(),
    }
}

/// Generates the run's SSH keys, saves its initial `pod.json`, then provisions
/// a reachable pod. The record exists before the first create request, so every
/// attempt can be tracked even when Runpod's answer is ambiguous.
async fn provision_run(
    ctx: &PodCtx<'_>,
    target: &RunpodTarget,
    run: &RunRecord,
    keep: bool,
    vram_floor_gb: Option<u32>,
) -> Result<(PodRecord, Provisioned), PodError> {
    let ssh_dir = ctx.runs.run_dir(&run.id)?.join(SSH_DIR);
    let client = process_client(target.ssh_client).map_err(PodError::SshClient)?;
    let keys = PodKeys::generate(&ssh_dir, &alias(&run.id), client)?;
    let mut record = PodRecord::new(&run.id, keep, target.gpu_count, &keys.host_public);
    record.max_cost_usd = target.max_cost_usd.filter(|_| !keep);
    record
        .network_volume_id
        .clone_from(&target.network_volume_id);
    record.max_volume_gb = target.max_volume_gb;
    record.save(ctx.runs)?;
    let plan = PodPlan {
        run_id: &run.id,
        target,
        keys: &keys,
        ssh_dir: &ssh_dir,
        workdir: target.workdir(),
        api_url: ctx.client.base_url(),
        vram_floor_gb,
        volume_gb: None,
    };
    let provisioned = provision(ctx, &plan, &mut record).await?;
    Ok((record, provisioned))
}

/// Hands the cost cap of `pod` to its watchdog before the job of `run` starts:
/// when to stop the job with a snapshot and when to delete the pod, in files
/// the watchdog reads every minute, with a warning when the cap buys too little
/// training ([`short_cap_warning`]). Nothing to do without a cap, or for a kept
/// pod.
///
/// # Errors
///
/// Returns [`PodError::CostCapUnset`] when a cap is set but cannot be applied:
/// Runpod gave no hourly rate for the pod, or the files cannot be written. No
/// job may start then, since nothing would stop it at its cap.
pub async fn arm_cost_cap<E: Executor>(
    executor: &E,
    pod: &PodRecord,
    run: &RunRecord,
) -> Result<(), PodError> {
    if let (Some(usd), Some(rate), false) = (pod.max_cost_usd, pod.cost_per_hour, pod.keep)
        && let Some(warning) = short_cap_warning(usd, rate)
    {
        tracing::warn!("{warning}");
    }
    match pod.cost_cap() {
        Some(cap) => write_cost_cap(executor, cap, &run.remote_dir)
            .await
            .map_err(|error| {
                PodError::CostCapUnset(format!("cannot write it on the pod: {error}"))
            }),
        None if pod.max_cost_usd.is_some() && !pod.keep => Err(PodError::CostCapUnset(
            "Runpod gave no hourly rate for the pod".to_string(),
        )),
        None => Ok(()),
    }
}

/// Writes the times of `cap` where the watchdog of the run in `remote_dir`
/// reads them.
async fn write_cost_cap<E: Executor>(
    executor: &E,
    cap: CostCap,
    remote_dir: &str,
) -> Result<(), ExecError> {
    for (file, at) in [
        (SNAPSHOT_AT_FILE, cap.snapshot_at),
        (COST_CAP_FILE, cap.delete_at),
    ] {
        executor
            .put_file(&format!("{remote_dir}/{file}"), &at.to_string())
            .await?;
    }
    Ok(())
}

/// Records that the job of the run was started on the pod.
///
/// # Errors
///
/// Returns [`PodError::Runs`] when `pod.json` cannot be saved.
pub fn job_started(runs: &Runs, pod: &mut PodRecord) -> Result<(), PodError> {
    pod.state = PodState::Running;
    pod.save(runs)?;
    Ok(())
}

/// Follows the started run `record` with `runs::watch`, until its job ends, with
/// the client's own guard: at the watchdog's deadline plus [`DEADLINE_MARGIN`], a
/// pod still there is deleted and the run saved `Failed` (never for a kept pod).
/// That is [`watch_on_pod`], the pod's logs kept meanwhile ([`with_pod_logs`]),
/// then [`settle_watch`].
///
/// # Errors
///
/// Returns [`PodError::DeadlineReached`] after that guard fired, or when the
/// target stays unreachable because the watchdog deleted the pod at its
/// deadline, [`PodError::PodGone`] when it stays unreachable because the pod no
/// longer exists for another reason (the run is saved `Failed` in both cases),
/// and [`PodError::Run`] for any other failure of the watch (the job may keep
/// running: attach again).
pub async fn follow<E: Executor, T: Trainer>(
    ctx: &PodCtx<'_>,
    run_ctx: &RunCtx<'_, E>,
    trainer: &T,
    record: RunRecord,
    pod: &mut PodRecord,
) -> Result<Outcome, PodError> {
    let id = record.id.clone();
    let watched = with_pod_logs(ctx, pod, watch_on_pod(ctx, run_ctx, trainer, record, pod)).await;
    settle_watch(ctx, pod, &id, watched).await
}

/// How [`watch_on_pod`] ended.
#[derive(Debug)]
pub enum Watched {
    /// The watch ended first, with its outcome or its error.
    Ended(Box<Result<Outcome, RunError>>),
    /// The client's deadline passed first; nothing was deleted yet.
    DeadlinePassed,
    /// The pod spent its cost cap first; nothing was deleted yet.
    CostCapPassed,
}

/// Watches the started run `record` until its job ends or the client's deadline
/// (the watchdog's plus [`DEADLINE_MARGIN`]; none for a kept pod) passes, and
/// warns once when the job's pace cannot end it before the watchdog's deadline.
/// The [disk policy](DiskWatch) runs beside it. It never deletes anything, so
/// dropping it (on Ctrl-C) loses nothing: what it found is acted on by
/// [`settle_watch`].
pub async fn watch_on_pod<E: Executor, T: Trainer>(
    ctx: &PodCtx<'_>,
    run_ctx: &RunCtx<'_, E>,
    trainer: &T,
    record: RunRecord,
    pod: &PodRecord,
) -> Watched {
    let run = record.clone();
    let disk = Box::pin(disk_watch(ctx, run_ctx.executor, &run, pod).run(run_ctx.bus.subscribe()));
    let Some(wait) = until_deadline(pod) else {
        return tokio::select! {
            outcome = watch(run_ctx, trainer, record) => Watched::Ended(Box::new(outcome)),
            never = disk => match never {},
        };
    };
    let events = run_ctx.bus.subscribe();
    tokio::select! {
        outcome = watch(run_ctx, trainer, record) => Watched::Ended(Box::new(outcome)),
        () = tokio::time::sleep(wait) => Watched::DeadlinePassed,
        // Never ends: it only warns.
        () = warn_overrun(events, pod) => Watched::DeadlinePassed,
        never = disk => match never {},
    }
}

/// The disk policy of the run `run` on `pod`.
fn disk_watch<'a, E: Executor>(
    ctx: &'a PodCtx<'_>,
    executor: &'a E,
    run: &'a RunRecord,
    pod: &'a PodRecord,
) -> DiskWatch<'a, E> {
    let volume = pod.network_volume_id.as_deref().map(|id| VolumeDisk {
        id,
        mount: VOLUME_MOUNT,
        max_gb: pod.max_volume_gb,
    });
    let watch = DiskWatch::new(executor, ctx.client, run, volume);
    // A run on a network volume whose ID `pod.json` does not hold (written
    // before overbrainer kept it) cannot be measured: `df` there is the whole
    // shared cluster.
    if volume.is_none() && run.remote_dir.starts_with(VOLUME_WORKDIR) {
        return watch.idle();
    }
    watch
}

/// [`watch_on_pod`] for a client that stays with the job: while the job keeps
/// making progress (a metric in the last 30 minutes), it renews the pod's lease
/// ([`LEASE_FILE`]) every few minutes, so `max_hours` deletes nothing, neither
/// the watchdog nor the client's own guard. Once the job stalls or the lease
/// cannot be renewed, the lease runs out after [`LEASE_TTL`] and the deadline
/// applies again. It warns once, as information, when the job ends after the
/// deadline. Once the pod has spent [`SNAPSHOT_SHARE`](super::SNAPSHOT_SHARE) of
/// its cost cap, it asks the job for a snapshot, as the watchdog does. A kept
/// pod has no deadline and needs no lease. The [disk policy](DiskWatch) runs
/// beside it, kept pod or not.
pub async fn watch_leased<E: Executor, T: Trainer>(
    ctx: &PodCtx<'_>,
    run_ctx: &RunCtx<'_, E>,
    trainer: &T,
    record: RunRecord,
    pod: &PodRecord,
) -> Watched {
    let run = record.clone();
    // Boxed: its state would otherwise weigh on every caller's future.
    let disk = Box::pin(disk_watch(ctx, run_ctx.executor, &run, pod).run(run_ctx.bus.subscribe()));
    if until_deadline(pod).is_none() {
        return tokio::select! {
            outcome = watch(run_ctx, trainer, record) => Watched::Ended(Box::new(outcome)),
            never = disk => match never {},
        };
    }
    let events = run_ctx.bus.subscribe();
    let remote_dir = record.remote_dir.clone();
    tokio::select! {
        outcome = watch(run_ctx, trainer, record) => Watched::Ended(Box::new(outcome)),
        // Boxed: its state would otherwise weigh on every caller's future.
        watched = Box::pin(hold_lease(run_ctx.executor, &remote_dir, events, pod)) => watched,
        never = disk => match never {},
    }
}

/// Renews the lease of the run in `remote_dir` while the job progresses, asks
/// for a snapshot at the cost cap, warns once when the job ends after the
/// deadline, and returns once the client's own deadline passed with no lease
/// held.
async fn hold_lease<E: Executor>(
    executor: &E,
    remote_dir: &str,
    mut events: Receiver<Event>,
    pod: &PodRecord,
) -> Watched {
    let mut lease = Lease::new(pod, remote_dir);
    let mut open = true;
    let mut tick = tokio::time::interval(LEASE_TICK);
    loop {
        tokio::select! {
            event = events.recv(), if open => match event {
                Ok(event) => lease.saw(&event),
                Err(RecvError::Lagged(_)) => {},
                Err(RecvError::Closed) => open = false,
            },
            _ = tick.tick() => {
                if let Some(guard) = lease.tick(executor).await {
                    return guard;
                }
            },
        }
    }
}

/// What [`hold_lease`] knows of the followed job, of its lease and of its cost
/// cap.
struct Lease<'a> {
    pod: &'a PodRecord,
    remote_dir: &'a str,
    pace: Pace,
    last_metric: SystemTime,
    renewed: Option<SystemTime>,
    warned: bool,
    /// Whether the snapshot for the cost cap was asked for.
    capped: bool,
}

impl<'a> Lease<'a> {
    fn new(pod: &'a PodRecord, remote_dir: &'a str) -> Self {
        Self {
            pod,
            remote_dir,
            pace: Pace::default(),
            last_metric: SystemTime::now(),
            renewed: None,
            warned: false,
            capped: false,
        }
    }

    /// Takes a training metric into account, and warns once when the job
    /// ends after the deadline.
    fn saw(&mut self, event: &Event) {
        let Event::Metric(metric) = event else {
            return;
        };
        self.pace.add(metric);
        self.last_metric = SystemTime::now();
        if self.warned {
            return;
        }
        if let Some(note) = overrun(&self.pace, self.pod, self.last_metric, true) {
            tracing::warn!("{note}");
            self.warned = true;
        }
    }

    /// Renews the lease when due and asks for the cost cap's snapshot once it
    /// is due; which of the client's guards fires, if one does: the cost cap
    /// (whatever the lease) or the deadline.
    async fn tick<E: Executor>(&mut self, executor: &E) -> Option<Watched> {
        let now = SystemTime::now();
        if renewal_due(self.renewed, self.last_metric, now) {
            self.renew(executor, now).await;
        }
        if !self.capped && cap_snapshot_due(self.pod, now) {
            self.cap(executor).await;
        }
        if past_cost_cap(self.pod, now) {
            return Some(Watched::CostCapPassed);
        }
        guard_fires(self.pod, self.renewed, now).then_some(Watched::DeadlinePassed)
    }

    /// Renews the lease at `now`.
    async fn renew<E: Executor>(&mut self, executor: &E, now: SystemTime) {
        let path = format!("{}/{LEASE_FILE}", self.remote_dir);
        match executor.put_file(&path, "").await {
            Ok(()) => self.renewed = Some(now),
            Err(error) => tracing::debug!("cannot renew the pod's lease: {error}"),
        }
    }

    /// Asks the job for the cost cap's snapshot; tried again next tick when
    /// the request cannot be written.
    async fn cap<E: Executor>(&mut self, executor: &E) {
        let path = format!("{}/{SNAPSHOT_REQUEST}", self.remote_dir);
        let written = executor.put_file(&path, SnapshotReason::Cost.name()).await;
        self.capped = written.is_ok();
        let note = match written {
            Ok(()) => cap_snapshot_note(self.pod, unix_secs(SystemTime::now())),
            Err(error) => format!("cannot ask for the cost cap's snapshot: {error}"),
        };
        tracing::warn!("{note}");
    }
}

/// Why the job on `pod` is stopped with a snapshot for its cost cap at `now`
/// (Unix seconds): it spent [`SNAPSHOT_SHARE`](super::SNAPSHOT_SHARE) of the
/// cap, or the cap is [`SNAPSHOT_LEAD`](super::SNAPSHOT_LEAD) away at most,
/// with what it spent so far.
fn cap_snapshot_note(pod: &PodRecord, now: u64) -> String {
    let (Some(cap), Some(usd), Some(spent)) = (pod.cost_cap(), pod.max_cost_usd, pod.spent_at(now))
    else {
        return "the pod neared max_cost_usd: its job is stopped with a snapshot".to_string();
    };
    let lead = super::SNAPSHOT_LEAD.as_secs();
    if cap.delete_at.saturating_sub(cap.snapshot_at) > lead {
        return format!(
            "the pod spent ${spent:.2} of max_cost_usd ${usd:.2} ({:.0}%): its job is stopped \
             with a snapshot",
            super::SNAPSHOT_SHARE * 100.0
        );
    }
    format!(
        "the pod reaches max_cost_usd ${usd:.2} in {} min (${spent:.2} spent so far): its job \
         is stopped with a snapshot {} min before the cap",
        cap.delete_at.saturating_sub(now).div_ceil(60),
        lead / 60
    )
}

/// Whether the job on `pod` is due for its cost cap's snapshot at `now`.
fn cap_snapshot_due(pod: &PodRecord, now: SystemTime) -> bool {
    pod.cost_cap()
        .is_some_and(|cap| unix_secs(now) >= cap.snapshot_at)
}

/// Seconds since the Unix epoch at `now`, 0 before it.
fn unix_secs(now: SystemTime) -> u64 {
    now.duration_since(UNIX_EPOCH)
        .map_or(0, |since| since.as_secs())
}

/// Whether the lease, last renewed at `renewed`, should be renewed at `now`:
/// once [`LEASE_RENEW`] passed, and only while the job made progress (its last
/// metric at `last_metric`) within [`STALL`].
fn renewal_due(renewed: Option<SystemTime>, last_metric: SystemTime, now: SystemTime) -> bool {
    let since = |at: SystemTime| now.duration_since(at).unwrap_or_default();
    since(last_metric) < STALL && renewed.is_none_or(|at| since(at) >= LEASE_RENEW)
}

/// Whether the client's own guard deletes `pod` at `now`: its deadline plus
/// [`DEADLINE_MARGIN`] passed, and no lease renewed within [`LEASE_TTL`].
fn guard_fires(pod: &PodRecord, renewed: Option<SystemTime>, now: SystemTime) -> bool {
    let held = renewed.is_some_and(|at| now.duration_since(at).unwrap_or_default() < LEASE_TTL);
    !held && until_deadline_at(pod, now).is_some_and(|left| left.is_zero())
}

/// Warns once, from the training metrics on the bus, when the job of `pod`
/// would need longer than its watchdog's deadline; then waits forever, so it
/// never ends the watch.
async fn warn_overrun(mut events: Receiver<Event>, pod: &PodRecord) {
    let mut pace = Pace::default();
    loop {
        match events.recv().await {
            Ok(Event::Metric(metric)) => {
                pace.add(&metric);
                if let Some(warning) = overrun(&pace, pod, SystemTime::now(), false) {
                    tracing::warn!("{warning}");
                    break;
                }
            },
            Ok(_) | Err(RecvError::Lagged(_)) => {},
            Err(RecvError::Closed) => break,
        }
    }
    std::future::pending::<()>().await;
}

/// The warning when, at `pace`, the job of `pod` still running at `now` would
/// end after its watchdog's deadline; `None` when it ends in time, for a kept
/// pod, or while the pace is unknown. With a `leased` pod the deadline only
/// applies once the client stops following the job.
fn overrun(pace: &Pace, pod: &PodRecord, now: SystemTime, leased: bool) -> Option<String> {
    if pod.keep {
        return None;
    }
    let deadline = pod.deadline_unix?;
    let left = pace.eta()?;
    let end = now.checked_add(left)?.duration_since(UNIX_EPOCH).ok()?;
    if end.as_secs() <= deadline {
        return None;
    }
    let hours = left.as_secs_f64() / 3600.0;
    let at = pod.deadline.as_deref().unwrap_or("its deadline");
    if leased {
        return Some(format!(
            "at this pace the job needs about {hours:.1}h more, past max_hours ({at}): \
             the pod stays while overbrainer follows the job; if it stops following, \
             the pod's watchdog deletes the pod"
        ));
    }
    Some(format!(
        "at this pace the job needs about {hours:.1}h more, but the pod's watchdog deletes it by {at} (max_hours): \
         raise the target's max_hours, or train with --keep-pod, to let it finish"
    ))
}

/// Acts on what [`watch_on_pod`] found for the run `run_id`: past the deadline
/// or the cost cap, deletes the pod and fails the run; after a failed watch, fails the run when
/// its pod is gone. Callers must not drop it midway (the CLI shields it from
/// Ctrl-C), since it may be deleting the pod.
///
/// # Errors
///
/// As [`follow`].
pub async fn settle_watch(
    ctx: &PodCtx<'_>,
    pod: &mut PodRecord,
    run_id: &str,
    watched: Watched,
) -> Result<Outcome, PodError> {
    match watched {
        Watched::Ended(ended) => match *ended {
            Ok(outcome) => Ok(outcome),
            Err(error) => Err(unreachable(ctx, pod, run_id, error.into()).await),
        },
        Watched::DeadlinePassed => Err(deadline_reached(ctx, pod, run_id).await),
        Watched::CostCapPassed => Err(cost_cap_reached(ctx, pod, run_id).await),
    }
}

/// Time left until the client's own deadline, `None` for a kept pod.
fn until_deadline(pod: &PodRecord) -> Option<Duration> {
    until_deadline_at(pod, SystemTime::now())
}

/// Time left at `now` until the client's own deadline (the watchdog's plus
/// [`DEADLINE_MARGIN`]), zero once it passed; `None` for a kept pod, or one
/// without a deadline.
pub(super) fn until_deadline_at(pod: &PodRecord, now: SystemTime) -> Option<Duration> {
    if pod.keep {
        return None;
    }
    let deadline = pod.deadline_unix?.saturating_add(DEADLINE_MARGIN.as_secs());
    let now = now
        .duration_since(UNIX_EPOCH)
        .map_or(0, |since| since.as_secs());
    Some(Duration::from_secs(deadline.saturating_sub(now)))
}

/// The pod outlived its deadline: deletes it (confirmed; recorded deleted by the
/// watchdog when it was already gone), forgets the run's keys (see
/// [`forget_keys`]) and fails the run.
async fn deadline_reached(ctx: &PodCtx<'_>, pod: &mut PodRecord, run_id: &str) -> PodError {
    if let Err(error) = remove(ctx, pod, DeleteReason::Deadline, DeletedBy::Client).await {
        return error;
    }
    forget_keys(ctx, run_id).await;
    fail_run(ctx.runs, run_id, MAX_HOURS_REACHED);
    PodError::DeadlineReached
}

/// The pod spent its cost cap: deletes it (confirmed), as its watchdog does
/// too, forgets the run's keys (see [`forget_keys`]) and fails the run.
async fn cost_cap_reached(ctx: &PodCtx<'_>, pod: &mut PodRecord, run_id: &str) -> PodError {
    if let Err(error) = remove(ctx, pod, DeleteReason::CostCap, DeletedBy::Client).await {
        return error;
    }
    forget_keys(ctx, run_id).await;
    fail_run(ctx.runs, run_id, MAX_COST_REACHED);
    PodError::CostCapReached
}

/// The watch failed with `error`: when the pod is gone, the run is failed with
/// that reason instead. A pod gone once its watchdog's deadline passed was
/// deleted by that watchdog: the run failed on `max_hours`, before the client's
/// own guard (the deadline plus [`DEADLINE_MARGIN`]) could fire.
async fn unreachable(
    ctx: &PodCtx<'_>,
    pod: &mut PodRecord,
    run_id: &str,
    error: PodError,
) -> PodError {
    match gone(ctx, pod).await {
        Ok(true) if past_cost_cap(pod, SystemTime::now()) => {
            fail_run(ctx.runs, run_id, MAX_COST_REACHED);
            PodError::CostCapReached
        },
        Ok(true) if past_deadline(pod, SystemTime::now()) => {
            fail_run(ctx.runs, run_id, MAX_HOURS_REACHED);
            PodError::DeadlineReached
        },
        Ok(true) => {
            let id = pod.pod_id.clone();
            if let Some(id) = &id {
                fail_run(ctx.runs, run_id, &format!("pod {id} no longer exists"));
            }
            id.map_or(error, PodError::PodGone)
        },
        Ok(false) => error,
        Err(check_error) => {
            tracing::warn!("cannot look the pod up: {check_error}");
            error
        },
    }
}

/// Whether the run's pod no longer exists; if so, `pod.json` records it deleted
/// by [`gone_by`]. Nothing is ever deleted here: see [`look_again`].
async fn gone(ctx: &PodCtx<'_>, pod: &mut PodRecord) -> Result<bool, PodError> {
    let Some(id) = pod.pod_id.clone() else {
        return Ok(true);
    };
    if pod.state == PodState::Deleted {
        return Ok(true);
    }
    if ctx.client.get_pod(&id).await?.is_some() || look_again(ctx, &id).await?.is_some() {
        return Ok(false);
    }
    mark_gone(ctx, pod, id, gone_by(pod, SystemTime::now())).await?;
    Ok(true)
}

/// Whether the watchdog's own deadline for `pod` has passed at `now`; never for
/// a kept pod, or one without a deadline. A pod found gone then was deleted by
/// its watchdog on `max_hours`.
#[must_use]
pub fn past_deadline(pod: &PodRecord, now: SystemTime) -> bool {
    let now = now
        .duration_since(UNIX_EPOCH)
        .map_or(0, |since| since.as_secs());
    !pod.keep && pod.deadline_unix.is_some_and(|deadline| now >= deadline)
}

/// Whether the cost cap of `pod` was spent at `now`: a pod found gone then was
/// deleted by its watchdog on `max_cost_usd`.
fn past_cost_cap(pod: &PodRecord, now: SystemTime) -> bool {
    pod.cost_cap()
        .is_some_and(|cap| unix_secs(now) >= cap.delete_at)
}

/// Why the run of `pod`, found gone at `now`, failed when a limit of its
/// watchdog had passed: [`MAX_COST_REACHED`] or [`MAX_HOURS_REACHED`].
#[must_use]
pub fn limit_reached(pod: &PodRecord, now: SystemTime) -> Option<&'static str> {
    if past_cost_cap(pod, now) {
        Some(MAX_COST_REACHED)
    } else if past_deadline(pod, now) {
        Some(MAX_HOURS_REACHED)
    } else {
        None
    }
}

/// Who deleted `pod`, found gone at `now` without overbrainer deleting it: its
/// watchdog once its deadline or its cost cap passed, otherwise nobody known.
fn gone_by(pod: &PodRecord, now: SystemTime) -> DeletedBy {
    if limit_reached(pod, now).is_some() {
        DeletedBy::Watchdog
    } else {
        DeletedBy::Unknown
    }
}

/// The pod `id` as the API shows it, or `None` when it is really gone: a 404
/// that [`look_again`] confirms. Nothing is ever deleted here.
///
/// # Errors
///
/// Returns [`PodError::Api`] when a look fails: the pod is then undetermined.
pub(super) async fn look_up(ctx: &PodCtx<'_>, id: &PodId) -> Result<Option<Pod>, PodError> {
    match ctx.client.get_pod(id).await? {
        Some(pod) => Ok(Some(pod)),
        None => look_again(ctx, id).await,
    }
}

/// Looks again at the pod `id`, whose first look just answered Runpod's 404:
/// `None` when it is really gone, that is [`GONE_LOOKS`] looks in a row answer
/// that 404, [`Timing::gone_interval`] apart, and the pod list no longer shows
/// it; otherwise the pod as last seen. A pod that only looks gone may be kept,
/// or hold results not yet retrieved, so it is never deleted: any answer showing
/// the pod means it is not gone, and any error leaves it undetermined (the error
/// is returned and nothing is recorded).
///
/// [`Timing::gone_interval`]: super::Timing::gone_interval
async fn look_again(ctx: &PodCtx<'_>, id: &PodId) -> Result<Option<Pod>, PodError> {
    for _ in 1..GONE_LOOKS {
        tokio::time::sleep(ctx.timing.gone_interval).await;
        if let Some(pod) = ctx.client.get_pod(id).await? {
            return Ok(Some(pod));
        }
    }
    Ok(ctx
        .client
        .list_pods()
        .await?
        .into_iter()
        .find(|listed| listed.id == *id))
}

/// What a batched look found of one pod (see [`look_up_all`]).
#[derive(Debug)]
pub(super) enum Look {
    /// The API shows the pod.
    Found(Box<Pod>),
    /// Confirmed gone: [`GONE_LOOKS`] 404s in a row and absent from the list.
    Gone,
    /// A look failed: the pod is undetermined. Holds the error's message.
    Failed(String),
}

/// [`look_up`] for several pods at once, with the same strength per pod but
/// one wait for all: every pod is looked at once; those answering Runpod's 404
/// are looked at again [`GONE_LOOKS`] - 1 times, [`Timing::gone_interval`]
/// apart, then checked against one pod list. No wait at all when every pod
/// answers. Nothing is ever deleted here. The looks are in the order of `ids`.
///
/// [`Timing::gone_interval`]: super::Timing::gone_interval
pub(super) async fn look_up_all(ctx: &PodCtx<'_>, ids: &[PodId]) -> Vec<Look> {
    let mut looks: Vec<Option<Look>> = ids.iter().map(|_| None).collect();
    let mut pending: Vec<usize> = (0..ids.len()).collect();
    for round in 0..GONE_LOOKS {
        if pending.is_empty() {
            break;
        }
        if round > 0 {
            tokio::time::sleep(ctx.timing.gone_interval).await;
        }
        pending = look_round(ctx, ids, &pending, &mut looks).await;
    }
    look_again_all(ctx, ids, &pending, &mut looks).await;
    looks
        .into_iter()
        .map(|look| {
            // Every pod gets a look above; one without is undetermined, never gone.
            look.unwrap_or_else(|| Look::Failed("the pod was not looked at".to_string()))
        })
        .collect()
}

/// Looks at the pods at `pending` once; returns those answering Runpod's 404.
async fn look_round(
    ctx: &PodCtx<'_>,
    ids: &[PodId],
    pending: &[usize],
    looks: &mut [Option<Look>],
) -> Vec<usize> {
    let mut missing = Vec::new();
    for &index in pending {
        match ctx.client.get_pod(&ids[index]).await {
            Ok(Some(pod)) => looks[index] = Some(Look::Found(Box::new(pod))),
            Ok(None) => missing.push(index),
            Err(error) => looks[index] = Some(Look::Failed(error.to_string())),
        }
    }
    missing
}

/// The last step of [`look_up_all`]: one pod list, which must not show the pods
/// at `pending` (each answered [`GONE_LOOKS`] 404s in a row) for them to be
/// gone.
async fn look_again_all(
    ctx: &PodCtx<'_>,
    ids: &[PodId],
    pending: &[usize],
    looks: &mut [Option<Look>],
) {
    if pending.is_empty() {
        return;
    }
    match ctx.client.list_pods().await {
        Ok(listed) => {
            for &index in pending {
                looks[index] = Some(match listed.iter().find(|pod| pod.id == ids[index]) {
                    Some(pod) => Look::Found(Box::new(pod.clone())),
                    None => Look::Gone,
                });
            }
        },
        Err(error) => {
            for &index in pending {
                looks[index] = Some(Look::Failed(error.to_string()));
            }
        },
    }
}

/// Records the pod `id` as deleted by `by`, then forgets the run's keys (see
/// [`forget_keys`]).
pub(super) async fn mark_gone(
    ctx: &PodCtx<'_>,
    pod: &mut PodRecord,
    id: PodId,
    by: DeletedBy,
) -> Result<(), PodError> {
    let now = SystemTime::now();
    let uptime = pod.uptime(now);
    pod.deleted(by, now);
    pod.save(ctx.runs)?;
    forget_keys(ctx, &pod.run_id).await;
    ctx.bus.publish(Event::PodStatus(PodStatus::Deleted {
        pod_id: id,
        uptime,
        estimated_spend: pod.estimated_spend,
    }));
    Ok(())
}

/// Saves the run `run_id` as `Failed` with `message`, best effort.
fn fail_run(runs: &Runs, run_id: &str, message: &str) {
    let result = runs.load(run_id).and_then(|mut run| {
        run.state = RunState::Failed;
        run.message = Some(message.to_string());
        runs.save(&run)
    });
    if let Err(error) = result {
        tracing::warn!("cannot record run {run_id} as failed: {error}");
    }
}

/// What became of a pod once its job ended.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Ending {
    /// Deleted after its results were retrieved.
    Deleted,
    /// Kept (`--keep-pod`), its results retrieved or not: only `overbrainer pod
    /// rm` deletes it.
    Kept,
    /// Its results were not retrieved and it is not kept: it stays for the
    /// watchdog's retrieve grace, or until `train attach` retrieves them.
    AwaitingRetrieval,
}

/// Ends the pod of the run `run` once its job ended, after copying the watchdog's
/// and the bootstrap's logs into `runs/<run-id>/.pod/`. Only when the results were `retrieved`
/// ([`Outcome::retrieved`]: every file the pod has was downloaded and verified)
/// is the "retrieved" marker written on the pod, first (so its watchdog deletes
/// it even if this client's delete fails), then the pod is deleted unless kept
/// and its deletion confirmed, other pods of the run are swept (any whose
/// deletion cannot be confirmed is recorded in `pod.json`), and, when none
/// is, the run's private client key is removed and its host key secret
/// deleted. Otherwise the pod is left to its watchdog's
/// retrieve grace (a kept pod stays kept), with a warning telling how to
/// retrieve the results again or remove the pod.
///
/// # Errors
///
/// Returns a [`PodError`] when the delete fails or cannot be confirmed (the pod
/// then stays in `pod.json`, as `Deleting`) or `pod.json` cannot be saved.
pub async fn end_pod(
    ctx: &PodCtx<'_>,
    pod: &mut PodRecord,
    executor: &SshExecutor,
    run: &RunRecord,
    retrieved: bool,
) -> Result<Ending, PodError> {
    fetch_watchdog_log(ctx, executor, run).await;
    if !retrieved {
        return unretrieved(ctx, pod, &run.id);
    }
    let marker = format!("{}/{RETRIEVED_MARKER}", run.remote_dir);
    if let Err(error) = executor.put_file(&marker, "").await {
        tracing::warn!(
            "cannot mark the results of run {} retrieved on the pod: {error}",
            run.id
        );
    }
    if pod.keep {
        pod.state = PodState::Kept;
        pod.save(ctx.runs)?;
        if let Some(id) = &pod.pod_id {
            ctx.bus
                .publish(Event::PodStatus(PodStatus::Kept { pod_id: id.clone() }));
        }
        return Ok(Ending::Kept);
    }
    remove(ctx, pod, DeleteReason::Retrieved, DeletedBy::Client).await?;
    sweep_then_forget(ctx, pod, &run.id).await;
    Ok(Ending::Deleted)
}

/// Sweeps the other pods of the run `run_id`, records in `pod.json` those whose
/// deletion cannot be confirmed, and forgets the run's keys when there is none
/// (see [`forget_keys`]): a stray keeps them even when `pod.json` cannot
/// record it, and so does a pod list that fails, since a stray may hide there.
async fn sweep_then_forget(ctx: &PodCtx<'_>, pod: &mut PodRecord, run_id: &str) {
    match sweep(ctx, run_id, None).await {
        Some(stray) if stray.is_empty() => forget_keys(ctx, run_id).await,
        Some(stray) => note_strays(ctx, pod, stray),
        None => tracing::debug!(
            "keeping the keys of run {run_id}: its pods could not be listed; the startup sweep deletes the secret later"
        ),
    }
}

/// The results of the run `run_id` were not retrieved: the pod stays, for its
/// watchdog's retrieve grace or, kept, until removed.
fn unretrieved(ctx: &PodCtx<'_>, pod: &mut PodRecord, run_id: &str) -> Result<Ending, PodError> {
    let name = pod
        .pod_id
        .as_ref()
        .map_or_else(|| "(none)".to_string(), ToString::to_string);
    let (stays, ending) = if pod.keep {
        pod.state = PodState::Kept;
        (
            "stays (--keep-pod), nothing deletes it automatically",
            Ending::Kept,
        )
    } else {
        pod.state = PodState::AwaitingRetrieval;
        (
            "stays until its watchdog's retrieve grace ends",
            Ending::AwaitingRetrieval,
        )
    };
    pod.save(ctx.runs)?;
    tracing::warn!("{}", unretrieved_line(ctx.runs, run_id, &name, stays));
    Ok(ending)
}

/// The warning when the results of `run_id`, a run or a job (an export or a
/// compare) of `runs`, were not retrieved from its pod `name`, which `stays`.
fn unretrieved_line(runs: &Runs, run_id: &str, name: &str, stays: &str) -> String {
    let subject = runs.subject(run_id);
    let retry = match runs.job_of() {
        Some((run, kind)) => kind.retry(&run),
        None => format!("retrieve them with `overbrainer train attach {run_id}`"),
    };
    format!(
        "the results of {subject} were not retrieved: pod {name} {stays}; {retry}, or remove the \
         pod with `overbrainer pod rm {run_id}`"
    )
}

/// The watchdog's log on the pod, in the run directory on the pod.
pub const WATCHDOG_LOG: &str = ".pod/watchdog.log";

/// Largest copy of the watchdog's or the bootstrap's log, in bytes.
const MAX_POD_FILE_LOG: u64 = 4 * 1024 * 1024;

/// Copies the watchdog's and the bootstrap's logs of the run into
/// `runs/<run-id>/.pod/`, best effort: the pod's own logs keep them too.
/// Each is read as bytes (never extracted from an archive the pod builds, so
/// the pod cannot make the copy land anywhere else), redacted line by line,
/// then written like the pod's log (see [`keep_file`]). A pod of an older
/// overbrainer has no bootstrap log: nothing is written for it.
async fn fetch_watchdog_log(ctx: &PodCtx<'_>, executor: &SshExecutor, run: &RunRecord) {
    let Ok(local) = ctx.runs.run_dir(&run.id) else {
        return;
    };
    for (entry, what) in [(WATCHDOG_LOG, "watchdog"), (BOOTSTRAP_LOG, "bootstrap")] {
        if let Err(error) = copy_pod_file(ctx, executor, &run.remote_dir, &local, entry).await {
            tracing::warn!("cannot copy the {what}'s log of run {}: {error}", run.id);
        }
    }
}

/// Copies `entry` of the run directory `remote_dir` on the pod to the run
/// directory `local`, redacted; nothing when it is empty or absent.
async fn copy_pod_file(
    ctx: &PodCtx<'_>,
    executor: &SshExecutor,
    remote_dir: &str,
    local: &Path,
    entry: &str,
) -> Result<(), PodError> {
    let remote = format!("{remote_dir}/{entry}");
    let bytes = executor.read_from(&remote, 0, MAX_POD_FILE_LOG).await?;
    if bytes.is_empty() {
        return Ok(());
    }
    let mut redactor = Redactor::new(ctx.client.log_secrets());
    let mut text = String::new();
    for line in String::from_utf8_lossy(&bytes).lines() {
        text.push_str(&redactor.line(line));
        text.push('\n');
    }
    keep_file(local, entry, text.as_bytes()).map_err(|source| PodError::Io {
        path: local.join(entry),
        source,
    })
}

/// Removes the private client key of a run whose pod is gone, best effort.
pub fn forget_client_key(runs: &Runs, run_id: &str) {
    let Ok(dir) = runs.run_dir(run_id) else {
        return;
    };
    let key = dir.join(SSH_DIR).join(CLIENT_KEY);
    match fs::remove_file(&key) {
        Ok(()) => {},
        Err(e) if e.kind() == io::ErrorKind::NotFound => {},
        Err(error) => tracing::warn!("cannot remove {}: {error}", key.display()),
    }
}

/// Connects again to the pod of the run `run` with `client`, for `train attach`
/// and `train cancel`: looks it up (its public port may have changed), rewrites
/// the run's ssh config and connects. `None` when the pod no longer exists (confirmed by
/// repeated looks, never by a delete), which `pod.json` then records.
///
/// # Errors
///
/// Returns [`PodError::PodStopped`] when the pod is stopped (`EXITED`),
/// [`PodError::NoEndpoint`] when it has no SSH endpoint yet (restarting), and
/// another [`PodError`] when the API, the files or the connection fail.
pub async fn reconnect(
    ctx: &PodCtx<'_>,
    pod: &mut PodRecord,
    run: &RunRecord,
    client: SshClient,
) -> Result<Option<SshExecutor>, PodError> {
    let Some(id) = pod.pod_id.clone() else {
        return Ok(None);
    };
    if pod.state == PodState::Deleted {
        return Ok(None);
    }
    let Some(remote) = look_up(ctx, &id).await? else {
        mark_gone(ctx, pod, id, gone_by(pod, SystemTime::now())).await?;
        return Ok(None);
    };
    if remote.status == RemoteStatus::Exited {
        return Err(PodError::PodStopped {
            pod_id: id,
            run_id: run.id.clone(),
        });
    }
    let direct = remote
        .direct()
        .ok_or_else(|| PodError::NoEndpoint(id.clone(), remote.status.name().to_string()))?;
    let endpoint = SshEndpoint {
        host: direct.host.clone(),
        port: direct.port,
        user: direct.username.clone(),
    };
    let ssh_dir = ctx.runs.run_dir(&run.id)?.join(SSH_DIR);
    let keys = PodKeys::new(
        ssh_dir.join(CLIENT_KEY),
        String::new(),
        pod.host_key.clone(),
        SecretString::from(String::new()),
    );
    let target = pod_endpoint(&ssh_dir, &alias(&run.id), &endpoint, &keys)?;
    let executor =
        SshExecutor::connect(&SshDestination::Pod(&target), pod_workdir(run), client).await?;
    pod.ssh = Some(endpoint);
    pod.save(ctx.runs)?;
    Ok(Some(executor))
}

/// Connects to the pod of the run `run` with `client` through the ssh config
/// the process following it wrote (the endpoint `pod.json` holds, for the
/// built-in client), for `train stop` while that process holds the project:
/// no Runpod API call, and nothing written locally, so the run directory stays
/// that process's alone. `None` when the run has no pod left or no ssh config
/// or endpoint yet.
///
/// # Errors
///
/// Returns a [`PodError`] when the run directory is invalid or the connection
/// fails.
pub async fn connect_followed(
    runs: &Runs,
    pod: &PodRecord,
    run: &RunRecord,
    client: SshClient,
) -> Result<Option<SshExecutor>, PodError> {
    if pod.pod_id.is_none() || pod.state == PodState::Deleted {
        return Ok(None);
    }
    let ssh_dir = runs.run_dir(&run.id)?.join(SSH_DIR);
    let config = ssh_dir.join(super::SSH_CONFIG);
    let (true, Some(endpoint)) = (config.is_file(), pod.ssh.as_ref()) else {
        return Ok(None);
    };
    let target = PodEndpoint {
        alias: alias(&run.id),
        config,
        host: endpoint.host.clone(),
        port: endpoint.port,
        user: endpoint.user.clone(),
        key: ssh_dir.join(CLIENT_KEY),
        host_key: pod.host_key.clone(),
    };
    let executor =
        SshExecutor::connect(&SshDestination::Pod(&target), pod_workdir(run), client).await?;
    Ok(Some(executor))
}

/// The work directory on the pod of the run `run`: the parent of its run
/// directory there.
fn pod_workdir(run: &RunRecord) -> &str {
    run.remote_dir
        .rsplit_once('/')
        .map_or(run.remote_dir.as_str(), |(parent, _)| parent)
}

/// The command reaching a kept pod of the run or export `run_id` of `runs`:
/// `ssh -F runs/<id>/ssh/config overbrainer-<id>`.
#[must_use]
pub fn ssh_command(runs: &Runs, run_id: &str) -> String {
    format!(
        "ssh -F {}/{SSH_DIR}/{} {}",
        runs.relative_dir(run_id),
        super::SSH_CONFIG,
        alias(run_id)
    )
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::events::EventBus;
    use crate::runpod::ApiError;
    use crate::runpod::{RunpodClient, Timing};
    use crate::train::TrainMetric;

    #[test]
    fn an_export_s_pod_messages_name_the_export_and_its_recovery()
    -> Result<(), Box<dyn std::error::Error>> {
        let project = tempfile::tempdir()?;
        let runs = Runs::new(project.path());
        let exports = runs.exports("r1")?;
        assert_eq!(
            unretrieved_line(&runs, "r1", "p1", "stays"),
            "the results of run r1 were not retrieved: pod p1 stays; retrieve them with \
             `overbrainer train attach r1`, or remove the pod with `overbrainer pod rm r1`"
        );
        assert_eq!(
            unretrieved_line(&exports, "e1", "p1", "stays"),
            "the results of export e1 of run r1 were not retrieved: pod p1 stays; export again \
             with `overbrainer export r1`, or remove the pod with `overbrainer pod rm e1`"
        );
        assert_eq!(
            ssh_command(&exports, "e1"),
            "ssh -F runs/r1/exports/e1/ssh/config overbrainer-e1"
        );
        assert_eq!(
            ssh_command(&runs, "r1"),
            "ssh -F runs/r1/ssh/config overbrainer-r1"
        );
        Ok(())
    }

    /// A pace of 1 step per 10 s at step 100 of 7300: 20 hours left.
    fn slow_pace() -> Pace {
        let mut pace = Pace::default();
        for (time, step) in [(0.0, 90), (100.0, 100)] {
            pace.add(&TrainMetric {
                time,
                step,
                epoch: None,
                max_steps: Some(7300),
                loss: Some(1.0),
                eval_loss: None,
                learning_rate: None,
                grad_norm: None,
            });
        }
        pace
    }

    /// A pod whose watchdog deletes it `hours` after `now`.
    fn pod_due_in(hours: u64, keep: bool, now: SystemTime) -> PodRecord {
        let mut pod = PodRecord::new("r1", keep, 1, "ssh-ed25519 AAAAhost");
        let deadline = now + Duration::from_secs(hours * 3600);
        pod.deadline_unix = deadline
            .duration_since(UNIX_EPOCH)
            .ok()
            .map(|since| since.as_secs());
        pod.deadline = Some("2026-09-29T19:54:50Z".into());
        pod
    }

    #[test]
    fn a_job_that_cannot_end_before_the_deadline_is_warned_about() {
        let now = SystemTime::now();
        assert_eq!(
            overrun(&slow_pace(), &pod_due_in(6, false, now), now, false).as_deref(),
            Some(
                "at this pace the job needs about 20.0h more, but the pod's watchdog deletes it by \
                 2026-09-29T19:54:50Z (max_hours): raise the target's max_hours, or train with \
                 --keep-pod, to let it finish"
            )
        );
    }

    #[test]
    fn a_followed_job_past_the_deadline_is_only_noted() {
        let now = SystemTime::now();
        let note = overrun(&slow_pace(), &pod_due_in(6, false, now), now, true);
        assert_eq!(
            note.as_deref(),
            Some(
                "at this pace the job needs about 20.0h more, past max_hours \
                 (2026-09-29T19:54:50Z): the pod stays while overbrainer follows the job; \
                 if it stops following, the pod's watchdog deletes the pod"
            )
        );
    }

    #[test]
    fn the_watchdog_holds_the_lease_as_long_as_the_client_counts_on() {
        let default = format!("OVERBRAINER_LEASE_TTL:-{}}}", LEASE_TTL.as_secs());
        assert!(super::super::watchdog_script().contains(&default));
    }

    #[test]
    fn the_lease_is_renewed_every_few_minutes_while_the_job_progresses() {
        let now = SystemTime::now();
        let ago = |minutes: u64| now - Duration::from_secs(minutes * 60);
        assert!(renewal_due(None, now, now));
        assert!(!renewal_due(Some(ago(1)), ago(1), now));
        assert!(renewal_due(Some(ago(5)), ago(1), now));
        // A job with no new metric for 30 minutes stalled: no renewal.
        assert!(!renewal_due(Some(ago(5)), ago(30), now));
        assert!(!renewal_due(None, ago(31), now));
    }

    #[test]
    fn the_client_guard_waits_for_the_lease_to_run_out() {
        let now = SystemTime::now();
        let ago = |minutes: u64| Some(now - Duration::from_secs(minutes * 60));
        // Deadline passed an hour ago, well beyond the margin.
        let late = pod_due_in(0, false, now - Duration::from_secs(3600));
        assert!(!guard_fires(&late, ago(1), now));
        assert!(!guard_fires(&late, ago(14), now));
        assert!(guard_fires(&late, ago(15), now));
        assert!(guard_fires(&late, None, now));
        // Before the deadline, or kept: never.
        assert!(!guard_fires(&pod_due_in(1, false, now), None, now));
        assert!(!guard_fires(
            &pod_due_in(0, true, now - Duration::from_secs(3600)),
            None,
            now
        ));
    }

    #[test]
    fn no_overrun_in_time_for_a_kept_pod_or_without_a_pace() {
        let now = SystemTime::now();
        assert_eq!(
            overrun(&slow_pace(), &pod_due_in(21, false, now), now, false),
            None
        );
        assert_eq!(
            overrun(&slow_pace(), &pod_due_in(6, true, now), now, false),
            None
        );
        assert_eq!(
            overrun(&Pace::default(), &pod_due_in(6, false, now), now, false),
            None
        );
    }

    #[test]
    fn the_cost_cap_asks_for_a_snapshot_then_names_why_the_pod_went() {
        let now = SystemTime::now();
        // Created an hour ago at $1/h, with a cap of $2: deleted at 2h, the
        // snapshot 15 min before (95% would leave only 6 min).
        let mut pod = pod_due_in(6, false, now);
        pod.created_unix = Some(unix_secs(now) - 3600);
        pod.cost_per_hour = Some(1.0);
        assert!(!cap_snapshot_due(&pod, now), "no cap");
        pod.max_cost_usd = Some(2.0);
        let hours = |h: f64| now + Duration::from_secs_f64(h * 3600.0);
        assert!(!cap_snapshot_due(&pod, hours(0.7)));
        assert!(cap_snapshot_due(&pod, hours(0.75)));
        assert_eq!(limit_reached(&pod, hours(0.95)), None);
        assert_eq!(limit_reached(&pod, hours(1.0)), Some(MAX_COST_REACHED));
        assert_eq!(gone_by(&pod, hours(1.0)), DeletedBy::Watchdog);
        assert_eq!(limit_reached(&pod, hours(6.0)), Some(MAX_COST_REACHED));
        pod.max_cost_usd = None;
        assert_eq!(limit_reached(&pod, hours(6.0)), Some(MAX_HOURS_REACHED));
        assert_eq!(gone_by(&pod, hours(1.0)), DeletedBy::Unknown);
    }

    #[test]
    fn the_cost_snapshot_says_why_it_is_due_with_the_amounts() {
        let now = SystemTime::now();
        let secs = unix_secs(now);
        // The live test: $0.05 at $0.24/h buys 12.5 min, less than the lead,
        // so the snapshot is due at once, at about 12% spent.
        let mut pod = billed_pod(now, Some(0.05));
        pod.cost_per_hour = Some(0.24);
        pod.created_unix = Some(secs - 90);
        assert_eq!(
            cap_snapshot_note(&pod, secs),
            "the pod reaches max_cost_usd $0.05 in 11 min ($0.01 spent so far): its job is \
             stopped with a snapshot 15 min before the cap"
        );
        // $20 at $1/h: 95% comes first, 1 hour before the cap.
        pod.max_cost_usd = Some(20.0);
        pod.cost_per_hour = Some(1.0);
        pod.created_unix = Some(secs - 68_400);
        assert_eq!(
            cap_snapshot_note(&pod, secs),
            "the pod spent $19.00 of max_cost_usd $20.00 (95%): its job is stopped with a \
             snapshot"
        );
        pod.cost_per_hour = None;
        assert_eq!(
            cap_snapshot_note(&pod, secs),
            "the pod neared max_cost_usd: its job is stopped with a snapshot"
        );
    }

    #[test]
    fn the_watchdog_asks_for_a_snapshot_as_early_as_the_client_counts_on() {
        let default = format!(
            "OVERBRAINER_SNAPSHOT_LEAD:-{}}}",
            super::super::SNAPSHOT_LEAD.as_secs()
        );
        assert!(super::super::watchdog_script().contains(&default));
    }

    /// A pod created an hour before `now` at $1/h.
    fn billed_pod(now: SystemTime, cap: Option<f64>) -> PodRecord {
        let mut pod = pod_due_in(6, false, now);
        pod.created_unix = Some(unix_secs(now) - 3600);
        pod.cost_per_hour = Some(1.0);
        pod.max_cost_usd = cap;
        pod
    }

    fn run_in(remote_dir: &str) -> RunRecord {
        RunRecord {
            id: "r1".into(),
            target: "gpu".into(),
            created: "2026-09-22T14:30:05Z".into(),
            remote_dir: remote_dir.into(),
            job: None,
            state: RunState::Preparing,
            message: None,
            snapshot: None,
            resumed_from: None,
            snapshots: true,
        }
    }

    #[tokio::test]
    async fn a_cost_cap_that_cannot_be_armed_refuses_the_run()
    -> Result<(), Box<dyn std::error::Error>> {
        let root = tempfile::tempdir()?;
        let executor = crate::exec::LocalExecutor::new(root.path())?;
        let now = SystemTime::now();
        let dir = root.path().join("r1");
        let run = run_in(&dir.to_string_lossy());
        // Armed: both times written where the watchdog reads them.
        arm_cost_cap(&executor, &billed_pod(now, Some(2.0)), &run).await?;
        let written = std::fs::read_to_string(dir.join(COST_CAP_FILE))?;
        assert_eq!(written, (unix_secs(now) + 3600).to_string());
        assert!(dir.join(SNAPSHOT_AT_FILE).is_file());
        // Nothing to arm without a cap.
        arm_cost_cap(&executor, &billed_pod(now, None), &run).await?;
        // No rate: refused.
        let mut unrated = billed_pod(now, Some(2.0));
        unrated.cost_per_hour = None;
        let refused = arm_cost_cap(&executor, &unrated, &run).await;
        assert!(
            matches!(refused, Err(PodError::CostCapUnset(_))),
            "{refused:?}"
        );
        // Files that cannot be written: refused.
        let blocked = root.path().join("blocked");
        std::fs::write(&blocked, "")?;
        let refused = arm_cost_cap(
            &executor,
            &billed_pod(now, Some(2.0)),
            &run_in(&blocked.to_string_lossy()),
        )
        .await;
        assert!(
            matches!(refused, Err(PodError::CostCapUnset(_))),
            "{refused:?}"
        );
        Ok(())
    }

    #[tokio::test]
    async fn the_client_deletes_the_pod_at_its_cost_cap_whatever_the_lease()
    -> Result<(), Box<dyn std::error::Error>> {
        let root = tempfile::tempdir()?;
        let executor = crate::exec::LocalExecutor::new(root.path())?;
        let now = SystemTime::now();
        let remote = root.path().join("r1").to_string_lossy().into_owned();
        let spent = billed_pod(now, Some(0.5));
        let mut lease = Lease::new(&spent, &remote);
        let fired = lease.tick(&executor).await;
        assert!(matches!(fired, Some(Watched::CostCapPassed)), "{fired:?}");
        // Its snapshot was asked for on the way.
        assert_eq!(
            std::fs::read_to_string(root.path().join("r1").join(SNAPSHOT_REQUEST))?,
            "cost"
        );
        let within = billed_pod(now, Some(20.0));
        let mut lease = Lease::new(&within, &remote);
        assert!(lease.tick(&executor).await.is_none());
        Ok(())
    }

    #[test]
    fn a_failed_run_keeps_only_the_status_of_an_api_answer() {
        let answered = PodError::Api(ApiError::Status {
            status: 500,
            message: "server text".into(),
            retry_after: None,
            capacity: false,
        });
        assert_eq!(run_message(&answered), "Runpod answered 500");
        assert_eq!(
            run_message(&PodError::Interrupted),
            "interrupted before the job started"
        );
    }

    /// The secret requests `server` received.
    async fn secret_calls(server: &wiremock::MockServer) -> Vec<String> {
        server
            .received_requests()
            .await
            .unwrap_or_default()
            .iter()
            .filter(|request| request.url.path().starts_with("/v2/account/secrets"))
            .map(|request| format!("{} {}", request.method, request.url.path()))
            .collect()
    }

    /// `sweep_then_forget` against a pod list answering `listed`: what it asked
    /// of the secrets API.
    async fn forget_after_sweep(
        listed: wiremock::ResponseTemplate,
    ) -> Result<Vec<String>, Box<dyn std::error::Error>> {
        use wiremock::matchers::{method, path};
        use wiremock::{Mock, MockServer, ResponseTemplate};
        let server = MockServer::start().await;
        Mock::given(method("GET"))
            .and(path("/v2/pods"))
            .respond_with(listed)
            .mount(&server)
            .await;
        Mock::given(method("GET"))
            .and(path("/v2/account/secrets"))
            .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
                "secrets": [{"id": "s1", "name": "overbrainer_host_key_r1",
                             "createdAt": "2026-10-01T00:00:00Z"}]
            })))
            .mount(&server)
            .await;
        Mock::given(method("DELETE"))
            .and(path("/v2/account/secrets/s1"))
            .respond_with(ResponseTemplate::new(204))
            .mount(&server)
            .await;
        let client = RunpodClient::new(
            &format!("{}/v2", server.uri()),
            &secrecy::SecretString::from("rp_key"),
        )?
        .with_policy(crate::retry::RetryPolicy {
            max_retries: 0,
            base: Duration::from_millis(1),
            cap: Duration::from_millis(2),
        });
        let project = tempfile::tempdir()?;
        let runs = Runs::new(project.path());
        let (bus, interrupted) = (EventBus::new(), std::sync::atomic::AtomicBool::new(false));
        let timing = Timing {
            poll: Duration::from_millis(5),
            ready_timeout: Duration::from_millis(400),
            preflight_timeout: Duration::from_millis(300),
            reconcile_waits: [Duration::from_millis(5); 2],
            delete_timeout: Duration::from_millis(200),
            gone_interval: Duration::from_millis(5),
        };
        let ctx = PodCtx {
            client: &client,
            runs: &runs,
            bus: &bus,
            timing: &timing,
            interrupted: &interrupted,
        };
        let mut pod = PodRecord::new("r1", false, 1, "ssh-ed25519 AAAAhost");
        sweep_then_forget(&ctx, &mut pod, "r1").await;
        Ok(secret_calls(&server).await)
    }

    #[tokio::test]
    async fn the_end_of_a_run_drops_its_secret_only_once_its_pods_are_listed()
    -> Result<(), Box<dyn std::error::Error>> {
        let empty =
            wiremock::ResponseTemplate::new(200).set_body_json(serde_json::json!({"pods": []}));
        let calls = forget_after_sweep(empty).await?;
        assert!(
            calls
                .iter()
                .any(|call| call == "DELETE /v2/account/secrets/s1"),
            "{calls:?}"
        );
        // A pod list that fails may hide a stray: the secret is left alone.
        let calls = forget_after_sweep(wiremock::ResponseTemplate::new(500)).await?;
        assert_eq!(calls, Vec::<String>::new());
        Ok(())
    }
}