1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
//! Batuta client for pricing and queue services.
use std::time::Duration;
use super::error::BatutaError;
use super::pricing::{FallbackPricing, GpuPricing};
use super::queue::QueueState;
/// Client for interacting with Batuta pricing and queue services.
#[derive(Debug, Clone)]
pub struct BatutaClient {
/// Base URL for Batuta API (None if using fallback only)
base_url: Option<String>,
/// Fallback pricing when Batuta is unavailable
fallback: FallbackPricing,
/// Connection timeout
timeout: Duration,
/// Whether Batuta service is available
service_available: bool,
}
impl Default for BatutaClient {
fn default() -> Self {
Self::new()
}
}
impl BatutaClient {
/// Create a new Batuta client with fallback pricing only.
pub fn new() -> Self {
Self {
base_url: None,
fallback: FallbackPricing::new(),
timeout: Duration::from_secs(5),
service_available: false,
}
}
/// Create a client connected to a Batuta instance.
pub fn with_url(url: impl Into<String>) -> Self {
Self {
base_url: Some(url.into()),
fallback: FallbackPricing::new(),
timeout: Duration::from_secs(5),
service_available: true,
}
}
/// Set connection timeout.
pub fn with_timeout(mut self, timeout: Duration) -> Self {
self.timeout = timeout;
self
}
/// Set custom fallback pricing.
pub fn with_fallback(mut self, fallback: FallbackPricing) -> Self {
self.fallback = fallback;
self
}
/// Check if connected to a live Batuta service.
pub fn is_connected(&self) -> bool {
self.base_url.is_some() && self.service_available
}
/// Get hourly rate for a GPU type.
///
/// Returns live pricing from Batuta if available, otherwise fallback pricing.
pub fn get_hourly_rate(&self, gpu_type: &str) -> Result<GpuPricing, BatutaError> {
// If we have a live connection, try to fetch from Batuta
if let Some(_url) = &self.base_url {
// In a real implementation, this would make an HTTP request
// For now, we simulate by returning fallback
// FUTURE(batuta-api): HTTP client integration pending API finalization
}
// Use fallback pricing
self.fallback
.get_rate(gpu_type)
.cloned()
.ok_or_else(|| BatutaError::UnknownGpuType(gpu_type.to_string()))
}
/// Get current queue depth.
///
/// Returns queue state from Batuta if available, otherwise returns
/// an optimistic default (no queue).
pub fn get_queue_depth(&self, gpu_type: &str) -> Result<QueueState, BatutaError> {
// Validate GPU type exists
if self.fallback.get_rate(gpu_type).is_none() {
return Err(BatutaError::UnknownGpuType(gpu_type.to_string()));
}
// If we have a live connection, try to fetch from Batuta
if let Some(_url) = &self.base_url {
// In a real implementation, this would make an HTTP request
// FUTURE(batuta-api): HTTP client integration pending API finalization
}
// Return optimistic default (no queue)
Ok(QueueState::new(0, 4, 4))
}
/// Get both pricing and queue state in one call.
pub fn get_status(&self, gpu_type: &str) -> Result<(GpuPricing, QueueState), BatutaError> {
let pricing = self.get_hourly_rate(gpu_type)?;
let queue = self.get_queue_depth(gpu_type)?;
Ok((pricing, queue))
}
/// Estimate total cost for a job.
pub fn estimate_cost(&self, gpu_type: &str, hours: f64) -> Result<f64, BatutaError> {
let pricing = self.get_hourly_rate(gpu_type)?;
Ok(pricing.hourly_rate * hours)
}
/// Get the cheapest GPU that meets memory requirements.
pub fn cheapest_gpu(&self, min_memory_gb: u32) -> Option<&GpuPricing> {
self.fallback.all_pricing().iter().filter(|p| p.memory_gb >= min_memory_gb).min_by(
|a, b| a.hourly_rate.partial_cmp(&b.hourly_rate).unwrap_or(std::cmp::Ordering::Equal),
)
}
}