entrenar/ecosystem/batuta/
client.rs

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

entrenar/ecosystem/batuta/
client.rs
