wafer-cli 0.2.25__tar.gz → 0.2.27__tar.gz

{wafer_cli-0.2.25 → wafer_cli-0.2.27}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: wafer-cli
-Version: 0.2.25
+Version: 0.2.27
 Summary: CLI for running GPU workloads, managing remote workspaces, and evaluating/optimizing kernels
 Requires-Python: >=3.11
 Description-Content-Type: text/markdown

{wafer_cli-0.2.25 → wafer_cli-0.2.27}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "wafer-cli"
-version = "0.2.25"
+version = "0.2.27"
 description = "CLI for running GPU workloads, managing remote workspaces, and evaluating/optimizing kernels"
 readme = "README.md"
 requires-python = ">=3.11"

{wafer_cli-0.2.25 → wafer_cli-0.2.27}/wafer/cli.py

@@ -268,6 +268,32 @@ Configure targets with: wafer config targets init ..."""
 )
 app.add_typer(targets_ops_app, name="targets", rich_help_panel="Infrastructure")
 
+# Specs management (new: local TOML configs)
+from wafer.specs_cli import specs_app
+
+app.add_typer(specs_app, name="specs", rich_help_panel="Configuration")
+
+# Live resource management (new: API-backed commands on `wafer targets`)
+# These become: wafer targets list, wafer targets terminate, etc.
+from wafer.targets_cli import (
+    targets_list as _targets_list_cmd,
+)
+from wafer.targets_cli import (
+    targets_provision as _targets_provision_cmd,
+)
+from wafer.targets_cli import (
+    targets_reconcile as _targets_reconcile_cmd,
+)
+from wafer.targets_cli import (
+    targets_terminate as _targets_terminate_cmd,
+)
+from wafer.targets_cli import (
+    targets_pools as _targets_pools_cmd,
+)
+from wafer.targets_cli import (
+    targets_probe as _targets_probe_cmd,
+)
+
 # Billing management - nested under config
 billing_app = typer.Typer(help="Manage billing, credits, and subscription")
 config_app.add_typer(billing_app, name="billing")
@@ -612,7 +638,9 @@ def skill_status() -> None:
 auth_app = typer.Typer(help="Authenticate with Wafer and cloud GPU providers")
 app.add_typer(auth_app, name="auth", rich_help_panel="Configuration")
 
-providers_app = typer.Typer(help="Manage API keys for cloud GPU providers (RunPod, DigitalOcean, etc.)")
+providers_app = typer.Typer(
+    help="Manage API keys for cloud GPU providers (RunPod, DigitalOcean, etc.)"
+)
 auth_app.add_typer(providers_app, name="providers")
 
 
@@ -1813,6 +1841,93 @@ def kernelbench_list_problems() -> None:
         raise typer.Exit(1) from None
 
 
+def _resolve_pool_query(pool: str, collector) -> tuple[str, object]:
+    """Resolve a PoolQuery pool to a target spec name + lock context.
+
+    Queries live providers, matches by pool query, locks one target,
+    returns (spec_name, lock_context) for the evaluator.
+    """
+    import trio
+    from wafer_core.targets.pool import resolve_pool
+
+    from .target_lock import acquire_from_pool
+
+    matched_targets = trio.run(resolve_pool, pool)
+
+    if not matched_targets:
+        collector.set_error("pool", "NoMatchingTargets", pool=pool)
+        collector.finalize()
+        raise typer.Exit(1)
+
+    # Filter to targets with a spec (evaluator needs spec fields)
+    spec_targets = [t for t in matched_targets if t.spec_name]
+    if not spec_targets:
+        collector.set_error(
+            "pool", "NoSpecTargets", pool=pool,
+            message="Matched targets have no spec binding — evaluator needs spec fields",
+        )
+        collector.finalize()
+        raise typer.Exit(1)
+
+    # Lock one by resource_id
+    resource_ids = [t.resource_id for t in spec_targets]
+    collector.emit("pool_acquire", pool=pool, count=len(resource_ids))
+
+    lock_ctx = acquire_from_pool(resource_ids)
+    acquired_id = lock_ctx.__enter__()
+
+    if acquired_id is None:
+        lock_ctx.__exit__(None, None, None)
+        collector.set_error("pool", "AllTargetsBusy", pool=pool, targets=resource_ids)
+        collector.finalize()
+        raise typer.Exit(1)
+
+    # Map resource_id back to spec_name
+    acquired_target = next(t for t in spec_targets if t.resource_id == acquired_id)
+    spec_name = acquired_target.spec_name
+
+    collector.emit("pool_acquired", target=spec_name, resource_id=acquired_id)
+    return spec_name, lock_ctx
+
+
+def _resolve_pool_legacy(pool: str, collector) -> tuple[str, object]:
+    """Resolve an old-style pool (static target name list) to a target name + lock context.
+
+    Old format: [pools.name] targets = ["t1", "t2"]
+    """
+    from .target_lock import acquire_from_pool
+    from .targets import filter_pool_by_auth, get_pool
+
+    try:
+        pool_targets = get_pool(pool)
+    except FileNotFoundError as e:
+        collector.set_error("pool", "PoolNotFound", pool=pool, message=str(e))
+        collector.finalize()
+        raise typer.Exit(1) from None
+
+    usable_targets, skipped = filter_pool_by_auth(pool_targets)
+    if skipped:
+        collector.emit("pool_auth_skip", targets=skipped)
+
+    if not usable_targets:
+        collector.set_error("pool", "NoUsableTargets", pool=pool)
+        collector.finalize()
+        raise typer.Exit(1) from None
+
+    collector.emit("pool_acquire", pool=pool, count=len(usable_targets))
+    lock_ctx = acquire_from_pool(usable_targets)
+    acquired_target = lock_ctx.__enter__()
+
+    if acquired_target is None:
+        lock_ctx.__exit__(None, None, None)
+        collector.set_error("pool", "AllTargetsBusy", pool=pool, targets=usable_targets)
+        collector.finalize()
+        raise typer.Exit(1)
+
+    collector.emit("pool_acquired", target=acquired_target)
+    return acquired_target, lock_ctx
+
+
 @kernelbench_app.callback(invoke_without_command=True)
 def kernelbench_evaluate(  # noqa: PLR0913, PLR0915
     ctx: typer.Context,
@@ -1943,39 +2058,12 @@ def kernelbench_evaluate(  # noqa: PLR0913, PLR0915
     pool_lock_context = None
 
     if pool:
-        from .target_lock import acquire_from_pool
-        from .targets import filter_pool_by_auth, get_pool
+        from wafer_core.targets.pool import is_query_pool
 
-        try:
-            pool_targets = get_pool(pool)
-        except FileNotFoundError as e:
-            collector.set_error("pool", "PoolNotFound", pool=pool, message=str(e))
-            collector.finalize()
-            raise typer.Exit(1) from None
-
-        # Filter to only targets with valid auth
-        usable_targets, skipped = filter_pool_by_auth(pool_targets)
-        if skipped:
-            collector.emit("pool_auth_skip", targets=skipped)
-
-        if not usable_targets:
-            collector.set_error("pool", "NoUsableTargets", pool=pool)
-            collector.finalize()
-            raise typer.Exit(1) from None
-
-        collector.emit("pool_acquire", pool=pool, count=len(usable_targets))
-        pool_lock_context = acquire_from_pool(usable_targets)
-        acquired_target = pool_lock_context.__enter__()
-
-        if acquired_target is None:
-            # Exit context manager before raising to avoid resource leak
-            pool_lock_context.__exit__(None, None, None)
-            collector.set_error("pool", "AllTargetsBusy", pool=pool, targets=usable_targets)
-            collector.finalize()
-            raise typer.Exit(1)
-
-        collector.emit("pool_acquired", target=acquired_target)
-        resolved_target = acquired_target
+        if is_query_pool(pool):
+            resolved_target, pool_lock_context = _resolve_pool_query(pool, collector)
+        else:
+            resolved_target, pool_lock_context = _resolve_pool_legacy(pool, collector)
 
     collector.target = resolved_target
 
@@ -5254,6 +5342,18 @@ def workspaces_pull(
         raise typer.Exit(1) from None
 
 
+# =============================================================================
+# Live resource commands (list/terminate/reconcile/provision)
+# =============================================================================
+
+targets_ops_app.command("list")(_targets_list_cmd)
+targets_ops_app.command("terminate")(_targets_terminate_cmd)
+targets_ops_app.command("reconcile")(_targets_reconcile_cmd)
+targets_ops_app.command("provision")(_targets_provision_cmd)
+targets_ops_app.command("pools")(_targets_pools_cmd)
+targets_ops_app.command("probe")(_targets_probe_cmd)
+
+
 # =============================================================================
 # Target operations commands (exec/ssh/sync)
 # =============================================================================
@@ -7787,6 +7887,9 @@ def compare_analyze(
     stack_traces: bool = typer.Option(
         False, "--stack-traces", help="Show Python stack traces for operations"
     ),
+    recommendations: bool = typer.Option(
+        False, "--recommendations", help="Generate prioritized recommendations for kernel team"
+    ),
     json: bool = typer.Option(
         False, "--json", hidden=True, help="Ignored (for compatibility with cliExecutor)"
     ),
@@ -7839,6 +7942,7 @@ def compare_analyze(
         show_layers=layers,
         show_all=all,
         show_stack_traces=stack_traces,
+        recommendations=recommendations,
     )
     _mark_command_success()
 
@@ -7883,14 +7987,69 @@ def compare_fusion_cmd(
         # CSV output to file
         wafer compare fusion amd_trace.json nvidia_trace.json --format csv -o fusion.csv
     """
