Remove Ray from integration tests; consolidate on Iris (#4601)

Replace the Ray-backed marin-itest with the existing Iris-backed
pipeline test, fold the iris-integration workflow into marin-itest
(which now starts its own local Iris controller), and drop the Ray
fixtures and tests from lib/fray/tests. Part of the Ray deprecation
effort.

Author: rjpower

.github/workflows/fray-unit-tests.yaml

@@ -49,51 +49,5 @@ jobs:
 
       - name: Test fray
         run: |
-          # NOTE: do not use frozen here because it makes ray sad when it sets up venvs with uv
           cd lib/fray && uv run --group=fray-test pytest --durations=5 --tb=short -m 'not slow and not tpu_ci' -v -s tests/
 
-  fray-tpu-tests:
-    needs: changes
-    runs-on: [tpu-ci]
-    timeout-minutes: 10
-    if: needs.changes.outputs.should_run == 'true' && (github.event_name == 'push' || github.event.pull_request.head.repo.full_name == github.repository)
-    steps:
-      - name: Checkout code
-        uses: actions/checkout@v3
-        with:
-          ref: ${{ (github.event_name == 'pull_request_review' && format('refs/pull/{0}/merge', github.event.pull_request.number)) || '' }}
-
-      - name: Run TPU tests in Docker
-        env:
-          HF_TOKEN: ${{ secrets.HF_TOKEN }}
-          WANDB_API_KEY: ${{ secrets.WANDB_API_KEY }}
-        run: |
-          bash infra/tpu-ci/clean-tpu.sh
-
-          DOCKER_IMAGE="ghcr.io/marin-community/marin/tpu-ci:latest"
-          echo "Using Docker image: $DOCKER_IMAGE"
-
-          # Create UV cache directory
-          mkdir -p /tmp/uv-cache
-          chmod 777 /tmp/uv-cache
-
-          docker run --rm \
-            --device /dev/vfio:/dev/vfio \
-            --shm-size=100g \
-            --stop-timeout=5 \
-            --cap-add=SYS_RESOURCE \
-            --ulimit memlock=68719476736:68719476736 \
-            -e HF_TOKEN \
-            -e JAX_COORDINATOR_ADDRESS=127.0.0.1 \
-            -e JAX_PLATFORMS=tpu,cpu \
-            -e PJRT_DEVICE=TPU \
-            -e TPU_MIN_LOG_LEVEL=3 \
-            -e TPU_STDERR_LOG_LEVEL=3 \
-            -e UV_CACHE_DIR=/tmp/uv-cache \
-            -e WANDB_API_KEY \
-            -e WANDB_MODE=offline \
-            -v ${{ github.workspace }}:/workspace-src:ro \
-            -v /tmp/uv-cache:/tmp/uv-cache:rw \
-            -w /workspace \
-            $DOCKER_IMAGE \
-            bash -c "cp -a /workspace-src/. /workspace/ && cd /workspace/lib/fray && timeout --kill-after=5 --signal=TERM 590 uv run --group=fray-tpu-test pytest tests/ -v --tb=short -s --log-cli-level=INFO -m tpu_ci"

.github/workflows/iris-coreweave-ci.yaml

@@ -174,7 +174,7 @@ jobs:
           FSSPEC_S3: '{"endpoint_url": "https://74981a43be0de7712369306c7b19133d.r2.cloudflarestorage.com"}'
         run: |
           IRIS_CONTROLLER_URL="http://localhost:${LOCAL_PORT}"
-          timeout 600 uv run tests/integration/iris/run_iris_full_integration.py \
+          timeout 600 uv run tests/integration_test.py \
             --controller-url "$IRIS_CONTROLLER_URL"
 
       - name: Stop port-forward

.github/workflows/iris-integration.yaml

@@ -12,18 +12,16 @@ jobs:
   marin-itest:
     if: github.event_name == 'push' || github.event.pull_request.head.repo.full_name == github.repository
     runs-on: ubuntu-latest
-    # uv has to resolve + download large binary wheels (jax/torch). 10 minutes wasn't
-    # enough for cache misses on GitHub's runners, so give ourselves more runway.
-    timeout-minutes: 25
+    timeout-minutes: 45
     concurrency:
       group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
       cancel-in-progress: true
 
     steps:
       - name: Checkout code
-        uses: actions/checkout@v3
+        uses: actions/checkout@v4
       - name: Set up Python 3.12
-        uses: actions/setup-python@v4
+        uses: actions/setup-python@v5
         with:
           python-version: "3.12"
 
@@ -41,21 +39,61 @@ jobs:
         # `--no-default-groups` keeps uv from also installing every workspace package's
         # dev/docs/test groups (e.g. `levanter[docs,test,dev]`), which saves multiple
         # minutes and avoids re-downloading CUDA wheels unnecessarily on cold caches.
-        run: uv sync --all-packages --extra=cpu --extra=dedup --no-default-groups
+        run: uv sync --all-packages --extra=cpu --extra=dedup --no-default-groups --group dev
 
       - name: Check df -h
         run: df -h
 
-      - name: Give Ray tmp space
-        run: sudo mkdir -p /mnt/ray && sudo chmod 777 /mnt/ray
+      - name: Start local Iris cluster
+        run: |
+          uv run iris --config lib/iris/examples/test.yaml \
+            cluster start --local > /tmp/iris-cluster.log 2>&1 &
+          CLUSTER_PID=$!
+          echo "CLUSTER_PID=$CLUSTER_PID" >> "$GITHUB_ENV"
 
-      - name: Run the quickstart script
-        shell: bash -l {0}
-        # N.B. You _must not_ use `uv run` here, as that triggers weird behavior from Ray
-        # https://github.com/ray-project/ray/issues/54344
-        run: .venv/bin/python tests/integration_test.py
+          # Wait for controller to print its URL
+          for i in $(seq 1 120); do
+            if grep -q "Controller started at" /tmp/iris-cluster.log 2>/dev/null; then
+              URL=$(grep "Controller started at" /tmp/iris-cluster.log | head -1 | sed -n 's/.*Controller started at //p')
+              echo "IRIS_CONTROLLER_URL=$URL" >> "$GITHUB_ENV"
+              echo "Cluster ready at $URL"
+              break
+            fi
+            sleep 1
+          done
+
+          if [ -z "${URL:-}" ]; then
+            echo "Cluster failed to start within timeout"
+            cat /tmp/iris-cluster.log
+            exit 1
+          fi
+
+      - name: Run iris integration tests
+        run: |
+          uv run pytest tests/integration/iris/ \
+            --controller-url "$IRIS_CONTROLLER_URL" \
+            -v -s --log-cli-level=INFO --tb=short --timeout=600 \
+            -o "addopts=" \
+            -x
+        env:
+          WANDB_MODE: disabled
+          WANDB_API_KEY: ""
+          JAX_TRACEBACK_FILTERING: off
+
+      - name: Run full marin integration pipeline
+        run: |
+          timeout 600 uv run tests/integration_test.py \
+            --controller-url "$IRIS_CONTROLLER_URL"
         env:
           HF_TOKEN: ${{ secrets.HF_TOKEN }}
+          WANDB_MODE: disabled
+          WANDB_API_KEY: ""
           JAX_TRACEBACK_FILTERING: off
-          WANDB_MODE: offline
-          RAY_TMPDIR: /mnt/ray/
+
+      - name: Stop cluster
+        if: always()
+        run: kill $CLUSTER_PID 2>/dev/null || true
+
+      - name: Show cluster logs on failure
+        if: failure()
+        run: cat /tmp/iris-cluster.log || true

.github/workflows/marin-itest.yaml

@@ -4,38 +4,15 @@
 """Pytest fixtures for fray tests."""
 
 
-import logging
-
 import pytest
-import ray
 from fray.v1.cluster.local_cluster import LocalCluster, LocalClusterConfig
 
 
-@pytest.fixture(scope="module")
-def ray_cluster():
-    from fray.v1.cluster.ray import RayCluster
-
-    if not ray.is_initialized():
-        logging.info("Initializing Ray cluster")
-        ray.init(
-            address="local",
-            num_cpus=8,
-            ignore_reinit_error=True,
-            logging_level="info",
-            log_to_driver=True,
-            resources={"head_node": 1},
-        )
-    yield RayCluster()
-
-
 @pytest.fixture(scope="module")
 def local_cluster():
     yield LocalCluster(LocalClusterConfig(use_isolated_env=False))
 
 
-@pytest.fixture(scope="module", params=["local", "ray"])
-def cluster(request, local_cluster, ray_cluster):
-    if request.param == "local":
-        return local_cluster
-    elif request.param == "ray":
-        return ray_cluster
+@pytest.fixture(scope="module")
+def cluster(local_cluster):
+    return local_cluster

lib/fray/tests/conftest.py

@@ -4,18 +4,15 @@
 """Tests for execution contexts."""
 
 import pytest
-from fray.v1.job import RayContext, SimpleActor, SyncContext, ThreadContext, create_job_ctx
+from fray.v1.job import SimpleActor, SyncContext, ThreadContext, create_job_ctx
 from fray.v1.job.context import _apply_default_jax_platforms
 
 
-@pytest.fixture(params=["sync", "threadpool", "ray"])
-def job_context(request, ray_cluster):
+@pytest.fixture(params=["sync", "threadpool"])
+def job_context(request):
     if request.param == "sync":
         return SyncContext()
-    elif request.param == "threadpool":
-        return ThreadContext(max_workers=2)
-
-    return RayContext()
+    return ThreadContext(max_workers=2)
 
 
 def test_context_put_get(job_context):