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

wafer/GUIDE.md

@@ -1,313 +1,107 @@
 # Wafer CLI Guide
 
-GPU development toolkit for LLM coding agents.
+GPU development primitives for LLM agents.
 
-## Quick Start
+## Quick Start: Cloud GPU (No Setup)
 
-Choose your path based on what GPU access you have:
-
-### Path A: I have my own GPU (SSH access)
+Run code on cloud GPUs instantly with workspaces:
 
 ```bash
-# 1. Set up your target
-wafer config targets init ssh \
-  --name my-gpu \
-  --host user@192.168.1.100:22 \
-  --gpu-type H100
-
-# 2. Test it works
-wafer evaluate make-template ./test-kernel
-wafer evaluate \
-  --impl ./test-kernel/kernel.py \
-  --reference ./test-kernel/reference.py \
-  --test-cases ./test-kernel/test_cases.json \
-  --target my-gpu
+wafer login                              # One-time auth
+wafer workspaces create dev --gpu B200   # Create workspace
+wafer workspaces exec dev -- python -c "import torch; print(torch.cuda.get_device_name(0))"
+wafer workspaces sync dev ./my-project   # Sync files
+wafer workspaces exec dev -- python train.py
 ```
 
-### Path B: I have RunPod/DigitalOcean API keys
-
-```bash
-# 1. Set your API key
-export WAFER_RUNPOD_API_KEY=your_key_here
-# or: export WAFER_AMD_DIGITALOCEAN_API_KEY=your_key_here
-
-# 2. Create a target (auto-provisions GPUs on demand)
-wafer config targets init runpod --gpu MI300X
-# or: wafer config targets init digitalocean
+## Documentation Lookup
 
-# 3. Run evaluation (provisions GPU, runs, cleans up)
-wafer evaluate \
-  --impl kernel.py \
-  --reference ref.py \
-  --test-cases tests.json \
-  --target runpod-mi300x
-```
-
-### Path C: I just want to ask questions / analyze traces
+Answer GPU programming questions from indexed documentation.
 
 ```bash
-# No GPU needed for these workflows
-
-# Download documentation
-wafer corpus download cuda
-
-# Ask questions
-wafer wevin -t ask-docs --corpus cuda "What is warp divergence?"
-
-# Analyze a PyTorch trace
-wafer nvidia perfetto query trace.json \
-  "SELECT name, dur/1e6 as ms FROM slice WHERE cat='kernel' ORDER BY dur DESC LIMIT 10"
-```
-
-## Core Commands
-
-| Command | Description |
-|---------|-------------|
-| `wafer wevin` | GPU programming assistant (delegates to specialized agent) |
-| `wafer evaluate` | Test kernel correctness and performance |
-| `wafer corpus` | Download/manage GPU documentation |
-| `wafer workspaces` | Cloud GPU environments |
-
-## Profiling Tools
-
-```bash
-# NVIDIA tools (under wafer nvidia)
-wafer nvidia ncu analyze profile.ncu-rep      # Nsight Compute analysis
-wafer nvidia nsys analyze profile.nsys-rep    # Nsight Systems analysis
-wafer nvidia perfetto query trace.json "SQL"  # Perfetto/PyTorch traces
-wafer nvidia tracelens report trace.json      # Performance reports
-
-# AMD tools (under wafer amd)
-wafer amd isa analyze kernel.co               # ISA analysis
-wafer amd rocprof-compute profile -- ./app    # ROCm profiling
-```
-
-## Common Workflows
-
-### 1. Research GPU Concepts
-
-```bash
-# Download documentation (one-time)
+# Download corpus (one-time)
 wafer corpus download cuda
 wafer corpus download cutlass
+wafer corpus download hip
 
-# Ask questions
-wafer wevin -t ask-docs --corpus cuda "How do I optimize shared memory usage?"
-wafer wevin -t ask-docs --corpus cutlass "What is a TiledMma?"
+# Query documentation
+wafer agent -t ask-docs --corpus cuda "What is warp divergence?"
+wafer agent -t ask-docs --corpus cutlass "What is a TiledMma?"
 ```
 
