use serde_json::Value;
use crate::backend::{
json_array, json_str, json_u64, DEVICE_TYPE_GPU, DEVICE_TYPE_IGPU, KEY_DEVICES, KEY_DEVICE_ID,
KEY_MEMORY_FREE_BYTES, KEY_NAME, KEY_TYPE,
};
use crate::error::Result;
use super::runtime_action_failed;
const UNKNOWN_DEVICE_COMPONENT: &str = "unknown";
const PARSE_BACKEND_RESIDENCY_FAILED: &str = "failed to parse backend residency";
#[derive(Debug, Clone, PartialEq, Eq)]
pub(super) struct GpuLeaseDevice {
pub key: String,
pub free_bytes: Option<u64>,
}
pub(super) fn parse_gpu_lease_devices(raw: &str) -> Result<Vec<GpuLeaseDevice>> {
let value: Value = serde_json::from_str(raw)
.map_err(|error| runtime_action_failed(PARSE_BACKEND_RESIDENCY_FAILED, error))?;
let Some(devices) = json_array(&value, KEY_DEVICES) else {
return Ok(Vec::new());
};
let mut out = Vec::with_capacity(devices.len());
for (index, device) in devices.iter().enumerate() {
if !is_gpu_lease_device(device) {
continue;
}
out.push(GpuLeaseDevice {
key: sanitize_lease_component(&gpu_lease_device_identity(device, index)),
free_bytes: json_u64(device, KEY_MEMORY_FREE_BYTES),
});
}
sort_unique_lease_devices(&mut out);
Ok(out)
}
fn is_gpu_lease_device(device: &Value) -> bool {
matches!(
json_str(device, KEY_TYPE),
Some(DEVICE_TYPE_GPU | DEVICE_TYPE_IGPU)
)
}
fn gpu_lease_device_identity(device: &Value, index: usize) -> String {
json_str(device, KEY_DEVICE_ID)
.or_else(|| json_str(device, KEY_NAME))
.map(str::to_string)
.unwrap_or_else(|| format!("gpu-{index}"))
}
fn sort_unique_lease_devices(devices: &mut Vec<GpuLeaseDevice>) {
devices.sort_by(|left, right| left.key.cmp(&right.key));
devices.dedup_by(|left, right| left.key == right.key);
}
fn sanitize_lease_component(value: &str) -> String {
let sanitized = value
.chars()
.map(|ch| {
if ch.is_ascii_alphanumeric() || ch == '-' || ch == '_' {
ch
} else {
'_'
}
})
.collect::<String>();
if sanitized.is_empty() {
UNKNOWN_DEVICE_COMPONENT.to_string()
} else {
sanitized
}
}
#[cfg(test)]
#[path = "../../tests/runtime/residency/devices_tests.rs"]
mod devices_tests;