wafer-cli 0.2.14__py3-none-any.whl → 0.2.30__py3-none-any.whl

wafer/tests/test_eval_cli_parity.py

@@ -0,0 +1,199 @@
+"""Test that eval and CLI agent configs are in sync.
+
+Run: python -m pytest apps/wafer-cli/wafer/tests/test_eval_cli_parity.py -v
+  or: python apps/wafer-cli/wafer/tests/test_eval_cli_parity.py  (standalone)
+  or: python apps/wafer-cli/wafer/tests/test_eval_cli_parity.py --dump  (side-by-side)
+
+Checks:
+- Same bash allowlist (identity check — must be the same object)
+- Same enabled tools
+- Same system prompt text (modulo runtime variables like pool name)
+- CLI instructions generated from same allowlist
+
+Coverage notes:
+  These tests verify config-level parity. After composition, both paths
+  flow through the same codepath (Actor → run_agent_step → rollout →
+  anthropic.py provider) with no further system prompt modifications.
+
+  Known infra-level differences NOT tested here:
+  - Claude Code identity prefix: anthropic.py prepends "You are Claude Code..."
+    when using OAuth/Claude Code API keys. Eval typically uses raw API keys,
+    so eval agents may not get this prefix. This is an auth concern, not a
+    prompt content concern.
+  - Skills layer: CLI path (wevin_cli.py) can append skill metadata if
+    template.include_skills is True. Currently False for optimize-kernelbench,
+    so this is a no-op. Eval path doesn't have this layer at all.
+"""
+
+from __future__ import annotations
+
+import difflib
+import string
+import sys
+
+
+def _normalize_prompt(prompt: str) -> str:
+    """Normalize runtime variables so eval and CLI prompts are comparable."""
+    # Replace pool/target specifics with placeholders
+    prompt = prompt.replace("--pool mi300x-pool", "--pool POOL")
+    prompt = prompt.replace("--pool kernelbench-pool", "--pool POOL")
+    prompt = prompt.replace("--target mi300x", "--target TARGET")
+    return prompt
+
+
+def test_bash_allowlist_parity() -> None:
+    from wafer.agent_defaults import KERNELBENCH_BASH_ALLOWLIST
+    from wafer.templates.optimize_kernelbench import template
+
+    # Template should use the shared allowlist (same object)
+    assert template.bash_allowlist is KERNELBENCH_BASH_ALLOWLIST, (
+        "Template bash_allowlist is not the shared KERNELBENCH_BASH_ALLOWLIST. "
+        "Import it from wafer.agent_defaults instead of defining a local copy."
+    )
+
+    # Eval should also alias to the shared allowlist
+    from optimize_kernelbench_eval.base_config import BASH_ALLOWLIST
+
+    assert BASH_ALLOWLIST is KERNELBENCH_BASH_ALLOWLIST, (
+        "Eval BASH_ALLOWLIST is not the shared KERNELBENCH_BASH_ALLOWLIST. "
+        "Import it from wafer.agent_defaults instead of defining a local copy."
+    )
+
+
+def test_enabled_tools_parity() -> None:
+    from wafer.agent_defaults import ENABLED_TOOLS
+    from wafer.templates.optimize_kernelbench import template
+
+    assert template.tools == ENABLED_TOOLS, (
+        f"Template tools {template.tools} != shared ENABLED_TOOLS {ENABLED_TOOLS}"
+    )
+
+
+def test_system_prompt_parity() -> None:
+    """The task-specific system prompt should be identical between eval and CLI
+    (after normalizing runtime variables like pool name)."""
+    from optimize_kernelbench_eval.base_config import SYSTEM_PROMPT as EVAL_PROMPT
+
+    from wafer.templates.optimize_kernelbench import template
+
+    # Format eval prompt with HIP defaults (most common)
+    eval_formatted = EVAL_PROMPT.format(
+        backend="HIP",
+        backend_lower="hip",
+        target_flag="--pool POOL",
+        reference_path="<reference_file>",
+    )
+
+    # Format CLI template prompt with matching defaults
+    params = dict(template.defaults)
+    params["target_flag"] = "--pool POOL"
+    cli_formatted = string.Template(template.system_prompt).safe_substitute(**params)
+
+    eval_normalized = _normalize_prompt(eval_formatted)
+    cli_normalized = _normalize_prompt(cli_formatted)
+
+    if eval_normalized != cli_normalized:
+        diff = "\n".join(
+            difflib.unified_diff(
+                eval_normalized.splitlines(),
+                cli_normalized.splitlines(),
+                fromfile="eval (base_config.py)",
+                tofile="cli (optimize_kernelbench.py template)",
+                lineterm="",
+                n=2,
+            )
+        )
+        raise AssertionError(
+            f"System prompts differ between eval and CLI template:\n\n{diff}\n\n"
+            "Both should define the same task instructions. "
+            "Edit one to match the other."
+        )
+
+
+def test_cli_instructions_identical() -> None:
+    """Both paths should generate the same CLI instructions
+    (since they use the same bash_allowlist)."""
+    from wafer.agent_defaults import KERNELBENCH_BASH_ALLOWLIST
+    from wafer.cli_instructions import build_cli_instructions
+    from wafer.templates.optimize_kernelbench import template
+
+    eval_instructions = build_cli_instructions(KERNELBENCH_BASH_ALLOWLIST)
+    cli_instructions = build_cli_instructions(template.bash_allowlist)
+
+    assert eval_instructions == cli_instructions, (
+        "CLI instructions differ — this means the bash allowlists diverged."
+    )
+    assert len(eval_instructions) > 0, "CLI instructions should not be empty"
+
+
+def _dump_full_prompts() -> None:
+    """Standalone: dump both composed prompts for manual comparison."""
+    from optimize_kernelbench_eval.base_config import SYSTEM_PROMPT as EVAL_PROMPT
+
+    from wafer.agent_defaults import KERNELBENCH_BASH_ALLOWLIST
+    from wafer.cli_instructions import build_cli_instructions
+    from wafer.templates.optimize_kernelbench import template
+
+    cli_instructions = build_cli_instructions(KERNELBENCH_BASH_ALLOWLIST)
+
+    # Eval
+    eval_sys = EVAL_PROMPT.format(
+        backend="HIP",
+        backend_lower="hip",
+        target_flag="--pool mi300x-pool",
+        reference_path="<reference_file>",
+    )
+    eval_sys += "\n\n" + cli_instructions
+
+    # CLI
+    params = dict(template.defaults)
+    cli_sys = string.Template(template.system_prompt).safe_substitute(**params)
+    cli_sys += "\n\n" + build_cli_instructions(template.bash_allowlist)
+
+    print("=" * 60)
+    print("EVAL SYSTEM PROMPT")
+    print("=" * 60)
+    print(eval_sys)
+    print()
+    print("=" * 60)
+    print("CLI SYSTEM PROMPT")
+    print("=" * 60)
+    print(cli_sys)
+    print()
+
+    # Diff
+    eval_norm = _normalize_prompt(eval_sys)
+    cli_norm = _normalize_prompt(cli_sys)
+    diff = list(
+        difflib.unified_diff(
+            eval_norm.splitlines(),
+            cli_norm.splitlines(),
+            fromfile="eval",
+            tofile="cli",
+            lineterm="",
+            n=1,
+        )
+    )
+    if diff:
+        print("=" * 60)
+        print("DIFFERENCES (after normalizing pool names)")
+        print("=" * 60)
+        for line in diff:
+            print(line)
+    else:
+        print("IDENTICAL (after normalizing pool names)")
+
+
+if __name__ == "__main__":
+    if "--dump" in sys.argv:
+        _dump_full_prompts()
+    else:
+        test_bash_allowlist_parity()
+        print("PASS: bash_allowlist_parity")
+        test_enabled_tools_parity()
+        print("PASS: enabled_tools_parity")
+        test_system_prompt_parity()
+        print("PASS: system_prompt_parity")
+        test_cli_instructions_identical()
+        print("PASS: cli_instructions_identical")
+        print("\nAll parity checks passed.")

