Add fault tolerance support during migration

.github/workflows/ci.yml

@@ -17,10 +17,10 @@ jobs:
 
     steps:
       - name: Checkout
-        uses: actions/checkout@v4
+        uses: actions/checkout@v6.0.2
 
       - name: Set up Python
-        uses: actions/setup-python@v5
+        uses: actions/setup-python@v6.2.0
         with:
           python-version: "3.14"
           cache: "pip"
@@ -37,7 +37,7 @@ jobs:
         run: pytest tests/unit/coordinator/ -v --cov=coordinator --cov-report=xml:coverage-coordinator.xml
 
       - name: Upload coordinator coverage to Codecov
-        uses: codecov/codecov-action@v5
+        uses: codecov/codecov-action@v5.5.3
         with:
           files: coverage-coordinator.xml
           flags: coordinator
@@ -51,10 +51,10 @@ jobs:
 
     steps:
       - name: Checkout
-        uses: actions/checkout@v4
+        uses: actions/checkout@v6.0.2
 
       - name: Set up Python
-        uses: actions/setup-python@v5
+        uses: actions/setup-python@v6.2.0
         with:
           python-version: "3.14"
           cache: "pip"
@@ -64,7 +64,7 @@ jobs:
           python -m pip install --upgrade pip setuptools Cython
           cd styx-package && python setup.py build_ext --inplace && cd ..
           python -m pip install styx-package/
-          python -m pip install msgspec pytest==9.1.1 pytest-cov pytest-asyncio psutil
+          python -m pip install msgspec pytest==9.1.1 pytest-cov pytest-asyncio psutil uvloop
 
       - name: Build Cython extensions
         run: python worker/setup.py build_ext --inplace
@@ -73,7 +73,7 @@ jobs:
         run: pytest tests/unit/worker/ -v --cov=worker --cov-report=xml:coverage-worker.xml
 
       - name: Upload worker coverage to Codecov
-        uses: codecov/codecov-action@v5
+        uses: codecov/codecov-action@v5.5.3
         with:
           files: coverage-worker.xml
           flags: worker
@@ -87,10 +87,10 @@ jobs:
 
     steps:
       - name: Checkout
-        uses: actions/checkout@v4
+        uses: actions/checkout@v6.0.2
 
       - name: Set up Python
-        uses: actions/setup-python@v5
+        uses: actions/setup-python@v6.2.0
         with:
           python-version: "3.14"
           cache: "pip"
@@ -106,7 +106,7 @@ jobs:
         run: pytest tests/unit/styx_package/ -v --cov=styx-package/styx --cov-report=xml:coverage-styx-package.xml
 
       - name: Upload styx-package coverage to Codecov
-        uses: codecov/codecov-action@v5
+        uses: codecov/codecov-action@v5.5.3
         with:
           files: coverage-styx-package.xml
           flags: styx-package
@@ -139,7 +139,7 @@ jobs:
           python -m pip install --upgrade pip setuptools Cython
           cd styx-package && python setup.py build_ext --inplace && cd ..
           python -m pip install -e styx-package/
-          python -m pip install msgspec pytest==9.1.1 pytest-asyncio psutil
+          python -m pip install msgspec pytest==9.1.1 pytest-asyncio psutil uvloop
           python worker/setup.py build_ext --inplace
 
       - name: Check the hot-path modules are the compiled ones
@@ -170,10 +170,10 @@ jobs:
 
     steps:
       - name: Checkout
-        uses: actions/checkout@v4
+        uses: actions/checkout@v6.0.2
 
       - name: Set up Python
-        uses: actions/setup-python@v5
+        uses: actions/setup-python@v6.2.0
         with:
           python-version: "3.14"
           cache: "pip"
@@ -190,7 +190,7 @@ jobs:
         run: pytest tests/integration/ -v --cov=styx-package/styx --cov=coordinator --cov=worker --cov-report=xml:coverage-integration.xml
 
       - name: Upload integration coverage to Codecov
-        uses: codecov/codecov-action@v5
+        uses: codecov/codecov-action@v5.5.3
         with:
           files: coverage-integration.xml
           flags: integration

.github/workflows/e2e.yml

@@ -10,21 +10,23 @@ concurrency:
   cancel-in-progress: true
 
 jobs:
-  e2e-tests:
+  e2e-base:
+    name: "E2E — base (non-migration)"
+    # Provides the legacy "e2e-tests" status check via the summary job below.
     runs-on: ubuntu-latest
-    timeout-minutes: 45
+    timeout-minutes: 60
 
     steps:
       - name: Checkout
-        uses: actions/checkout@v4
+        uses: actions/checkout@v6.0.2
 
       - name: Docker info
         run: |
           docker version
           docker compose version
 
       - name: Set up Python
-        uses: actions/setup-python@v5
+        uses: actions/setup-python@v6.2.0
         with:
           python-version: "3.14"
           cache: "pip"
@@ -34,37 +36,100 @@ jobs:
           python -m pip install --upgrade pip setuptools Cython
           cd styx-package && python setup.py build_ext --inplace && cd ..
           python -m pip install -r requirements.txt
-          # Install Styx package from repo root
           python -m pip install -e styx-package/
 
       - name: Make scripts executable
         run: |
           chmod +x scripts/start_styx_cluster.sh scripts/stop_styx_cluster.sh || true
 
-      # Needed so Kafka advertises a reachable external listener
       - name: Set DOCKER_HOST_IP
         run: echo "DOCKER_HOST_IP=127.0.0.1" >> $GITHUB_ENV
 
-      - name: Run e2e tests
+      - name: Run non-migration e2e tests
         env:
           PYTHONUNBUFFERED: "1"
         run: |
-          pytest tests/e2e/ -m e2e -q
+          pytest tests/e2e/ -m "e2e" -q --log-level=DEBUG
 
