modelexpress_common/
envs.rs

1// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Centralized registry of every environment variable the workspace reads.
5//!
6//! Inspired by vLLM's `envs.py`: this module is the single source of truth for
7//! env-var *names* (the `pub const` block) and provides typed *getters* that
8//! encapsulate defaults, fallback chains, and parsing. Both ModelExpress-owned
9//! variables (`MODEL_EXPRESS_*`, `MX_*`) and third-party variables the code
10//! depends on (`HF_*`, `NGC_*`, `REDIS_*`, `POD_NAMESPACE`, `HOME`) live here.
11//!
12//! # Conventions
13//! - Getters read `std::env` on **every** call and never cache. Some callers
14//!   (and tests via [`crate::test_support::EnvVarGuard`]) mutate the process
15//!   environment at runtime, so a cached value would go stale.
16//! - The name constants are referenced directly from clap `#[arg(env = ...)]`
17//!   attributes so the CLI and the getters can never drift apart.
18//!
19//! # Not covered here
20//! - `CARGO_PKG_VERSION` is read at compile time via the `env!` macro in
21//!   `modelexpress_server::services`; it is not a runtime variable.
22//! - `HF_ENDPOINT` is read directly by the `hf_hub` crate (via
23//!   `ApiBuilder::from_env`); ModelExpress only sets it in tests. Its name is
24//!   registered below for reference.
25
26use crate::Error;
27use crate::constants;
28use std::env;
29use std::path::PathBuf;
30
31// ── Config-loader prefix ────────────────────────────────────────────────────
32/// Prefix consumed by the `config` crate's `Environment` source in
33/// [`crate::config::load_layered_config`] (env vars like `MODEL_EXPRESS_*`
34/// override matching config-file fields).
35pub const MODEL_EXPRESS_PREFIX: &str = "MODEL_EXPRESS";
36
37// ── ModelExpress-owned variables ────────────────────────────────────────────
38/// Client server endpoint (`ClientArgs::endpoint`).
39pub const MODEL_EXPRESS_ENDPOINT: &str = "MODEL_EXPRESS_ENDPOINT";
40/// Client request timeout in seconds (`ClientArgs::timeout`).
41pub const MODEL_EXPRESS_TIMEOUT: &str = "MODEL_EXPRESS_TIMEOUT";
42/// Local model cache directory (client, server, and both providers).
43pub const MODEL_EXPRESS_CACHE_DIRECTORY: &str = "MODEL_EXPRESS_CACHE_DIRECTORY";
44/// Log level (client and server).
45pub const MODEL_EXPRESS_LOG_LEVEL: &str = "MODEL_EXPRESS_LOG_LEVEL";
46/// Log output format (client and server).
47pub const MODEL_EXPRESS_LOG_FORMAT: &str = "MODEL_EXPRESS_LOG_FORMAT";
48/// Maximum connection/request retries (`ClientArgs::max_retries`).
49pub const MODEL_EXPRESS_MAX_RETRIES: &str = "MODEL_EXPRESS_MAX_RETRIES";
50/// Delay between retries in seconds (`ClientArgs::retry_delay`).
51pub const MODEL_EXPRESS_RETRY_DELAY: &str = "MODEL_EXPRESS_RETRY_DELAY";
52/// Disable shared-storage mode (`ClientArgs::no_shared_storage`).
53pub const MODEL_EXPRESS_NO_SHARED_STORAGE: &str = "MODEL_EXPRESS_NO_SHARED_STORAGE";
54/// File-transfer chunk size in bytes (`ClientArgs::transfer_chunk_size`).
55pub const MODEL_EXPRESS_TRANSFER_CHUNK_SIZE: &str = "MODEL_EXPRESS_TRANSFER_CHUNK_SIZE";
56/// gRPC server listen port (`ServerArgs::port`).
57pub const MODEL_EXPRESS_SERVER_PORT: &str = "MODEL_EXPRESS_SERVER_PORT";
58/// Server host/bind address (`ServerArgs::host`).
59pub const MODEL_EXPRESS_SERVER_HOST: &str = "MODEL_EXPRESS_SERVER_HOST";
60/// Toggle the background cache-eviction sweeper (`ServerArgs::cache_eviction_enabled`).
61pub const MODEL_EXPRESS_CACHE_EVICTION_ENABLED: &str = "MODEL_EXPRESS_CACHE_EVICTION_ENABLED";
62/// Server endpoint used by the cache module's default-endpoint helper.
63pub const MODEL_EXPRESS_SERVER_ENDPOINT: &str = "MODEL_EXPRESS_SERVER_ENDPOINT";
64
65// ── HuggingFace ─────────────────────────────────────────────────────────────
66/// HuggingFace Hub auth token.
67pub const HF_TOKEN: &str = "HF_TOKEN";
68/// HuggingFace Hub cache directory.
69pub const HF_HUB_CACHE: &str = "HF_HUB_CACHE";
70/// Enables HuggingFace offline mode.
71pub const HF_HUB_OFFLINE: &str = "HF_HUB_OFFLINE";
72/// HuggingFace Hub endpoint override. Read directly by the `hf_hub` crate;
73/// registered here for reference (ModelExpress only sets it in tests).
74pub const HF_ENDPOINT: &str = "HF_ENDPOINT";
75
76// ── NGC ─────────────────────────────────────────────────────────────────────
77/// Base URL for the NGC artifact/download API.
78pub const NGC_API_ENDPOINT: &str = "NGC_API_ENDPOINT";
79/// Base URL for the NGC authentication endpoint.
80pub const NGC_AUTH_ENDPOINT: &str = "NGC_AUTH_ENDPOINT";
81/// NGC API key.
82pub const NGC_API_KEY: &str = "NGC_API_KEY";
83/// Alternate NGC CLI API key.
84pub const NGC_CLI_API_KEY: &str = "NGC_CLI_API_KEY";
85/// Root directory used to locate the NGC CLI config file (`~/.ngc/config`).
86pub const NGC_CLI_HOME: &str = "NGC_CLI_HOME";
87
88/// Default NGC API base URL when [`NGC_API_ENDPOINT`] is unset.
89pub const DEFAULT_NGC_API_BASE: &str = "https://api.ngc.nvidia.com";
90/// Default NGC auth base URL when [`NGC_AUTH_ENDPOINT`] is unset.
91pub const DEFAULT_NGC_AUTHN_BASE: &str = "https://authn.nvidia.com";
92
93// ── Redis / metadata backend (server) ───────────────────────────────────────
94/// Selects the metadata backend implementation (`redis`, `kubernetes`, `memory`).
95pub const MX_METADATA_BACKEND: &str = "MX_METADATA_BACKEND";
96/// Full Redis connection URL for the redis metadata backend.
97pub const REDIS_URL: &str = "REDIS_URL";
98/// Redis host (preferred) when building the URL from host + port.
99pub const MX_REDIS_HOST: &str = "MX_REDIS_HOST";
100/// Redis host alias for charts predating the `MX_` prefix.
101pub const REDIS_HOST: &str = "REDIS_HOST";
102/// Redis port (preferred) when building the URL from host + port.
103pub const MX_REDIS_PORT: &str = "MX_REDIS_PORT";
104/// Redis port alias for charts predating the `MX_` prefix.
105pub const REDIS_PORT: &str = "REDIS_PORT";
106/// Kubernetes namespace for ModelCacheEntry CRs (overrides [`POD_NAMESPACE`]).
107pub const MX_METADATA_NAMESPACE: &str = "MX_METADATA_NAMESPACE";
108/// Kubernetes namespace injected via the downward API for in-cluster pods.
109pub const POD_NAMESPACE: &str = "POD_NAMESPACE";
110/// Kubernetes pod name injected via the downward API (used by clients).
111pub const POD_NAME: &str = "POD_NAME";
112/// Kubernetes pod UID injected via the downward API (used by clients).
113pub const POD_UID: &str = "POD_UID";
114
115// ── Reaper (server) ─────────────────────────────────────────────────────────
116/// Interval (seconds) between reaper scans for stale/GC worker sweeps.
117pub const MX_REAPER_SCAN_INTERVAL_SECS: &str = "MX_REAPER_SCAN_INTERVAL_SECS";
118/// Age (seconds) after which an active worker's heartbeat is considered stale.
119pub const MX_HEARTBEAT_TIMEOUT_SECS: &str = "MX_HEARTBEAT_TIMEOUT_SECS";
120/// Age (seconds) after which a STALE worker is garbage-collected.
121pub const MX_GC_TIMEOUT_SECS: &str = "MX_GC_TIMEOUT_SECS";
122
123// ── System ──────────────────────────────────────────────────────────────────
124/// Primary source for the user's home directory.
125pub const HOME: &str = "HOME";
126/// Windows fallback for the home directory when [`HOME`] is unset.
127pub const USERPROFILE: &str = "USERPROFILE";
128/// Path to a kubeconfig file (consumed by the k8s integration tests).
129pub const KUBECONFIG: &str = "KUBECONFIG";
130
131// ── Default reaper timings ───────────────────────────────────────────────────
132const DEFAULT_REAPER_SCAN_INTERVAL_SECS: u64 = 30;
133const DEFAULT_HEARTBEAT_TIMEOUT_SECS: u64 = 90;
134const DEFAULT_GC_TIMEOUT_SECS: u64 = 3600;
135
136// ── Getters ───────────────────────────────────────────────────────────────
137
138/// Resolve the user's home directory: [`HOME`], then [`USERPROFILE`].
139///
140/// # Errors
141/// Returns an error when neither variable is set.
142pub fn home_dir() -> std::result::Result<String, Box<Error>> {
143    env::var(HOME)
144        .or_else(|_| env::var(USERPROFILE))
145        .map_err(|e| Error::Generic(format!("Failed to get home directory: {e}")).into())
146}
147
148/// Home directory as a `PathBuf`, falling back to `.` when unresolved.
149pub fn home_dir_or_cwd() -> PathBuf {
150    PathBuf::from(home_dir().unwrap_or_else(|_| ".".to_string()))
151}
152
153/// Model cache directory override from [`MODEL_EXPRESS_CACHE_DIRECTORY`].
154pub fn cache_directory() -> Option<PathBuf> {
155    env::var(MODEL_EXPRESS_CACHE_DIRECTORY)
156        .ok()
157        .map(PathBuf::from)
158}
159
160/// Default server gRPC endpoint: [`MODEL_EXPRESS_SERVER_ENDPOINT`] or
161/// `http://localhost:{DEFAULT_GRPC_PORT}`. Not normalized.
162pub fn server_endpoint_or_default() -> String {
163    env::var(MODEL_EXPRESS_SERVER_ENDPOINT)
164        .unwrap_or_else(|_| format!("http://localhost:{}", constants::DEFAULT_GRPC_PORT))
165}
166
167/// HuggingFace Hub token from [`HF_TOKEN`].
168pub fn hf_token() -> Option<String> {
169    env::var(HF_TOKEN).ok()
170}
171
172/// HuggingFace Hub cache directory from [`HF_HUB_CACHE`].
173pub fn hf_hub_cache() -> Option<PathBuf> {
174    env::var(HF_HUB_CACHE).ok().map(PathBuf::from)
175}
176
177/// Whether HuggingFace offline mode is enabled via [`HF_HUB_OFFLINE`].
178/// Enabled when the value is one of `1`, `ON`, `YES`, `TRUE` (case-insensitive).
179pub fn hf_offline() -> bool {
180    env::var(HF_HUB_OFFLINE)
181        .map(|v| matches!(v.to_uppercase().as_str(), "1" | "ON" | "YES" | "TRUE"))
182        .unwrap_or(false)
183}
184
185/// NGC API base URL: [`NGC_API_ENDPOINT`] or [`DEFAULT_NGC_API_BASE`].
186pub fn ngc_api_base() -> String {
187    env::var(NGC_API_ENDPOINT).unwrap_or_else(|_| DEFAULT_NGC_API_BASE.to_string())
188}
189
190/// NGC auth base URL: [`NGC_AUTH_ENDPOINT`] or [`DEFAULT_NGC_AUTHN_BASE`].
191pub fn ngc_authn_base() -> String {
192    env::var(NGC_AUTH_ENDPOINT).unwrap_or_else(|_| DEFAULT_NGC_AUTHN_BASE.to_string())
193}
194
195/// NGC API key from [`NGC_API_KEY`], then [`NGC_CLI_API_KEY`].
196/// Returns the first non-empty, trimmed value found.
197pub fn ngc_api_key() -> Option<String> {
198    for var in [NGC_API_KEY, NGC_CLI_API_KEY] {
199        if let Ok(v) = env::var(var) {
200            let trimmed = v.trim().to_string();
201            if !trimmed.is_empty() {
202                return Some(trimmed);
203            }
204        }
205    }
206    None
207}
208
209/// Root directory for the NGC CLI config from [`NGC_CLI_HOME`].
210pub fn ngc_cli_home() -> Option<PathBuf> {
211    env::var(NGC_CLI_HOME).ok().map(PathBuf::from)
212}
213
214/// Raw value of [`MX_METADATA_BACKEND`] (empty string when unset).
215pub fn metadata_backend() -> String {
216    env::var(MX_METADATA_BACKEND).unwrap_or_default()
217}
218
219/// Full Redis URL from [`REDIS_URL`].
220pub fn redis_url() -> Option<String> {
221    env::var(REDIS_URL).ok()
222}
223
224/// Redis host from [`MX_REDIS_HOST`], then [`REDIS_HOST`].
225pub fn redis_host() -> Option<String> {
226    env::var(MX_REDIS_HOST)
227        .or_else(|_| env::var(REDIS_HOST))
228        .ok()
229}
230
231/// Redis port from [`MX_REDIS_PORT`], then [`REDIS_PORT`].
232pub fn redis_port() -> Option<String> {
233    env::var(MX_REDIS_PORT)
234        .or_else(|_| env::var(REDIS_PORT))
235        .ok()
236}
237
238/// Kubernetes namespace from [`MX_METADATA_NAMESPACE`], then [`POD_NAMESPACE`].
239pub fn metadata_namespace() -> Option<String> {
240    env::var(MX_METADATA_NAMESPACE)
241        .or_else(|_| env::var(POD_NAMESPACE))
242        .ok()
243}
244
245/// Reaper scan interval in seconds ([`MX_REAPER_SCAN_INTERVAL_SECS`], default 30).
246pub fn reaper_scan_interval_secs() -> u64 {
247    env_u64(
248        MX_REAPER_SCAN_INTERVAL_SECS,
249        DEFAULT_REAPER_SCAN_INTERVAL_SECS,
250    )
251}
252
253/// Heartbeat staleness timeout in seconds ([`MX_HEARTBEAT_TIMEOUT_SECS`], default 90).
254pub fn heartbeat_timeout_secs() -> u64 {
255    env_u64(MX_HEARTBEAT_TIMEOUT_SECS, DEFAULT_HEARTBEAT_TIMEOUT_SECS)
256}
257
258/// Garbage-collection timeout in seconds ([`MX_GC_TIMEOUT_SECS`], default 3600).
259pub fn gc_timeout_secs() -> u64 {
260    env_u64(MX_GC_TIMEOUT_SECS, DEFAULT_GC_TIMEOUT_SECS)
261}
262
263/// Read an environment variable as `u64`, falling back to `default`.
264fn env_u64(name: &str, default: u64) -> u64 {
265    env::var(name)
266        .ok()
267        .and_then(|v| v.parse().ok())
268        .unwrap_or(default)
269}
270
271#[cfg(test)]
272#[allow(clippy::expect_used)]
273mod tests {
274    use super::*;
275    use crate::test_support::{EnvVarGuard, acquire_env_mutex};
276
277    #[test]
278    fn name_constants_match_their_literals() {
279        assert_eq!(MODEL_EXPRESS_PREFIX, "MODEL_EXPRESS");
280        assert_eq!(MODEL_EXPRESS_ENDPOINT, "MODEL_EXPRESS_ENDPOINT");
281        assert_eq!(MODEL_EXPRESS_TIMEOUT, "MODEL_EXPRESS_TIMEOUT");
282        assert_eq!(
283            MODEL_EXPRESS_CACHE_DIRECTORY,
284            "MODEL_EXPRESS_CACHE_DIRECTORY"
285        );
286        assert_eq!(MODEL_EXPRESS_LOG_LEVEL, "MODEL_EXPRESS_LOG_LEVEL");
287        assert_eq!(MODEL_EXPRESS_LOG_FORMAT, "MODEL_EXPRESS_LOG_FORMAT");
288        assert_eq!(MODEL_EXPRESS_MAX_RETRIES, "MODEL_EXPRESS_MAX_RETRIES");
289        assert_eq!(MODEL_EXPRESS_RETRY_DELAY, "MODEL_EXPRESS_RETRY_DELAY");
290        assert_eq!(
291            MODEL_EXPRESS_NO_SHARED_STORAGE,
292            "MODEL_EXPRESS_NO_SHARED_STORAGE"
293        );
294        assert_eq!(
295            MODEL_EXPRESS_TRANSFER_CHUNK_SIZE,
296            "MODEL_EXPRESS_TRANSFER_CHUNK_SIZE"
297        );
298        assert_eq!(MODEL_EXPRESS_SERVER_PORT, "MODEL_EXPRESS_SERVER_PORT");
299        assert_eq!(MODEL_EXPRESS_SERVER_HOST, "MODEL_EXPRESS_SERVER_HOST");
300        assert_eq!(
301            MODEL_EXPRESS_CACHE_EVICTION_ENABLED,
302            "MODEL_EXPRESS_CACHE_EVICTION_ENABLED"
303        );
304        assert_eq!(
305            MODEL_EXPRESS_SERVER_ENDPOINT,
306            "MODEL_EXPRESS_SERVER_ENDPOINT"
307        );
308        assert_eq!(HF_TOKEN, "HF_TOKEN");
309        assert_eq!(HF_HUB_CACHE, "HF_HUB_CACHE");
310        assert_eq!(HF_HUB_OFFLINE, "HF_HUB_OFFLINE");
311        assert_eq!(HF_ENDPOINT, "HF_ENDPOINT");
312        assert_eq!(NGC_API_ENDPOINT, "NGC_API_ENDPOINT");
313        assert_eq!(NGC_AUTH_ENDPOINT, "NGC_AUTH_ENDPOINT");
314        assert_eq!(NGC_API_KEY, "NGC_API_KEY");
315        assert_eq!(NGC_CLI_API_KEY, "NGC_CLI_API_KEY");
316        assert_eq!(NGC_CLI_HOME, "NGC_CLI_HOME");
317        assert_eq!(MX_METADATA_BACKEND, "MX_METADATA_BACKEND");
318        assert_eq!(REDIS_URL, "REDIS_URL");
319        assert_eq!(MX_REDIS_HOST, "MX_REDIS_HOST");
320        assert_eq!(REDIS_HOST, "REDIS_HOST");
321        assert_eq!(MX_REDIS_PORT, "MX_REDIS_PORT");
322        assert_eq!(REDIS_PORT, "REDIS_PORT");
323        assert_eq!(MX_METADATA_NAMESPACE, "MX_METADATA_NAMESPACE");
324        assert_eq!(POD_NAMESPACE, "POD_NAMESPACE");
325        assert_eq!(MX_REAPER_SCAN_INTERVAL_SECS, "MX_REAPER_SCAN_INTERVAL_SECS");
326        assert_eq!(MX_HEARTBEAT_TIMEOUT_SECS, "MX_HEARTBEAT_TIMEOUT_SECS");
327        assert_eq!(MX_GC_TIMEOUT_SECS, "MX_GC_TIMEOUT_SECS");
328        assert_eq!(HOME, "HOME");
329        assert_eq!(USERPROFILE, "USERPROFILE");
330        assert_eq!(KUBECONFIG, "KUBECONFIG");
331        assert_eq!(POD_NAME, "POD_NAME");
332        assert_eq!(POD_UID, "POD_UID");
333    }
334
335    #[test]
336    fn hf_offline_parses_truthy_values() {
337        let lock = acquire_env_mutex();
338        for truthy in ["1", "on", "YES", "true", "True"] {
339            let _g = EnvVarGuard::set(&lock, HF_HUB_OFFLINE, truthy);
340            assert!(hf_offline(), "expected {truthy} to enable offline mode");
341        }
342        for falsey in ["0", "off", "no", "maybe"] {
343            let _g = EnvVarGuard::set(&lock, HF_HUB_OFFLINE, falsey);
344            assert!(!hf_offline(), "expected {falsey} to disable offline mode");
345        }
346        let _g = EnvVarGuard::remove(&lock, HF_HUB_OFFLINE);
347        assert!(!hf_offline(), "unset should disable offline mode");
348    }
349
350    #[test]
351    fn ngc_bases_default_then_override() {
352        let lock = acquire_env_mutex();
353        let _api = EnvVarGuard::remove(&lock, NGC_API_ENDPOINT);
354        let _authn = EnvVarGuard::remove(&lock, NGC_AUTH_ENDPOINT);
355        assert_eq!(ngc_api_base(), DEFAULT_NGC_API_BASE);
356        assert_eq!(ngc_authn_base(), DEFAULT_NGC_AUTHN_BASE);
357
358        let _api = EnvVarGuard::set(&lock, NGC_API_ENDPOINT, "https://api.example.com");
359        let _authn = EnvVarGuard::set(&lock, NGC_AUTH_ENDPOINT, "https://authn.example.com");
360        assert_eq!(ngc_api_base(), "https://api.example.com");
361        assert_eq!(ngc_authn_base(), "https://authn.example.com");
362    }
363
364    #[test]
365    fn ngc_api_key_prefers_primary_then_falls_back() {
366        let lock = acquire_env_mutex();
367        let _p = EnvVarGuard::set(&lock, NGC_API_KEY, "  primary  ");
368        let _s = EnvVarGuard::set(&lock, NGC_CLI_API_KEY, "secondary");
369        assert_eq!(ngc_api_key().as_deref(), Some("primary"));
370
371        let _p = EnvVarGuard::remove(&lock, NGC_API_KEY);
372        assert_eq!(ngc_api_key().as_deref(), Some("secondary"));
373
374        let _s = EnvVarGuard::remove(&lock, NGC_CLI_API_KEY);
375        assert_eq!(ngc_api_key(), None);
376    }
377
378    #[test]
379    fn redis_and_namespace_fallbacks() {
380        let lock = acquire_env_mutex();
381        let _h1 = EnvVarGuard::remove(&lock, MX_REDIS_HOST);
382        let _h2 = EnvVarGuard::set(&lock, REDIS_HOST, "legacy-host");
383        assert_eq!(redis_host().as_deref(), Some("legacy-host"));
384        let _h1 = EnvVarGuard::set(&lock, MX_REDIS_HOST, "mx-host");
385        assert_eq!(redis_host().as_deref(), Some("mx-host"));
386
387        let _n1 = EnvVarGuard::remove(&lock, MX_METADATA_NAMESPACE);
388        let _n2 = EnvVarGuard::set(&lock, POD_NAMESPACE, "pod-ns");
389        assert_eq!(metadata_namespace().as_deref(), Some("pod-ns"));
390    }
391
392    #[test]
393    fn reaper_getters_default_parse_and_fallback() {
394        let lock = acquire_env_mutex();
395        let _g = EnvVarGuard::remove(&lock, MX_REAPER_SCAN_INTERVAL_SECS);
396        assert_eq!(
397            reaper_scan_interval_secs(),
398            DEFAULT_REAPER_SCAN_INTERVAL_SECS
399        );
400
401        let _g = EnvVarGuard::set(&lock, MX_HEARTBEAT_TIMEOUT_SECS, "120");
402        assert_eq!(heartbeat_timeout_secs(), 120);
403
404        let _g = EnvVarGuard::set(&lock, MX_GC_TIMEOUT_SECS, "not-a-number");
405        assert_eq!(gc_timeout_secs(), DEFAULT_GC_TIMEOUT_SECS);
406    }
407}
modelexpress_common/envs.rs

modelexpress_common/
envs.rs
