modelexpress_common/
client_config.rs

1// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4use crate::cache::CacheConfig;
5use crate::config::{
6    ConnectionConfig, LogFormat, LogLevel, load_layered_config, normalize_grpc_endpoint,
7};
8use anyhow::Result;
9use clap::Parser;
10use config::ConfigError;
11use serde::{Deserialize, Serialize};
12use std::path::PathBuf;
13
14/// Shared command line arguments for the ModelExpress client.
15///
16/// # Adding New Arguments
17///
18/// This struct is the **single source of truth** for client CLI arguments and environment
19/// variables. It is shared between:
20/// - The `modelexpress-cli` binary (via `#[command(flatten)]` in the `Cli` struct)
21/// - Any other client binaries that need these arguments
22/// - The `ClientConfig::load()` function which applies these values
23///
24/// When adding a new argument:
25/// 1. Add the field here with appropriate `#[arg(...)]` attributes
26/// 2. Include `env = "MODEL_EXPRESS_..."` for environment variable support
27/// 3. Update `ClientConfig::load()` to apply the new argument to the config
28/// 4. Add tests in the `tests` module below
29/// 5. Update CLI.md documentation if applicable
30///
31/// # Short Flags
32///
33/// Avoid using `-v` as a short flag here - it's reserved for the CLI's `--verbose` flag
34/// which uses `-v`, `-vv`, `-vvv` counting. The CLI embeds this struct via flatten,
35/// so short flag conflicts will cause runtime panics.
36#[derive(Parser, Debug, Clone)]
37#[command(author, version, about, long_about = None)]
38pub struct ClientArgs {
39    /// Configuration file path
40    #[arg(short, long, value_name = "FILE")]
41    pub config: Option<PathBuf>,
42
43    /// Server endpoint
44    #[arg(short, long, env = crate::envs::MODEL_EXPRESS_ENDPOINT)]
45    pub endpoint: Option<String>,
46
47    /// Request timeout in seconds
48    #[arg(short, long, env = crate::envs::MODEL_EXPRESS_TIMEOUT)]
49    pub timeout: Option<u64>,
50
51    /// Cache path override
52    #[arg(long, env = crate::envs::MODEL_EXPRESS_CACHE_DIRECTORY)]
53    pub cache_path: Option<PathBuf>,
54
55    /// Log level (no short flag to avoid conflict with CLI's -v/--verbose)
56    #[arg(long, env = crate::envs::MODEL_EXPRESS_LOG_LEVEL, value_enum)]
57    pub log_level: Option<LogLevel>,
58
59    /// Log format
60    #[arg(long, env = crate::envs::MODEL_EXPRESS_LOG_FORMAT, value_enum)]
61    pub log_format: Option<LogFormat>,
62
63    /// Quiet mode (suppress all output except errors)
64    #[arg(long, short = 'q')]
65    pub quiet: bool,
66
67    /// Maximum number of retries
68    #[arg(long, env = crate::envs::MODEL_EXPRESS_MAX_RETRIES)]
69    pub max_retries: Option<u32>,
70
71    /// Retry delay in seconds
72    #[arg(long, env = crate::envs::MODEL_EXPRESS_RETRY_DELAY)]
73    pub retry_delay: Option<u64>,
74
75    /// Disable shared storage mode (will transfer files from server to client)
76    #[arg(long, env = crate::envs::MODEL_EXPRESS_NO_SHARED_STORAGE)]
77    pub no_shared_storage: bool,
78
79    /// Chunk size in bytes for file transfer when shared storage is disabled
80    #[arg(long, env = crate::envs::MODEL_EXPRESS_TRANSFER_CHUNK_SIZE)]
81    pub transfer_chunk_size: Option<usize>,
82}
83
84/// Complete client configuration
85#[derive(Debug, Clone, Serialize, Deserialize, Default)]
86pub struct ClientConfig {
87    /// Connection settings
88    pub connection: ConnectionConfig,
89    /// Cache configuration
90    pub cache: CacheConfig,
91    /// Logging configuration
92    pub logging: LoggingConfig,
93}
94
95/// Logging configuration for the client
96#[derive(Debug, Clone, Serialize, Deserialize, Default)]
97pub struct LoggingConfig {
98    /// Log level
99    #[serde(default)]
100    pub level: LogLevel,
101    /// Log format
102    #[serde(default)]
103    pub format: LogFormat,
104    /// Quiet mode
105    pub quiet: bool,
106}
107
108impl ClientConfig {
109    /// Load configuration from multiple sources in order of precedence:
110    /// 1. Command line arguments (highest priority)
111    /// 2. Environment variables (handled by clap's `env` attribute on `ClientArgs`)
112    /// 3. Configuration file
113    /// 4. Default values (lowest priority)
114    ///
115    /// # Adding New Arguments
116    ///
117    /// When you add a new field to `ClientArgs`:
118    /// 1. Add the corresponding override logic below in the "Apply CLI argument overrides" section
119    /// 2. Map the `ClientArgs` field to the appropriate `ClientConfig` field
120    /// 3. Add a test in the `tests` module to verify the override works
121    pub fn load(args: ClientArgs) -> Result<Self, ConfigError> {
122        // Start with layered config loading (file + env + defaults)
123        let mut config = load_layered_config(
124            args.config.clone(),
125            crate::envs::MODEL_EXPRESS_PREFIX,
126            Self::default(),
127        )?;
128
129        // ==================== APPLY CLI ARGUMENT OVERRIDES ====================
130        // When adding a new field to ClientArgs, add the override logic here.
131        // These overrides apply CLI arguments (which include env vars via clap)
132        // on top of the config file values.
133
134        // Connection settings
135        if let Some(endpoint) = args.endpoint {
136            config.connection.endpoint = normalize_grpc_endpoint(endpoint);
137        }
138
139        if let Some(timeout) = args.timeout {
140            config.connection.timeout_secs = Some(timeout);
141        }
142
143        if let Some(max_retries) = args.max_retries {
144            config.connection.max_retries = Some(max_retries);
145        }
146
147        if let Some(retry_delay) = args.retry_delay {
148            config.connection.retry_delay_secs = Some(retry_delay);
149        }
150
151        // Cache settings
152        if let Some(cache_path) = args.cache_path {
153            config.cache.local_path = cache_path;
154        }
155
156        if args.no_shared_storage {
157            config.cache.shared_storage = false;
158        }
159
160        if let Some(chunk_size) = args.transfer_chunk_size {
161            config.cache.transfer_chunk_size = chunk_size;
162        }
163
164        // Logging settings
165        if let Some(log_level) = args.log_level {
166            config.logging.level = log_level;
167        }
168
169        if let Some(log_format) = args.log_format {
170            config.logging.format = log_format;
171        }
172
173        if args.quiet {
174            config.logging.quiet = true;
175        }
176
177        // ==================== END CLI ARGUMENT OVERRIDES ====================
178
179        config.connection.endpoint =
180            normalize_grpc_endpoint(std::mem::take(&mut config.connection.endpoint));
181        config.cache.server_endpoint =
182            normalize_grpc_endpoint(std::mem::take(&mut config.cache.server_endpoint));
183
184        // Validate configuration
185        config.validate()?;
186
187        Ok(config)
188    }
189
190    /// Validate the configuration
191    pub fn validate(&self) -> Result<(), ConfigError> {
192        // Validate endpoint
193        if self.connection.endpoint.is_empty() {
194            return Err(ConfigError::Message(
195                "Server endpoint cannot be empty".to_string(),
196            ));
197        }
198
199        // Validate timeout
200        if let Some(timeout) = self.connection.timeout_secs
201            && timeout == 0
202        {
203            return Err(ConfigError::Message(
204                "Timeout must be greater than 0".to_string(),
205            ));
206        }
207
208        // Validate cache path exists or can be created
209        if !self.cache.local_path.exists()
210            && let Err(e) = std::fs::create_dir_all(&self.cache.local_path)
211        {
212            return Err(ConfigError::Message(format!(
213                "Cannot create cache directory {:?}: {}",
214                self.cache.local_path, e
215            )));
216        }
217
218        Ok(())
219    }
220
221    /// Get the gRPC endpoint for backward compatibility
222    pub fn grpc_endpoint(&self) -> &str {
223        &self.connection.endpoint
224    }
225
226    /// Get the timeout in seconds for backward compatibility
227    pub fn timeout_secs(&self) -> Option<u64> {
228        self.connection.timeout_secs
229    }
230
231    /// Create a simple client config for testing
232    pub fn for_testing(endpoint: impl Into<String>) -> Self {
233        Self {
234            connection: ConnectionConfig::new(endpoint),
235            cache: CacheConfig::default(),
236            logging: LoggingConfig::default(),
237        }
238    }
239
240    /// Apply cache path override if provided
241    pub fn with_cache_path(mut self, cache_path: Option<PathBuf>) -> Self {
242        if let Some(path) = cache_path {
243            self.cache.local_path = path;
244        }
245        self
246    }
247
248    /// Set timeout for the connection
249    pub fn with_timeout(mut self, timeout_secs: u64) -> Self {
250        self.connection.timeout_secs = Some(timeout_secs);
251        self
252    }
253
254    /// Set the server endpoint for both connection and cache
255    pub fn with_endpoint(mut self, endpoint: String) -> Self {
256        let endpoint = normalize_grpc_endpoint(endpoint);
257        self.connection.endpoint = endpoint.clone();
258        self.cache.server_endpoint = endpoint;
259        self
260    }
261}
262
263#[cfg(test)]
264#[allow(clippy::expect_used)]
265mod tests {
266    use super::*;
267    use crate::constants;
268
269    #[test]
270    fn test_client_config_default() {
271        let config = ClientConfig::default();
272        assert!(config.connection.endpoint.contains("8001"));
273        assert_eq!(config.connection.timeout_secs, Some(30));
274        assert!(!config.logging.quiet);
275    }
276
277    #[test]
278    fn test_client_config_for_testing() {
279        let config = ClientConfig::for_testing("http://test.example.com:1234");
280        assert_eq!(config.connection.endpoint, "http://test.example.com:1234");
281    }
282
283    #[test]
284    fn test_client_config_with_endpoint() {
285        let config =
286            ClientConfig::default().with_endpoint("http://custom.example.com:5678".to_string());
287
288        assert_eq!(config.connection.endpoint, "http://custom.example.com:5678");
289        assert_eq!(
290            config.cache.server_endpoint,
291            "http://custom.example.com:5678"
292        );
293    }
294
295    #[test]
296    fn test_client_config_with_endpoint_accepts_bare_host_port() {
297        let config = ClientConfig::default().with_endpoint("custom.example.com:5678".to_string());
298
299        assert_eq!(config.connection.endpoint, "http://custom.example.com:5678");
300        assert_eq!(
301            config.cache.server_endpoint,
302            "http://custom.example.com:5678"
303        );
304    }
305
306    #[test]
307    fn test_client_config_validation() {
308        let mut config = ClientConfig::default();
309        assert!(config.validate().is_ok());
310
311        config.connection.endpoint = String::new();
312        assert!(config.validate().is_err());
313    }
314
315    #[test]
316    fn test_client_config_backward_compatibility() {
317        let config = ClientConfig::for_testing("http://test.com:8080");
318        assert_eq!(config.grpc_endpoint(), "http://test.com:8080");
319        assert_eq!(config.timeout_secs(), Some(30));
320    }
321
322    #[test]
323    fn test_client_config_shared_storage_defaults() {
324        let config = ClientConfig::default();
325        assert!(config.cache.shared_storage);
326        assert_eq!(
327            config.cache.transfer_chunk_size,
328            constants::DEFAULT_TRANSFER_CHUNK_SIZE
329        );
330    }
331
332    #[test]
333    fn test_client_config_shared_storage_override() {
334        let mut config = ClientConfig::default();
335        config.cache.shared_storage = false;
336        config.cache.transfer_chunk_size = 64 * 1024;
337
338        assert!(!config.cache.shared_storage);
339        assert_eq!(config.cache.transfer_chunk_size, 64 * 1024);
340    }
341
342    #[test]
343    fn test_client_args_parse_defaults() {
344        // Test that ClientArgs can be parsed with no arguments (uses defaults)
345        let args = ClientArgs::try_parse_from(["test"]).expect("Failed to parse empty args");
346
347        assert!(args.endpoint.is_none());
348        assert!(args.timeout.is_none());
349        assert!(args.cache_path.is_none());
350        assert!(!args.quiet);
351        assert!(!args.no_shared_storage);
352        assert!(args.transfer_chunk_size.is_none());
353    }
354
355    #[test]
356    fn test_client_args_parse_cli_flags() {
357        // Test parsing various CLI flags
358        let args = ClientArgs::try_parse_from([
359            "test",
360            "--endpoint",
361            "http://custom:9000",
362            "--timeout",
363            "60",
364            "--quiet",
365            "--no-shared-storage",
366            "--transfer-chunk-size",
367            "1048576",
368        ])
369        .expect("Failed to parse CLI args");
370
371        assert_eq!(args.endpoint, Some("http://custom:9000".to_string()));
372        assert_eq!(args.timeout, Some(60));
373        assert!(args.quiet);
374        assert!(args.no_shared_storage);
375        assert_eq!(args.transfer_chunk_size, Some(1048576));
376    }
377
378    #[test]
379    fn test_client_args_short_flags() {
380        // Test short flag variants (-e for endpoint, -t for timeout, -q for quiet)
381        let args =
382            ClientArgs::try_parse_from(["test", "-e", "http://short:8000", "-t", "45", "-q"])
383                .expect("Failed to parse short flags");
384
385        assert_eq!(args.endpoint, Some("http://short:8000".to_string()));
386        assert_eq!(args.timeout, Some(45));
387        assert!(args.quiet);
388    }
389
390    #[test]
391    fn test_client_args_log_level() {
392        // Test --log-level flag (no short flag to avoid conflict with CLI's -v)
393        let args = ClientArgs::try_parse_from(["test", "--log-level", "debug"])
394            .expect("Failed to parse log level");
395
396        assert_eq!(args.log_level, Some(LogLevel::Debug));
397    }
398
399    #[test]
400    fn test_client_config_load_applies_cli_args() {
401        // Test that ClientConfig::load() properly applies CLI arguments
402        let args = ClientArgs {
403            config: None,
404            endpoint: Some("cli-override:7777".to_string()),
405            timeout: Some(120),
406            cache_path: None,
407            log_level: None,
408            log_format: None,
409            quiet: true,
410            max_retries: Some(5),
411            retry_delay: Some(10),
412            no_shared_storage: true,
413            transfer_chunk_size: Some(2097152),
414        };
415
416        let config = ClientConfig::load(args).expect("Failed to load config");
417
418        assert_eq!(config.connection.endpoint, "http://cli-override:7777");
419        assert_eq!(config.connection.timeout_secs, Some(120));
420        assert!(config.logging.quiet);
421        assert_eq!(config.connection.max_retries, Some(5));
422        assert_eq!(config.connection.retry_delay_secs, Some(10));
423        assert!(!config.cache.shared_storage);
424        assert_eq!(config.cache.transfer_chunk_size, 2097152);
425    }
426}
modelexpress_common/client_config.rs

modelexpress_common/
client_config.rs
