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(|r| {
114 r.id == forced && backend_dispatches(r.id) && !is_reference_oracle_backend(r.id)
115 });
116 return match hit {
117 Some(reg) => Ok(RouterDecision {
118 backend: reg.id,
119 reason: Reason::EnvOverride,
120 }),
121 None => Err(BackendError::new(format!(
122 "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."
123 ))),
124 };
125 }
126 }
127
128 for reg in registered_backends_by_precedence_slice() {
132 if registered.iter().any(|r| r.id == reg.id)
133 && backend_dispatches(reg.id)
134 && !is_reference_oracle_backend(reg.id)
135 {
136 return Ok(RouterDecision {
137 backend: reg.id,
138 reason: Reason::Precedence,
139 });
140 }
141 }
142
143 Err(BackendError::new(
144 "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.",
145 ))
146 }
147
148 #[must_use]
152 pub fn enumerate_by_precedence() -> Vec<&'static BackendRegistration> {
153 registered_backends_by_precedence_slice().to_vec()
154 }
155}
156
157fn is_reference_oracle_backend(id: &str) -> bool {
158 matches!(id, "cpu-ref" | "reference")
159}
160
161#[cfg(test)]
162mod tests {
163 use super::*;
164 use vyre_driver::backend::backend_precedence;
165
166 fn noop_program() -> Program {
167 Program::wrapped(Vec::new(), [1, 1, 1], Vec::new())
170 }
171
172 #[test]
173 fn enumerate_by_precedence_puts_wgpu_before_reference() {
174 let wgpu_rank = backend_precedence("wgpu");
178 let ref_rank = backend_precedence("cpu-ref");
179 assert!(
180 wgpu_rank < ref_rank || ref_rank == u32::MAX,
181 "wgpu (rank {wgpu_rank}) must take precedence over the CPU reference oracle (rank {ref_rank})"
182 );
183 }
184
185 #[test]
186 fn enumerate_by_precedence_is_inventory_driven() {
187 let ranked = BackendRouter::enumerate_by_precedence();
189 let wgpu = ranked.iter().find(|r| r.id == "wgpu").expect(
191 "Fix: wgpu backend registered in this crate; restore this invariant before continuing.",
192 );
193 assert_eq!(backend_precedence(wgpu.id), 30);
194 }
195
196 #[test]
197 fn explicit_override_with_unknown_backend_surfaces_error() {
198 let router = BackendRouter::new();
199 let err = router
200 .pick_with_override(
201 &noop_program(),
202 Override::Explicit("does-not-exist-backend"),
203 )
204 .expect_err("unknown backend must error");
205 let msg = format!("{err}");
206 assert!(msg.contains("does-not-exist-backend"));
207 assert!(msg.contains("Fix:"));
208 }
209
210 #[test]
211 fn explicit_override_picks_the_named_backend_when_registered() {
212 let router = BackendRouter::new();
213 let decision = router
215 .pick_with_override(&noop_program(), Override::Explicit("wgpu"))
216 .expect("Fix: wgpu backend is registered in this crate");
217 assert_eq!(decision.backend, "wgpu");
218 assert_eq!(decision.reason, Reason::EnvOverride);
219 }
220
221 #[test]
222 fn precedence_picks_wgpu_when_registered() {
223 let router = BackendRouter::new();
224 let decision = router
225 .pick_with_override(&noop_program(), Override::None)
226 .expect("Fix: at least one backend must register");
227 assert_eq!(decision.reason, Reason::Precedence);
228 assert!(
231 backend_precedence(decision.backend) < u32::MAX,
232 "picked backend {} did not submit a BackendPrecedence inventory entry",
233 decision.backend
234 );
235 }
236}