wafer-cli 0.1.0__tar.gz

wafer_cli-0.1.0/PKG-INFO

@@ -0,0 +1,9 @@
+Metadata-Version: 2.4
+Name: wafer-cli
+Version: 0.1.0
+Summary: CLI tool for running commands on remote GPUs and GPU kernel optimization agent
+Requires-Python: >=3.11
+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

wafer_cli-0.1.0/README.md

@@ -0,0 +1,221 @@
+# Wafer CLI
+
+Run commands on remote GPUs in Docker containers.
+
+## Installation
+
+```bash
+cd apps/wafer-cli
+uv sync
+```
+
+## Quick Start
+
+```bash
+# Login with GitHub
+wafer login
+
+# Run a command on remote GPU
+wafer remote-run -- nvidia-smi
+
+# Run Python script with file upload
+wafer remote-run --upload-dir ./my_project -- python3 train.py
+
+# Check who you're logged in as
+wafer whoami
+```
+
+## Commands
+
+### `wafer login`
+
+Authenticate with GitHub OAuth. Opens browser for login flow.
+
+```bash
+wafer login
+```
+
+Credentials are stored in `~/.wafer/credentials.json`.
+
+### `wafer remote-run`
+
+Run any command on a remote GPU inside a Docker container.
+
+```bash
+# Basic: run a command
+wafer remote-run -- nvidia-smi
+
+# Upload files first, then run
+wafer remote-run --upload-dir ./my_project -- python3 main.py
+
+# Custom Docker image
+wafer remote-run --image pytorch/pytorch:2.5.1-cuda12.4-cudnn9-devel -- python3 -c "import torch; print(torch.cuda.get_device_name())"
+
+# With custom entrypoint (for images with non-shell defaults)
+wafer remote-run --image vllm/vllm-openai:latest --docker-entrypoint bash -- python3 script.py
+```
+
+Options:
+- `--upload-dir <path>` - Upload local directory before running command
+- `--image <image>` - Docker image to use (default: vllm/vllm-openai:latest)
+- `--docker-entrypoint <cmd>` - Override container entrypoint
+- `--pull-image` - Pull image even if it exists locally
+- `--require-hwc` - Require hardware counters (for NCU profiling)
+
+### `wafer push`
+
+Upload files to a remote workspace (for multi-command workflows).
+
+```bash
+# Push files
+wafer push ./my_project
+
+# Returns workspace_id for use with remote-run
+# Example: ws_abc123
+```
+
+### `wafer logout`
+
+Remove stored credentials.
+
+```bash
+wafer logout
+```
+
+### `wafer whoami`
+
+Show current authenticated user.
+
+```bash
+wafer whoami
+# Output: Logged in as user@example.com
+```
+
+## Examples
+
+### Run PyTorch training
+
+```bash
+wafer remote-run --upload-dir ./training -- python3 train.py --epochs 10
+```
+
+### Profile with NCU
+
+```bash
+wafer remote-run --require-hwc --upload-dir ./kernel -- ncu --set full python3 benchmark.py
+```
+
+### Interactive debugging
+
+```bash
+# Upload once
+WORKSPACE=$(wafer push ./project)
+
+# Run multiple commands against same workspace
+wafer remote-run --workspace-id $WORKSPACE -- python3 test1.py
+wafer remote-run --workspace-id $WORKSPACE -- python3 test2.py
+```
+
+### Custom CUDA image
+
+```bash
+wafer remote-run \
+  --image nvidia/cuda:12.8.0-cudnn-devel-ubuntu22.04 \
+  --upload-dir ./cuda_kernels \
+  -- nvcc -o kernel kernel.cu && ./kernel
+```
+
+---
+
+## Architecture
+
+```
+CLI                              wafer-api                         GPU
+ |                                   |                               |
+ |-- POST /v1/gpu/jobs ------------->|                               |
+ |   { command, files[], image }     |                               |
+ |                                   |-- SSH upload + docker ------->|
+ |<-- SSE stream: stdout/stderr -----|<-- stream output -------------|
+ |                                   |                               |
+```
+
+### Components
+
+- **wafer-cli** (`apps/wafer-cli/`) - Thin client that calls wafer-api
+- **wafer-api** (`services/wafer-api/`) - Backend that owns GPU targets and SSH credentials
+- **wafer-core** (`packages/wafer-core/`) - Internal SSH client for file upload and command execution
+
+### Why API-backed?
+
+- **Security**: SSH credentials stay on server, not in client config
+- **Routing**: Backend picks best available GPU based on requirements
+- **Multi-tenant**: Multiple users share GPU pool without credential management
+
+## Local Target Management (Advanced)
+
+For direct SSH access (bypassing wafer-api), you can configure local targets:
+
+```bash
+wafer targets list
+wafer targets add examples/targets/my-gpu.toml
+wafer targets default my-gpu
+```
+
+See `examples/targets/` for TOML format.
+
+## Testing
+
+### Prerequisites
+
+1. Start the wafer-api server:
+   ```bash
+   cd services/wafer-api && uv run uvicorn src.main:app --port 8000
+   ```
+
+2. Authenticate (if not already logged in):
+   ```bash
+   cd apps/wafer-cli && uv run wafer login
+   ```
+
+### Run Integration Tests
+
+```bash
+# GPU API tests (push, jobs, cancellation)
+python scripts/test_gpu_api.py --no-server
+
+# CLI remote-run tests (basic, upload-dir, nested)
+python scripts/test_cli_remote_run.py --no-server
+```
+
+### Manual Command Tests
+
+```bash
+cd apps/wafer-cli
+
+# wafer remote-run (via API)
+uv run wafer remote-run -- nvidia-smi
+
+# wafer remote-run with file upload
+uv run wafer remote-run --upload-dir ./my_project -- python script.py
+
+# wafer ncu-analyze (via API)
+uv run wafer ncu-analyze path/to/profile.ncu-rep --remote
+
+# wafer ask-docs (requires docs-tool running on port 8002)
+WAFER_DOCS_URL=http://localhost:8002 uv run wafer ask-docs "What is a Triton kernel?"
+
+# wafer wevin with --tools and --json flags
+uv run wafer wevin --ref kernel.py --desc "Optimize" --test "n=128" --tools read,write,edit --json --no-tui --max-turns 1
+```
+
+### Expected Results
+
+- All integration tests should pass
+- `wafer remote-run` executes commands on remote GPU via API
+- `wafer ncu-analyze --remote` uploads profile and returns analysis
+- All SSH operations use internal `wafer_core.ssh.SSHClient`
+
+## Requirements
+
+- Python 3.10+
+- GitHub account (for authentication)

