diff --git a/README.md b/README.md index 472d99c..f6c440e 100644 --- a/README.md +++ b/README.md @@ -89,7 +89,7 @@ Interested in new kernel competitions? Join [discord.gg/gpumode](https://discord ## Discover Problems -The CLI supports (almost) everything Discord does, so you can also discover which leaderboards are available. To make discovery more pleasant we also offer a TUI experience. +The CLI supports everything Discord does, so you can also discover which leaderboards are available. To make discovery more pleasant we also offer a TUI experience. ```bash popcorn-cli submit diff --git a/src/cmd/admin.rs b/src/cmd/admin.rs new file mode 100644 index 0000000..338f373 --- /dev/null +++ b/src/cmd/admin.rs @@ -0,0 +1,179 @@ +use anyhow::{anyhow, Result}; +use clap::Subcommand; +use std::env; + +use crate::service; + +#[derive(Subcommand, Debug)] +pub enum AdminAction { + /// Start accepting jobs on the server + Start, + /// Stop accepting jobs on the server + Stop, + /// Get server statistics + Stats { + /// Only show stats for the last 24 hours + #[arg(long)] + last_day: bool, + }, + /// Get a submission by ID + GetSubmission { + /// The submission ID to retrieve + id: i64, + }, + /// Delete a submission by ID + DeleteSubmission { + /// The submission ID to delete + id: i64, + }, + /// Create a dev leaderboard from a problem directory (requires gpus in task.yml) + CreateLeaderboard { + /// Problem directory name (e.g., "identity_py") + directory: String, + }, + /// Delete a leaderboard + DeleteLeaderboard { + /// Name of the leaderboard to delete + name: String, + /// Force deletion even if there are submissions + #[arg(long)] + force: bool, + }, + /// Update problems from a GitHub repository (mirrors Discord /admin update-problems) + UpdateProblems { + /// Problem set name (e.g., "nvidia", "pmpp_v2"). If not specified, updates all. + #[arg(long)] + problem_set: Option, + + /// Repository in format "owner/repo" (default: gpu-mode/reference-kernels) + #[arg(long, default_value = "gpu-mode/reference-kernels")] + repository: String, + + /// Branch to pull from (default: main) + #[arg(long, default_value = "main")] + branch: String, + + /// Force update even if task definition changed significantly + #[arg(long)] + force: bool, + }, +} + +fn get_admin_token() -> Result { + env::var("POPCORN_ADMIN_TOKEN").map_err(|_| { + anyhow!( + "POPCORN_ADMIN_TOKEN environment variable is not set.\n\ + Set it to your admin token to use admin commands:\n\ + export POPCORN_ADMIN_TOKEN=your_token_here" + ) + }) +} + +pub async fn handle_admin(action: AdminAction) -> Result<()> { + let admin_token = get_admin_token()?; + let client = service::create_admin_client(&admin_token)?; + + match action { + AdminAction::Start => { + let result = service::admin_start(&client).await?; + println!("Server started accepting jobs"); + println!("{}", serde_json::to_string_pretty(&result)?); + } + AdminAction::Stop => { + let result = service::admin_stop(&client).await?; + println!("Server stopped accepting jobs"); + println!("{}", serde_json::to_string_pretty(&result)?); + } + AdminAction::Stats { last_day } => { + let result = service::admin_stats(&client, last_day).await?; + println!("{}", serde_json::to_string_pretty(&result)?); + } + AdminAction::GetSubmission { id } => { + let result = service::admin_get_submission(&client, id).await?; + println!("{}", serde_json::to_string_pretty(&result)?); + } + AdminAction::DeleteSubmission { id } => { + let result = service::admin_delete_submission(&client, id).await?; + println!("Deleted submission {}", id); + println!("{}", serde_json::to_string_pretty(&result)?); + } + AdminAction::CreateLeaderboard { directory } => { + let result = service::admin_create_leaderboard(&client, &directory).await?; + let name = result["leaderboard"].as_str().unwrap_or(&directory); + println!("Created leaderboard '{}'", name); + println!("{}", serde_json::to_string_pretty(&result)?); + } + AdminAction::DeleteLeaderboard { name, force } => { + let result = service::admin_delete_leaderboard(&client, &name, force).await?; + println!("Deleted leaderboard '{}'", name); + println!("{}", serde_json::to_string_pretty(&result)?); + } + AdminAction::UpdateProblems { + problem_set, + repository, + branch, + force, + } => { + println!( + "Updating problems from {}/tree/{}{}...", + repository, + branch, + problem_set + .as_ref() + .map(|ps| format!(" (problem set: {})", ps)) + .unwrap_or_default() + ); + let result = service::admin_update_problems( + &client, + problem_set.as_deref(), + &repository, + &branch, + force, + ) + .await?; + + // Pretty print the results + if let Some(created) = result.get("created").and_then(|v| v.as_array()) { + if !created.is_empty() { + println!("\nCreated {} leaderboard(s):", created.len()); + for name in created { + println!(" + {}", name.as_str().unwrap_or("unknown")); + } + } + } + if let Some(updated) = result.get("updated").and_then(|v| v.as_array()) { + if !updated.is_empty() { + println!("\nUpdated {} leaderboard(s):", updated.len()); + for name in updated { + println!(" ~ {}", name.as_str().unwrap_or("unknown")); + } + } + } + if let Some(skipped) = result.get("skipped").and_then(|v| v.as_array()) { + if !skipped.is_empty() { + println!("\nSkipped {} leaderboard(s):", skipped.len()); + for item in skipped { + let name = item.get("name").and_then(|n| n.as_str()).unwrap_or("unknown"); + let reason = item + .get("reason") + .and_then(|r| r.as_str()) + .unwrap_or("no changes"); + println!(" - {} ({})", name, reason); + } + } + } + if let Some(errors) = result.get("errors").and_then(|v| v.as_array()) { + if !errors.is_empty() { + println!("\nErrors ({}):", errors.len()); + for item in errors { + let name = item.get("name").and_then(|n| n.as_str()).unwrap_or("unknown"); + let error = item.get("error").and_then(|e| e.as_str()).unwrap_or("unknown"); + println!(" ! {}: {}", name, error); + } + } + } + } + } + + Ok(()) +} diff --git a/src/cmd/mod.rs b/src/cmd/mod.rs index 9de3637..6da81bc 100644 --- a/src/cmd/mod.rs +++ b/src/cmd/mod.rs @@ -6,9 +6,12 @@ use serde_yaml; use std::fs::File; use std::path::PathBuf; +mod admin; mod auth; mod submit; +pub use admin::AdminAction; + #[derive(Serialize, Deserialize, Debug, Default)] struct Config { cli_id: Option, @@ -105,6 +108,11 @@ enum Commands { #[arg(long)] no_tui: bool, }, + /// Admin commands (requires POPCORN_ADMIN_TOKEN env var) + Admin { + #[command(subcommand)] + action: AdminAction, + }, } pub async fn execute(cli: Cli) -> Result<()> { @@ -142,7 +150,7 @@ pub async fn execute(cli: Cli) -> Result<()> { // Use filepath from Submit command first, fallback to top-level filepath let final_filepath = filepath.or(cli.filepath); - + if no_tui { submit::run_submit_plain( final_filepath, // Resolved filepath @@ -165,6 +173,9 @@ pub async fn execute(cli: Cli) -> Result<()> { .await } } + Some(Commands::Admin { action }) => { + admin::handle_admin(action).await + } None => { // Check if any of the submission-related flags were used at the top level if cli.gpu.is_some() || cli.leaderboard.is_some() || cli.mode.is_some() { diff --git a/src/service/mod.rs b/src/service/mod.rs index 264b428..b5b6409 100644 --- a/src/service/mod.rs +++ b/src/service/mod.rs @@ -34,6 +34,187 @@ pub fn create_client(cli_id: Option) -> Result { .map_err(|e| anyhow!("Failed to create HTTP client: {}", e)) } +/// Create an HTTP client with admin token authentication +pub fn create_admin_client(admin_token: &str) -> Result { + let mut default_headers = HeaderMap::new(); + + let auth_value = format!("Bearer {}", admin_token); + match HeaderValue::from_str(&auth_value) { + Ok(val) => { + default_headers.insert("Authorization", val); + } + Err(_) => { + return Err(anyhow!("Invalid admin token format for HTTP header")); + } + } + + Client::builder() + .timeout(Duration::from_secs(60)) + .default_headers(default_headers) + .build() + .map_err(|e| anyhow!("Failed to create HTTP client: {}", e)) +} + +/// Start accepting jobs on the server +pub async fn admin_start(client: &Client) -> Result { + let base_url = env::var("POPCORN_API_URL").map_err(|_| anyhow!("POPCORN_API_URL is not set"))?; + + let resp = client + .post(format!("{}/admin/start", base_url)) + .header("Content-Length", "0") + .timeout(Duration::from_secs(30)) + .send() + .await?; + + handle_admin_response(resp).await +} + +/// Stop accepting jobs on the server +pub async fn admin_stop(client: &Client) -> Result { + let base_url = env::var("POPCORN_API_URL").map_err(|_| anyhow!("POPCORN_API_URL is not set"))?; + + let resp = client + .post(format!("{}/admin/stop", base_url)) + .header("Content-Length", "0") + .timeout(Duration::from_secs(30)) + .send() + .await?; + + handle_admin_response(resp).await +} + +/// Get server stats +pub async fn admin_stats(client: &Client, last_day_only: bool) -> Result { + let base_url = env::var("POPCORN_API_URL").map_err(|_| anyhow!("POPCORN_API_URL is not set"))?; + + let url = if last_day_only { + format!("{}/admin/stats?last_day_only=true", base_url) + } else { + format!("{}/admin/stats", base_url) + }; + + let resp = client + .get(url) + .timeout(Duration::from_secs(30)) + .send() + .await?; + + handle_admin_response(resp).await +} + +/// Get a submission by ID +pub async fn admin_get_submission(client: &Client, submission_id: i64) -> Result { + let base_url = env::var("POPCORN_API_URL").map_err(|_| anyhow!("POPCORN_API_URL is not set"))?; + + let resp = client + .get(format!("{}/admin/submissions/{}", base_url, submission_id)) + .timeout(Duration::from_secs(30)) + .send() + .await?; + + handle_admin_response(resp).await +} + +/// Delete a submission by ID +pub async fn admin_delete_submission(client: &Client, submission_id: i64) -> Result { + let base_url = env::var("POPCORN_API_URL").map_err(|_| anyhow!("POPCORN_API_URL is not set"))?; + + let resp = client + .delete(format!("{}/admin/submissions/{}", base_url, submission_id)) + .timeout(Duration::from_secs(30)) + .send() + .await?; + + handle_admin_response(resp).await +} + +/// Create a dev leaderboard from a problem directory +pub async fn admin_create_leaderboard( + client: &Client, + directory: &str, +) -> Result { + let base_url = env::var("POPCORN_API_URL").map_err(|_| anyhow!("POPCORN_API_URL is not set"))?; + + let payload = serde_json::json!({ + "directory": directory + }); + + let resp = client + .post(format!("{}/admin/leaderboards", base_url)) + .json(&payload) + .timeout(Duration::from_secs(30)) + .send() + .await?; + + handle_admin_response(resp).await +} + +/// Delete a leaderboard +pub async fn admin_delete_leaderboard(client: &Client, leaderboard_name: &str, force: bool) -> Result { + let base_url = env::var("POPCORN_API_URL").map_err(|_| anyhow!("POPCORN_API_URL is not set"))?; + + let url = if force { + format!("{}/admin/leaderboards/{}?force=true", base_url, leaderboard_name) + } else { + format!("{}/admin/leaderboards/{}", base_url, leaderboard_name) + }; + + let resp = client + .delete(url) + .timeout(Duration::from_secs(30)) + .send() + .await?; + + handle_admin_response(resp).await +} + +/// Update problems from a GitHub repository +pub async fn admin_update_problems( + client: &Client, + problem_set: Option<&str>, + repository: &str, + branch: &str, + force: bool, +) -> Result { + let base_url = env::var("POPCORN_API_URL").map_err(|_| anyhow!("POPCORN_API_URL is not set"))?; + + let mut payload = serde_json::json!({ + "repository": repository, + "branch": branch, + "force": force + }); + + if let Some(ps) = problem_set { + payload["problem_set"] = serde_json::Value::String(ps.to_string()); + } + + let resp = client + .post(format!("{}/admin/update-problems", base_url)) + .json(&payload) + .timeout(Duration::from_secs(120)) // Longer timeout for repo download + .send() + .await?; + + handle_admin_response(resp).await +} + +/// Helper to handle admin API responses +async fn handle_admin_response(resp: reqwest::Response) -> Result { + let status = resp.status(); + if !status.is_success() { + let error_text = resp.text().await?; + let detail = serde_json::from_str::(&error_text) + .ok() + .and_then(|v| v.get("detail").and_then(|d| d.as_str()).map(str::to_string)); + return Err(anyhow!( + "Server returned status {}: {}", + status, + detail.unwrap_or(error_text) + )); + } + resp.json().await.map_err(|e| anyhow!("Failed to parse response: {}", e)) +} + pub async fn fetch_leaderboards(client: &Client) -> Result> { let base_url = env::var("POPCORN_API_URL").map_err(|_| anyhow!("POPCORN_API_URL is not set"))?;