wafer-cli 0.2.3__py3-none-any.whl → 0.2.4__py3-none-any.whl

wafer/gpu_run.py

@@ -19,7 +19,10 @@ CONTAINER_WORKSPACE = "/workspace"
 class PushResult:
     """Result of pushing a directory to remote target."""
 
-    workspace_path: str  # Absolute path on remote (tilde-expanded)
+    workspace_name: str  # Just the workspace name (e.g., "project")
+    workspace_path: (
+        str  # Full absolute path on remote (e.g., "/home/user/.wafer/workspaces/project")
+    )
     files_uploaded: list[str]  # Relative paths of uploaded files
 
 
@@ -71,6 +74,7 @@ def push_directory(
             files_uploaded.append(str(file.relative_to(local_path)))
 
     return PushResult(
+        workspace_name=workspace_name,
         workspace_path=expanded_workspace,
         files_uploaded=files_uploaded,
     )

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

wafer/wevin_cli.py

@@ -86,8 +86,21 @@ class StreamingChunkFrontend:
             })
 
         elif isinstance(event, ToolResultReceived):
-            # Emit tool_result event
-            self._emit({"type": "tool_result", "is_error": event.is_error})
+            # Emit tool_result event with error details
+            result_event = {"type": "tool_result", "is_error": event.is_error}
+            # Include error message and content if available
+            if event.error:
+                result_event["error"] = event.error
+            if event.content:
+                # Convert content to string if it's a list
+                if isinstance(event.content, list):
+                    result_event["content"] = "\n".join(
+                        str(item) if not isinstance(item, dict) else item.get("text", str(item))
+                        for item in event.content
+                    )
+                else:
+                    result_event["content"] = str(event.content)
+            self._emit(result_event)
 
         elif isinstance(event, StreamDone):
             # Will be handled by stop()
@@ -240,6 +253,7 @@ def _build_environment(
 ) -> Environment:
     """Build a CodingEnvironment from template config."""
     from wafer_core.environments.coding import CodingEnvironment
+    from wafer_core.rollouts.templates import DANGEROUS_BASH_COMMANDS
 
     working_dir = Path(corpus_path) if corpus_path else Path.cwd()
     resolved_tools = tools_override or tpl.tools
@@ -247,6 +261,7 @@ def _build_environment(
         working_dir=working_dir,
         enabled_tools=resolved_tools,
         bash_allowlist=tpl.bash_allowlist,
+        bash_denylist=DANGEROUS_BASH_COMMANDS,
     )  # type: ignore[assignment]
     return env
 
