wafer-cli 0.2.14__py3-none-any.whl

wafer/output.py

@@ -0,0 +1,248 @@
+"""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)
+        # Handle None (correctness not run), True (pass), False (fail)
+        if result.all_correct is None:
+            status = "OK"  # Correctness wasn't checked (e.g., compile-only or prepare-only)
+        elif result.all_correct:
+            status = "PASS"
+        else:
+            status = "FAIL"
+        typer.echo(f"Result: {status}")
+        if result.total_tests > 0:
+            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/problems.py

@@ -0,0 +1,357 @@
+"""Problem set management for Wafer CLI.
+
+Download and manage kernel optimization problem sets for evaluation.
+Follows the same pattern as corpus.py for consistency.
+"""
+
+import shutil
+import tarfile
+import tempfile
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Literal
+
+import httpx
+
+PROBLEMS_CACHE_DIR = Path.home() / ".cache" / "wafer" / "problems"
+
+ProblemSetName = Literal["kernelbench", "gpumode"]
+
+
+@dataclass
+class ProblemSetConfig:
+    """Configuration for a downloadable problem set."""
+
+    name: ProblemSetName
+    description: str
+    repo: str  # GitHub repo in "owner/repo" format
+    repo_paths: list[str]  # Paths within repo to download
+    format_description: str  # Brief description of the format
+
+
+PROBLEM_SETS: dict[ProblemSetName, ProblemSetConfig] = {
+    "kernelbench": ProblemSetConfig(
+        name="kernelbench",
+        description="KernelBench GPU kernel optimization problems (level1-4)",
+        repo="ScalingIntelligence/KernelBench",
+        repo_paths=["KernelBench"],
+        format_description="Class-based: Model/ModelNew with get_inputs/get_init_inputs",
+    ),
+    "gpumode": ProblemSetConfig(
+        name="gpumode",
+        description="GPU Mode reference kernels (pmpp, amd, nvidia, bioml)",
+        repo="gpu-mode/reference-kernels",
+        repo_paths=["problems"],
+        format_description="Functional: ref_kernel/custom_kernel with generate_input",
+    ),
+}
+
+
+def _problems_path(name: ProblemSetName) -> Path:
+    """Get local path for problem set."""
+    return PROBLEMS_CACHE_DIR / name
+
+
+def _ensure_cache_dir() -> None:
+    """Ensure cache directory exists."""
+    PROBLEMS_CACHE_DIR.mkdir(parents=True, exist_ok=True)
+
+
+def _download_github_repo(config: ProblemSetConfig, dest: Path, verbose: bool = True) -> int:
+    """Download specific paths from GitHub repo.
+
+    Args:
+        config: Problem set configuration
+        dest: Destination directory
+        verbose: Print progress
+
+    Returns:
+        Number of files downloaded
+    """
+    # Fetch tarball from GitHub
+    resp = _fetch_github_tarball(config.repo, verbose)
+
+    # Save to temp file
+    with tempfile.NamedTemporaryFile(suffix=".tar.gz", delete=False) as tmp:
+        tmp.write(resp.content)
+        tmp_path = Path(tmp.name)
+
+    # Extract matching files
+    try:
+        downloaded = _extract_tarball(tmp_path, dest, config.repo_paths, verbose)
+    finally:
+        tmp_path.unlink()
+
+    return downloaded
+
+
+def _fetch_github_tarball(repo: str, verbose: bool) -> httpx.Response:
+    """Fetch tarball from GitHub, trying main then master branch."""
+    with httpx.Client(timeout=120.0, follow_redirects=True) as client:
+        for branch in ["main", "master"]:
+            tarball_url = f"https://api.github.com/repos/{repo}/tarball/{branch}"
+            if verbose:
+                print(f"  Fetching {repo} ({branch} branch)...")
+            try:
+                resp = client.get(tarball_url)
+                resp.raise_for_status()
+                return resp
+            except httpx.HTTPStatusError:
+                if branch == "master":
+                    raise
+    raise RuntimeError(f"Failed to fetch tarball from {repo}")  # Should not reach
+
+
+def _extract_tarball(tmp_path: Path, dest: Path, repo_paths: list[str], verbose: bool) -> int:
+    """Extract files from tarball matching repo_paths."""
+    downloaded = 0
+    with tarfile.open(tmp_path, "r:gz") as tar:
+        for member in tar.getmembers():
+            if not member.isfile():
+                continue
+            # Strip the root directory (e.g., "ScalingIntelligence-KernelBench-abc123/")
+            rel_path = "/".join(member.name.split("/")[1:])
+            if not _matches_repo_paths(rel_path, repo_paths):
+                continue
+            target = dest / rel_path
+            target.parent.mkdir(parents=True, exist_ok=True)
+            extracted = tar.extractfile(member)
+            if extracted:
+                target.write_bytes(extracted.read())
+                downloaded += 1
+                if verbose and downloaded <= 10:
+                    print(f"  ✓ {rel_path}")
+        if verbose and downloaded > 10:
+            print(f"  ... and {downloaded - 10} more files")
+    return downloaded
+
+
+def _matches_repo_paths(rel_path: str, repo_paths: list[str]) -> bool:
+    """Check if rel_path starts with any of the repo_paths."""
+    return any(rel_path.startswith(rp) for rp in repo_paths)
+
+
+def download_problems(name: ProblemSetName, force: bool = False, verbose: bool = True) -> Path:
+    """Download a problem set to local cache.
+
+    Args:
+        name: Problem set name
+        force: Re-download even if exists
+        verbose: Print progress
+
+    Returns:
+        Path to downloaded problem set
+
+    Raises:
+        ValueError: If problem set name is unknown
+        httpx.HTTPError: If download fails
+    """
+    if name not in PROBLEM_SETS:
+        raise ValueError(f"Unknown problem set: {name}. Available: {list(PROBLEM_SETS.keys())}")
+
+    config = PROBLEM_SETS[name]
+    dest = _problems_path(name)
+
+    if dest.exists() and not force:
+        if verbose:
+            print(f"Problem set '{name}' already exists at {dest}")
+            print("Use --force to re-download")
+        return dest
+
+    _ensure_cache_dir()
+
+    if dest.exists():
+        shutil.rmtree(dest)
+    dest.mkdir(parents=True)
+
+    if verbose:
+        print(f"Downloading {name}: {config.description}")
+
+    try:
+        count = _download_github_repo(config, dest, verbose)
+    except Exception:
+        # Clean up partial download so next run doesn't skip with stale cache
+        if dest.exists():
+            shutil.rmtree(dest)
+        raise
+
+    if verbose:
+        print(f"Downloaded {count} files to {dest}")
+
+    return dest
+
+
+def get_problems_path(name: ProblemSetName) -> Path | None:
+    """Get path to downloaded problem set, or None if not downloaded.
+
+    Args:
+        name: Problem set name
+
+    Returns:
+        Path if downloaded, None otherwise
+    """
+    if name not in PROBLEM_SETS:
+        return None
+    path = _problems_path(name)
+    return path if path.exists() else None
+
+
+def list_problem_sets(verbose: bool = True) -> dict[ProblemSetName, bool]:
+    """List available problem sets and their download status.
+
+    Returns:
+        Dict of problem set name -> is_downloaded
+    """
+    result: dict[ProblemSetName, bool] = {}
+    for name, config in PROBLEM_SETS.items():
+        path = _problems_path(name)
+        exists = path.exists()
+        result[name] = exists
+        if verbose:
+            status = "✓" if exists else " "
+            print(f"[{status}] {name}: {config.description}")
+            print(f"    Format: {config.format_description}")
+            if exists:
+                file_count = sum(1 for _ in path.rglob("*.py") if _.is_file())
+                print(f"    Location: {path} ({file_count} Python files)")
+    return result
+
+
+def list_problems(name: ProblemSetName, verbose: bool = True) -> list[str]:
+    """List available problems in a problem set.
+
+    Args:
+        name: Problem set name
+        verbose: Print to stdout
+
+    Returns:
+        List of problem IDs
+
+    Raises:
+        ValueError: If problem set not downloaded
+    """
+    path = get_problems_path(name)
+    if path is None:
+        raise ValueError(
+            f"Problem set '{name}' is not downloaded. Run:\n  wafer evaluate {name} download"
+        )
+
+    if name == "kernelbench":
+        problems = _list_kernelbench_problems(path)
+    elif name == "gpumode":
+        problems = _list_gpumode_problems(path)
+    else:
+        problems = []
+
+    if verbose:
+        if not problems:
+            print(f"No problems found in {name}")
+        else:
+            print(f"Available problems in {name} ({len(problems)} total):\n")
+            for p in problems:
+                print(f"  {p}")
+
+    return problems
+
+
+def _list_kernelbench_problems(path: Path) -> list[str]:
+    """List KernelBench problems: level1/1_Name.py format."""
+    problems: list[str] = []
+    kb_root = path / "KernelBench"
+    if not kb_root.exists():
+        kb_root = path  # In case structure is flat
+
+    for level_dir in sorted(kb_root.iterdir()):
+        if not (level_dir.is_dir() and level_dir.name.startswith("level")):
+            continue
+        for problem_file in sorted(level_dir.glob("*.py")):
+            if problem_file.name.startswith("__"):
+                continue
+            problem_id = f"{level_dir.name}/{problem_file.stem}"
+            problems.append(problem_id)
+    return problems
+
+
+def _list_gpumode_problems(path: Path) -> list[str]:
+    """List GPUMode problems: category/problem_name format."""
+    problems: list[str] = []
+    problems_root = path / "problems"
+    if not problems_root.exists():
+        problems_root = path
+
+    for category_dir in sorted(problems_root.iterdir()):
+        if not _is_valid_problem_dir(category_dir):
+            continue
+        for problem_dir in sorted(category_dir.iterdir()):
+            if not _is_valid_problem_dir(problem_dir):
+                continue
+            # Check if it has the expected files
+            has_reference = (problem_dir / "reference.py").exists()
+            has_task = (problem_dir / "task.yml").exists()
+            if has_reference or has_task:
+                problem_id = f"{category_dir.name}/{problem_dir.name}"
+                problems.append(problem_id)
+    return problems
+
+
+def _is_valid_problem_dir(path: Path) -> bool:
+    """Check if path is a valid problem directory (not hidden/special)."""
+    return path.is_dir() and not path.name.startswith((".", "_"))
+
+
+def get_problem_path(name: ProblemSetName, problem_id: str) -> Path | None:
+    """Get path to a specific problem.
+
+    Args:
+        name: Problem set name
+        problem_id: Problem ID (e.g., "level4/103" or "pmpp/vectoradd_py")
+
+    Returns:
+        Path to problem file/directory, or None if not found
+    """
+    base_path = get_problems_path(name)
+    if base_path is None:
+        return None
+
+    if name == "kernelbench":
+        # Parse problem_id like "level4/103" or "level4/103_GroupedQueryAttention"
+        parts = problem_id.split("/")
+        if len(parts) != 2:
+            return None
+
+        level_str, problem_part = parts
+        if not level_str.startswith("level"):
+            level_str = f"level{level_str}"
+
+        kb_root = base_path / "KernelBench"
+        if not kb_root.exists():
+            kb_root = base_path
+
+        problem_dir = kb_root / level_str
+        if not problem_dir.exists():
+            return None
+
+        # Find matching problem file
+        problem_files = list(problem_dir.glob(f"{problem_part}*.py"))
+        if not problem_files:
+            # Try exact match
+            exact = problem_dir / f"{problem_part}.py"
+            if exact.exists():
+                return exact
+            return None
+
+        return problem_files[0]
+
+    elif name == "gpumode":
+        # Parse problem_id like "pmpp/vectoradd_py"
+        problems_root = base_path / "problems"
+        if not problems_root.exists():
+            problems_root = base_path
+
+        problem_path = problems_root / problem_id
+        if problem_path.exists() and problem_path.is_dir():
+            return problem_path
+
+        return None
+
+    return None