Add GitHub Action for running GPU tests on Modal

.github/workflows/gpu_test.yml

@@ -0,0 +1,33 @@
+name: GPU Tests on Modal
+
+on:
+  push:
+    branches: [ "main" ]
+  pull_request:
+    branches: [ "main" ]
+  workflow_dispatch:  # Allow manual triggering
+
+jobs:
+  gpu-tests:
+    runs-on: ubuntu-latest
+
+    steps:
+    - uses: actions/checkout@v4
+
+    - name: Install uv
+      uses: astral-sh/setup-uv@v5
+      with:
+        enable-cache: true
+
+    - name: Set up Python
+      run: uv python install 3.12
+
+    - name: Install dependencies
+      run: uv sync --extra dev
+
+    - name: Run GPU tests on Modal
+      env:
+        MODAL_TOKEN_ID: ${{ secrets.MODAL_TOKEN_ID }}
+        MODAL_TOKEN_SECRET: ${{ secrets.MODAL_TOKEN_SECRET }}
+      run: |
+        uv run modal run -m llamax.ci_gpu

llamax/ci_gpu.py

@@ -0,0 +1,143 @@
+"""Modal CI configuration for running GPU tests."""
+
+from pathlib import Path
+
+import modal
+
+# Get the project root directory
+project_root = Path(__file__).parent.parent
+
+# Create a Modal image with all necessary dependencies using uv
+image = (
+    modal.Image.debian_slim(python_version="3.12")
+    .uv_sync(
+        uv_project_dir=project_root,
+        extras=["dev"],
+    )
+    .add_local_python_source("llamax", str(project_root / "llamax"))
+    .run_commands("cd /root && uv pip install -e .")
+    .env({"JAX_ENABLE_X64": "True"})
+)
+
+app = modal.App("llamax-gpu-tests", image=image)
+
+
+def create_gpu_test_function(gpu_type: str = "any"):
+    """Create a Modal function with the specified GPU type."""
+
+    @app.function(
+        gpu=gpu_type,
+        timeout=3600,  # 1 hour timeout for tests
+    )
+    def run_tests(test_filter: str = "gpu"):
+        """Run pytest on GPU tests with optional filter."""
+        import subprocess
+
+        # Build pytest command
+        cmd = [
+            "pytest",
+            "-v",
+            "llamax",
+            "-k",
+            test_filter,
+            "--tb=short",
+        ]
+
+        # Run pytest
+        result = subprocess.run(
+            cmd,
+            check=False,
+            capture_output=True,
+            text=True,
+        )
+
+        # Print output
+        print(result.stdout)
+        if result.stderr:
+            print(result.stderr)
+
+        # Exit with the same code as pytest
+        if result.returncode != 0:
+            raise SystemExit(result.returncode)
+
+        return result.returncode
+
+    return run_tests
+
+
+# Default function for backward compatibility (used by GitHub Actions)
+@app.function(
+    gpu="any",
+    timeout=3600,
+)
+def run_gpu_tests():
+    """Run pytest on all tests with 'gpu' in the name (backward compatible)."""
+    import subprocess
+    import sys
+
+    # Print debugging information
+    print("=== Modal GPU Test Environment ===")
+    print(f"Python version: {sys.version}")
+    print(f"Python executable: {sys.executable}")
+
+    # Check if pytest is available
+    try:
+        import pytest
+        print(f"pytest version: {pytest.__version__}")
+    except ImportError as e:
+        print(f"ERROR: pytest not found: {e}")
+        return 1
+
+    # Check if jax is available
+    try:
+        import jax
+        print(f"jax version: {jax.__version__}")
+        print(f"jax devices: {jax.devices()}")
+    except ImportError as e:
+        print(f"ERROR: jax not found: {e}")
+        return 1
+
+    result = subprocess.run(
+        [
+            "pytest",
+            "-v",
+            "llamax",
+            "-k",
+            "gpu",
+            "--tb=short",
+        ],
+        check=False,
+        capture_output=True,
+        text=True,
+    )
+
+    print("\n=== Pytest Output ===")
+    print(result.stdout)
+    if result.stderr:
+        print("\n=== Pytest Errors ===")
+        print(result.stderr)
+
+    if result.returncode != 0:
+        print(f"\n=== Test failed with exit code {result.returncode} ===")
+        raise SystemExit(result.returncode)
+
+    return result.returncode
+
+
+@app.local_entrypoint()
+def main(gpu_type: str = "any", filter: str = "gpu"):
+    """
+    Main entrypoint for running GPU tests locally via Modal.
+
+    Args:
+        gpu_type: GPU type to use (any, h100, a100, etc.)
+        filter: Test filter pattern for pytest -k option
+    """
+    print(f"Running GPU tests with GPU type: {gpu_type}")
+    print(f"Test filter: {filter}")
+
+    # Create function with specified GPU type
+    test_func = create_gpu_test_function(gpu_type)
+
+    # Run the tests
+    return test_func.remote(test_filter=filter)

llamax/example_gpu.py

