From b883deaab23bcb189ee3e9e1b037babf090983d2 Mon Sep 17 00:00:00 2001 From: Willem Jiang Date: Thu, 12 Nov 2026 11:08:31 +0800 Subject: [PATCH] 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 | 6 +++--- backend/tests/test_atomic_admission_worker.py | 15 ++++++++++----- backend/tests/test_client_live_policy.py | 16 +++++++++------- 3 files changed, 22 insertions(+), 15 deletions(-) diff --git a/.github/workflows/backend-unit-tests.yml b/.github/workflows/backend-unit-tests.yml index 7c1a564693..3ea2c75fef 100644 --- a/.github/workflows/backend-unit-tests.yml +++ b/.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 }} diff --git a/backend/tests/test_atomic_admission_worker.py b/backend/tests/test_atomic_admission_worker.py index 59207bc9a0..9737da7e89 100644 --- a/backend/tests/test_atomic_admission_worker.py +++ b/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 diff --git a/backend/tests/test_client_live_policy.py b/backend/tests/test_client_live_policy.py index beb0098f34..dc38f8a101 100644 --- a/backend/tests/test_client_live_policy.py +++ b/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