1use crate::{
5 Utils,
6 config::normalize_grpc_endpoint,
7 constants,
8 models::ModelProvider,
9 providers::{
10 gcs::GcsProviderCache, huggingface::HuggingFaceProviderCache, ngc::NgcProviderCache,
11 s3::S3ProviderCache,
12 },
13};
14use anyhow::{Context, Result};
15use serde::{Deserialize, Serialize};
16use std::env;
17use std::fs;
18use std::path::{Path, PathBuf};
19use tracing::{debug, info, warn};
20
21#[derive(Debug, Clone, Serialize, Deserialize)]
23pub struct CacheConfig {
24 pub local_path: PathBuf,
26 pub server_endpoint: String,
28 pub timeout_secs: Option<u64>,
30 #[serde(default = "default_shared_storage")]
33 pub shared_storage: bool,
34 #[serde(default = "default_transfer_chunk_size")]
36 pub transfer_chunk_size: usize,
37}
38
39fn default_shared_storage() -> bool {
40 constants::DEFAULT_SHARED_STORAGE
41}
42
43fn default_transfer_chunk_size() -> usize {
44 constants::DEFAULT_TRANSFER_CHUNK_SIZE
45}
46
47impl Default for CacheConfig {
48 fn default() -> Self {
49 let home = Utils::get_home_dir().unwrap_or_else(|_| ".".to_string());
50 Self {
51 local_path: PathBuf::from(home).join(constants::DEFAULT_CACHE_PATH),
52 server_endpoint: format!("http://localhost:{}", constants::DEFAULT_GRPC_PORT),
53 timeout_secs: None,
54 shared_storage: constants::DEFAULT_SHARED_STORAGE,
55 transfer_chunk_size: constants::DEFAULT_TRANSFER_CHUNK_SIZE,
56 }
57 }
58}
59
60impl CacheConfig {
61 pub fn discover() -> Result<Self> {
63 if let Some(path) = Self::get_cache_path_from_args() {
72 return Self::from_path(path);
73 }
74
75 if let Some(path) = crate::envs::cache_directory() {
77 return Self::from_path(path);
78 }
79
80 if let Ok(config) = Self::from_config_file() {
82 return Ok(config);
83 }
84
85 if let Ok(config) = Self::auto_detect() {
87 return Ok(config);
88 }
89
90 debug!("Using default cache configuration");
92 Ok(Self::default())
93 }
94
95 pub fn new(local_path: PathBuf, server_endpoint: Option<String>) -> Result<Self> {
97 fs::create_dir_all(&local_path)
99 .with_context(|| format!("Failed to create cache directory: {local_path:?}"))?;
100
101 Ok(Self {
102 local_path,
103 server_endpoint: normalize_grpc_endpoint(
104 server_endpoint.unwrap_or_else(Self::get_default_server_endpoint),
105 ),
106 timeout_secs: None,
107 shared_storage: constants::DEFAULT_SHARED_STORAGE,
108 transfer_chunk_size: constants::DEFAULT_TRANSFER_CHUNK_SIZE,
109 })
110 }
111
112 pub fn from_path<P: AsRef<Path>>(path: P) -> Result<Self> {
114 let local_path = path.as_ref().to_path_buf();
115
116 fs::create_dir_all(&local_path)
118 .with_context(|| format!("Failed to create cache directory: {local_path:?}"))?;
119
120 Ok(Self {
121 local_path,
122 server_endpoint: Self::get_default_server_endpoint(),
123 timeout_secs: None,
124 shared_storage: constants::DEFAULT_SHARED_STORAGE,
125 transfer_chunk_size: constants::DEFAULT_TRANSFER_CHUNK_SIZE,
126 })
127 }
128
129 pub fn from_config_file() -> Result<Self> {
131 let config_path = Self::get_config_path()?;
132
133 if !config_path.exists() {
134 return Err(anyhow::anyhow!("Config file not found: {:?}", config_path));
135 }
136
137 let content = fs::read_to_string(&config_path)
138 .with_context(|| format!("Failed to read config file: {config_path:?}"))?;
139
140 let mut config: Self = serde_yaml::from_str(&content)
141 .with_context(|| format!("Failed to parse config file: {config_path:?}"))?;
142 config.server_endpoint =
143 normalize_grpc_endpoint(std::mem::take(&mut config.server_endpoint));
144
145 Ok(config)
146 }
147
148 pub fn save_to_config_file(&self) -> Result<()> {
150 let config_path = Self::get_config_path()?;
151
152 if let Some(parent) = config_path.parent() {
154 fs::create_dir_all(parent)
155 .with_context(|| format!("Failed to create config directory: {parent:?}"))?;
156 }
157
158 let content = serde_yaml::to_string(self).context("Failed to serialize config")?;
159
160 fs::write(&config_path, content)
161 .with_context(|| format!("Failed to write config file: {config_path:?}"))?;
162
163 Ok(())
164 }
165
166 pub fn auto_detect() -> Result<Self> {
168 let home = Utils::get_home_dir().unwrap_or_else(|_| ".".to_string());
169 let common_paths = vec![
170 PathBuf::from(&home).join(constants::DEFAULT_CACHE_PATH),
171 PathBuf::from(&home).join(constants::DEFAULT_HF_CACHE_PATH),
172 PathBuf::from("/cache"),
173 PathBuf::from("/app/models"),
174 PathBuf::from("./cache"),
175 PathBuf::from("./models"),
176 ];
177
178 for path in common_paths {
179 if path.exists() && path.is_dir() {
180 return Ok(Self {
181 local_path: path,
182 server_endpoint: Self::get_default_server_endpoint(),
183 timeout_secs: None,
184 shared_storage: constants::DEFAULT_SHARED_STORAGE,
185 transfer_chunk_size: constants::DEFAULT_TRANSFER_CHUNK_SIZE,
186 });
187 }
188 }
189
190 Err(anyhow::anyhow!(
191 "No cache directory found in common locations"
192 ))
193 }
194
195 fn get_cache_path_from_args() -> Option<String> {
197 let args: Vec<String> = env::args().collect();
198
199 for (i, arg) in args.iter().enumerate() {
200 if arg == "--cache-path"
201 && let Some(next_arg) = args.get(i.saturating_add(1))
202 {
203 return Some(next_arg.clone());
204 }
205 }
206
207 None
208 }
209
210 fn get_default_server_endpoint() -> String {
212 normalize_grpc_endpoint(crate::envs::server_endpoint_or_default())
213 }
214
215 fn get_config_path() -> Result<PathBuf> {
217 let home = Utils::get_home_dir().unwrap_or_else(|_| ".".to_string());
218
219 Ok(PathBuf::from(home).join(constants::DEFAULT_CONFIG_PATH))
220 }
221
222 pub fn get_cache_stats(&self) -> Result<CacheStats> {
224 let mut models = Vec::new();
225
226 if !self.local_path.exists() {
227 return Ok(CacheStats {
228 total_models: 0,
229 total_size: 0,
230 models,
231 });
232 }
233
234 for provider in [
235 ModelProvider::HuggingFace,
236 ModelProvider::Ngc,
237 ModelProvider::Gcs,
238 ModelProvider::S3,
239 ] {
240 models.extend(cache_for_provider(provider).list_models(&self.local_path)?);
241 }
242
243 models.sort_by(|left, right| {
244 provider_sort_key(left.provider)
245 .cmp(&provider_sort_key(right.provider))
246 .then_with(|| left.name.cmp(&right.name))
247 });
248
249 let total_size = models.iter().map(|model| model.size).sum();
250
251 Ok(CacheStats {
252 total_models: models.len(),
253 total_size,
254 models,
255 })
256 }
257
258 pub fn clear_model(&self, model_name: &str, provider: ModelProvider) -> Result<()> {
260 cache_for_provider(provider).clear_model(&self.local_path, model_name)
261 }
262
263 pub fn clear_all(&self) -> Result<()> {
265 if self.local_path.exists() {
266 for entry in fs::read_dir(&self.local_path)
267 .with_context(|| format!("Failed to read cache directory: {:?}", self.local_path))?
268 {
269 let entry = entry
270 .with_context(|| format!("Failed to read entry in: {:?}", self.local_path))?;
271 let path = entry.path();
272 if path.is_dir() {
273 fs::remove_dir_all(&path)
274 .with_context(|| format!("Failed to remove directory: {:?}", path))?;
275 } else {
276 fs::remove_file(&path)
277 .with_context(|| format!("Failed to remove file: {:?}", path))?;
278 }
279 }
280 info!("Cleared entire cache");
281 } else {
282 warn!("Cache directory does not exist");
283 }
284
285 Ok(())
286 }
287}
288
289#[derive(Debug, Clone)]
291pub struct CacheStats {
292 pub total_models: usize,
293 pub total_size: u64,
294 pub models: Vec<ModelInfo>,
295}
296
297#[derive(Debug, Clone)]
299pub struct ModelInfo {
300 pub provider: ModelProvider,
301 pub name: String,
302 pub size: u64,
303 pub path: PathBuf,
304}
305
306impl CacheStats {
307 fn format_bytes(bytes: u64) -> String {
309 const KB: u64 = 1024;
310 const MB: u64 = KB * 1024;
311 const GB: u64 = MB * 1024;
312
313 match bytes {
314 size if size >= GB => format!("{:.2} GB", size as f64 / GB as f64),
315 size if size >= MB => format!("{:.2} MB", size as f64 / MB as f64),
316 size if size >= KB => format!("{:.2} KB", size as f64 / KB as f64),
317 size => format!("{size} bytes"),
318 }
319 }
320
321 pub fn format_total_size(&self) -> String {
323 Self::format_bytes(self.total_size)
324 }
325
326 pub fn format_model_size(&self, model: &ModelInfo) -> String {
328 Self::format_bytes(model.size)
329 }
330}
331
332pub(crate) trait ProviderCache: Send + Sync {
333 fn clear_model(&self, cache_root: &Path, model_name: &str) -> Result<()>;
334 fn resolve_model_path(
335 &self,
336 cache_root: &Path,
337 model_name: &str,
338 revision: Option<&str>,
339 ) -> Result<PathBuf>;
340 fn list_models(&self, cache_root: &Path) -> Result<Vec<ModelInfo>>;
341}
342
343pub(crate) fn cache_for_provider(provider: ModelProvider) -> &'static dyn ProviderCache {
344 match provider {
345 ModelProvider::HuggingFace => &HuggingFaceProviderCache,
346 ModelProvider::Ngc => &NgcProviderCache,
347 ModelProvider::Gcs => &GcsProviderCache,
348 ModelProvider::S3 => &S3ProviderCache,
349 }
350}
351
352pub fn resolve_model_path(
353 cache_root: &Path,
354 provider: ModelProvider,
355 model_name: &str,
356 revision: Option<&str>,
357) -> Result<PathBuf> {
358 cache_for_provider(provider).resolve_model_path(cache_root, model_name, revision)
359}
360
361pub(crate) fn directory_size(path: &Path) -> Result<u64> {
362 let mut size: u64 = 0;
363
364 for entry in fs::read_dir(path)? {
365 let entry = entry?;
366 let path = entry.path();
367
368 if path.is_file() {
369 size = size.saturating_add(fs::metadata(&path)?.len());
370 } else if path.is_dir() {
371 size = size.saturating_add(directory_size(&path)?);
372 }
373 }
374
375 Ok(size)
376}
377
378fn provider_sort_key(provider: ModelProvider) -> u8 {
379 match provider {
380 ModelProvider::HuggingFace => 0,
381 ModelProvider::Ngc => 1,
382 ModelProvider::Gcs => 2,
383 ModelProvider::S3 => 3,
384 }
385}
386
387#[cfg(test)]
388#[allow(clippy::expect_used)]
389mod tests {
390 use super::*;
391 use crate::Utils;
392 use tempfile::TempDir;
393
394 #[test]
395 #[allow(clippy::expect_used)]
396 fn test_cache_config_from_path() {
397 let temp_dir = TempDir::new().expect("Failed to create temp directory");
398 let config =
399 CacheConfig::from_path(temp_dir.path()).expect("Failed to create config from path");
400
401 assert_eq!(config.local_path, temp_dir.path());
402 }
403
404 #[test]
405 #[allow(clippy::expect_used)]
406 fn test_cache_config_save_and_load() {
407 let temp_dir = TempDir::new().expect("Failed to create temp directory");
408 let original_config = CacheConfig {
409 local_path: temp_dir.path().join("cache"),
410 server_endpoint: "http://localhost:8001".to_string(),
411 timeout_secs: Some(30),
412 shared_storage: false,
413 transfer_chunk_size: 64 * 1024,
414 };
415
416 original_config
418 .save_to_config_file()
419 .expect("Failed to save config");
420
421 let loaded_config = CacheConfig::from_config_file().expect("Failed to load config");
423
424 assert_eq!(loaded_config.local_path, original_config.local_path);
425 assert_eq!(
426 loaded_config.server_endpoint,
427 original_config.server_endpoint
428 );
429 assert_eq!(loaded_config.timeout_secs, original_config.timeout_secs);
430 assert_eq!(loaded_config.shared_storage, original_config.shared_storage);
431 assert_eq!(
432 loaded_config.transfer_chunk_size,
433 original_config.transfer_chunk_size
434 );
435 }
436
437 #[test]
438 fn test_cache_stats_formatting() {
439 let stats = CacheStats {
440 total_models: 2,
441 total_size: 1024 * 1024 * 5, models: vec![
443 ModelInfo {
444 provider: ModelProvider::HuggingFace,
445 name: "model1".to_string(),
446 size: 1024 * 1024 * 2, path: PathBuf::from("/test/model1"),
448 },
449 ModelInfo {
450 provider: ModelProvider::Gcs,
451 name: "gs://bucket/model2".to_string(),
452 size: 1024 * 1024 * 3, path: PathBuf::from("/test/model2"),
454 },
455 ],
456 };
457
458 assert_eq!(stats.format_total_size(), "5.00 MB");
459 assert_eq!(stats.format_model_size(&stats.models[0]), "2.00 MB");
460 assert_eq!(stats.format_model_size(&stats.models[1]), "3.00 MB");
461 }
462
463 #[test]
464 fn test_cache_config_default() {
465 let config = CacheConfig::default();
466
467 let home = Utils::get_home_dir().unwrap_or_else(|_| ".".to_string());
468 assert_eq!(
469 config.local_path,
470 PathBuf::from(&home).join(constants::DEFAULT_CACHE_PATH)
471 );
472 assert_eq!(
473 config.server_endpoint,
474 String::from("http://localhost:8001")
475 );
476 assert_eq!(config.timeout_secs, None);
477 assert!(config.shared_storage);
478 assert_eq!(
479 config.transfer_chunk_size,
480 constants::DEFAULT_TRANSFER_CHUNK_SIZE
481 );
482 }
483
484 #[test]
485 #[allow(clippy::expect_used)]
486 fn test_cache_config_new_accepts_bare_host_port() {
487 let temp_dir = TempDir::new().expect("Failed to create temp directory");
488 let config = CacheConfig::new(
489 temp_dir.path().join("cache"),
490 Some("modelexpress-server:8001".to_string()),
491 )
492 .expect("Failed to create cache config");
493
494 assert_eq!(config.server_endpoint, "http://modelexpress-server:8001");
495 }
496
497 #[test]
498 #[allow(clippy::expect_used)]
499 fn test_get_config_path() {
500 let config_path = CacheConfig::get_config_path().expect("Failed to get config path");
501
502 let home = Utils::get_home_dir().unwrap_or_else(|_| ".".to_string());
503 assert_eq!(
504 config_path,
505 PathBuf::from(&home).join(constants::DEFAULT_CONFIG_PATH)
506 );
507 }
508
509 #[test]
510 fn test_resolve_model_path_huggingface_uses_snapshot_layout() {
511 let cache_root = Path::new("/tmp/cache");
512
513 assert_eq!(
514 resolve_model_path(
515 cache_root,
516 ModelProvider::HuggingFace,
517 "google/t5-small",
518 Some("abc123"),
519 )
520 .expect("Expected HF model path"),
521 PathBuf::from("/tmp/cache/models--google--t5-small/snapshots/abc123")
522 );
523 }
524
525 #[test]
526 fn test_resolve_model_path_gcs_uses_full_url_layout() {
527 let cache_root = Path::new("/tmp/cache");
528
529 assert_eq!(
530 resolve_model_path(
531 cache_root,
532 ModelProvider::Gcs,
533 "gs://envbucket/dev/bake/qwen/rev123",
534 None,
535 )
536 .expect("Expected GCS model path"),
537 PathBuf::from("/tmp/cache/gcs/envbucket/dev/bake/qwen/rev123")
538 );
539 }
540
541 #[test]
542 fn test_resolve_model_path_gcs_full_url_trailing_slash_normalizes() {
543 let cache_root = Path::new("/tmp/cache");
544
545 assert_eq!(
546 resolve_model_path(
547 cache_root,
548 ModelProvider::Gcs,
549 "gs://sourcebucket/dev/bake/qwen/rev123/",
550 None,
551 )
552 .expect("Expected GCS model path"),
553 PathBuf::from("/tmp/cache/gcs/sourcebucket/dev/bake/qwen/rev123")
554 );
555 }
556
557 fn create_test_cache_config(local_path: PathBuf) -> CacheConfig {
558 CacheConfig {
559 local_path,
560 server_endpoint: "http://localhost:8001".to_string(),
561 timeout_secs: None,
562 shared_storage: false,
563 transfer_chunk_size: 64 * 1024,
564 }
565 }
566
567 #[test]
568 fn test_get_cache_stats_supports_hf_and_gcs_layouts() {
569 let temp_dir = TempDir::new().expect("Failed to create temp directory");
570 let cache_path = temp_dir.path().join("cache");
571 fs::create_dir_all(&cache_path).expect("Failed to create cache directory");
572
573 let hf_model_dir = cache_path.join("models--google--t5-small");
574 fs::create_dir_all(&hf_model_dir).expect("Failed to create HF model directory");
575 fs::write(hf_model_dir.join("config.json"), b"{}").expect("Failed to write HF file");
576
577 let gcs_model_dir = resolve_model_path(
578 &cache_path,
579 ModelProvider::Gcs,
580 "gs://envbucket/dev/bake/qwen/rev123",
581 None,
582 )
583 .expect("Failed to resolve GCS path");
584 fs::create_dir_all(gcs_model_dir.join("weights"))
585 .expect("Failed to create GCS model directory");
586 fs::write(gcs_model_dir.join("tokenizer.json"), b"{}")
587 .expect("Failed to write GCS tokenizer");
588 fs::write(gcs_model_dir.join("weights/model.bin"), b"abcd")
589 .expect("Failed to write GCS weights");
590 let gcs_metadata_dir = gcs_model_dir.join(".mx");
591 fs::create_dir_all(&gcs_metadata_dir).expect("Failed to create GCS metadata directory");
592 fs::write(
593 gcs_metadata_dir.join("manifest.json"),
594 r#"{"version":1,"model":"gs://envbucket/dev/bake/qwen/rev123","files":[{"path":"tokenizer.json","size":2,"crc32c":"00000000","generation":null},{"path":"weights/model.bin","size":4,"crc32c":"00000000","generation":null}]}
595"#,
596 )
597 .expect("Failed to write GCS manifest");
598
599 let ignored_dir = cache_path.join("tmp");
600 fs::create_dir_all(&ignored_dir).expect("Failed to create ignored directory");
601 fs::write(ignored_dir.join("scratch.txt"), b"ignore")
602 .expect("Failed to write ignored file");
603
604 let stats = create_test_cache_config(cache_path)
605 .get_cache_stats()
606 .expect("Failed to get cache stats");
607
608 assert_eq!(stats.total_models, 2);
609 assert_eq!(stats.total_size, 8);
610 assert_eq!(stats.models.len(), 2);
611
612 assert_eq!(stats.models[0].provider, ModelProvider::HuggingFace);
613 assert_eq!(stats.models[0].name, "google/t5-small");
614 assert_eq!(stats.models[0].size, 2);
615 assert_eq!(stats.models[0].path, hf_model_dir);
616 assert_eq!(stats.models[1].provider, ModelProvider::Gcs);
617 assert_eq!(stats.models[1].name, "gs://envbucket/dev/bake/qwen/rev123");
618 assert_eq!(stats.models[1].size, 6);
619 assert_eq!(stats.models[1].path, gcs_model_dir);
620 assert!(stats.models.iter().all(|model| model.name != "tmp"));
621 }
622
623 #[test]
624 fn test_clear_model_removes_only_requested_layout() {
625 let temp_dir = TempDir::new().expect("Failed to create temp directory");
626 let cache_path = temp_dir.path().join("cache");
627 fs::create_dir_all(&cache_path).expect("Failed to create cache directory");
628
629 let hf_model_dir = cache_path.join("models--google--t5-small");
630 fs::create_dir_all(&hf_model_dir).expect("Failed to create HF model directory");
631 fs::write(hf_model_dir.join("config.json"), b"{}").expect("Failed to write HF file");
632
633 let gcs_model_dir = resolve_model_path(
634 &cache_path,
635 ModelProvider::Gcs,
636 "gs://envbucket/org/model/rev1",
637 None,
638 )
639 .expect("Failed to resolve GCS path");
640 fs::create_dir_all(&gcs_model_dir).expect("Failed to create GCS model directory");
641 fs::write(gcs_model_dir.join("tokenizer.json"), b"{}").expect("Failed to write GCS file");
642
643 let config = create_test_cache_config(cache_path);
644
645 config
646 .clear_model("gs://envbucket/org/model/rev1", ModelProvider::Gcs)
647 .expect("Failed to clear GCS model");
648 assert!(hf_model_dir.exists(), "HF model should remain");
649 assert!(!gcs_model_dir.exists(), "GCS model should be removed");
650
651 config
652 .clear_model("google/t5-small", ModelProvider::HuggingFace)
653 .expect("Failed to clear HF model");
654 assert!(!hf_model_dir.exists(), "HF model should be removed");
655 }
656
657 #[test]
658 fn test_clear_all_removes_contents_but_keeps_directory() {
659 let temp_dir = TempDir::new().expect("Failed to create temp directory");
660 let cache_path = temp_dir.path().join("cache");
661 fs::create_dir_all(&cache_path).expect("Failed to create cache directory");
662
663 let model_dir = cache_path.join("models--test--model");
665 fs::create_dir_all(&model_dir).expect("Failed to create model directory");
666 fs::write(model_dir.join("config.json"), "{}").expect("Failed to write file");
667 fs::write(cache_path.join("test_file.txt"), "test").expect("Failed to write file");
668
669 let config = create_test_cache_config(cache_path.clone());
670
671 config.clear_all().expect("Failed to clear cache");
673
674 assert!(cache_path.exists(), "Cache directory should still exist");
676 assert!(
677 fs::read_dir(&cache_path)
678 .expect("Failed to read dir")
679 .next()
680 .is_none(),
681 "Cache directory should be empty"
682 );
683 }
684
685 #[test]
686 fn test_clear_all_handles_nonexistent_directory() {
687 let temp_dir = TempDir::new().expect("Failed to create temp directory");
688 let cache_path = temp_dir.path().join("nonexistent_cache");
689
690 let config = create_test_cache_config(cache_path.clone());
691
692 config
694 .clear_all()
695 .with_context(|| format!("Failed to clear cache: {cache_path:?}"))
696 .expect("Failed to clear cache");
697 assert!(!cache_path.exists());
698 }
699
700 #[test]
701 fn test_clear_all_removes_nested_directories() {
702 let temp_dir = TempDir::new().expect("Failed to create temp directory");
703 let cache_path = temp_dir.path().join("cache");
704 fs::create_dir_all(&cache_path).expect("Failed to create cache directory");
705
706 let deep_path = cache_path.join("a").join("b").join("c");
708 fs::create_dir_all(&deep_path).expect("Failed to create nested directories");
709 fs::write(deep_path.join("deep_file.txt"), "deep").expect("Failed to write file");
710
711 let config = create_test_cache_config(cache_path.clone());
712
713 config.clear_all().expect("Failed to clear cache");
714
715 assert!(cache_path.exists(), "Cache directory should still exist");
716 assert!(
717 fs::read_dir(&cache_path)
718 .expect("Failed to read dir")
719 .next()
720 .is_none(),
721 "Cache directory should be empty after clearing nested content"
722 );
723 }
724}