wafer-cli 0.2.29__py3-none-any.whl → 0.2.31__py3-none-any.whl

wafer/evaluate.py

@@ -78,9 +78,10 @@ def _build_docker_run_command(
         for cap in cap_add:
             parts.extend(["--cap-add", cap])
 
-    # GPU access - use single quotes for the device spec to avoid shell escaping issues
+    # GPU access - use --runtime=nvidia alongside --gpus for compatibility
+    # with newer NVIDIA drivers (580+) where --gpus alone may not initialize CUDA
     if gpus:
-        parts.extend(["--gpus", f"'{gpus}'"])
+        parts.extend(["--runtime=nvidia", "--gpus", f"'{gpus}'"])
 
     # Volume mounts
     if volumes:
@@ -379,18 +380,6 @@ def _build_docker_pip_install_cmd(target: BaremetalTarget | VMTarget) -> str:
     return " && ".join(commands)
 
 
-def _get_wafer_root() -> Path:
-    """Get wafer monorepo root directory.
-
-    Walks up from this file to find the wafer repo root (contains apps/, packages/).
-    """
-    current = Path(__file__).resolve()
-    for parent in [current] + list(current.parents):
-        if (parent / "apps").is_dir() and (parent / "packages").is_dir():
-            return parent
-    raise RuntimeError(f"Could not find wafer root from {__file__}")
-
-
 async def run_evaluate_docker(
     args: EvaluateArgs,
     target: BaremetalTarget | VMTarget,
@@ -2033,54 +2022,13 @@ async def run_evaluate_runpod(
                         error_message=f"Failed to setup Python environment: {e}",
                     )
 
-                # Upload wafer-core to remote
-                try:
-                    wafer_root = _get_wafer_root()
-                    wafer_core_path = wafer_root / "packages" / "wafer-core"
-                    print(f"Uploading wafer-core from {wafer_core_path}...")
-
-                    wafer_core_remote = f"{REMOTE_WORKSPACE}/wafer-core"
-                    await client.exec(f"mkdir -p {wafer_core_remote}")
-                    wafer_core_workspace = await client.expand_path(wafer_core_remote)
-
-                    upload_result = await client.upload_files(
-                        str(wafer_core_path), wafer_core_workspace, recursive=True
-                    )
-
-                    # Wide event logging for upload result
-                    upload_event = {
-                        "event": "wafer_core_upload",
-                        "target": target.name,
-                        "target_type": "runpod",
-                        "ssh_host": f"{client.user}@{client.host}:{client.port}",
-                        "local_path": str(wafer_core_path),
-                        "remote_path": wafer_core_workspace,
-                        "success": upload_result.success,
-                        "files_copied": upload_result.files_copied,
-                        "duration_seconds": upload_result.duration_seconds,
-                        "error_message": upload_result.error_message,
-                    }
-                    if upload_result.debug_info:
-                        upload_event["debug_info"] = upload_result.debug_info
-                    logger.info(json.dumps(upload_event))
-
-                    # Fail fast if upload failed
-                    if not upload_result.success:
-                        print(f"ERROR: Upload failed: {upload_result.error_message}")
-                        if upload_result.debug_info:
-                            print(f"Debug info: {json.dumps(upload_result.debug_info, indent=2)}")
-                        return EvaluateResult(
-                            success=False,
-                            all_correct=False,
-                            correctness_score=0.0,
-                            geomean_speedup=0.0,
-                            passed_tests=0,
-                            total_tests=0,
-                            error_message=f"Failed to upload wafer-core: {upload_result.error_message}",
-                        )
-
-                    print(f"Uploaded {upload_result.files_copied} files")
-                except Exception as e:
+                # Install wafer-core in remote venv
+                print("Installing wafer-core...")
+                install_result = await client.exec(
+                    f'export PATH="$HOME/.local/bin:$HOME/.cargo/bin:$PATH" && '
+                    f"uv pip install --python {python_exe} wafer-core"
+                )
+                if install_result.exit_code != 0:
                     return EvaluateResult(
                         success=False,
                         all_correct=False,
@@ -2088,7 +2036,7 @@ async def run_evaluate_runpod(
                         geomean_speedup=0.0,
                         passed_tests=0,
                         total_tests=0,
-                        error_message=f"Failed to upload wafer-core: {e}",
+                        error_message=f"Failed to install wafer-core: {install_result.stderr}",
                     )
 
                 # Select GPU (RunPod pods typically have GPU 0)
@@ -2229,11 +2177,18 @@ async def run_evaluate_runpod(
                         error_message=f"Evaluation timed out after {target.eval_timeout}s",
                     )
 
-                # Parse output
+                # Show output to user
                 stdout = result.stdout
                 stderr = result.stderr
+                if stdout:
+                    print(stdout)
 
                 if result.exit_code != 0:
+                    error_parts = [f"Evaluation failed (exit code {result.exit_code}):"]
+                    if stdout:
+                        error_parts.append(f"stdout: {stdout}")
+                    if stderr:
+                        error_parts.append(f"stderr: {stderr}")
                     return EvaluateResult(
                         success=False,
                         all_correct=False,
@@ -2241,20 +2196,27 @@ async def run_evaluate_runpod(
                         geomean_speedup=0.0,
                         passed_tests=0,
                         total_tests=0,
-                        error_message=f"Evaluation failed:\nstdout: {stdout}\nstderr: {stderr}",
+                        error_message="\n".join(error_parts),
                     )
 
-                # Find JSON result in output
-                result_json = None
-                for line in reversed(stdout.strip().split("\n")):
-                    if line.startswith("{"):
-                        try:
-                            result_json = json.loads(line)
-                            break
-                        except json.JSONDecodeError:
-                            continue
-
-                if result_json is None:
+                # Read results from results.json file written by evaluate module
+                results_path = f"{run_path}/results.json"
+                cat_result = await client.exec(f"cat {results_path}")
+
+                if cat_result.exit_code != 0:
+                    return EvaluateResult(
+                        success=False,
+                        all_correct=False,
+                        correctness_score=0.0,
+                        geomean_speedup=0.0,
+                        passed_tests=0,
+                        total_tests=0,
+                        error_message=f"Failed to read results: {cat_result.stderr}",
+                    )
+
+                try:
+                    results_data = json.loads(cat_result.stdout)
+                except json.JSONDecodeError as e:
                     return EvaluateResult(
                         success=False,
                         all_correct=False,
@@ -2262,10 +2224,12 @@ async def run_evaluate_runpod(
                         geomean_speedup=0.0,
                         passed_tests=0,
                         total_tests=0,
-                        error_message=f"No JSON result in output:\n{stdout}",
+                        error_message=f"Invalid JSON in results: {e}",
                     )
 
-                if "error" in result_json:
+                # Extract backend results (same format as DigitalOcean/SSH path)
+                backends = results_data.get("backends", [])
+                if not backends:
                     return EvaluateResult(
                         success=False,
                         all_correct=False,
@@ -2273,18 +2237,20 @@ async def run_evaluate_runpod(
                         geomean_speedup=0.0,
                         passed_tests=0,
                         total_tests=0,
-                        error_message=result_json["error"],
+                        error_message="No backend results found",
                     )
 
-                passed = result_json.get("passed", 0)
-                total = result_json.get("total", 0)
+                backend = backends[0]
+                correctness_tests = backend.get("correctness_tests", [])
+                passed = sum(1 for t in correctness_tests if t.get("is_correct", False))
+                total = len(correctness_tests)
                 correctness = passed / total if total > 0 else 0.0
 
                 return EvaluateResult(
                     success=True,
-                    all_correct=result_json.get("all_correct", False),
+                    all_correct=backend.get("all_correct", False),
                     correctness_score=correctness,
-                    geomean_speedup=result_json.get("speedup", 0.0),
+                    geomean_speedup=backend.get("geomean_speedup", 0.0),
                     passed_tests=passed,
                     total_tests=total,
                 )
@@ -2385,61 +2351,13 @@ async def run_evaluate_digitalocean(
                             error_message=f"Failed to setup Python environment: {e}",
                         )
 
-                    # Upload wafer-core to remote
-                    try:
-                        wafer_root = _get_wafer_root()
-                        wafer_core_path = wafer_root / "packages" / "wafer-core"
-                        print(f"Uploading wafer-core from {wafer_core_path}...")
-
-                        wafer_core_remote = f"{REMOTE_WORKSPACE}/wafer-core"
-                        await client.exec(f"mkdir -p {wafer_core_remote}")
-                        wafer_core_workspace = await client.expand_path(wafer_core_remote)
-
-                        # Use SFTP instead of rsync to avoid SSH subprocess timeout issues
-                        # (DigitalOcean may rate-limit new SSH connections)
-                        upload_result = await client.upload_files(
-                            str(wafer_core_path),
-                            wafer_core_workspace,
-                            recursive=True,
-                            use_sftp=True,
-                        )
-
-                        # Wide event logging for upload result
-                        upload_event = {
-                            "event": "wafer_core_upload",
-                            "target": target.name,
-                            "target_type": "digitalocean",
-                            "ssh_host": f"{client.user}@{client.host}:{client.port}",
-                            "local_path": str(wafer_core_path),
-                            "remote_path": wafer_core_workspace,
-                            "success": upload_result.success,
-                            "files_copied": upload_result.files_copied,
-                            "duration_seconds": upload_result.duration_seconds,
-                            "error_message": upload_result.error_message,
-                        }
-                        if upload_result.debug_info:
-                            upload_event["debug_info"] = upload_result.debug_info
-                        logger.info(json.dumps(upload_event))
-
-                        # Fail fast if upload failed
-                        if not upload_result.success:
-                            print(f"ERROR: Upload failed: {upload_result.error_message}")
-                            if upload_result.debug_info:
-                                print(
-                                    f"Debug info: {json.dumps(upload_result.debug_info, indent=2)}"
-                                )
-                            return EvaluateResult(
-                                success=False,
-                                all_correct=False,
-                                correctness_score=0.0,
-                                geomean_speedup=0.0,
-                                passed_tests=0,
-                                total_tests=0,
-                                error_message=f"Failed to upload wafer-core: {upload_result.error_message}",
-                            )
-
-                        print(f"Uploaded {upload_result.files_copied} files")
-                    except Exception as e:
+                    # Install wafer-core in remote venv
+                    print("Installing wafer-core...")
+                    install_result = await client.exec(
+                        f'export PATH="$HOME/.local/bin:$HOME/.cargo/bin:$PATH" && '
+                        f"uv pip install --python {python_exe} wafer-core"
+                    )
+                    if install_result.exit_code != 0:
                         return EvaluateResult(
                             success=False,
                             all_correct=False,
@@ -2447,7 +2365,7 @@ async def run_evaluate_digitalocean(
                             geomean_speedup=0.0,
                             passed_tests=0,
                             total_tests=0,
-                            error_message=f"Failed to upload wafer-core: {e}",
+                            error_message=f"Failed to install wafer-core: {install_result.stderr}",
                         )
 
                     # Select GPU (DigitalOcean droplets typically have GPU 0)
@@ -3242,15 +3160,35 @@ def main():
             inputs = [x.cuda() if isinstance(x, torch.Tensor) else x for x in inputs]
 
             if run_defense and defense_module is not None:
-                # Use full defense suite
+                # Use extended defense suite (Makora taxonomy + CUDA-L2)
                 print("[KernelBench] Running defense checks on implementation...")
-                run_all_defenses = defense_module.run_all_defenses
+                run_extended = defense_module.run_all_defenses_extended
                 time_with_defenses = defense_module.time_execution_with_defenses
 
-                # Run defense checks on implementation
-                all_passed, defense_results, _ = run_all_defenses(
+                # Read source code for LLM adversarial evaluator
+                _problem_code = None
+                _kernel_code = None
+                try:
+                    _problem_code = Path(args.reference).read_text()
+                    _kernel_code = Path(args.impl).read_text()
+                except Exception:
+                    pass
+
+                # Input generator for caching/multi-input checks
+                def _input_generator():
+                    _ins = get_inputs()
+                    return tuple(x.cuda() if isinstance(x, torch.Tensor) else x for x in _ins)
+
+                # Run all defense checks (original + extended)
+                all_passed, defense_results, _ = run_extended(
                     lambda *x: new_model(*x),
                     *inputs,
+                    reference_fn=lambda *x: ref_model(*x),
+                    input_generator=_input_generator,
+                    test_shapes=[(128, 128), (256, 256), (512, 512)],
+                    check_precision_ulp=True,
+                    problem_code=_problem_code,
+                    kernel_code=_kernel_code,
                 )
                 results["defense_results"] = {
                     name: {"passed": passed, "message": msg}

wafer/templates/optimize_kernel.py

@@ -35,7 +35,8 @@ Strategy:
 Commands:
 - `wafer evaluate --impl <file> --reference <ref> --test-cases <tests>` - Run evaluation
 - `wafer evaluate --impl <file> --reference <ref> --test-cases <tests> --profile` - With NCU profiling
-- `wafer remote-run "<command>"` - Run arbitrary commands on remote GPU
+- `wafer workspaces exec -- <command>` - Run arbitrary commands on remote GPU
+- `wafer targets exec <target> -- <command>` - Run commands on a configured target via SSH
 
 Output:
 - Summary of optimizations applied
@@ -48,7 +49,8 @@ IMPORTANT: Always verify correctness with wafer evaluate before claiming success
     tools=["read", "write", "edit", "glob", "grep", "bash"],
     bash_allowlist=[
         "wafer evaluate",
-        "wafer remote-run",
+        "wafer workspaces exec",
+        "wafer targets exec",
         "wafer nvidia ncu",
         "wafer nvidia nsys",
         "wafer nvidia perfetto",

{wafer_cli-0.2.29.dist-info → wafer_cli-0.2.31.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: wafer-cli
-Version: 0.2.29
+Version: 0.2.31
 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.29.dist-info → wafer_cli-0.2.31.dist-info}/RECORD

@@ -5,12 +5,13 @@ wafer/analytics.py,sha256=qLY6Z16usVHFD8TCv7XBuz7l47vXVdXk-qhOzA-hW_8,8179
 wafer/api_client.py,sha256=i_Az2b2llC3DSW8yOL-BKqa7LSKuxOr8hSN40s-oQXY,6313
 wafer/auth.py,sha256=dwss_se5P-FFc9IN38q4kh_dBrA6k-CguDBkivgcdj0,14003
 wafer/autotuner.py,sha256=41WYP41pTDvMijv2h42vm89bcHtDMJXObDlWmn6xpFU,44416
+wafer/baseline.py,sha256=OrGCAut_xtkH9Ogx4mMU5-94Q0oClIXqac94YRwqERY,21534
 wafer/billing.py,sha256=hEEwtrtIsbPQ3lLJNcyTLMsapUbcuvcVW_e9_0SxzVo,7199
-wafer/cli.py,sha256=zuVZhPdML5AOBtLUqLwAwjl8XMNe9EwQkffZxtBGLx4,282748
+wafer/cli.py,sha256=9wpLZlrKC3_DtkYuNWH42H1x9gJJOkJ32X2bNECG0rY,273013
 wafer/cli_instructions.py,sha256=bziUKDNDAXABVMvKPLEMXm-hFSD2TcFSh-FKRYa949k,4693
 wafer/config.py,sha256=h5Eo9_yfWqWGoPNdVQikI9GoZVUeysunSYiixf1mKcw,3411
 wafer/corpus.py,sha256=CY9T7wXENNDJxnrtI-XsQmXeptrFfKG4x-lngrc9_3s,24748
-wafer/evaluate.py,sha256=HMFQD-uwC6Wky1t_0JxYZaoHWgLaTBkjxOxgpZVnGrc,190519
+wafer/evaluate.py,sha256=i15PliAVI3W04_4eju46PBDdh2BwSToLME5n7yGu7dU,187355
 wafer/global_config.py,sha256=iu1HbTDr1695tSeDG2NfkK7PiY7XD6vjCk37w1wHbgk,11920
 wafer/gpu_run.py,sha256=TwqXy72T7f2I7e6n5WWod3xgxCPnDhU0BgLsB4CUoQY,9716
 wafer/inference.py,sha256=tZCO5i05FKY27ewis3CSBHFBeFbXY3xwj0DSjdoMY9s,4314
@@ -36,12 +37,12 @@ wafer/workspaces.py,sha256=J-TXGwHXSZlzRWCew63KNvk6HLJ-zTSELRgzjryTkMk,35710
 wafer/skills/wafer-guide/SKILL.md,sha256=UDsXCD5Kb-lDParKCTf2WkE3kodVs-rja8XeumSBO5U,3934
 wafer/templates/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 wafer/templates/ask_docs.py,sha256=15t1Aa4WBMwMox8XmFdzyosOZfBLMdXyaxo3GDb7nTE,2254
-wafer/templates/optimize_kernel.py,sha256=4-MaKm_C9BQHQEllrNLLYkcdhJpcj6D-8zbJ4FdLUEY,2444
+wafer/templates/optimize_kernel.py,sha256=Q4FA_8ECEegW_3DS51mkLCX6Vk1dcWWzY3A_RQ4NW8U,2576
 wafer/templates/optimize_kernelbench.py,sha256=T3co9Y9eSLWDrZG66gwQVFMdnGVoyUQos-TxnMMBLL8,3747
 wafer/templates/trace_analyze.py,sha256=B7CiRlsokERzBjLL-k49kGjpU2zlJZqzTE05xbRS1WI,2878
 wafer/tests/test_eval_cli_parity.py,sha256=SGmaj2NGBZ7GdDF53bXsECvQbV21iHZw8YeL_MJOLk0,7206
-wafer_cli-0.2.29.dist-info/METADATA,sha256=Cqwe6AdROCtQG3Xo6BzyICS2n5Rvtk-ex85P_s0PSIU,2799
-wafer_cli-0.2.29.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
-wafer_cli-0.2.29.dist-info/entry_points.txt,sha256=WqB7hB__WhtPY8y1cO2sZiUz7fCq6Ik-usAigpeFvWE,41
-wafer_cli-0.2.29.dist-info/top_level.txt,sha256=2MK1IVMWfpLL8BZCQ3E9aG6L6L666gSA_teYlwan4fs,6
-wafer_cli-0.2.29.dist-info/RECORD,,
+wafer_cli-0.2.31.dist-info/METADATA,sha256=rrIOyDKT02oelCRNJTTdlCvidNgVXmNSlE2vXmRsa-U,2799
+wafer_cli-0.2.31.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
+wafer_cli-0.2.31.dist-info/entry_points.txt,sha256=WqB7hB__WhtPY8y1cO2sZiUz7fCq6Ik-usAigpeFvWE,41
+wafer_cli-0.2.31.dist-info/top_level.txt,sha256=2MK1IVMWfpLL8BZCQ3E9aG6L6L666gSA_teYlwan4fs,6
+wafer_cli-0.2.31.dist-info/RECORD,,