-    from .trace_compare import compare_fusion
+    from .trace_compare import compare_align
+    
+    compare_align(
+        trace1=trace1,
+        trace2=trace2,
+        output=output,
+        output_format=format,
+        phase="all",
+    )
+    _mark_command_success()
 
-    compare_fusion(
+
+@compare_app.command("align")
+def compare_align_cmd(
+    trace1: Path = typer.Argument(..., help="First trace file (AMD or NVIDIA)", exists=True),
+    trace2: Path = typer.Argument(..., help="Second trace file (AMD or NVIDIA)", exists=True),
+    format: str = typer.Option(
+        "json",
+        "--format",
+        "-f",
+        help="Output format: json",
+    ),
+    output: Path | None = typer.Option(
+        None, "--output", "-o", help="Output file (default: stdout)"
+    ),
+    phase: str = typer.Option(
+        "all",
+        "--phase",
+        help="Filter by phase: all, prefill, decode",
+    ),
+    layer: int | None = typer.Option(
+        None,
+        "--layer",
+        help="Focus on specific layer number",
+    ),
+) -> None:
+    """Align kernels at layer level for exact kernel-to-kernel comparison.
+
+    Provides kernel-to-kernel mapping across AMD and NVIDIA platforms,
+    showing which kernels correspond to each other at each layer position.
+
+    Examples:
+        # Basic alignment (stdout JSON)
+        wafer compare align amd_trace.json nvidia_trace.json
+
+        # Save to file
+        wafer compare align amd_trace.json nvidia_trace.json -o alignment.json
+
+        # Focus on decode phase only
+        wafer compare align amd_trace.json nvidia_trace.json --phase decode
+
+        # Focus on specific layer
+        wafer compare align amd_trace.json nvidia_trace.json --layer 5
+    """
+    from .trace_compare import compare_align
+    
+    compare_align(
         trace1=trace1,
         trace2=trace2,
         output=output,
-        format_type=format,
-        min_group_size=min_group_size,
+        output_format=format,
+        phase=phase,
+        layer=layer,
     )
     _mark_command_success()
 

{wafer_cli-0.2.25 → wafer_cli-0.2.27}/wafer/corpus.py

@@ -109,14 +109,34 @@ CORPORA: dict[CorpusName, CorpusConfig] = {
     ),
     "hip": CorpusConfig(
         name="hip",
-        description="HIP programming guide and API reference",
-        source_type="github_repo",
-        repo="ROCm/HIP",
-        repo_paths=["docs"],
+        description="HIP programming guide, API reference, and examples",
+        source_type="github_multi_repo",
+        repos=[
+            # HIP - main documentation and API
+            RepoSource(
+                repo="ROCm/HIP",
+                paths=["docs"],
+            ),
+            # HIP examples - code samples
+            RepoSource(
+                repo="ROCm/HIP-Examples",
+                paths=["HIP-Examples-Applications", "mini-nbody"],
+            ),
+            # clr - HIP/OpenCL runtime (low-level)
+            RepoSource(
+                repo="ROCm/clr",
+                paths=["hipamd/include", "rocclr/device/gpu"],
+            ),
+            # ROCm docs - official documentation
+            RepoSource(
+                repo="ROCm/ROCm",
+                paths=["docs"],
+            ),
+        ],
     ),
     "amd": CorpusConfig(
         name="amd",
-        description="AMD GPU kernel development (rocWMMA, CK, AITER, rocBLAS, HipKittens, vLLM)",
+        description="AMD GPU kernel development (rocWMMA, CK, AITER, rocBLAS, HipKittens, vLLM, FlashAttention)",
         source_type="github_multi_repo",
         repos=[
             # rocWMMA - wave matrix multiply-accumulate (WMMA) intrinsics
@@ -186,6 +206,46 @@ CORPORA: dict[CorpusName, CorpusConfig] = {
                 repo="huggingface/hf-rocm-kernels",
                 paths=["csrc", "hf_rocm_kernels", "docs"],
             ),
+            # ROCm/flash-attention - FlashAttention for AMD GPUs
+            RepoSource(
+                repo="ROCm/flash-attention",
+                paths=["csrc", "docs"],
+            ),
+            # ROCm/triton - Triton compiler for AMD GPUs
+            RepoSource(
+                repo="ROCm/triton",
+                paths=["python/tutorials", "third_party/amd"],
+            ),
+            # ROCm/rccl - ROCm Communication Collectives Library (multi-GPU)
+            RepoSource(
+                repo="ROCm/rccl",
+                paths=["docs"],
+            ),
+            # ROCm/rocprofiler-sdk - AMD GPU profiling SDK
+            RepoSource(
+                repo="ROCm/rocprofiler-sdk",
+                paths=["docs", "samples"],
+            ),
+            # ROCm/omniperf - AMD GPU profiling tool
+            RepoSource(
+                repo="ROCm/omniperf",
+                paths=["docs", "src/omniperf_analyze"],
+            ),
+            # ROCm/omnitrace - Application tracing for AMD
+            RepoSource(
+                repo="ROCm/omnitrace",
+                paths=["docs"],
+            ),
+            # AMD GPUOpen Performance Guides
+            RepoSource(
+                repo="GPUOpen-Tools/gpu_performance_api",
+                paths=["docs"],
+            ),
+            # AMD LLVM - AMD GPU compiler backend
+            RepoSource(
+                repo="ROCm/llvm-project",
+                paths=["amd/device-libs/README.md", "llvm/docs/AMDGPUUsage.rst"],
+            ),
         ],
     ),
 }