wafer_cli-0.1.0/pyproject.toml

@@ -0,0 +1,73 @@
+[project]
+name = "wafer-cli"
+version = "0.1.0"
+description = "CLI tool for running commands on remote GPUs and GPU kernel optimization agent"
+requires-python = ">=3.11"
+dependencies = [
+    "typer>=0.12.0",
+    "trio>=0.24.0",
+    "trio-asyncio>=0.15.0",  # Bridge asyncssh (asyncio) to trio for async SSH
+    # Wafer core for environments and utils (includes rollouts)
+    "wafer-core>=0.1.0",
+]
+
+[project.scripts]
+wafer = "wafer.cli:main"
+
+[tool.uv.sources]
+wafer-core = { workspace = true }
+
+[build-system]
+requires = ["setuptools>=61.0"]
+build-backend = "setuptools.build_meta"
+
+[tool.setuptools.packages.find]
+where = ["."]
+include = ["wafer*"]
+
+[tool.ruff]
+line-length = 100
+target-version = "py311"
+preview = true  # Required for PLR1702 (too-many-nested-blocks)
+
+[tool.ruff.lint]
+select = [
+    "E",       # pycodestyle errors
+    "F",       # pyflakes
+    "I",       # isort (import sorting)
+    "ANN",     # flake8-annotations (enforce type annotations)
+    "ASYNC",   # flake8-async (trio/asyncio best practices)
+    "B",       # flake8-bugbear (common bugs)
+    "UP",      # pyupgrade (modern Python patterns)
+    "PLR0913", # too-many-arguments (Tiger Style: "hourglass shape: few parameters")
+    "PLR0915", # too-many-statements (Tiger Style: 70 line limit)
+    "PLR1702", # too-many-nested-blocks (Tiger Style: "centralize control flow")
+    "PLW2901", # redefined-loop-variable (Carmack SSA: single assignment)
+    "RET506",  # superfluous-else-raise (explicit control flow)
+    "RET507",  # superfluous-else-continue (explicit control flow)
+    "A",       # flake8-builtins (shadowing builtins like list, str)
+    "RUF018",  # assignment-in-assert (catches typos)
+    "TRY002",  # raise-vanilla-exception (use specific exceptions)
+    "TRY003",  # raise-vanilla-args (proper exception messages)
+    "TRY004",  # type-check-without-type-error (use TypeError for type checks)
+    "TRY201",  # verbose-raise (use bare raise)
+    "TRY300",  # try-consider-else (clear control flow)
+    "TRY400",  # error-instead-of-exception (use logging.exception)
+]
+ignore = [
+    "E501",    # Line too long (handled by formatter)
+    "B008",    # Typer uses function calls in defaults
+    "TRY003",  # Avoid specifying long messages outside exception class (too opinionated - revisit later)
+    "TRY300",  # Consider moving to else block (too opinionated about try/except style - revisit later)
+    "TRY400",  # Use logging.exception instead of logging.error (we prefer explicit error logging - revisit later)
+]
+
+[tool.ruff.lint.per-file-ignores]
+"tests/**/*.py" = ["ANN201"]  # Don't require return type annotations in tests
+"wafer/evaluate.py" = ["PLR0915"]  # run_evaluate_docker/ssh have complex deployment flows (80/77 statements) - TODO: refactor
+
+[tool.ruff.lint.pylint]
+max-args = 7             # Max function arguments (Tiger Style: few parameters)
+max-statements = 70      # Max statements per function (Tiger Style: 70 line limit)
+max-branches = 12        # Max if/elif branches
+max-nested-blocks = 5    # Max nesting depth (Linus rule: "if you need more than 3 levels, you're screwed")

