Skip to main content

mesh_llm_native_runtime/
host.rs

1use crate::NativeRuntimeBackendKind;
2use serde::{Deserialize, Serialize};
3use std::collections::{BTreeMap, BTreeSet};
4
5#[derive(Clone, Debug, Default, Eq, PartialEq, Serialize, Deserialize)]
6pub struct HostGpuProbe {
7    pub source: String,
8    #[serde(default, skip_serializing_if = "Option::is_none")]
9    pub path: Option<String>,
10    #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
11    pub fields: BTreeMap<String, String>,
12    #[serde(default, skip_serializing_if = "Vec::is_empty")]
13    pub raw_lines: Vec<String>,
14}
15
16#[derive(Clone, Debug, Default, Eq, PartialEq, Serialize, Deserialize)]
17pub struct HostGpuProfile {
18    pub display_name: String,
19    pub backend_device: Option<String>,
20    pub stable_id: Option<String>,
21    pub vram_bytes: Option<u64>,
22    pub unified_memory: bool,
23    #[serde(default, skip_serializing_if = "Option::is_none")]
24    pub probe: Option<HostGpuProbe>,
25    #[serde(default, skip_serializing_if = "Option::is_none")]
26    pub cuda_sm: Option<String>,
27    #[serde(default, skip_serializing_if = "Option::is_none")]
28    pub rocm_gfx: Option<String>,
29}
30
31#[derive(Clone, Debug, Default, Eq, PartialEq, Serialize, Deserialize)]
32pub struct HostCudaProfile {
33    #[serde(default)]
34    pub toolkit_majors: BTreeSet<u32>,
35    #[serde(default, skip_serializing_if = "Option::is_none")]
36    pub driver_version: Option<String>,
37    #[serde(default)]
38    pub gpu_arches: BTreeSet<String>,
39}
40
41#[derive(Clone, Debug, Default, Eq, PartialEq, Serialize, Deserialize)]
42pub struct HostRocmProfile {
43    #[serde(default, skip_serializing_if = "Option::is_none")]
44    pub version: Option<String>,
45    #[serde(default)]
46    pub gpu_arches: BTreeSet<String>,
47}
48
49#[derive(Clone, Debug, Default, Eq, PartialEq, Serialize, Deserialize)]
50pub struct HostVulkanProfile {
51    #[serde(default, skip_serializing_if = "Option::is_none")]
52    pub api_version: Option<String>,
53}
54
55#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
56pub struct HostRuntimeProfile {
57    pub os: String,
58    pub arch: String,
59    pub target_triple: Option<String>,
60    pub available_flavors: BTreeSet<NativeRuntimeBackendKind>,
61    pub gpus: Vec<HostGpuProfile>,
62    #[serde(default, skip_serializing_if = "Option::is_none")]
63    pub cuda: Option<HostCudaProfile>,
64    #[serde(default, skip_serializing_if = "Option::is_none")]
65    pub rocm: Option<HostRocmProfile>,
66    #[serde(default, skip_serializing_if = "Option::is_none")]
67    pub vulkan: Option<HostVulkanProfile>,
68}
69
70impl HostRuntimeProfile {
71    pub fn current_without_gpu_probe() -> Self {
72        let mut available_flavors = BTreeSet::from([NativeRuntimeBackendKind::Cpu]);
73        if cfg!(target_os = "macos") {
74            available_flavors.insert(NativeRuntimeBackendKind::Metal);
75        }
76        Self {
77            os: std::env::consts::OS.to_string(),
78            arch: std::env::consts::ARCH.to_string(),
79            target_triple: option_env!("TARGET").map(str::to_string),
80            available_flavors,
81            gpus: Vec::new(),
82            cuda: None,
83            rocm: None,
84            vulkan: None,
85        }
86    }
87
88    pub fn supports_flavor(&self, flavor: &NativeRuntimeBackendKind) -> bool {
89        self.available_flavors.contains(flavor)
90    }
91
92    pub fn has_gpu_name_matching(&self, needle: &str) -> bool {
93        let needle = needle.trim().to_ascii_lowercase();
94        !needle.is_empty()
95            && self
96                .gpus
97                .iter()
98                .any(|gpu| gpu.display_name.to_ascii_lowercase().contains(&needle))
99    }
100}