wafer/trace_compare.py

@@ -0,0 +1,274 @@
+"""CLI wrapper for trace comparison commands.
+
+This module provides the CLI interface for the `wafer compare` commands.
+All core logic is in wafer_core.lib.trace_compare.
+"""
+
+import sys
+from pathlib import Path
+from typing import Any
+
+import typer
+
+import json
+import sys
+
+from wafer_core.lib.trace_compare import (
+    analyze_trace_pair,
+    format_csv,
+    format_json,
+    format_text,
+    ArchitectureType,
+    detect_architecture,
+)
+from wafer_core.lib.trace_compare.loader import StreamingMetadata
+
+
+def compare_traces(
+    trace1: Path,
+    trace2: Path,
+    output: Path | None = None,
+    output_format: str = "text",
+    phase: str = "all",
+    show_layers: bool = False,
+    show_all: bool = False,
+    show_stack_traces: bool = False,
+    recommendations: bool = False,
+) -> None:
+    """Compare two GPU traces and generate performance report.
+
+    Args:
+        trace1: Path to first trace file (AMD or NVIDIA)
+        trace2: Path to second trace file (AMD or NVIDIA)
+        output: Optional output file path (default: stdout)
+        output_format: Output format ('text', 'text-layers', 'csv', 'csv-layers', or 'json')
+        phase: Filter by phase ('all', 'prefill', or 'decode')
+        show_layers: Show layer-wise performance breakdown (text format only)
+        show_all: Show all items without truncation (applies to layers, operations, kernels)
+        show_stack_traces: Show Python stack traces for operations
+    """
+    # Validate files exist
+    if not trace1.exists():
+        typer.secho(f"❌ File not found: {trace1}", fg=typer.colors.RED, err=True)
+        raise typer.Exit(1)
+
+    if not trace2.exists():
+        typer.secho(f"❌ File not found: {trace2}", fg=typer.colors.RED, err=True)
+        raise typer.Exit(1)
+
+    # Progress callback for JSON format (emits NDJSON to stdout)
+    def progress_callback(stage: str, fraction: float) -> None:
+        if output_format == 'json':
+            progress_msg = json.dumps({"type": "progress", "stage": stage, "fraction": fraction})
+            print(progress_msg, file=sys.stdout, flush=True)
+        elif output_format != 'json':
+            percent = int(fraction * 100)
+            typer.echo(f"📊 {stage}: {percent}%", err=True)
+
+    # Metadata callback for JSON format (emits NDJSON with early GPU info)
+    def metadata_callback(meta1: StreamingMetadata, meta2: StreamingMetadata) -> None:
+        if output_format == 'json':
+            metadata_msg = json.dumps({
+                "type": "metadata",
+                "trace1": {
+                    "platform": meta1.platform,
+                    "gpu": meta1.gpu_name,
+                    "file_size_mb": round(meta1.file_size_mb, 1),
+                },
+                "trace2": {
+                    "platform": meta2.platform,
+                    "gpu": meta2.gpu_name,
+                    "file_size_mb": round(meta2.file_size_mb, 1),
+                },
+            })
+            print(metadata_msg, file=sys.stdout, flush=True)
+        else:
+            typer.echo(f"📊 Trace 1: {meta1.platform} - {meta1.gpu_name} ({meta1.file_size_mb:.1f}MB)", err=True)
+            typer.echo(f"📊 Trace 2: {meta2.platform} - {meta2.gpu_name} ({meta2.file_size_mb:.1f}MB)", err=True)
+
+    # Analyze traces using unified API
+    if output_format != 'json':
+        typer.echo("📊 Loading traces...")
+
+    try:
+        result_obj = analyze_trace_pair(
+            trace1,
+            trace2,
+            phase=phase,
+            include_stacks=True,
+            on_progress=progress_callback,
+            on_metadata=metadata_callback,
+        )
+        
+        results = {
+            "metadata": result_obj.metadata,
+            "operations": result_obj.operations,
+            "layers": result_obj.layers,
+            "warnings": [{"code": w.code, "severity": w.severity, "message": w.message, "suggestion": w.suggestion} for w in result_obj.warnings],
+            "architecture": result_obj.architecture.value,
+            "layer_alignments": result_obj.layer_alignments,
+            "fusion_analysis": result_obj.fusion_analysis,
+            "same_kernel_analysis": result_obj.same_kernel_analysis,
+        }
+    except ValueError as e:
+        typer.secho(f"❌ {e}", fg=typer.colors.RED, err=True)
+        raise typer.Exit(1)
+    except Exception as e:
+        typer.secho(f"❌ Error analyzing traces: {e}", fg=typer.colors.RED, err=True)
+        raise typer.Exit(1)
+
+    if output_format != 'json':
+        meta = results["metadata"]
+        if meta['trace1_platform'] == 'AMD':
+            amd_gpu, nvidia_gpu = meta['trace1_gpu'], meta['trace2_gpu']
+        else:
+            amd_gpu, nvidia_gpu = meta['trace2_gpu'], meta['trace1_gpu']
+        typer.echo(f"✅ Loaded: AMD ({amd_gpu}) vs NVIDIA ({nvidia_gpu})")
+        
+        # Display warnings
+        warnings = results.get("warnings", [])
+        if warnings:
+            typer.echo()
+            for warning in warnings:
+                icon = "❌" if warning["severity"] == "error" else "⚠️" if warning["severity"] == "warning" else "ℹ️"
+                typer.secho(f"{icon}  {warning['message']}", fg=typer.colors.YELLOW if warning["severity"] == "warning" else typer.colors.BLUE)
+                if warning.get("suggestion"):
+                    typer.secho(f"   Suggestion: {warning['suggestion']}", fg=typer.colors.BLUE)
+    typer.echo()
+
+
+    # Generate output based on format
+    if output_format == "text":
+        output_str = format_text(results, show_layers=show_layers, show_all=show_all, show_stack_traces=show_stack_traces)
+    elif output_format == "text-layers":
+        output_str = format_text(results, show_layers=True, show_all=show_all, show_stack_traces=show_stack_traces)
+    elif output_format == "csv":
+        output_str = format_csv(results, report_type="operations")
+    elif output_format == "csv-layers":
+        output_str = format_csv(results, report_type="layers")
+    elif output_format == "json":
+        output_str = format_json(results)
+    else:
+        typer.secho(f"❌ Unknown format: {output_format}", fg=typer.colors.RED, err=True)
+        raise typer.Exit(1)
+
+    # Write output
+    if output:
+        output.write_text(output_str)
+        typer.secho(f"✅ Report saved to {output}", fg=typer.colors.GREEN)
+    else:
+        typer.echo(output_str)
+
+
+def compare_align(
+    trace1: Path,
+    trace2: Path,
+    output: Path | None = None,
+    output_format: str = "json",
+    phase: str = "all",
+    layer: int | None = None,
+) -> None:
+    """Align kernels at layer level for exact kernel-to-kernel comparison.
+
+    Args:
+        trace1: Path to first trace file (AMD or NVIDIA)
+        trace2: Path to second trace file (AMD or NVIDIA)
+        output: Optional output file path (default: stdout)
+        output_format: Output format ('json' only for now)
+        phase: Filter by phase ('all', 'prefill', or 'decode')
+        layer: Focus on specific layer number (optional)
+    """
+    # Validate files exist
+    if not trace1.exists():
+        typer.secho(f"❌ File not found: {trace1}", fg=typer.colors.RED, err=True)
+        raise typer.Exit(1)
+
+    if not trace2.exists():
+        typer.secho(f"❌ File not found: {trace2}", fg=typer.colors.RED, err=True)
+        raise typer.Exit(1)
+
+    # Progress callback for JSON format (emits NDJSON to stdout)
+    def progress_callback(stage: str, fraction: float) -> None:
+        if output_format == 'json':
+            progress_msg = json.dumps({"type": "progress", "stage": stage, "fraction": fraction})
+            print(progress_msg, file=sys.stdout, flush=True)
+        else:
+            percent = int(fraction * 100)
+            typer.echo(f"📊 {stage}: {percent}%", err=True)
+
+    # Metadata callback for JSON format
+    def metadata_callback(meta1: StreamingMetadata, meta2: StreamingMetadata) -> None:
+        if output_format == 'json':
+            metadata_msg = json.dumps({
+                "type": "metadata",
+                "trace1": {
+                    "platform": meta1.platform,
+                    "gpu": meta1.gpu_name,
+                    "file_size_mb": round(meta1.file_size_mb, 1),
+                },
+                "trace2": {
+                    "platform": meta2.platform,
+                    "gpu": meta2.gpu_name,
+                    "file_size_mb": round(meta2.file_size_mb, 1),
+                },
+            })
+            print(metadata_msg, file=sys.stdout, flush=True)
+        else:
+            typer.echo(f"📊 Trace 1: {meta1.platform} - {meta1.gpu_name} ({meta1.file_size_mb:.1f}MB)", err=True)
+            typer.echo(f"📊 Trace 2: {meta2.platform} - {meta2.gpu_name} ({meta2.file_size_mb:.1f}MB)", err=True)
+
+    # Analyze traces using unified API
+    if output_format != 'json':
+        typer.echo("📊 Loading traces...")
+
+    try:
+        result_obj = analyze_trace_pair(
+            trace1,
+            trace2,
+            phase=phase,
+            include_stacks=True,
+            on_progress=progress_callback,
+            on_metadata=metadata_callback,
+        )
+        
+        results = {
+            "metadata": result_obj.metadata,
+            "layer_alignments": result_obj.layer_alignments or [],
+            "fusion_analysis": result_obj.fusion_analysis or {},
+            "same_kernel_analysis": result_obj.same_kernel_analysis or {},
+            "operations": result_obj.operations,
+            "layers": result_obj.layers,
+            "warnings": [{"code": w.code, "severity": w.severity, "message": w.message, "suggestion": w.suggestion} for w in result_obj.warnings],
+            "architecture": result_obj.architecture.value,
+        }
+        
+        if layer is not None:
+            results["layer_alignments"] = [
+                la for la in results["layer_alignments"] if la.get("layer") == layer
+            ]
+    except ValueError as e:
+        typer.secho(f"❌ {e}", fg=typer.colors.RED, err=True)
+        raise typer.Exit(1)
+    except Exception as e:
+        typer.secho(f"❌ Error analyzing traces: {e}", fg=typer.colors.RED, err=True)
+        import traceback
+        traceback.print_exc()
+        raise typer.Exit(1)
+
+    if output_format != 'json':
+        meta = results["metadata"]
+        typer.echo(f"✅ Loaded: {meta.get('amd_gpu', 'Unknown')} vs {meta.get('nvidia_gpu', 'Unknown')}")
+        typer.echo(f"✅ Found {len(results['layer_alignments'])} layers")
+        typer.echo()
+
+    if output_format == "json":
+        output_str = format_json(results)
+    else:
+        typer.secho(f"❌ Format {output_format} not yet supported for align command. Use 'json'.", fg=typer.colors.RED, err=True)
+        raise typer.Exit(1)
+
+    if output:
+        output.write_text(output_str)
+        typer.secho(f"✅ Report saved to {output}", fg=typer.colors.GREEN)
+    else:
+        typer.echo(output_str)