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

wafer/cli.py

@@ -18,9 +18,11 @@ Setup:
   config                CLI configuration and local GPU targets
 """
 
+import atexit
 import json
 import os
 import sys
+import time
 from pathlib import Path
 
 import trio
@@ -28,12 +30,112 @@ 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",
     no_args_is_help=True,
 )
 
+# =============================================================================
+# Analytics tracking
+# =============================================================================
+
+# Track command start time for duration calculation
+_command_start_time: float | None = None
+# Track command outcome (defaults to failure, set to success on clean exit)
+_command_outcome: str = "failure"
+
+
+def _get_command_path(ctx: typer.Context) -> tuple[str, str | None]:
+    """Extract command and subcommand from Typer context.
+
+    Returns:
+        Tuple of (command, subcommand). subcommand may be None.
+    """
+    # Build command path from invoked subcommand chain
+    invoked = ctx.invoked_subcommand
+    info_name = ctx.info_name or ""
+
+    # Get parent command if exists
+    parent_cmd = None
+    if ctx.parent and ctx.parent.info_name and ctx.parent.info_name != "wafer":
+        parent_cmd = ctx.parent.info_name
+
+    if parent_cmd:
+        return parent_cmd, info_name
+    return info_name or "unknown", invoked
+
+
+def _mark_command_success() -> None:
+    """Mark the current command as successful.
+
+    Call this at the end of successful command execution.
+    Commands that raise typer.Exit(1) or exceptions will remain marked as failures.
+    """
+    global _command_outcome
+    _command_outcome = "success"
+
+
+@app.callback()
+def main_callback(ctx: typer.Context) -> None:
+    """Initialize analytics and track command execution."""
+    global _command_start_time, _command_outcome
+    _command_start_time = time.time()
+    _command_outcome = "success"  # Default to success, mark failure on exceptions
+
+    # Initialize analytics (lazy import to avoid slowing down --help)
+    from . import analytics
+
+    analytics.init_analytics()
+
+    # Install exception hook to catch SystemExit and mark failures
+    original_excepthook = sys.excepthook
+
+    def custom_excepthook(exc_type, exc_value, exc_traceback):
+        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
+            if exit_code != 0 and exit_code is not None:
+                _command_outcome = "failure"
+        else:
+            _command_outcome = "failure"
+        # Call original excepthook
+        original_excepthook(exc_type, exc_value, exc_traceback)
+
+    sys.excepthook = custom_excepthook
+
+    # Register tracking at exit to capture command outcome
+    def track_on_exit() -> None:
+        command, subcommand = _get_command_path(ctx)
+
+        # Skip tracking for --help and --version
+        if ctx.resilient_parsing:
+            return
+
+        # Calculate duration
+        duration_ms = None
+        if _command_start_time is not None:
+            duration_ms = int((time.time() - _command_start_time) * 1000)
+
+        # Track the command execution with the recorded outcome
+        analytics.track_command(
+            command=command,
+            subcommand=subcommand,
+            outcome=_command_outcome,
+            duration_ms=duration_ms,
+        )
+
+    atexit.register(track_on_exit)
+
 
 # =============================================================================
 # Autocompletion helpers
@@ -106,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)
 # =============================================================================
@@ -302,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"),
@@ -1195,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:
@@ -1216,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
@@ -1409,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)
@@ -1436,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"
     ),
@@ -1500,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,
@@ -1561,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:
@@ -1572,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
@@ -1642,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
 # =============================================================================
@@ -1773,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)
@@ -1945,17 +2494,34 @@ def login(
     token: str | None = typer.Option(
         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:
     """Authenticate CLI with wafer-api via GitHub OAuth.
 
     Opens browser for GitHub authentication. Use --token to skip browser.
     Uses the API environment from config (see 'wafer config show').
 
+    SSH Users:
+    - Automatically uses port 8765 (just set up port forwarding once)
+    - On local machine: ssh -L 8765:localhost:8765 user@host
+    - On remote machine: wafer login
+    - Browser opens locally, redirect works through tunnel
+
+    Manual token option:
+    - Visit auth.wafer.ai, authenticate, copy token from URL
+    - Run: wafer login --token <paste-token>
+
     Examples:
-        wafer login              # opens browser for GitHub OAuth
-        wafer login --token xyz  # use existing token
+        wafer login                    # auto-detects SSH, uses appropriate port
+        wafer login --port 9000        # override port
+        wafer login --token xyz        # manual token (no browser)
 
