Revert "Use Modal for CI" (#314)

* Revert "Use Modal for CI (#311)"

This reverts commit e413950.

* Trigger CI

* update

Author: Mark Saroufim

.github/workflows/runner_ci.yml

@@ -11,25 +11,59 @@ on:
       - dev
 
 jobs:
-  check-modal:
-    runs-on: ubuntu-latest
-    timeout-minutes: 30
+  check-cuda:
+    runs-on: [gpumode-nvidia-arc]
+    timeout-minutes: 10
+    container:
+      image: nvidia/cuda:12.4.0-devel-ubuntu22.04
     steps:
       - uses: actions/checkout@v3
 
       - name: Setup Python
         uses: actions/setup-python@v5
         with:
-          python-version: '3.13'
+          python-version: '3.10'
 
-      - name: Install dependencies
-        run: |
-          pip install modal
-          pip install -r requirements.txt
+      - name: Install pytest
+        shell: bash
+        run: pip install pytest
+
+      - name: Run script
+        shell: bash
+        run: pytest scripts/ci_test_cuda.py
+
+    env:
+      CUDA_VISIBLE_DEVICES: 0
+
+  check-pytorch:
+    runs-on: [gpumode-nvidia-arc]
+    timeout-minutes: 10
+    container:
+      image: nvidia/cuda:12.4.0-devel-ubuntu22.04
+    steps:
+      - uses: actions/checkout@v3
 
-      - name: Run Modal tests
+      - name: Setup Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: '3.10'
+
+      - name: Install uv
+        uses: astral-sh/setup-uv@v3
+        with:
+          version: "latest"
+
+      - name: Setup Python environment
         run: |
-          python scripts/modal_ci_test.py
-        env:
-          MODAL_TOKEN_ID: ${{ secrets.MODAL_TOKEN_ID }}
-          MODAL_TOKEN_SECRET: ${{ secrets.MODAL_TOKEN_SECRET }}
+          uv venv .venv
+          echo "VIRTUAL_ENV=$PWD/.venv" >> $GITHUB_ENV
+          echo "$PWD/.venv/bin" >> $GITHUB_PATH
+          uv pip install numpy torch setuptools ninja pytest
+
+      - name: Run script
+        shell: bash
+        run: pytest scripts/ci_test_python.py
+
+    env:
+      CUDA_VISIBLE_DEVICES: 0
+

requirements.txt

@@ -10,7 +10,6 @@ yoyo-migrations
 ruff
 pre-commit
 better_profanity
-pytest
 
 # api
 fastapi[all] # install all to avoid random bugs

scripts/ci_test_python.py

@@ -59,7 +59,7 @@ def custom_kernel(input):
     assert run.result["test.0.status"] == "fail"
     assert (
         run.result["test.0.error"]
-        == "mismatch found! custom implementation doesn't match reference.:"
+        == "mismatch found! custom implementation doesn't match reference:"
         " Number of mismatched elements: 256"
     )
     assert run.exit_code == ExitCode.VALIDATE_FAIL

scripts/modal_ci_test.py

@@ -28,18 +28,16 @@
         "requests~=2.32.4",
         "packaging~=25.0",
         "numpy~=2.3",
-        "pytest",
-
     )
     .pip_install(
         "torch~=2.7",
         "torchvision~=0.22",
         "torchaudio~=2.7",
-        index_url="https://download.pytorch.org/whl/cu128",
+        index_url="https://download.pytorch.org/whl/cu128"
     )
     # other frameworks
     .pip_install(
-        "jax[cuda12]==0.5.3",  # 0.6 want's cudnn 9.8 in conflict with torch 2.7
+        "jax[cuda12]==0.5.3",   # 0.6 want's cudnn 9.8 in conflict with torch 2.7
         "jax2torch==0.0.7",
         "tinygrad~=0.10",
     )
@@ -49,8 +47,8 @@
         "nvidia-cutlass-dsl~=4.0",
         "cuda-core[cu12]~=0.3",
         "cuda-python[all]==12.8",
-        # "nvmath-python[cu12]~=0.4",
-        # "numba-cuda[cu12]~=0.15",
+        #"nvmath-python[cu12]~=0.4",
+        #"numba-cuda[cu12]~=0.15",
     )
 )
 

src/discord-cluster-manager/modal_runner.py

@@ -191,5 +191,4 @@ def build_task_config(
             "sources": sources,
             "headers": headers,
             "include_dirs": task.config.include_dirs,
-            **common,
         }