-### 2. Analyze a PyTorch Trace
+## Trace Analysis
+
+Analyze performance traces from NCU, NSYS, or PyTorch profiler.
 
 ```bash
-# PyTorch profiler outputs Chrome trace format JSON
-# Wafer's perfetto tool can analyze these
+# AI-assisted analysis
+wafer agent -t trace-analyze --args trace=./profile.ncu-rep "Why is this kernel slow?"
+wafer agent -t trace-analyze --args trace=./trace.json "What's the bottleneck?"
 
-# List available event categories
+# Direct trace queries (PyTorch/Perfetto JSON)
 wafer nvidia perfetto tables trace.json
-
-# Find slowest GPU kernels
 wafer nvidia perfetto query trace.json \
-  "SELECT name, dur/1e6 as ms FROM slice WHERE cat='kernel' ORDER BY dur DESC LIMIT 20"
-
-# Get kernel time breakdown
-wafer nvidia perfetto query trace.json \
-  "SELECT name, SUM(dur)/1e6 as total_ms, COUNT(*) as calls
-   FROM slice WHERE cat='kernel'
-   GROUP BY name ORDER BY total_ms DESC LIMIT 20"
-
-# Or use the trace-analyze template for AI-assisted analysis
-wafer wevin -t trace-analyze --args trace=./trace.json "What's the bottleneck?"
-```
-
-### 3. Analyze NCU Profiles
+  "SELECT name, dur/1e6 as ms FROM slice WHERE cat='kernel' ORDER BY dur DESC LIMIT 10"
 
-```bash
-# Basic analysis with recommendations
+# NCU/NSYS analysis
 wafer nvidia ncu analyze profile.ncu-rep
-
-# JSON output for programmatic use
-wafer nvidia ncu analyze profile.ncu-rep --json
-
-# With SASS source correlation
-wafer nvidia ncu analyze profile.ncu-rep --include-source
+wafer nvidia nsys analyze profile.nsys-rep
 ```
 
-### 4. Evaluate Kernel Performance
+## Kernel Evaluation
 
-```bash
-# Generate template files
-wafer evaluate make-template ./my-kernel
+Test kernel correctness and measure speedup against a reference.
 
-# Edit the generated files:
-#   kernel.py    - Your optimized implementation
-#   reference.py - Ground truth + input generator
-#   test_cases.json - Test parameters
+```bash
+# Using workspaces (no target setup required):
+wafer workspaces create dev --gpu B200
+wafer workspaces exec --sync ./my-kernel dev -- python test_kernel.py
 
-# Run evaluation
+# Or using configured targets (for your own hardware):
+wafer evaluate make-template ./my-kernel
 wafer evaluate \
   --impl ./my-kernel/kernel.py \
   --reference ./my-kernel/reference.py \
   --test-cases ./my-kernel/test_cases.json \
-  --target vultr-b200
-
-# With profiling
-wafer evaluate ... --profile
+  --target <target-name>
 ```
 
-### 5. Optimize a Kernel (AI-assisted)
+For target setup, see `wafer config targets --help`.
+
+## Kernel Optimization (AI-assisted)
+
+Iteratively optimize a kernel with evaluation feedback.
 
 ```bash
-# Use the optimize-kernel template
-wafer wevin -t optimize-kernel \
+wafer agent -t optimize-kernel \
   --args kernel=./my_kernel.cu \
   --args target=H100 \
   "Optimize this GEMM for memory bandwidth"
 ```
 
-## Configuration
+## Workspaces
 
-### Targets (GPU Access)
+Cloud GPU environments with no setup required.
 
-Targets define how to access GPUs. Three types:
-
-**1. SSH (your own hardware)**
-```bash
-wafer config targets init ssh \
-  --name my-gpu \
-  --host user@hostname:22 \
-  --gpu-type H100 \
-  --ncu  # if NCU profiling is available
-```
-
-**2. RunPod (on-demand cloud GPUs)**
 ```bash
-export WAFER_RUNPOD_API_KEY=your_key  # from runpod.io/console/user/settings
-wafer config targets init runpod --gpu MI300X  # or H100, A100
+wafer workspaces create dev --gpu B200     # Create
+wafer workspaces list                       # List all
+wafer workspaces sync dev ./project         # Sync files
+wafer workspaces exec dev -- ./run.sh       # Run commands
+wafer workspaces ssh dev                    # Interactive SSH
+wafer workspaces delete dev                 # Cleanup
 ```
 
-**3. DigitalOcean (AMD MI300X)**
-```bash
-export WAFER_AMD_DIGITALOCEAN_API_KEY=your_key
-wafer config targets init digitalocean
-```
+See `wafer workspaces --help` for details.
 
-**Managing targets:**
-```bash
-wafer config targets list              # List all targets
-wafer config targets show my-gpu       # Show details
-wafer config targets default my-gpu    # Set default
-wafer config targets remove my-gpu     # Delete
-```
+## Command Reference
 
-**Advanced: Manual TOML** (for custom configurations)
 ```bash