wafer_cli-0.1.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

wafer_cli-0.1.0/tests/test_compiler_analyze_integration.py

@@ -0,0 +1,320 @@
+"""Integration tests for wafer compiler-analyze CLI command.
+
+Tests the compiler-analyze command end-to-end, verifying:
+- CLI command execution
+- JSON output format
+- Error handling
+- Various input combinations
+"""
+
+import json
+import subprocess
+
+import pytest
+
+# Constants
+CLI_TIMEOUT_SECONDS = 30
+CLI_COMMAND = "wafer"
+COMPILER_ANALYZE_SUBCOMMAND = "compiler-analyze"
+JSON_FLAG = "--json"
+MLIR_FLAG = "--mlir-text"
+PTX_FLAG = "--ptx-text"
+SASS_FLAG = "--sass-text"
+SOURCE_FLAG = "--source-text"
+KERNEL_NAME_FLAG = "--kernel-name"
+
+SUCCESS_KEY = "success"
+DATA_KEY = "data"
+ERROR_KEY = "error"
+KERNEL_NAME_KEY = "kernel_name"
+MLIR_TEXT_KEY = "mlir_text"
+PTX_TEXT_KEY = "ptx_text"
+SASS_TEXT_KEY = "sass_text"
+SOURCE_CODE_KEY = "source_code"
+PARSED_MLIR_KEY = "parsed_mlir"
+PARSED_PTX_KEY = "parsed_ptx"
+PARSED_SASS_KEY = "parsed_sass"
+LAYOUTS_KEY = "layouts"
+MEMORY_PATHS_KEY = "memory_paths"
+PIPELINE_STAGES_KEY = "pipeline_stages"
+
+
+def test_compiler_analyze_basic() -> None:
+    """Test basic compiler-analyze with minimal inputs."""
+    mlir_text = "module { func @kernel() { } }"
+    ptx_text = ".version 8.0\n.target sm_80\n.entry kernel() { ret; }"
+    sass_text = "// SASS code"
+    
+    args = [
+        CLI_COMMAND,
+        COMPILER_ANALYZE_SUBCOMMAND,
+        MLIR_FLAG, mlir_text,
+        PTX_FLAG, ptx_text,
+        SASS_FLAG, sass_text,
+        JSON_FLAG
+    ]
+    
+    result = subprocess.run(
+        args,
+        capture_output=True,
+        text=True,
+        timeout=CLI_TIMEOUT_SECONDS,
+    )
+    
+    assert result.returncode == 0
+    assert result.stderr is not None
+    
+    output_json = json.loads(result.stdout)
+    assert isinstance(output_json, dict)
+    assert output_json.get(SUCCESS_KEY) is True
+    assert DATA_KEY in output_json
+    
+    data = output_json[DATA_KEY]
+    assert isinstance(data, dict)
+    assert KERNEL_NAME_KEY in data
+    assert MLIR_TEXT_KEY in data
+    assert PTX_TEXT_KEY in data
+    assert SASS_TEXT_KEY in data
+
+
+def test_compiler_analyze_with_source_code() -> None:
+    """Test compiler-analyze with optional source code."""
+    mlir_text = "module { func @test_kernel() { } }"
+    ptx_text = ".version 8.0\n.target sm_80"
+    sass_text = "// SASS"
+    source_code = "__global__ void test_kernel() { }"
+    
+    args = [
+        CLI_COMMAND,
+        COMPILER_ANALYZE_SUBCOMMAND,
+        MLIR_FLAG, mlir_text,
+        PTX_FLAG, ptx_text,
+        SASS_FLAG, sass_text,
+        SOURCE_FLAG, source_code,
+        JSON_FLAG
+    ]
+    
+    result = subprocess.run(
+        args,
+        capture_output=True,
+        text=True,
+        timeout=CLI_TIMEOUT_SECONDS,
+    )
+    
+    assert result.returncode == 0
+    assert result.stderr is not None
+    
+    output_json = json.loads(result.stdout)
+    assert isinstance(output_json, dict)
+    assert output_json.get(SUCCESS_KEY) is True
+    assert DATA_KEY in output_json
+    
+    data = output_json[DATA_KEY]
+    assert isinstance(data, dict)
+    assert data.get(SOURCE_CODE_KEY) == source_code
+
+
+def test_compiler_analyze_with_kernel_name() -> None:
+    """Test compiler-analyze with explicit kernel name."""
+    mlir_text = "module { func @my_kernel() { } }"
+    ptx_text = ".version 8.0\n.target sm_80"
+    sass_text = "// SASS"
+    kernel_name = "my_custom_kernel"
+    
+    args = [
+        CLI_COMMAND,
+        COMPILER_ANALYZE_SUBCOMMAND,
+        MLIR_FLAG, mlir_text,
+        PTX_FLAG, ptx_text,
+        SASS_FLAG, sass_text,
+        KERNEL_NAME_FLAG, kernel_name,
+        JSON_FLAG
+    ]
+    
+    result = subprocess.run(
+        args,
+        capture_output=True,
+        text=True,
+        timeout=CLI_TIMEOUT_SECONDS,
+    )
+    
+    assert result.returncode == 0
+    assert result.stderr is not None
+    
+    output_json = json.loads(result.stdout)
+    assert isinstance(output_json, dict)
+    assert output_json.get(SUCCESS_KEY) is True
+    assert DATA_KEY in output_json
+    
+    data = output_json[DATA_KEY]
+    assert isinstance(data, dict)
+    assert data.get(KERNEL_NAME_KEY) == kernel_name
+
+
+def test_compiler_analyze_missing_required_args() -> None:
+    """Test compiler-analyze fails gracefully when required args are missing."""
+    args = [
+        CLI_COMMAND,
+        COMPILER_ANALYZE_SUBCOMMAND,
+        JSON_FLAG
+    ]
+    
+    result = subprocess.run(
+        args,
+        capture_output=True,
+        text=True,
+        timeout=CLI_TIMEOUT_SECONDS,
+    )
+    
+    stdout_has_content = len(result.stdout.strip()) > 0
+    
+    if stdout_has_content:
+        output_json = json.loads(result.stdout)
+        assert isinstance(output_json, dict)
+        success_value = output_json.get(SUCCESS_KEY)
+        has_error = ERROR_KEY in output_json
+        assert success_value is False or has_error
+    else:
+        has_nonzero_exit = result.returncode != 0
+        has_stderr = len(result.stderr) > 0
+        assert has_nonzero_exit or has_stderr
+
+
+def test_compiler_analyze_partial_inputs() -> None:
+    """Test compiler-analyze with only some inputs provided."""
+    mlir_text = "module { func @kernel() { } }"
+    ptx_text = ".version 8.0\n.target sm_80"
+    empty_sass = ""
+    
+    args = [
+        CLI_COMMAND,
+        COMPILER_ANALYZE_SUBCOMMAND,
+        MLIR_FLAG, mlir_text,
+        PTX_FLAG, ptx_text,
+        SASS_FLAG, empty_sass,
+        JSON_FLAG
+    ]
+    
+    result = subprocess.run(
+        args,
+        capture_output=True,
+        text=True,
+        timeout=CLI_TIMEOUT_SECONDS,
+    )
+    
+    assert result.returncode == 0
+    assert result.stderr is not None
+    
+    output_json = json.loads(result.stdout)
+    assert isinstance(output_json, dict)
+    has_success = SUCCESS_KEY in output_json
+    has_error = ERROR_KEY in output_json
+    assert has_success or has_error
+
+
+def test_compiler_analyze_output_structure() -> None:
+    """Test compiler-analyze output has expected structure."""
+    mlir_text = "module { func @kernel() { } }"
+    ptx_text = ".version 8.0\n.target sm_80"
+    sass_text = "// SASS"
+    
+    args = [
+        CLI_COMMAND,
+        COMPILER_ANALYZE_SUBCOMMAND,
+        MLIR_FLAG, mlir_text,
+        PTX_FLAG, ptx_text,
+        SASS_FLAG, sass_text,
+        JSON_FLAG
+    ]
+    
+    result = subprocess.run(
+        args,
+        capture_output=True,
+        text=True,
+        timeout=CLI_TIMEOUT_SECONDS,
+    )
+    
+    assert result.returncode == 0
+    assert result.stdout is not None
+    
+    output_json = json.loads(result.stdout)
+    assert isinstance(output_json, dict)
+    
+    data = output_json.get(DATA_KEY, {})
+    assert isinstance(data, dict)
+    
+    expected_fields = [
+        KERNEL_NAME_KEY,
+        MLIR_TEXT_KEY,
+        PTX_TEXT_KEY,
+        SASS_TEXT_KEY,
+        PARSED_MLIR_KEY,
+        PARSED_PTX_KEY,
+        PARSED_SASS_KEY,
+        LAYOUTS_KEY,
+        MEMORY_PATHS_KEY,
+        PIPELINE_STAGES_KEY,
+    ]
+    
+    for field in expected_fields:
+        assert field in data
+
+
+def test_compiler_analyze_error_handling() -> None:
+    """Test compiler-analyze handles invalid inputs gracefully."""
+    invalid_mlir = "{ invalid json }"
+    ptx_text = ".version 8.0"
+    sass_text = "// SASS"
+    
+    args = [
+        CLI_COMMAND,
+        COMPILER_ANALYZE_SUBCOMMAND,
+        MLIR_FLAG, invalid_mlir,
+        PTX_FLAG, ptx_text,
+        SASS_FLAG, sass_text,
+        JSON_FLAG
+    ]
+    
+    result = subprocess.run(
+        args,
+        capture_output=True,
+        text=True,
+        timeout=CLI_TIMEOUT_SECONDS,
+    )
+    
+    assert result.stdout is not None
+    
+    output_json = json.loads(result.stdout)
+    assert isinstance(output_json, dict)
+    
+    has_success = SUCCESS_KEY in output_json
+    has_error = ERROR_KEY in output_json
+    assert has_success or has_error
+
+
+def test_compiler_analyze_no_json_flag() -> None:
+    """Test compiler-analyze without --json flag (should still work)."""
+    mlir_text = "module { func @kernel() { } }"
+    ptx_text = ".version 8.0\n.target sm_80"
+    sass_text = "// SASS"
+    
+    args = [
+        CLI_COMMAND,
+        COMPILER_ANALYZE_SUBCOMMAND,
+        MLIR_FLAG, mlir_text,
+        PTX_FLAG, ptx_text,
+        SASS_FLAG, sass_text,
+    ]
+    
+    result = subprocess.run(
+        args,
+        capture_output=True,
+        text=True,
+        timeout=CLI_TIMEOUT_SECONDS,
+    )
+    
+    assert result.returncode == 0
+    assert result.stdout is not None
+    
+    output_json = json.loads(result.stdout)
+    assert isinstance(output_json, dict)

