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    /// Disable shared storage mode (will transfer files from server to client)
68    #[arg(long, env = crate::envs::MODEL_EXPRESS_NO_SHARED_STORAGE)]
69    pub no_shared_storage: bool,
70
71    /// Chunk size in bytes for file transfer when shared storage is disabled
72    #[arg(long, env = crate::envs::MODEL_EXPRESS_TRANSFER_CHUNK_SIZE)]
73    pub transfer_chunk_size: Option<usize>,
74}
75
76/// Complete client configuration
77#[derive(Debug, Clone, Serialize, Deserialize, Default)]
78pub struct ClientConfig {
79    /// Connection settings
80    pub connection: ConnectionConfig,
81    /// Cache configuration
82    pub cache: CacheConfig,
83    /// Logging configuration
84    pub logging: LoggingConfig,
85}
86
87/// Logging configuration for the client
88#[derive(Debug, Clone, Serialize, Deserialize, Default)]
89pub struct LoggingConfig {
90    /// Log level
91    #[serde(default)]
92    pub level: LogLevel,
93    /// Log format
94    #[serde(default)]
95    pub format: LogFormat,
96    /// Quiet mode
97    pub quiet: bool,
98}
99
100impl ClientConfig {
101    /// Load configuration from multiple sources in order of precedence:
102    /// 1. Command line arguments (highest priority)
103    /// 2. Environment variables (handled by clap's `env` attribute on `ClientArgs`)
104    /// 3. Configuration file
105    /// 4. Default values (lowest priority)
106    ///
107    /// # Adding New Arguments
108    ///
109    /// When you add a new field to `ClientArgs`:
110    /// 1. Add the corresponding override logic below in the "Apply CLI argument overrides" section
111    /// 2. Map the `ClientArgs` field to the appropriate `ClientConfig` field
112    /// 3. Add a test in the `tests` module to verify the override works
113    pub fn load(args: ClientArgs) -> Result<Self, ConfigError> {
114        // Start with layered config loading (file + env + defaults)
115        let mut config = load_layered_config(
116            args.config.clone(),
117            crate::envs::MODEL_EXPRESS_PREFIX,
118            Self::default(),
119        )?;
120
121        // ==================== APPLY CLI ARGUMENT OVERRIDES ====================
122        // When adding a new field to ClientArgs, add the override logic here.
123        // These overrides apply CLI arguments (which include env vars via clap)
124        // on top of the config file values.
125
126        // Connection settings
127        if let Some(endpoint) = args.endpoint {
128            config.connection.endpoint = normalize_grpc_endpoint(endpoint);
129        }
130
131        if let Some(timeout) = args.timeout {
132            config.connection.timeout_secs = Some(timeout);
133        }
134
135        // Cache settings
136        if let Some(cache_path) = args.cache_path {
137            config.cache.local_path = cache_path;
138        }
139
140        if args.no_shared_storage {
141            config.cache.shared_storage = false;
142        }
143
144        if let Some(chunk_size) = args.transfer_chunk_size {
145            config.cache.transfer_chunk_size = chunk_size;
146        }
147
148        // Logging settings
149        if let Some(log_level) = args.log_level {
150            config.logging.level = log_level;
151        }
152
153        if let Some(log_format) = args.log_format {
154            config.logging.format = log_format;
155        }
156
157        if args.quiet {
158            config.logging.quiet = true;
159        }
160
161        // ==================== END CLI ARGUMENT OVERRIDES ====================
162
163        config.connection.endpoint =
164            normalize_grpc_endpoint(std::mem::take(&mut config.connection.endpoint));
165        config.cache.server_endpoint =
166            normalize_grpc_endpoint(std::mem::take(&mut config.cache.server_endpoint));
167
168        // Validate configuration
169        config.validate()?;
170
171        Ok(config)
172    }
173
174    /// Validate the configuration
175    pub fn validate(&self) -> Result<(), ConfigError> {
176        // Validate endpoint
177        if self.connection.endpoint.is_empty() {
178            return Err(ConfigError::Message(
179                "Server endpoint cannot be empty".to_string(),
180            ));
181        }
182
183        // Validate timeout
184        if let Some(timeout) = self.connection.timeout_secs
185            && timeout == 0
186        {
187            return Err(ConfigError::Message(
188                "Timeout must be greater than 0".to_string(),
189            ));
190        }
191
192        // Validate cache path exists or can be created
193        if !self.cache.local_path.exists()
194            && let Err(e) = std::fs::create_dir_all(&self.cache.local_path)
195        {
196            return Err(ConfigError::Message(format!(
197                "Cannot create cache directory {:?}: {}",
198                self.cache.local_path, e
199            )));
200        }
201
202        Ok(())
203    }
204
205    /// Get the gRPC endpoint for backward compatibility
206    pub fn grpc_endpoint(&self) -> &str {
207        &self.connection.endpoint
208    }
209
210    /// Get the timeout in seconds for backward compatibility
211    pub fn timeout_secs(&self) -> Option<u64> {
212        self.connection.timeout_secs
213    }
214
215    /// Create a simple client config for testing
216    pub fn for_testing(endpoint: impl Into<String>) -> Self {
217        Self {
218            connection: ConnectionConfig::new(endpoint),
219            cache: CacheConfig::default(),
220            logging: LoggingConfig::default(),
221        }
222    }
223
224    /// Set timeout for the connection
225    pub fn with_timeout(mut self, timeout_secs: u64) -> Self {
226        self.connection.timeout_secs = Some(timeout_secs);
227        self
228    }
229
230    /// Set the server endpoint for both connection and cache
231    pub fn with_endpoint(mut self, endpoint: String) -> Self {
232        let endpoint = normalize_grpc_endpoint(endpoint);
233        self.connection.endpoint = endpoint.clone();
234        self.cache.server_endpoint = endpoint;
235        self
236    }
237}
238
239#[cfg(test)]
240#[allow(clippy::expect_used)]
241mod tests {
242    use super::*;
243    use crate::constants;
244
245    #[test]
246    fn test_client_config_default() {
247        let config = ClientConfig::default();
248        assert!(config.connection.endpoint.contains("8001"));
249        assert_eq!(config.connection.timeout_secs, Some(30));
250        assert!(!config.logging.quiet);
251    }
252
253    #[test]
254    fn test_client_config_for_testing() {
255        let config = ClientConfig::for_testing("http://test.example.com:1234");
256        assert_eq!(config.connection.endpoint, "http://test.example.com:1234");
257    }
258
259    #[test]
260    fn test_client_config_with_endpoint() {
261        let config =
262            ClientConfig::default().with_endpoint("http://custom.example.com:5678".to_string());
263
264        assert_eq!(config.connection.endpoint, "http://custom.example.com:5678");
265        assert_eq!(
266            config.cache.server_endpoint,
267            "http://custom.example.com:5678"
268        );
269    }
270
271    #[test]
272    fn test_client_config_with_endpoint_accepts_bare_host_port() {
273        let config = ClientConfig::default().with_endpoint("custom.example.com:5678".to_string());
274
275        assert_eq!(config.connection.endpoint, "http://custom.example.com:5678");
276        assert_eq!(
277            config.cache.server_endpoint,
278            "http://custom.example.com:5678"
279        );
280    }
281
282    #[test]
283    fn test_client_config_validation() {
284        let mut config = ClientConfig::default();
285        assert!(config.validate().is_ok());
286
287        config.connection.endpoint = String::new();
288        assert!(config.validate().is_err());
289    }
290
291    #[test]
292    fn test_client_config_backward_compatibility() {
293        let config = ClientConfig::for_testing("http://test.com:8080");
294        assert_eq!(config.grpc_endpoint(), "http://test.com:8080");
295        assert_eq!(config.timeout_secs(), Some(30));
296    }
297
298    #[test]
299    fn test_client_config_shared_storage_defaults() {
300        let config = ClientConfig::default();
301        assert!(config.cache.shared_storage);
302        assert_eq!(
303            config.cache.transfer_chunk_size,
304            constants::DEFAULT_TRANSFER_CHUNK_SIZE
305        );
306    }
307
308    #[test]
309    fn test_client_config_shared_storage_override() {
310        let mut config = ClientConfig::default();
311        config.cache.shared_storage = false;
312        config.cache.transfer_chunk_size = 64 * 1024;
313
314        assert!(!config.cache.shared_storage);
315        assert_eq!(config.cache.transfer_chunk_size, 64 * 1024);
316    }
317
318    #[test]
319    fn test_client_args_parse_defaults() {
320        // Test that ClientArgs can be parsed with no arguments (uses defaults)
321        let args = ClientArgs::try_parse_from(["test"]).expect("Failed to parse empty args");
322
323        assert!(args.endpoint.is_none());
324        assert!(args.timeout.is_none());
325        assert!(args.cache_path.is_none());
326        assert!(!args.quiet);
327        assert!(!args.no_shared_storage);
328        assert!(args.transfer_chunk_size.is_none());
329    }
330
331    #[test]
332    fn test_client_args_parse_cli_flags() {
333        // Test parsing various CLI flags
334        let args = ClientArgs::try_parse_from([
335            "test",
336            "--endpoint",
337            "http://custom:9000",
338            "--timeout",
339            "60",
340            "--quiet",
341            "--no-shared-storage",
342            "--transfer-chunk-size",
343            "1048576",
344        ])
345        .expect("Failed to parse CLI args");
346
347        assert_eq!(args.endpoint, Some("http://custom:9000".to_string()));
348        assert_eq!(args.timeout, Some(60));
349        assert!(args.quiet);
350        assert!(args.no_shared_storage);
351        assert_eq!(args.transfer_chunk_size, Some(1048576));
352    }
353
354    #[test]
355    fn test_client_args_short_flags() {
356        // Test short flag variants (-e for endpoint, -t for timeout, -q for quiet)
357        let args =
358            ClientArgs::try_parse_from(["test", "-e", "http://short:8000", "-t", "45", "-q"])
359                .expect("Failed to parse short flags");
360
361        assert_eq!(args.endpoint, Some("http://short:8000".to_string()));
362        assert_eq!(args.timeout, Some(45));
363        assert!(args.quiet);
364    }
365
366    #[test]
367    fn test_client_args_log_level() {
368        // Test --log-level flag (no short flag to avoid conflict with CLI's -v)
369        let args = ClientArgs::try_parse_from(["test", "--log-level", "debug"])
370            .expect("Failed to parse log level");
371
372        assert_eq!(args.log_level, Some(LogLevel::Debug));
373    }
374
375    #[test]
376    fn test_client_config_load_applies_cli_args() {
377        // Test that ClientConfig::load() properly applies CLI arguments
378        let args = ClientArgs {
379            config: None,
380            endpoint: Some("cli-override:7777".to_string()),
381            timeout: Some(120),
382            cache_path: None,
383            log_level: None,
384            log_format: None,
385            quiet: true,
386            no_shared_storage: true,
387            transfer_chunk_size: Some(2097152),
388        };
389
390        let config = ClientConfig::load(args).expect("Failed to load config");
391
392        assert_eq!(config.connection.endpoint, "http://cli-override:7777");
393        assert_eq!(config.connection.timeout_secs, Some(120));
394        assert!(config.logging.quiet);
395        assert!(!config.cache.shared_storage);
396        assert_eq!(config.cache.transfer_chunk_size, 2097152);
397    }
398}
modelexpress_common/client_config.rs

modelexpress_common/
client_config.rs
