diff --git a/cloud_pipelines_backend/backend_types_sql.py b/cloud_pipelines_backend/backend_types_sql.py index e061739..9381a54 100644 --- a/cloud_pipelines_backend/backend_types_sql.py +++ b/cloud_pipelines_backend/backend_types_sql.py @@ -12,8 +12,15 @@ class ContainerExecutionStatus(str, enum.Enum): + """The lifecycle status of an execution node. + + `UNINITIALIZED` is the parked state: an execution a `QueuedExecutionInterceptor` + took off the launch path. The queued sweep deliberately does not select it, so a + parked execution stays invisible until whoever parked it puts it back to `QUEUED`. + """ + INVALID = "INVALID" # Compatibility with Vertex AI CustomJob - UNINITIALIZED = "UNINITIALIZED" # Remove + UNINITIALIZED = "UNINITIALIZED" # Parked by an interceptor; not swept QUEUED = "QUEUED" # Before WAITING_FOR_UPSTREAM or STARTING # READY_TO_START = "READY_TO_START" # Input artifacts ready, but no job ID WAITING_FOR_UPSTREAM = "WAITING_FOR_UPSTREAM" diff --git a/cloud_pipelines_backend/orchestrator_sql.py b/cloud_pipelines_backend/orchestrator_sql.py index fc470a8..4d05178 100644 --- a/cloud_pipelines_backend/orchestrator_sql.py +++ b/cloud_pipelines_backend/orchestrator_sql.py @@ -39,6 +39,21 @@ class OrchestratorError(RuntimeError): pass +class QueuedExecutionInterceptor(typing.Protocol): + """Given a chance to take a queued execution off the launch path. + + Implemented downstream. Called on the orchestrator's session once the execution is + known to be launchable -- inputs present, not conditionally skipped, no cache hit, not + cancelled. An implementation that returns True owns the execution from that point: it + sets whatever status it wants and commits. The orchestrator makes no assumption about + which status that is. + """ + + def intercept(self, *, session: orm.Session, execution: bts.ExecutionNode) -> bool: + """True if this execution was taken over and must not launch; False to continue.""" + ... + + class OrchestratorService_Sql: def __init__( self, @@ -57,6 +72,7 @@ def __init__( _max_container_execution_refresh_error_retries: int = 3, _max_queue_batch_size: int = 1, _max_queue_batch_duration: datetime.timedelta = datetime.timedelta(), + queued_execution_interceptor: QueuedExecutionInterceptor | None = None, ): self._session_factory = session_factory self._launcher = launcher @@ -75,6 +91,7 @@ def __init__( self._max_queue_batch_size = _max_queue_batch_size self._max_queue_batch_duration = _max_queue_batch_duration + self._queued_execution_interceptor = queued_execution_interceptor def run_loop(self): while True: @@ -124,12 +141,8 @@ def internal_process_queued_executions_queue(self, session: orm.Session): query_start_timestamp = time.monotonic_ns() query = ( sql.select(bts.ExecutionNode).where( - bts.ExecutionNode.container_execution_status.in_( - ( - bts.ContainerExecutionStatus.UNINITIALIZED, - bts.ContainerExecutionStatus.QUEUED, - ) - ) + bts.ExecutionNode.container_execution_status + == bts.ContainerExecutionStatus.QUEUED ) # TODO: Maybe add last_processed_at # .order_by(bts.ExecutionNode.last_processed_at) @@ -610,6 +623,28 @@ def internal_process_one_queued_execution( session.commit() return + # Give the interceptor a chance to take this execution off the launch path. + # If it returns True it has taken ownership: it decided what state the execution is + # in and committed that itself. We stop here and do not launch. + if self._queued_execution_interceptor is not None: + try: + intercepted = self._queued_execution_interceptor.intercept( + session=session, execution=execution + ) + except Exception: + # Fail open. An optional gate must not be able to stop the fleet: the + # failure mode of a broken interceptor is no gating, not no launches. + # No rollback here on purpose: a gate that raised mid-write leaves the + # session dirty, and the launch path below opens a new transaction before + # it writes anything, which discards it. + _logger.exception( + f"Queued-execution interceptor raised on execution {execution.id}; " + f"launching ungated." + ) + intercepted = False + if intercepted: + return + # Creating new container execution container_execution_uuid = _generate_random_id() diff --git a/tests/test_orchestrator_sql.py b/tests/test_orchestrator_sql.py index 554fc46..6a405c6 100644 --- a/tests/test_orchestrator_sql.py +++ b/tests/test_orchestrator_sql.py @@ -1,6 +1,6 @@ -"""Tests for ``orchestrator_sql``. -""" +"""Tests for ``orchestrator_sql``.""" +import datetime from typing import Callable from unittest import mock @@ -76,6 +76,19 @@ def _get_execution_node(session: orm.Session, task_id: str) -> bts.ExecutionNode return node +def _request_termination(*, session: orm.Session, task_id: str) -> None: + """Cancel a node the way a caller does -- by asking, not by writing the status. + + `orchestrator_sql.py:606` keys the cancellation branch off the *request* in + `extra_data`, and reaches it before the interceptor. A test that set + `container_execution_status = CANCELLED` directly would never take that branch, because + the sweep only selects rows that are still QUEUED. + """ + node = _get_execution_node(session, task_id) + node.extra_data = {**(node.extra_data or {}), "desired_state": "TERMINATED"} + session.commit() + + def _make_launched_container_mock() -> mock.MagicMock: launched_container_mock = mock.MagicMock( status=launcher_interfaces.ContainerStatus.PENDING, @@ -84,17 +97,33 @@ def _make_launched_container_mock() -> mock.MagicMock: return mock.MagicMock(return_value=launched_container_mock) -def _process_queued_executions( +def _make_orchestrator( + *, session_factory: Callable[[], orm.Session], launched_container_mock: mock.MagicMock, - max_number_of_executions: int = 20, -) -> None: - orchestrator = orchestrator_sql.OrchestratorService_Sql( + queued_execution_interceptor: ( + orchestrator_sql.QueuedExecutionInterceptor | None + ) = None, +) -> orchestrator_sql.OrchestratorService_Sql: + """An orchestrator wired to mocks, launching through `launched_container_mock`.""" + return orchestrator_sql.OrchestratorService_Sql( session_factory=session_factory, launcher=mock.MagicMock(launch_container_task=launched_container_mock), storage_provider=mock.MagicMock(), data_root_uri="file:///tmp/artifacts", logs_root_uri="file:///tmp/logs", + queued_execution_interceptor=queued_execution_interceptor, + ) + + +def _process_queued_executions( + session_factory: Callable[[], orm.Session], + launched_container_mock: mock.MagicMock, + max_number_of_executions: int = 20, +) -> None: + orchestrator = _make_orchestrator( + session_factory=session_factory, + launched_container_mock=launched_container_mock, ) session = session_factory() # Process the queued queue until it is drained. A bound guards against the @@ -119,7 +148,7 @@ def _output_argument(task_id: str, output_name: str) -> structures.TaskOutputArg class TestQueuedExecutionSystemErrorSkipsDownstream: """Test orphans with SYSTEM_ERROR and WAITING_FOR_UPSTREAM. - + Currently covers the queued-execution failure handler (``OrchestratorService_Sql.internal_process_queued_executions_queue``): when processing a queued execution raises, the execution is marked ``SYSTEM_ERROR`` @@ -328,3 +357,348 @@ def test_failing_downstream_skip_still_marks_system_error(self) -> None: downstream.container_execution_status == bts.ContainerExecutionStatus.WAITING_FOR_UPSTREAM ) + + +# --------------------------------------------------------------------------- # +# The sweep must not select parked (UNINITIALIZED) executions. +# --------------------------------------------------------------------------- # + + +class TestSweepIgnoresUninitialized: + """`UNINITIALIZED` is off the launch path, not merely behind it. + + Downstream (Oasis quota groups) parks an execution by setting it back to + `UNINITIALIZED`. That only hides the node if the sweep stops selecting the + status: were it still selected, the node would be picked again on the next + tick, redo everything above the gate, re-park -- and with no `ORDER BY` the + same low-id node would be chosen every time, spending the whole sweep budget + on one parked execution. + """ + + def test_uninitialized_execution_is_not_selected(self) -> None: + root_task = _make_graph_task_spec( + tasks={ + "parked": structures.TaskSpec( + component_ref=structures.ComponentReference( + spec=_make_container_component() + ), + ), + }, + ) + session_factory = _create_session_factory() + _create_pipeline_run(session_factory, root_task) + launched_container_mock = _make_launched_container_mock() + + # Park it, exactly as the downstream interceptor will. + session = session_factory() + _get_execution_node(session, "parked").container_execution_status = ( + bts.ContainerExecutionStatus.UNINITIALIZED + ) + session.commit() + + orchestrator = _make_orchestrator( + session_factory=session_factory, + launched_container_mock=launched_container_mock, + ) + selected = orchestrator.internal_process_queued_executions_queue( + session=session_factory() + ) + + assert selected is False, "the sweep selected a parked execution" + launched_container_mock.assert_not_called() + assert ( + _get_execution_node(session_factory(), "parked").container_execution_status + == bts.ContainerExecutionStatus.UNINITIALIZED + ), "a parked execution must be left exactly as it was found" + + def test_queued_execution_is_still_selected(self) -> None: + """The other half: narrowing the selection set did not break the sweep.""" + root_task = _make_graph_task_spec( + tasks={ + "runnable": structures.TaskSpec( + component_ref=structures.ComponentReference( + spec=_make_container_component() + ), + ), + }, + ) + session_factory = _create_session_factory() + _create_pipeline_run(session_factory, root_task) + launched_container_mock = _make_launched_container_mock() + + orchestrator = _make_orchestrator( + session_factory=session_factory, + launched_container_mock=launched_container_mock, + ) + selected = orchestrator.internal_process_queued_executions_queue( + session=session_factory() + ) + + assert selected is True + launched_container_mock.assert_called_once() + + +# --------------------------------------------------------------------------- # +# The interceptor seam: a downstream implementation can take an execution over. +# --------------------------------------------------------------------------- # + + +# A row in a table the orchestrator never writes on this path, so finding it on disk means +# an uncommitted write survived when it should not have. Stands in for the claim row the +# real gate inserts before it decides -- `execution.extra_data` cannot serve, because the +# orchestrator rewrites that field itself with the status history. +_SCRIBBLE = "half_written_by_a_broken_gate" + + +class _StubInterceptor: + """Records what it was called with and answers with a fixed decision. + + Stands in for the downstream (Oasis) quota gate, and owns its transaction the way the + protocol requires, so the tests exercise the contract and not just the branch: + + * `intercepted=False` -- writes nothing and lets the launch proceed. + * `intercepted=True` -- takes the execution over: moves it off QUEUED and commits that + itself, so something other than the sweep has to bring it back. + * `intercepted=True, leaves_queued=True` -- declines to launch without reaching a + decision. It writes a status, then rolls itself back, which is what the real gate + does when its compare-and-set budget runs out and the row must stay sweepable. + + `raises=True` makes it blow up after writing and before committing -- the one case + where the orchestrator, not the gate, has to clean the session up. + """ + + def __init__( + self, + *, + intercepted: bool, + leaves_queued: bool = False, + raises: bool = False, + ) -> None: + self._intercepted = intercepted + self._leaves_queued = leaves_queued + self._raises = raises + self.calls: list[str] = [] + + def intercept(self, *, session: orm.Session, execution: bts.ExecutionNode) -> bool: + self.calls.append(execution.id) + if self._raises: + now = datetime.datetime.now(tz=datetime.timezone.utc) + session.add( + bts.Secret( + user_id=_SCRIBBLE, + secret_name=_SCRIBBLE, + secret_value="", + created_at=now, + updated_at=now, + ) + ) + raise RuntimeError("the gate is broken") + if not self._intercepted: + return False + execution.container_execution_status = ( + bts.ContainerExecutionStatus.UNINITIALIZED + ) + if self._leaves_queued: + session.rollback() + else: + session.commit() + return True + + +def _single_task_pipeline() -> structures.TaskSpec: + return _make_graph_task_spec( + tasks={ + "task": structures.TaskSpec( + component_ref=structures.ComponentReference( + spec=_make_container_component() + ), + ), + }, + ) + + +class TestQueuedExecutionInterceptor: + """One test per arm of the seam: intercepted, deferred, launched, no interceptor.""" + + def test_true_takes_the_execution_off_the_launch_path(self) -> None: + session_factory = _create_session_factory() + _create_pipeline_run(session_factory, _single_task_pipeline()) + launched_container_mock = _make_launched_container_mock() + interceptor = _StubInterceptor(intercepted=True) + + orchestrator = _make_orchestrator( + session_factory=session_factory, + launched_container_mock=launched_container_mock, + queued_execution_interceptor=interceptor, + ) + orchestrator.internal_process_queued_executions_queue(session=session_factory()) + + assert len(interceptor.calls) == 1 + launched_container_mock.assert_not_called() + node = _get_execution_node(session_factory(), "task") + assert ( + node.container_execution_status + == bts.ContainerExecutionStatus.UNINITIALIZED + ), "the interceptor committed this itself and the orchestrator left it alone" + assert node.container_execution is None, "no container may have been created" + + def test_true_can_decline_and_still_leave_the_row_sweepable(self) -> None: + """Declining without deciding stops the launch and nothing else.""" + session_factory = _create_session_factory() + _create_pipeline_run(session_factory, _single_task_pipeline()) + launched_container_mock = _make_launched_container_mock() + interceptor = _StubInterceptor(intercepted=True, leaves_queued=True) + + orchestrator = _make_orchestrator( + session_factory=session_factory, + launched_container_mock=launched_container_mock, + queued_execution_interceptor=interceptor, + ) + orchestrator.internal_process_queued_executions_queue(session=session_factory()) + + assert len(interceptor.calls) == 1 + launched_container_mock.assert_not_called() + node = _get_execution_node(session_factory(), "task") + assert ( + node.container_execution_status == bts.ContainerExecutionStatus.QUEUED + ), "the gate rolled its own write back; the row must still be sweepable" + assert node.container_execution is None, "no container may have been created" + + def test_false_launches_exactly_as_before(self) -> None: + session_factory = _create_session_factory() + _create_pipeline_run(session_factory, _single_task_pipeline()) + launched_container_mock = _make_launched_container_mock() + interceptor = _StubInterceptor(intercepted=False) + + orchestrator = _make_orchestrator( + session_factory=session_factory, + launched_container_mock=launched_container_mock, + queued_execution_interceptor=interceptor, + ) + orchestrator.internal_process_queued_executions_queue(session=session_factory()) + + assert len(interceptor.calls) == 1 + launched_container_mock.assert_called_once() + + def test_a_raising_interceptor_fails_open_and_its_write_is_rolled_back( + self, + ) -> None: + """A broken gate must cost gating, not launches -- and leave no trace. + + The stub raises *after* writing and *before* committing, which is the only way + the orchestrator's session can be left dirty: on every ordinary exit the gate has + already committed or rolled back for itself. What discards that write is the + `session.rollback()` the launch path runs to open its own transaction -- delete it + and this test goes red. + """ + session_factory = _create_session_factory() + _create_pipeline_run(session_factory, _single_task_pipeline()) + launched_container_mock = _make_launched_container_mock() + interceptor = _StubInterceptor(intercepted=False, raises=True) + + orchestrator = _make_orchestrator( + session_factory=session_factory, + launched_container_mock=launched_container_mock, + queued_execution_interceptor=interceptor, + ) + orchestrator.internal_process_queued_executions_queue(session=session_factory()) + + assert len(interceptor.calls) == 1 + launched_container_mock.assert_called_once() + leftover = session_factory().get(bts.Secret, (_SCRIBBLE, _SCRIBBLE)) + assert leftover is None, "what the broken gate wrote must not have reached disk" + node = _get_execution_node(session_factory(), "task") + assert ( + node.container_execution is not None + ), "the launch must have been recorded" + + def test_no_interceptor_launches_exactly_as_before(self) -> None: + """The default. Every existing caller passes nothing and must be unaffected.""" + session_factory = _create_session_factory() + _create_pipeline_run(session_factory, _single_task_pipeline()) + launched_container_mock = _make_launched_container_mock() + + _process_queued_executions(session_factory, launched_container_mock) + + launched_container_mock.assert_called_once() + + +class TestInterceptorIsOfferedOnlyLaunchableExecutions: + """The seam's *position*, not its return value. + + The interceptor sits at `orchestrator_sql.py:626`, after every earlier exit from + `internal_process_one_queued_execution`:: + + cache hit? --yes--> reuse the cached execution, return :584 + |no + cancelled? --yes--> CANCELLED, skip downstream, return :624 + |no + intercept() <-- only executions that would otherwise launch reach here + | + launch + + Moving the seam one branch earlier would offer implementations executions that are + never going to run, and whatever work an implementation does per offer would be spent on + them. Both tests use an interceptor that intercepts everything, so a call that should not + happen shows up as a launch that did not. + """ + + def test_a_cancelled_execution_is_never_offered(self) -> None: + session_factory = _create_session_factory() + _create_pipeline_run(session_factory, _single_task_pipeline()) + _request_termination(session=session_factory(), task_id="task") + launched_container_mock = _make_launched_container_mock() + interceptor = _StubInterceptor(intercepted=True) + + orchestrator = _make_orchestrator( + session_factory=session_factory, + launched_container_mock=launched_container_mock, + queued_execution_interceptor=interceptor, + ) + orchestrator.internal_process_queued_executions_queue(session=session_factory()) + + assert interceptor.calls == [], "cancellation is decided before the seam" + launched_container_mock.assert_not_called() + node = _get_execution_node(session_factory(), "task") + assert node.container_execution_status == bts.ContainerExecutionStatus.CANCELLED + + def test_a_cache_hit_is_never_offered(self) -> None: + """Caching is on by default; only `max_cache_staleness == "P0D"` turns it off.""" + session_factory = _create_session_factory() + launched_container_mock = _make_launched_container_mock() + + # The first run has nothing to reuse, so it launches and leaves a PENDING container + # execution behind -- the cache candidate the second run finds. + _create_pipeline_run(session_factory, _single_task_pipeline()) + _make_orchestrator( + session_factory=session_factory, + launched_container_mock=launched_container_mock, + ).internal_process_queued_executions_queue(session=session_factory()) + launched_container_mock.assert_called_once() + + _create_pipeline_run(session_factory, _single_task_pipeline()) + interceptor = _StubInterceptor(intercepted=True) + _make_orchestrator( + session_factory=session_factory, + launched_container_mock=launched_container_mock, + queued_execution_interceptor=interceptor, + ).internal_process_queued_executions_queue(session=session_factory()) + + assert interceptor.calls == [], "the cache hit is decided before the seam" + assert ( + launched_container_mock.call_count == 1 + ), "the second run must have reused, not launched" + # Pin that a cache hit is what happened, so the test cannot pass because the second + # node was never processed at all. + reused = [ + node + for node in session_factory() + .scalars( + sql.select(bts.ExecutionNode).where( + bts.ExecutionNode.task_id_in_parent_execution == "task" + ) + ) + .all() + if (node.extra_data or {}).get("reused_from_execution_node_id") + ] + assert len(reused) == 1, "exactly one of the two nodes must be a cache reuse"