rlx_runtime/
device_router.rs

1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3// SPDX-License-Identifier: MIT OR Apache-2.0
4
5//! Serving-oriented wrapper: warm-all backends, fallback execution, throttle-aware re-warm.
6
7use rlx_driver::Device;
8use rlx_ir::{DType, Graph};
9
10use crate::device_parse::device_label;
11use crate::device_policy::{DeviceFallbackError, DevicePolicy};
12use crate::graph_devices::GraphDevices;
13use crate::hwinfo::HwSnapshot;
14
15/// Production helper: compile all viable backends up front, run with fallback chain.
16pub struct DeviceRouter {
17    runner: GraphDevices,
18    rebench_on_throttle: bool,
19}
20
21impl DeviceRouter {
22    pub fn new(graph: Graph, policy: DevicePolicy) -> Result<Self, String> {
23        let mut runner = GraphDevices::with_policy(graph, policy);
24        runner.warm_all()?;
25        Ok(Self {
26            runner,
27            rebench_on_throttle: true,
28        })
29    }
30
31    pub fn from_env(graph: Graph) -> Result<Self, String> {
32        Self::new(graph, DevicePolicy::from_env())
33    }
34
35    pub fn with_rebench_on_throttle(mut self, enabled: bool) -> Self {
36        self.rebench_on_throttle = enabled;
37        self
38    }
39
40    pub fn set_rebench_on_throttle(&mut self, enabled: bool) {
41        self.rebench_on_throttle = enabled;
42    }
43
44    pub fn runner(&self) -> &GraphDevices {
45        &self.runner
46    }
47
48    pub fn runner_mut(&mut self) -> &mut GraphDevices {
49        &mut self.runner
50    }
51
52    fn maybe_rewarm(&mut self) -> Result<(), String> {
53        if self.rebench_on_throttle && HwSnapshot::collect().is_throttled() {
54            self.runner.invalidate_cache();
55            self.runner.warm_all()?;
56        }
57        Ok(())
58    }
59
60    /// Run on explicit device.
61    pub fn run_on(
62        &mut self,
63        device: Device,
64        inputs: &[(&str, &[f32])],
65    ) -> Result<Vec<Vec<f32>>, String> {
66        self.maybe_rewarm()?;
67        self.runner.run(device, inputs)
68    }
69
70    /// Resolve hint / env / benchmark policy, then execute.
71    pub fn run(
72        &mut self,
73        inputs: &[(&str, &[f32])],
74        hint: Option<Device>,
75    ) -> Result<(Device, Vec<Vec<f32>>), String> {
76        self.maybe_rewarm()?;
77        let device = self.runner.resolve_with_inputs(hint, inputs)?;
78        let out = self.runner.run(device, inputs)?;
79        Ok((device, out))
80    }
81
82    /// Fallback chain (`RLX_DEVICE_CHAIN` or explicit).
83    pub fn run_chain(
84        &mut self,
85        inputs: &[(&str, &[f32])],
86        hint: Option<Device>,
87    ) -> Result<(Device, Vec<Vec<f32>>), DeviceFallbackError> {
88        self.maybe_rewarm()?;
89        self.runner.run_chain(hint, inputs)
90    }
91
92    pub fn set_param(&mut self, name: &str, data: &[f32]) {
93        self.runner.set_param(name, data);
94    }
95
96    pub fn set_param_typed(&mut self, name: &str, data: &[u8], dtype: DType) {
97        self.runner.set_param_typed(name, data, dtype);
98    }
99
100    pub fn devices(&self) -> Vec<&'static str> {
101        self.runner
102            .devices()
103            .iter()
104            .map(|d| device_label(*d))
105            .collect()
106    }
107}
108
109#[cfg(test)]
110mod tests {
111    use super::*;
112    use rlx_ir::{DType, Shape};
113
114    #[test]
115    fn router_runs_identity_on_cpu() {
116        let mut g = Graph::new("id");
117        let x = g.input("x", Shape::new(&[4], DType::F32));
118        g.set_outputs(vec![x]);
119        let mut router = DeviceRouter::new(g, DevicePolicy::only([Device::Cpu])).expect("router");
120        let (dev, out) = router
121            .run(&[("x", &[1.0, 2.0, 3.0, 4.0])], None)
122            .expect("run");
123        assert_eq!(dev, Device::Cpu);
124        assert_eq!(out[0], vec![1.0, 2.0, 3.0, 4.0]);
125    }
126}
rlx_runtime/device_router.rs

rlx_runtime/
device_router.rs
