use serde::{Deserialize, Serialize};
use std::time::Duration;
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct QueueState {
pub queue_depth: u32,
pub avg_wait_seconds: u64,
pub available_gpus: u32,
pub total_gpus: u32,
pub eta_seconds: Option<u64>,
}
impl QueueState {
pub fn new(queue_depth: u32, available_gpus: u32, total_gpus: u32) -> Self {
Self { queue_depth, avg_wait_seconds: 0, available_gpus, total_gpus, eta_seconds: None }
}
pub fn with_avg_wait(mut self, seconds: u64) -> Self {
self.avg_wait_seconds = seconds;
self
}
pub fn with_eta(mut self, seconds: u64) -> Self {
self.eta_seconds = Some(seconds);
self
}
pub fn is_available(&self) -> bool {
self.available_gpus > 0
}
pub fn utilization(&self) -> f64 {
if self.total_gpus == 0 {
return 1.0;
}
1.0 - (f64::from(self.available_gpus) / f64::from(self.total_gpus))
}
}
pub fn adjust_eta(base_eta_seconds: u64, queue_state: &QueueState) -> Duration {
let mut adjusted = base_eta_seconds;
if !queue_state.is_available() {
let queue_wait = if queue_state.avg_wait_seconds > 0 {
u64::from(queue_state.queue_depth) * queue_state.avg_wait_seconds
} else {
u64::from(queue_state.queue_depth) * 300
};
adjusted += queue_wait;
}
if let Some(eta) = queue_state.eta_seconds {
adjusted = adjusted.max(eta);
}
let utilization = queue_state.utilization();
if utilization > 0.8 {
let multiplier = 1.0 + (utilization - 0.8) * 2.0; adjusted = (adjusted as f64 * multiplier) as u64;
}
Duration::from_secs(adjusted)
}