stormchaser_engine/
resource_utils.rs

1/// Parse cpu.
2pub fn parse_cpu(cpu_str: &str) -> Option<f64> {
3    if let Some(m_idx) = cpu_str.find('m') {
4        if let Ok(m_cores) = cpu_str[..m_idx].parse::<f64>() {
5            return Some(m_cores / 1000.0);
6        }
7    } else if let Ok(cores) = cpu_str.parse::<f64>() {
8        return Some(cores);
9    }
10    None
11}
12
13/// Parse memory.
14pub fn parse_memory(memory_str: &str) -> Option<i64> {
15    let mem = memory_str.to_lowercase();
16    if let Some(idx) = mem.find(|c: char| c.is_alphabetic()) {
17        let val = mem[..idx].trim().parse::<i64>().ok()?;
18        let suffix = &mem[idx..];
19        match suffix {
20            "k" | "ki" => Some(val * 1024),
21            "m" | "mi" => Some(val * 1024 * 1024),
22            "g" | "gi" => Some(val * 1024 * 1024 * 1024),
23            "t" | "ti" => Some(val * 1024 * 1024 * 1024 * 1024),
24            _ => None,
25        }
26    } else {
27        mem.trim().parse::<i64>().ok()
28    }
29}
30
31use serde_json::Value;
32
33/// Extracts CPU (in cores) and memory (in bytes) requirements from a step's specification.
34pub fn get_step_resource_requirements(step_type: &str, spec: &Value) -> (f64, i64) {
35    let mut cpu_req = 0.0;
36    let mut mem_req = 0;
37
38    if step_type == "RunContainer" || step_type == "RunK8sJob" {
39        if let Some(cpu) = spec.get("cpu").and_then(|v| v.as_str()) {
40            cpu_req = parse_cpu(cpu).unwrap_or(0.0);
41        }
42        if let Some(mem) = spec.get("memory").and_then(|v| v.as_str()) {
43            mem_req = parse_memory(mem).unwrap_or(0);
44        }
45
46        if step_type == "RunK8sJob" {
47            // Also check resources block for K8sJobSpec
48            if let Some(resources) = spec.get("resources") {
49                if let Some(requests) = resources.get("requests") {
50                    if let Some(cpu) = requests.get("cpu").and_then(|v| v.as_str()) {
51                        cpu_req = parse_cpu(cpu).unwrap_or(cpu_req);
52                    }
53                    if let Some(mem) = requests.get("memory").and_then(|v| v.as_str()) {
54                        mem_req = parse_memory(mem).unwrap_or(mem_req);
55                    }
56                }
57            }
58        }
59    }
60
61    (cpu_req, mem_req)
62}
stormchaser_engine/resource_utils.rs

stormchaser_engine/
resource_utils.rs
