[iris] Redesign integration test suite (#4009)

Redesigns the Iris integration test suite (`tests/integration/iris/`):

- **Full marin pipeline on Iris**: `test_marin_on_iris.py` now runs the
real `create_steps()` pipeline from `tests/integration_test.py` via
`FrayIrisClient`, replacing the fake JSON-shuffling placeholder.
- **No local cluster fallback**: `conftest.py` stripped to ~43 lines.
`--controller-url` is always required; no hidden 10-worker local
cluster.
- **Flat job functions**: `IntegrationJobs` class replaced with
top-level functions in `jobs.py`.
- **Auth/checkpoint/GPU tests moved to `lib/iris/tests/`**: These create
their own `LocalCluster` and don't belong in the generic integration
suite.
- **CLI dispatch tests**: `test_cli_dispatch.py` exercises `iris job
run` via subprocess.
- **`is_remote` removed**: No local/remote specialization. All tests
assume a real cluster.
- **GitHub workflow**: `.github/workflows/iris-integration.yaml` starts
a local Iris cluster and runs the full suite including the marin
pipeline test.

Author: claude[bot]

.github/workflows/iris-integration.yaml

@@ -0,0 +1,86 @@
+name: Iris - Integration Tests
+
+on:
+  push:
+    branches: [main]
+  pull_request:
+  workflow_dispatch:
+
+jobs:
+  iris-itest:
+    if: github.event_name == 'push' || github.event.pull_request.head.repo.full_name == github.repository
+    runs-on: ubuntu-latest
+    timeout-minutes: 45
+    concurrency:
+      group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
+      cancel-in-progress: true
+
+    strategy:
+      matrix:
+        python-version: ["3.12"]
+
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
+
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python-version }}
+
+      - name: Set up Node.js
+        uses: actions/setup-node@v4
+        with:
+          node-version: "22"
+
+      - name: Install uv
+        uses: astral-sh/setup-uv@v7
+        with:
+          enable-cache: true
+
+      - name: Install dependencies
+        run: uv sync --all-packages --extra=cpu --extra=dedup --no-default-groups --group dev
+
+      - 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"
+
+          # 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 integration tests
+        run: |
+          uv run pytest tests/integration/iris/ \
+            --controller-url "$IRIS_CONTROLLER_URL" \
+            -v --tb=short --timeout=600 \
+            -o "addopts=" \
+            -x
+        env:
+          WANDB_MODE: disabled
+          WANDB_API_KEY: ""
+          JAX_TRACEBACK_FILTERING: off
+
+      - 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

lib/iris/examples/test.yaml

@@ -33,7 +33,7 @@ controller:
     port: 10000
 
 scale_groups:
-  # Active scale group with capacity
+  # Active scale group with capacity (non-preemptible so executor-style jobs land here)
   tpu_v5e_4:
     num_vms: 1
     resources:
@@ -43,7 +43,7 @@ scale_groups:
       device_type: tpu
       device_variant: v5litepod-4
       device_count: 4
-      preemptible: true
+      preemptible: false
     min_slices: 2
     max_slices: 10
     slice_template:

lib/iris/tests/conftest.py

@@ -7,15 +7,38 @@
 import os
 import subprocess
 import sys
+from pathlib import Path
 import threading
 import time
 import traceback
 import warnings
 
 import pytest
+from iris.cluster.config import load_config, make_local_config
+from iris.rpc import config_pb2
 from iris.test_util import SentinelFile
 from iris.time_utils import Deadline, Duration
 
+IRIS_ROOT = Path(__file__).resolve().parents[1]
+DEFAULT_CONFIG = IRIS_ROOT / "examples" / "test.yaml"
+
+
+def _make_controller_only_config() -> config_pb2.IrisClusterConfig:
+    """Build a local config with no auto-scaled workers."""
+    config = load_config(DEFAULT_CONFIG)
+    config.scale_groups.clear()
+    sg = config.scale_groups["placeholder"]
+    sg.name = "placeholder"
+    sg.num_vms = 1
+    sg.min_slices = 0
+    sg.max_slices = 0
+    sg.resources.cpu_millicores = 1000
+    sg.resources.memory_bytes = 1 * 1024**3
+    sg.resources.disk_bytes = 10 * 1024**3
+    sg.resources.device_type = config_pb2.ACCELERATOR_TYPE_CPU
+    sg.slice_template.local.SetInParent()
+    return make_local_config(config)
+
 
 def _docker_image_exists(tag: str) -> bool:
     try:

lib/iris/tests/test_auth.py

