entrenar/ecosystem/batuta/
pricing.rs1use serde::{Deserialize, Serialize};
4
5#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
7pub struct GpuPricing {
8 pub gpu_type: String,
10 pub hourly_rate: f64,
12 pub memory_gb: u32,
14 pub is_spot: bool,
16 pub provider: String,
18 pub region: String,
20}
21
22impl GpuPricing {
23 pub fn new(gpu_type: impl Into<String>, hourly_rate: f64, memory_gb: u32) -> Self {
25 Self {
26 gpu_type: gpu_type.into(),
27 hourly_rate,
28 memory_gb,
29 is_spot: false,
30 provider: "unknown".to_string(),
31 region: "unknown".to_string(),
32 }
33 }
34
35 pub fn with_spot(mut self, is_spot: bool) -> Self {
37 self.is_spot = is_spot;
38 self
39 }
40
41 pub fn with_provider(mut self, provider: impl Into<String>) -> Self {
43 self.provider = provider.into();
44 self
45 }
46
47 pub fn with_region(mut self, region: impl Into<String>) -> Self {
49 self.region = region.into();
50 self
51 }
52}
53
54#[derive(Debug, Clone)]
58pub struct FallbackPricing {
59 pricing: Vec<GpuPricing>,
61}
62
63impl Default for FallbackPricing {
64 fn default() -> Self {
65 Self::new()
66 }
67}
68
69impl FallbackPricing {
70 pub fn new() -> Self {
72 Self {
73 pricing: vec![
74 GpuPricing::new("a100-80gb", 3.00, 80)
75 .with_provider("generic")
76 .with_region("us-east-1"),
77 GpuPricing::new("a100-40gb", 2.50, 40)
78 .with_provider("generic")
79 .with_region("us-east-1"),
80 GpuPricing::new("v100", 2.00, 16).with_provider("generic").with_region("us-east-1"),
81 GpuPricing::new("t4", 0.50, 16).with_provider("generic").with_region("us-east-1"),
82 GpuPricing::new("l4", 0.75, 24).with_provider("generic").with_region("us-east-1"),
83 GpuPricing::new("a10g", 1.00, 24).with_provider("generic").with_region("us-east-1"),
84 GpuPricing::new("h100-80gb", 4.50, 80)
85 .with_provider("generic")
86 .with_region("us-east-1"),
87 ],
88 }
89 }
90
91 pub fn get_rate(&self, gpu_type: &str) -> Option<&GpuPricing> {
93 let normalized = gpu_type.to_lowercase().replace(['-', '_'], "");
94 self.pricing.iter().find(|p| {
95 let p_normalized = p.gpu_type.to_lowercase().replace(['-', '_'], "");
96 p_normalized == normalized
97 })
98 }
99
100 pub fn all_pricing(&self) -> &[GpuPricing] {
102 &self.pricing
103 }
104
105 pub fn set_rate(&mut self, pricing: GpuPricing) {
107 let normalized = pricing.gpu_type.to_lowercase().replace(['-', '_'], "");
108 if let Some(existing) = self.pricing.iter_mut().find(|p| {
109 let p_normalized = p.gpu_type.to_lowercase().replace(['-', '_'], "");
110 p_normalized == normalized
111 }) {
112 *existing = pricing;
113 } else {
114 self.pricing.push(pricing);
115 }
116 }
117}