wafer_cli-0.2.27/wafer/specs_cli.py

@@ -0,0 +1,157 @@
+"""CLI commands for wafer specs — TargetSpec TOML management.
+
+These are the local config commands (no API calls).
+Registered as: wafer specs list|show|add|remove|default|init
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import typer
+
+specs_app = typer.Typer(
+    help="""Manage GPU target specs (provisioning blueprints).
+
+Specs define how to access or provision GPUs. They are TOML files in ~/.wafer/specs/.
+
+  wafer specs list                        # List all specs
+  wafer specs show runpod-mi300x          # Show one spec
+  wafer specs add /path/to/spec.toml      # Add from file
+  wafer specs remove old-target           # Remove a spec
+  wafer specs default runpod-mi300x       # Set default
+
+To create a new spec interactively:
+  wafer config targets init ssh           # (legacy, still works)
+  wafer config targets init runpod
+"""
+)
+
+
+@specs_app.command("list")
+def specs_list() -> None:
+    """List all configured specs.
+
+    Example:
+        wafer specs list
+    """
+    from wafer_core.targets.spec_store import list_spec_names, load_spec
+
+    from .targets import get_default_target
+
+    names = list_spec_names()
+    default = get_default_target()
+
+    if not names:
+        typer.echo("No specs configured.")
+        typer.echo("Add one with: wafer specs add <path/to/spec.toml>")
+        typer.echo("Or interactively: wafer config targets init ssh")
+        return
+
+    typer.echo("Configured specs:")
+    for name in names:
+        marker = " (default)" if name == default else ""
+        try:
+            spec = load_spec(name)
+            type_name = type(spec).__name__.replace("Target", "")
+            typer.echo(f"  {name}{marker}  [{type_name}]  gpu={spec.gpu_type}")
+        except Exception as e:
+            typer.echo(f"  {name}{marker}  [error: {e}]")
+
+
+@specs_app.command("show")
+def specs_show(
+    name: str = typer.Argument(..., help="Spec name"),
+) -> None:
+    """Show details for a spec.
+
+    Example:
+        wafer specs show runpod-mi300x
+    """
+    from wafer_core.targets.spec_store import load_spec
+
+    from .targets import get_target_info
+
+    try:
+        spec = load_spec(name)
+    except FileNotFoundError:
+        typer.echo(f"Spec not found: {name}", err=True)
+        raise typer.Exit(1) from None
+
+    typer.echo(f"Spec: {name}")
+    for key, value in get_target_info(spec).items():
+        typer.echo(f"  {key}: {value}")
+
+
+@specs_app.command("add")
+def specs_add(
+    file_path: Path = typer.Argument(..., help="Path to TOML spec file"),
+) -> None:
+    """Add a spec from a TOML file.
+
+    Example:
+        wafer specs add ./my-target.toml
+    """
+    import tomllib
+
+    from wafer_core.targets.spec_store import parse_spec, save_spec
+
+    if not file_path.exists():
+        typer.echo(f"File not found: {file_path}", err=True)
+        raise typer.Exit(1) from None
+
+    try:
+        with open(file_path, "rb") as f:
+            data = tomllib.load(f)
+        spec = parse_spec(data)
+        save_spec(spec)
+        typer.echo(f"Added spec: {spec.name}")
+    except Exception as e:
+        typer.echo(f"Error: {e}", err=True)
+        raise typer.Exit(1) from None
+
+
+@specs_app.command("remove")
+def specs_remove(
+    name: str = typer.Argument(..., help="Spec name to remove"),
+    force: bool = typer.Option(False, "--force", "-f", help="Skip confirmation"),
+) -> None:
+    """Remove a spec.
+
+    Example:
+        wafer specs remove old-target
+    """
+    from wafer_core.targets.spec_store import remove_spec
+
+    if not force:
+        confirm = typer.confirm(f"Remove spec '{name}'?")
+        if not confirm:
+            return
+
+    try:
+        remove_spec(name)
+        typer.echo(f"Removed spec: {name}")
+    except FileNotFoundError:
+        typer.echo(f"Spec not found: {name}", err=True)
+        raise typer.Exit(1) from None
+
+
+@specs_app.command("default")
+def specs_default(
+    name: str = typer.Argument(..., help="Spec name to set as default"),
+) -> None:
+    """Set the default spec.
+
+    Example:
+        wafer specs default runpod-mi300x
+    """
+    from wafer_core.targets.spec_store import list_spec_names
+
+    from .targets import set_default_target
+
+    if name not in list_spec_names():
+        typer.echo(f"Spec not found: {name}", err=True)
+        raise typer.Exit(1) from None
+
+    set_default_target(name)
+    typer.echo(f"Default spec set to: {name}")