From d04a85bac6df847a842d871c0361124fc37199e6 Mon Sep 17 00:00:00 2001 From: kerthcet Date: Sun, 5 Jul 2026 17:13:46 +0100 Subject: [PATCH 1/3] support puma run Signed-off-by: kerthcet --- src/cli/commands.rs | 76 +++++++++++++++++++++++++++++-------------- src/downloader/mod.rs | 25 ++++++++++++++ 2 files changed, 77 insertions(+), 24 deletions(-) diff --git a/src/cli/commands.rs b/src/cli/commands.rs index 6b92e77..d2ffa19 100644 --- a/src/cli/commands.rs +++ b/src/cli/commands.rs @@ -2,8 +2,7 @@ use clap::{Parser, Subcommand}; use prettytable::{format, row, Table}; use crate::cli::{inspect, ls, rm}; -use crate::downloader::huggingface::HuggingFaceDownloader; -use crate::downloader::Downloader; +use crate::downloader::{self, Provider}; use crate::registry::model_registry::ModelRegistry; use crate::system::system_info::SystemInfo; use crate::utils::format::{format_size_decimal, format_time_ago}; @@ -26,7 +25,7 @@ enum Commands { /// Download a model from a model provider PULL(PullArgs), /// Create and run a new model - RUN, + RUN(RunArgs), /// Stop one running model STOP, /// Remove one model @@ -41,6 +40,22 @@ enum Commands { SERVE(ServeArgs), } +#[derive(Parser)] +struct RunArgs { + /// Model name to run (e.g., inftyai/tiny-random-gpt2) + model: String, + + /// Provider to download from if model not found + #[arg( + short = 'p', + long, + value_name = "model provider", + value_enum, + default_value = "huggingface" + )] + provider: Provider, +} + #[derive(Parser)] struct ServeArgs { /// Model name to serve (e.g., inftyai/tiny-random-gpt2) @@ -91,14 +106,6 @@ struct InspectArgs { model: String, } -#[derive(Debug, Clone, Default, clap::ValueEnum)] -pub enum Provider { - #[default] - #[value(alias = "hf")] - Huggingface, - #[value(alias = "ms")] - Modelscope, -} // Support commands like: pull, ls, run, ps, stop, rm, info, inspect, show. pub async fn run(cli: Cli) { @@ -171,22 +178,43 @@ pub async fn run(cli: Cli) { table.printstd(); } - Commands::PULL(args) => match args.provider { - Provider::Huggingface => { - let downloader = HuggingFaceDownloader::new(); - // Make sure to use lowercase for model name to ensure consistent caching and registry entries. - if let Err(e) = downloader.download_model(&args.model.to_lowercase()).await { - eprintln!("❌ Error downloading model: {}", e); + Commands::PULL(args) => { + if let Err(e) = downloader::download_model(&args.model, args.provider).await { + eprintln!("❌ Error: {}", e); + std::process::exit(1); + } + } + + Commands::RUN(args) => { + let registry = ModelRegistry::new(None); + + // Check if model exists + match registry.get_model(&args.model) { + Ok(Some(_)) => { + // Model exists, proceed + println!("Running model: {}", args.model); + // TODO: Implement actual model execution + println!("Model execution not yet implemented"); + } + Ok(None) => { + // Model not found, download it first + println!("Model '{}' not found locally. Downloading...", args.model); + + if let Err(e) = downloader::download_model(&args.model, args.provider).await { + eprintln!("❌ Error: {}", e); + std::process::exit(1); + } + + // Now run the model + println!("Running model: {}", args.model); + // TODO: Implement actual model execution + println!("Model execution not yet implemented"); + } + Err(e) => { + eprintln!("❌ Error checking model: {}", e); std::process::exit(1); } } - Provider::Modelscope => { - println!("Downloading model from Modelscope..."); - } - }, - - Commands::RUN => { - println!("Creating and running a new model..."); } Commands::STOP => { diff --git a/src/downloader/mod.rs b/src/downloader/mod.rs index 2389b29..2ff0f57 100644 --- a/src/downloader/mod.rs +++ b/src/downloader/mod.rs @@ -29,3 +29,28 @@ impl fmt::Display for DownloadError { pub trait Downloader { async fn download_model(&self, name: &str) -> Result<(), DownloadError>; } + +/// Provider for downloading models +#[derive(Debug, Clone, Copy, Default, clap::ValueEnum)] +pub enum Provider { + #[default] + #[value(alias = "hf")] + Huggingface, + #[value(alias = "ms")] + Modelscope, +} + +/// Download a model from the specified provider +pub async fn download_model(model_name: &str, provider: Provider) -> Result<(), DownloadError> { + match provider { + Provider::Huggingface => { + let downloader = huggingface::HuggingFaceDownloader::new(); + downloader.download_model(&model_name.to_lowercase()).await + } + Provider::Modelscope => { + Err(DownloadError::ApiError( + "Modelscope provider not yet implemented".to_string(), + )) + } + } +} From 3550de8a992b741b7f3854c131529bcb01f00bb4 Mon Sep 17 00:00:00 2001 From: kerthcet Date: Sun, 5 Jul 2026 17:30:27 +0100 Subject: [PATCH 2/3] fix lint Signed-off-by: kerthcet --- src/cli/commands.rs | 9 ++++++--- src/downloader/mod.rs | 8 +++----- 2 files changed, 9 insertions(+), 8 deletions(-) diff --git a/src/cli/commands.rs b/src/cli/commands.rs index d2ffa19..de24edf 100644 --- a/src/cli/commands.rs +++ b/src/cli/commands.rs @@ -1,4 +1,5 @@ use clap::{Parser, Subcommand}; +use colored::Colorize; use prettytable::{format, row, Table}; use crate::cli::{inspect, ls, rm}; @@ -106,7 +107,6 @@ struct InspectArgs { model: String, } - // Support commands like: pull, ls, run, ps, stop, rm, info, inspect, show. pub async fn run(cli: Cli) { match cli.command { @@ -198,7 +198,10 @@ pub async fn run(cli: Cli) { } Ok(None) => { // Model not found, download it first - println!("Model '{}' not found locally. Downloading...", args.model); + println!( + "Model {} not found locally. Downloading...", + args.model.cyan().bold() + ); if let Err(e) = downloader::download_model(&args.model, args.provider).await { eprintln!("❌ Error: {}", e); @@ -206,7 +209,7 @@ pub async fn run(cli: Cli) { } // Now run the model - println!("Running model: {}", args.model); + println!("Running model: {}", args.model.cyan().bold()); // TODO: Implement actual model execution println!("Model execution not yet implemented"); } diff --git a/src/downloader/mod.rs b/src/downloader/mod.rs index 2ff0f57..0067561 100644 --- a/src/downloader/mod.rs +++ b/src/downloader/mod.rs @@ -47,10 +47,8 @@ pub async fn download_model(model_name: &str, provider: Provider) -> Result<(), let downloader = huggingface::HuggingFaceDownloader::new(); downloader.download_model(&model_name.to_lowercase()).await } - Provider::Modelscope => { - Err(DownloadError::ApiError( - "Modelscope provider not yet implemented".to_string(), - )) - } + Provider::Modelscope => Err(DownloadError::ApiError( + "Modelscope provider not yet implemented".to_string(), + )), } } From cc08825ee9f08f5da93bfe10b8689e7e9d3d4b29 Mon Sep 17 00:00:00 2001 From: kerthcet Date: Sun, 5 Jul 2026 17:39:15 +0100 Subject: [PATCH 3/3] fix test Signed-off-by: kerthcet --- src/cli/commands.rs | 42 ++++++++++++++++++++++++++++++++++++++++++ tests/cli_test.rs | 7 ++++--- 2 files changed, 46 insertions(+), 3 deletions(-) diff --git a/src/cli/commands.rs b/src/cli/commands.rs index de24edf..8542ca7 100644 --- a/src/cli/commands.rs +++ b/src/cli/commands.rs @@ -497,4 +497,46 @@ mod tests { ]); assert!(result.is_ok()); } + + #[test] + fn test_run_args_parsing() { + use clap::CommandFactory; + let app = Cli::command(); + + // This should fail without model argument + let result = app.clone().try_get_matches_from(vec!["puma", "run"]); + assert!(result.is_err()); + + // This should succeed with model argument (default provider) + let result = app + .clone() + .try_get_matches_from(vec!["puma", "run", "test/model"]); + assert!(result.is_ok()); + + // This should succeed with explicit huggingface provider + let result = app.clone().try_get_matches_from(vec![ + "puma", + "run", + "test/model", + "-p", + "huggingface", + ]); + assert!(result.is_ok()); + + // This should succeed with hf alias + let result = + app.clone() + .try_get_matches_from(vec!["puma", "run", "test/model", "--provider", "hf"]); + assert!(result.is_ok()); + + // This should succeed with modelscope provider + let result = + app.clone() + .try_get_matches_from(vec!["puma", "run", "test/model", "-p", "modelscope"]); + assert!(result.is_ok()); + + // This should succeed with ms alias + let result = app.try_get_matches_from(vec!["puma", "run", "test/model", "-p", "ms"]); + assert!(result.is_ok()); + } } diff --git a/tests/cli_test.rs b/tests/cli_test.rs index a2e1593..50fc86e 100644 --- a/tests/cli_test.rs +++ b/tests/cli_test.rs @@ -172,9 +172,10 @@ fn test_run_command() { let temp_dir = TempDir::new().unwrap(); let home = temp_dir.path().to_str().unwrap(); - let output = run_puma(home, &["run"]); - assert!(output.status.success()); - assert!(output_contains(&output, "Creating and running a new model")); + // RUN command requires a model argument + let output = run_puma(home, &["run", "test/model"]); + // Will fail because test/model doesn't exist on HuggingFace, but at least it attempts + assert!(output_contains(&output, "not found locally") || output_contains(&output, "Error")); } #[test]