xet_runtime/utils/
configuration_utils.rs

1use std::str::FromStr;
2use std::sync::LazyLock;
3
4use tracing::{Level, event, info, warn};
5
6use super::ByteSize;
7#[cfg(not(target_family = "wasm"))]
8use super::TemplatedPathBuf;
9
10#[cfg(not(feature = "elevated_information_level"))]
11pub const INFORMATION_LOG_LEVEL: Level = Level::DEBUG;
12#[cfg(feature = "elevated_information_level")]
13pub const INFORMATION_LOG_LEVEL: Level = Level::INFO;
14
15/// A trait to control how a value is parsed from an environment string or other config source
16/// if it's present. Also provides serialization back to a string representation that
17/// roundtrips with `parse_user_value`.
18pub trait ParsableConfigValue: std::fmt::Debug + Sized {
19    fn parse_user_value(value: &str) -> Option<Self>;
20
21    /// Serialize this value to a string that can be parsed back via `parse_user_value`.
22    fn to_config_string(&self) -> String;
23
24    /// Try to update this value in place from a string. Returns true on success.
25    /// The default implementation delegates to `parse_user_value`, but types like
26    /// `ConfigEnum` override this to use context-aware validation.
27    fn try_update_in_place(&mut self, value: &str) -> bool {
28        if let Some(v) = Self::parse_user_value(value) {
29            *self = v;
30            true
31        } else {
32            false
33        }
34    }
35
36    /// Parse the value, returning the default if it can't be parsed or the string is empty.  
37    /// Issue a warning if it can't be parsed.
38    fn parse_config_value(variable_name: &str, value: Option<String>, default: Self) -> Self {
39        match value {
40            Some(v) => match Self::parse_user_value(&v) {
41                Some(v) => {
42                    info!("Config: {variable_name} = {v:?} (user set)");
43                    v
44                },
45                None => {
46                    warn!(
47                        "Configuration value {v} for {variable_name} cannot be parsed into correct type; reverting to default."
48                    );
49                    info!("Config: {variable_name} = {default:?} (default due to parse error)");
50                    default
51                },
52            },
53            None => {
54                event!(INFORMATION_LOG_LEVEL, "Config: {variable_name} = {default:?} (default)");
55                default
56            },
57        }
58    }
59}
60
61/// Most values work with the FromStr implementation, but we want to override the behavior for some types
62/// (e.g. Option<T> and bool) to have custom parsing behavior.
63pub trait FromStrParseable: FromStr + std::fmt::Debug + std::fmt::Display {}
64
65impl<T: FromStrParseable> ParsableConfigValue for T {
66    fn parse_user_value(value: &str) -> Option<Self> {
67        value.parse::<T>().ok()
68    }
69
70    fn to_config_string(&self) -> String {
71        self.to_string()
72    }
73}
74
75// Implement FromStrParseable for all the base types where the FromStr parsing method just works.
76impl FromStrParseable for usize {}
77impl FromStrParseable for u8 {}
78impl FromStrParseable for u16 {}
79impl FromStrParseable for u32 {}
80impl FromStrParseable for u64 {}
81impl FromStrParseable for isize {}
82impl FromStrParseable for i8 {}
83impl FromStrParseable for i16 {}
84impl FromStrParseable for i32 {}
85impl FromStrParseable for i64 {}
86impl FromStrParseable for f32 {}
87impl FromStrParseable for f64 {}
88impl FromStrParseable for String {}
89impl FromStrParseable for ByteSize {}
90
91/// Special handling for bool:
92/// - true: "1","true","yes","y","on"  -> true
93/// - false: "0","false","no","n","off" -> false
94pub(crate) fn parse_bool_value(value: &str) -> Option<bool> {
95    let t = value.trim().to_ascii_lowercase();
96
97    match t.as_str() {
98        "0" | "false" | "no" | "n" | "off" => Some(false),
99        "1" | "true" | "yes" | "y" | "on" => Some(true),
100        _ => None,
101    }
102}
103
104impl ParsableConfigValue for bool {
105    fn parse_user_value(value: &str) -> Option<Self> {
106        parse_bool_value(value)
107    }
108
109    fn to_config_string(&self) -> String {
110        if *self { "true" } else { "false" }.to_owned()
111    }
112}
113
114/// Enable Option<T> to allow the default value to be None if nothing is set and appear as
115/// Some(Value) if the user specifies the value.
116impl<T: ParsableConfigValue> ParsableConfigValue for Option<T> {
117    fn parse_user_value(value: &str) -> Option<Self> {
118        if value.trim().is_empty() {
119            return Some(None);
120        }
121        T::parse_user_value(value).map(Some)
122    }
123
124    fn to_config_string(&self) -> String {
125        match self {
126            Some(v) => v.to_config_string(),
127            None => String::new(),
128        }
129    }
130}
131
132/// Implement proper parsing for Duration types as well.
133///
134/// Now the following suffixes are supported: s, ms, us, ns, m, h, d, etc.;
135/// see the humantime crate for the full list.
136impl ParsableConfigValue for std::time::Duration {
137    fn parse_user_value(value: &str) -> Option<Self> {
138        humantime::parse_duration(value).ok()
139    }
140
141    fn to_config_string(&self) -> String {
142        let total_ms = self.as_millis();
143        if self.subsec_nanos() == 0 {
144            format!("{}s", self.as_secs())
145        } else if self.subsec_nanos().is_multiple_of(1_000_000) {
146            format!("{total_ms}ms")
147        } else if self.subsec_nanos().is_multiple_of(1_000) {
148            format!("{}us", self.as_micros())
149        } else {
150            format!("{}ns", self.as_nanos())
151        }
152    }
153}
154
155#[cfg(not(target_family = "wasm"))]
156impl ParsableConfigValue for TemplatedPathBuf {
157    fn parse_user_value(value: &str) -> Option<Self> {
158        Some(Self::new(value))
159    }
160
161    fn to_config_string(&self) -> String {
162        self.template_string()
163    }
164}
165
166#[macro_export]
167macro_rules! test_configurable_constants {
168    ($(
169        $(#[$meta:meta])*
170        ref $name:ident : $type:ty = $value:expr;
171    )+) => {
172        $(
173            $(#[$meta])*
174            pub static $name: ::std::sync::LazyLock<$type> = ::std::sync::LazyLock::new(|| {
175                #[cfg(debug_assertions)]
176                {
177                    let default_value = $value;
178                    let maybe_env_value = ::std::env::var(concat!("HF_XET_", stringify!($name))).ok();
179                    <$type as $crate::configuration_utils::ParsableConfigValue>::parse_config_value(
180                        stringify!($name),
181                        maybe_env_value,
182                        default_value,
183                    )
184                }
185                #[cfg(not(debug_assertions))]
186                {
187                    $value
188                }
189            });
190        )+
191    };
192}
193
194pub use ctor as ctor_reexport;
195
196#[cfg(not(doctest))]
197/// A macro for **tests** that sets `HF_XET_<CONSTANT_NAME>` to `$value` **before**
198/// the constant is initialized, and then checks that the constant actually picks up
199/// that value. If the constant was already accessed (thus initialized), or if it
200/// doesn't match after being set, this macro panics.
201///
202/// Typically you would document *the macro itself* here, rather than placing
203/// doc comments above each call to `test_set_constants!`, because it doesn't
204/// define a new item.
205///
206/// # Example
207/// ```rust
208/// use xet_runtime::{test_configurable_constants, test_set_constants};
209/// test_configurable_constants! {
210///    /// Target chunk size
211///    ref CHUNK_TARGET_SIZE: u64 = 1024;
212///
213///    /// Max Chunk size, only adjustable in testing mode.
214///    ref MAX_CHUNK_SIZE: u64 = 4096;
215/// }
216///
217/// test_set_constants! {
218///    CHUNK_TARGET_SIZE = 2048;
219/// }
220/// assert_eq!(*CHUNK_TARGET_SIZE, 2048);
221/// ```
222#[macro_export]
223macro_rules! test_set_constants {
224    ($(
225        $var_name:ident = $val:expr;
226    )+) => {
227        use $crate::configuration_utils::ctor_reexport as ctor;
228
229        #[ctor::ctor(unsafe)]
230        fn set_constants_on_load() {
231            $(
232                let val = $val;
233                let val_str = format!("{val:?}");
234
235                // Construct the environment variable name, e.g. "HF_XET_MAX_NUM_CHUNKS"
236                let env_name = concat!("HF_XET_", stringify!($var_name));
237
238                // Set the environment
239                unsafe {
240                    std::env::set_var(env_name, &val_str);
241                }
242
243                // Force the lazy initialization now:
244                let actual_value = *$var_name;
245
246                if format!("{actual_value:?}") != val_str {
247                    panic!(
248                        "test_set_constants! failed: wanted {} to be {:?}, but got {:?}",
249                        stringify!($var_name),
250                        val,
251                        actual_value
252                    );
253                }
254                eprintln!("> Set {} to {:?}",
255                        stringify!($var_name),
256                        val);
257            )+
258        }
259    }
260}
261
262#[cfg(not(doctest))]
263/// A macro for **tests** that sets config group environment variables **before**
264/// XetContext is initialized. The environment variables follow the pattern
265/// `HF_XET_{GROUP_NAME}_{FIELD_NAME}`.
266///
267/// This macro uses `ctor` to run on module load, ensuring environment variables
268/// are set before any config values are read.
269///
270/// # Example
271/// ```rust
272/// use xet_runtime::config::XetConfig;
273/// use xet_runtime::test_set_config;
274///
275/// test_set_config! {
276///     data {
277///         max_concurrent_uploads = 16;
278///         max_concurrent_downloads = 20;
279///     }
280///     client {
281///         upload_reporting_block_size = 1024000;
282///     }
283/// }
284///
285/// // Now XetConfig::new() will pick up these values
286/// let config = XetConfig::new();
287/// assert_eq!(config.data.max_concurrent_uploads, 16);
288/// ```
289#[macro_export]
290macro_rules! test_set_config {
291    ($(
292        $group_name:ident {
293            $(
294                $field_name:ident = $val:expr;
295            )+
296        }
297    )+) => {
298        use $crate::configuration_utils::ctor_reexport as config_ctor;
299
300        #[config_ctor::ctor(unsafe)]
301        fn set_config_on_load() {
302            $(
303                let group_name_upper = stringify!($group_name).to_uppercase();
304                $(
305                    let val = $val;
306                    let val_str = format!("{val:?}");
307                    let field_name_upper = stringify!($field_name).to_uppercase();
308
309                    // Construct the environment variable name: HF_XET_{GROUP_NAME}_{FIELD_NAME}
310                    let env_name = format!("HF_XET_{}_{}", group_name_upper, field_name_upper);
311
312                    // Set the environment
313                    unsafe {
314                        std::env::set_var(&env_name, &val_str);
315                    }
316
317                    eprintln!("> Set config {}.{} to {:?} (env: {})",
318                            stringify!($group_name),
319                            stringify!($field_name),
320                            val,
321                            env_name);
322                )+
323            )+
324        }
325    }
326}
327
328fn get_high_performance_flag() -> bool {
329    if let Ok(val) = std::env::var("HF_XET_HIGH_PERFORMANCE") {
330        parse_bool_value(&val).unwrap_or(false)
331    } else if let Ok(val) = std::env::var("HF_XET_HP") {
332        parse_bool_value(&val).unwrap_or(false)
333    } else {
334        false
335    }
336}
337
338/// To set the high performance mode to true, set either of the following environment variables to 1 or true:
339///  - HF_XET_HIGH_PERFORMANCE
340///  - HF_XET_HP
341pub static HIGH_PERFORMANCE: LazyLock<bool> = LazyLock::new(get_high_performance_flag);
342
343#[inline]
344pub fn is_high_performance() -> bool {
345    *HIGH_PERFORMANCE
346}
347
348/// `huggingface_hub` environment variables that suppress client telemetry regardless of
349/// `HF_XET_TELEMETRY_ENABLED`: the same variables `huggingface_hub` checks before sending its own
350/// telemetry (`constants.HF_HUB_DISABLE_TELEMETRY` and `constants.HF_HUB_OFFLINE`).
351///
352/// Temporary: lets an hf-xet release honor these opt-outs without waiting on a `huggingface_hub`
353/// release. Remove once `huggingface_hub` sets `telemetry.enabled` on the `XetConfig` it passes to
354/// `XetSession`.
355pub(crate) const TELEMETRY_OPT_OUT_VARS: &[&str] = &[
356    "HF_HUB_DISABLE_TELEMETRY",
357    "DISABLE_TELEMETRY",
358    "DO_NOT_TRACK",
359    "HF_HUB_OFFLINE",
360    "TRANSFORMERS_OFFLINE",
361];
362
363/// Whether any of [`TELEMETRY_OPT_OUT_VARS`] is set to a truthy value.
364///
365/// The value must parse as truthy, so `HF_HUB_OFFLINE=0` leaves telemetry on. Not memoized: it is
366/// read each time a `XetConfig` is built, so a process that changes the variables sees the change.
367pub(crate) fn telemetry_opted_out() -> bool {
368    TELEMETRY_OPT_OUT_VARS
369        .iter()
370        .any(|name| std::env::var(name).ok().and_then(|v| parse_bool_value(&v)).unwrap_or(false))
371}
372
373#[cfg(test)]
374mod tests {
375    use std::time::Duration;
376
377    use super::*;
378
379    fn assert_roundtrip<T: ParsableConfigValue + PartialEq + std::fmt::Debug>(value: T) {
380        let s = value.to_config_string();
381        let restored = T::parse_user_value(&s).unwrap_or_else(|| {
382            panic!("Failed to parse config string '{s}' back into {:?}", std::any::type_name::<T>())
383        });
384        assert_eq!(value, restored);
385    }
386
387    macro_rules! assert_roundtrips {
388        ($($value:expr),+ $(,)?) => {
389            $(assert_roundtrip($value);)+
390        };
391    }
392
393    #[test]
394    fn test_roundtrip_numeric_primitives() {
395        assert_roundtrips!(
396            0usize,
397            42usize,
398            usize::MAX,
399            0u8,
400            255u8,
401            0u16,
402            65535u16,
403            0u32,
404            123456u32,
405            0u64,
406            u64::MAX,
407            0isize,
408            -42isize,
409            isize::MAX,
410            -128i8,
411            127i8,
412            -32768i16,
413            32767i16,
414            0i32,
415            -123456i32,
416            0i64,
417            i64::MIN,
418            0.0f32,
419            3.14f32,
420            -1.5f32,
421            0.0f64,
422            std::f64::consts::PI,
423            -1e10f64
424        );
425    }
426
427    #[test]
428    fn test_roundtrip_bool_and_string() {
429        assert_roundtrips!(true, false, String::new(), "hello world".to_owned(), "http://localhost:8080".to_owned());
430    }
431
432    #[test]
433    fn test_roundtrip_byte_size() {
434        assert_roundtrips!(
435            ByteSize::new(0),
436            ByteSize::new(1000),
437            ByteSize::new(1_000_000),
438            ByteSize::new(8_000_000),
439            ByteSize::new(10_000_000_000)
440        );
441    }
442
443    #[test]
444    fn test_roundtrip_duration() {
445        assert_roundtrips!(
446            Duration::from_secs(0),
447            Duration::from_secs(60),
448            Duration::from_secs(120),
449            Duration::from_millis(200),
450            Duration::from_millis(3000),
451            Duration::from_secs(360),
452            Duration::from_micros(123),
453            Duration::from_nanos(123)
454        );
455    }
456
457    #[test]
458    fn test_roundtrip_option_some() {
459        assert_roundtrips!(Some(42usize), Some("hello".to_owned()));
460    }
461
462    #[test]
463    fn test_roundtrip_option_none_string() {
464        let none_val: Option<String> = None;
465        let s = none_val.to_config_string();
466        assert_eq!(s, "");
467    }
468
469    #[test]
470    fn test_parse_option_empty_string_as_none() {
471        assert_eq!(Option::<String>::parse_user_value(""), Some(None));
472        assert_eq!(Option::<String>::parse_user_value("   "), Some(None));
473        assert_eq!(Option::<usize>::parse_user_value(""), Some(None));
474    }
475
476    #[test]
477    fn test_parse_bool_empty_string_as_none() {
478        assert_eq!(bool::parse_user_value(""), None);
479        assert_eq!(bool::parse_user_value("   "), None);
480    }
481
482    #[cfg(not(target_family = "wasm"))]
483    #[test]
484    fn test_roundtrip_templated_path_buf() {
485        let path = TemplatedPathBuf::new("/some/simple/path");
486        let s = path.to_config_string();
487        assert_eq!(s, "/some/simple/path");
488        let restored = TemplatedPathBuf::parse_user_value(&s).unwrap();
489        assert_eq!(path.template_string(), restored.template_string());
490    }
491}
xet_runtime/utils/configuration_utils.rs

xet_runtime/utils/
configuration_utils.rs
