wafer-cli 0.2.9__py3-none-any.whl → 0.2.10__py3-none-any.whl

wafer/output.py

@@ -0,0 +1,241 @@
+"""Structured output formatting for CLI commands.
+
+Provides JSON and JSONL output formats for machine-readable CLI output.
+"""
+
+from __future__ import annotations
+
+import json
+import re
+from dataclasses import asdict, dataclass, field
+from datetime import UTC, datetime
+from enum import StrEnum
+from typing import TYPE_CHECKING, Any, Literal
+
+import typer
+
+if TYPE_CHECKING:
+    from .evaluate import EvaluateResult
+
+
+class OutputFormat(StrEnum):
+    """Output format for CLI commands."""
+
+    TEXT = "text"  # Human-readable (default)
+    JSON = "json"  # Single JSON object at end
+    JSONL = "jsonl"  # Streaming JSON Lines
+
+
+@dataclass
+class EvalOutput:
+    """Structured evaluation result for JSON output."""
+
+    status: Literal["success", "failure", "error"]
+    target: str | None = None
+    phase: str | None = None
+    correctness: dict[str, Any] | None = None
+    benchmark: dict[str, Any] | None = None
+    profile: dict[str, Any] | None = None
+    error: dict[str, Any] | None = None
+    raw_compiler_output: str | None = None
+
+    def to_json(self, indent: int | None = 2) -> str:
+        """Serialize to JSON, excluding None values."""
+        data = {k: v for k, v in asdict(self).items() if v is not None}
+        return json.dumps(data, indent=indent, default=str)
+
+
+@dataclass
+class OutputCollector:
+    """Collects output events and formats them according to the output format."""
+
+    format: OutputFormat
+    target: str | None = None
+    _result: EvalOutput = field(default_factory=lambda: EvalOutput(status="success"))
+
+    def emit(self, event: str, **data: Any) -> None:
+        """Emit an event.
+
+        For JSONL format, prints immediately. For TEXT, prints human-readable.
+        For JSON, events are collected and output at the end.
+        """
+        if self.format == OutputFormat.JSONL:
+            obj = {
+                "event": event,
+                "timestamp": datetime.now(UTC).isoformat(),
+                **data,
+            }
+            print(json.dumps(obj, default=str), flush=True)
+        elif self.format == OutputFormat.TEXT:
+            status = data.get("status", "")
+            if status:
+                typer.echo(f"[wafer] {event}: {status}")
+            else:
+                typer.echo(f"[wafer] {event}")
+
+    def set_result(
+        self,
+        *,
+        status: Literal["success", "failure", "error"],
+        phase: str | None = None,
+        correctness: dict[str, Any] | None = None,
+        benchmark: dict[str, Any] | None = None,
+        profile: dict[str, Any] | None = None,
+        error: dict[str, Any] | None = None,
+        raw_compiler_output: str | None = None,
+    ) -> None:
+        """Set the final result data."""
+        self._result = EvalOutput(
+            status=status,
+            target=self.target,
+            phase=phase,
+            correctness=correctness,
+            benchmark=benchmark,
+            profile=profile,
+            error=error,
+            raw_compiler_output=raw_compiler_output,
+        )
+
+    def set_error(self, phase: str, error_type: str, **details: Any) -> None:
+        """Set an error result."""
+        self._result.status = "error"
+        self._result.phase = phase
+        self._result.error = {"type": error_type, **details}
+        self._result.target = self.target
+
+    def finalize(self) -> None:
+        """Print final output based on format."""
+        if self.format == OutputFormat.JSON:
+            print(self._result.to_json())
+        elif self.format == OutputFormat.JSONL:
+            print(
+                json.dumps(
+                    {
+                        "event": "completed",
+                        "timestamp": datetime.now(UTC).isoformat(),
+                        "result": {k: v for k, v in asdict(self._result).items() if v is not None},
+                    },
+                    default=str,
+                )
+            )
+        # TEXT format already printed incrementally
+
+    def output_text_result(self, result: EvaluateResult) -> None:
+        """Print human-readable result summary (TEXT format only)."""
+        if self.format != OutputFormat.TEXT:
+            return
+
+        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")
+        typer.echo("=" * 60)
+
+    def output_text_error(self, error_message: str) -> None:
+        """Print error message (TEXT format only)."""
+        if self.format == OutputFormat.TEXT:
+            typer.echo(f"Error: {error_message}", err=True)
+
+
+def format_evaluate_result(result: EvaluateResult, target: str | None = None) -> EvalOutput:
+    """Convert EvaluateResult to structured EvalOutput."""
+    if not result.success:
+        # Error case
+        error_info = parse_error_message(result.error_message or "Unknown error")
+        return EvalOutput(
+            status="error",
+            target=target,
+            phase=error_info.get("phase", "unknown"),
+            error=error_info,
+        )
+
+    if not result.all_correct:
+        # Correctness failure
+        return EvalOutput(
+            status="failure",
+            target=target,
+            phase="correctness",
+            correctness={
+                "passed": False,
+                "tests_run": result.total_tests,
+                "tests_passed": result.passed_tests,
+            },
+        )
+
+    # Success
+    output = EvalOutput(
+        status="success",
+        target=target,
+        correctness={
+            "passed": True,
+            "tests_run": result.total_tests,
+            "tests_passed": result.passed_tests,
+        },
+    )
+
+    if result.geomean_speedup > 0:
+        output.benchmark = {"speedup": result.geomean_speedup}
+
+    return output
+
+
+def parse_error_message(error_message: str) -> dict[str, Any]:
+    """Parse error message to extract structured information."""
+    error_info: dict[str, Any] = {"message": error_message}
+
+    # Try to identify the phase and type from common patterns
+    error_lower = error_message.lower()
+
+    if "compilation" in error_lower or "compile" in error_lower:
+        error_info["phase"] = "compilation"
+        error_info["type"] = "CompilationError"
+        # Try to parse compiler error format: file:line:col: error: message
+        parsed = parse_compilation_error(error_message)
+        if parsed:
+            error_info.update(parsed)
+    elif "hsa_status" in error_lower or "memory" in error_lower or "segfault" in error_lower:
+        error_info["phase"] = "runtime"
+        error_info["type"] = "MemoryViolation"
+    elif "timeout" in error_lower:
+        error_info["phase"] = "runtime"
+        error_info["type"] = "Timeout"
+    elif "correctness" in error_lower:
+        error_info["phase"] = "correctness"
+        error_info["type"] = "CorrectnessError"
+    else:
+        error_info["phase"] = "unknown"
+        error_info["type"] = "UnknownError"
+
+    return error_info
+
+
+def parse_compilation_error(raw_output: str) -> dict[str, Any] | None:
+    """Extract structured info from compiler error output.
+
+    Matches patterns like: file.hip:10:14: error: message
+    """
+    match = re.search(
+        r"(?P<file>[\w./]+):(?P<line>\d+):(?P<col>\d+): error: (?P<message>.+)",
+        raw_output,
+    )
+    if match:
+        return {
+            "file": match.group("file"),
+            "line": int(match.group("line")),
+            "column": int(match.group("col")),
+            "message": match.group("message").strip(),
+        }
+    return None
+
+
+def get_output_format(json_flag: bool, jsonl_flag: bool) -> OutputFormat:
+    """Determine output format from CLI flags."""
+    if jsonl_flag:
+        return OutputFormat.JSONL
+    if json_flag:
+        return OutputFormat.JSON
+    return OutputFormat.TEXT

