wafer-cli 0.2.19__tar.gz → 0.2.36__tar.gz

wafer_cli-0.2.36/PKG-INFO

@@ -0,0 +1,260 @@
+Metadata-Version: 2.4
+Name: wafer-cli
+Version: 0.2.36
+Summary: CLI for running GPU workloads, managing remote workspaces, and evaluating/optimizing kernels
+Requires-Python: >=3.11
+Description-Content-Type: text/markdown
+Requires-Dist: typer>=0.12.0
+Requires-Dist: trio>=0.24.0
+Requires-Dist: trio-asyncio>=0.15.0
+Requires-Dist: wafer-core>=0.1.0
+Requires-Dist: perfetto>=0.16.0
+Requires-Dist: posthog>=3.0.0
+Provides-Extra: dev
+Requires-Dist: pytest>=8.0.0; extra == "dev"
+Requires-Dist: pytest-cov>=4.1.0; extra == "dev"
+Requires-Dist: diff-cover>=8.0.0; extra == "dev"
+Requires-Dist: ruff>=0.4.0; extra == "dev"
+
+# Wafer CLI
+
+Run GPU workloads, optimize kernels, and query GPU documentation.
+
+## Getting Started
+
+```bash
+# Install
+cd apps/wafer-cli && uv sync
+
+# Use staging (workspaces and other features require staging)
+wafer config set api.environment staging
+
+# Login
+wafer login
+
+# Run a command on a remote GPU
+wafer remote-run -- nvidia-smi
+```
+
+## Commands
+
+### `wafer login` / `wafer logout` / `wafer whoami`
+
+Authenticate with GitHub OAuth.
+
+```bash
+wafer login          # Opens browser for GitHub OAuth
+wafer whoami         # Show current user
+wafer logout         # Remove credentials
+```
+
+### `wafer remote-run`
+
+Run any command on a remote GPU.
+
+```bash
+wafer remote-run -- nvidia-smi
+wafer remote-run --upload-dir ./my_code -- python3 train.py
+```
+
+### `wafer workspaces`
+
+Create and manage persistent GPU environments.
+
+**Available GPUs:**
+
+- `MI300X` - AMD Instinct MI300X (192GB HBM3, ROCm)
+- `B200` - NVIDIA Blackwell B200 (180GB HBM3e, CUDA) - default
+
+```bash
+wafer workspaces list
+wafer workspaces create my-workspace --gpu B200 --wait   # NVIDIA B200
+wafer workspaces create amd-dev --gpu MI300X             # AMD MI300X
+wafer workspaces ssh <workspace-id>
+wafer workspaces delete <workspace-id>
+```
+
+### `wafer agent`
+
+AI assistant for GPU kernel development. Helps with CUDA/Triton optimization, documentation queries, and performance analysis.
+
+```bash
+wafer agent "What is TMEM in CuTeDSL?"
+wafer agent -s "optimize this kernel" < kernel.py
+```
+
+### `wafer evaluate`
+
+Evaluate kernel correctness and performance against a reference implementation.
+
+**Functional format** (default):
+```bash
+# Generate template files
+wafer evaluate make-template ./my-kernel
+
+# Run evaluation
+wafer evaluate --impl kernel.py --reference ref.py --test-cases tests.json --benchmark
+```
+
+The implementation must define `custom_kernel(inputs)`, the reference must define `ref_kernel(inputs)` and `generate_input(**params)`.
+
+**KernelBench format** (ModelNew class):
+```bash
+# Extract a KernelBench problem as template
+wafer evaluate kernelbench make-template level1/1
+
+# Run evaluation
+wafer evaluate kernelbench --impl my_kernel.py --reference problem.py --benchmark
+```
+
+The implementation must define `class ModelNew(nn.Module)`, the reference must define `class Model`, `get_inputs()`, and `get_init_inputs()`.
+
+### `wafer wevin -t ask-docs`
+
+Query GPU documentation using the docs template.
+
+```bash
+wafer wevin -t ask-docs --json -s "What causes bank conflicts in shared memory?"
+```
+
+### `wafer corpus`
+
+Download documentation to local filesystem for agents to search.
+
+```bash
+wafer corpus list
+wafer corpus download cuda-programming-guide
+```
+
+---
+
+## Customization
+
+### `wafer remote-run` options
+
+```bash
+wafer remote-run --image pytorch/pytorch:2.5.1-cuda12.4-cudnn9-devel -- python3 script.py
+wafer remote-run --require-hwc -- ncu --set full python3 bench.py   # Hardware counters for NCU
+```
+
+### `wafer evaluate` options
+
+```bash
+wafer evaluate --impl k.py --reference r.py --test-cases t.json \
+    --target vultr-b200 \    # Specific GPU target
+    --benchmark \            # Measure performance
+    --profile                # Enable torch.profiler + NCU
+```
+
+### `wafer push` for multi-command workflows
+
+```bash
+WORKSPACE=$(wafer push ./project)
+wafer remote-run --workspace-id $WORKSPACE -- python3 test1.py
+wafer remote-run --workspace-id $WORKSPACE -- python3 test2.py
+```
+
+### Profile analysis
+
+```bash
+wafer nvidia ncu analyze profile.ncu-rep
+wafer nvidia nsys analyze profile.nsys-rep
+```
+
+---
+
+## Advanced
+
+### Local targets
+
+Bypass the API and SSH directly to your own GPUs:
+
+```bash
+wafer targets list
+wafer targets add ./my-gpu.toml
+wafer targets default my-gpu
+```
+
+### Defensive evaluation
+
+Detect evaluation hacking (stream injection, lazy evaluation, etc.):
+
+```bash
+wafer evaluate --impl k.py --reference r.py --test-cases t.json --benchmark --defensive
+```
+
+### Other tools
+
+```bash
+wafer perfetto <trace.json> --query "SELECT * FROM slice"   # Perfetto SQL queries
+wafer capture ./script.py                                    # Capture execution snapshot
+wafer compiler-analyze kernel.ptx                            # Analyze PTX/SASS
+```
+
+### ROCm profiling (AMD GPUs)
+
+```bash
+wafer rocprof-sdk ...
+wafer rocprof-systems ...
+wafer rocprof-compute ...
+```
+
+---
+
+## Shell Completion
+
+Enable tab completion for commands, options, and target names:
+
+```bash
+# Install completion (zsh/bash/fish)
+wafer --install-completion
+
+# Then restart your terminal, or source your shell config:
+source ~/.zshrc  # or ~/.bashrc
+```
+
+Now you can tab-complete:
+- Commands: `wafer eva<TAB>` → `wafer evaluate`
+- Options: `wafer evaluate --<TAB>`
+- Target names: `wafer evaluate --target v<TAB>` → `wafer evaluate --target vultr-b200`
+- File paths: `wafer evaluate --impl ./<TAB>`
+
+---
+
+## AI Assistant Skills
+
+Install the Wafer CLI skill to make wafer commands discoverable by your AI coding assistant:
+
+```bash
+# Install for all supported tools (Claude Code, Codex CLI, Cursor)
+wafer skill install
+
+# Install for a specific tool
+wafer skill install -t cursor    # Cursor
+wafer skill install -t claude    # Claude Code
+wafer skill install -t codex     # Codex CLI
+
+# Check installation status
+wafer skill status
+
+# Uninstall
+wafer skill uninstall
+```
+
+### Installing from GitHub (Cursor)
+
+You can also install the skill directly from GitHub in Cursor:
+
+1. Open Cursor Settings (Cmd+Shift+J / Ctrl+Shift+J)
+2. Navigate to **Rules** → **Add Rule** → **Remote Rule (Github)**
+3. Enter: `https://github.com/wafer-ai/skills`
+4. Cursor will automatically discover skills in `.cursor/skills/`
+
+The skill provides comprehensive guidance for GPU kernel development, including documentation lookup, trace analysis, kernel evaluation, and optimization workflows.
+
+---
+
+## Requirements
+
+- Python 3.10+
+- GitHub account (for authentication)

