use std::path::Path;
use std::sync::atomic::{AtomicBool, Ordering};
use std::time::{Duration, Instant, SystemTime};
use crate::events::{Event, EventBus};
use crate::exec::{Executor, SshExecutor};
use crate::runs::Runs;
use super::{
ApiError, Attempt, AttemptResult, CreateEnv, CreatePod, DeleteReason, DeletedBy, GpuRequest,
HOST_KEY_ENV, MIN_CUDA_VERSION, Mounts, NetworkMount, Pod, PodError, PodId, PodKeys, PodRecord,
PodSettings, PodStatus, RunpodClient, RunpodTarget, SshEndpoint, VOLUME_MOUNT, alias,
pod_command, pod_env, write_config,
};
const MAX_AMBIGUOUS: u32 = 2;
const VERDICT_BYTES: u64 = 4096;
const VERDICT_CHARS: usize = 200;
const BOOTSTRAP_FAILED: &str = "bootstrap: ";
const VERDICT_WARN_EVERY: Duration = Duration::from_secs(60);
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Timing {
pub poll: Duration,
pub ready_timeout: Duration,
pub preflight_timeout: Duration,
pub reconcile_waits: [Duration; 2],
pub delete_timeout: Duration,
pub gone_interval: Duration,
}
impl Timing {
#[must_use]
pub fn standard() -> Self {
Self {
poll: Duration::from_secs(5),
ready_timeout: Duration::from_secs(15 * 60),
preflight_timeout: Duration::from_secs(3 * 60),
reconcile_waits: [Duration::from_secs(5), Duration::from_secs(15)],
delete_timeout: Duration::from_secs(60),
gone_interval: Duration::from_secs(10),
}
}
}
#[derive(Debug, Clone, Copy)]
pub struct PodCtx<'a> {
pub client: &'a RunpodClient,
pub runs: &'a Runs,
pub bus: &'a EventBus,
pub timing: &'a Timing,
pub interrupted: &'a AtomicBool,
}
impl PodCtx<'_> {
fn publish(&self, status: PodStatus) {
self.bus.publish(Event::PodStatus(status));
}
fn check(&self) -> Result<(), PodError> {
if self.interrupted.load(Ordering::SeqCst) {
Err(PodError::Interrupted)
} else {
Ok(())
}
}
}
#[derive(Debug, Clone, Copy)]
pub struct PodPlan<'a> {
pub run_id: &'a str,
pub target: &'a RunpodTarget,
pub keys: &'a PodKeys,
pub ssh_dir: &'a Path,
pub workdir: &'a str,
pub api_url: &'a str,
}
#[derive(Debug)]
pub struct Provisioned {
pub pod_id: PodId,
pub executor: SshExecutor,
}
enum Wait {
NotReady(String),
Refused(String),
Failed(PodError),
}
enum Created {
Pod(Box<Pod>, AttemptResult),
Unavailable(String),
Ambiguous,
}
pub async fn provision(
ctx: &PodCtx<'_>,
plan: &PodPlan<'_>,
record: &mut PodRecord,
) -> Result<Provisioned, PodError> {
let result = walk(ctx, plan, record).await;
if result.is_err() {
after_failure(ctx, record).await;
}
result
}
async fn after_failure(ctx: &PodCtx<'_>, record: &mut PodRecord) {
let unclear: Vec<String> = record
.attempts
.iter()
.filter(|attempt| attempt.result == AttemptResult::Ambiguous)
.map(|attempt| attempt.name.clone())
.collect();
if unclear.is_empty() {
return;
}
let stray = sweep(ctx, &record.run_id.clone(), None).await;
note_strays(ctx, record, stray);
warn(&format!(
"the create calls {} got no clear answer from Runpod: a pod may still appear; check `overbrainer pod ls`",
unclear.join(", ")
));
}
pub(super) fn note_strays(ctx: &PodCtx<'_>, record: &mut PodRecord, stray: Vec<PodId>) {
if stray.is_empty() {
return;
}
for id in stray {
record.note_stray(id);
}
if let Err(error) = record.save(ctx.runs) {
warn(&format!("cannot record the stray pods: {}", chain(&error)));
}
}
async fn walk(
ctx: &PodCtx<'_>,
plan: &PodPlan<'_>,
record: &mut PodRecord,
) -> Result<Provisioned, PodError> {
let first = record.attempts.len();
let mut last_detail = None;
for gpu in &plan.target.gpu_types {
let mut ambiguous = 0;
loop {
ctx.check()?;
match create(ctx, plan, record, gpu).await? {
Created::Pod(pod, result) => match settle(ctx, plan, record, &pod, result).await? {
Some(provisioned) => return Ok(provisioned),
None => break,
},
Created::Unavailable(detail) => {
last_detail = Some(detail);
break;
},
Created::Ambiguous => {
ambiguous += 1;
if ambiguous >= MAX_AMBIGUOUS {
break;
}
},
}
}
}
let tried = &record.attempts[first..];
if !tried.is_empty()
&& tried
.iter()
.all(|attempt| attempt.result == AttemptResult::Ambiguous)
{
return Err(PodError::Unanswered);
}
Err(no_capacity(
tried,
last_detail,
plan.target.network_volume_id.is_some(),
))
}
fn no_capacity(tried: &[Attempt], detail: Option<String>, volume: bool) -> PodError {
let not_ready: Vec<&Attempt> = tried
.iter()
.filter(|attempt| attempt.result == AttemptResult::NotReady)
.collect();
let mut message = match not_ready.last() {
None => match detail {
Some(detail) => format!("no gpu_types entry could be placed: {detail}"),
None => "no gpu_types entry could be placed".to_string(),
},
Some(last) => {
let pods = match not_ready.len() {
1 => "1 pod was".to_string(),
count => format!("{count} pods were"),
};
let because = last.detail.as_deref().map_or_else(String::new, |reason| {
format!(", the last one because {reason}")
});
let mut message = format!(
"no gpu_types entry gave a ready pod: {pods} created but never became ready{because}"
);
if let Some(detail) = detail {
message = format!("{message}; the other types could not be placed: {detail}");
}
message
},
};
if volume {
message = format!("{message} (check network_volume_id and data_center_ids)");
}
PodError::NoCapacity(message)
}
async fn create(
ctx: &PodCtx<'_>,
plan: &PodPlan<'_>,
record: &mut PodRecord,
gpu: &str,
) -> Result<Created, PodError> {
let attempt = record
.begin_attempt(gpu, SystemTime::now(), plan.target.max_hours)
.clone();
record.save(ctx.runs)?;
ctx.publish(PodStatus::Creating {
name: attempt.name.clone(),
gpu_type: gpu.to_string(),
});
let request = request(plan, record.keep, &attempt);
let error = match ctx.client.create_pod(&request).await {
Ok(pod) => return Ok(Created::Pod(Box::new(pod), AttemptResult::Created)),
Err(error) => error,
};
let detail = error.to_string();
match skipped(&error) {
Some(result) => {
record.end_attempt(result, Some(detail.clone()));
record.save(ctx.runs)?;
ctx.publish(PodStatus::Unavailable {
gpu_type: gpu.to_string(),
reason: detail.clone(),
});
Ok(Created::Unavailable(detail))
},
None if error.is_ambiguous() => {
record.end_attempt(AttemptResult::Ambiguous, Some(detail));
record.save(ctx.runs)?;
tracing::warn!("no clear answer to the create call ({error}): looking for the pod");
Ok(match reconcile(ctx, record, &attempt.name).await? {
Some(pod) => Created::Pod(Box::new(pod), AttemptResult::Adopted),
None => Created::Ambiguous,
})
},
None => {
record.end_attempt(AttemptResult::Rejected, Some(detail));
record.save(ctx.runs)?;
Err(stop(error))
},
}
}
fn skipped(error: &ApiError) -> Option<AttemptResult> {
if error.is_capacity() {
Some(AttemptResult::Unavailable)
} else if error.status() == Some(403) {
Some(AttemptResult::Forbidden)
} else {
None
}
}
fn stop(error: ApiError) -> PodError {
match error.status() {
Some(402) => PodError::NoCredits,
Some(400 | 422) => PodError::Rejected(error.to_string()),
_ => error.into(),
}
}
fn request(plan: &PodPlan<'_>, keep: bool, attempt: &Attempt) -> CreatePod {
let target = plan.target;
let env = pod_env(&PodSettings {
run_id: plan.run_id,
workdir: plan.workdir,
deadline_unix: attempt.deadline_unix,
boot_grace: target.boot_grace,
retrieve_grace: target.retrieve_grace,
keep,
api_url: plan.api_url,
authorized_key: &plan.keys.client_public,
});
CreatePod {
name: attempt.name.clone(),
image: target.image.clone(),
cloud: "SECURE",
gpu: GpuRequest {
id: attempt.gpu_type.clone(),
count: target.gpu_count,
min_cuda_version: MIN_CUDA_VERSION,
},
disk: target.container_disk_gb,
ports: vec!["22/tcp".to_string()],
start_ssh: false,
data_center_ids: (!target.data_center_ids.is_empty())
.then(|| target.data_center_ids.clone()),
mounts: target.network_volume_id.as_ref().map(|volume| Mounts {
network: vec![NetworkMount {
volume_id: volume.clone(),
path: VOLUME_MOUNT.to_string(),
}],
}),
env: CreateEnv {
plain: env,
host_key_name: HOST_KEY_ENV,
host_key: plan.keys.host_private().clone(),
},
cmd: pod_command(),
}
}
async fn settle(
ctx: &PodCtx<'_>,
plan: &PodPlan<'_>,
record: &mut PodRecord,
pod: &Pod,
result: AttemptResult,
) -> Result<Option<Provisioned>, PodError> {
record.created(pod, result, SystemTime::now());
let saved = record.save(ctx.runs);
if saved.is_err() {
let removed = remove(ctx, record, DeleteReason::Requested, DeletedBy::Client).await;
saved_then_removed(saved, removed)?;
}
ctx.publish(PodStatus::Created {
pod_id: pod.id.clone(),
gpu_type: record.gpu_type.clone().unwrap_or_default(),
data_center: record.data_center_id.clone(),
cost_per_hour: record.cost_per_hour,
});
match ready(ctx, plan, record, &pod.id).await {
Ok(provisioned) => Ok(Some(provisioned)),
Err(Wait::NotReady(reason)) => {
not_ready(ctx, record, &pod.id, reason).await?;
Ok(None)
},
Err(Wait::Refused(reason)) => Err(refuse(ctx, record, &pod.id, reason).await),
Err(Wait::Failed(error)) => Err(abandon(ctx, record, error).await),
}
}
async fn not_ready(
ctx: &PodCtx<'_>,
record: &mut PodRecord,
id: &PodId,
reason: String,
) -> Result<(), PodError> {
warn(&format!("pod {id} is not ready: {reason}"));
record.end_attempt(AttemptResult::NotReady, Some(reason));
let saved = record.save(ctx.runs);
let removed = remove(ctx, record, DeleteReason::NotReady, DeletedBy::Client).await;
saved_then_removed(saved, removed)?;
record.forget_pod();
record.save(ctx.runs)?;
Ok(())
}
fn saved_then_removed(
saved: Result<(), crate::runs::RunsError>,
removed: Result<(), PodError>,
) -> Result<(), PodError> {
match (saved, removed) {
(Err(save), Err(delete)) => {
warn(&format!("cannot save pod.json: {}", chain(&save)));
Err(delete)
},
(_, Err(delete)) => Err(delete),
(Err(save), Ok(())) => Err(save.into()),
(Ok(()), Ok(())) => Ok(()),
}
}
async fn refuse(ctx: &PodCtx<'_>, record: &mut PodRecord, id: &PodId, reason: String) -> PodError {
record.end_attempt(AttemptResult::Refused, Some(reason.clone()));
let saved = record.save(ctx.runs);
let (delete_reason, refusal) = refusal(id, reason);
let removed = remove(ctx, record, delete_reason, DeletedBy::Client).await;
if let Err(error) = saved_then_removed(saved, removed) {
return error;
}
refusal
}
fn refusal(id: &PodId, reason: String) -> (DeleteReason, PodError) {
match reason.strip_prefix(BOOTSTRAP_FAILED) {
Some(bootstrap) => (
DeleteReason::BootstrapFailed,
PodError::BootstrapFailed {
pod_id: id.clone(),
reason: bootstrap.to_string(),
},
),
None => (
DeleteReason::Refused,
PodError::WatchdogRefused {
pod_id: id.clone(),
reason,
},
),
}
}
async fn abandon(ctx: &PodCtx<'_>, record: &mut PodRecord, error: PodError) -> PodError {
let reason = if matches!(error, PodError::Interrupted) {
DeleteReason::Interrupted
} else {
DeleteReason::Requested
};
if let Err(delete_error) = remove(ctx, record, reason, DeletedBy::Client).await {
warn(&delete_error.to_string());
}
error
}
fn warn(message: &str) {
tracing::warn!("{message}");
}
async fn ready(
ctx: &PodCtx<'_>,
plan: &PodPlan<'_>,
record: &mut PodRecord,
id: &PodId,
) -> Result<Provisioned, Wait> {
let (executor, endpoint) = reach(ctx, plan, record, id).await?;
let verdict = format!("{}/{}/.pod/watchdog", executor.workdir(), plan.run_id);
preflight(ctx, &executor, &verdict, id).await?;
let now = SystemTime::now();
record.ready(endpoint, now);
record
.save(ctx.runs)
.map_err(|error| Wait::Failed(error.into()))?;
ctx.publish(PodStatus::Ready {
pod_id: id.clone(),
after: record.uptime(now).unwrap_or_default(),
deadline: record.deadline.clone(),
});
let stray = sweep(ctx, plan.run_id, Some(id)).await;
note_strays(ctx, record, stray);
Ok(Provisioned {
pod_id: id.clone(),
executor,
})
}
async fn reach(
ctx: &PodCtx<'_>,
plan: &PodPlan<'_>,
record: &mut PodRecord,
id: &PodId,
) -> Result<(SshExecutor, SshEndpoint), Wait> {
let started = Instant::now();
let alias = alias(plan.run_id);
let mut last = "no SSH endpoint yet".to_string();
loop {
ctx.check().map_err(Wait::Failed)?;
if started.elapsed() >= ctx.timing.ready_timeout {
return Err(Wait::NotReady(format!(
"not reachable over SSH after {}s: {last}",
ctx.timing.ready_timeout.as_secs()
)));
}
let pod = alive(ctx, id).await?;
record.note_rate(&pod);
if let Some(direct) = pod.direct() {
let endpoint = SshEndpoint {
host: direct.host.clone(),
port: direct.port,
user: direct.username.clone(),
};
let config =
write_config(plan.ssh_dir, &alias, &endpoint, plan.keys).map_err(Wait::Failed)?;
match SshExecutor::connect(&alias, plan.workdir, Some(&config)).await {
Ok(executor) => return Ok((executor, endpoint)),
Err(error) => last = chain(&error),
}
}
tokio::time::sleep(ctx.timing.poll).await;
}
}
async fn alive(ctx: &PodCtx<'_>, id: &PodId) -> Result<Pod, Wait> {
let pod = ctx
.client
.get_pod(id)
.await
.map_err(|error| Wait::Failed(error.into()))?
.ok_or_else(|| Wait::NotReady("the pod disappeared".to_string()))?;
if pod.status.is_dead() {
return Err(Wait::NotReady(format!("the pod is {}", pod.status.name())));
}
Ok(pod)
}
async fn preflight(
ctx: &PodCtx<'_>,
executor: &SshExecutor,
path: &str,
id: &PodId,
) -> Result<(), Wait> {
let started = Instant::now();
let mut warned: Option<Instant> = None;
loop {
ctx.check().map_err(Wait::Failed)?;
match executor.read_from(path, 0, VERDICT_BYTES).await {
Ok(bytes) => {
let text = one_line(&String::from_utf8_lossy(&bytes), VERDICT_CHARS);
if text == "ready" {
return Ok(());
}
if let Some(reason) = text.strip_prefix("failed ") {
return Err(Wait::Refused(reason.to_string()));
}
},
Err(error) => {
if warned.is_none_or(|at| at.elapsed() >= VERDICT_WARN_EVERY) {
tracing::warn!("cannot read the watchdog's verdict: {}", chain(&error));
warned = Some(Instant::now());
}
},
}
alive(ctx, id).await?;
if started.elapsed() >= ctx.timing.preflight_timeout {
return Err(Wait::Refused(format!(
"no verdict within {}s",
ctx.timing.preflight_timeout.as_secs()
)));
}
tokio::time::sleep(ctx.timing.poll).await;
}
}
pub(super) fn one_line(text: &str, max: usize) -> String {
text.chars()
.filter(|c| c.is_ascii_graphic() || *c == ' ')
.take(max)
.collect::<String>()
.trim()
.to_string()
}
pub async fn remove(
ctx: &PodCtx<'_>,
record: &mut PodRecord,
reason: DeleteReason,
by: DeletedBy,
) -> Result<(), PodError> {
let Some(id) = record.pod_id.clone() else {
return Ok(());
};
ctx.publish(PodStatus::Deleting {
pod_id: id.clone(),
reason,
});
record.state = super::PodState::Deleting;
let saved = record.save(ctx.runs);
let gone = matches!(ctx.client.get_pod(&id).await, Ok(None));
let found = match delete_confirmed(ctx, &id).await {
Ok(found) => found,
Err(error) => return saved_then_removed(saved, Err(error)),
};
let now = SystemTime::now();
let uptime = record.uptime(now);
record.deleted(
if gone && !found {
DeletedBy::Watchdog
} else {
by
},
now,
);
let recorded = record.save(ctx.runs);
ctx.publish(PodStatus::Deleted {
pod_id: id,
uptime,
estimated_spend: record.estimated_spend,
});
recorded?;
saved?;
Ok(())
}
pub(super) async fn delete_confirmed(ctx: &PodCtx<'_>, id: &PodId) -> Result<bool, PodError> {
let found = ctx.client.delete_pod(id).await?;
wait_gone(ctx, id).await?;
Ok(found)
}
pub(crate) async fn wait_gone(ctx: &PodCtx<'_>, id: &PodId) -> Result<(), PodError> {
let started = Instant::now();
loop {
if ctx.client.get_pod(id).await?.is_none() {
return Ok(());
}
if started.elapsed() >= ctx.timing.delete_timeout {
return Err(PodError::NotDeleted(id.clone()));
}
tokio::time::sleep(ctx.timing.poll).await;
}
}
async fn reconcile(
ctx: &PodCtx<'_>,
record: &mut PodRecord,
name: &str,
) -> Result<Option<Pod>, PodError> {
for wait in ctx.timing.reconcile_waits {
nap(ctx, wait).await;
let interrupted = ctx.check().is_err();
let mine: Vec<Pod> = ctx
.client
.list_pods()
.await?
.into_iter()
.filter(|pod| pod.run_id() == Some(record.run_id.as_str()))
.collect();
if let Some(found) = mine.iter().find(|pod| pod.name == name) {
let mut stray = Vec::new();
for duplicate in mine.iter().filter(|pod| pod.id != found.id) {
if !delete_duplicate(ctx, &duplicate.id).await {
stray.push(duplicate.id.clone());
}
}
note_strays(ctx, record, stray);
return Ok(Some(found.clone()));
}
if interrupted {
return Err(PodError::Interrupted);
}
}
Ok(None)
}
async fn nap(ctx: &PodCtx<'_>, wait: Duration) {
let started = Instant::now();
while ctx.check().is_ok() {
let left = wait.saturating_sub(started.elapsed());
if left.is_zero() {
return;
}
tokio::time::sleep(left.min(ctx.timing.poll)).await;
}
}
pub async fn sweep(ctx: &PodCtx<'_>, run_id: &str, keep: Option<&PodId>) -> Vec<PodId> {
let pods = match ctx.client.list_pods().await {
Ok(pods) => pods,
Err(error) => {
tracing::warn!("cannot look for duplicate pods of run {run_id}: {error}");
return Vec::new();
},
};
let mut stray = Vec::new();
for pod in pods {
if pod.run_id() == Some(run_id)
&& Some(&pod.id) != keep
&& !delete_duplicate(ctx, &pod.id).await
{
stray.push(pod.id);
}
}
stray
}
async fn delete_duplicate(ctx: &PodCtx<'_>, id: &PodId) -> bool {
ctx.publish(PodStatus::Deleting {
pod_id: id.clone(),
reason: DeleteReason::Duplicate,
});
match delete_confirmed(ctx, id).await {
Ok(_) => true,
Err(error) => {
tracing::warn!(
"cannot confirm the deletion of duplicate pod {id}: {error}; its watchdog deletes it after its boot grace"
);
false
},
}
}
#[must_use]
pub fn chain(error: &dyn std::error::Error) -> String {
let mut parts = vec![error.to_string()];
let mut source = error.source();
while let Some(cause) = source {
parts.push(cause.to_string());
source = cause.source();
}
parts.join(": ")
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn a_verdict_is_one_short_line_of_printable_ascii() {
assert_eq!(one_line("ready\n", VERDICT_CHARS), "ready");
assert_eq!(
one_line(
"failed http_403\u{1b}[31m\r\nsecond line\u{7}",
VERDICT_CHARS
),
"failed http_403[31msecond line"
);
assert_eq!(one_line("failed caf\u{e9}", VERDICT_CHARS), "failed caf");
assert_eq!(one_line("run\u{1b}[31m\nx\u{e9}", 64), "run[31mx");
let long = format!("failed {}", "x".repeat(5000));
assert_eq!(one_line(&long, VERDICT_CHARS).len(), VERDICT_CHARS);
assert_eq!(one_line(&long, 64).len(), 64);
}
#[test]
fn a_failed_bootstrap_has_its_own_delete_reason() -> Result<(), crate::runpod::InvalidPodId> {
let id = PodId::new("p1")?;
let (reason, error) = refusal(&id, "bootstrap: cannot start sshd".to_string());
assert_eq!(reason, DeleteReason::BootstrapFailed);
assert_eq!(reason.describe(), "its bootstrap failed");
assert!(matches!(error, PodError::BootstrapFailed { .. }), "{error}");
let (reason, error) = refusal(&id, "http_403".to_string());
assert_eq!(reason, DeleteReason::Refused);
assert!(matches!(error, PodError::WatchdogRefused { .. }), "{error}");
Ok(())
}
}