Skip to content

Commit adbb9e0

Browse files
kerthcetkerthcet
authored
Support filter in puma ls (#33)
* support label filtering Signed-off-by: kerthcet <kerthcet@gmail.com> * add tests Signed-off-by: kerthcet <kerthcet@gmail.com> * better organize the structure Signed-off-by: kerthcet <kerthcet@gmail.com> * fix task type error Signed-off-by: kerthcet <kerthcet@gmail.com> * support all Signed-off-by: kerthcet <kerthcet@gmail.com> * fix lint Signed-off-by: kerthcet <kerthcet@gmail.com> --------- Signed-off-by: kerthcet <kerthcet@gmail.com>
1 parent 1e6282b commit adbb9e0

14 files changed

Lines changed: 950 additions & 168 deletions

File tree

‎Cargo.lock‎

Lines changed: 1 addition & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎Cargo.toml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@ serde_json = "1.0"
2424
sysinfo = "0.32"
2525
rusqlite = { version = "0.32", features = ["bundled"] }
2626
rusqlite_migration = "1.3"
27+
regex = "1.11"
2728

2829
[dev-dependencies]
2930
tempfile = "3.12"

‎src/cli/commands.rs‎

Lines changed: 42 additions & 129 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
use clap::{Parser, Subcommand};
22
use prettytable::{format, row, Table};
33

4+
use crate::cli::{inspect, ls, rm};
45
use crate::downloader::downloader::Downloader;
56
use crate::downloader::huggingface::HuggingFaceDownloader;
67
use crate::registry::model_registry::ModelRegistry;
@@ -21,7 +22,7 @@ enum Commands {
2122
/// List running models
2223
PS,
2324
/// List local models
24-
LS,
25+
LS(LsArgs),
2526
/// Download a model from a model provider
2627
PULL(PullArgs),
2728
/// Create and run a new model
@@ -38,9 +39,19 @@ enum Commands {
3839
VERSION,
3940
}
4041

42+
#[derive(Parser)]
43+
struct LsArgs {
44+
/// Optional model name pattern to filter (e.g., qwen, openai/*)
45+
pattern: Option<String>,
46+
47+
/// Advanced filter using SQL WHERE conditions (e.g., author=inftyai,license=mit)
48+
#[arg(short = 'l', long, value_name = "KEY=VALUE,...")]
49+
query: Option<String>,
50+
}
51+
4152
#[derive(Parser)]
4253
struct PullArgs {
43-
/// Model name to download (e.g., InftyAI/tiny-random-gpt2)
54+
/// Model name to download (e.g., inftyai/tiny-random-gpt2)
4455
model: String,
4556
#[arg(
4657
short = 'p',
@@ -54,13 +65,13 @@ struct PullArgs {
5465

5566
#[derive(Parser)]
5667
struct RmArgs {
57-
/// Model name to remove (e.g., InftyAI/tiny-random-gpt2)
68+
/// Model name to remove (e.g., inftyai/tiny-random-gpt2)
5869
model: String,
5970
}
6071

6172
#[derive(Parser)]
6273
struct InspectArgs {
63-
/// Model name to inspect (e.g., InftyAI/tiny-random-gpt2)
74+
/// Model name to inspect (e.g., inftyai/tiny-random-gpt2)
6475
model: String,
6576
}
6677

@@ -94,9 +105,17 @@ pub async fn run(cli: Cli) {
94105
table.printstd();
95106
}
96107

97-
Commands::LS => {
108+
Commands::LS(args) => {
98109
let registry = ModelRegistry::new(None);
99-
let models = registry.load_models().unwrap_or_default();
110+
111+
let models =
112+
match ls::execute(&registry, args.pattern.as_deref(), args.query.as_deref()) {
113+
Ok(models) => models,
114+
Err(e) => {
115+
eprintln!("{}", e);
116+
std::process::exit(1);
117+
}
118+
};
100119

101120
let mut table = Table::new();
102121
table.set_format(
@@ -105,7 +124,9 @@ pub async fn run(cli: Cli) {
105124
.padding(0, 1)
106125
.build(),
107126
);
108-
table.add_row(row!["MODEL", "PROVIDER", "REVISION", "SIZE", "AGE"]);
127+
table.add_row(row![
128+
"MODEL", "TASK", "PROVIDER", "REVISION", "SIZE", "CREATED"
129+
]);
109130
for model in models {
110131
let size_str = format_size_decimal(model.metadata.artifact.size);
111132

@@ -117,8 +138,11 @@ pub async fn run(cli: Cli) {
117138

118139
let created_str = format_time_ago(&model.created_at);
119140

141+
let model_task = model.task.as_deref().unwrap_or("N/A");
142+
120143
table.add_row(row![
121144
model.name,
145+
model_task,
122146
model.provider,
123147
revision_short,
124148
size_str,
@@ -153,23 +177,9 @@ pub async fn run(cli: Cli) {
153177
Commands::RM(args) => {
154178
let registry = ModelRegistry::new(None);
155179

156-
// Check if model exists first
157-
match registry.get_model(&args.model) {
158-
Ok(Some(_)) => {
159-
// Delete model (cache + registry)
160-
if let Err(e) = registry.remove_model(&args.model) {
161-
eprintln!("Failed to remove model: {}", e);
162-
std::process::exit(1);
163-
}
164-
}
165-
Ok(None) => {
166-
eprintln!("Model not found: {}", args.model);
167-
std::process::exit(1);
168-
}
169-
Err(e) => {
170-
eprintln!("Failed to load registry: {}", e);
171-
std::process::exit(1);
172-
}
180+
if let Err(e) = rm::execute(&registry, &args.model) {
181+
eprintln!("{}", e);
182+
std::process::exit(1);
173183
}
174184
}
175185

@@ -181,81 +191,10 @@ pub async fn run(cli: Cli) {
181191
Commands::INSPECT(args) => {
182192
let registry = ModelRegistry::new(None);
183193

184-
match registry.get_model(&args.model) {
185-
Ok(Some(model)) => {
186-
println!("Name: {}", model.name);
187-
println!("Kind: Model");
188-
println!("Spec:");
189-
println!(
190-
" Author: {}",
191-
model.author.as_deref().unwrap_or("N/A")
192-
);
193-
println!(
194-
" Type: {}",
195-
model.r#type.as_deref().unwrap_or("N/A")
196-
);
197-
println!(
198-
" License: {}",
199-
model
200-
.license
201-
.as_ref()
202-
.map(|s| s.to_uppercase())
203-
.unwrap_or_else(|| "N/A".to_string())
204-
);
205-
println!(
206-
" Model Series: {}",
207-
model.model_series.as_deref().unwrap_or("N/A")
208-
);
209-
println!(
210-
" Context Window: {}",
211-
model
212-
.metadata
213-
.context_window
214-
.map(|w| crate::utils::format::format_parameters(w as u64))
215-
.unwrap_or_else(|| "N/A".to_string())
216-
);
217-
if let Some(st) = &model.metadata.safetensors {
218-
println!(" Safetensors:");
219-
if let Some(total) = st.get("total").and_then(|v| v.as_u64()) {
220-
println!(
221-
" Total: {}",
222-
crate::utils::format::format_parameters(total)
223-
);
224-
}
225-
if let Some(params) = st.get("parameters").and_then(|v| v.as_object()) {
226-
println!(" Parameters:");
227-
for (dtype, count) in params {
228-
if let Some(num) = count.as_u64() {
229-
println!(
230-
" {:<12} {}",
231-
format!("{}:", dtype),
232-
crate::utils::format::format_parameters(num)
233-
);
234-
}
235-
}
236-
}
237-
} else {
238-
println!(" Safetensors: N/A");
239-
}
240-
// Artifact section
241-
println!(" Artifact:");
242-
println!(" Provider: {}", model.provider);
243-
println!(" Revision: {}", model.metadata.artifact.revision);
244-
println!(
245-
" Size: {}",
246-
format_size_decimal(model.metadata.artifact.size)
247-
);
248-
println!(" Cache Path: {}", model.metadata.artifact.path);
249-
println!("Status:");
250-
println!(" Created: {}", format_time_ago(&model.created_at));
251-
println!(" Updated: {}", format_time_ago(&model.updated_at));
252-
}
253-
Ok(None) => {
254-
eprintln!("Model not found: {}", args.model);
255-
std::process::exit(1);
256-
}
194+
match inspect::execute(&registry, &args.model) {
195+
Ok(model) => inspect::display(&model),
257196
Err(e) => {
258-
eprintln!("Failed to load registry: {}", e);
197+
eprintln!("{}", e);
259198
std::process::exit(1);
260199
}
261200
}
@@ -286,7 +225,7 @@ mod tests {
286225
uuid: revision.to_string(),
287226
name: name.to_string(),
288227
author: Some("test-author".to_string()),
289-
r#type: Some("text-generation".to_string()),
228+
task: Some("text-generation".to_string()),
290229
model_series: Some("gpt2".to_string()),
291230
provider: "huggingface".to_string(),
292231
license: Some("mit".to_string()),
@@ -309,7 +248,7 @@ mod tests {
309248
let temp_dir = TempDir::new().unwrap();
310249
let registry = ModelRegistry::new(Some(temp_dir.path().to_path_buf()));
311250

312-
let models = registry.load_models().unwrap_or_default();
251+
let models = registry.load_models(None).unwrap_or_default();
313252
assert_eq!(models.len(), 0);
314253
}
315254

@@ -322,7 +261,7 @@ mod tests {
322261

323262
registry.register_model(model).unwrap();
324263

325-
let models = registry.load_models().unwrap();
264+
let models = registry.load_models(None).unwrap();
326265
assert_eq!(models.len(), 1);
327266
assert_eq!(models[0].name, "test/model");
328267
assert_eq!(models[0].provider, "huggingface");
@@ -335,7 +274,7 @@ mod tests {
335274

336275
let mut model = create_test_model("test/gpt-model", "abc123def456");
337276
model.author = Some("test-org".to_string());
338-
model.r#type = Some("text-generation".to_string());
277+
model.task = Some("text-generation".to_string());
339278
model.license = Some("mit".to_string());
340279
model.updated_at = "2025-01-02T00:00:00Z".to_string();
341280

@@ -349,7 +288,7 @@ mod tests {
349288
assert_eq!(model_info.created_at, "2025-01-01T00:00:00Z");
350289
assert_eq!(model_info.updated_at, "2025-01-02T00:00:00Z");
351290
assert_eq!(model_info.author, Some("test-org".to_string()));
352-
assert_eq!(model_info.r#type, Some("text-generation".to_string()));
291+
assert_eq!(model_info.task, Some("text-generation".to_string()));
353292
assert_eq!(model_info.license, Some("mit".to_string()));
354293
assert_eq!(model_info.model_series, Some("gpt2".to_string()));
355294
assert_eq!(model_info.metadata.context_window, Some(2048));
@@ -386,32 +325,6 @@ mod tests {
386325
assert!(model_info.metadata.safetensors.is_none());
387326
}
388327

389-
#[test]
390-
fn test_rm_command() {
391-
let temp_dir = TempDir::new().unwrap();
392-
let registry = ModelRegistry::new(Some(temp_dir.path().to_path_buf()));
393-
394-
let model = create_test_model("test/remove-model", "abc123");
395-
396-
registry.register_model(model).unwrap();
397-
assert!(registry.get_model("test/remove-model").unwrap().is_some());
398-
399-
// Simulate RM command
400-
let result = registry.get_model("test/remove-model");
401-
assert!(result.is_ok());
402-
assert!(result.unwrap().is_some());
403-
}
404-
405-
#[test]
406-
fn test_rm_command_nonexistent() {
407-
let temp_dir = TempDir::new().unwrap();
408-
let registry = ModelRegistry::new(Some(temp_dir.path().to_path_buf()));
409-
410-
let result = registry.get_model("nonexistent/model");
411-
assert!(result.is_ok());
412-
assert!(result.unwrap().is_none());
413-
}
414-
415328
#[test]
416329
fn test_revision_truncation() {
417330
let long_revision = "abc123def456ghi789jkl012";

0 commit comments

Comments
 (0)