{wafer_cli-0.2.19 → wafer_cli-0.2.36}/pyproject.toml

@@ -1,7 +1,8 @@
 [project]
 name = "wafer-cli"
-version = "0.2.19"
-description = "CLI tool for running commands on remote GPUs and GPU kernel optimization agent"
+version = "0.2.36"
+description = "CLI for running GPU workloads, managing remote workspaces, and evaluating/optimizing kernels"
+readme = "README.md"
 requires-python = ">=3.11"
 dependencies = [
     "typer>=0.12.0",

{wafer_cli-0.2.19 → wafer_cli-0.2.36}/tests/test_cli_coverage.py

@@ -719,3 +719,17 @@ class TestWorkspacesExecFlagPassthrough:
             "workspaces", "exec", "test-ws", "--", "cmd", "--output=/tmp/out"
         ])
         assert "no such option" not in result.output.lower()
+
+
+class TestAgentNoSandboxOption:
+    """Test --no-sandbox option in wafer agent command."""
+
+    def test_agent_no_sandbox_option_exists(self) -> None:
+        """Test that --no-sandbox option is accepted by wafer agent command."""
+        result = runner.invoke(app, ["agent", "--help"])
+        assert result.exit_code == 0
+        # Strip ANSI escape codes before checking (help output may contain color codes)
+        ansi_escape = re.compile(r'\x1b\[[0-9;]*m')
+        clean_output = ansi_escape.sub('', result.stdout)
+        assert "--no-sandbox" in clean_output
+        assert "liability" in clean_output.lower()  # Warning text should be in help

