tract_data/
knobs.rs

1//! Runtime configuration knobs.
2//!
3//! A knob is a named, typed runtime setting — a kernel-selection toggle or a
4//! heuristic threshold. Its value is resolved, in priority order, from an
5//! explicit programmatic override, then the calling thread's [`KnobScope`],
6//! then the process environment, then the global values a deployment config
7//! installed ([`set_config_global`]), then a compiled-in default.
8//!
9//! A knob declared `scopable` is read while a model is transformed or
10//! prepared, so a [`KnobScope`] entered around those calls tunes it for that
11//! one model. Every other knob is read elsewhere (machine detection, eval) and
12//! only takes a process-wide value.
13//!
14//! Knobs are declared with [`declare_knob!`] and register themselves into a
15//! stack-wide inventory (this crate is the lowest in the stack, so `linalg` and
16//! everything above can declare). The inventory lets every knob be listed and
17//! documented from one place, and set through the API on targets where
18//! environment variables are unavailable (e.g. wasm).
19//!
20//! A knob is the lowest, deployer-facing layer of configuration; it is not a
21//! substitute for settings that must travel with a model.
22
23use std::cell::RefCell;
24use std::collections::HashMap;
25use std::str::FromStr;
26use std::sync::Arc;
27
28use anyhow::bail;
29use parking_lot::RwLock;
30
31use crate::TractResult;
32
33pub use inventory;
34
35/// A type usable as a knob value: parseable from the string sources (env / API)
36/// and renderable for listing.
37pub trait KnobValue: Sized + Clone + Send + Sync + 'static {
38    const TYPE_NAME: &'static str;
39    /// Parse from a source string. `None` means "unset/invalid": the resolver
40    /// falls through to the next source rather than overriding with garbage.
41    fn parse_knob(s: &str) -> Option<Self>;
42    fn render_knob(&self) -> String;
43}
44
45impl KnobValue for bool {
46    const TYPE_NAME: &'static str = "bool";
47    fn parse_knob(s: &str) -> Option<Self> {
48        match s.trim().to_ascii_lowercase().as_str() {
49            "1" | "true" | "on" | "yes" => Some(true),
50            "0" | "false" | "off" | "no" => Some(false),
51            _ => None,
52        }
53    }
54    fn render_knob(&self) -> String {
55        if *self { "true" } else { "false" }.to_string()
56    }
57}
58
59macro_rules! impl_knob_value_fromstr {
60    ($ty:ty, $name:expr) => {
61        impl KnobValue for $ty {
62            const TYPE_NAME: &'static str = $name;
63            fn parse_knob(s: &str) -> Option<Self> {
64                <$ty>::from_str(s.trim()).ok()
65            }
66            fn render_knob(&self) -> String {
67                self.to_string()
68            }
69        }
70    };
71}
72impl_knob_value_fromstr!(i64, "i64");
73impl_knob_value_fromstr!(usize, "usize");
74impl_knob_value_fromstr!(f32, "f32");
75
76impl KnobValue for String {
77    const TYPE_NAME: &'static str = "string";
78    fn parse_knob(s: &str) -> Option<Self> {
79        Some(s.to_string())
80    }
81    fn render_knob(&self) -> String {
82        self.clone()
83    }
84}
85
86/// Optional knob: `None` is the unset/default state (e.g. "autodetect"), `Some`
87/// is an explicit override. A source string that does not parse as `T` is
88/// treated as unset, so the resolver falls through to the next source.
89impl<T: KnobValue> KnobValue for Option<T> {
90    const TYPE_NAME: &'static str = "option";
91    fn parse_knob(s: &str) -> Option<Self> {
92        T::parse_knob(s).map(Some)
93    }
94    fn render_knob(&self) -> String {
95        match self {
96            Some(v) => v.render_knob(),
97            None => "(unset)".to_string(),
98        }
99    }
100}
101
102static OVERRIDES: RwLock<Option<HashMap<&'static str, String>>> = RwLock::new(None);
103
104static CONFIG_GLOBALS: RwLock<Option<HashMap<&'static str, String>>> = RwLock::new(None);
105
106thread_local! {
107    static SCOPE: RefCell<Option<KnobScope>> = const { RefCell::new(None) };
108}
109
110fn with_override<R>(name: &str, f: impl FnOnce(Option<&String>) -> R) -> R {
111    let guard = OVERRIDES.read();
112    f(guard.as_ref().and_then(|m| m.get(name)))
113}
114
115fn scoped_value<T: KnobValue>(name: &str) -> Option<T> {
116    SCOPE.with_borrow(|s| s.as_ref().and_then(|s| s.0.get(name)).and_then(|v| T::parse_knob(v)))
117}
118
119fn config_global_value<T: KnobValue>(name: &str) -> Option<T> {
120    CONFIG_GLOBALS.read().as_ref().and_then(|m| m.get(name)).and_then(|v| T::parse_knob(v))
121}
122
123fn find(name: &str) -> TractResult<&'static KnobInfo> {
124    match all().into_iter().find(|k| k.name == name) {
125        Some(info) => Ok(info),
126        None => bail!("No knob named {name}"),
127    }
128}
129
130fn checked(name: &str, value: &str) -> TractResult<&'static KnobInfo> {
131    let info = find(name)?;
132    if !(info.parses)(value) {
133        bail!("{value:?} is not a valid {} for knob {name}", info.type_name);
134    }
135    Ok(info)
136}
137
138/// Knob values tuned for one model: entered around the calls that transform
139/// and prepare it, they reach every [`scopable`](KnobInfo::scopable) knob read
140/// on this thread meanwhile.
141#[derive(Clone, Debug, Default)]
142pub struct KnobScope(Arc<HashMap<&'static str, String>>);
143
144impl KnobScope {
145    /// Fails on a name no knob carries, a knob that is not scopable, or a value
146    /// that does not parse as the knob's type.
147    pub fn new<'a>(values: impl IntoIterator<Item = (&'a str, &'a str)>) -> TractResult<Self> {
148        let mut map = HashMap::new();
149        for (name, value) in values {
150            let info = checked(name, value)?;
151            if !info.scopable {
152                bail!(
153                    "Knob {name} is not read while a model is prepared: it only takes a global value"
154                );
155            }
156            map.insert(info.name, value.to_string());
157        }
158        Ok(KnobScope(Arc::new(map)))
159    }
160
161    /// Run `f` with this scope active on the calling thread. The previously
162    /// active scope, if any, is restored afterwards.
163    pub fn enter<R>(&self, f: impl FnOnce() -> R) -> R {
164        let previous = SCOPE.replace(Some(self.clone()));
165        struct Restore(Option<KnobScope>);
166        impl Drop for Restore {
167            fn drop(&mut self) {
168                SCOPE.set(self.0.take());
169            }
170        }
171        let _restore = Restore(previous);
172        f()
173    }
174}
175
176/// Install a process-wide value below the environment, for a deployment
177/// config. Setting a knob again to a different value fails: two configs in one
178/// process must agree on what they both set globally.
179pub fn set_config_global(name: &str, value: &str) -> TractResult<()> {
180    let info = checked(name, value)?;
181    let mut globals = CONFIG_GLOBALS.write();
182    let globals = globals.get_or_insert_with(HashMap::new);
183    if let Some(previous) = globals.get(info.name)
184        && previous != value
185    {
186        bail!("Knob {name} already set to {previous:?} by another config, refusing {value:?}");
187    }
188    globals.insert(info.name, value.to_string());
189    Ok(())
190}
191
192/// A declared knob. Construct via [`declare_knob!`] rather than directly so it
193/// is registered for listing.
194pub struct Knob<T: KnobValue> {
195    pub name: &'static str,
196    pub doc: &'static str,
197    default: fn() -> T,
198}
199
200impl<T: KnobValue> Knob<T> {
201    pub const fn new(name: &'static str, default: fn() -> T, doc: &'static str) -> Self {
202        Knob { name, doc, default }
203    }
204
205    /// The compiled-in default, ignoring override and environment.
206    pub fn default_value(&self) -> T {
207        (self.default)()
208    }
209
210    /// Resolve the current value: programmatic override, then the thread's
211    /// [`KnobScope`], then environment, then config globals, then default.
212    pub fn get(&self) -> T {
213        if let Some(v) = with_override(self.name, |o| o.and_then(|s| T::parse_knob(s))) {
214            return v;
215        }
216        if let Some(v) = scoped_value(self.name) {
217            return v;
218        }
219        if let Ok(s) = std::env::var(self.name)
220            && let Some(v) = T::parse_knob(&s)
221        {
222            return v;
223        }
224        if let Some(v) = config_global_value(self.name) {
225            return v;
226        }
227        self.default_value()
228    }
229
230    /// Set a programmatic override (highest priority), effective for subsequent
231    /// [`get`](Self::get) calls. Prefer threading config through build options
232    /// where possible; this is global mutable state.
233    pub fn set(&self, value: T) {
234        OVERRIDES.write().get_or_insert_with(HashMap::new).insert(self.name, value.render_knob());
235    }
236
237    /// Drop the programmatic override, reverting to environment/default.
238    pub fn clear(&self) {
239        if let Some(m) = OVERRIDES.write().as_mut() {
240            m.remove(self.name);
241        }
242    }
243}
244
245/// Type-erased description of a declared knob, collected stack-wide for listing
246/// and documentation.
247pub struct KnobInfo {
248    pub name: &'static str,
249    pub doc: &'static str,
250    pub type_name: &'static str,
251    /// Read while a model is transformed or prepared, so a [`KnobScope`] may
252    /// set it for one model.
253    pub scopable: bool,
254    /// Whether a source string is a valid value.
255    pub parses: fn(&str) -> bool,
256    /// The default value, rendered.
257    pub default: fn() -> String,
258    /// The currently-resolved value, rendered.
259    pub current: fn() -> String,
260}
261
262inventory::collect!(KnobInfo);
263
264/// All declared knobs, sorted by name.
265pub fn all() -> Vec<&'static KnobInfo> {
266    let mut v: Vec<&'static KnobInfo> = inventory::iter::<KnobInfo>.into_iter().collect();
267    v.sort_by_key(|k| k.name);
268    v
269}
270
271/// Set a knob by name from a raw string — the environment-style source, for
272/// hosts where environment variables are unavailable (e.g. wasm). Returns
273/// `false` if no knob by that name is declared.
274pub fn set_str(name: &str, value: &str) -> bool {
275    match all().into_iter().find(|k| k.name == name) {
276        Some(info) => {
277            OVERRIDES.write().get_or_insert_with(HashMap::new).insert(info.name, value.to_string());
278            true
279        }
280        None => false,
281    }
282}
283
284/// Drop a knob's programmatic override by name. Returns `false` if not declared.
285pub fn clear_str(name: &str) -> bool {
286    match all().into_iter().find(|k| k.name == name) {
287        Some(info) => {
288            if let Some(m) = OVERRIDES.write().as_mut() {
289                m.remove(info.name);
290            }
291            true
292        }
293        None => false,
294    }
295}
296
297/// Render all declared knobs as a human-readable list.
298pub fn list() -> String {
299    use std::fmt::Write;
300    let shown = |s: String| if s.is_empty() { "(unset)".to_string() } else { s };
301    let mut out = String::new();
302    for k in all() {
303        let current = shown((k.current)());
304        let default = shown((k.default)());
305        write!(out, "{} = {}  [{}", k.name, current, k.type_name).unwrap();
306        if k.scopable {
307            write!(out, ", scopable").unwrap();
308        }
309        if current != default {
310            write!(out, ", default {default}").unwrap();
311        }
312        writeln!(out, "]\n    {}\n", k.doc).unwrap();
313    }
314    out
315}
316
317/// Declare a knob: a typed runtime setting resolved from override / environment
318/// / default, registered for listing.
319///
320/// The static is named exactly like the environment variable (the `TRACT_`
321/// prefix included), so one grep finds both the declaration and any use, and
322/// the env-var name is derived from the identifier — there is no separate
323/// string to drift.
324///
325/// Prefix the name with `scopable` for a knob read only while a model is
326/// transformed or prepared, never at eval nor cached for the process: such a
327/// knob honours a [`KnobScope`].
328///
329/// ```ignore
330/// declare_knob!(TRACT_MY_FLAG, bool, false, "Enable the thing.");
331/// declare_knob!(scopable TRACT_MY_RULE_OFF, bool, false, "Skip my rewrite rule.");
332/// // ... later ...
333/// if TRACT_MY_FLAG.get() { /* ... */ }
334/// ```
335#[macro_export]
336macro_rules! declare_knob {
337    (scopable $ident:ident, $ty:ty, $default:expr, $doc:literal) => {
338        $crate::declare_knob!(@ $ident, $ty, $default, $doc, true);
339    };
340    ($ident:ident, $ty:ty, $default:expr, $doc:literal) => {
341        $crate::declare_knob!(@ $ident, $ty, $default, $doc, false);
342    };
343    (@ $ident:ident, $ty:ty, $default:expr, $doc:literal, $scopable:literal) => {
344        #[doc = $doc]
345        pub static $ident: $crate::knobs::Knob<$ty> =
346            $crate::knobs::Knob::new(stringify!($ident), || $default, $doc);
347
348        $crate::knobs::inventory::submit! {
349            $crate::knobs::KnobInfo {
350                name: stringify!($ident),
351                doc: $doc,
352                type_name: <$ty as $crate::knobs::KnobValue>::TYPE_NAME,
353                scopable: $scopable,
354                parses: |s| <$ty as $crate::knobs::KnobValue>::parse_knob(s).is_some(),
355                default: || $crate::knobs::KnobValue::render_knob(&$ident.default_value()),
356                current: || $crate::knobs::KnobValue::render_knob(&$ident.get()),
357            }
358        }
359    };
360}
361
362#[cfg(test)]
363mod tests {
364    use super::*;
365
366    declare_knob!(TRACT_TEST_KNOB_BOOL, bool, false, "Test bool knob.");
367    declare_knob!(TRACT_TEST_KNOB_INT, usize, 7, "Test int knob.");
368    declare_knob!(TRACT_TEST_KNOB_SET_INT, usize, 7, "Test int knob for set-by-name.");
369    declare_knob!(scopable TRACT_TEST_KNOB_SCOPED, usize, 7, "Test scopable knob.");
370    declare_knob!(TRACT_TEST_KNOB_GLOBAL, usize, 7, "Test global knob.");
371
372    #[test]
373    fn scope_applies_inside_enter_only() {
374        let scope = KnobScope::new([("TRACT_TEST_KNOB_SCOPED", "3")]).unwrap();
375        assert_eq!(scope.enter(|| TRACT_TEST_KNOB_SCOPED.get()), 3);
376        assert_eq!(TRACT_TEST_KNOB_SCOPED.get(), 7);
377        let other = std::thread::spawn(move || scope.enter(|| TRACT_TEST_KNOB_SCOPED.get()));
378        assert_eq!(other.join().unwrap(), 3);
379    }
380
381    #[test]
382    fn scope_rejects_bad_entries() {
383        assert!(KnobScope::new([("TRACT_TEST_KNOB_GLOBAL", "3")]).is_err());
384        assert!(KnobScope::new([("TRACT_TEST_KNOB_SCOPED", "three")]).is_err());
385        assert!(KnobScope::new([("TRACT_TEST_KNOB_DOES_NOT_EXIST", "3")]).is_err());
386    }
387
388    #[test]
389    fn config_global_below_override_and_scope() {
390        set_config_global("TRACT_TEST_KNOB_GLOBAL", "5").unwrap();
391        assert_eq!(TRACT_TEST_KNOB_GLOBAL.get(), 5);
392        set_config_global("TRACT_TEST_KNOB_GLOBAL", "5").unwrap();
393        assert!(set_config_global("TRACT_TEST_KNOB_GLOBAL", "6").is_err());
394        TRACT_TEST_KNOB_GLOBAL.set(9);
395        assert_eq!(TRACT_TEST_KNOB_GLOBAL.get(), 9);
396        TRACT_TEST_KNOB_GLOBAL.clear();
397    }
398
399    #[test]
400    fn default_then_override_then_clear() {
401        assert!(!TRACT_TEST_KNOB_BOOL.get());
402        assert_eq!(TRACT_TEST_KNOB_INT.get(), 7);
403
404        TRACT_TEST_KNOB_BOOL.set(true);
405        assert!(TRACT_TEST_KNOB_BOOL.get());
406        TRACT_TEST_KNOB_BOOL.clear();
407        assert!(!TRACT_TEST_KNOB_BOOL.get());
408    }
409
410    #[test]
411    fn set_by_name() {
412        assert!(set_str("TRACT_TEST_KNOB_SET_INT", "42"));
413        assert_eq!(TRACT_TEST_KNOB_SET_INT.get(), 42);
414        assert!(clear_str("TRACT_TEST_KNOB_SET_INT"));
415        assert_eq!(TRACT_TEST_KNOB_SET_INT.get(), 7);
416        assert!(!set_str("TRACT_TEST_KNOB_DOES_NOT_EXIST", "1"));
417    }
418
419    #[test]
420    fn registered_in_inventory() {
421        let names: Vec<_> = all().iter().map(|k| k.name).collect();
422        assert!(names.contains(&"TRACT_TEST_KNOB_BOOL"));
423        assert!(names.contains(&"TRACT_TEST_KNOB_INT"));
424    }
425
426    #[test]
427    fn bool_parsing() {
428        assert_eq!(bool::parse_knob("1"), Some(true));
429        assert_eq!(bool::parse_knob("Off"), Some(false));
430        assert_eq!(bool::parse_knob("maybe"), None);
431    }
432}
tract_data/knobs.rs

tract_data/
knobs.rs
