leash_harness/
accelerator.rs

1use anyhow::{bail, Result};
2use serde::{Deserialize, Serialize};
3
4use crate::config::AcceleratorBackend;
5
6pub trait AcceleratorProvider: Send + Sync {
7    fn backend(&self) -> AcceleratorBackend;
8    fn compiled(&self) -> bool;
9    fn available(&self) -> bool;
10    fn message(&self) -> &'static str;
11
12    fn probe(&self, selected: bool) -> AcceleratorProbe {
13        AcceleratorProbe {
14            backend: self.backend(),
15            compiled: self.compiled(),
16            available: self.available(),
17            selected,
18            message: self.message().to_string(),
19        }
20    }
21}
22
23#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
24#[cfg_attr(feature = "mcp", derive(schemars::JsonSchema))]
25pub struct AcceleratorProbe {
26    pub backend: AcceleratorBackend,
27    pub compiled: bool,
28    pub available: bool,
29    pub selected: bool,
30    pub message: String,
31}
32
33#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
34#[cfg_attr(feature = "mcp", derive(schemars::JsonSchema))]
35pub struct AcceleratorStatus {
36    pub requested: AcceleratorBackend,
37    pub active: AcceleratorBackend,
38    pub available: bool,
39    pub required: bool,
40    pub message: String,
41    pub probes: Vec<AcceleratorProbe>,
42}
43
44#[derive(Debug)]
45struct CpuAccelerator;
46
47impl AcceleratorProvider for CpuAccelerator {
48    fn backend(&self) -> AcceleratorBackend {
49        AcceleratorBackend::Cpu
50    }
51
52    fn compiled(&self) -> bool {
53        true
54    }
55
56    fn available(&self) -> bool {
57        true
58    }
59
60    fn message(&self) -> &'static str {
61        "CPU accelerator backend available"
62    }
63}
64
65#[derive(Debug)]
66struct CudaAccelerator;
67
68impl AcceleratorProvider for CudaAccelerator {
69    fn backend(&self) -> AcceleratorBackend {
70        AcceleratorBackend::Cuda
71    }
72
73    fn compiled(&self) -> bool {
74        cfg!(feature = "cuda")
75    }
76
77    fn available(&self) -> bool {
78        false
79    }
80
81    fn message(&self) -> &'static str {
82        if cfg!(feature = "cuda") {
83            "CUDA feature compiled; device probe not implemented"
84        } else {
85            "CUDA feature not compiled"
86        }
87    }
88}
89
90pub fn resolve_accelerator(
91    requested: AcceleratorBackend,
92    required: bool,
93) -> Result<AcceleratorStatus> {
94    let probes = probe_inventory(AcceleratorBackend::None);
95    match requested {
96        AcceleratorBackend::None => {
97            if required {
98                bail!("accelerator is required but no accelerator backend was selected");
99            }
100            Ok(AcceleratorStatus {
101                requested,
102                active: AcceleratorBackend::None,
103                available: true,
104                required,
105                message: "no accelerator requested".to_string(),
106                probes,
107            })
108        }
109        AcceleratorBackend::Cpu => {
110            let probes = probe_inventory(AcceleratorBackend::Cpu);
111            Ok(AcceleratorStatus {
112                requested,
113                active: AcceleratorBackend::Cpu,
114                available: true,
115                required,
116                message: "CPU accelerator backend active".to_string(),
117                probes,
118            })
119        }
120        AcceleratorBackend::Cuda => cuda_status(required),
121    }
122}
123
124pub fn probe_inventory(selected: AcceleratorBackend) -> Vec<AcceleratorProbe> {
125    let cpu = CpuAccelerator;
126    let cuda = CudaAccelerator;
127    vec![
128        cpu.probe(selected == AcceleratorBackend::Cpu),
129        cuda.probe(selected == AcceleratorBackend::Cuda),
130    ]
131}
132
133fn cuda_status(required: bool) -> Result<AcceleratorStatus> {
134    let cuda_probe = CudaAccelerator.probe(false);
135    if required {
136        bail!(
137            "CUDA accelerator requested as required but unavailable: {}",
138            cuda_probe.message
139        );
140    }
141    Ok(AcceleratorStatus {
142        requested: AcceleratorBackend::Cuda,
143        active: AcceleratorBackend::Cpu,
144        available: false,
145        required,
146        message: "CUDA accelerator unavailable; CPU fallback active".to_string(),
147        probes: probe_inventory(AcceleratorBackend::Cpu),
148    })
149}
150
151#[cfg(test)]
152mod tests {
153    use super::*;
154
155    #[test]
156    fn cpu_backend_is_available_without_hardware() {
157        let status = resolve_accelerator(AcceleratorBackend::Cpu, true).unwrap();
158        assert_eq!(status.active, AcceleratorBackend::Cpu);
159        assert!(status.available);
160        assert!(status.required);
161        assert_eq!(
162            probe_for(&status, AcceleratorBackend::Cpu),
163            Some(&AcceleratorProbe {
164                backend: AcceleratorBackend::Cpu,
165                compiled: true,
166                available: true,
167                selected: true,
168                message: "CPU accelerator backend available".to_string(),
169            })
170        );
171    }
172
173    #[test]
174    fn required_none_backend_is_rejected() {
175        let err = resolve_accelerator(AcceleratorBackend::None, true)
176            .unwrap_err()
177            .to_string();
178        assert!(err.contains("required"));
179    }
180
181    #[test]
182    fn cuda_selection_has_ci_safe_behavior() {
183        let status = resolve_accelerator(AcceleratorBackend::Cuda, false).unwrap();
184        assert_eq!(status.requested, AcceleratorBackend::Cuda);
185        assert_eq!(status.active, AcceleratorBackend::Cpu);
186        assert!(!status.available);
187        assert_eq!(
188            probe_for(&status, AcceleratorBackend::Cuda).map(|probe| probe.compiled),
189            Some(cfg!(feature = "cuda"))
190        );
191        assert_eq!(
192            probe_for(&status, AcceleratorBackend::Cuda).map(|probe| probe.available),
193            Some(false)
194        );
195        assert_eq!(
196            probe_for(&status, AcceleratorBackend::Cpu).map(|probe| probe.selected),
197            Some(true)
198        );
199    }
200
201    #[test]
202    fn required_cuda_requires_available_backend() {
203        let err = resolve_accelerator(AcceleratorBackend::Cuda, true)
204            .unwrap_err()
205            .to_string();
206        assert!(err.contains("unavailable"));
207    }
208
209    #[test]
210    fn default_inventory_reports_compiled_backends() {
211        let status = resolve_accelerator(AcceleratorBackend::None, false).unwrap();
212        assert_eq!(
213            probe_for(&status, AcceleratorBackend::Cpu).map(|probe| probe.compiled),
214            Some(true)
215        );
216        assert_eq!(
217            probe_for(&status, AcceleratorBackend::Cuda).map(|probe| probe.compiled),
218            Some(cfg!(feature = "cuda"))
219        );
220        assert!(status.probes.iter().all(|probe| !probe.selected));
221    }
222
223    fn probe_for(
224        status: &AcceleratorStatus,
225        backend: AcceleratorBackend,
226    ) -> Option<&AcceleratorProbe> {
227        status.probes.iter().find(|probe| probe.backend == backend)
228    }
229}
leash_harness/accelerator.rs

leash_harness/
accelerator.rs
