wafer-core 0.1.38__py3-none-any.whl → 0.1.40__py3-none-any.whl

wafer_core/tools/compile/compiler.py

@@ -0,0 +1,301 @@
+"""CUDA compiler wrapper.
+
+This module provides the core compilation logic that can be used locally
+or through Modal. It handles request/response conversion and provides
+a clean interface for the API and CLI.
+"""
+
+import subprocess
+from dataclasses import asdict
+
+from wafer_core.tools.compile.types import (
+    CompileRequest,
+    CompileResponse,
+    OutputFormat,
+)
+
+
+def request_to_dict(request: CompileRequest) -> dict:
+    """Convert a CompileRequest to a dict for Modal invocation.
+
+    Args:
+        request: The compile request
+
+    Returns:
+        Dict suitable for passing to Modal function
+    """
+    return {
+        "files": dict(request.files),
+        "arch": request.arch,
+        "flags": list(request.flags),
+        "output": [fmt.value for fmt in request.output],
+    }
+
+
+def response_from_dict(data: dict) -> CompileResponse:
+    """Convert a Modal response dict to a CompileResponse.
+
+    Args:
+        data: Dict returned from Modal function
+
+    Returns:
+        CompileResponse object
+    """
+    return CompileResponse(
+        success=data.get("success", False),
+        ptx=data.get("ptx"),
+        sass=data.get("sass"),
+        stderr=data.get("stderr", ""),
+        compilation_time_ms=data.get("compilation_time_ms", 0),
+    )
+
+
+async def compile_cuda_remote(
+    request: CompileRequest,
+    *,
+    modal_token_id: str | None = None,
+    modal_token_secret: str | None = None,
+) -> CompileResponse:
+    """Compile CUDA code using Modal (remote execution).
+
+    This function calls the deployed Modal function directly using asyncio.to_thread
+    to avoid blocking the event loop.
+
+    Args:
+        request: The compile request
+        modal_token_id: Optional Modal token ID (uses env var if not provided)
+        modal_token_secret: Optional Modal token secret
+
+    Returns:
+        CompileResponse with PTX/SASS or error
+    """
+    import asyncio
+    import os
+    import time
+    from contextlib import contextmanager
+
+    @contextmanager
+    def temporary_env_vars(env_updates: dict[str, str]):
+        """Context manager to temporarily set environment variables.
+
+        Saves original values, sets new values, yields, then restores originals.
+        This ensures we don't leak credentials between concurrent requests.
+        """
+        original_values: dict[str, str | None] = {}
+        for key, value in env_updates.items():
+            original_values[key] = os.environ.get(key)
+            os.environ[key] = value
+
+        try:
+            yield
+        finally:
+            for key, original in original_values.items():
+                if original is None:
+                    os.environ.pop(key, None)
+                else:
+                    os.environ[key] = original
+
+    start_time = time.time()
+    request_dict = request_to_dict(request)
+
+    # Build env updates for credentials (only if provided)
+    env_updates: dict[str, str] = {}
+    if modal_token_id:
+        env_updates["MODAL_TOKEN_ID"] = modal_token_id
+    if modal_token_secret:
+        env_updates["MODAL_TOKEN_SECRET"] = modal_token_secret
+
+    def call_modal() -> dict:
+        """Call Modal function synchronously (runs in thread pool)."""
+        import modal
+
+        # Look up the deployed function
+        compile_fn = modal.Function.from_name("cuda-compile", "compile_cuda")
+
+        # Call the function remotely
+        return compile_fn.remote(request_dict)
+
+    try:
+        # Run Modal call in thread pool with temporary credentials
+        # The context manager ensures env vars are restored after the call
+        def call_modal_with_env() -> dict:
+            with temporary_env_vars(env_updates):
+                return call_modal()
+
+        result = await asyncio.to_thread(call_modal_with_env)
+        return response_from_dict(result)
+
+    except Exception as e:
+        error_str = str(e)
+        # Check for common Modal auth errors
+        if "MODAL_TOKEN" in error_str or "AuthError" in error_str or "not authenticated" in error_str.lower():
+            return CompileResponse.error(
+                "Modal not configured. Set MODAL_TOKEN_ID and MODAL_TOKEN_SECRET environment variables, "
+                "or run 'modal token new' to authenticate.",
+                compilation_time_ms=int((time.time() - start_time) * 1000),
+            )
+        return CompileResponse.error(
+            f"Compilation failed: {error_str}",
+            compilation_time_ms=int((time.time() - start_time) * 1000),
+        )
+
+
+def compile_cuda_local(request: CompileRequest) -> CompileResponse:
+    """Compile CUDA code locally using nvcc.
+
+    This function requires nvcc to be installed locally.
+    Primarily useful for testing without Modal.
+
+    Args:
+        request: The compile request
+
+    Returns:
+        CompileResponse with PTX/SASS or error
+    """
+    import os
+    import tempfile
+    import time
+    from pathlib import Path
+
+    start_time = time.time()
+
+    # Check if nvcc is available
+    try:
+        subprocess.run(
+            ["nvcc", "--version"],
+            capture_output=True,
+            check=True,
+            timeout=10,
+        )
+    except (subprocess.CalledProcessError, FileNotFoundError, subprocess.TimeoutExpired):
+        return CompileResponse.error(
+            "nvcc not found. Install CUDA toolkit or use Modal for remote compilation.",
+            compilation_time_ms=int((time.time() - start_time) * 1000),
+        )
+
+    try:
+        with tempfile.TemporaryDirectory() as tmpdir:
+            tmp_path = Path(tmpdir)
+
+            # Write all files to temp directory
+            for filename, content in request.files.items():
+                file_path = (tmp_path / filename).resolve()
+                if not file_path.is_relative_to(tmp_path):
+                    return CompileResponse.error(
+                        f"Invalid filename: {filename}",
+                        compilation_time_ms=int((time.time() - start_time) * 1000),
+                    )
+                file_path.parent.mkdir(parents=True, exist_ok=True)
+                file_path.write_text(content)
+
+            main_cu_path = tmp_path / request.main_cu_file
+            include_dir = main_cu_path.parent
+
+            results: dict[str, str | None] = {"ptx": None, "sass": None}
+
+            # Build base nvcc command
+            base_cmd = [
+                "nvcc",
+                "-arch",
+                request.arch,
+                f"-I{include_dir}",
+            ]
+            base_cmd.extend(request.flags)
+
+            # Generate PTX if requested
+            if OutputFormat.PTX in request.output:
+                ptx_output = tmp_path / "output.ptx"
+                ptx_cmd = base_cmd + [
+                    "--ptx",
+                    "-o",
+                    str(ptx_output),
+                    str(main_cu_path),
+                ]
+
+                ptx_result = subprocess.run(
+                    ptx_cmd,
+                    capture_output=True,
+                    text=True,
+                    timeout=60,
+                    cwd=tmpdir,
+                )
+
+                if ptx_result.returncode != 0:
+                    return CompileResponse.error(
+                        ptx_result.stderr or ptx_result.stdout,
+                        compilation_time_ms=int((time.time() - start_time) * 1000),
+                    )
+
+                if ptx_output.exists():
+                    results["ptx"] = ptx_output.read_text()
+
+            # Generate SASS if requested
+            if OutputFormat.SASS in request.output:
+                cubin_output = tmp_path / "output.cubin"
+                cubin_cmd = base_cmd + [
+                    "--cubin",
+                    "-o",
+                    str(cubin_output),
+                    str(main_cu_path),
+                ]
+
+                cubin_result = subprocess.run(
+                    cubin_cmd,
+                    capture_output=True,
+                    text=True,
+                    timeout=60,
+                    cwd=tmpdir,
+                )
+
+                if cubin_result.returncode != 0:
+                    if results["ptx"]:
+                        return CompileResponse(
+                            success=True,
+                            ptx=results["ptx"],
+                            sass=None,
+                            stderr=f"SASS generation failed: {cubin_result.stderr}",
+                            compilation_time_ms=int((time.time() - start_time) * 1000),
+                        )
+                    return CompileResponse.error(
+                        cubin_result.stderr or cubin_result.stdout,
+                        compilation_time_ms=int((time.time() - start_time) * 1000),
+                    )
+
+                if cubin_output.exists():
+                    sass_result = subprocess.run(
+                        ["cuobjdump", "--dump-sass", str(cubin_output)],
+                        capture_output=True,
+                        text=True,
+                        timeout=30,
+                        cwd=tmpdir,
+                    )
+
+                    if sass_result.returncode == 0:
+                        results["sass"] = sass_result.stdout
+
+            if not results["ptx"] and not results["sass"]:
+                return CompileResponse.error(
+                    "No output generated",
+                    compilation_time_ms=int((time.time() - start_time) * 1000),
+                )
+
+            return CompileResponse(
+                success=True,
+                ptx=results["ptx"],
+                sass=results["sass"],
+                stderr="",
+                compilation_time_ms=int((time.time() - start_time) * 1000),
+            )
+
+    except subprocess.TimeoutExpired:
+        return CompileResponse.error(
+            "Compilation timed out",
+            compilation_time_ms=int((time.time() - start_time) * 1000),
+        )
+    except Exception as e:
+        import traceback
+
+        return CompileResponse.error(
+            f"Internal error: {e}\n{traceback.format_exc()}",
+            compilation_time_ms=int((time.time() - start_time) * 1000),
+        )

