ci: split backend unit tests into six shards (#8164)

* ci: split backend unit tests into six shards

* test(runtime): synchronize rollback lease renewal

.github/workflows/backend-unit-tests.yml

@@ -57,7 +57,7 @@ jobs:
       # its peers.
       fail-fast: false
       matrix:
-        shard: [1, 2, 3, 4]
+        shard: [1, 2, 3, 4, 5, 6]
 
     services:
       postgres:
@@ -105,7 +105,7 @@ jobs:
         working-directory: backend
         run: uv sync --group dev --extra postgres
 
-      - name: Run unit tests of backend (shard ${{ matrix.shard }} of 4)
+      - name: Run unit tests of backend (shard ${{ matrix.shard }} of ${{ strategy.job-total }})
         working-directory: backend
         env:
           # Expose the job's Postgres service to the cross-pod dedupe integration
@@ -114,4 +114,4 @@ jobs:
           # exercised in CI.
           DEERFLOW_TEST_WORKFLOW_POSTGRES: ${{ env.TEST_POSTGRES_URI }}
           DEDUPE_TEST_POSTGRES_URL: ${{ env.TEST_POSTGRES_URI }}
-        run: make test-shard SPLITS=4 GROUP=${{ matrix.shard }}
+        run: make test-shard SPLITS=${{ strategy.job-total }} GROUP=${{ matrix.shard }}

backend/tests/test_atomic_admission_worker.py

@@ -32,10 +32,12 @@ def anyio_backend():
 )
 @pytest.mark.parametrize("with_event_store", [True, False])
 async def test_real_worker_atomic_replacement_retains_finalization_authority(store, action, with_event_store, host_finalization_cancel, stale_renewal_race):
-    manager = RunManager(store=store, worker_id="owner", run_ownership_config=RunOwnershipConfig(heartbeat_enabled=True, lease_seconds=10, grace_seconds=0))
+    ownership = RunOwnershipConfig(heartbeat_enabled=True, lease_seconds=10, grace_seconds=0)
+    manager = RunManager(store=store, worker_id="owner", run_ownership_config=ownership)
     record = await manager.create_or_reject("thread", user_id="review")
     started, rollback_entered, release_rollback = (asyncio.Event() for _ in range(3))
     error_committed, release_error_ack = asyncio.Event(), asyncio.Event()
+    final_renewal_entered = asyncio.Event()
     hold_rollback = False
 
     class Saver(InMemorySaver):
@@ -79,14 +81,16 @@ async def test_real_worker_atomic_replacement_retains_finalization_authority(sto
     original_status = store.update_status
     original_renew = store.renew_lease
 
-    async def hold_first_active_renewal(run_id, **kwargs):
+    async def observe_renewal(run_id, **kwargs):
         stale_renewal_calls.append(kwargs.get("allow_interrupted", False))
         if stale_renewal_race and len(stale_renewal_calls) == 1:
             stale_renewal_entered.set()
             await release_stale_renewal.wait()
+        if run_id == record.run_id and error_committed.is_set():
+            final_renewal_entered.set()
         return await original_renew(run_id, **kwargs)
 
-    store.renew_lease = hold_first_active_renewal
+    store.renew_lease = observe_renewal
 
     async def hold_final_ack(run_id, status, **kwargs):
         result = await original_status(run_id, status, **kwargs)
@@ -143,7 +147,7 @@ async def test_real_worker_atomic_replacement_retains_finalization_authority(sto
             await asyncio.wait_for(error_committed.wait(), 2)
             assert not record.terminal_status_persisted
             renewal = asyncio.create_task(manager._renew_leases())
-            await asyncio.sleep(0)
+            await asyncio.wait_for(final_renewal_entered.wait(), 2)
             if host_finalization_cancel:
                 record.task.cancel("host cancellation during rollback terminal ACK")
                 await asyncio.sleep(0)
@@ -160,7 +164,8 @@ async def test_real_worker_atomic_replacement_retains_finalization_authority(sto
         else:
             await asyncio.wait_for(record.task, 2)
         if renewal:
-            await asyncio.wait_for(renewal, 1)
+            # Let the manager's lease deadline govern the in-flight store call.
+            await asyncio.wait_for(renewal, ownership.lease_seconds)
         row = await store.get(record.run_id, user_id="review")
         expected = "error" if action == "rollback" else "interrupted"
         assert row["status"] == expected

backend/tests/test_client_live_policy.py

@@ -242,18 +242,20 @@ def test_default_ci_workflow_does_not_opt_in_to_live_tests() -> None:
 
 
 def test_ci_unit_test_workflow_runs_duration_aware_shards() -> None:
-    workflow = (REPO_ROOT / ".github" / "workflows" / "backend-unit-tests.yml").read_text(encoding="utf-8")
+    import yaml
 
-    assert "shard: [1, 2, 3, 4]" in workflow
-    assert "make test-shard SPLITS=4" in workflow
-    assert "GROUP=${{ matrix.shard }}" in workflow
+    workflow = yaml.safe_load((REPO_ROOT / ".github" / "workflows" / "backend-unit-tests.yml").read_text(encoding="utf-8"))
+    job = workflow["jobs"]["backend-unit-tests"]
+    assert job["strategy"]["matrix"] == {"shard": [1, 2, 3, 4, 5, 6]}
+    command = "make test-shard SPLITS=${{ strategy.job-total }} GROUP=${{ matrix.shard }}"
+    assert any(step.get("run") == command for step in job["steps"])
 
     # The sharded command must select the offline suite and balance the shards by
     # real wall-clock cost (duration-aware), not fall back to an even split.
-    command = _dry_run_make_target("test-shard", "SPLITS=4", "GROUP=2")
+    command = _dry_run_make_target("test-shard", "SPLITS=6", "GROUP=6")
 
-    assert "--splits 4" in command
-    assert "--group 2" in command
+    assert "--splits 6" in command
+    assert "--group 6" in command
     assert 'pytest -m "not live"' in command
     assert "--ignore=tests/blocking_io" in command
     assert "--splitting-algorithm least_duration" in command