Publish curated ITWorx-ModelForge source from 990401737f96c8d46a9f7b0d84994516cadcb6a0

backend/tests/test_lab_execution_races_postgres.py

@@ -0,0 +1,198 @@
+"""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