wafer_cli-0.1.0/tests/test_config_integration.py

@@ -0,0 +1,50 @@
+"""Integration tests for config loading and environment resolution.
+
+Tests WaferConfig loading and environment resolution.
+"""
+
+from pathlib import Path
+
+import pytest
+
+from wafer.config import WaferConfig
+from wafer.inference import resolve_environment
+
+# Constants
+CONFIG_DIR = ".wafer"
+CONFIG_FILENAME = "config.toml"
+
+
+def test_config_and_environment_resolution() -> None:
+    """Test config loading and environment resolution."""
+    home_dir = Path.home()
+    assert home_dir.exists()
+    
+    config_path = home_dir / CONFIG_DIR / CONFIG_FILENAME
+    
+    if not config_path.exists():
+        pytest.skip(f"Config not found: {config_path}")
+
+    config = WaferConfig.from_toml(config_path)
+    assert config is not None
+
+    assert config.target is not None
+    assert len(config.target) > 0
+    assert config.ssh_key is not None
+    assert len(config.ssh_key) > 0
+    assert len(config.environments) > 0
+
+    default_env = resolve_environment(config, None)
+    assert default_env is not None
+    assert default_env.docker is not None
+    assert len(default_env.docker) > 0
+    assert default_env.name is not None
+    assert len(default_env.name) > 0
+
+    env_names = list(config.environments.keys())
+    assert len(env_names) > 0
+    
+    explicit_env_name = env_names[0]
+    explicit_env = resolve_environment(config, explicit_env_name)
+    assert explicit_env is not None
+    assert explicit_env.name == explicit_env_name