wafer_core/tools/compile/modal_compile.py

@@ -0,0 +1,369 @@
+"""Modal function for cloud CUDA compilation.
+
+This module provides a CPU-only Modal function that compiles CUDA code
+and returns PTX/SASS assembly. No GPU is required - nvcc can generate
+PTX/SASS for any architecture without a physical GPU.
+
+Usage:
+    # From Python
+    from wafer_core.tools.compile.modal_compile import compile_cuda
+    result = compile_cuda.remote(request_dict)
+
+    # Deploy Modal app
+    modal deploy modal_compile.py
+"""
+
+from enum import Enum
+
+import modal
+
+
+# Inline OutputFormat to avoid import chain issues with Modal deployment
+class OutputFormat(str, Enum):
+    """Supported output formats."""
+
+    PTX = "ptx"
+    SASS = "sass"
+
+# ══════════════════════════════════════════════════════════════════════════════
+# Modal Image Configuration
+# ══════════════════════════════════════════════════════════════════════════════
+
+# Build a CPU-only image with nvcc for CUDA compilation
+# No GPU driver needed - nvcc can generate PTX/SASS without a GPU
+compile_image = (
+    # Start with NVIDIA's CUDA development image (includes nvcc)
+    # Use debian-slim + manual CUDA install for smaller image
+    modal.Image.from_registry(
+        "nvidia/cuda:13.0.1-devel-ubuntu22.04",
+        add_python="3.12",
+    )
+    # System dependencies
+    .apt_install("git", "build-essential", "wget")
+    # Install PyTorch headers (CPU-only, we just need the headers)
+    .pip_install(
+        "torch>=2.5.0",
+        index_url="https://download.pytorch.org/whl/cpu",
+        extra_index_url="https://pypi.org/simple",
+    )
+    # Install CUTLASS headers (v4.3.5)
+    .run_commands(
+        "git clone --depth 1 --branch v4.3.5 https://github.com/NVIDIA/cutlass.git /usr/local/cutlass",
+    )
+    # Set environment variables for CUDA and libraries
+    .env(
+        {
+            "CUDA_HOME": "/usr/local/cuda",
+            "PATH": "/usr/local/cuda/bin:$PATH",
+            "CUTLASS_PATH": "/usr/local/cutlass/include",
+        }
+    )
+    # Verify nvcc is working
+    .run_commands(
+        "nvcc --version",
+        "ls -la /usr/local/cuda/bin/nvcc",
+    )
+)
+
+# Create Modal app
+app = modal.App(name="cuda-compile", image=compile_image)
+
+
+# ══════════════════════════════════════════════════════════════════════════════
+# Modal Function - CUDA Compilation
+# ══════════════════════════════════════════════════════════════════════════════
+
+
+@app.function(
+    # CPU only - no GPU needed for compilation!
+    cpu=4,
+    memory=8192,  # 8GB RAM
+    timeout=120,  # 2 minute timeout
+    # Keep one container warm to avoid cold starts (~5-10s savings)
+    min_containers=1,
+)
+@modal.concurrent(max_inputs=4)  # Allow concurrent compilations for better throughput
+def compile_cuda(request: dict) -> dict:
+    """Compile CUDA code and return PTX/SASS.
+
+    This function runs on a CPU-only container with nvcc installed.
+    No GPU is required because nvcc can cross-compile for any target architecture.
+
+    Args:
+        request: Dict with the following fields:
+            - files: dict[str, str] - Mapping of filename to content
+            - arch: str - Target architecture (e.g., "sm_90a")
+            - flags: list[str] - Additional nvcc flags
+            - output: list[str] - Output formats ("ptx", "sass")
+
+    Returns:
+        Dict with the following fields:
+            - success: bool - Whether compilation succeeded
+            - ptx: str | None - Generated PTX code
+            - sass: str | None - Generated SASS code
+            - stderr: str - Compiler warnings/errors
+            - compilation_time_ms: int - Time taken in ms
+    """
+    import os
+    import subprocess
+    import tempfile
+    import time
+    from concurrent.futures import ThreadPoolExecutor, as_completed
+    from pathlib import Path
+
+    start_time = time.time()
+
+    # Extract request fields
+    files: dict[str, str] = request.get("files", {})
+    arch: str = request.get("arch", "sm_90a")
+    flags: list[str] = request.get("flags", [])
+    output_formats: list[str] = request.get("output", ["ptx", "sass"])
+
+    # Validate request
+    if not files:
+        return {
+            "success": False,
+            "ptx": None,
+            "sass": None,
+            "stderr": "No files provided",
+            "compilation_time_ms": int((time.time() - start_time) * 1000),
+        }
+
+    # Find the main .cu file
+    cu_files = [f for f in files.keys() if f.endswith(".cu")]
+    if not cu_files:
+        return {
+            "success": False,
+            "ptx": None,
+            "sass": None,
+            "stderr": "No .cu file found",
+            "compilation_time_ms": int((time.time() - start_time) * 1000),
+        }
+
+    main_cu_file = cu_files[0]
+
+    # Build environment for nvcc
+    nvcc_env = {
+        **os.environ,
+        "CUDA_HOME": "/usr/local/cuda",
+        "PATH": f"/usr/local/cuda/bin:{os.environ.get('PATH', '')}",
+    }
+
+    def compile_ptx(tmpdir: str, base_cmd: list[str], main_cu_path: Path) -> tuple[str | None, str | None]:
+        """Compile to PTX. Returns (ptx_content, error_message)."""
+        ptx_output = Path(tmpdir) / "output.ptx"
+        ptx_cmd = base_cmd + [
+            "--ptx",
+            "-o",
+            str(ptx_output),
+            str(main_cu_path),
+        ]
+
+        ptx_result = subprocess.run(
+            ptx_cmd,
+            capture_output=True,
+            text=True,
+            timeout=60,
+            cwd=tmpdir,
+            env=nvcc_env,
+        )
+
+        if ptx_result.returncode != 0:
+            return None, ptx_result.stderr or ptx_result.stdout
+
+        if ptx_output.exists():
+            return ptx_output.read_text(), None
+        return None, "PTX output file not created"
+
+    def compile_sass(tmpdir: str, base_cmd: list[str], main_cu_path: Path) -> tuple[str | None, str | None]:
+        """Compile to SASS (via cubin). Returns (sass_content, error_message)."""
+        cubin_output = Path(tmpdir) / "output.cubin"
+        cubin_cmd = base_cmd + [
+            "--cubin",
+            "-o",
+            str(cubin_output),
+            str(main_cu_path),
+        ]
+
+        cubin_result = subprocess.run(
+            cubin_cmd,
+            capture_output=True,
+            text=True,
+            timeout=60,
+            cwd=tmpdir,
+            env=nvcc_env,
+        )
+
+        if cubin_result.returncode != 0:
+            return None, cubin_result.stderr or cubin_result.stdout
+
+        if not cubin_output.exists():
+            return None, "cubin output file not created"
+
+        # Disassemble cubin to SASS
+        sass_result = subprocess.run(
+            ["cuobjdump", "--dump-sass", str(cubin_output)],
+            capture_output=True,
+            text=True,
+            timeout=30,
+            cwd=tmpdir,
+        )
+
+        if sass_result.returncode == 0:
+            return sass_result.stdout, None
+        return None, f"SASS disassembly failed: {sass_result.stderr}"
+
+    try:
+        with tempfile.TemporaryDirectory() as tmpdir:
+            tmp_path = Path(tmpdir)
+
+            # Write all files to temp directory, preserving subdirectory structure
+            for filename, content in files.items():
+                file_path = (tmp_path / filename).resolve()
+                if not file_path.is_relative_to(tmp_path):
+                    return {
+                        "success": False,
+                        "ptx": None,
+                        "sass": None,
+                        "stderr": f"Invalid filename: {filename}",
+                        "compilation_time_ms": int((time.time() - start_time) * 1000),
+                    }
+                file_path.parent.mkdir(parents=True, exist_ok=True)
+                file_path.write_text(content)
+
+            # Determine the directory of the main .cu file for include paths
+            main_cu_path = tmp_path / main_cu_file
+            include_dir = main_cu_path.parent
+
+            # Build base nvcc command with common flags
+            base_cmd = [
+                "nvcc",
+                "-arch",
+                arch,
+                f"-I{include_dir}",
+                "-I/usr/local/lib/python3.12/site-packages/torch/include",
+                "-I/usr/local/lib/python3.12/site-packages/torch/include/torch/csrc/api/include",
+                "-I/usr/local/cutlass/include",
+            ]
+            base_cmd.extend(flags)
+
+            # Determine what to compile
+            want_ptx = OutputFormat.PTX.value in output_formats
+            want_sass = OutputFormat.SASS.value in output_formats
+
+            results: dict[str, str | None] = {"ptx": None, "sass": None}
+            errors: list[str] = []
+
+            # Run compilations in parallel if both are requested
+            if want_ptx and want_sass:
+                with ThreadPoolExecutor(max_workers=2) as executor:
+                    futures = {
+                        executor.submit(compile_ptx, tmpdir, base_cmd, main_cu_path): "ptx",
+                        executor.submit(compile_sass, tmpdir, base_cmd, main_cu_path): "sass",
+                    }
+
+                    for future in as_completed(futures):
+                        output_type = futures[future]
+                        try:
+                            content, error = future.result()
+                            if content:
+                                results[output_type] = content
+                            if error:
+                                errors.append(f"{output_type.upper()}: {error}")
+                        except Exception as e:
+                            errors.append(f"{output_type.upper()} compilation error: {e}")
+
+            elif want_ptx:
+                content, error = compile_ptx(tmpdir, base_cmd, main_cu_path)
+                if content:
+                    results["ptx"] = content
+                if error:
+                    errors.append(error)
+
+            elif want_sass:
+                content, error = compile_sass(tmpdir, base_cmd, main_cu_path)
+                if content:
+                    results["sass"] = content
+                if error:
+                    errors.append(error)
+
+            # Check results
+            if not results["ptx"] and not results["sass"]:
+                return {
+                    "success": False,
+                    "ptx": None,
+                    "sass": None,
+                    "stderr": "\n".join(errors) if errors else "No output generated",
+                    "compilation_time_ms": int((time.time() - start_time) * 1000),
+                }
+
+            # Partial success if we got at least one output
+            stderr = ""
+            if errors and (results["ptx"] or results["sass"]):
+                stderr = "\n".join(errors)
+
+            return {
+                "success": True,
+                "ptx": results["ptx"],
+                "sass": results["sass"],
+                "stderr": stderr,
+                "compilation_time_ms": int((time.time() - start_time) * 1000),
+            }
+
+    except subprocess.TimeoutExpired:
+        return {
+            "success": False,
+            "ptx": None,
+            "sass": None,
+            "stderr": "Compilation timed out",
+            "compilation_time_ms": int((time.time() - start_time) * 1000),
+        }
+    except Exception as e:
+        import traceback
+
+        return {
+            "success": False,
+            "ptx": None,
+            "sass": None,
+            "stderr": f"Internal error: {e}\n{traceback.format_exc()}",
+            "compilation_time_ms": int((time.time() - start_time) * 1000),
+        }
+
+
+# ══════════════════════════════════════════════════════════════════════════════
+# Health Check
+# ══════════════════════════════════════════════════════════════════════════════
+
+
+@app.function(cpu=1, memory=1024, timeout=30)
+def health_check() -> dict:
+    """Verify the compilation environment is working.
+
+    Returns:
+        Dict with status and version info
+    """
+    import subprocess
+
+    try:
+        # Get nvcc version
+        nvcc_result = subprocess.run(
+            ["nvcc", "--version"],
+            capture_output=True,
+            text=True,
+            timeout=10,
+        )
+
+        # Parse version from output
+        version_line = nvcc_result.stdout.strip().split("\n")[-1]
+
+        return {
+            "status": "ok",
+            "nvcc_version": version_line,
+            "nvcc_available": nvcc_result.returncode == 0,
+        }
+    except Exception as e:
+        return {
+            "status": "error",
+            "error": str(e),
+            "nvcc_available": False,
+        }

wafer_core/tools/compile/tests/__init__.py

@@ -0,0 +1 @@
+"""Tests for the cloud CUDA compiler."""