vyre_driver_wgpu/runtime/
router.rs1use std::env;
23
24use vyre_driver::backend::{backend_dispatches, registered_backends_by_precedence_slice};
25use vyre_driver::{BackendError, BackendRegistration};
26use vyre_foundation::ir::Program;
27
28#[derive(Debug, Clone, Copy)]
30#[non_exhaustive]
31pub enum Override<'a> {
32 FromEnv,
34 Explicit(&'a str),
36 None,
38}
39
40const OVERRIDE_ENV: &str = "VYRE_BACKEND";
41
42#[derive(Debug, Clone)]
44pub struct RouterDecision {
45 pub backend: &'static str,
47 pub reason: Reason,
49}
50
51#[derive(Debug, Clone, Copy, PartialEq, Eq)]
53#[non_exhaustive]
54pub enum Reason {
55 EnvOverride,
57 Precedence,
60}
61
62#[derive(Default)]
68pub struct BackendRouter;
69
70impl BackendRouter {
71 #[must_use]
73 pub fn new() -> Self {
74 Self
75 }
76
77 pub fn pick(&self, program: &Program) -> Result<RouterDecision, BackendError> {
89 self.pick_with_override(program, Override::FromEnv)
90 }
91
92 pub fn pick_with_override(
99 &self,
100 _program: &Program,
101 source: Override<'_>,
102 ) -> Result<RouterDecision, BackendError> {
103 let registered = vyre_driver::backend::registered_backends()?;
104
105 let forced: Option<String> = match source {
106 Override::FromEnv => env::var(OVERRIDE_ENV).ok(),
107 Override::Explicit(s) => Some(s.to_owned()),
108 Override::None => None,
109 };
110 if let Some(forced) = forced {
111 let forced = forced.trim();
112 if !forced.is_empty() {
113 let hit = registered.iter().find_map(|registration| {
114 (registration.id == forced && !registration.reference_oracle)
115 .then_some(registration)
116 });
117 let hit = match hit {
118 Some(registration) if backend_dispatches(registration.id)? => {
119 Some(registration)
120 }
121 _ => None,
122 };
123 return match hit {
124 Some(reg) => Ok(RouterDecision {
125 backend: reg.id,
126 reason: Reason::EnvOverride,
127 }),
128 None => Err(BackendError::new(format!(
129 "VYRE_BACKEND={forced} is not an executable registered GPU backend. Fix: link CUDA/WGPU or unset VYRE_BACKEND; cpu-ref/reference are explicit conformance oracles, not runtime router targets."
130 ))),
131 };
132 }
133 }
134
135 for registration in registered_backends_by_precedence_slice()? {
139 if registered
140 .iter()
141 .any(|candidate| candidate.id == registration.id)
142 && backend_dispatches(registration.id)?
143 && !registration.reference_oracle
144 {
145 return Ok(RouterDecision {
146 backend: registration.id,
147 reason: Reason::Precedence,
148 });
149 }
150 }
151
152 Err(BackendError::new(
153 "no executable GPU backend is registered. Fix: link vyre-driver-cuda or vyre-driver-wgpu into the binary and verify the GPU driver probe succeeds.",
154 ))
155 }
156
157 pub fn enumerate_by_precedence() -> Result<Vec<&'static BackendRegistration>, BackendError> {
164 Ok(registered_backends_by_precedence_slice()?.to_vec())
165 }
166}
167
168#[cfg(test)]
169mod tests {
170 use super::*;
171 use vyre_driver::backend::backend_precedence;
172
173 fn noop_program() -> Program {
174 Program::wrapped(Vec::new(), [1, 1, 1], Vec::new())
177 }
178
179 #[test]
180 fn enumerate_by_precedence_puts_wgpu_before_reference() {
181 let wgpu_rank = backend_precedence("wgpu").expect("valid backend registry");
185 let ref_rank = backend_precedence("cpu-ref").expect("valid backend registry");
186 assert!(
187 wgpu_rank < ref_rank || ref_rank == u32::MAX,
188 "wgpu (rank {wgpu_rank}) must take precedence over the CPU reference oracle (rank {ref_rank})"
189 );
190 }
191
192 #[test]
193 fn enumerate_by_precedence_is_inventory_driven() {
194 let ranked = BackendRouter::enumerate_by_precedence().expect("valid backend registry");
196 let wgpu = ranked.iter().find(|r| r.id == "wgpu").expect(
198 "Fix: wgpu backend registered in this crate; restore this invariant before continuing.",
199 );
200 assert_eq!(
201 backend_precedence(wgpu.id).expect("valid backend registry"),
202 30
203 );
204 }
205
206 #[test]
207 fn explicit_override_with_unknown_backend_surfaces_error() {
208 let router = BackendRouter::new();
209 let err = router
210 .pick_with_override(
211 &noop_program(),
212 Override::Explicit("does-not-exist-backend"),
213 )
214 .expect_err("unknown backend must error");
215 let msg = format!("{err}");
216 assert!(msg.contains("does-not-exist-backend"));
217 assert!(msg.contains("Fix:"));
218 }
219
220 #[test]
221 fn explicit_override_picks_the_named_backend_when_registered() {
222 let router = BackendRouter::new();
223 let decision = router
225 .pick_with_override(&noop_program(), Override::Explicit("wgpu"))
226 .expect("Fix: wgpu backend is registered in this crate");
227 assert_eq!(decision.backend, "wgpu");
228 assert_eq!(decision.reason, Reason::EnvOverride);
229 }
230
231 #[test]
232 fn precedence_picks_wgpu_when_registered() {
233 let router = BackendRouter::new();
234 let decision = router
235 .pick_with_override(&noop_program(), Override::None)
236 .expect("Fix: at least one backend must register");
237 assert_eq!(decision.reason, Reason::Precedence);
238 assert!(
241 backend_precedence(decision.backend).expect("valid backend registry") < u32::MAX,
242 "picked backend {} did not submit a BackendPrecedence inventory entry",
243 decision.backend
244 );
245 }
246}