@@ -0,0 +1,35 @@
+"""Example GPU test file to demonstrate the GPU testing pattern."""
+
+import unittest
+
+import jax
+import jax.numpy as jnp
+
+
+class ExampleGPUTest(unittest.TestCase):
+    """Example test class to verify GPU functionality."""
+
+    def test_gpu_available(self):
+        """Test that JAX can see GPU devices."""
+        devices = jax.devices()
+        # This test will pass on both CPU and GPU, but on Modal it should see GPUs
+        self.assertGreater(len(devices), 0)
+        print(f"Available devices: {devices}")
+
+    def test_simple_gpu_computation(self):
+        """Test a simple computation that would benefit from GPU."""
+        # Create a simple matrix multiplication
+        key = jax.random.PRNGKey(0)
+        a = jax.random.normal(key, (1000, 1000))
+        b = jax.random.normal(key, (1000, 1000))
+
+        # Perform computation
+        c = jnp.dot(a, b)
+
+        # Verify shape
+        self.assertEqual(c.shape, (1000, 1000))
+        print(f"Computation device: {c.device()}")
+
+
+if __name__ == "__main__":
+    unittest.main()

pyproject.toml

@@ -25,6 +25,7 @@ dependencies = [
 dev = [
     "ruff>=0.8.0",
     "pre-commit>=3.5.0",
+    "modal>=0.64.0",
 ]
 
 [build-system]
@@ -48,3 +49,11 @@ ignore = ["N812"]
 
 [tool.ruff.lint.per-file-ignores]
 "__init__.py" = ["F401"]
+
+[tool.pytest.ini_options]
+python_files = ["test_*.py", "*_test.py", "*_gpu.py"]
+python_classes = ["Test*"]
+python_functions = ["test_*"]
+testpaths = ["llamax"]
+# Exclude Modal CI and example files which require modal package
+addopts = "--ignore=llamax/ci_gpu.py --ignore=llamax/example_gpu.py"

scripts/README.md

@@ -0,0 +1,53 @@
+# Scripts
+
+## GPU Tests
+
+### `run_gpu_tests.sh`
+
+Run GPU tests on Modal using H100 or A100 instances.
+
+#### Usage
+
+```bash
+# Run all GPU tests on any available GPU
+./scripts/run_gpu_tests.sh
+
+# Run all GPU tests on H100
+./scripts/run_gpu_tests.sh --gpu=h100
+
+# Run all GPU tests on A100
+./scripts/run_gpu_tests.sh --gpu=a100
+
+# Run specific test pattern
+./scripts/run_gpu_tests.sh test_pattern
+
+# Combine GPU type and test pattern
+./scripts/run_gpu_tests.sh --gpu=h100 test_pattern
+```
+
+#### Examples
+
+```bash
+# Run all GPU tests on H100
+./scripts/run_gpu_tests.sh --gpu=h100
+
+# Run only tests with "integration" in the name on A100
+./scripts/run_gpu_tests.sh --gpu=a100 integration
+
+# Run tests matching "model" pattern
+./scripts/run_gpu_tests.sh model
+```
+
+#### Requirements
+
+- Modal account and authentication (`modal setup`)
+- `uv` package manager installed
+- Modal secrets configured (optional, for HuggingFace access)
+
+#### GPU Types
+
+- `any`: Use any available GPU (default)
+- `h100`: NVIDIA H100 GPU
+- `a100`: NVIDIA A100 GPU
+
+The test filter uses pytest's `-k` option, so you can use any pytest filter expression.

scripts/run_gpu_tests.sh

@@ -0,0 +1,51 @@
+#!/bin/bash
+# Run GPU tests on Modal using H100 or A100 instances
+#
+# Usage:
+#   ./scripts/run_gpu_tests.sh                    # Run all GPU tests
+#   ./scripts/run_gpu_tests.sh test_pattern       # Run tests matching pattern
+#   ./scripts/run_gpu_tests.sh --gpu=h100         # Specify GPU type (h100 or a100)
+#   ./scripts/run_gpu_tests.sh --gpu=a100 pattern # Combine GPU type and pattern
+
+set -e
+
+# Default values
+GPU_TYPE="any"
+TEST_FILTER=""
+
+# Parse arguments
+while [[ $# -gt 0 ]]; do
+    case $1 in
+        --gpu=*)
+            GPU_TYPE="${1#*=}"
+            shift
+            ;;
+        --gpu)
+            GPU_TYPE="$2"
+            shift 2
+            ;;
+        *)
+            TEST_FILTER="$1"
+            shift
+            ;;
+    esac
+done
+
+# Validate GPU type
+if [[ ! "$GPU_TYPE" =~ ^(any|h100|a100|H100|A100)$ ]]; then
+    echo "Error: Invalid GPU type '$GPU_TYPE'. Must be one of: any, h100, a100"
+    exit 1
+fi
+
+# Convert to lowercase for Modal
+GPU_TYPE=$(echo "$GPU_TYPE" | tr '[:upper:]' '[:lower:]')
+
+echo "Running GPU tests on Modal..."
+echo "GPU Type: $GPU_TYPE"
+if [ -n "$TEST_FILTER" ]; then
+    echo "Test Filter: $TEST_FILTER"
+    uv run modal run -m llamax.ci_gpu --gpu-type "$GPU_TYPE" --filter "$TEST_FILTER"
+else
+    echo "Test Filter: all GPU tests"
+    uv run modal run -m llamax.ci_gpu --gpu-type "$GPU_TYPE"
+fi