wafer-cli 0.2.52__tar.gz → 0.2.53__tar.gz

{wafer_cli-0.2.52 → wafer_cli-0.2.53}/PKG-INFO

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

@@ -1,6 +1,6 @@
 [project]
 name = "wafer-cli"
-version = "0.2.52"
+version = "0.2.53"
 description = "CLI for running GPU workloads, managing remote workspaces, and evaluating/optimizing kernels"
 readme = "README.md"
 requires-python = ">=3.11"
@@ -78,6 +78,7 @@ ignore = [
 
 [tool.ruff.lint.per-file-ignores]
 "tests/**/*.py" = ["ANN001", "ANN201", "ANN202", "ANN204"]  # Don't require type annotations in tests
+"tests/test_ncu_run_local_e2e.py" = ["PLR0915"]  # E2E test has a long sequential flow by design
 "wafer/evaluate.py" = ["PLR0915", "PLR1702", "E402", "PLW2901", "ASYNC221"]  # complex deployment flows - TODO: refactor
 "wafer/output.py" = ["ANN401"]  # Output collector uses **kwargs for flexible event data
 "wafer/autotuner.py" = ["PLR0915", "PLR1702", "B007", "B904"]  # complex sweep logic - TODO: refactor

wafer_cli-0.2.53/tests/test_ncu_run.py

@@ -0,0 +1,337 @@
+"""Unit tests for NCU remote profiling CLI module.
+
+Tests directory packaging, .gitignore filtering, tar.gz creation,
+and SSE stream processing — all locally, no B200 needed.
+"""
+
+import os
+import subprocess
+import tarfile
+import tempfile
+from pathlib import Path
+
+import pytest
+
+
+# --- Directory Packaging ---
+
+
+class TestCollectFilesForUpload:
+    """Tests for _collect_files_for_upload — the directory scanning logic."""
+
+    def test_collects_all_files_in_flat_dir(self) -> None:
+        """All files in a simple directory are collected."""
+        from wafer.ncu_run import _collect_files_for_upload
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            d = Path(tmpdir)
+            (d / "run.py").write_text("print('hello')")
+            (d / "kernel.py").write_text("# kernel")
+            (d / "data.txt").write_text("data")
+
+            files = _collect_files_for_upload(d)
+
+            names = {f.name for f in files}
+            assert "run.py" in names
+            assert "kernel.py" in names
+            assert "data.txt" in names
+
+    def test_collects_nested_files(self) -> None:
+        """Files in subdirectories are collected."""
+        from wafer.ncu_run import _collect_files_for_upload
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            d = Path(tmpdir)
+            (d / "run.py").write_text("import utils")
+            (d / "utils").mkdir()
+            (d / "utils" / "helper.py").write_text("# helper")
+
+            files = _collect_files_for_upload(d)
+            rel_paths = {str(f.relative_to(d)) for f in files}
+
+            assert "run.py" in rel_paths
+            assert "utils/helper.py" in rel_paths
+
+    def test_excludes_pycache(self) -> None:
+        """__pycache__ directories are excluded."""
+        from wafer.ncu_run import _collect_files_for_upload
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            d = Path(tmpdir)
+            (d / "run.py").write_text("print('hello')")
+            (d / "__pycache__").mkdir()
+            (d / "__pycache__" / "run.cpython-312.pyc").write_bytes(b"\x00")
+
+            files = _collect_files_for_upload(d)
+            names = {f.name for f in files}
+
+            assert "run.py" in names
+            assert "run.cpython-312.pyc" not in names
+
+    def test_excludes_git_dir(self) -> None:
+        """.git directory is excluded."""
+        from wafer.ncu_run import _collect_files_for_upload
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            d = Path(tmpdir)
+            (d / "run.py").write_text("print('hello')")
+            (d / ".git").mkdir()
+            (d / ".git" / "HEAD").write_text("ref: refs/heads/main")
+
+            files = _collect_files_for_upload(d)
+            names = {f.name for f in files}
+
+            assert "run.py" in names
+            assert "HEAD" not in names
+
+    def test_excludes_ncu_rep_files(self) -> None:
+        """Existing .ncu-rep files are excluded from upload."""
+        from wafer.ncu_run import _collect_files_for_upload
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            d = Path(tmpdir)
+            (d / "run.py").write_text("print('hello')")
+            (d / "old_profile.ncu-rep").write_bytes(b"\x00" * 100)
+
+            files = _collect_files_for_upload(d)
+            names = {f.name for f in files}
+
+            assert "run.py" in names
+            assert "old_profile.ncu-rep" not in names
+
+    def test_excludes_venv(self) -> None:
+        """Virtual environment directories are excluded."""
+        from wafer.ncu_run import _collect_files_for_upload
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            d = Path(tmpdir)
+            (d / "run.py").write_text("print('hello')")
+            (d / ".venv").mkdir()
+            (d / ".venv" / "bin").mkdir()
+            (d / ".venv" / "bin" / "python").write_text("#!/usr/bin/env python")
+
+            files = _collect_files_for_upload(d)
+            names = {f.name for f in files}
+
+            assert "run.py" in names
+            assert "python" not in names
+
+    def test_custom_excludes(self) -> None:
+        """Extra exclude patterns are applied."""
+        from wafer.ncu_run import _collect_files_for_upload
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            d = Path(tmpdir)
+            (d / "run.py").write_text("print('hello')")
+            (d / "big_data.bin").write_bytes(b"\x00" * 100)
+
+            files = _collect_files_for_upload(d, extra_excludes={"*.bin"})
+            names = {f.name for f in files}
+
+            assert "run.py" in names
+            assert "big_data.bin" not in names
+
+    def test_git_ls_files_used_in_git_repo(self) -> None:
+        """In a git repo, git ls-files is used (respects .gitignore)."""
+        from wafer.ncu_run import _collect_files_for_upload
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            d = Path(tmpdir)
+            # Initialize a git repo
+            subprocess.run(["git", "init"], cwd=d, capture_output=True, check=True)
+            subprocess.run(
+                ["git", "config", "user.email", "test@test.com"],
+                cwd=d, capture_output=True, check=True,
+            )
+            subprocess.run(
+                ["git", "config", "user.name", "Test"],
+                cwd=d, capture_output=True, check=True,
+            )
+
+            # Create files
+            (d / "tracked.py").write_text("# tracked")
+            (d / "ignored.log").write_text("logs")
+
+            # Create .gitignore
+            (d / ".gitignore").write_text("*.log\n")
+
+            # Git add tracked file
+            subprocess.run(["git", "add", "tracked.py", ".gitignore"], cwd=d, capture_output=True, check=True)
+
+            files = _collect_files_for_upload(d)
+            names = {f.name for f in files}
+
+            assert "tracked.py" in names
+            assert ".gitignore" in names
+            assert "ignored.log" not in names
+
+
+# --- Tar.gz Packaging ---
+
+
+class TestPackageDirectory:
+    """Tests for package_directory — creating uploadable tar.gz archives."""
+
+    def test_creates_valid_tarball(self) -> None:
+        """package_directory produces a valid tar.gz."""
+        from wafer.ncu_run import package_directory
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            d = Path(tmpdir)
+            (d / "run.py").write_text("print('hello')")
+            (d / "kernel.py").write_text("# kernel code")
+
+            tarball_bytes, count = package_directory(d)
+
+            assert count == 2
+            assert len(tarball_bytes) > 0
+
+            # Verify it's valid tar.gz
+            import io
+            buf = io.BytesIO(tarball_bytes)
+            with tarfile.open(fileobj=buf, mode="r:gz") as tar:
+                names = tar.getnames()
+                assert "run.py" in names
+                assert "kernel.py" in names
+
+    def test_preserves_relative_paths(self) -> None:
+        """Tarball entries have paths relative to the directory root."""
+        from wafer.ncu_run import package_directory
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            d = Path(tmpdir)
+            (d / "src").mkdir()
+            (d / "src" / "model.py").write_text("# model")
+            (d / "run.py").write_text("from src.model import *")
+
+            tarball_bytes, count = package_directory(d)
+
+            import io
+            buf = io.BytesIO(tarball_bytes)
+            with tarfile.open(fileobj=buf, mode="r:gz") as tar:
+                names = tar.getnames()
+                # Paths should be relative, not absolute
+                assert "run.py" in names
+                assert "src/model.py" in names
+                # No absolute paths
+                for name in names:
+                    assert not name.startswith("/")
+
+    def test_empty_directory_raises(self) -> None:
+        """Empty directory raises AssertionError."""
+        from wafer.ncu_run import package_directory
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            d = Path(tmpdir)
+            with pytest.raises(AssertionError, match="No files found"):
+                package_directory(d)
+
+    def test_realistic_gpu_mode_project(self) -> None:
+        """Packages a realistic GPU Mode kernel project correctly.
+
+        This mimics what a GPU Mode community member would have:
+        - A Python entry script
+        - A Triton kernel file
+        - A utils module
+        - Some files that should be excluded
+        """
+        from wafer.ncu_run import package_directory
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            d = Path(tmpdir)
+
+            # Realistic GPU Mode project structure
+            (d / "run.py").write_text(
+                "import torch\n"
+                "from kernel import my_kernel\n"
+                "x = torch.randn(1024, device='cuda')\n"
+                "my_kernel(x)\n"
+            )
+            (d / "kernel.py").write_text(
+                "import triton\n"
+                "import triton.language as tl\n"
+                "@triton.jit\n"
+                "def my_kernel(x_ptr, n: tl.constexpr):\n"
+                "    pid = tl.program_id(0)\n"
+                "    offsets = pid * 128 + tl.arange(0, 128)\n"
+                "    x = tl.load(x_ptr + offsets)\n"
+                "    tl.store(x_ptr + offsets, x * 2)\n"
+            )
+            (d / "utils.py").write_text("# utility functions\n")
+            (d / "requirements.txt").write_text("torch>=2.0\ntriton>=2.1\n")
+
+            # Files that SHOULD be excluded
+            (d / "__pycache__").mkdir()
+            (d / "__pycache__" / "kernel.cpython-312.pyc").write_bytes(b"\x00")
+            (d / "old_profile.ncu-rep").write_bytes(b"\x00" * 50)
+
+            tarball_bytes, count = package_directory(d)
+
+            # Should include the 4 real files (run.py, kernel.py, utils.py, requirements.txt)
+            # Should NOT include __pycache__ or .ncu-rep
+            import io
+            buf = io.BytesIO(tarball_bytes)
+            with tarfile.open(fileobj=buf, mode="r:gz") as tar:
+                names = set(tar.getnames())
+                assert "run.py" in names
+                assert "kernel.py" in names
+                assert "utils.py" in names
+                assert "requirements.txt" in names
+                # Excluded files
+                assert not any("__pycache__" in n for n in names)
+                assert not any(".ncu-rep" in n for n in names)
+
+
+# --- SSE Stream Processing ---
+
+
+class TestHandleSseLine:
+    """Tests for _handle_sse_line — CLI-side SSE event processing."""
+
+    def test_done_sets_done_flag(self) -> None:
+        """[DONE] event sets the done flag."""
+        from wafer.ncu_run import _handle_sse_line
+
+        state: dict = {"exit_code": 0, "job_id": None, "done": False}
+        _handle_sse_line("[DONE]", state)
+
+        assert state["done"] is True
+
+    def test_error_sets_exit_code(self) -> None:
+        """[ERROR] event sets exit_code to 1."""
+        from wafer.ncu_run import _handle_sse_line
+
+        state: dict = {"exit_code": 0, "job_id": None, "done": False}
+        _handle_sse_line("[ERROR] something broke", state)
+
+        assert state["exit_code"] == 1
+        assert state["done"] is True
+
+    def test_exit_code_parsed(self) -> None:
+        """[EXIT:N] event parses the exit code."""
+        from wafer.ncu_run import _handle_sse_line
+
+        state: dict = {"exit_code": 0, "job_id": None, "done": False}
+        _handle_sse_line("[EXIT:42]", state)
+
+        assert state["exit_code"] == 42
+
+    def test_report_ready_captures_job_id(self) -> None:
+        """[REPORT_READY:id] event captures the job ID."""
+        from wafer.ncu_run import _handle_sse_line
+
+        state: dict = {"exit_code": 0, "job_id": None, "done": False}
+        _handle_sse_line("[REPORT_READY:abc123def]", state)
+
+        assert state["job_id"] == "abc123def"
+
+    def test_status_events_dont_affect_state(self) -> None:
+        """[STATUS:...] events are informational only."""
+        from wafer.ncu_run import _handle_sse_line
+
+        state: dict = {"exit_code": 0, "job_id": None, "done": False}
+        _handle_sse_line("[STATUS:UPLOADING]", state)
+
+        assert state["exit_code"] == 0
+        assert state["job_id"] is None
+        assert state["done"] is False

wafer_cli-0.2.53/tests/test_ncu_run_e2e.py

@@ -0,0 +1,225 @@
+"""End-to-end tests for NCU remote profiling.
+
+These tests require:
+1. Authentication: `wafer login` (or WAFER_API_URL pointed at staging)
+2. B200 GPU access: A running B200 target in the pool
+3. Credits: Sufficient credits for a short NCU run
+
+Run manually:
+    cd apps/wafer-cli
+    uv run pytest tests/test_ncu_run_e2e.py -v -s
+
+Or with staging:
+    WAFER_API_URL=https://wafer-api-staging.onrender.com uv run pytest tests/test_ncu_run_e2e.py -v -s
+"""
+
+import os
+import subprocess
+import tempfile
+from pathlib import Path
+
+import pytest
+
+# Mark all tests in this file as E2E (skip in CI by default)
+pytestmark = [
+    pytest.mark.skipif(
+        os.environ.get("RUN_E2E_TESTS") != "1",
+        reason="Set RUN_E2E_TESTS=1 to run E2E tests (requires B200 + auth)",
+    ),
+]
+
+# Path to the realistic GPU Mode test fixture
+FIXTURE_DIR = Path(__file__).parent / "fixtures" / "gpu_mode_kernel"
+
+
+class TestNcuRunE2E:
+    """End-to-end tests for the full `wafer ncu run` flow."""
+
+    def _run_wafer(self, *args: str, cwd: str | None = None) -> subprocess.CompletedProcess:
+        """Run a wafer CLI command and return the result."""
+        cmd = ["uv", "run", "wafer", *args]
+        env = {**os.environ}
+        if "WAFER_API_URL" not in env:
+            # Default to staging for E2E tests
+            env["WAFER_API_URL"] = "https://wafer-api-staging.onrender.com"
+
+        return subprocess.run(
+            cmd,
+            capture_output=True,
+            text=True,
+            cwd=cwd or str(Path(__file__).parent.parent),
+            env=env,
+            timeout=300,  # 5 min timeout for NCU runs
+        )
+
+    def test_ncu_run_help(self) -> None:
+        """wafer nvidia ncu run --help shows usage."""
+        result = self._run_wafer("nvidia", "ncu", "run", "--help")
+
+        assert result.returncode == 0
+        assert "Run NCU profiling remotely" in result.stdout
+        assert "--dir" in result.stdout
+        assert "--ncu-args" in result.stdout
+
+    def test_ncu_run_no_command_fails(self) -> None:
+        """wafer nvidia ncu run (no command) fails with error."""
+        result = self._run_wafer("nvidia", "ncu", "run")
+
+        # Should fail because no command provided
+        assert result.returncode != 0
+
+    def test_ncu_run_simple_python(self) -> None:
+        """Full E2E: profile a simple Python script on B200.
+
+        This is THE critical test — it validates the entire flow:
+        1. Directory packaging (tar.gz)
+        2. Upload to API
+        3. SFTP to B200
+        4. Docker execution with NCU
+        5. SSE streaming
+        6. Report download
+        """
+        with tempfile.TemporaryDirectory() as tmpdir:
+            d = Path(tmpdir)
+            # Minimal script that NCU can profile
+            (d / "run.py").write_text(
+                "import torch\n"
+                "x = torch.randn(1024, device='cuda')\n"
+                "y = x * 2  # Simple operation for NCU to profile\n"
+                "print(f'Result shape: {y.shape}')\n"
+            )
+
+            result = self._run_wafer(
+                "nvidia", "ncu", "run",
+                "--dir", str(d),
+                "--timeout", "120",
+                "python", "run.py",
+            )
+
+            print("STDOUT:", result.stdout)
+            print("STDERR:", result.stderr)
+
+            # Should upload and start running
+            assert "Packaging" in result.stdout
+            assert "Uploading" in result.stdout
+
+            # Should get NCU output (==PROF== lines)
+            assert "==PROF==" in result.stdout or "RUNNING" in result.stdout
+
+    def test_ncu_run_gpu_mode_fixture(self) -> None:
+        """Full E2E: profile the GPU Mode test fixture (Triton kernel).
+
+        This tests the realistic GPU Mode use case:
+        - Multi-file project (run.py + requirements.txt)
+        - Triton kernel that needs profiling
+        - Requirements installation
+        """
+        if not FIXTURE_DIR.exists():
+            pytest.skip(f"Fixture not found: {FIXTURE_DIR}")
+
+        with tempfile.TemporaryDirectory() as output_dir:
+            output_file = Path(output_dir) / "profile.ncu-rep"
+
+            result = self._run_wafer(
+                "nvidia", "ncu", "run",
+                "--dir", str(FIXTURE_DIR),
+                "--output", str(output_file),
+                "--timeout", "180",
+                "python", "run.py",
+            )
+
+            print("STDOUT:", result.stdout)
+            print("STDERR:", result.stderr)
+
+            # Should complete successfully
+            assert "Packaging" in result.stdout
+            assert "Uploading" in result.stdout
+
+            # If NCU ran successfully, report should be downloaded
+            if result.returncode == 0 and output_file.exists():
+                # Report should be a valid NCU report (starts with NVR magic)
+                header = output_file.read_bytes()[:3]
+                assert header == b"NVR", f"Invalid report header: {header}"
+                print(f"Report downloaded: {output_file.stat().st_size / 1024:.1f} KB")
+
+    def test_ncu_run_with_ncu_args(self) -> None:
+        """E2E: custom NCU flags are passed through."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            d = Path(tmpdir)
+            (d / "run.py").write_text(
+                "import torch\n"
+                "x = torch.randn(1024, device='cuda')\n"
+                "y = x * 2\n"
+            )
+
+            result = self._run_wafer(
+                "nvidia", "ncu", "run",
+                "--dir", str(d),
+                "--ncu-args", "--set full",
+                "--no-download",
+                "--timeout", "120",
+                "python", "run.py",
+            )
+
+            print("STDOUT:", result.stdout)
+            print("STDERR:", result.stderr)
+
+            # Should show the ncu args in the command
+            assert "--set full" in result.stdout
+
+    def test_ncu_run_no_download(self) -> None:
+        """E2E: --no-download skips report download."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            d = Path(tmpdir)
+            (d / "run.py").write_text(
+                "import torch\n"
+                "x = torch.randn(1024, device='cuda')\n"
+                "y = x * 2\n"
+            )
+
+            result = self._run_wafer(
+                "nvidia", "ncu", "run",
+                "--dir", str(d),
+                "--no-download",
+                "--timeout", "120",
+                "python", "run.py",
+            )
+
+            print("STDOUT:", result.stdout)
+
+            # Should NOT contain download messages
+            assert "Downloading report" not in result.stdout
+
+
+class TestNcuRunDirPackaging:
+    """E2E tests focused on directory packaging behavior."""
+
+    def test_large_directory_with_exclusions(self) -> None:
+        """Directories with lots of excludable content are packaged efficiently."""
+        from wafer.ncu_run import package_directory
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            d = Path(tmpdir)
+
+            # Create a realistic project with lots of junk
+            (d / "run.py").write_text("print('hello')")
+            (d / "kernel.cu").write_text("__global__ void k() {}")
+
+            # Lots of excludable stuff
+            (d / "__pycache__").mkdir()
+            for i in range(50):
+                (d / "__pycache__" / f"module_{i}.cpython-312.pyc").write_bytes(b"\x00" * 1000)
+
+            (d / "node_modules").mkdir()
+            (d / "node_modules" / "big_dep").mkdir()
+            (d / "node_modules" / "big_dep" / "index.js").write_text("// big")
+
+            (d / ".git").mkdir()
+            (d / ".git" / "objects").mkdir()
+
+            tarball, count = package_directory(d)
+
+            # Should only include the 2 real files
+            assert count == 2
+            # Tarball should be tiny (just run.py + kernel.cu)
+            assert len(tarball) < 1024  # < 1KB