"""PostgreSQL qualification for durable LAB cancellation versus worker completion.""" from __future__ import annotations import hashlib import importlib import threading from concurrent.futures import ThreadPoolExecutor from concurrent.futures import TimeoutError as FutureTimeoutError from datetime import UTC, datetime, timedelta import pytest from sqlalchemy import event, select from sqlalchemy.orm import Session pytest_plugins = ("tests.test_managed_identities_postgres",) from modelforge_api.domain.lab import LabExecutionStartRequest from modelforge_api.domain.serving import AgentServingJobComplete from modelforge_api.persistence.models import GatewayRequest, LabExecutionRun, ServingJob from modelforge_api.services.lab_runs import LabRunService from modelforge_api.services.serving import ServingService from tests.test_dispatch_context import operator_fixture from tests.test_managed_identities_postgres import IdentityDatabase @pytest.mark.parametrize("first_writer", ["cancel", "complete"]) def test_lab_cancel_and_worker_completion_linearize_on_job_then_run_locks( identity_database: IdentityDatabase, first_writer: str ) -> None: """Hold the winning transaction open and prove the other writer observes its result.""" importlib.import_module("alembic.command").upgrade(identity_database.migration(), "head") engine = identity_database.runtime with Session(engine) as session: serving, preflight, selected = operator_fixture(session) request = LabExecutionStartRequest( target=preflight.target, input=preflight.input, limits=preflight.limits, deployment_id=selected.capability_deployment_id, experiment_route_id=selected.id, ) payload_store = serving.payload_store run = LabRunService(serving).start( request, actor_subject=f"postgres-lab-race-{first_writer}", idempotency_key=f"worker-completion-{first_writer}", ) persisted_run = session.get(LabExecutionRun, run.id) assert persisted_run is not None gateway = session.get(GatewayRequest, persisted_run.gateway_request_id) assert gateway is not None token = f"lab-race-token-{first_writer}-" + "t" * 32 now = datetime.now(UTC) job = ServingJob( capability_deployment_id=selected.capability_deployment_id, compute_node_id=selected.deployment.compute_node_id, gateway_request_id=gateway.id, lab_execution_run_id=persisted_run.id, operation="invoke", status="running", priority="lab", idempotency_key=f"worker-job-{first_writer}", payload_reference=persisted_run.payload_reference, attempt_count=1, lease_token_hash=hashlib.sha256(token.encode()).hexdigest(), lease_expires_at=now + timedelta(minutes=2), result_summary={}, started_at=now, ) persisted_run.status = "running" persisted_run.started_at = now persisted_run.heartbeat_at = now gateway.status = "running" session.add(job) session.commit() run_id = run.id job_id = job.id deployment_id = selected.capability_deployment_id node_id = selected.deployment.compute_node_id embedding_space = selected.deployment.embedding_space assert embedding_space is not None dimension = embedding_space.dimension deployment_settings = serving._deployment_settings manifests = serving.manifests actor_subject = f"postgres-lab-race-{first_writer}" first_writer_has_locks = threading.Event() allow_first_writer_to_commit = threading.Event() second_writer_requested_job_lock = threading.Event() def make_serving(session: Session) -> ServingService: return ServingService(session, deployment_settings, manifests, payload_store) def complete(service: ServingService) -> str: deployment = service.repo.deployment(deployment_id) node = service.repo.node(node_id) assert deployment is not None and node is not None vector = [1.0] + [0.0] * (dimension - 1) service.complete_job( job_id, node, AgentServingJobComplete( lease_token=token, worker_instance_id="postgres-race-worker", state="warm", result={"vectors": [vector], "count": 1, "normalized": True}, metrics={"inference_time_ms": 8.0}, health={}, runtime_facts={}, ), ) return "completed" def write(kind: str, *, hold_commit: bool) -> str: with Session(engine) as session: service = make_serving(session) original_commit = session.commit def commit_after_gate() -> None: if hold_commit: first_writer_has_locks.set() if not allow_first_writer_to_commit.wait(timeout=20): raise RuntimeError("timed out holding the first LAB writer locks") original_commit() session.commit = commit_after_gate # type: ignore[method-assign] if kind == "cancel": return LabRunService(service).cancel( run_id, actor_subject=actor_subject ).status return complete(service) @event.listens_for(engine, "before_cursor_execute") def observe_second_job_lock( _connection: object, _cursor: object, statement: str, _parameters: object, _context: object, _executemany: bool, ) -> None: normalized = statement.lower() if ( threading.current_thread().name.startswith("lab-race-second") and "serving_jobs" in normalized and "for update" in normalized ): second_writer_requested_job_lock.set() second_writer = "complete" if first_writer == "cancel" else "cancel" try: with ThreadPoolExecutor(max_workers=1, thread_name_prefix="lab-race-first") as first_pool: first_future = first_pool.submit(write, first_writer, hold_commit=True) if not first_writer_has_locks.wait(timeout=10): first_future.result(timeout=1) pytest.fail("first writer never reached commit") with ThreadPoolExecutor( max_workers=1, thread_name_prefix="lab-race-second" ) as second_pool: second_future = second_pool.submit(write, second_writer, hold_commit=False) if not second_writer_requested_job_lock.wait(timeout=10): second_future.result(timeout=1) pytest.fail("second writer never requested the serving-job lock") with pytest.raises(FutureTimeoutError): second_future.result(timeout=0.25) allow_first_writer_to_commit.set() first_future.result(timeout=30) second_future.result(timeout=30) finally: allow_first_writer_to_commit.set() event.remove(engine, "before_cursor_execute", observe_second_job_lock) expected_status = "cancelled" if first_writer == "cancel" else "completed" with Session(engine) as session: persisted_run = session.get(LabExecutionRun, run_id) persisted_job = session.get(ServingJob, job_id) gateway = session.scalar( select(GatewayRequest).where(GatewayRequest.id == persisted_run.gateway_request_id) ) if persisted_run else None assert persisted_run is not None and persisted_job is not None and gateway is not None assert persisted_run.status == expected_status assert persisted_job.status == expected_status assert gateway.status == expected_status assert persisted_job.lease_token_hash is None if first_writer == "cancel": assert persisted_run.cancellation_requested_at is not None assert persisted_run.cancellation_acknowledged_at is not None assert persisted_run.output_metadata is None assert persisted_job.result_reference is None assert persisted_job.result_summary["cancelled"] is True else: assert persisted_run.cancellation_requested_at is None assert persisted_run.cancellation_acknowledged_at is None assert persisted_run.output_metadata is not None assert persisted_job.result_reference is None