wafer-cli 0.2.3__py3-none-any.whl → 0.2.4__py3-none-any.whl

wafer/cli.py

@@ -30,6 +30,14 @@ import typer
 
 from .config import WaferConfig, WaferEnvironment
 from .inference import infer_upload_files, resolve_environment
+from .problems import (
+    download_problems,
+    get_problem_path,
+    get_problems_path,
+)
+from .problems import (
+    list_problems as list_problems_fn,
+)
 
 app = typer.Typer(
     help="GPU development toolkit for LLM coding agents",
@@ -95,7 +103,7 @@ def main_callback(ctx: typer.Context) -> None:
         global _command_outcome
         # Mark as failure if SystemExit with non-zero code, or any other exception
         if exc_type is SystemExit:
-            exit_code = exc_value.code if hasattr(exc_value, 'code') else 1
+            exit_code = exc_value.code if hasattr(exc_value, "code") else 1
             if exit_code != 0 and exit_code is not None:
                 _command_outcome = "failure"
         else:
@@ -200,6 +208,13 @@ kernelbench_app = typer.Typer(
 )
 evaluate_app.add_typer(kernelbench_app, name="kernelbench")
 
+# Nested subcommand for gpumode format
+gpumode_app = typer.Typer(
+    help="Evaluate kernels in GPUMode format (custom_kernel/ref_kernel functions)",
+    invoke_without_command=True,
+)
+evaluate_app.add_typer(gpumode_app, name="gpumode")
+
 # =============================================================================
 # Dev commands (internal, used by web app proxy)
 # =============================================================================
@@ -396,6 +411,124 @@ def skill_status() -> None:
             typer.echo(f"{tool_name}: Not installed")
 
 
+# =============================================================================
+# Provider auth management (wafer auth ...)
+# =============================================================================
+
+provider_auth_app = typer.Typer(help="Manage API keys for cloud GPU providers")
+app.add_typer(provider_auth_app, name="auth")
+
+
+@provider_auth_app.command("login")
+def provider_auth_login(
+    provider: str = typer.Argument(
+        ...,
+        help="Provider name: runpod, digitalocean, or modal",
+    ),
+    api_key: str | None = typer.Option(
+        None,
+        "--api-key",
+        "-k",
+        help="API key (if not provided, reads from stdin)",
+    ),
+) -> None:
+    """Save API key for a cloud GPU provider.
+
+    Stores the key in ~/.wafer/auth.json. Environment variables
+    (e.g., WAFER_RUNPOD_API_KEY) take precedence over stored keys.
+
+    Examples:
+        wafer auth login runpod --api-key rp_xxx
+        wafer auth login digitalocean --api-key dop_v1_xxx
+        echo $API_KEY | wafer auth login runpod
+    """
+    import sys
+
+    from wafer_core.auth import PROVIDERS, save_api_key
+
+    # Validate provider
+    if provider not in PROVIDERS:
+        typer.echo(f"Error: Unknown provider '{provider}'", err=True)
+        typer.echo(f"Valid providers: {', '.join(PROVIDERS.keys())}", err=True)
+        raise typer.Exit(1)
+
+    # Get API key from option or stdin
+    if api_key is None:
+        if sys.stdin.isatty():
+            typer.echo(f"Enter API key for {PROVIDERS[provider]['display_name']}:")
+            api_key = typer.prompt("API key", hide_input=True)
+        else:
+            api_key = sys.stdin.read().strip()
+
+    if not api_key:
+        typer.echo("Error: No API key provided", err=True)
+        raise typer.Exit(1)
+
+    # Save the key
+    save_api_key(provider, api_key)
+    typer.echo(f"API key saved for {PROVIDERS[provider]['display_name']}")
+    typer.echo(f"Stored in: ~/.wafer/auth.json")
+
+
+@provider_auth_app.command("logout")
+def provider_auth_logout(
+    provider: str = typer.Argument(
+        ...,
+        help="Provider name: runpod, digitalocean, or modal",
+    ),
+) -> None:
+    """Remove stored API key for a cloud GPU provider.
+
+    Examples:
+        wafer auth logout runpod
+        wafer auth logout digitalocean
+    """
+    from wafer_core.auth import PROVIDERS, remove_api_key
+
+    # Validate provider
+    if provider not in PROVIDERS:
+        typer.echo(f"Error: Unknown provider '{provider}'", err=True)
+        typer.echo(f"Valid providers: {', '.join(PROVIDERS.keys())}", err=True)
+        raise typer.Exit(1)
+
+    if remove_api_key(provider):
+        typer.echo(f"API key removed for {PROVIDERS[provider]['display_name']}")
+    else:
+        typer.echo(f"No stored API key found for {PROVIDERS[provider]['display_name']}")
+
+
+@provider_auth_app.command("status")
+def provider_auth_status() -> None:
+    """Show authentication status for all cloud GPU providers.
+
+    Displays which providers have API keys configured and where
+    the keys are coming from (environment variable or auth.json).
+
+    Example:
+        wafer auth status
+    """
+    from wafer_core.auth import get_all_auth_status
+
+    statuses = get_all_auth_status()
+
+    typer.echo("Cloud GPU Provider Authentication Status")
+    typer.echo("=" * 45)
+
+    for status in statuses:
+        if status.is_authenticated:
+            source_str = f"({status.source})" if status.source else ""
+            typer.echo(
+                f"  {status.display_name}: ✓ {status.key_preview} {source_str}"
+            )
+        else:
+            typer.echo(f"  {status.display_name}: ✗ Not configured")
+            typer.echo(f"      Run: wafer auth login {status.provider}")
+            typer.echo(f"      Or set: {status.key_url}")
+
+    typer.echo("")
+    typer.echo("Note: Environment variables take precedence over stored keys.")
+
+
 @app.command(hidden=True)
 def run(
     command: str = typer.Argument(..., help="Command to run in Docker container"),
@@ -1289,13 +1422,25 @@ def evaluate(  # noqa: PLR0913
             --benchmark --defensive
 
     Subcommands:
-        make-template  Generate template files for this format
+        gpumode        Use GPUMode format (functional) - RECOMMENDED
         kernelbench    Use KernelBench format (ModelNew class)
+        make-template  Generate template files for this format (deprecated)
     """
     # If a subcommand is being invoked, skip the main evaluation logic
     if ctx.invoked_subcommand is not None:
         return
 
+    # Deprecation warning for bare evaluate
+    typer.echo(
+        "⚠️  Deprecation warning: 'wafer evaluate' will be removed in a future version.",
+        err=True,
+    )
+    typer.echo(
+        "   Use 'wafer evaluate gpumode' instead for the functional format.",
+        err=True,
+    )
+    typer.echo("", err=True)
+
     # Validate required args when running evaluation (not subcommands)
     missing_args = []
     if implementation is None:
@@ -1310,12 +1455,12 @@ def evaluate(  # noqa: PLR0913
         typer.echo(f"  Required: {', '.join(missing_args)}", err=True)
         typer.echo("", err=True)
         typer.echo(
-            "Usage: wafer evaluate --impl KERNEL.py --reference REF.py --test-cases TESTS.json",
+            "Usage: wafer evaluate gpumode --impl KERNEL.py --reference REF.py --test-cases TESTS.json",
             err=True,
         )
         typer.echo("", err=True)
-        typer.echo("Run 'wafer evaluate --help' for full options.", err=True)
-        typer.echo("Run 'wafer evaluate make-template DIR' to generate starter files.", err=True)
+        typer.echo("Run 'wafer evaluate gpumode --help' for full options.", err=True)
+        typer.echo("Run 'wafer evaluate gpumode download' to download problem sets.", err=True)
         raise typer.Exit(1)
 
     from .evaluate import EvaluateArgs, run_evaluate
@@ -1503,8 +1648,59 @@ def evaluate_make_template(
 # KernelBench format evaluation
 # =============================================================================
 
-# Path to KernelBench problems (relative to wafer root)
-KERNELBENCH_ROOT = Path(__file__).parent.parent.parent.parent / "research" / "KernelBench"
+
+def _get_kernelbench_root() -> Path | None:
+    """Get KernelBench problems root, preferring downloaded location."""
+    # First check downloaded location
+    downloaded = get_problems_path("kernelbench")
+    if downloaded is not None:
+        kb_root = downloaded / "KernelBench"
+        if kb_root.exists():
+            return kb_root
+        return downloaded
+
+    # Fall back to legacy location (for development)
+    legacy = Path(__file__).parent.parent.parent.parent / "research" / "KernelBench" / "KernelBench"
+    if legacy.exists():
+        return legacy
+
+    return None
+
+
+@kernelbench_app.command("download")
+def kernelbench_download(
+    force: bool = typer.Option(False, "--force", "-f", help="Re-download even if exists"),
+) -> None:
+    """Download KernelBench problems from GitHub.
+
+    Downloads the problem set to ~/.cache/wafer/problems/kernelbench/
+
+    Examples:
+        wafer evaluate kernelbench download
+        wafer evaluate kernelbench download --force  # Re-download
+    """
+    try:
+        path = download_problems("kernelbench", force=force, verbose=True)
+        typer.echo("")
+        typer.echo(f"Problems available at: {path}")
+        typer.echo("Run 'wafer evaluate kernelbench list-problems' to see available problems.")
+    except Exception as e:
+        typer.echo(f"Error downloading problems: {e}", err=True)
+        raise typer.Exit(1) from None
+
+
+@kernelbench_app.command("list-problems")
+def kernelbench_list_problems() -> None:
+    """List available KernelBench problems.
+
+    Examples:
+        wafer evaluate kernelbench list-problems
+    """
+    try:
+        list_problems_fn("kernelbench", verbose=True)
+    except ValueError as e:
+        typer.echo(str(e), err=True)
+        raise typer.Exit(1) from None
 
 
 @kernelbench_app.callback(invoke_without_command=True)
@@ -1530,6 +1726,10 @@ def kernelbench_evaluate(  # noqa: PLR0913
     ),
     benchmark: bool = typer.Option(False, "--benchmark", help="Run performance benchmarks"),
     profile: bool = typer.Option(False, "--profile", help="Enable profiling"),
+    inputs: Path | None = typer.Option(
+        None, "--inputs", help="Custom inputs file to override get_inputs()"
+    ),
+    seed: int = typer.Option(42, "--seed", help="Random seed for weight initialization"),
     defensive: bool = typer.Option(
         False, "--defensive", help="Enable defensive timing to detect evaluation hacking"
     ),
@@ -1594,6 +1794,8 @@ def kernelbench_evaluate(  # noqa: PLR0913
         target_name=target or "",
         benchmark=benchmark,
         profile=profile,
+        inputs=inputs,
+        seed=seed,
         defensive=defensive,
         sync_artifacts=sync_artifacts,
         gpu_id=gpu_id,
@@ -1655,6 +1857,13 @@ def kernelbench_make_template(
         # Overwrite existing
         wafer evaluate kernelbench make-template level1/1 --force
     """
+    # Get problems root (downloaded or legacy)
+    kb_root = _get_kernelbench_root()
+    if kb_root is None:
+        typer.echo("Error: KernelBench problems not found.", err=True)
+        typer.echo("Run 'wafer evaluate kernelbench download' to download problems.", err=True)
+        raise typer.Exit(1)
+
     # Parse problem ID
     parts = problem.split("/")
     if len(parts) != 2:
@@ -1666,10 +1875,10 @@ def kernelbench_make_template(
         level_str = f"level{level_str}"
 
     # Find the problem file
-    problem_dir = KERNELBENCH_ROOT / "KernelBench" / level_str
+    problem_dir = kb_root / level_str
     if not problem_dir.exists():
         typer.echo(f"Error: KernelBench level directory not found: {problem_dir}", err=True)
-        typer.echo(f"Make sure KernelBench is at: {KERNELBENCH_ROOT}", err=True)
+        typer.echo("Run 'wafer evaluate kernelbench download' to download problems.", err=True)
         raise typer.Exit(1)
 
     # Find matching problem file
@@ -1736,6 +1945,252 @@ def kernelbench_make_template(
     typer.echo(f"     wafer evaluate kernelbench --impl my_kernel.py --reference {output}")
 
 
+# =============================================================================
+# GPUMode format evaluation
+# =============================================================================
+
+
+@gpumode_app.command("download")
+def gpumode_download(
+    force: bool = typer.Option(False, "--force", "-f", help="Re-download even if exists"),
+) -> None:
+    """Download GPUMode reference kernels from GitHub.
+
+    Downloads the problem set to ~/.cache/wafer/problems/gpumode/
+
+    Examples:
+        wafer evaluate gpumode download
+        wafer evaluate gpumode download --force  # Re-download
+    """
+    try:
+        path = download_problems("gpumode", force=force, verbose=True)
+        typer.echo("")
+        typer.echo(f"Problems available at: {path}")
+        typer.echo("Run 'wafer evaluate gpumode list-problems' to see available problems.")
+    except Exception as e:
+        typer.echo(f"Error downloading problems: {e}", err=True)
+        raise typer.Exit(1) from None
+
+
+@gpumode_app.command("list-problems")
+def gpumode_list_problems() -> None:
+    """List available GPUMode problems.
+
+    Examples:
+        wafer evaluate gpumode list-problems
+    """
+    try:
+        list_problems_fn("gpumode", verbose=True)
+    except ValueError as e:
+        typer.echo(str(e), err=True)
+        raise typer.Exit(1) from None
+
+
+@gpumode_app.command("make-template")
+def gpumode_make_template(
+    problem: str = typer.Option(
+        ...,
+        "--problem",
+        "-p",
+        help="Problem ID (e.g., 'pmpp/vectoradd_py' or 'amd/fp8-mm')",
+    ),
+    output: Path = typer.Option(
+        None, "--output", "-o", help="Output directory (default: ./<problem_name>/)"
+    ),
+    force: bool = typer.Option(False, "--force", "-f", help="Overwrite existing files"),
+) -> None:
+    """Extract a GPUMode problem as template files.
+
+    Creates a directory with reference.py, task.yml, and other problem files.
+    You then create kernel.py with your custom_kernel implementation.
+
+    Examples:
+        # Extract pmpp vectoradd problem
+        wafer evaluate gpumode make-template --problem pmpp/vectoradd_py
+
+        # Extract to specific directory
+        wafer evaluate gpumode make-template --problem pmpp/vectoradd_py --output ./my-kernel/
+    """
+    import shutil
+
+    # Get problem path
+    problem_path = get_problem_path("gpumode", problem)
+    if problem_path is None:
+        # Check if problems are downloaded
+        if get_problems_path("gpumode") is None:
+            typer.echo("Error: GPUMode problems not downloaded.", err=True)
+            typer.echo("Run 'wafer evaluate gpumode download' first.", err=True)
+        else:
+            typer.echo(f"Error: Problem '{problem}' not found.", err=True)
+            typer.echo(
+                "Run 'wafer evaluate gpumode list-problems' to see available problems.", err=True
+            )
+        raise typer.Exit(1)
+
+    # Determine output path
+    if output is None:
+        output = Path.cwd() / problem.replace("/", "_")
+
+    output = output.resolve()
+
+    # Check if exists
+    if output.exists() and not force:
+        typer.echo(f"Error: {output} already exists. Use --force to overwrite.", err=True)
+        raise typer.Exit(1)
+
+    # Copy the problem directory
+    if output.exists():
+        shutil.rmtree(output)
+    shutil.copytree(problem_path, output)
+
+    typer.echo(f"Created {output}/")
+    typer.echo("")
+    typer.echo("Contents:")
+    for f in sorted(output.iterdir()):
+        if not f.name.startswith("."):
+            typer.echo(f"  {f.name}")
+    typer.echo("")
+    typer.echo("Next steps:")
+    typer.echo("  1. Read reference.py to understand the kernel interface")
+    typer.echo("  2. Create kernel.py with your custom_kernel implementation:")
+    typer.echo("")
+    typer.echo("     def custom_kernel(data):")
+    typer.echo("         # Your optimized implementation")
+    typer.echo("         ...")
+    typer.echo("")
+    typer.echo("  3. Run evaluation:")
+    typer.echo(
+        f"     wafer evaluate gpumode --impl {output}/kernel.py --reference {output}/reference.py \\"
+    )
+    typer.echo(f"         --test-cases {output}/test_cases.json --target <target>")
+
+
+@gpumode_app.callback(invoke_without_command=True)
+def gpumode_evaluate(  # noqa: PLR0913
+    ctx: typer.Context,
+    implementation: Path | None = typer.Option(
+        None, "--impl", "-i", help="Path to implementation kernel file"
+    ),
+    reference: Path | None = typer.Option(
+        None, "--reference", help="Path to reference kernel file"
+    ),
+    test_cases: Path | None = typer.Option(
+        None, "--test-cases", help="Path to test cases JSON file"
+    ),
+    target: str | None = typer.Option(
+        None,
+        "--target",
+        "-t",
+        help="GPU target name. See 'wafer config targets list' for available targets.",
+        autocompletion=complete_target_name,
+    ),
+    benchmark: bool = typer.Option(False, "--benchmark", help="Run performance benchmarks"),
+    profile: bool = typer.Option(False, "--profile", help="Enable profiling"),
+    defensive: bool = typer.Option(
+        False, "--defensive", help="Enable defensive timing to detect evaluation hacking"
+    ),
+    sync_artifacts: bool = typer.Option(
+        True, "--sync-artifacts/--no-sync-artifacts", help="Download artifacts"
+    ),
+    gpu_id: int | None = typer.Option(None, "--gpu-id", help="Override GPU ID"),
+) -> None:
+    """Run kernel evaluation in GPUMode format (functional).
+
+    This format expects:
+    - Implementation: Python file with `custom_kernel(inputs)` function
+    - Reference: Python file with `ref_kernel(inputs)` and `generate_input(**kwargs)` functions
+    - Test cases: JSON file with test parameters
+
+    Examples:
+        # Basic correctness check
+        wafer evaluate gpumode --impl kernel.py --reference ref.py --test-cases tests.json
+
+        # With benchmarking
+        wafer evaluate gpumode --impl kernel.py --reference ref.py --test-cases tests.json \\
+            --target vultr-b200 --benchmark
+
+    Subcommands:
+        download       Download GPUMode problems from GitHub
+        list-problems  List available problems
+        make-template  Extract a problem as template files
+    """
+    # If a subcommand is being invoked, skip the main evaluation logic
+    if ctx.invoked_subcommand is not None:
+        return
+
+    # Validate required args when running evaluation (not subcommands)
+    missing_args = []
+    if implementation is None:
+        missing_args.append("--impl/-i")
+    if reference is None:
+        missing_args.append("--reference")
+    if test_cases is None:
+        missing_args.append("--test-cases")
+
+    if missing_args:
+        typer.echo("Error: Missing required arguments", err=True)
+        typer.echo(f"  Required: {', '.join(missing_args)}", err=True)
+        typer.echo("", err=True)
+        typer.echo(
+            "Usage: wafer evaluate gpumode --impl KERNEL.py --reference REF.py --test-cases TESTS.json",
+            err=True,
+        )
+        typer.echo("", err=True)
+        typer.echo("Run 'wafer evaluate gpumode --help' for full options.", err=True)
+        typer.echo("Run 'wafer evaluate gpumode download' to download problem sets.", err=True)
+        raise typer.Exit(1)
+
+    # Reuse the existing evaluate logic (same format)
+    from .evaluate import EvaluateArgs, run_evaluate
+
+    args = EvaluateArgs(
+        implementation=implementation,
+        reference=reference,
+        test_cases=test_cases,
+        target_name=target or "",
+        benchmark=benchmark,
+        profile=profile,
+        defensive=defensive,
+        sync_artifacts=sync_artifacts,
+        gpu_id=gpu_id,
+    )
+
+    try:
+        import trio_asyncio
+
+        result = trio_asyncio.run(run_evaluate, args)
+    except KeyboardInterrupt:
+        typer.echo("\nInterrupted by user", err=True)
+        raise typer.Exit(130) from None
+    except Exception as e:
+        if hasattr(e, "exceptions") and e.exceptions:
+            for exc in e.exceptions:
+                typer.echo(f"Error: {type(exc).__name__}: {exc}", err=True)
+        else:
+            typer.echo(f"Error: {e}", err=True)
+        raise typer.Exit(1) from None
+
+    # Print results
+    if result.success:
+        typer.echo("")
+        typer.echo("=" * 60)
+        status = "PASS" if result.all_correct else "FAIL"
+        typer.echo(f"Result: {status}")
+        score_pct = f"{result.correctness_score:.1%}"
+        typer.echo(f"Correctness: {result.passed_tests}/{result.total_tests} ({score_pct})")
+        if result.geomean_speedup > 0:
+            typer.echo(f"Speedup: {result.geomean_speedup:.2f}x")
+        if result.artifact_path:
+            typer.echo(f"Artifacts: {result.artifact_path}")
+        typer.echo("=" * 60)
+
+        if not result.all_correct:
+            raise typer.Exit(1)
+    else:
+        typer.echo(f"Error: {result.error_message}", err=True)
+        raise typer.Exit(1)
+
+
 # =============================================================================
 # Push and Remote-Run commands
 # =============================================================================
@@ -1867,7 +2322,7 @@ def _run_direct_mode(
         typer.echo(f"Uploading {upload_dir.name}...")
         try:
             push_result = push_direct(upload_dir, target)
-            workspace_name = push_result.workspace_path
+            workspace_name = push_result.workspace_name
             typer.echo(f"Uploaded {len(push_result.files_uploaded)} files")
         except Exception as e:
             typer.echo(f"Error uploading: {e}", err=True)
@@ -2040,7 +2495,10 @@ def login(
         None, "--token", "-t", help="Access token (skip browser OAuth)"
     ),
     port: int | None = typer.Option(
-        None, "--port", "-p", help="Port for OAuth callback server (default: 8765 for SSH, random for local)"
+        None,
+        "--port",
+        "-p",
+        help="Port for OAuth callback server (default: 8765 for SSH, random for local)",
     ),
 ) -> None:
     """Authenticate CLI with wafer-api via GitHub OAuth.
@@ -2142,9 +2600,8 @@ def login(
 @app.command("logout")
 def logout() -> None:
     """Remove stored credentials."""
-    from .auth import clear_credentials
-
     from . import analytics
+    from .auth import clear_credentials
 
     # Track logout event first (while credentials still exist for user identification)
     # Note: track_logout() handles the case where user is not logged in
@@ -3115,7 +3572,9 @@ def billing_usage(
 @billing_app.command("topup")
 def billing_topup(
     amount: int = typer.Argument(25, help="Amount in dollars ($10-$500)"),
-    no_browser: bool = typer.Option(False, "--no-browser", help="Print URL instead of opening browser"),
+    no_browser: bool = typer.Option(
+        False, "--no-browser", help="Print URL instead of opening browser"
+    ),
 ) -> None:
     """Add credits to your account.
 
@@ -3161,7 +3620,9 @@ def billing_topup(
 
 @billing_app.command("portal")
 def billing_portal(
-    no_browser: bool = typer.Option(False, "--no-browser", help="Print URL instead of opening browser"),
+    no_browser: bool = typer.Option(
+        False, "--no-browser", help="Print URL instead of opening browser"
+    ),
 ) -> None:
     """Open Stripe billing portal.
 
@@ -4437,8 +4898,8 @@ def _setup_wafer_core_env() -> None:
     - WAFER_API_URL: If already set, uses that instead of config
     - WAFER_AUTH_TOKEN: If already set, uses that instead of cached token
     """
-    from .global_config import get_api_url
     from .auth import get_valid_token
+    from .global_config import get_api_url
 
     # Set API URL (get_api_url already respects WAFER_API_URL env var)
     os.environ["WAFER_API_URL"] = get_api_url()
@@ -4742,8 +5203,8 @@ def capture_command(  # noqa: PLR0915
     import os
     import tomllib
 
-    from .global_config import get_api_url
     from .auth import get_valid_token
+    from .global_config import get_api_url
 
     # Set environment variables for wafer-core BEFORE importing it
     # wafer-core backend.py reads WAFER_API_URL and WAFER_AUTH_TOKEN from env
@@ -4947,8 +5408,8 @@ def capture_list_command(
     """
     import os
 
-    from .global_config import get_api_url
     from .auth import get_valid_token
+    from .global_config import get_api_url
 
     # Set environment variables for wafer-core BEFORE importing it
     os.environ["WAFER_API_URL"] = get_api_url()