wafer/skills/wafer-guide/SKILL.md

@@ -7,6 +7,19 @@ description: GPU kernel development with the Wafer CLI. Use when working on CUDA
 
 GPU development primitives for optimizing CUDA and HIP kernels.
 
+## Installation
+
+Before using Wafer CLI commands, install the tool:
+
+```bash
+# Install wafer-cli using uv (recommended)
+uv tool install wafer-cli
+
+# Authenticate (one-time setup)
+wafer login
+
+```
+
 ## When to Use This Skill
 
 Activate this skill when:

wafer/ssh_keys.py

@@ -0,0 +1,261 @@
+"""SSH Keys CLI - Manage SSH public keys for workspace access.
+
+This module provides the implementation for the `wafer ssh-keys` subcommand.
+Users register their SSH public keys here, which are then installed in all
+workspaces they attach to (BYOK - Bring Your Own Key model).
+"""
+
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass
+from pathlib import Path
+
+import httpx
+
+from .api_client import get_api_url
+from .auth import get_auth_headers
+
+
+@dataclass(frozen=True)
+class SshKey:
+    """Registered SSH key info."""
+
+    id: str
+    public_key: str
+    name: str | None
+    created_at: str
+
+
+def _get_client() -> tuple[str, dict[str, str]]:
+    """Get API URL and auth headers."""
+    api_url = get_api_url()
+    headers = get_auth_headers()
+
+    assert api_url, "API URL must be configured"
+    assert api_url.startswith("http"), "API URL must be a valid HTTP(S) URL"
+
+    return api_url, headers
+
+
+def _get_key_fingerprint(public_key: str) -> str:
+    """Extract a short fingerprint from a public key for display.
+
+    Returns the first 12 characters of the base64 data portion.
+    """
+    parts = public_key.strip().split()
+    if len(parts) >= 2:
+        return parts[1][:12] + "..."
+    return public_key[:12] + "..."
+
+
+def _get_key_type(public_key: str) -> str:
+    """Extract the key type from a public key."""
+    parts = public_key.strip().split()
+    if parts:
+        return parts[0]
+    return "unknown"
+
+
+def _detect_ssh_keys() -> list[Path]:
+    """Detect existing SSH public keys on disk.
+
+    Returns list of paths to found public key files, in preference order.
+    """
+    ssh_dir = Path.home() / ".ssh"
+    candidates = [
+        "id_ed25519.pub",  # Preferred (modern, secure, fast)
+        "id_rsa.pub",  # Legacy but common
+        "id_ecdsa.pub",  # Less common
+        "id_dsa.pub",  # Deprecated
+    ]
+
+    found = []
+    for filename in candidates:
+        key_path = ssh_dir / filename
+        if key_path.exists():
+            found.append(key_path)
+
+    return found
+
+
+def list_ssh_keys(json_output: bool = False) -> str:
+    """List all registered SSH keys.
+
+    Returns:
+        Formatted output string
+    """
+    api_url, headers = _get_client()
+
+    try:
+        with httpx.Client(timeout=30.0, headers=headers) as client:
+            response = client.get(f"{api_url}/v1/user/ssh-keys")
+            response.raise_for_status()
+            keys = response.json()
+    except httpx.HTTPStatusError as e:
+        if e.response.status_code == 401:
+            raise RuntimeError("Not authenticated. Run: wafer login") from e
+        raise RuntimeError(f"API error: {e.response.status_code} - {e.response.text}") from e
+    except httpx.RequestError as e:
+        raise RuntimeError(f"Could not reach API: {e}") from e
+
+    if json_output:
+        return json.dumps(keys, indent=2)
+
+    if not keys:
+        return (
+            "No SSH keys registered.\n"
+            "\n"
+            "Add your SSH key:\n"
+            "  wafer ssh-keys add\n"
+            "\n"
+            "This will auto-detect your key from ~/.ssh/"
+        )
+
+    lines = ["SSH Keys:"]
+    for key in keys:
+        key_type = _get_key_type(key["public_key"])
+        fingerprint = _get_key_fingerprint(key["public_key"])
+        name = key.get("name") or "(no name)"
+        lines.append(f"  • {name}: {key_type} {fingerprint}")
+        lines.append(f"    ID: {key['id']}")
+
+    return "\n".join(lines)
+
+
+def add_ssh_key(
+    pubkey_path: Path | None = None,
+    name: str | None = None,
+    json_output: bool = False,
+) -> str:
+    """Add an SSH public key.
+
+    Args:
+        pubkey_path: Path to public key file (auto-detects if None)
+        name: Optional friendly name for the key
+        json_output: Return JSON instead of formatted output
+
+    Returns:
+        Formatted output string
+    """
+    # Auto-detect if no path provided
+    if pubkey_path is None:
+        detected = _detect_ssh_keys()
+        if not detected:
+            raise RuntimeError(
+                "No SSH key found in ~/.ssh/\n"
+                "\n"
+                "Generate one with:\n"
+                "  ssh-keygen -t ed25519\n"
+                "\n"
+                "Or specify a path:\n"
+                "  wafer ssh-keys add /path/to/key.pub"
+            )
+        pubkey_path = detected[0]
+
+    # Validate path
+    if not pubkey_path.exists():
+        raise RuntimeError(f"File not found: {pubkey_path}")
+
+    if not pubkey_path.suffix == ".pub" and "pub" not in pubkey_path.name:
+        raise RuntimeError(
+            f"Expected a public key file (.pub), got: {pubkey_path}\n"
+            "\n"
+            "Make sure you're adding the PUBLIC key, not the private key."
+        )
+
+    # Read key content
+    try:
+        public_key = pubkey_path.read_text().strip()
+    except Exception as e:
+        raise RuntimeError(f"Could not read key file: {e}") from e
+
+    # Validate basic format
+    if not public_key.startswith(("ssh-", "ecdsa-", "sk-")):
+        raise RuntimeError(
+            f"Invalid SSH public key format in {pubkey_path}\n"
+            "\n"
+            "Expected OpenSSH format (e.g., 'ssh-ed25519 AAAAC3... user@host')"
+        )
+
+    # Auto-generate name from key type and filename if not provided
+    if name is None:
+        key_type = _get_key_type(public_key)
+        # Use key type without prefix (e.g., "ed25519" instead of "ssh-ed25519")
+        short_type = key_type.replace("ssh-", "").replace("ecdsa-sha2-", "")
+        name = short_type
+
+    # Call API
+    api_url, headers = _get_client()
+    request_body = {
+        "public_key": public_key,
+        "name": name,
+    }
+
+    try:
+        with httpx.Client(timeout=30.0, headers=headers) as client:
+            response = client.post(
+                f"{api_url}/v1/user/ssh-keys",
+                json=request_body,
+            )
+            response.raise_for_status()
+            key_data = response.json()
+    except httpx.HTTPStatusError as e:
+        if e.response.status_code == 401:
+            raise RuntimeError("Not authenticated. Run: wafer login") from e
+        if e.response.status_code == 400:
+            # Parse error detail
+            try:
+                detail = e.response.json().get("detail", e.response.text)
+            except Exception:
+                detail = e.response.text
+            raise RuntimeError(f"Invalid key: {detail}") from e
+        raise RuntimeError(f"API error: {e.response.status_code} - {e.response.text}") from e
+    except httpx.RequestError as e:
+        raise RuntimeError(f"Could not reach API: {e}") from e
+
+    if json_output:
+        return json.dumps(key_data, indent=2)
+
+    key_type = _get_key_type(public_key)
+    fingerprint = _get_key_fingerprint(public_key)
+
+    return (
+        f"✓ SSH key registered: {name}\n"
+        f"  Type: {key_type}\n"
+        f"  Fingerprint: {fingerprint}\n"
+        f"  Source: {pubkey_path}\n"
+        f"\n"
+        f"Your key will be installed in all workspaces you attach to."
+    )
+
+
+def remove_ssh_key(key_id: str, json_output: bool = False) -> str:
+    """Remove an SSH key.
+
+    Args:
+        key_id: UUID of the key to remove
+        json_output: Return JSON instead of formatted output
+
+    Returns:
+        Formatted output string
+    """
+    api_url, headers = _get_client()
+
+    try:
+        with httpx.Client(timeout=30.0, headers=headers) as client:
+            response = client.delete(f"{api_url}/v1/user/ssh-keys/{key_id}")
+            response.raise_for_status()
+    except httpx.HTTPStatusError as e:
+        if e.response.status_code == 401:
+            raise RuntimeError("Not authenticated. Run: wafer login") from e
+        if e.response.status_code == 404:
+            raise RuntimeError(f"SSH key not found: {key_id}") from e
+        raise RuntimeError(f"API error: {e.response.status_code} - {e.response.text}") from e
+    except httpx.RequestError as e:
+        raise RuntimeError(f"Could not reach API: {e}") from e
+
+    if json_output:
+        return json.dumps({"status": "deleted", "key_id": key_id}, indent=2)
+
+    return f"✓ SSH key removed: {key_id}"