@@ -347,6 +362,11 @@ def main(  # noqa: PLR0913, PLR0915
                 print(f"  {s.session_id}  {preview}")
         return
 
+    # Emit early event for JSON mode before heavy imports
+    # This gives immediate feedback that the CLI started correctly
+    if json_output:
+        print(json.dumps({"type": "initializing"}), flush=True)
+
     import trio
     from wafer_core.rollouts import FileSessionStore, Message, Trajectory
     from wafer_core.rollouts.frontends import NoneFrontend, RunnerConfig, run_interactive

{wafer_cli-0.2.3.dist-info → wafer_cli-0.2.4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: wafer-cli
-Version: 0.2.3
+Version: 0.2.4
 Summary: CLI tool for running commands on remote GPUs and GPU kernel optimization agent
 Requires-Python: >=3.11
 Requires-Dist: typer>=0.12.0

{wafer_cli-0.2.3.dist-info → wafer_cli-0.2.4.dist-info}/RECORD

@@ -5,29 +5,30 @@ wafer/api_client.py,sha256=cPULiTxqOAYYSfDTNJgd-6Pqrt3IM4Gm9903U7yGIwY,6163
 wafer/auth.py,sha256=ZLsXZ73GDLD8GL7Rij1ELtuLqyJ5EU_uPBUMPVKwExA,10703
 wafer/autotuner.py,sha256=6gH0Ho7T58EFerMQcHQxshWe3DF4qU7fb5xthAh5SPM,44364
 wafer/billing.py,sha256=jbLB2lI4_9f2KD8uEFDi_ixLlowe5hasC0TIZJyIXRg,7163
-wafer/cli.py,sha256=kscK-JBiAu0H7roMdcaZ_lT7xlgR6mUJ1nBruidUdsE,184745
+wafer/cli.py,sha256=19BwYcZo6QfJWfGwkSXNFx__63ZjZM22woL66fNQhIM,200512
 wafer/config.py,sha256=h5Eo9_yfWqWGoPNdVQikI9GoZVUeysunSYiixf1mKcw,3411
 wafer/corpus.py,sha256=yTF3UA5bOa8BII2fmcXf-3WsIsM5DX4etysv0AzVknE,8912
-wafer/evaluate.py,sha256=7E1BqVl0um7kfU2_4FkXZrNc-PQJIHe09G1-L07PXmg,114971
+wafer/evaluate.py,sha256=sGtEkLmZ8wEeC2BRYMoCbCPIlQvuv6tG77910gbIygc,135553
 wafer/global_config.py,sha256=fhaR_RU3ufMksDmOohH1OLeQ0JT0SDW1hEip_zaP75k,11345
-wafer/gpu_run.py,sha256=gUbzMsMPsw3UHcn00bI-zTSHrU8c2FEpDvbcsczlDPo,9557
+wafer/gpu_run.py,sha256=TwqXy72T7f2I7e6n5WWod3xgxCPnDhU0BgLsB4CUoQY,9716
 wafer/inference.py,sha256=tZCO5i05FKY27ewis3CSBHFBeFbXY3xwj0DSjdoMY9s,4314
 wafer/ncu_analyze.py,sha256=rAWzKQRZEY6E_CL3gAWUaW3uZ4kvQVZskVCPDpsFJuE,24633
 wafer/nsys_analyze.py,sha256=dRsYNYp1IqzGSPrQuEMW5vRbIxr-VrQwQbotLSrPvlY,6795
+wafer/problems.py,sha256=ce2sy10A1nnNUG3VGsseTS8jL7LZsku4dE8zVf9JHQ4,11296
 wafer/rocprof_compute.py,sha256=Tu16Vb05b2grvheFWi1XLGlAr6m48NEDeZoDyw_4Uzw,19885
 wafer/rocprof_sdk.py,sha256=fAYCxpfJa5BZTTkIMBOXg4KsYK4i_wNOKrJJn1ZfypM,10086
 wafer/rocprof_systems.py,sha256=4IWbMcbYk1x_8iS7P3FC_u5sgH6EXADCtR2lV9id80M,18629
 wafer/targets.py,sha256=WE5TJgFPGtEIh7VaTQHZ4wB2t4kW0c5K8-UmQ_39Ock,10254
 wafer/tracelens.py,sha256=g9ZIeFyNojZn4uTd3skPqIrRiL7aMJOz_-GOd3aiyy4,7998
-wafer/wevin_cli.py,sha256=tuvHM_wsJ0qo4uUV8rj93vNBL7UgDIMY6f3YXCPf4lg,15821
+wafer/wevin_cli.py,sha256=1_o2P47namZmPkbt47TnyYDmwhEzQYbSg5zjHffu2JQ,16802
 wafer/workspaces.py,sha256=92LG1mtkzNz-ap3XzcqY6KnQ9SUCFG8VBIOUj1Who64,25757
 wafer/skills/wafer-guide/SKILL.md,sha256=UfBeIe5GKFzOYcbPmcs8U2nrjbfr-jSMRwg0jQDBfb0,3058
 wafer/templates/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 wafer/templates/ask_docs.py,sha256=Lxs-faz9v5m4Qa4NjF2X_lE8KwM9ES9MNJkxo7ep56o,2256
 wafer/templates/optimize_kernel.py,sha256=u6AL7Q3uttqlnBLzcoFdsiPq5lV2TV3bgqwCYYlK9gk,2357
 wafer/templates/trace_analyze.py,sha256=XE1VqzVkIUsZbXF8EzQdDYgg-AZEYAOFpr6B_vnRELc,2880
-wafer_cli-0.2.3.dist-info/METADATA,sha256=MKzRmMwc1zhKUVeu_NE4h5_17EFnX1bjLE4Jcox0AAk,559
-wafer_cli-0.2.3.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-wafer_cli-0.2.3.dist-info/entry_points.txt,sha256=WqB7hB__WhtPY8y1cO2sZiUz7fCq6Ik-usAigpeFvWE,41
-wafer_cli-0.2.3.dist-info/top_level.txt,sha256=2MK1IVMWfpLL8BZCQ3E9aG6L6L666gSA_teYlwan4fs,6
-wafer_cli-0.2.3.dist-info/RECORD,,
+wafer_cli-0.2.4.dist-info/METADATA,sha256=OqNyd57hfOUZvBWFd8iMdrfH1VICqcqQf_K1sHZVhfA,559
+wafer_cli-0.2.4.dist-info/WHEEL,sha256=qELbo2s1Yzl39ZmrAibXA2jjPLUYfnVhUNTlyF1rq0Y,92
+wafer_cli-0.2.4.dist-info/entry_points.txt,sha256=WqB7hB__WhtPY8y1cO2sZiUz7fCq6Ik-usAigpeFvWE,41
+wafer_cli-0.2.4.dist-info/top_level.txt,sha256=2MK1IVMWfpLL8BZCQ3E9aG6L6L666gSA_teYlwan4fs,6
+wafer_cli-0.2.4.dist-info/RECORD,,

{wafer_cli-0.2.3.dist-info → wafer_cli-0.2.4.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (80.9.0)
+Generator: setuptools (80.10.1)
 Root-Is-Purelib: true
 Tag: py3-none-any