wafer-cli 0.2.14__py3-none-any.whl

wafer/GUIDE.md

@@ -0,0 +1,118 @@
+# Wafer CLI Guide
+
+GPU development primitives for LLM agents.
+
+## Quick Start: Cloud GPU (No Setup)
+
+Run code on cloud GPUs instantly with workspaces:
+
+```bash
+wafer login                              # One-time auth
+wafer workspaces create dev --gpu B200   # Create workspace (NVIDIA B200)
+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
+```
+
+**Available GPUs:**
+
+- `MI300X` - AMD Instinct MI300X (192GB HBM3, ROCm)
+- `B200` - NVIDIA Blackwell B200 (180GB HBM3e, CUDA) - default
+
+## Documentation Lookup
+
+Answer GPU programming questions from indexed documentation.
+
+```bash
+# Download corpus (one-time)
+wafer corpus download cuda
+wafer corpus download cutlass
+wafer corpus download hip
+
+# Query documentation
+wafer agent -t ask-docs --corpus cuda "What is warp divergence?"
+wafer agent -t ask-docs --corpus cutlass "What is a TiledMma?"
+```
+
+## Trace Analysis
+
+Analyze performance traces from NCU, NSYS, or PyTorch profiler.
+
+```bash
+# 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?"
+
+# Direct trace queries (PyTorch/Perfetto JSON)
+wafer nvidia perfetto tables trace.json
+wafer nvidia perfetto query trace.json \
+  "SELECT name, dur/1e6 as ms FROM slice WHERE cat='kernel' ORDER BY dur DESC LIMIT 10"
+
+# NCU/NSYS analysis
+wafer nvidia ncu analyze profile.ncu-rep
+wafer nvidia nsys analyze profile.nsys-rep
+```
+
+## Kernel Evaluation
+
+Test kernel correctness and measure speedup against a reference.
+
+```bash
+# Using workspaces (no target setup required):
+wafer workspaces create dev --gpu B200
+wafer workspaces exec --sync ./my-kernel dev -- python test_kernel.py
+
+# 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 <target-name>
+```
+
+For target setup, see `wafer config targets --help`.
+
+## Kernel Optimization (AI-assisted)
+
+Iteratively optimize a kernel with evaluation feedback.
+
+```bash
+wafer agent -t optimize-kernel \
+  --args kernel=./my_kernel.cu \
+  --args target=H100 \
+  "Optimize this GEMM for memory bandwidth"
+```
+
+## Workspaces
+
+Cloud GPU environments with no setup required.
+
+**Available GPUs:**
+
+- `MI300X` - AMD Instinct MI300X (192GB HBM3, ROCm)
+- `B200` - NVIDIA Blackwell B200 (180GB HBM3e, CUDA) - default
+
+```bash
+wafer workspaces create dev --gpu B200 --wait     # NVIDIA B200
+wafer workspaces create amd-dev --gpu MI300X      # AMD MI300X
+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
+```
+
+See `wafer workspaces --help` for details.
+
+## Command Reference
+
+```bash
+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/__init__.py

@@ -0,0 +1,3 @@
+"""Wafer - Standalone tool for Docker execution on remote GPUs."""
+
+__version__ = "0.1.0"

wafer/analytics.py

@@ -0,0 +1,306 @@
+"""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 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/api_client.py

@@ -0,0 +1,195 @@
+"""Wafer API client for remote GPU operations.
+
+Thin client that calls wafer-api endpoints instead of direct SSH.
+"""
+
+import base64
+import sys
+from dataclasses import dataclass
+from pathlib import Path
+
+import httpx
+
+from .global_config import get_api_url  # noqa: F401 - re-exported for backwards compat
+
+
+@dataclass(frozen=True)
+class PushResult:
+    """Result of pushing files to GPU."""
+
+    workspace_id: str
+    workspace_path: str
+    files_uploaded: list[str]
+
+
+@dataclass(frozen=True)
+class ApiConfig:
+    """API client configuration."""
+
+    base_url: str = "http://localhost:8000"  # Only used if ApiConfig is instantiated directly
+    timeout: float = 60.0
+
+
+def _get_auth_headers() -> dict[str, str]:
+    """Get auth headers from stored credentials (lazy import to avoid circular)."""
+    from .auth import get_auth_headers
+
+    return get_auth_headers()
+
+
+def push_directory(local_path: Path, workspace_name: str | None = None) -> PushResult:
+    """Push local directory to GPU via wafer-api.
+
+    Args:
+        local_path: Local directory to upload
+        workspace_name: Optional workspace name (defaults to directory name)
+
+    Returns:
+        PushResult with workspace_id and uploaded files
+
+    Raises:
+        FileNotFoundError: If local_path doesn't exist
+        ValueError: If local_path is not a directory
+        httpx.HTTPError: If API request fails
+    """
+    if not local_path.exists():
+        raise FileNotFoundError(f"Path not found: {local_path}")
+    if not local_path.is_dir():
+        raise ValueError(f"Not a directory: {local_path}")
+
+    # Collect files and encode as base64
+    files = []
+    for file_path in local_path.rglob("*"):
+        if file_path.is_file():
+            relative_path = file_path.relative_to(local_path)
+            content = file_path.read_bytes()
+            files.append({
+                "path": str(relative_path),
+                "content": base64.b64encode(content).decode(),
+            })
+
+    # Build request
+    request_body = {
+        "files": files,
+        "workspace_name": workspace_name or local_path.name,
+    }
+
+    # Call API
+    api_url = get_api_url()
+    headers = _get_auth_headers()
+    with httpx.Client(timeout=60.0, headers=headers) as client:
+        response = client.post(f"{api_url}/v1/gpu/push", json=request_body)
+        response.raise_for_status()
+        data = response.json()
+
+    return PushResult(
+        workspace_id=data["workspace_id"],
+        workspace_path=data["workspace_path"],
+        files_uploaded=data["files_uploaded"],
+    )
+
+
+def _collect_files(local_path: Path) -> list[dict]:
+    """Collect files from directory as base64-encoded dicts."""
+    files = []
+    for file_path in local_path.rglob("*"):
+        if file_path.is_file():
+            relative_path = file_path.relative_to(local_path)
+            content = file_path.read_bytes()
+            files.append({
+                "path": str(relative_path),
+                "content": base64.b64encode(content).decode(),
+            })
+    return files
+
+
+def run_command_stream(
+    command: str,
+    upload_dir: Path | None = None,
+    workspace_id: str | None = None,
+    gpu_id: int | None = None,
+    gpu_count: int = 1,
+    docker_image: str | None = None,
+    docker_entrypoint: str | None = None,
+    pull_image: bool = False,
+    require_hardware_counters: bool = False,
+    target: str | None = None,
+) -> int:
+    """Run command on GPU via wafer-api, streaming output.
+
+    Two modes (mutually exclusive):
+    - upload_dir: Upload files and run (stateless, high-level)
+    - workspace_id: Use existing workspace (low-level)
+
+    Args:
+        command: Command to execute inside container
+        upload_dir: Directory to upload (stateless mode)
+        workspace_id: Workspace ID from push (low-level mode)
+        gpu_id: GPU ID to use (optional)
+        gpu_count: Number of GPUs needed (1-8, default 1)
+        docker_image: Docker image override (optional)
+        docker_entrypoint: Docker entrypoint override (optional, e.g., "bash")
+        pull_image: Pull image if not available (optional, default False)
+        require_hardware_counters: Require baremetal for ncu profiling (optional)
+        target: Target name to use (optional, defaults to user's default)
+
+    Returns:
+        Exit code (0 = success, non-zero = failure)
+
+    Raises:
+        httpx.HTTPError: If API request fails
+    """
+    request_body: dict = {
+        "command": command,
+    }
+
+    # Add files or workspace_id (mutually exclusive)
+    if upload_dir is not None:
+        files = _collect_files(upload_dir)
+        request_body["files"] = files
+        request_body["workspace_name"] = upload_dir.name
+    elif workspace_id is not None:
+        request_body["workspace_id"] = workspace_id
+    # else: no files, no workspace (run command in temp workspace)
+
+    if gpu_id is not None:
+        request_body["gpu_id"] = gpu_id
+    if gpu_count > 1:
+        request_body["gpu_count"] = gpu_count
+    if docker_image is not None:
+        request_body["docker_image"] = docker_image
+    if docker_entrypoint is not None:
+        request_body["docker_entrypoint"] = docker_entrypoint
+    if pull_image:
+        request_body["pull_image"] = True
+    if require_hardware_counters:
+        request_body["require_hardware_counters"] = True
+    if target is not None:
+        request_body["target"] = target
+
+    api_url = get_api_url()
+    headers = _get_auth_headers()
+    exit_code = 0
+
+    with httpx.Client(timeout=None, headers=headers) as client:  # No timeout for streaming
+        with client.stream("POST", f"{api_url}/v1/gpu/jobs", json=request_body) as response:
+            response.raise_for_status()
+
+            for line in response.iter_lines():
+                if not line:
+                    continue
+
+                # Parse SSE format: "data: <content>"
+                if line.startswith("data: "):
+                    content = line[6:]  # Strip "data: " prefix
+
+                    if content == "[DONE]":
+                        break
+                    elif content.startswith("[ERROR]"):
+                        print(content[8:], file=sys.stderr)  # Strip "[ERROR] "
+                        exit_code = 1
+                        break
+                    else:
+                        print(content)
+
+    return exit_code