-        # To login to a different environment:
+        # Change environment:
         wafer config set api.environment staging
         wafer login
     """
@@ -1971,11 +2537,21 @@ def login(
     typer.echo(f"Auth: {get_supabase_url()}")
     typer.echo("")
 
+    # Auto-detect SSH and use fixed port
+    if port is None:
+        is_ssh = bool(os.environ.get("SSH_CONNECTION") or os.environ.get("SSH_CLIENT"))
+        if is_ssh:
+            port = 8765
+            typer.echo("🔒 SSH session detected - using port 8765 for OAuth callback")
+            typer.echo("   Make sure you have port forwarding set up:")
+            typer.echo("   ssh -L 8765:localhost:8765 user@host")
+            typer.echo("")
+
     # Browser OAuth if no token provided
     refresh_token = None
     if token is None:
         try:
-            token, refresh_token = browser_login()
+            token, refresh_token = browser_login(port=port)
         except TimeoutError as e:
             typer.echo(f"Error: {e}", err=True)
             raise typer.Exit(1) from None
@@ -2009,6 +2585,11 @@ def login(
     # Save credentials (with refresh token if available)
     save_credentials(token, refresh_token, user_info.email)
 
+    # Track login event with analytics
+    from . import analytics
+
+    analytics.track_login(user_info.user_id, user_info.email)
+
     if user_info.email:
         typer.echo(f"Logged in as {user_info.email}")
     else:
@@ -2019,8 +2600,14 @@ def login(
 @app.command("logout")
 def logout() -> None:
     """Remove stored 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
+    analytics.track_logout()
+
+    # Clear credentials and report result
     if clear_credentials():
         typer.echo("Logged out. Credentials removed.")
     else:
@@ -2985,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.
 
@@ -3031,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.
 
@@ -4298,6 +4889,29 @@ autotuner_app = typer.Typer(help="Hyperparameter sweep for performance engineeri
 app.add_typer(autotuner_app, name="autotuner", hidden=True)
 
 
+def _setup_wafer_core_env() -> None:
+    """Set environment variables for wafer-core to use.
+
+    Call this before using any wafer-core functions that need API access.
+
+    Respects explicit environment variable overrides:
+    - WAFER_API_URL: If already set, uses that instead of config
+    - WAFER_AUTH_TOKEN: If already set, uses that instead of cached token
+    """
+    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()
+
+    # Only set auth token if not explicitly provided in environment
+    # This allows CI/service accounts to override with their own tokens
+    if "WAFER_AUTH_TOKEN" not in os.environ:
+        token = get_valid_token()
+        if token:
+            os.environ["WAFER_AUTH_TOKEN"] = token
+
+
 @autotuner_app.command("list")
 def autotuner_list(
     show_all: bool = typer.Option(
@@ -4313,6 +4927,7 @@ def autotuner_list(
         wafer autotuner list
         wafer autotuner list --all
     """
+    _setup_wafer_core_env()
     from .autotuner import list_command
 
     try:
@@ -4353,6 +4968,7 @@ def autotuner_delete(
         wafer autotuner delete --all --status pending
         wafer autotuner delete --all --status failed --yes
     """
+    _setup_wafer_core_env()
     from .autotuner import delete_all_command, delete_command
 
     # Validate arguments
@@ -4419,6 +5035,7 @@ def autotuner_run(
         wafer autotuner run --resume <sweep-id>
         wafer autotuner run --resume <sweep-id> --parallel 8
     """
+    _setup_wafer_core_env()
     from .autotuner import run_sweep_command
 
     # Validate arguments
@@ -4583,8 +5200,23 @@ def capture_command(  # noqa: PLR0915
         wafer capture grid-search "python train.py --lr {LR} --bs {BS}" --sweep "LR=0.001,0.01,0.1" --sweep "BS=16,32"
     """
     import itertools
+    import os
     import tomllib
 
+    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
+    os.environ["WAFER_API_URL"] = get_api_url()
+
+    # Only set auth token if not explicitly provided in environment
+    # This allows CI/service accounts to override with their own tokens
+    if "WAFER_AUTH_TOKEN" not in os.environ:
+        token = get_valid_token()
+        if token:
+            os.environ["WAFER_AUTH_TOKEN"] = token
+
     import trio
     from wafer_core.tools.capture_tool import (  # pragma: no cover
         CaptureConfig,
@@ -4774,6 +5406,20 @@ def capture_list_command(
         # Pagination
         wafer capture-list --limit 20 --offset 20
     """
+    import os
+
+    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()
+
+    # Only set auth token if not explicitly provided in environment
+    # This allows CI/service accounts to override with their own tokens
+    if "WAFER_AUTH_TOKEN" not in os.environ:
+        token = get_valid_token()
+        if token:
+            os.environ["WAFER_AUTH_TOKEN"] = token
 
     import trio
     from wafer_core.utils.backend import list_captures  # pragma: no cover