From 3b0b8ade1c5b3ffb391c586f436c3e8b33055b42 Mon Sep 17 00:00:00 2001 From: kerthcet Date: Sun, 26 Apr 2026 17:12:50 +0100 Subject: [PATCH 1/5] optimize the modelInfo Signed-off-by: kerthcet --- src/api/tests.rs | 6 +++--- src/cli/commands.rs | 20 ++++++++++---------- src/cli/inspect.rs | 12 ++++++------ src/cli/ls.rs | 6 +++--- src/cli/rm.rs | 8 ++++---- src/downloader/huggingface.rs | 8 ++++---- src/registry/model_registry.rs | 24 ++++++++++++------------ src/storage/sqlite.rs | 14 +++++++------- 8 files changed, 49 insertions(+), 49 deletions(-) diff --git a/src/api/tests.rs b/src/api/tests.rs index f5a44f2..512e55a 100644 --- a/src/api/tests.rs +++ b/src/api/tests.rs @@ -14,7 +14,7 @@ use tower::util::ServiceExt; // for `oneshot` and `ready` use super::routes::create_router; use crate::backend::mock::MockEngine; -use crate::registry::model_registry::{ArtifactInfo, ModelInfo, ModelMetadata, ModelRegistry}; +use crate::registry::model_registry::{CacheInfo, ModelInfo, ModelMetadata, ModelRegistry}; /// Helper to create test app with a pre-registered test model /// Returns the router and the temp directory (which must be kept alive) @@ -27,15 +27,15 @@ fn create_test_app() -> (axum::Router, TempDir) { let test_model = ModelInfo { uuid: "test-uuid".to_string(), name: "test-model".to_string(), + provider: "test".to_string(), author: Some("test-author".to_string()), task: Some("text-generation".to_string()), model_series: Some("test-series".to_string()), - provider: "test".to_string(), license: Some("MIT".to_string()), created_at: chrono::Utc::now().to_rfc3339(), updated_at: chrono::Utc::now().to_rfc3339(), metadata: ModelMetadata { - artifact: ArtifactInfo { + cache: CacheInfo { revision: "test-rev".to_string(), size: 1000, path: "/tmp/test-model".to_string(), diff --git a/src/cli/commands.rs b/src/cli/commands.rs index ae37255..641e95d 100644 --- a/src/cli/commands.rs +++ b/src/cli/commands.rs @@ -141,12 +141,12 @@ pub async fn run(cli: Cli) { "MODEL", "TASK", "PROVIDER", "REVISION", "SIZE", "CREATED" ]); for model in models { - let size_str = format_size_decimal(model.metadata.artifact.size); + let size_str = format_size_decimal(model.metadata.cache.size); - let revision_short = if model.metadata.artifact.revision.len() > 8 { - &model.metadata.artifact.revision[..8] + let revision_short = if model.metadata.cache.revision.len() > 8 { + &model.metadata.cache.revision[..8] } else { - &model.metadata.artifact.revision + &model.metadata.cache.revision }; let created_str = format_time_ago(&model.created_at); @@ -230,7 +230,7 @@ pub async fn run(cli: Cli) { #[cfg(test)] mod tests { use super::*; - use crate::registry::model_registry::{ArtifactInfo, ModelInfo, ModelMetadata}; + use crate::registry::model_registry::{CacheInfo, ModelInfo, ModelMetadata}; use tempfile::TempDir; // Helper to create a test model @@ -245,15 +245,15 @@ mod tests { ModelInfo { uuid: revision.to_string(), name: name.to_string(), + provider: "huggingface".to_string(), author: Some("test-author".to_string()), task: Some("text-generation".to_string()), model_series: Some("gpt2".to_string()), - provider: "huggingface".to_string(), license: Some("mit".to_string()), created_at: "2025-01-01T00:00:00Z".to_string(), updated_at: "2025-01-01T00:00:00Z".to_string(), metadata: ModelMetadata { - artifact: ArtifactInfo { + cache: CacheInfo { revision: revision.to_string(), size: 1000, path: "/tmp/test".to_string(), @@ -375,7 +375,7 @@ mod tests { // Update the model let mut updated_model = create_test_model("test/updated-model", "v2"); - updated_model.metadata.artifact.size = 2000; + updated_model.metadata.cache.size = 2000; updated_model.created_at = "2025-01-05T00:00:00Z".to_string(); updated_model.updated_at = "2025-01-05T00:00:00Z".to_string(); @@ -387,7 +387,7 @@ mod tests { // updated_at should be new assert_eq!(result.updated_at, "2025-01-05T00:00:00Z"); // Other fields should be updated - assert_eq!(result.metadata.artifact.revision, "v2"); - assert_eq!(result.metadata.artifact.size, 2000); + assert_eq!(result.metadata.cache.revision, "v2"); + assert_eq!(result.metadata.cache.size, 2000); } } diff --git a/src/cli/inspect.rs b/src/cli/inspect.rs index e8253a6..2610434 100644 --- a/src/cli/inspect.rs +++ b/src/cli/inspect.rs @@ -68,12 +68,12 @@ pub fn display(model: &ModelInfo) { // Artifact section println!(" artifact:"); println!(" provider: {}", model.provider); - println!(" revision: {}", model.metadata.artifact.revision); + println!(" revision: {}", model.metadata.cache.revision); println!( " size: {}", - format_size_decimal(model.metadata.artifact.size) + format_size_decimal(model.metadata.cache.size) ); - println!(" cache_path: {}", model.metadata.artifact.path); + println!(" cache_path: {}", model.metadata.cache.path); println!("status:"); println!(" created: {}", format_time_ago(&model.created_at)); println!(" updated: {}", format_time_ago(&model.updated_at)); @@ -82,7 +82,7 @@ pub fn display(model: &ModelInfo) { #[cfg(test)] mod tests { use super::*; - use crate::registry::model_registry::{ArtifactInfo, ModelInfo, ModelMetadata}; + use crate::registry::model_registry::{CacheInfo, ModelInfo, ModelMetadata}; use tempfile::TempDir; fn create_test_model(name: &str, uuid: &str) -> ModelInfo { @@ -94,15 +94,15 @@ mod tests { ModelInfo { uuid: uuid.to_string(), name: name.to_string(), + provider: "huggingface".to_string(), author: Some("test-author".to_string()), task: Some("text-generation".to_string()), model_series: Some("gpt2".to_string()), - provider: "huggingface".to_string(), license: Some("mit".to_string()), created_at: "2025-01-01T00:00:00Z".to_string(), updated_at: "2025-01-01T00:00:00Z".to_string(), metadata: ModelMetadata { - artifact: ArtifactInfo { + cache: CacheInfo { revision: uuid.to_string(), size: 1000, path: "/tmp/test".to_string(), diff --git a/src/cli/ls.rs b/src/cli/ls.rs index 28a19f1..21f32e3 100644 --- a/src/cli/ls.rs +++ b/src/cli/ls.rs @@ -52,7 +52,7 @@ pub fn execute( #[cfg(test)] mod tests { use super::*; - use crate::registry::model_registry::{ArtifactInfo, ModelInfo, ModelMetadata}; + use crate::registry::model_registry::{CacheInfo, ModelInfo, ModelMetadata}; use tempfile::TempDir; fn create_test_model(name: &str, uuid: &str, author: &str) -> ModelInfo { @@ -64,15 +64,15 @@ mod tests { ModelInfo { uuid: uuid.to_string(), name: name.to_string(), + provider: "huggingface".to_string(), author: Some(author.to_string()), task: Some("text-generation".to_string()), model_series: Some("gpt2".to_string()), - provider: "huggingface".to_string(), license: Some("mit".to_string()), created_at: "2025-01-01T00:00:00Z".to_string(), updated_at: "2025-01-01T00:00:00Z".to_string(), metadata: ModelMetadata { - artifact: ArtifactInfo { + cache: CacheInfo { revision: uuid.to_string(), size: 1000, path: "/tmp/test".to_string(), diff --git a/src/cli/rm.rs b/src/cli/rm.rs index 818f490..e8800f1 100644 --- a/src/cli/rm.rs +++ b/src/cli/rm.rs @@ -14,7 +14,7 @@ pub fn execute(registry: &ModelRegistry, model_name: &str) -> Result<(), String> #[cfg(test)] mod tests { use super::*; - use crate::registry::model_registry::{ArtifactInfo, ModelInfo, ModelMetadata}; + use crate::registry::model_registry::{CacheInfo, ModelInfo, ModelMetadata}; use tempfile::TempDir; fn create_test_model(name: &str, uuid: &str) -> ModelInfo { @@ -26,15 +26,15 @@ mod tests { ModelInfo { uuid: uuid.to_string(), name: name.to_string(), + provider: "huggingface".to_string(), author: Some("test-author".to_string()), task: Some("text-generation".to_string()), model_series: Some("gpt2".to_string()), - provider: "huggingface".to_string(), license: Some("mit".to_string()), created_at: "2025-01-01T00:00:00Z".to_string(), updated_at: "2025-01-01T00:00:00Z".to_string(), metadata: ModelMetadata { - artifact: ArtifactInfo { + cache: CacheInfo { revision: uuid.to_string(), size: 1000, path: "/tmp/test".to_string(), @@ -55,7 +55,7 @@ mod tests { std::fs::write(cache_dir.join("model.safetensors"), "fake data").unwrap(); let mut model = create_test_model("test/remove-model", "abc123"); - model.metadata.artifact.path = cache_dir.to_string_lossy().to_string(); + model.metadata.cache.path = cache_dir.to_string_lossy().to_string(); registry.register_model(model).unwrap(); assert!(registry.get_model("test/remove-model").unwrap().is_some()); diff --git a/src/downloader/huggingface.rs b/src/downloader/huggingface.rs index 8d7ec73..0a789f1 100644 --- a/src/downloader/huggingface.rs +++ b/src/downloader/huggingface.rs @@ -6,7 +6,7 @@ use indicatif::{ProgressBar, ProgressStyle}; use crate::downloader::downloader::{DownloadError, Downloader}; use crate::downloader::progress::{DownloadProgressManager, FileProgress}; -use crate::registry::model_registry::{ArtifactInfo, ModelInfo, ModelMetadata, ModelRegistry}; +use crate::registry::model_registry::{CacheInfo, ModelInfo, ModelMetadata, ModelRegistry}; use crate::utils::file::{self, format_model_name}; /// Adapter to bridge HuggingFace's Progress trait with our FileProgress @@ -284,14 +284,14 @@ impl Downloader for HuggingFaceDownloader { let model_size = storage_from_api.unwrap_or_else(|| progress_manager.total_downloaded_bytes()); - let artifact = ArtifactInfo { + let cache = CacheInfo { revision: sha.clone(), size: model_size, path: model_cache_path.to_string_lossy().to_string(), }; let metadata = ModelMetadata { - artifact, + cache, context_window, safetensors: safetensors_from_api, }; @@ -300,10 +300,10 @@ impl Downloader for HuggingFaceDownloader { let model_info_record = ModelInfo { uuid: sha, // Use revision SHA as UUID for now name: name.to_string(), + provider: "huggingface".to_string(), author: author_from_api, task: task_from_api, model_series: model_series_from_api, - provider: "huggingface".to_string(), license: license_from_api, created_at: now.clone(), updated_at: now, diff --git a/src/registry/model_registry.rs b/src/registry/model_registry.rs index 773fe68..f0c2727 100644 --- a/src/registry/model_registry.rs +++ b/src/registry/model_registry.rs @@ -8,7 +8,7 @@ use crate::storage::{ModelStorage, SqliteStorage}; use crate::utils::file; #[derive(Debug, Serialize, Deserialize, Clone)] -pub struct ArtifactInfo { +pub struct CacheInfo { pub revision: String, pub size: u64, pub path: String, @@ -16,7 +16,7 @@ pub struct ArtifactInfo { #[derive(Debug, Serialize, Deserialize, Clone)] pub struct ModelMetadata { - pub artifact: ArtifactInfo, + pub cache: CacheInfo, #[serde(skip_serializing_if = "Option::is_none")] pub context_window: Option, #[serde(skip_serializing_if = "Option::is_none")] @@ -27,12 +27,12 @@ pub struct ModelMetadata { pub struct ModelInfo { pub uuid: String, pub name: String, + pub provider: String, pub author: Option, #[serde(skip_serializing_if = "Option::is_none")] pub task: Option, // Task type (image-text-to-text, text-generation) #[serde(skip_serializing_if = "Option::is_none")] pub model_series: Option, // Architecture series (qwen3_5, gpt2, llama3) - pub provider: String, #[serde(skip_serializing_if = "Option::is_none")] pub license: Option, pub metadata: ModelMetadata, @@ -82,7 +82,7 @@ impl ModelRegistry { if let Some(info) = model_info { // Delete artifact directory if it exists - let artifact_path = std::path::Path::new(&info.metadata.artifact.path); + let artifact_path = std::path::Path::new(&info.metadata.cache.path); if artifact_path.exists() { fs::remove_dir_all(artifact_path)?; } @@ -119,15 +119,15 @@ mod tests { ModelInfo { uuid: revision.to_string(), name: name.to_string(), + provider: "huggingface".to_string(), author: Some("test-author".to_string()), task: Some("text-generation".to_string()), model_series: Some("gpt2".to_string()), - provider: "huggingface".to_string(), license: Some("mit".to_string()), created_at: "2025-01-01T00:00:00Z".to_string(), updated_at: "2025-01-01T00:00:00Z".to_string(), metadata: ModelMetadata { - artifact: ArtifactInfo { + cache: CacheInfo { revision: revision.to_string(), size: 1000, path: "/tmp/test".to_string(), @@ -202,8 +202,8 @@ mod tests { registry.register_model(model1).unwrap(); let mut model2 = create_test_model("test/model", "def456"); - model2.metadata.artifact.size = 2000; - model2.metadata.artifact.path = "/tmp/test2".to_string(); + model2.metadata.cache.size = 2000; + model2.metadata.cache.path = "/tmp/test2".to_string(); model2.created_at = "2025-01-02T00:00:00Z".to_string(); model2.updated_at = "2025-01-02T00:00:00Z".to_string(); @@ -211,8 +211,8 @@ mod tests { let models = registry.load_models(None).unwrap(); assert_eq!(models.len(), 1); - assert_eq!(models[0].metadata.artifact.revision, "def456"); - assert_eq!(models[0].metadata.artifact.size, 2000); + assert_eq!(models[0].metadata.cache.revision, "def456"); + assert_eq!(models[0].metadata.cache.size, 2000); // created_at should be preserved from model1 assert_eq!(models[0].created_at, "2025-01-01T00:00:00Z"); // updated_at should be from model2 @@ -230,7 +230,7 @@ mod tests { fs::write(cache_dir.join("test.txt"), "test data").unwrap(); let mut model = create_test_model("test/model", "abc123"); - model.metadata.artifact.path = cache_dir.to_string_lossy().to_string(); + model.metadata.cache.path = cache_dir.to_string_lossy().to_string(); registry.register_model(model).unwrap(); assert_eq!(registry.load_models(None).unwrap().len(), 1); @@ -271,7 +271,7 @@ mod tests { let model_info = retrieved.unwrap(); assert_eq!(model_info.name, "test/gpt-model"); assert_eq!(model_info.provider, "huggingface"); - assert_eq!(model_info.metadata.artifact.revision, "abc123def456"); + assert_eq!(model_info.metadata.cache.revision, "abc123def456"); assert_eq!(model_info.model_series, Some("gpt2".to_string())); assert_eq!(model_info.metadata.context_window, Some(2048)); assert_eq!( diff --git a/src/storage/sqlite.rs b/src/storage/sqlite.rs index 23a6923..49a8792 100644 --- a/src/storage/sqlite.rs +++ b/src/storage/sqlite.rs @@ -112,10 +112,10 @@ impl ModelStorage for SqliteStorage { Ok(ModelInfo { uuid: row.get(0)?, name: row.get(1)?, + provider: row.get(5)?, author: row.get(2)?, task: row.get(3)?, model_series: row.get(4)?, - provider: row.get(5)?, license: row.get(6)?, metadata, created_at: row.get(8)?, @@ -202,10 +202,10 @@ impl ModelStorage for SqliteStorage { Ok(ModelInfo { uuid: row.get(0)?, name: row.get(1)?, + provider: row.get(5)?, author: row.get(2)?, task: row.get(3)?, model_series: row.get(4)?, - provider: row.get(5)?, license: row.get(6)?, created_at: row.get(8)?, updated_at: row.get(9)?, @@ -225,7 +225,7 @@ impl ModelStorage for SqliteStorage { #[cfg(test)] mod tests { use super::*; - use crate::registry::model_registry::{ArtifactInfo, ModelMetadata}; + use crate::registry::model_registry::{CacheInfo, ModelMetadata}; use tempfile::TempDir; fn create_test_model(name: &str, uuid: &str) -> ModelInfo { @@ -237,15 +237,15 @@ mod tests { ModelInfo { uuid: uuid.to_string(), name: name.to_string(), + provider: "huggingface".to_string(), author: Some("test-author".to_string()), task: Some("text-generation".to_string()), model_series: Some("gpt2".to_string()), - provider: "huggingface".to_string(), license: Some("mit".to_string()), created_at: "2025-01-01T00:00:00Z".to_string(), updated_at: "2025-01-01T00:00:00Z".to_string(), metadata: ModelMetadata { - artifact: ArtifactInfo { + cache: CacheInfo { revision: "abc123".to_string(), size: 1000, path: "/tmp/test".to_string(), @@ -336,8 +336,8 @@ mod tests { storage.register_model(model).unwrap(); let retrieved = storage.get_model("test/model").unwrap().unwrap(); - assert_eq!(retrieved.metadata.artifact.revision, "abc123"); - assert_eq!(retrieved.metadata.artifact.size, 1000); + assert_eq!(retrieved.metadata.cache.revision, "abc123"); + assert_eq!(retrieved.metadata.cache.size, 1000); assert_eq!(retrieved.metadata.context_window, Some(2048)); assert!(retrieved.metadata.safetensors.is_some()); From f6da3a5c71c87488c8554852389721c2f664c7de Mon Sep 17 00:00:00 2001 From: kerthcet Date: Sun, 26 Apr 2026 17:19:53 +0100 Subject: [PATCH 2/5] optimize the layout Signed-off-by: kerthcet --- src/cli/inspect.rs | 16 ++++++++-------- src/main.rs | 2 +- 2 files changed, 9 insertions(+), 9 deletions(-) diff --git a/src/cli/inspect.rs b/src/cli/inspect.rs index 2610434..2eb5064 100644 --- a/src/cli/inspect.rs +++ b/src/cli/inspect.rs @@ -19,6 +19,11 @@ pub fn display(model: &ModelInfo) { " author: {}", model.author.as_deref().unwrap_or("N/A") ); + println!(" provider: {}", model.provider); + println!( + " model_series: {}", + model.model_series.as_deref().unwrap_or("N/A") + ); println!( " task: {}", model.task.as_deref().unwrap_or("N/A") @@ -31,10 +36,6 @@ pub fn display(model: &ModelInfo) { .map(|s| s.to_uppercase()) .unwrap_or_else(|| "N/A".to_string()) ); - println!( - " model_series: {}", - model.model_series.as_deref().unwrap_or("N/A") - ); println!( " context_window: {}", model @@ -65,15 +66,14 @@ pub fn display(model: &ModelInfo) { println!(" safetensors: N/A"); } - // Artifact section - println!(" artifact:"); - println!(" provider: {}", model.provider); + // Cache section + println!(" cache:"); println!(" revision: {}", model.metadata.cache.revision); println!( " size: {}", format_size_decimal(model.metadata.cache.size) ); - println!(" cache_path: {}", model.metadata.cache.path); + println!(" path: {}", model.metadata.cache.path); println!("status:"); println!(" created: {}", format_time_ago(&model.created_at)); println!(" updated: {}", format_time_ago(&model.updated_at)); diff --git a/src/main.rs b/src/main.rs index 614a698..2359efb 100644 --- a/src/main.rs +++ b/src/main.rs @@ -18,7 +18,7 @@ fn main() { tracing_subscriber::fmt() .with_env_filter( tracing_subscriber::EnvFilter::try_from_default_env() - .unwrap_or_else(|_| "info,hf_hub=warn,tower_http=info".into()), + .unwrap_or_else(|_| "info,hf_hub=warn,tower_http=info,rusqlite_migration=warn".into()), ) .init(); From 371922b5a5d5584eefb3223d12b61df8c2a7cc99 Mon Sep 17 00:00:00 2001 From: kerthcet Date: Sun, 26 Apr 2026 17:25:40 +0100 Subject: [PATCH 3/5] add alias to providers Signed-off-by: kerthcet --- src/cli/commands.rs | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/cli/commands.rs b/src/cli/commands.rs index 641e95d..ea0a87c 100644 --- a/src/cli/commands.rs +++ b/src/cli/commands.rs @@ -91,7 +91,9 @@ struct InspectArgs { #[derive(Debug, Clone, Default, clap::ValueEnum)] pub enum Provider { #[default] + #[value(alias = "hf")] Huggingface, + #[value(alias = "ms")] Modelscope, } From 5ed311d1b803303e61565c250c62543624b5677b Mon Sep 17 00:00:00 2001 From: kerthcet Date: Sun, 26 Apr 2026 17:27:09 +0100 Subject: [PATCH 4/5] polish comments Signed-off-by: kerthcet --- README.md | 4 ++-- src/registry/model_registry.rs | 8 ++++---- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/README.md b/README.md index 4c05803..2a6c989 100644 --- a/README.md +++ b/README.md @@ -228,6 +228,7 @@ name: inftyai/tiny-random-gpt2 kind: Model spec: author: inftyai + provider: huggingface task: text-generation license: MIT model_series: gpt2 @@ -236,8 +237,7 @@ spec: total: 7.00B parameters: f32: 7.00B - artifact: - provider: huggingface + cache: revision: abc123de size: 1.24 GB cache_path: ~/.puma/cache/... diff --git a/src/registry/model_registry.rs b/src/registry/model_registry.rs index f0c2727..26ebcc7 100644 --- a/src/registry/model_registry.rs +++ b/src/registry/model_registry.rs @@ -81,10 +81,10 @@ impl ModelRegistry { let model_info = self.get_model(name)?; if let Some(info) = model_info { - // Delete artifact directory if it exists - let artifact_path = std::path::Path::new(&info.metadata.cache.path); - if artifact_path.exists() { - fs::remove_dir_all(artifact_path)?; + // Delete cache directory if it exists + let cache_path = std::path::Path::new(&info.metadata.cache.path); + if cache_path.exists() { + fs::remove_dir_all(cache_path)?; } // Remove from registry From bdb73c4d5cf50d1f6213dd0b22d8010b5290c8cd Mon Sep 17 00:00:00 2001 From: kerthcet Date: Sun, 26 Apr 2026 17:28:22 +0100 Subject: [PATCH 5/5] polish comments Signed-off-by: kerthcet --- src/main.rs | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/main.rs b/src/main.rs index 2359efb..02929ea 100644 --- a/src/main.rs +++ b/src/main.rs @@ -17,8 +17,9 @@ fn main() { // Setup tracing subscriber for tower-http TraceLayer tracing_subscriber::fmt() .with_env_filter( - tracing_subscriber::EnvFilter::try_from_default_env() - .unwrap_or_else(|_| "info,hf_hub=warn,tower_http=info,rusqlite_migration=warn".into()), + tracing_subscriber::EnvFilter::try_from_default_env().unwrap_or_else(|_| { + "info,hf_hub=warn,tower_http=info,rusqlite_migration=warn".into() + }), ) .init();