-wafer config targets add ~/my-target.toml
-```
-
-Example TOML:
-```toml
-name = "my-gpu"
-type = "baremetal"
-ssh_target = "user@hostname:22"
-ssh_key = "~/.ssh/id_ed25519"
-gpu_ids = [0]
-gpu_type = "H100"
-compute_capability = "9.0"
-ncu_available = true
-docker_image = "nvcr.io/nvidia/pytorch:25.01-py3"
-```
-
-### Workspaces (Cloud GPUs)
-
-Workspaces provide on-demand cloud GPU access.
-
-```bash
-# List workspaces
-wafer workspaces list
-
-# Create a workspace
-wafer workspaces create my-workspace --gpu H100
-
-# Get SSH credentials
-wafer workspaces attach <workspace-id>
-
-# Delete when done
-wafer workspaces delete <workspace-id>
-```
-
-### CLI Configuration
-
-```bash
-# Initialize config
-wafer config init
-
-# Show current config
-wafer config show
-
-# Set API environment
-wafer config set api.environment prod
-```
-
-## Templates
-
-Wafer includes templates for common tasks:
-
-| Template | Usage |
-|----------|-------|
-| `ask-docs` | Answer GPU programming questions from documentation |
-| `trace-analyze` | Analyze performance traces (NCU, NSYS, PyTorch) |
-| `optimize-kernel` | Iteratively optimize a kernel implementation |
-
-```bash
-# Use a template
-wafer wevin -t <template-name> [--args key=value] "your prompt"
-
-# Examples
-wafer wevin -t ask-docs --corpus cuda "What is occupancy?"
-wafer wevin -t trace-analyze --args trace=./profile.ncu-rep "Why is this slow?"
-wafer wevin -t optimize-kernel --args kernel=./gemm.cu "Make this 2x faster"
-```
-
-## Getting Help
-
-```bash
-# General help
-wafer --help
-
-# Command-specific help
-wafer evaluate --help
-wafer nvidia ncu --help
-wafer wevin --help
-
-# This guide
-wafer guide
-```
-
-## Quick Reference
-
-```bash
-# Authentication
-wafer login                    # Login via GitHub OAuth
-wafer logout                   # Remove credentials
-wafer whoami                   # Show current user
-wafer whoami --verify          # Verify token is valid
-
-# Documentation
-wafer corpus list              # Show available/downloaded corpora
-wafer corpus download cuda     # Download CUDA docs
-wafer corpus path cuda         # Get path to corpus
-
-# Analysis
-wafer nvidia ncu analyze *.ncu-rep     # NCU profile analysis
-wafer nvidia nsys analyze *.nsys-rep   # NSYS timeline analysis
-wafer nvidia perfetto query f.json "SQL"  # Perfetto SQL queries
-
-# Evaluation
-wafer evaluate make-template ./dir     # Generate template files
-wafer evaluate --impl ... --reference ... --test-cases ...
-
-# AI Assistant
-wafer wevin "your question"            # Interactive assistant
-wafer wevin -t ask-docs --corpus cuda "question"  # Docs lookup
-wafer wevin -s "single turn question"  # Non-interactive
+wafer corpus list|download|path   # Manage documentation corpora
+wafer workspaces                  # Cloud GPU environments (no setup)
+wafer evaluate                    # Test kernel correctness/performance
+wafer nvidia ncu|nsys|perfetto    # NVIDIA profiling tools
+wafer amd isa|rocprof-compute     # AMD profiling tools
+wafer agent -t <template>         # AI-assisted workflows
+wafer config targets              # Configure your own GPU targets
 ```

wafer/analytics.py

@@ -0,0 +1,307 @@
+"""PostHog analytics for Wafer CLI.
+
+Tracks CLI command usage and user activity for product analytics.
+Mirrors the analytics implementation in apps/wevin-extension/src/services/analytics.ts.
+
+Usage:
+    from .analytics import track_command, identify_user, shutdown_analytics
+
+    # Track a command execution
+    track_command("evaluate", {"subcommand": "kernelbench", "outcome": "success"})
+
+    # Identify user after login
+    identify_user("user-id", "user@example.com")
+"""
+
+import atexit
+import platform
+import sys
+import uuid
+from pathlib import Path
+from typing import Any
+
+# PostHog configuration - same as wevin-extension
+POSTHOG_API_KEY = "phc_9eDjkY72ud9o4l1mA1Gr1dnRT1yx71rP3XY9z66teFh"
+POSTHOG_HOST = "https://us.i.posthog.com"
+
+# Anonymous ID storage
+ANONYMOUS_ID_FILE = Path.home() / ".wafer" / ".analytics_id"
+
+# Global state
+_posthog_client: Any = None
+_distinct_id: str | None = None
+_initialized: bool = False
+
+
+def _get_anonymous_id() -> str:
+    """Get or create anonymous ID for users who aren't logged in."""
+    if ANONYMOUS_ID_FILE.exists():
+        return ANONYMOUS_ID_FILE.read_text().strip()
+
+    # Generate new anonymous ID
+    anonymous_id = f"anon_{uuid.uuid4().hex}"
+    ANONYMOUS_ID_FILE.parent.mkdir(parents=True, exist_ok=True)
+    ANONYMOUS_ID_FILE.write_text(anonymous_id)
+    return anonymous_id
+
+
+def _get_user_id_from_credentials() -> tuple[str | None, str | None]:
+    """Get user ID and email from stored credentials.
+
+    Returns:
+        Tuple of (user_id, email), both may be None if not logged in.
+    """
+    # Import here to avoid circular imports
+    from .auth import load_credentials, verify_token
+
+    creds = load_credentials()
+    if not creds:
+        return None, None
+
+    # Try to get user info from token
+    try:
+        user_info = verify_token(creds.access_token)
+        return user_info.user_id, user_info.email or creds.email
+    except Exception:
+        # Token verification failed, use email from credentials if available
+        return None, creds.email
+
+
+def _is_analytics_enabled() -> bool:
+    """Check if analytics is enabled via preferences.
+
+    Returns True by default, respects user preference in config.
+    """
+    from .global_config import get_preferences
+
+    try:
+        prefs = get_preferences()
+        return getattr(prefs, "analytics_enabled", True)
+    except Exception:
+        # Default to enabled if we can't read preferences
+        return True
+
+
+def init_analytics() -> bool:
+    """Initialize PostHog client.
+
+    Returns:
+        True if initialization succeeded, False otherwise.
+    """
+    global _posthog_client, _distinct_id, _initialized
+
+    if _initialized:
+        return _posthog_client is not None
+
+    _initialized = True
+
+    # Check if analytics is enabled
+    if not _is_analytics_enabled():
+        return False
+
+    try:
+        from posthog import Posthog
+
+        _posthog_client = Posthog(
+            api_key=POSTHOG_API_KEY,
+            host=POSTHOG_HOST,
+            # Flush immediately for CLI - commands are short-lived
+            flush_at=1,
+            flush_interval=1,
+            # Disable debug logging
+            debug=False,
+        )
+
+        # Set up distinct ID - prefer authenticated user, fall back to anonymous
+        user_id, email = _get_user_id_from_credentials()
+        if user_id:
+            _distinct_id = user_id
+            # Identify the user with their email
+            if email:
+                _posthog_client.identify(
+                    distinct_id=user_id,
+                    properties={
+                        "email": email,
+                        "auth_provider": "github",
+                    },
+                )
+        else:
+            _distinct_id = _get_anonymous_id()
+
+        # Register shutdown handler to flush events
+        atexit.register(shutdown_analytics)
+
+        return True
+
+    except ImportError:
+        # PostHog not installed - analytics disabled
+        return False
+    except Exception:
+        # Any other error - fail silently, don't break CLI
+        return False
+
+
+def shutdown_analytics() -> None:
+    """Shutdown PostHog client and flush pending events."""
+    global _posthog_client
+
+    if _posthog_client is not None:
+        try:
+            _posthog_client.flush()
+            _posthog_client.shutdown()
+        except Exception:
+            pass  # Fail silently on shutdown
+        _posthog_client = None
+
+
+def identify_user(user_id: str, email: str | None = None) -> None:
+    """Identify a user after login.
+
+    Args:
+        user_id: Supabase user ID
+        email: User's email address
+    """
+    global _distinct_id
+
+    if not init_analytics():
+        return
+
+    if _posthog_client is None:
+        return
+
+    _distinct_id = user_id
+
+    try:
+        properties: dict[str, Any] = {"auth_provider": "github"}
+        if email:
+            properties["email"] = email
+
+        _posthog_client.identify(
+            distinct_id=user_id,
+            properties=properties,
+        )
+        _posthog_client.flush()
+    except Exception:
+        pass  # Fail silently
+
+
+def reset_user_identity() -> None:
+    """Reset user identity after logout."""
+    global _distinct_id
+
+    _distinct_id = _get_anonymous_id()
+
+
+def get_distinct_id() -> str:
+    """Get current distinct ID for tracking."""
+    global _distinct_id
+
+    if _distinct_id is None:
+        user_id, _ = _get_user_id_from_credentials()
+        _distinct_id = user_id or _get_anonymous_id()
+
+    return _distinct_id
+
+
+def _get_cli_version() -> str:
+    """Get CLI version from package metadata."""
+    try:
+        from importlib.metadata import version
+
+        return version("wafer-cli")
+    except Exception:
+        return "unknown"
+
+
+def _get_base_properties() -> dict[str, Any]:
+    """Get base properties included with all events."""
+    return {
+        "platform": "cli",
+        "tool_id": "cli",
+        "cli_version": _get_cli_version(),
+        "os": platform.system().lower(),
+        "os_version": platform.release(),
+        "python_version": platform.python_version(),
+    }
+
+
+def track_event(event_name: str, properties: dict[str, Any] | None = None) -> None:
+    """Track a generic event.
+
+    Args:
+        event_name: Name of the event to track
+        properties: Additional properties to include
+    """
+    if not init_analytics():
+        return
+
+    if _posthog_client is None:
+        return
+
+    try:
+        event_properties = _get_base_properties()
+        if properties:
+            event_properties.update(properties)
+
+        _posthog_client.capture(
+            distinct_id=get_distinct_id(),
+            event=event_name,
+            properties=event_properties,
+        )
+    except Exception:
+        pass  # Fail silently
+
+
+def track_command(
+    command: str,
+    subcommand: str | None = None,
+    outcome: str = "success",
+    duration_ms: int | None = None,
+    properties: dict[str, Any] | None = None,
+) -> None:
+    """Track a CLI command execution.
+
+    This event counts towards DAU in the internal dashboard.
+
+    Args:
+        command: The main command name (e.g., "evaluate", "agent")
+        subcommand: Optional subcommand (e.g., "kernelbench")
+        outcome: "success" or "error"
+        duration_ms: Command execution time in milliseconds
+        properties: Additional properties to include
+    """
+    event_properties: dict[str, Any] = {
+        "command": command,
+        "outcome": outcome,
+    }
+
+    if subcommand:
+        event_properties["subcommand"] = subcommand
+
+    if duration_ms is not None:
+        event_properties["duration_ms"] = duration_ms
+
+    if properties:
+        event_properties.update(properties)
+
+    track_event("cli_command_executed", event_properties)
+
+
+def track_login(user_id: str, email: str | None = None) -> None:
+    """Track user login event.
+
+    Args:
+        user_id: Supabase user ID
+        email: User's email address
+    """
+    # First identify the user
+    identify_user(user_id, email)
+
+    # Then track the login event
+    track_event("cli_user_signed_in", {"user_id": user_id})
+
+
+def track_logout() -> None:
+    """Track user logout event."""
+    track_event("cli_user_signed_out")
+    reset_user_identity()

wafer/auth.py

@@ -287,7 +287,7 @@ if (accessToken) {
             self.end_headers()
 
 
-def browser_login(timeout: int = 120) -> tuple[str, str | None]:
+def browser_login(timeout: int = 120, port: int | None = None) -> tuple[str, str | None]:
     """Open browser for GitHub OAuth and return tokens.
 
     Starts a local HTTP server, opens browser to Supabase OAuth,
@@ -295,6 +295,7 @@ def browser_login(timeout: int = 120) -> tuple[str, str | None]:
 
     Args:
         timeout: Seconds to wait for callback (default 120)
+        port: Port for callback server. If None, finds a free port (default None)
 
     Returns:
         Tuple of (access_token, refresh_token). refresh_token may be None.
@@ -303,7 +304,8 @@ def browser_login(timeout: int = 120) -> tuple[str, str | None]:
         TimeoutError: If no callback received within timeout
         RuntimeError: If OAuth flow failed
     """
-    port = _find_free_port()
+    if port is None:
+        port = _find_free_port()
     redirect_uri = f"http://localhost:{port}/callback"
     supabase_url = get_supabase_url()