From 49b27808a2a5113acf11a5167efe4e7f757922f3 Mon Sep 17 00:00:00 2001 From: Yue Chao Qin Date: Wed, 26 Aug 2026 15:50:50 -0700 Subject: [PATCH] feat(orchestrator): a seam for taking queued executions off the launch path Adds `QueuedExecutionInterceptor`, a Protocol the orchestrator consults after the cancellation check and before creating a container. Returning True means the implementation owns the execution: it sets whatever status it wants and commits, and the orchestrator does not launch. `OrchestratorService_Sql` gains one keyword-only `queued_execution_interceptor` parameter defaulting to None, so every existing caller is unaffected. The queued sweep now selects QUEUED only, not UNINITIALIZED too, which makes UNINITIALIZED a parked state that is actually hidden. Without this a parked execution is re-selected on the next tick, redoes the work above the gate and re-parks -- and with no ORDER BY the same low-id row is picked every time, spending the whole sweep budget on one parked execution. Assisted-By: devx/20d7f01c-ddc9-41c5-8b3e-5e921c5b7717 --- cloud_pipelines_backend/backend_types_sql.py | 9 +- cloud_pipelines_backend/orchestrator_sql.py | 34 +++- tests/test_orchestrator_sql.py | 197 ++++++++++++++++++- 3 files changed, 228 insertions(+), 12 deletions(-) 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..ff5519c 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,15 @@ 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: + if self._queued_execution_interceptor.intercept( + session=session, execution=execution + ): + 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..7d4e2a7 100644 --- a/tests/test_orchestrator_sql.py +++ b/tests/test_orchestrator_sql.py @@ -84,17 +84,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 +135,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 +344,174 @@ 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. +# --------------------------------------------------------------------------- # + + +class _StubInterceptor: + """Records what it was called with and answers with a fixed verdict. + + Stands in for the downstream (Oasis) quota gate. When it claims an execution it + behaves as the protocol requires -- sets a status of its own choosing and commits -- + so the test exercises the contract, not just the branch. + """ + + def __init__(self, *, take_over: bool) -> None: + self._take_over = take_over + self.calls: list[str] = [] + + def intercept(self, *, session: orm.Session, execution: bts.ExecutionNode) -> bool: + self.calls.append(execution.id) + if not self._take_over: + return False + execution.container_execution_status = ( + bts.ContainerExecutionStatus.UNINITIALIZED + ) + 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: + """`intercept` returning True must stop the launch, and False must change nothing.""" + + 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(take_over=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 status the interceptor committed must survive" + 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(take_over=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_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()