Skip to content

feat(orchestrator): a seam for taking queued executions off the launch path #346

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
yuechao-qin wants to merge 1 commit into
base: master
Choose a base branch
from
Open

feat(orchestrator): a seam for taking queued executions off the launch path #346

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 8 additions & 1 deletion cloud_pipelines_backend/backend_types_sql.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down
34 changes: 28 additions & 6 deletions cloud_pipelines_backend/orchestrator_sql.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
"""
Comment on lines +43 to +50

@Volv-G Volv-G Sep 2, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

(AI-assisted)

The Protocol licenses commit() on the orchestrator's own session, but doesn't state the invariant that makes that safe — or what happens if the implementation raises.

The commit() is only safe because of something invisible from here: at the call site (line 626) the session has no pending orchestrator writes. session.rollback() at line 597 clears everything, and the only work between that rollback and intercept is the two extra_data reads for the cancellation check. Nothing pins that invariant. Add a last_processed_at bump or a status-history append between 597 and 626 later, and an implementation's commit() silently flushes orchestrator state that was never meant to be durable — presenting as half-applied execution rows far from this diff.

Separately, raising out of intercept is a much bigger deal than the docstring implies: it propagates to the handler in internal_process_queued_executions_queue (lines 165-185), which marks the execution SYSTEM_ERROR and _mark_all_downstream_executions_as_skipped. So a transient error in downstream code — which by design talks to its own tables — permanently kills the run.

Both belong in this docstring:

    The session is the orchestrator's own and has no uncommitted orchestrator
    changes at call time, so an implementation may commit freely; it must not
    assume that remains true if it holds the session past `intercept`.

    Raising propagates: the orchestrator marks the execution SYSTEM_ERROR and
    skips everything downstream. Swallow transient errors and return False.

Worth a one-line comment above session.rollback() at line 597 too, noting the clean-session invariant is load-bearing for the seam.


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,
Expand All @@ -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
Expand All @@ -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:
Expand Down Expand Up @@ -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
Comment on lines 142 to +145

@Volv-G Volv-G Sep 2, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

(AI-assisted)

Overloading UNINITIALIZED as the parked state is the part of this design I'd push back on — I think this wants a dedicated status, which is where the design discussion landed too ("look into creating a new state", left unresolved).

Three reasons the reuse bothers me:

  1. The two meanings are opposites. "Never initialized" and "deliberately taken off the launch path by a gate that intends to put it back" are different facts about a node, and after this PR nothing distinguishes them. A reader of a row — or of a dashboard — cannot tell which one they're looking at.

  2. Legacy rows change meaning retroactively. 3a2173b API server - Changed the initial status from UNINITIALIZED to QUEUED means this was the initial status for new nodes. Any surviving row at UNINITIALIZED is drained by the sweep today and becomes permanently invisible after this change. Probably zero rows in practice — SELECT COUNT(*) FROM execution_node WHERE container_execution_status = 'UNINITIALIZED' settles it — but with a distinct state the question wouldn't arise at all.

  3. It's user-visible. A quota-parked node will render as UNINITIALIZED in the UI, which is meaningless to the person whose pipeline it is. That label mapping does not live in this repo, so nothing here can soften it.

The honest counter-argument, which I don't think is fatal: container_execution_status compiles to a MySQL ENUM(...), so a new member is ALTER TABLE execution_node MODIFY COLUMN … on a very hot table — and at least one downstream consumer builds its schema with metadata.create_all() and no migration framework, so it would need a hand-written migration there too. That is real cost. But it is a one-time cost paid at the bottom of a six-PR stack, and it only gets more expensive once parked rows exist in production and the ambiguity is load-bearing.

So: either add the state now, or — if the cost wins — please record the decision in the enum comment (backend_types_sql.py:17-23) as an explicit, priced trade-off rather than a reuse of a spare member, so the next person to touch this knows it was chosen and not inherited.

)
# TODO: Maybe add last_processed_at
# .order_by(bts.ExecutionNode.last_processed_at)
Expand Down Expand Up @@ -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
Comment on lines +626 to +627

@Volv-G Volv-G Sep 2, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

(AI-assisted)

The seam's contract is positional, and no test pins the position — which matters more than usual because a consumer vendoring this repo at a pinned SHA cannot detect it moving.

The docstring's promise is where this runs: "once the execution is known to be launchable — inputs present, not conditionally skipped, no cache hit, not cancelled." TestQueuedExecutionInterceptor covers True / False / no-interceptor, all on a plain launchable single-task pipeline. Every one of those would still pass if a later refactor moved this block above the cache lookup or the cancel check.

The failure that allows is quiet. Move it above the cache lookup and a cache hit now consults the interceptor — an implementation that gates on capacity would charge a slot for work that is about to be satisfied from cache, and could park a node waiting for a slot it never needed. Move it above the cancel check and cancelled executions do the same. Neither errors; both surface much later as "the gate is mysteriously saturated."

Worth pinning here specifically because this is a library seam. A downstream implementation typically vendors this repo at a fixed revision, so its own CI compiles against a frozen copy — a move here cannot fail any test it owns until someone advances the pin, at which point the behaviour change looks like it came from the bump rather than the refactor. These tests are the only place the guarantee can be stated where it runs on every push.

Cheap, given _make_orchestrator and _StubInterceptor already exist — each case is one assertion that intercept was not called:

# interceptor.calls == [] for each of:
#  1. missing input artifact      -> WAITING_FOR_UPSTREAM
#  2. cache hit                   -> reuses cached execution
#  3. desired_state = TERMINATED  -> CANCELLED   (node-level and run-level)
#  4. is_enabled: false           -> SKIPPED

If only one is worth it, take the cache-hit case: it is the one whose failure mode is silent. The cancel case at least ends in a terminal status somebody notices.

# 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
Comment on lines +626 to +631

@Volv-G Volv-G Sep 2, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

(AI-assisted)

An implementation that returns True but leaves the row QUEUED silently re-creates the single-row sweep starvation this PR narrows the query to remove — and nothing errors or logs when it happens.

The contract is prose only: "it sets whatever status it wants and commits." If an implementation returns True but leaves the status as QUEUED, or sets a status and forgets to commit, then:

  • this method returns without committing;
  • process_each_queue_once exits its with self._session_factory() as session: block, which rolls back;
  • the row is still QUEUED;
  • the next sweep runs SELECT ... WHERE status = QUEUED LIMIT 1 with no ORDER BY (lines 141-150 — .order_by(...) is still commented out), so it deterministically re-selects the same low-id row;
  • the orchestrator spends its whole sweep budget re-parsing the task spec, re-deriving the cache key and re-querying candidates for one execution, indefinitely.

That is the exact failure your own description gives as the reason for narrowing the selector, reintroduced through the seam — but silent this time, because a buggy implementation looks identical to a working one from here.

Cheap fix — make the violation loud rather than fatal:

if self._queued_execution_interceptor.intercept(session=session, execution=execution):
    session.refresh(execution)
    if execution.container_execution_status == bts.ContainerExecutionStatus.QUEUED:
        _logger.error(
            "Interceptor claimed execution %s but left it QUEUED; it will be re-swept.",
            execution.id,
        )
    return

Stronger alternative: make the contract enforceable by construction — have intercept return ContainerExecutionStatus | None, where None means "not mine, carry on" and a status means "park it here", and let the orchestrator do the write and the commit. That removes an implementation's ability to leave the row unchanged or the session half-committed, and it also settles the question of what session state the implementation is allowed to assume. If the boolean form was a deliberate choice — to let an implementation pick a status the orchestrator has no opinion about — a sentence saying so would help; it is currently the load-bearing assumption of the whole seam.

):
Comment on lines +629 to +632

@Volv-G Volv-G Sep 2, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

(AI-assisted)

Building on the docstring comment above (#discussion_r3919337398) rather than repeating it — that one asks you to document that raising here is fatal. This is the argument that the call site should instead make it survivable.

To restate only the conclusion: an exception out of intercept propagates to the handler in internal_process_queued_executions_queue (lines 165-185), which commits SYSTEM_ERROR on the node and then skips every downstream execution. Unrecoverable.

The reason that is worth more than a docstring: this is an optional seam whose implementation is, by construction, foreign code talking to storage the orchestrator knows nothing about. As written, installing an admission gate silently couples every pipeline's survival to that gate's availability — a lock-wait timeout or a deploy-time blip during a sweep is enough to kill a run outright. That is a much stronger commitment than an opt-in hook implies, and an implementer reading the signature has no way to infer it.

Fail-open at the call site:

if self._queued_execution_interceptor is not None:
    try:
        if self._queued_execution_interceptor.intercept(session=session, execution=execution):
            return
    except Exception as exc:
        _logger.exception("Queued execution interceptor raised; launching anyway.")
        bugsnag_instrumentation.notify(exception=exc)
        session.rollback()

A broken gate then degrades to "no gating" rather than "no pipelines". The cost is a bounded, self-correcting overshoot while the implementation is down; the benefit is that an optional component cannot take the orchestrator's core job with it.

If you would rather fail closed, one caveat worth stating explicitly: swallowing the exception and falling through without launching leaves the row QUEUED, which re-enters the sweep next tick and re-creates the single-row starvation described in the adjacent comment. Failing closed safely would mean parking the execution — and the orchestrator cannot do that, because the whole premise of this seam is that it has no opinion about which status a claimed execution should hold. That asymmetry is itself an argument for fail-open.

return

# Creating new container execution
container_execution_uuid = _generate_random_id()

Expand Down
197 changes: 192 additions & 5 deletions tests/test_orchestrator_sql.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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``
Expand Down Expand Up @@ -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()
Loading