@@ -0,0 +1,105 @@
+# Copyright The Marin Authors
+# SPDX-License-Identifier: Apache-2.0
+"""Auth tests for Iris controller with static token authentication."""
+
+import pytest
+from iris.cluster.providers.local.cluster import LocalCluster
+from iris.cluster.types import Entrypoint, ResourceSpec
+from iris.rpc import cluster_pb2
+from iris.rpc.cluster_connect import ControllerServiceClientSync
+
+from .conftest import _make_controller_only_config
+
+_AUTH_TOKEN = "e2e-test-token"
+_AUTH_USER = "test-user"
+
+
+def _login_for_jwt(url: str, identity_token: str) -> str:
+    """Exchange a raw identity token for a JWT via the Login RPC."""
+    client = ControllerServiceClientSync(address=url, timeout_ms=10000)
+    try:
+        resp = client.login(cluster_pb2.LoginRequest(identity_token=identity_token))
+        return resp.token
+    finally:
+        client.close()
+
+
+def _quick():
+    return 1
+
+
+def test_static_auth_rpc_access():
+    """Static auth rejects unauthenticated and wrong-token RPCs, accepts valid JWT."""
+    from connectrpc.errors import ConnectError
+    from iris.rpc.auth import AuthTokenInjector, StaticTokenProvider
+
+    config = _make_controller_only_config()
+    config.auth.static.tokens[_AUTH_TOKEN] = _AUTH_USER
+    controller = LocalCluster(config)
+    url = controller.start()
+
+    try:
+        list_req = cluster_pb2.Controller.ListWorkersRequest()
+
+        unauth_client = ControllerServiceClientSync(address=url, timeout_ms=5000)
+        with pytest.raises(ConnectError, match=r"(?i)(authorization|authenticat)"):
+            unauth_client.list_workers(list_req)
+        unauth_client.close()
+
+        wrong_injector = AuthTokenInjector(StaticTokenProvider("wrong-token"))
+        wrong_client = ControllerServiceClientSync(address=url, timeout_ms=5000, interceptors=[wrong_injector])
+        with pytest.raises(ConnectError, match=r"(?i)authenticat"):
+            wrong_client.list_workers(list_req)
+        wrong_client.close()
+
+        jwt_token = _login_for_jwt(url, _AUTH_TOKEN)
+        valid_injector = AuthTokenInjector(StaticTokenProvider(jwt_token))
+        valid_client = ControllerServiceClientSync(address=url, timeout_ms=5000, interceptors=[valid_injector])
+        response = valid_client.list_workers(list_req)
+        assert response is not None
+        valid_client.close()
+    finally:
+        controller.close()
+
+
+def test_static_auth_job_ownership():
+    """Job ownership: user A cannot terminate user B's job."""
+    from connectrpc.errors import ConnectError
+    from iris.rpc.auth import AuthTokenInjector, StaticTokenProvider
+
+    _TOKEN_A = "token-user-a"
+    _TOKEN_B = "token-user-b"
+
+    config = _make_controller_only_config()
+    config.auth.static.tokens[_TOKEN_A] = "user-a"
+    config.auth.static.tokens[_TOKEN_B] = "user-b"
+    controller = LocalCluster(config)
+    url = controller.start()
+
+    try:
+        jwt_a = _login_for_jwt(url, _TOKEN_A)
+        jwt_b = _login_for_jwt(url, _TOKEN_B)
+
+        injector_a = AuthTokenInjector(StaticTokenProvider(jwt_a))
+        client_a = ControllerServiceClientSync(address=url, timeout_ms=10000, interceptors=[injector_a])
+
+        entrypoint = Entrypoint.from_callable(_quick)
+        launch_req = cluster_pb2.Controller.LaunchJobRequest(
+            name="/user-a/auth-owned-job",
+            entrypoint=entrypoint.to_proto(),
+            resources=ResourceSpec(cpu=1, memory="1g").to_proto(),
+        )
+        resp = client_a.launch_job(launch_req)
+        job_id = resp.job_id
+
+        injector_b = AuthTokenInjector(StaticTokenProvider(jwt_b))
+        client_b = ControllerServiceClientSync(address=url, timeout_ms=10000, interceptors=[injector_b])
+        with pytest.raises(ConnectError, match="cannot access resources owned by"):
+            client_b.terminate_job(cluster_pb2.Controller.TerminateJobRequest(job_id=job_id))
+
+        client_a.terminate_job(cluster_pb2.Controller.TerminateJobRequest(job_id=job_id))
+
+        client_a.close()
+        client_b.close()
+    finally:
+        controller.close()