{wafer_cli-0.2.19 → wafer_cli-0.2.36}/tests/test_wevin_cli.py

@@ -634,35 +634,114 @@ def test_streaming_frontend_session_start_state_without_session_id():
 
 def test_streaming_frontend_session_start_resumed_then_new():
     """Test session_start emission when resuming but states have different session_id.
-    
+
     Edge case: --resume used but states return different session_id (should use states one).
     """
     import trio
 
     from wafer.wevin_cli import StreamingChunkFrontend
-    
+
     async def run_test() -> None:
         # Start with resumed session_id
         frontend = StreamingChunkFrontend(
             session_id="resumed-session-123",
             model="claude-sonnet-4.5"
         )
-        
+
         emitted_events = []
 
         def mock_emit(obj) -> None:
             emitted_events.append(obj)
-        
+
         frontend._emit = mock_emit
-        
+
         # start() emits session_start for resumed session
         await frontend.start()
         assert len(emitted_events) == 1
         assert emitted_events[0]["session_id"] == "resumed-session-123"
-        
+
         # If states have different session_id (shouldn't happen, but handle gracefully)
         # The logic in main() checks `if first_session_id and not session_id`
         # So if session_id was set, it won't emit again
         # This is correct behavior - use the one from --resume
-    
+
     trio.run(run_test)
+
+
+# =============================================================================
+# --no-sandbox flag tests
+# =============================================================================
+
+
+def test_no_sandbox_parameter_accepted():
+    """Test that no_sandbox parameter exists in wevin_main signature."""
+    import inspect
+
+    from wafer.wevin_cli import main as wevin_main
+
+    sig = inspect.signature(wevin_main)
+    params = sig.parameters
+
+    # Verify parameter exists
+    assert 'no_sandbox' in params
+
+    # Verify type and default
+    assert str(params['no_sandbox'].annotation) in ('bool', "<class 'bool'>")
+    assert params['no_sandbox'].default is False
+
+
+def test_build_environment_accepts_no_sandbox():
+    """Test that _build_environment accepts no_sandbox parameter."""
+    import inspect
+
+    from wafer.wevin_cli import _build_environment
+
+    sig = inspect.signature(_build_environment)
+    params = sig.parameters
+
+    assert 'no_sandbox' in params
+    assert params['no_sandbox'].default is False
+
+
+def test_build_environment_with_no_sandbox_false():
+    """Test _build_environment creates env with sandbox ENABLED when no_sandbox=False."""
+    from wafer_core.rollouts.templates import TemplateConfig
+    from wafer_core.sandbox import SandboxMode
+
+    from wafer.wevin_cli import _build_environment
+
+    tpl = TemplateConfig(
+        name="test",
+        description="Test template",
+        system_prompt="Test",
+        tools=["read"],
+    )
+
+    # This will raise RuntimeError if sandbox is unavailable on this system
+    # That's expected - we're testing that sandbox is ENABLED by default
+    try:
+        env = _build_environment(tpl, None, None, no_sandbox=False)
+        # If we get here, sandbox is available - verify it's enabled
+        assert env.sandbox_mode == SandboxMode.ENABLED
+    except RuntimeError as e:
+        # Sandbox unavailable - that's OK, the error proves ENABLED is set
+        assert "sandboxing is not available" in str(e)
+
+
+def test_build_environment_with_no_sandbox_true():
+    """Test _build_environment creates env with sandbox DISABLED when no_sandbox=True."""
+    from wafer_core.rollouts.templates import TemplateConfig
+    from wafer_core.sandbox import SandboxMode
+
+    from wafer.wevin_cli import _build_environment
+
+    tpl = TemplateConfig(
+        name="test",
+        description="Test template",
+        system_prompt="Test",
+        tools=["read"],
+    )
+
+    # This should NOT raise - sandbox is disabled
+    env = _build_environment(tpl, None, None, no_sandbox=True)
+    assert env.sandbox_mode == SandboxMode.DISABLED

