wafer-cli 0.2.8__tar.gz → 0.2.9__tar.gz

{wafer_cli-0.2.8 → wafer_cli-0.2.9}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: wafer-cli
-Version: 0.2.8
+Version: 0.2.9
 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.8 → wafer_cli-0.2.9}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "wafer-cli"
-version = "0.2.8"
+version = "0.2.9"
 description = "CLI tool for running commands on remote GPUs and GPU kernel optimization agent"
 requires-python = ">=3.11"
 dependencies = [

wafer_cli-0.2.9/tests/test_kernel_scope_cli.py

@@ -0,0 +1,467 @@
+"""Unit tests for Kernel Scope CLI commands.
+
+Tests the wafer amd kernel-scope command using CliRunner.
+
+Run with: PYTHONPATH=apps/wafer-cli uv run pytest apps/wafer-cli/tests/test_kernel_scope_cli.py -v
+"""
+
+import json
+import re
+from pathlib import Path
+from unittest.mock import patch, MagicMock
+
+import pytest
+from typer.testing import CliRunner
+
+from wafer.cli import app
+from wafer.kernel_scope import (
+    analyze_command,
+    metrics_command,
+    targets_command,
+)
+
+
+runner = CliRunner()
+
+
+def strip_ansi(text: str) -> str:
+    """Remove ANSI escape codes from text."""
+    ansi_escape = re.compile(r"\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])")
+    return ansi_escape.sub("", text)
+
+
+# Sample ISA for testing
+SAMPLE_ISA = '''
+.amdgcn_target "amdgcn-amd-amdhsa--gfx90a"
+
+.text
+.globl test_kernel
+
+test_kernel:
+    s_load_dwordx4 s[0:3], s[4:5], 0x0
+    s_waitcnt lgkmcnt(0)
+    global_load_dwordx4 v[0:3], v[4:5], off
+    ds_read_b128 v[8:11], v12
+    s_waitcnt vmcnt(0) lgkmcnt(0)
+    v_mfma_f32_32x32x8f16 a[0:15], v[0:1], v[2:3], a[0:15]
+    v_mfma_f32_32x32x8f16 a[16:31], v[4:5], v[6:7], a[16:31]
+    v_add_f32 v0, v1, v2
+    v_fma_f32 v3, v4, v5, v6
+    global_store_dwordx4 v[20:21], v[0:3], off
+    s_barrier
+    s_endpgm
+
+.amdhsa_kernel test_kernel
+    .amdhsa_next_free_vgpr 64
+    .amdhsa_next_free_sgpr 32
+    .amdhsa_group_segment_fixed_size 16384
+    .amdhsa_private_segment_fixed_size 0
+.end_amdhsa_kernel
+'''
+
+SAMPLE_ISA_WITH_SPILLS = '''
+.amdgcn_target "amdgcn-amd-amdhsa--gfx942"
+
+.text
+.globl spilling_kernel
+
+spilling_kernel:
+    s_load_dwordx4 s[0:3], s[4:5], 0x0
+    s_waitcnt lgkmcnt(0)
+    scratch_store_dwordx4 off, v[0:3], s0
+    scratch_store_dwordx4 off, v[4:7], s0
+    scratch_load_dwordx4 v[8:11], off, s0
+    v_add_f32 v0, v1, v2
+    s_waitcnt 0
+    s_endpgm
+
+.amdhsa_kernel spilling_kernel
+    .amdhsa_next_free_vgpr 256
+    .amdhsa_next_free_sgpr 100
+    .amdhsa_group_segment_fixed_size 0
+    .amdhsa_private_segment_fixed_size 1024
+.end_amdhsa_kernel
+'''
+
+
+# ============================================================================
+# Direct Function Tests
+# ============================================================================
+
+class TestAnalyzeCommandFunction:
+    """Tests for analyze_command function directly."""
+
+    def test_analyze_single_file(self, tmp_path: Path) -> None:
+        """Should analyze a single ISA file."""
+        isa_file = tmp_path / "kernel.s"
+        isa_file.write_text(SAMPLE_ISA)
+
+        output = analyze_command(str(isa_file))
+
+        assert "test_kernel" in output
+        assert "gfx90a" in output
+        assert "VGPRs:" in output or "vgpr" in output.lower()
+
+    def test_analyze_file_not_found(self, tmp_path: Path) -> None:
+        """Should raise error for missing file."""
+        with pytest.raises(FileNotFoundError):
+            analyze_command(str(tmp_path / "missing.s"))
+
+    def test_analyze_json_output(self, tmp_path: Path) -> None:
+        """Should output valid JSON when json_output=True."""
+        isa_file = tmp_path / "kernel.s"
+        isa_file.write_text(SAMPLE_ISA)
+
+        output = analyze_command(str(isa_file), json_output=True)
+        data = json.loads(output)
+
+        assert data["success"] is True
+        assert data["isa_analysis"]["kernel_name"] == "test_kernel"
+        assert data["isa_analysis"]["architecture"] == "gfx90a"
+        assert data["isa_analysis"]["vgpr_count"] == 64
+        assert data["isa_analysis"]["mfma_count"] == 2
+
+    def test_analyze_csv_output(self, tmp_path: Path) -> None:
+        """Should output CSV when csv_output=True."""
+        isa_file = tmp_path / "kernel.s"
+        isa_file.write_text(SAMPLE_ISA)
+
+        output = analyze_command(str(isa_file), csv_output=True)
+
+        # Should have header and data row
+        lines = output.strip().split("\n")
+        assert len(lines) == 2
+        assert "kernel_name" in lines[0]
+        assert "vgpr_count" in lines[0]
+        assert "test_kernel" in lines[1]
+
+    def test_analyze_directory(self, tmp_path: Path) -> None:
+        """Should analyze all files in directory."""
+        (tmp_path / "kernel1.s").write_text(SAMPLE_ISA)
+        (tmp_path / "kernel2.s").write_text(SAMPLE_ISA_WITH_SPILLS)
+
+        output = analyze_command(str(tmp_path))
+
+        assert "Analyzed 2 files" in output
+
+    def test_analyze_with_filter(self, tmp_path: Path) -> None:
+        """Should filter results based on expression."""
+        (tmp_path / "kernel1.s").write_text(SAMPLE_ISA)
+        (tmp_path / "kernel2.s").write_text(SAMPLE_ISA_WITH_SPILLS)
+
+        output = analyze_command(str(tmp_path), filter_expr="spills > 0")
+
+        # Should only show spilling kernel
+        assert "1 files" in output or "kernel2" in output or "spilling" in output.lower()
+
+    def test_analyze_output_to_file(self, tmp_path: Path) -> None:
+        """Should write output to file."""
+        isa_file = tmp_path / "kernel.s"
+        isa_file.write_text(SAMPLE_ISA)
+        output_file = tmp_path / "output.json"
+
+        analyze_command(
+            str(isa_file),
+            json_output=True,
+            output_file=str(output_file)
+        )
+
+        assert output_file.exists()
+        data = json.loads(output_file.read_text())
+        assert data["success"] is True
+
+
+class TestMetricsCommandFunction:
+    """Tests for metrics_command function."""
+
+    def test_lists_metrics(self) -> None:
+        """Should list available metrics."""
+        output = metrics_command()
+
+        assert "vgpr_count" in output
+        assert "sgpr_count" in output
+        assert "spill_count" in output
+        assert "mfma_count" in output
+        assert "mfma_density_pct" in output
+        assert "theoretical_occupancy" in output
+
+    def test_includes_instruction_categories(self) -> None:
+        """Should include instruction category descriptions."""
+        output = metrics_command()
+
+        assert "VALU" in output
+        assert "SALU" in output
+        assert "VMEM" in output
+        assert "MFMA" in output
+        assert "LDS" in output
+        assert "SPILL" in output
+
+
+class TestTargetsCommandFunction:
+    """Tests for targets_command function."""
+
+    def test_lists_targets(self) -> None:
+        """Should list supported GPU targets."""
+        output = targets_command()
+
+        assert "gfx90a" in output
+        assert "gfx942" in output
+        assert "gfx908" in output
+
+    def test_includes_specs(self) -> None:
+        """Should include hardware specs for targets."""
+        output = targets_command()
+
+        assert "MI200" in output or "MI300" in output
+        assert "VGPRs" in output or "VGPR" in output
+
+
+# ============================================================================
+# CLI Integration Tests
+# ============================================================================
+
+class TestKernelScopeCliCommands:
+    """Tests for wafer amd kernel-scope CLI commands."""
+
+    def test_analyze_command_help(self) -> None:
+        """Should display help for analyze command."""
+        result = runner.invoke(app, ["amd", "kernel-scope", "analyze", "--help"])
+
+        assert result.exit_code == 0
+        output = strip_ansi(result.stdout)
+        assert "Analyze" in output or "analyze" in output
+
+    def test_analyze_file_via_cli(self, tmp_path: Path) -> None:
+        """Should analyze file via CLI."""
+        isa_file = tmp_path / "kernel.s"
+        isa_file.write_text(SAMPLE_ISA)
+
+        result = runner.invoke(app, [
+            "amd", "kernel-scope", "analyze", str(isa_file)
+        ])
+
+        assert result.exit_code == 0, f"Failed: {result.output}"
+        assert "test_kernel" in result.stdout
+
+    def test_analyze_json_via_cli(self, tmp_path: Path) -> None:
+        """Should output JSON via CLI."""
+        isa_file = tmp_path / "kernel.s"
+        isa_file.write_text(SAMPLE_ISA)
+
+        result = runner.invoke(app, [
+            "amd", "kernel-scope", "analyze", str(isa_file), "--json"
+        ])
+
+        assert result.exit_code == 0, f"Failed: {result.output}"
+        data = json.loads(result.stdout)
+        assert data["success"] is True
+
+    def test_analyze_csv_via_cli(self, tmp_path: Path) -> None:
+        """Should output CSV via CLI."""
+        isa_file = tmp_path / "kernel.s"
+        isa_file.write_text(SAMPLE_ISA)
+
+        result = runner.invoke(app, [
+            "amd", "kernel-scope", "analyze", str(isa_file), "--csv"
+        ])
+
+        assert result.exit_code == 0, f"Failed: {result.output}"
+        assert "kernel_name" in result.stdout
+
+    def test_analyze_missing_file_via_cli(self, tmp_path: Path) -> None:
+        """Should fail for missing file."""
+        result = runner.invoke(app, [
+            "amd", "kernel-scope", "analyze", str(tmp_path / "missing.s")
+        ])
+
+        assert result.exit_code != 0
+
+    def test_metrics_via_cli(self) -> None:
+        """Should list metrics via CLI."""
+        result = runner.invoke(app, ["amd", "kernel-scope", "metrics"])
+
+        assert result.exit_code == 0
+        assert "vgpr_count" in result.stdout
+
+    def test_targets_via_cli(self) -> None:
+        """Should list targets via CLI."""
+        result = runner.invoke(app, ["amd", "kernel-scope", "targets"])
+
+        assert result.exit_code == 0
+        assert "gfx90a" in result.stdout or "gfx942" in result.stdout
+
+
+class TestKernelScopeCliHelp:
+    """Tests for kernel-scope command help text."""
+
+    def test_kernel_scope_help(self) -> None:
+        """Should display help for kernel-scope command group."""
+        result = runner.invoke(app, ["amd", "kernel-scope", "--help"])
+
+        assert result.exit_code == 0
+        output = strip_ansi(result.stdout)
+        assert "analyze" in output.lower() or "Analyze" in output
+        assert "metrics" in output.lower()
+        assert "targets" in output.lower()
+
+    def test_amd_help_includes_kernel_scope(self) -> None:
+        """AMD help should mention kernel-scope."""
+        result = runner.invoke(app, ["amd", "--help"])
+
+        assert result.exit_code == 0
+        assert "kernel-scope" in result.stdout
+
+
+# ============================================================================
+# Output Format Tests
+# ============================================================================
+
+class TestOutputFormats:
+    """Tests for output formatting."""
+
+    def test_text_output_spills_warning(self, tmp_path: Path) -> None:
+        """Text output should show spills warning."""
+        isa_file = tmp_path / "kernel.s"
+        isa_file.write_text(SAMPLE_ISA_WITH_SPILLS)
+
+        output = analyze_command(str(isa_file))
+
+        assert "SPILLS" in output or "spills" in output.lower()
+
+    def test_text_output_registers_section(self, tmp_path: Path) -> None:
+        """Text output should have registers section."""
+        isa_file = tmp_path / "kernel.s"
+        isa_file.write_text(SAMPLE_ISA)
+
+        output = analyze_command(str(isa_file))
+
+        assert "Registers" in output or "VGPRs" in output
+
+    def test_text_output_memory_section(self, tmp_path: Path) -> None:
+        """Text output should have memory section."""
+        isa_file = tmp_path / "kernel.s"
+        isa_file.write_text(SAMPLE_ISA)
+
+        output = analyze_command(str(isa_file))
+
+        assert "Memory" in output or "LDS" in output
+
+    def test_text_output_instructions_section(self, tmp_path: Path) -> None:
+        """Text output should have instructions section."""
+        isa_file = tmp_path / "kernel.s"
+        isa_file.write_text(SAMPLE_ISA)
+
+        output = analyze_command(str(isa_file))
+
+        assert "Instructions" in output or "MFMA" in output
+
+    def test_text_output_occupancy_section(self, tmp_path: Path) -> None:
+        """Text output should have occupancy section."""
+        isa_file = tmp_path / "kernel.s"
+        isa_file.write_text(SAMPLE_ISA)
+
+        output = analyze_command(str(isa_file))
+
+        assert "Occupancy" in output or "waves" in output.lower()
+
+    def test_json_output_has_all_fields(self, tmp_path: Path) -> None:
+        """JSON output should include all analysis fields."""
+        isa_file = tmp_path / "kernel.s"
+        isa_file.write_text(SAMPLE_ISA)
+
+        output = analyze_command(str(isa_file), json_output=True)
+        data = json.loads(output)
+        analysis = data["isa_analysis"]
+
+        # Check required fields
+        assert "kernel_name" in analysis
+        assert "architecture" in analysis
+        assert "vgpr_count" in analysis
+        assert "sgpr_count" in analysis
+        assert "spill_count" in analysis
+        assert "mfma_count" in analysis
+        assert "mfma_density_pct" in analysis
+        assert "instruction_mix" in analysis
+        assert "theoretical_occupancy" in analysis
+        assert "warnings" in analysis
+
+
+# ============================================================================
+# Filter Tests
+# ============================================================================
+
+class TestFiltering:
+    """Tests for result filtering."""
+
+    def test_filter_spills_greater_than_zero(self, tmp_path: Path) -> None:
+        """Filter 'spills > 0' should only show files with spills."""
+        (tmp_path / "no_spills.s").write_text(SAMPLE_ISA)
+        (tmp_path / "has_spills.s").write_text(SAMPLE_ISA_WITH_SPILLS)
+
+        output = analyze_command(str(tmp_path), filter_expr="spills > 0")
+
+        # Should filter to only spilling kernel
+        assert "1 files" in output or "Analyzed 1" in output
+
+    def test_filter_vgpr_count(self, tmp_path: Path) -> None:
+        """Filter on VGPR count should work."""
+        (tmp_path / "small_vgpr.s").write_text(SAMPLE_ISA)  # 64 VGPRs
+        (tmp_path / "large_vgpr.s").write_text(SAMPLE_ISA_WITH_SPILLS)  # 256 VGPRs
+
+        output = analyze_command(str(tmp_path), filter_expr="vgpr_count > 128")
+
+        # Should only show high-VGPR kernel
+        assert "1 files" in output or "Analyzed 1" in output
+
+    def test_filter_mfma_count(self, tmp_path: Path) -> None:
+        """Filter on MFMA count should work."""
+        (tmp_path / "has_mfma.s").write_text(SAMPLE_ISA)  # 2 MFMAs
+        (tmp_path / "no_mfma.s").write_text(SAMPLE_ISA_WITH_SPILLS)  # 0 MFMAs
+
+        output = analyze_command(str(tmp_path), filter_expr="mfma > 0")
+
+        # Should show kernel with MFMA
+        assert "1 files" in output or "has_mfma" in output
+
+    def test_filter_invalid_expression(self, tmp_path: Path, capsys) -> None:
+        """Invalid filter expression should warn."""
+        (tmp_path / "kernel.s").write_text(SAMPLE_ISA)
+
+        output = analyze_command(str(tmp_path), filter_expr="invalid filter")
+
+        # Should still analyze, just warn about invalid filter
+        # The function prints warning to stderr
+        captured = capsys.readouterr()
+        assert "Invalid" in captured.err or "Warning" in captured.err or "2 files" not in output
+
+
+# ============================================================================
+# Edge Cases
+# ============================================================================
+
+class TestEdgeCases:
+    """Tests for edge cases and error handling."""
+
+    def test_analyze_unsupported_file_type(self, tmp_path: Path) -> None:
+        """Should handle unsupported file types gracefully."""
+        txt_file = tmp_path / "file.xyz"
+        txt_file.write_text("not ISA content")
+
+        with pytest.raises(RuntimeError, match="Unsupported"):
+            analyze_command(str(txt_file))
+
+    def test_analyze_empty_directory(self, tmp_path: Path) -> None:
+        """Should handle empty directories."""
+        output = analyze_command(str(tmp_path))
+
+        assert "0 files" in output or "No supported files" in output.lower()
+
+    def test_analyze_directory_with_subdirs(self, tmp_path: Path) -> None:
+        """Should scan subdirectories when recursive."""
+        subdir = tmp_path / "subdir"
+        subdir.mkdir()
+        (subdir / "kernel.s").write_text(SAMPLE_ISA)
+
+        output = analyze_command(str(tmp_path), recursive=True)
+
+        assert "1 files" in output or "test_kernel" in output

{wafer_cli-0.2.8 → wafer_cli-0.2.9}/wafer/auth.py

@@ -345,3 +345,88 @@ def browser_login(timeout: int = 120, port: int | None = None) -> tuple[str, str
 
     server.server_close()
     raise TimeoutError(f"No response within {timeout} seconds")
+
+
+def device_code_login(timeout: int = 600) -> tuple[str, str | None]:
+    """Authenticate using state-based flow (no browser/port forwarding needed).
+
+    This is the SSH-friendly auth flow similar to GitHub CLI:
+    1. Request a state token from the API
+    2. Display the auth URL with state parameter
+    3. User visits URL on any device and signs in normally
+    4. Poll API until user completes authentication
+
+    Args:
+        timeout: Seconds to wait for authentication (default 600 = 10 minutes)
+
+    Returns:
+        Tuple of (access_token, refresh_token). refresh_token may be None.
+
+    Raises:
+        TimeoutError: If user doesn't authenticate within timeout
+        RuntimeError: If auth flow failed
+    """
+    api_url = get_api_url()
+
+    # Request state and auth URL
+    with httpx.Client(timeout=10.0) as client:
+        response = client.post(f"{api_url}/v1/auth/cli-auth/start", json={})
+        response.raise_for_status()
+        data = response.json()
+
+    state = data["state"]
+    auth_url = data["auth_url"]
+    expires_in = data["expires_in"]
+
+    # Display instructions to user
+    print("\n" + "=" * 60)
+    print("  WAFER CLI - Authentication")
+    print("=" * 60)
+    print(f"\n  Visit: {auth_url}")
+    print("\n  Sign in with GitHub to complete authentication")
+    print("\n" + "=" * 60 + "\n")
+
+    # Poll for authentication
+    start = time.time()
+    poll_interval = 5  # Poll every 5 seconds
+    last_poll = 0.0
+
+    print("Waiting for authentication", end="", flush=True)
+
+    while time.time() - start < min(timeout, expires_in):
+        # Show progress dots
+        if time.time() - last_poll >= poll_interval:
+            print(".", end="", flush=True)
+
+            # Poll the API
+            with httpx.Client(timeout=10.0) as client:
+                try:
+                    response = client.post(f"{api_url}/v1/auth/cli-auth/token", json={"state": state})
+
+                    if response.status_code == 200:
+                        # Success!
+                        data = response.json()
+                        print(f" {CHECK}\n")
+                        return data["access_token"], data.get("refresh_token")
+
+                    if response.status_code == 428:
+                        # Still waiting
+                        last_poll = time.time()
+                        time.sleep(1)
+                        continue
+
+                    # Some other error
+                    print(f" {CROSS}\n")
+                    raise RuntimeError(f"CLI auth flow failed: {response.status_code} {response.text}")
+
+                except httpx.RequestError as e:
+                    # Network error, retry
+                    print("!", end="", flush=True)
+                    last_poll = time.time()
+                    time.sleep(1)
+                    continue
+
+        time.sleep(0.5)  # Small sleep to avoid busy loop
+
+    print(f" {CROSS}\n")
+    raise TimeoutError(f"Authentication not completed within {expires_in} seconds")