lib/iris/tests/test_checkpoint_restore.py

@@ -0,0 +1,103 @@
+# Copyright The Marin Authors
+# SPDX-License-Identifier: Apache-2.0
+"""Checkpoint/restore test for Iris controller."""
+
+from pathlib import Path
+
+import time
+
+from iris.client.client import IrisClient, Job
+from iris.cluster.config import load_config, make_local_config
+from iris.cluster.providers.local.cluster import LocalCluster
+from iris.cluster.types import Entrypoint, EnvironmentSpec, ResourceSpec, is_job_finished
+from iris.rpc import cluster_pb2
+from iris.rpc.cluster_connect import ControllerServiceClientSync
+
+IRIS_ROOT = Path(__file__).resolve().parents[1]
+DEFAULT_CONFIG = IRIS_ROOT / "examples" / "test.yaml"
+
+
+def _quick():
+    return 1
+
+
+class _IrisTestHelper:
+    """Minimal helper to submit and wait for jobs (standalone, no integration fixtures)."""
+
+    def __init__(self, url: str, client: IrisClient, controller_client: ControllerServiceClientSync):
+        self.url = url
+        self.client = client
+        self.controller_client = controller_client
+
+    def wait_for_workers(self, count: int, timeout: float = 30):
+        deadline = time.monotonic() + timeout
+        while time.monotonic() < deadline:
+            resp = self.controller_client.list_workers(cluster_pb2.Controller.ListWorkersRequest())
+            healthy = [w for w in resp.workers if w.healthy]
+            if len(healthy) >= count:
+                return
+            time.sleep(0.5)
+        raise TimeoutError(f"Expected {count} healthy workers, timed out")
+
+    def submit(self, fn, name: str) -> Job:
+        return self.client.submit(
+            entrypoint=Entrypoint.from_callable(fn),
+            name=name,
+            resources=ResourceSpec(cpu=1, memory="1g"),
+            environment=EnvironmentSpec(),
+        )
+
+    def wait(self, job: Job, timeout: float = 30) -> cluster_pb2.JobStatus:
+        deadline = time.monotonic() + timeout
+        while time.monotonic() < deadline:
+            resp = self.controller_client.get_job_status(
+                cluster_pb2.Controller.GetJobStatusRequest(job_id=job.job_id.to_wire())
+            )
+            if is_job_finished(resp.job.state):
+                return resp.job
+            time.sleep(0.5)
+        raise TimeoutError(f"Job {job.job_id} did not finish within {timeout}s")
+
+
+def test_checkpoint_restore():
+    """Controller restart resumes from checkpoint: completed jobs visible, cluster functional."""
+    config = load_config(DEFAULT_CONFIG)
+    config = make_local_config(config)
+
+    cluster = LocalCluster(config)
+    url = cluster.start()
+    try:
+        client = IrisClient.remote(url, workspace=IRIS_ROOT)
+        controller_client = ControllerServiceClientSync(address=url, timeout_ms=30000)
+        tc = _IrisTestHelper(url=url, client=client, controller_client=controller_client)
+        tc.wait_for_workers(1, timeout=30)
+
+        job = tc.submit(_quick, "pre-restart")
+        tc.wait(job, timeout=30)
+        saved_job_id = job.job_id.to_wire()
+
+        ckpt = controller_client.begin_checkpoint(cluster_pb2.Controller.BeginCheckpointRequest())
+        assert ckpt.checkpoint_path, "begin_checkpoint returned empty path"
+        assert ckpt.job_count >= 1
+        controller_client.close()
+
+        url = cluster.restart()
+
+        controller_client = ControllerServiceClientSync(address=url, timeout_ms=30000)
+        tc = _IrisTestHelper(
+            url=url, client=IrisClient.remote(url, workspace=IRIS_ROOT), controller_client=controller_client
+        )
+
+        resp = controller_client.get_job_status(cluster_pb2.Controller.GetJobStatusRequest(job_id=saved_job_id))
+        assert (
+            resp.job.state == cluster_pb2.JOB_STATE_SUCCEEDED
+        ), f"Pre-restart job has state {resp.job.state} after restore"
+
+        tc.wait_for_workers(1, timeout=30)
+        post_job = tc.submit(_quick, "post-restart")
+        status = tc.wait(post_job, timeout=30)
+        assert status.state == cluster_pb2.JOB_STATE_SUCCEEDED
+
+        controller_client.close()
+    finally:
+        cluster.close()