wafer_cli-0.2.36/wafer/agent_defaults.py

@@ -0,0 +1,197 @@
+"""Shared agent defaults for kernel optimization tasks.
+
+Single source of truth for bash allowlists and enabled tools used by both:
+- CLI templates (apps/wafer-cli/wafer/templates/*.py)
+- Eval configs (research/evals/*_eval/*.py)
+
+Import from here instead of defining your own copy.
+"""
+
+from __future__ import annotations
+
+# Tools available to the agent (coding environment tools)
+ENABLED_TOOLS: list[str] = ["read", "write", "edit", "glob", "grep", "bash"]
+
+# vLLM-specific tools (same as ENABLED_TOOLS for now)
+VLLM_ENABLED_TOOLS: list[str] = ["read", "write", "edit", "glob", "grep", "bash"]
+
+# Bash commands allowed for kernel optimization agents.
+# Uses prefix matching — "wafer evaluate" also allows "wafer evaluate kernelbench".
+KERNELBENCH_BASH_ALLOWLIST: list[str] = [
+    # Kernel evaluation
+    "wafer evaluate",
+    # Profiling — AMD
+    "wafer amd rocprof-compute",
+    "wafer amd rocprof-sdk",
+    "wafer amd rocprof-systems",
+    # Profiling — NVIDIA
+    "wafer nvidia ncu",
+    "wafer nvidia nsys",
+    # Analysis
+    "wafer compiler-analyze",
+    # Sub-agents
+    "wafer agent -t ask-docs",
+    # General utilities
+    "python",
+    "python3",
+    "timeout",
+    "ls",
+    "cat",
+    "head",
+    "tail",
+    "wc",
+    "pwd",
+    "which",
+]
+
+# Tools available to aiter optimization agents (full coding environment)
+AITER_ENABLED_TOOLS: list[str] = ["read", "write", "edit", "glob", "grep", "bash"]
+
+# System prompt for aiter optimization (shared between eval and template)
+# Uses {op_name}, {test_file}, {target_flag} placeholders
+AITER_SYSTEM_PROMPT = """\
+You are a GPU kernel optimization expert specializing in AMD MI300X and the aiter library.
+
+## Context
+
+aiter (ROCm/aiter) is AMD's centralized repository for high-performance AI operators.
+Operators are implemented using Triton kernels, Composable Kernel (CK), or HIP/ROCm.
+
+Each operator has a test in `op_tests/test_{{op}}.py` that validates correctness and
+measures performance against a reference implementation.
+
+## Your Task
+
+1. **Understand the operator**: Read the test file and trace imports to find implementation
+2. **Establish baseline**: Run the evaluation to measure current performance
+   ```bash
+   # Quick check with one shape (fast iteration)
+   wafer evaluate aiter --aiter-dir . --cmd "python op_tests/test_{{op}}.py --mnk 128,32,8192" {target_flag}
+
+   # Full test suite (final validation)
+   wafer evaluate aiter --aiter-dir . --cmd "python op_tests/test_{{op}}.py" {target_flag}
+   ```
+3. **Identify optimizations**: Look for memory access patterns, occupancy, instruction selection
+4. **Implement changes**: Modify the operator to improve performance
+5. **Validate**: Re-run evaluation to verify correctness and measure speedup
+6. **Iterate**: Use quick checks during development, full suite for final validation
+
+## Finding Source Files
+
+The aiter codebase structure varies by operator. To find implementation files:
+
+1. **Start with the test file**: `op_tests/test_{{op}}.py`
+   - Read imports to see what modules are used
+   - Look for the main function being tested
+
+2. **Check common locations** (not all ops have all of these):
+   - `aiter/ops/{{op}}.py` — High-level Python API (some ops)
+   - `aiter/triton_kernels/` — Triton kernel implementations
+   - `csrc/kernels/` — CUDA/HIP kernel implementations
+   - `csrc/py_itfs_cu/` — Python interface CUDA files
+   - `csrc/cktile_*/` — Composable Kernel tile implementations
+
+3. **Search for the op name**:
+   ```bash
+   find . -name "*{{op}}*" -type f | grep -v __pycache__
+   grep -r "def {{function_name}}" aiter/ csrc/ --include="*.py" --include="*.cu"
+   ```
+
+## Key Directories
+
+- `aiter/` — Main package with operator implementations
+- `aiter/ops/` — High-level operator APIs (some ops)
+- `aiter/triton_kernels/` — Triton kernel implementations
+- `csrc/` — C++/CUDA/HIP implementations
+- `op_tests/` — Tests for each operator
+- `aiter/configs/` — Tuned configurations (CSV files)
+
+## Output
+
+Your goal is to produce:
+1. Modified operator code with optimizations
+2. Benchmark results showing correctness and speedup
+3. A summary of what you changed and why
+
+The optimization should be correct (pass the op_test) and faster than baseline."""
+
+# Bash commands allowed for aiter optimization agents.
+AITER_BASH_ALLOWLIST: list[str] = [
+    # Read-only
+    "ls",
+    "cat",
+    "head",
+    "tail",
+    "wc",
+    "find",
+    "grep",
+    "rg",
+    "pwd",
+    "tree",
+    "which",
+    "diff",
+    "sort",
+    # Filesystem
+    "mkdir",
+    "cp",
+    "mv",
+    # Git
+    "git diff",
+    "git status",
+    "git log",
+    # Compilation
+    "hipcc",
+    "g++",
+    "gcc",
+    "clang",
+    "python",
+    "python3",
+    "pip",
+    "pytest",
+    # Execution — allows running compiled binaries and python scripts
+    "./",
+    # Kernel evaluation
+    "wafer evaluate aiter",
+    # Profiling — AMD
+    "wafer amd rocprof-compute",
+    "wafer amd rocprof-sdk",
+    "wafer amd rocprof-systems",
+    "wafer amd isa",
+    # Sub-agents
+    "wafer agent -t ask-docs",
+    # Misc
+    "timeout",
+]
+
+# Bash commands allowed for vLLM kernel optimization agents.
+VLLM_BASH_ALLOWLIST: list[str] = [
+    # vLLM evaluation
+    "wafer evaluate vllm",
+    # vLLM's own test and benchmark commands (run inside vllm dir)
+    "pytest",
+    # Profiling — AMD
+    "wafer amd rocprof-compute",
+    "wafer amd rocprof-sdk",
+    "wafer amd rocprof-systems",
+    # Profiling — NVIDIA
+    "wafer nvidia ncu",
+    "wafer nvidia nsys",
+    # Analysis
+    "wafer compiler-analyze",
+    # Sub-agents
+    "wafer agent -t ask-docs",
+    # General utilities
+    "python",
+    "python3",
+    "pip",
+    "timeout",
+    "ls",
+    "cat",
+    "head",
+    "tail",
+    "wc",
+    "pwd",
+    "which",
+    "cd",
+    "git",
+]