Skip to content

Commit 89cbd85

Browse files
kerthcetkerthcet
authored
optimize the modelInfo (#45)
* optimize the modelInfo Signed-off-by: kerthcet <kerthcet@gmail.com> * optimize the layout Signed-off-by: kerthcet <kerthcet@gmail.com> * add alias to providers Signed-off-by: kerthcet <kerthcet@gmail.com> * polish comments Signed-off-by: kerthcet <kerthcet@gmail.com> * polish comments Signed-off-by: kerthcet <kerthcet@gmail.com> --------- Signed-off-by: kerthcet <kerthcet@gmail.com>
1 parent ff7a96d commit 89cbd85

10 files changed

Lines changed: 66 additions & 63 deletions

File tree

‎README.md‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -228,6 +228,7 @@ name: inftyai/tiny-random-gpt2
228228
kind: Model
229229
spec:
230230
author: inftyai
231+
provider: huggingface
231232
task: text-generation
232233
license: MIT
233234
model_series: gpt2
@@ -236,8 +237,7 @@ spec:
236237
total: 7.00B
237238
parameters:
238239
f32: 7.00B
239-
artifact:
240-
provider: huggingface
240+
cache:
241241
revision: abc123de
242242
size: 1.24 GB
243243
cache_path: ~/.puma/cache/...

‎src/api/tests.rs‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@ use tower::util::ServiceExt; // for `oneshot` and `ready`
1414

1515
use super::routes::create_router;
1616
use crate::backend::mock::MockEngine;
17-
use crate::registry::model_registry::{ArtifactInfo, ModelInfo, ModelMetadata, ModelRegistry};
17+
use crate::registry::model_registry::{CacheInfo, ModelInfo, ModelMetadata, ModelRegistry};
1818

1919
/// Helper to create test app with a pre-registered test model
2020
/// Returns the router and the temp directory (which must be kept alive)
@@ -27,15 +27,15 @@ fn create_test_app() -> (axum::Router, TempDir) {
2727
let test_model = ModelInfo {
2828
uuid: "test-uuid".to_string(),
2929
name: "test-model".to_string(),
30+
provider: "test".to_string(),
3031
author: Some("test-author".to_string()),
3132
task: Some("text-generation".to_string()),
3233
model_series: Some("test-series".to_string()),
33-
provider: "test".to_string(),
3434
license: Some("MIT".to_string()),
3535
created_at: chrono::Utc::now().to_rfc3339(),
3636
updated_at: chrono::Utc::now().to_rfc3339(),
3737
metadata: ModelMetadata {
38-
artifact: ArtifactInfo {
38+
cache: CacheInfo {
3939
revision: "test-rev".to_string(),
4040
size: 1000,
4141
path: "/tmp/test-model".to_string(),

‎src/cli/commands.rs‎

Lines changed: 12 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -91,7 +91,9 @@ struct InspectArgs {
9191
#[derive(Debug, Clone, Default, clap::ValueEnum)]
9292
pub enum Provider {
9393
#[default]
94+
#[value(alias = "hf")]
9495
Huggingface,
96+
#[value(alias = "ms")]
9597
Modelscope,
9698
}
9799

@@ -141,12 +143,12 @@ pub async fn run(cli: Cli) {
141143
"MODEL", "TASK", "PROVIDER", "REVISION", "SIZE", "CREATED"
142144
]);
143145
for model in models {
144-
let size_str = format_size_decimal(model.metadata.artifact.size);
146+
let size_str = format_size_decimal(model.metadata.cache.size);
145147

146-
let revision_short = if model.metadata.artifact.revision.len() > 8 {
147-
&model.metadata.artifact.revision[..8]
148+
let revision_short = if model.metadata.cache.revision.len() > 8 {
149+
&model.metadata.cache.revision[..8]
148150
} else {
149-
&model.metadata.artifact.revision
151+
&model.metadata.cache.revision
150152
};
151153

152154
let created_str = format_time_ago(&model.created_at);
@@ -230,7 +232,7 @@ pub async fn run(cli: Cli) {
230232
#[cfg(test)]
231233
mod tests {
232234
use super::*;
233-
use crate::registry::model_registry::{ArtifactInfo, ModelInfo, ModelMetadata};
235+
use crate::registry::model_registry::{CacheInfo, ModelInfo, ModelMetadata};
234236
use tempfile::TempDir;
235237

236238
// Helper to create a test model
@@ -245,15 +247,15 @@ mod tests {
245247
ModelInfo {
246248
uuid: revision.to_string(),
247249
name: name.to_string(),
250+
provider: "huggingface".to_string(),
248251
author: Some("test-author".to_string()),
249252
task: Some("text-generation".to_string()),
250253
model_series: Some("gpt2".to_string()),
251-
provider: "huggingface".to_string(),
252254
license: Some("mit".to_string()),
253255
created_at: "2025-01-01T00:00:00Z".to_string(),
254256
updated_at: "2025-01-01T00:00:00Z".to_string(),
255257
metadata: ModelMetadata {
256-
artifact: ArtifactInfo {
258+
cache: CacheInfo {
257259
revision: revision.to_string(),
258260
size: 1000,
259261
path: "/tmp/test".to_string(),
@@ -375,7 +377,7 @@ mod tests {
375377

376378
// Update the model
377379
let mut updated_model = create_test_model("test/updated-model", "v2");
378-
updated_model.metadata.artifact.size = 2000;
380+
updated_model.metadata.cache.size = 2000;
379381
updated_model.created_at = "2025-01-05T00:00:00Z".to_string();
380382
updated_model.updated_at = "2025-01-05T00:00:00Z".to_string();
381383

@@ -387,7 +389,7 @@ mod tests {
387389
// updated_at should be new
388390
assert_eq!(result.updated_at, "2025-01-05T00:00:00Z");
389391
// Other fields should be updated
390-
assert_eq!(result.metadata.artifact.revision, "v2");
391-
assert_eq!(result.metadata.artifact.size, 2000);
392+
assert_eq!(result.metadata.cache.revision, "v2");
393+
assert_eq!(result.metadata.cache.size, 2000);
392394
}
393395
}

‎src/cli/inspect.rs‎

Lines changed: 13 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,11 @@ pub fn display(model: &ModelInfo) {
1919
" author: {}",
2020
model.author.as_deref().unwrap_or("N/A")
2121
);
22+
println!(" provider: {}", model.provider);
23+
println!(
24+
" model_series: {}",
25+
model.model_series.as_deref().unwrap_or("N/A")
26+
);
2227
println!(
2328
" task: {}",
2429
model.task.as_deref().unwrap_or("N/A")
@@ -31,10 +36,6 @@ pub fn display(model: &ModelInfo) {
3136
.map(|s| s.to_uppercase())
3237
.unwrap_or_else(|| "N/A".to_string())
3338
);
34-
println!(
35-
" model_series: {}",
36-
model.model_series.as_deref().unwrap_or("N/A")
37-
);
3839
println!(
3940
" context_window: {}",
4041
model
@@ -65,15 +66,14 @@ pub fn display(model: &ModelInfo) {
6566
println!(" safetensors: N/A");
6667
}
6768

68-
// Artifact section
69-
println!(" artifact:");
70-
println!(" provider: {}", model.provider);
71-
println!(" revision: {}", model.metadata.artifact.revision);
69+
// Cache section
70+
println!(" cache:");
71+
println!(" revision: {}", model.metadata.cache.revision);
7272
println!(
7373
" size: {}",
74-
format_size_decimal(model.metadata.artifact.size)
74+
format_size_decimal(model.metadata.cache.size)
7575
);
76-
println!(" cache_path: {}", model.metadata.artifact.path);
76+
println!(" path: {}", model.metadata.cache.path);
7777
println!("status:");
7878
println!(" created: {}", format_time_ago(&model.created_at));
7979
println!(" updated: {}", format_time_ago(&model.updated_at));
@@ -82,7 +82,7 @@ pub fn display(model: &ModelInfo) {
8282
#[cfg(test)]
8383
mod tests {
8484
use super::*;
85-
use crate::registry::model_registry::{ArtifactInfo, ModelInfo, ModelMetadata};
85+
use crate::registry::model_registry::{CacheInfo, ModelInfo, ModelMetadata};
8686
use tempfile::TempDir;
8787

8888
fn create_test_model(name: &str, uuid: &str) -> ModelInfo {
@@ -94,15 +94,15 @@ mod tests {
9494
ModelInfo {
9595
uuid: uuid.to_string(),
9696
name: name.to_string(),
97+
provider: "huggingface".to_string(),
9798
author: Some("test-author".to_string()),
9899
task: Some("text-generation".to_string()),
99100
model_series: Some("gpt2".to_string()),
100-
provider: "huggingface".to_string(),
101101
license: Some("mit".to_string()),
102102
created_at: "2025-01-01T00:00:00Z".to_string(),
103103
updated_at: "2025-01-01T00:00:00Z".to_string(),
104104
metadata: ModelMetadata {
105-
artifact: ArtifactInfo {
105+
cache: CacheInfo {
106106
revision: uuid.to_string(),
107107
size: 1000,
108108
path: "/tmp/test".to_string(),

‎src/cli/ls.rs‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -52,7 +52,7 @@ pub fn execute(
5252
#[cfg(test)]
5353
mod tests {
5454
use super::*;
55-
use crate::registry::model_registry::{ArtifactInfo, ModelInfo, ModelMetadata};
55+
use crate::registry::model_registry::{CacheInfo, ModelInfo, ModelMetadata};
5656
use tempfile::TempDir;
5757

5858
fn create_test_model(name: &str, uuid: &str, author: &str) -> ModelInfo {
@@ -64,15 +64,15 @@ mod tests {
6464
ModelInfo {
6565
uuid: uuid.to_string(),
6666
name: name.to_string(),
67+
provider: "huggingface".to_string(),
6768
author: Some(author.to_string()),
6869
task: Some("text-generation".to_string()),
6970
model_series: Some("gpt2".to_string()),
70-
provider: "huggingface".to_string(),
7171
license: Some("mit".to_string()),
7272
created_at: "2025-01-01T00:00:00Z".to_string(),
7373
updated_at: "2025-01-01T00:00:00Z".to_string(),
7474
metadata: ModelMetadata {
75-
artifact: ArtifactInfo {
75+
cache: CacheInfo {
7676
revision: uuid.to_string(),
7777
size: 1000,
7878
path: "/tmp/test".to_string(),

‎src/cli/rm.rs‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@ pub fn execute(registry: &ModelRegistry, model_name: &str) -> Result<(), String>
1414
#[cfg(test)]
1515
mod tests {
1616
use super::*;
17-
use crate::registry::model_registry::{ArtifactInfo, ModelInfo, ModelMetadata};
17+
use crate::registry::model_registry::{CacheInfo, ModelInfo, ModelMetadata};
1818
use tempfile::TempDir;
1919

2020
fn create_test_model(name: &str, uuid: &str) -> ModelInfo {
@@ -26,15 +26,15 @@ mod tests {
2626
ModelInfo {
2727
uuid: uuid.to_string(),
2828
name: name.to_string(),
29+
provider: "huggingface".to_string(),
2930
author: Some("test-author".to_string()),
3031
task: Some("text-generation".to_string()),
3132
model_series: Some("gpt2".to_string()),
32-
provider: "huggingface".to_string(),
3333
license: Some("mit".to_string()),
3434
created_at: "2025-01-01T00:00:00Z".to_string(),
3535
updated_at: "2025-01-01T00:00:00Z".to_string(),
3636
metadata: ModelMetadata {
37-
artifact: ArtifactInfo {
37+
cache: CacheInfo {
3838
revision: uuid.to_string(),
3939
size: 1000,
4040
path: "/tmp/test".to_string(),
@@ -55,7 +55,7 @@ mod tests {
5555
std::fs::write(cache_dir.join("model.safetensors"), "fake data").unwrap();
5656

5757
let mut model = create_test_model("test/remove-model", "abc123");
58-
model.metadata.artifact.path = cache_dir.to_string_lossy().to_string();
58+
model.metadata.cache.path = cache_dir.to_string_lossy().to_string();
5959

6060
registry.register_model(model).unwrap();
6161
assert!(registry.get_model("test/remove-model").unwrap().is_some());

‎src/downloader/huggingface.rs‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@ use indicatif::{ProgressBar, ProgressStyle};
66

77
use crate::downloader::downloader::{DownloadError, Downloader};
88
use crate::downloader::progress::{DownloadProgressManager, FileProgress};
9-
use crate::registry::model_registry::{ArtifactInfo, ModelInfo, ModelMetadata, ModelRegistry};
9+
use crate::registry::model_registry::{CacheInfo, ModelInfo, ModelMetadata, ModelRegistry};
1010
use crate::utils::file::{self, format_model_name};
1111

1212
/// Adapter to bridge HuggingFace's Progress trait with our FileProgress
@@ -284,14 +284,14 @@ impl Downloader for HuggingFaceDownloader {
284284
let model_size =
285285
storage_from_api.unwrap_or_else(|| progress_manager.total_downloaded_bytes());
286286

287-
let artifact = ArtifactInfo {
287+
let cache = CacheInfo {
288288
revision: sha.clone(),
289289
size: model_size,
290290
path: model_cache_path.to_string_lossy().to_string(),
291291
};
292292

293293
let metadata = ModelMetadata {
294-
artifact,
294+
cache,
295295
context_window,
296296
safetensors: safetensors_from_api,
297297
};
@@ -300,10 +300,10 @@ impl Downloader for HuggingFaceDownloader {
300300
let model_info_record = ModelInfo {
301301
uuid: sha, // Use revision SHA as UUID for now
302302
name: name.to_string(),
303+
provider: "huggingface".to_string(),
303304
author: author_from_api,
304305
task: task_from_api,
305306
model_series: model_series_from_api,
306-
provider: "huggingface".to_string(),
307307
license: license_from_api,
308308
created_at: now.clone(),
309309
updated_at: now,

‎src/main.rs‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -17,8 +17,9 @@ fn main() {
1717
// Setup tracing subscriber for tower-http TraceLayer
1818
tracing_subscriber::fmt()
1919
.with_env_filter(
20-
tracing_subscriber::EnvFilter::try_from_default_env()
21-
.unwrap_or_else(|_| "info,hf_hub=warn,tower_http=info".into()),
20+
tracing_subscriber::EnvFilter::try_from_default_env().unwrap_or_else(|_| {
21+
"info,hf_hub=warn,tower_http=info,rusqlite_migration=warn".into()
22+
}),
2223
)
2324
.init();
2425

0 commit comments

Comments
 (0)