rlx_ir/
env.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//! Unified `RLX_*` configuration — readable from **code overrides** or process env.
6//!
7//! Code overrides (via [`set`], [`RlxEnv::apply`], or [`RuntimeOverrides::install`])
8//! take precedence over `std::env` for the same key.
9//!
10//! ```rust
11//! use rlx_ir::env::{self, RlxEnv};
12//!
13//! // Single knob
14//! env::set("RLX_VERBOSE", "1");
15//! assert!(env::flag("RLX_VERBOSE"));
16//!
17//! // Bulk
18//! RlxEnv::new()
19//!     .set("RLX_DISABLE_MPSGRAPH", "1")
20//!     .set("RLX_MPSGRAPH_MIN_FLOPS", "100000")
21//!     .apply();
22//! ```
23
24use std::collections::HashMap;
25use std::ffi::OsString;
26use std::str::FromStr;
27use std::sync::{OnceLock, RwLock};
28
29static OVERRIDES: OnceLock<RwLock<HashMap<String, String>>> = OnceLock::new();
30
31fn map() -> &'static RwLock<HashMap<String, String>> {
32    OVERRIDES.get_or_init(|| RwLock::new(HashMap::new()))
33}
34
35/// Normalize to `RLX_*` form.
36pub fn normalize_key(key: &str) -> String {
37    if key.starts_with("RLX_") {
38        key.to_string()
39    } else {
40        format!("RLX_{key}")
41    }
42}
43
44/// Set a code-side override. Pass `"0"` / `"false"` to force a flag off even
45/// when the process environment has it enabled.
46pub fn set(key: impl AsRef<str>, value: impl Into<String>) {
47    let key = normalize_key(key.as_ref());
48    if let Ok(mut g) = map().write() {
49        g.insert(key, value.into());
50    }
51}
52
53/// Remove a code override for `key`; subsequent reads fall back to process env.
54pub fn unset(key: impl AsRef<str>) {
55    let key = normalize_key(key.as_ref());
56    if let Ok(mut g) = map().write() {
57        g.remove(&key);
58    }
59}
60
61/// Drop every code override.
62pub fn clear_overrides() {
63    if let Ok(mut g) = map().write() {
64        g.clear();
65    }
66}
67
68/// Read configuration: code override first, then `std::env::var`.
69pub fn var(key: &str) -> Option<String> {
70    let key = normalize_key(key);
71    if let Ok(g) = map().read() {
72        if let Some(v) = g.get(&key) {
73            return Some(v.clone());
74        }
75    }
76    std::env::var(&key).ok()
77}
78
79/// Like [`var`] but returns an `OsString` (mirrors `std::env::var_os`).
80///
81/// **Not** a drop-in for `std::env::var_os` on non-UTF-8 values: this routes
82/// through [`var`], which uses `std::env::var(..).ok()` and therefore drops a
83/// value that is not valid UTF-8, where `std::env::var_os` would return it.
84/// That is deliberate — an override is stored as a `String`, so there is no
85/// representation for non-UTF-8 here — but it means a path-valued `RLX_*`
86/// variable containing invalid UTF-8 (possible on Linux, not on macOS) is
87/// treated as unset rather than used.
88pub fn var_os(key: &str) -> Option<OsString> {
89    var(key).map(Into::into)
90}
91
92/// True when the variable is set to a truthy value (`1`, `true`, `yes`, `on`, …).
93/// False when unset or set to `0` / `false` / `off` / `no` / empty.
94pub fn flag(key: &str) -> bool {
95    match var(key) {
96        Some(v) => truthy(&v),
97        None => false,
98    }
99}
100
101/// Like [`flag`], but returns `default` when the key is unset. Use for
102/// failsafe defaults that must stay on unless explicitly disabled (`=0`).
103pub fn flag_or(key: &str, default: bool) -> bool {
104    match var(key) {
105        Some(v) => truthy(&v),
106        None => default,
107    }
108}
109
110/// True when neither a code override nor process env provides the key.
111pub fn is_unset(key: &str) -> bool {
112    var(key).is_none()
113}
114
115/// Should a device-gated test skip, given whether the backend is built and
116/// whether it could be instantiated? Panics instead under `RLX_REQUIRE_DEVICE`.
117///
118/// A skipped test still prints `ok`, which is indistinguishable from a real
119/// pass in the summary line. On a rig whose card has fallen off the PCIe bus,
120/// a whole suite goes green having executed nothing — this repo has watched
121/// that happen. `RLX_REQUIRE_DEVICE=1` (set by `rig.sh` and the `Justfile` GPU
122/// recipes) turns the skip into a failure.
123///
124/// The `compiled` argument is what keeps that honest in both directions: a
125/// backend the build does not contain is *absent*, not *broken*, and failing
126/// for it would make every host fail for every backend it wasn't built with.
127/// Only a backend that IS built and still could not be instantiated is what the
128/// flag is about.
129///
130/// This lives in `rlx-ir` rather than in `rlx-runtime`'s test helpers because
131/// the backend crates need it too and cannot depend on `rlx-runtime` — it
132/// depends on them. Taking the two facts as plain booleans keeps it free of any
133/// backend or `Device` type:
134///
135/// ```no_run
136/// # fn is_available() -> bool { true }
137/// if rlx_ir::env::skip_unless_device("metal", cfg!(feature = "metal"), is_available()) {
138///     return;
139/// }
140/// ```
141#[must_use = "the caller must actually return when this is true, or the test runs without a device"]
142pub fn skip_unless_device(label: &str, compiled: bool, available: bool) -> bool {
143    if available {
144        return false;
145    }
146    assert!(
147        !flag("RLX_REQUIRE_DEVICE") || !compiled,
148        "RLX_REQUIRE_DEVICE=1, the `{label}` backend is compiled in, and it still \
149         could not be instantiated — this run would have reported `ok` without \
150         executing anything on it"
151    );
152    if compiled {
153        eprintln!("{label} unavailable on this host — skipping");
154    }
155    true
156}
157
158/// Parse an integer/bool/string knob, falling back to `default`.
159pub fn parse_or<T: FromStr>(key: &str, default: T) -> T {
160    var(key).and_then(|s| s.parse().ok()).unwrap_or(default)
161}
162
163fn truthy(v: &str) -> bool {
164    let s = v.trim();
165    if s.is_empty() {
166        return false;
167    }
168    match s.to_ascii_lowercase().as_str() {
169        "0" | "false" | "off" | "no" => false,
170        "1" | "true" | "yes" | "on" => true,
171        _ if s.chars().all(|c| c.is_ascii_digit()) => s != "0",
172        _ => true, // any other non-empty value counts as enabled
173    }
174}
175
176/// Bulk builder for code-side `RLX_*` overrides.
177#[derive(Debug, Clone, Default)]
178pub struct RlxEnv {
179    pairs: Vec<(String, String)>,
180}
181
182impl RlxEnv {
183    pub fn new() -> Self {
184        Self::default()
185    }
186
187    pub fn set(mut self, key: impl AsRef<str>, value: impl Into<String>) -> Self {
188        self.pairs.push((normalize_key(key.as_ref()), value.into()));
189        self
190    }
191
192    pub fn flag(mut self, key: impl AsRef<str>, on: bool) -> Self {
193        self.pairs.push((
194            normalize_key(key.as_ref()),
195            if on { "1" } else { "0" }.into(),
196        ));
197        self
198    }
199
200    /// Apply all pairs to the global override map.
201    pub fn apply(self) {
202        for (k, v) in self.pairs {
203            set(&k, v);
204        }
205    }
206}
207
208/// RAII guard: installs overrides on construction, restores previous values on drop.
209pub struct RuntimeOverrides {
210    saved: Vec<(String, Option<String>)>,
211}
212
213impl RuntimeOverrides {
214    /// Install `pairs` for the lifetime of the returned guard.
215    pub fn install(pairs: impl IntoIterator<Item = (impl AsRef<str>, impl Into<String>)>) -> Self {
216        let mut saved = Vec::new();
217        for (key, value) in pairs {
218            let key = normalize_key(key.as_ref());
219            let prev = map().read().ok().and_then(|g| g.get(&key).cloned());
220            saved.push((key.clone(), prev));
221            set(&key, value);
222        }
223        Self { saved }
224    }
225}
226
227impl Drop for RuntimeOverrides {
228    fn drop(&mut self) {
229        for (key, prev) in self.saved.drain(..) {
230            match prev {
231                Some(v) => set(&key, v),
232                None => unset(&key),
233            }
234        }
235    }
236}
237
238#[cfg(test)]
239mod tests {
240    use super::*;
241    use std::sync::Mutex;
242
243    /// Env tests share the process-global override map.
244    static ENV_TEST_LOCK: Mutex<()> = Mutex::new(());
245
246    fn with_clean_overrides(f: impl FnOnce()) {
247        let _guard = ENV_TEST_LOCK.lock().expect("env test lock poisoned");
248        clear_overrides();
249        f();
250        clear_overrides();
251    }
252
253    #[test]
254    fn code_override_wins_over_process_env() {
255        with_clean_overrides(|| {
256            let _g = RuntimeOverrides::install([("VERBOSE", "2")]);
257            assert_eq!(var("RLX_VERBOSE"), Some("2".into()));
258            assert!(flag("RLX_VERBOSE"));
259        });
260    }
261
262    #[test]
263    fn flag_parses_falsy_override() {
264        with_clean_overrides(|| {
265            set("RLX_DISABLE_MPSGRAPH", "0");
266            assert!(!flag("RLX_DISABLE_MPSGRAPH"));
267        });
268    }
269
270    #[test]
271    fn rlx_env_bulk_apply() {
272        with_clean_overrides(|| {
273            RlxEnv::new()
274                .set("MPSGRAPH_MIN_FLOPS", "42")
275                .flag("USE_ICB", true)
276                .apply();
277            assert_eq!(parse_or("RLX_MPSGRAPH_MIN_FLOPS", 0u64), 42);
278            assert!(flag("RLX_USE_ICB"));
279        });
280    }
281}
rlx_ir/env.rs

rlx_ir/
env.rs