-      - name: Collect worker and coordinator logs (always)
+      - name: Collect logs
         if: always()
         run: |
           mkdir -p artifacts
-          if [ -d logs ]; then
-            cp -r logs artifacts/
-          else
-            echo "No logs folder found to copy."
-          fi
+          if [ -d logs ]; then cp -r logs artifacts/; fi
+
+      - name: Upload artifacts
+        if: always()
+        uses: actions/upload-artifact@v7.0.0
+        with:
+          name: e2e-base-artifacts
+          path: artifacts
+          if-no-files-found: ignore
+
+  e2e-migration:
+    name: "E2E — migration"
+    runs-on: ubuntu-latest
+    timeout-minutes: 90
+
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v6.0.2
+
+      - name: Docker info
+        run: |
+          docker version
+          docker compose version
+
+      - name: Set up Python
+        uses: actions/setup-python@v6.2.0
+        with:
+          python-version: "3.14"
+          cache: "pip"
+
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip setuptools Cython
+          cd styx-package && python setup.py build_ext --inplace && cd ..
+          python -m pip install -r requirements.txt
+          python -m pip install -e styx-package/
 
-      - name: Upload artifacts (always)
+      - name: Make scripts executable
+        run: |
+          chmod +x scripts/start_styx_cluster.sh scripts/stop_styx_cluster.sh || true
+
+      - name: Set DOCKER_HOST_IP
+        run: echo "DOCKER_HOST_IP=127.0.0.1" >> $GITHUB_ENV
+
+      - name: Run migration e2e tests
+        env:
+          PYTHONUNBUFFERED: "1"
+        run: |
+          pytest tests/e2e/ -m "e2e_migration" -q --log-level=DEBUG
+
+      - name: Collect logs
+        if: always()
+        run: |
+          mkdir -p artifacts
+          if [ -d logs ]; then cp -r logs artifacts/; fi
+
+      - name: Upload artifacts
         if: always()
-        uses: actions/upload-artifact@v4
+        uses: actions/upload-artifact@v7.0.0
         with:
-          name: e2e-artifacts
+          name: e2e-migration-artifacts
           path: artifacts
           if-no-files-found: ignore
+
+  # Legacy summary job — provides the "e2e-tests" status check that branch
+  # protection rules may reference from before the split into two jobs.
+  e2e-tests:
+    name: "e2e-tests"
+    if: always()
+    needs: [e2e-base, e2e-migration]
+    runs-on: ubuntu-latest
+    steps:
+      - name: Check results
+        run: |
+          if [[ "${{ needs.e2e-base.result }}" == "failure" || "${{ needs.e2e-migration.result }}" == "failure" ]]; then
+            echo "One or more E2E jobs failed"
+            exit 1
+          fi

.github/workflows/publish-helm-chart.yml

@@ -15,10 +15,10 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - name: Checkout
-        uses: actions/checkout@v4
+        uses: actions/checkout@v6.0.2
 
       - name: Set up Helm
-        uses: azure/setup-helm@v4
+        uses: azure/setup-helm@v5.0.0
 
       - name: Log in to GHCR
         run: echo "${{ secrets.GITHUB_TOKEN }}" | helm registry login ghcr.io --username ${{ github.repository_owner }} --password-stdin

.github/workflows/publish-images.yml

@@ -25,21 +25,21 @@ jobs:
             image_suffix: worker
     steps:
       - name: Checkout
-        uses: actions/checkout@v4
+        uses: actions/checkout@v6.0.2
 
       - name: Set up Build
-        uses: docker/setup-buildx-action@v3
+        uses: docker/setup-buildx-action@v4.0.0
 
       - name: Log in to GHCR
-        uses: docker/login-action@v3
+        uses: docker/login-action@v4.0.0
         with:
           registry: ghcr.io
           username: ${{ github.repository_owner }}
           password: ${{ secrets.GITHUB_TOKEN }}
 
       - name: Extract metadata
         id: meta
-        uses: docker/metadata-action@v5
+        uses: docker/metadata-action@v6.0.0
         with:
           images: ghcr.io/${{ github.repository_owner }}/${{ github.event.repository.name }}-${{ matrix.image_suffix }}
           tags: |
@@ -49,7 +49,7 @@ jobs:
             type=raw,value=latest,enable=${{ github.ref == format('refs/heads/{0}', github.event.repository.default_branch) }}
 
       - name: Build and push
-        uses: docker/build-push-action@v6
+        uses: docker/build-push-action@v7.0.0
         with:
           context: .
           file: ${{ matrix.dockerfile }}

.github/workflows/ruff.yml

@@ -16,10 +16,10 @@ jobs:
 
     steps:
       - name: Checkout
-        uses: actions/checkout@v4
+        uses: actions/checkout@v6.0.2
 
       - name: Set up Python
-        uses: actions/setup-python@v5
+        uses: actions/setup-python@v6.2.0
         with:
           python-version: "3.14"
           cache: "pip"

coordinator/aria_sync_metadata.py

@@ -15,6 +15,8 @@ def __init__(self, n_workers: int) -> None:
         self.global_read_set: dict | None = None
         self.stop_next_epoch: bool = False
         self.take_snapshot: bool = False
+        # Barriers completed so far, so a stalled one can be told from a new one.
+        self.completed_barriers: int = 0
 
     def check_distributed_barrier(self) -> bool:
         return len(self.arrived) == self.n_workers
@@ -122,6 +124,7 @@ def __merge_rw_reservations(
         return output_dict
 
     def cleanup(self, epoch_end: bool = False, take_snapshot: bool = False) -> None:
+        self.completed_barriers += 1
         self.arrived.clear()
         self.logic_aborts_everywhere: set[int] = set()
         self.sent_proceed_msg: bool = False