Skip to content

Commit 9c2c93f

Browse files
kerthcetkerthcet
authored
feat: add metadata (#30)
* Support HF downloading models (#16) * Add HF downloader support Signed-off-by: kerthcet <kerthcet@gmail.com> * add bars Signed-off-by: kerthcet <kerthcet@gmail.com> * fix color Signed-off-by: kerthcet <kerthcet@gmail.com> * fix color Signed-off-by: kerthcet <kerthcet@gmail.com> * add download successfully message Signed-off-by: kerthcet <kerthcet@gmail.com> * change the color Signed-off-by: kerthcet <kerthcet@gmail.com> * change the rending shape Signed-off-by: kerthcet <kerthcet@gmail.com> --------- Signed-off-by: kerthcet <kerthcet@gmail.com> * Support `puma rm <model>` (#17) * support new cache structure Signed-off-by: kerthcet <kerthcet@gmail.com> * support puma rm Signed-off-by: kerthcet <kerthcet@gmail.com> * use readable format Signed-off-by: kerthcet <kerthcet@gmail.com> * remove requests.rs Signed-off-by: kerthcet <kerthcet@gmail.com> * fix lint Signed-off-by: kerthcet <kerthcet@gmail.com> --------- Signed-off-by: kerthcet <kerthcet@gmail.com> * support puma info (#18) Signed-off-by: kerthcet <kerthcet@gmail.com> * Reuse the model cache to avoid duplicate download (#19) * polish the format of the ls command Signed-off-by: kerthcet <kerthcet@gmail.com> * Have a progress manager Signed-off-by: kerthcet <kerthcet@gmail.com> * Reuse caches Signed-off-by: kerthcet <kerthcet@gmail.com> * rename util to utils Signed-off-by: kerthcet <kerthcet@gmail.com> * polish the layout of the download progress Signed-off-by: kerthcet <kerthcet@gmail.com> * revert change Signed-off-by: kerthcet <kerthcet@gmail.com> * add make format Signed-off-by: kerthcet <kerthcet@gmail.com> --------- Signed-off-by: kerthcet <kerthcet@gmail.com> * remove available mem (#22) Signed-off-by: kerthcet <kerthcet@gmail.com> * add speed at the end (#23) * add speed at the end Signed-off-by: kerthcet <kerthcet@gmail.com> * fix lint Signed-off-by: kerthcet <kerthcet@gmail.com> --------- Signed-off-by: kerthcet <kerthcet@gmail.com> * fix: do no register model once cached (#26) Signed-off-by: kerthcet <kerthcet@gmail.com> * Support GPU detect (#27) * support GPU detect Signed-off-by: kerthcet <kerthcet@gmail.com> * fix lint Signed-off-by: kerthcet <kerthcet@gmail.com> --------- Signed-off-by: kerthcet <kerthcet@gmail.com> * update readme.md (#28) Signed-off-by: kerthcet <kerthcet@gmail.com> * Support inspect command (#29) * add support for inspect Signed-off-by: kerthcet <kerthcet@gmail.com> * add support for inspect Signed-off-by: kerthcet <kerthcet@gmail.com> * add pull progress bar Signed-off-by: kerthcet <kerthcet@gmail.com> * polish the download progress Signed-off-by: kerthcet <kerthcet@gmail.com> * reorganize the structure Signed-off-by: kerthcet <kerthcet@gmail.com> * optimize the structure Signed-off-by: kerthcet <kerthcet@gmail.com> * fix test Signed-off-by: kerthcet <kerthcet@gmail.com> * fix lint Signed-off-by: kerthcet <kerthcet@gmail.com> --------- Signed-off-by: kerthcet <kerthcet@gmail.com> * add metadata Signed-off-by: kerthcet <kerthcet@gmail.com> --------- Signed-off-by: kerthcet <kerthcet@gmail.com>
1 parent a0ebe7e commit 9c2c93f

3 files changed

Lines changed: 243 additions & 18 deletions

File tree

‎src/cli/commands.rs‎

Lines changed: 205 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -105,8 +105,7 @@ pub async fn run(cli: Cli) {
105105
.padding(0, 1)
106106
.build(),
107107
);
108-
table.add_row(row!["MODEL", "PROVIDER", "REVISION", "SIZE", "MODIFIED"]);
109-
108+
table.add_row(row!["MODEL", "PROVIDER", "REVISION", "SIZE", "AGE"]);
110109
for model in models {
111110
let size_str = format_size_decimal(model.size);
112111

@@ -116,7 +115,7 @@ pub async fn run(cli: Cli) {
116115
&model.revision
117116
};
118117

119-
let created_str = format_time_ago(&model.modified_at);
118+
let created_str = format_time_ago(&model.created_at);
120119

121120
table.add_row(row![
122121
model.name,
@@ -186,6 +185,9 @@ pub async fn run(cli: Cli) {
186185
Ok(Some(model)) => {
187186
println!("Name: {}", model.name);
188187
println!("Kind: Model");
188+
println!("Metadata:");
189+
println!(" Created: {}", format_time_ago(&model.created_at));
190+
println!(" Updated: {}", format_time_ago(&model.updated_at));
189191

190192
println!("Spec:");
191193
// Architecture section (only if info is available)
@@ -209,10 +211,6 @@ pub async fn run(cli: Cli) {
209211
println!(" Provider: {}", model.provider);
210212
println!(" Revision: {}", model.revision);
211213
println!(" Size: {}", format_size_decimal(model.size));
212-
println!(
213-
" Modified: {}",
214-
format_time_ago(&model.modified_at)
215-
);
216214
println!(" Cache Path: {}", model.cache_path);
217215
}
218216
Ok(None) => {
@@ -231,3 +229,203 @@ pub async fn run(cli: Cli) {
231229
}
232230
}
233231
}
232+
233+
#[cfg(test)]
234+
mod tests {
235+
use super::*;
236+
use crate::registry::model_registry::{ModelArchitecture, ModelInfo};
237+
use tempfile::TempDir;
238+
239+
#[test]
240+
fn test_ls_command_empty() {
241+
let temp_dir = TempDir::new().unwrap();
242+
let registry = ModelRegistry::new(Some(temp_dir.path().to_path_buf()));
243+
244+
let models = registry.load_models().unwrap_or_default();
245+
assert_eq!(models.len(), 0);
246+
}
247+
248+
#[test]
249+
fn test_ls_command_with_models() {
250+
let temp_dir = TempDir::new().unwrap();
251+
let registry = ModelRegistry::new(Some(temp_dir.path().to_path_buf()));
252+
253+
let model = ModelInfo {
254+
name: "test/model".to_string(),
255+
provider: "huggingface".to_string(),
256+
revision: "abc123def456".to_string(),
257+
size: 1_000_000,
258+
created_at: "2025-01-01T00:00:00Z".to_string(),
259+
updated_at: "2025-01-01T00:00:00Z".to_string(),
260+
cache_path: "/tmp/test".to_string(),
261+
arch: None,
262+
};
263+
264+
registry.register_model(model).unwrap();
265+
266+
let models = registry.load_models().unwrap();
267+
assert_eq!(models.len(), 1);
268+
assert_eq!(models[0].name, "test/model");
269+
assert_eq!(models[0].provider, "huggingface");
270+
}
271+
272+
#[test]
273+
fn test_inspect_command_with_metadata() {
274+
let temp_dir = TempDir::new().unwrap();
275+
let registry = ModelRegistry::new(Some(temp_dir.path().to_path_buf()));
276+
277+
let model = ModelInfo {
278+
name: "test/gpt-model".to_string(),
279+
provider: "huggingface".to_string(),
280+
revision: "abc123def456".to_string(),
281+
size: 7_000_000_000,
282+
created_at: "2025-01-01T00:00:00Z".to_string(),
283+
updated_at: "2025-01-02T00:00:00Z".to_string(),
284+
cache_path: "/tmp/test/gpt".to_string(),
285+
arch: Some(ModelArchitecture {
286+
model_type: Some("gpt2".to_string()),
287+
classes: Some(vec!["GPT2LMHeadModel".to_string()]),
288+
context_window: Some(2048),
289+
parameters: Some("7.00B".to_string()),
290+
}),
291+
};
292+
293+
registry.register_model(model.clone()).unwrap();
294+
295+
let retrieved = registry.get_model("test/gpt-model").unwrap();
296+
assert!(retrieved.is_some());
297+
298+
let model_info = retrieved.unwrap();
299+
assert_eq!(model_info.name, "test/gpt-model");
300+
assert_eq!(model_info.created_at, "2025-01-01T00:00:00Z");
301+
assert_eq!(model_info.updated_at, "2025-01-02T00:00:00Z");
302+
303+
let arch = model_info.arch.unwrap();
304+
assert_eq!(arch.model_type, Some("gpt2".to_string()));
305+
assert_eq!(arch.classes, Some(vec!["GPT2LMHeadModel".to_string()]));
306+
assert_eq!(arch.context_window, Some(2048));
307+
assert_eq!(arch.parameters, Some("7.00B".to_string()));
308+
}
309+
310+
#[test]
311+
fn test_inspect_command_without_architecture() {
312+
let temp_dir = TempDir::new().unwrap();
313+
let registry = ModelRegistry::new(Some(temp_dir.path().to_path_buf()));
314+
315+
let model = ModelInfo {
316+
name: "test/simple-model".to_string(),
317+
provider: "huggingface".to_string(),
318+
revision: "xyz789".to_string(),
319+
size: 500_000,
320+
created_at: "2025-01-01T00:00:00Z".to_string(),
321+
updated_at: "2025-01-01T00:00:00Z".to_string(),
322+
cache_path: "/tmp/test/simple".to_string(),
323+
arch: None,
324+
};
325+
326+
registry.register_model(model).unwrap();
327+
328+
let retrieved = registry.get_model("test/simple-model").unwrap();
329+
assert!(retrieved.is_some());
330+
331+
let model_info = retrieved.unwrap();
332+
assert_eq!(model_info.name, "test/simple-model");
333+
assert!(model_info.arch.is_none());
334+
}
335+
336+
#[test]
337+
fn test_rm_command() {
338+
let temp_dir = TempDir::new().unwrap();
339+
let registry = ModelRegistry::new(Some(temp_dir.path().to_path_buf()));
340+
341+
let model = ModelInfo {
342+
name: "test/remove-model".to_string(),
343+
provider: "huggingface".to_string(),
344+
revision: "abc123".to_string(),
345+
size: 1000,
346+
created_at: "2025-01-01T00:00:00Z".to_string(),
347+
updated_at: "2025-01-01T00:00:00Z".to_string(),
348+
cache_path: "/tmp/test/remove".to_string(),
349+
arch: None,
350+
};
351+
352+
registry.register_model(model).unwrap();
353+
assert!(registry.get_model("test/remove-model").unwrap().is_some());
354+
355+
// Simulate RM command
356+
let result = registry.get_model("test/remove-model");
357+
assert!(result.is_ok());
358+
assert!(result.unwrap().is_some());
359+
}
360+
361+
#[test]
362+
fn test_rm_command_nonexistent() {
363+
let temp_dir = TempDir::new().unwrap();
364+
let registry = ModelRegistry::new(Some(temp_dir.path().to_path_buf()));
365+
366+
let result = registry.get_model("nonexistent/model");
367+
assert!(result.is_ok());
368+
assert!(result.unwrap().is_none());
369+
}
370+
371+
#[test]
372+
fn test_revision_truncation() {
373+
let long_revision = "abc123def456ghi789jkl012";
374+
let short = if long_revision.len() > 8 {
375+
&long_revision[..8]
376+
} else {
377+
long_revision
378+
};
379+
assert_eq!(short, "abc123de");
380+
381+
let short_revision = "abc123";
382+
let short = if short_revision.len() > 8 {
383+
&short_revision[..8]
384+
} else {
385+
short_revision
386+
};
387+
assert_eq!(short, "abc123");
388+
}
389+
390+
#[test]
391+
fn test_metadata_timestamps_differ() {
392+
let temp_dir = TempDir::new().unwrap();
393+
let registry = ModelRegistry::new(Some(temp_dir.path().to_path_buf()));
394+
395+
let model = ModelInfo {
396+
name: "test/updated-model".to_string(),
397+
provider: "huggingface".to_string(),
398+
revision: "v1".to_string(),
399+
size: 1000,
400+
created_at: "2025-01-01T00:00:00Z".to_string(),
401+
updated_at: "2025-01-01T00:00:00Z".to_string(),
402+
cache_path: "/tmp/test".to_string(),
403+
arch: None,
404+
};
405+
406+
registry.register_model(model).unwrap();
407+
408+
// Update the model
409+
let updated_model = ModelInfo {
410+
name: "test/updated-model".to_string(),
411+
provider: "huggingface".to_string(),
412+
revision: "v2".to_string(),
413+
size: 2000,
414+
created_at: "2025-01-05T00:00:00Z".to_string(),
415+
updated_at: "2025-01-05T00:00:00Z".to_string(),
416+
cache_path: "/tmp/test".to_string(),
417+
arch: None,
418+
};
419+
420+
registry.register_model(updated_model).unwrap();
421+
422+
let result = registry.get_model("test/updated-model").unwrap().unwrap();
423+
// created_at should remain the same
424+
assert_eq!(result.created_at, "2025-01-01T00:00:00Z");
425+
// updated_at should be new
426+
assert_eq!(result.updated_at, "2025-01-05T00:00:00Z");
427+
// Other fields should be updated
428+
assert_eq!(result.revision, "v2");
429+
assert_eq!(result.size, 2000);
430+
}
431+
}

‎src/downloader/huggingface.rs‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -213,12 +213,14 @@ impl Downloader for HuggingFaceDownloader {
213213
None
214214
};
215215

216+
let now = chrono::Local::now().to_rfc3339();
216217
let model_info_record = ModelInfo {
217218
name: name.to_string(),
218219
provider: "huggingface".to_string(),
219220
revision: sha,
220221
size: downloaded_size,
221-
modified_at: chrono::Local::now().to_rfc3339(),
222+
created_at: now.clone(),
223+
updated_at: now,
222224
cache_path: model_cache_path.to_string_lossy().to_string(),
223225
arch,
224226
};

‎src/registry/model_registry.rs‎

Lines changed: 35 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -98,7 +98,8 @@ pub struct ModelInfo {
9898
pub provider: String,
9999
pub revision: String,
100100
pub size: u64,
101-
pub modified_at: String,
101+
pub created_at: String,
102+
pub updated_at: String,
102103
pub cache_path: String,
103104
#[serde(skip_serializing_if = "Option::is_none")]
104105
pub arch: Option<ModelArchitecture>,
@@ -148,10 +149,22 @@ impl ModelRegistry {
148149
pub fn register_model(&self, model: ModelInfo) -> Result<(), std::io::Error> {
149150
let mut models = self.load_models()?;
150151

152+
// Check if model already exists to preserve created_at
153+
let existing_created_at = models
154+
.iter()
155+
.find(|m| m.name == model.name)
156+
.map(|m| m.created_at.clone());
157+
151158
// Remove existing model with same name if exists
152159
models.retain(|m| m.name != model.name);
153160

154-
models.push(model);
161+
// Use existing created_at if this is an update, otherwise use the provided one
162+
let mut final_model = model;
163+
if let Some(created_at) = existing_created_at {
164+
final_model.created_at = created_at;
165+
}
166+
167+
models.push(final_model);
155168
self.save_models(&models)?;
156169

157170
Ok(())
@@ -211,7 +224,8 @@ mod tests {
211224
provider: "huggingface".to_string(),
212225
revision: "abc123".to_string(),
213226
size: 1000,
214-
modified_at: "2025-01-01T00:00:00Z".to_string(),
227+
created_at: "2025-01-01T00:00:00Z".to_string(),
228+
updated_at: "2025-01-01T00:00:00Z".to_string(),
215229
cache_path: "/tmp/test".to_string(),
216230
arch: None,
217231
};
@@ -233,7 +247,8 @@ mod tests {
233247
provider: "huggingface".to_string(),
234248
revision: "abc123".to_string(),
235249
size: 1000,
236-
modified_at: "2025-01-01T00:00:00Z".to_string(),
250+
created_at: "2025-01-01T00:00:00Z".to_string(),
251+
updated_at: "2025-01-01T00:00:00Z".to_string(),
237252
cache_path: "/tmp/test".to_string(),
238253
arch: None,
239254
};
@@ -255,7 +270,8 @@ mod tests {
255270
provider: "huggingface".to_string(),
256271
revision: "abc123".to_string(),
257272
size: 1000,
258-
modified_at: "2025-01-01T00:00:00Z".to_string(),
273+
created_at: "2025-01-01T00:00:00Z".to_string(),
274+
updated_at: "2025-01-01T00:00:00Z".to_string(),
259275
cache_path: "/tmp/test".to_string(),
260276
arch: None,
261277
};
@@ -290,7 +306,8 @@ mod tests {
290306
provider: "huggingface".to_string(),
291307
revision: "abc123".to_string(),
292308
size: 1000,
293-
modified_at: "2025-01-01T00:00:00Z".to_string(),
309+
created_at: "2025-01-01T00:00:00Z".to_string(),
310+
updated_at: "2025-01-01T00:00:00Z".to_string(),
294311
cache_path: "/tmp/test".to_string(),
295312
arch: None,
296313
};
@@ -302,7 +319,8 @@ mod tests {
302319
provider: "huggingface".to_string(),
303320
revision: "def456".to_string(),
304321
size: 2000,
305-
modified_at: "2025-01-02T00:00:00Z".to_string(),
322+
created_at: "2025-01-02T00:00:00Z".to_string(),
323+
updated_at: "2025-01-02T00:00:00Z".to_string(),
306324
cache_path: "/tmp/test2".to_string(),
307325
arch: None,
308326
};
@@ -313,6 +331,10 @@ mod tests {
313331
assert_eq!(models.len(), 1);
314332
assert_eq!(models[0].revision, "def456");
315333
assert_eq!(models[0].size, 2000);
334+
// created_at should be preserved from model1
335+
assert_eq!(models[0].created_at, "2025-01-01T00:00:00Z");
336+
// updated_at should be from model2
337+
assert_eq!(models[0].updated_at, "2025-01-02T00:00:00Z");
316338
}
317339

318340
#[test]
@@ -330,7 +352,8 @@ mod tests {
330352
provider: "huggingface".to_string(),
331353
revision: "abc123".to_string(),
332354
size: 1000,
333-
modified_at: "2025-01-01T00:00:00Z".to_string(),
355+
created_at: "2025-01-01T00:00:00Z".to_string(),
356+
updated_at: "2025-01-01T00:00:00Z".to_string(),
334357
cache_path: cache_dir.to_string_lossy().to_string(),
335358
arch: None,
336359
};
@@ -369,7 +392,8 @@ mod tests {
369392
provider: "huggingface".to_string(),
370393
revision: "abc123def456".to_string(),
371394
size: 7_000_000_000,
372-
modified_at: "2025-01-01T00:00:00Z".to_string(),
395+
created_at: "2025-01-01T00:00:00Z".to_string(),
396+
updated_at: "2025-01-01T00:00:00Z".to_string(),
373397
cache_path: "/tmp/test/gpt".to_string(),
374398
arch: Some(ModelArchitecture {
375399
model_type: Some("gpt2".to_string()),
@@ -407,7 +431,8 @@ mod tests {
407431
provider: "huggingface".to_string(),
408432
revision: "legacy123".to_string(),
409433
size: 1_000_000,
410-
modified_at: "2024-01-01T00:00:00Z".to_string(),
434+
created_at: "2024-01-01T00:00:00Z".to_string(),
435+
updated_at: "2024-01-01T00:00:00Z".to_string(),
411436
cache_path: "/tmp/test/legacy".to_string(),
412437
arch: None,
413438
};

0 commit comments

Comments
 (0)