Skip to content

feat(orchestrator): Prefetch a finished container's log and outputs off the loop thread - #427

Open
morgan-wowk wants to merge 1 commit into
masterfrom
orchestrator-completion-prefetch
Open

morgan-wowk wants to merge 1 commit into
masterfrom
orchestrator-completion-prefetch

Conversation

@morgan-wowk

@morgan-wowk morgan-wowk commented Nov 3, 2026 •

Copy link
Copy Markdown
Collaborator

Problem

When the running sweep sees a container that ended, it commits SUCCEEDED or FAILED on the loop thread. Before that, it uploads the log and reads every output (exists, get_info, small values), each under _retry. That is the slow tail of running-queue visits (up to seconds per visit), and every other container execution waits behind it.

Changes

New internal kwarg _completion_prefetch_max_workers (default 0, off). When it is set:

  • The first visit that sees an ended container submits a prefetch to a ThreadPoolExecutor and returns without changing the status. At most one prefetch runs per container execution.
    • The prefetch uploads the log under _retry.
    • For SUCCEEDED, it also checks every output with exists, then runs get_info if they all exist, then reads small values.
  • When the prefetch finishes, its done callback only records the container execution id. Before its next selection, the running queue handler, on the loop thread, sets last_processed_at to 1970-01-01 for the recorded ids that are still PENDING or RUNNING, in one UPDATE. That moves them to the front of the running queue.
  • The next visit takes the prefetch and runs the existing completion code. It reuses the ended container the prefetch started from instead of refreshing it again, since an ended status is final. Each result the prefetch got is used; anything it could not get is fetched inline, exactly as today. If the prefetch itself fails, every call runs inline.
  • The prefetch keeps only an exists() that found the output. A miss is asked again inline when the execution is committed, so an output that becomes visible a moment after the container ends does not fail the execution.
  • Cancelling a container execution, or marking it SYSTEM_ERROR, forgets its prefetch.
  • New shutdown_completion_prefetches(timeout_seconds=...) cancels prefetches that have not started and waits a bounded time for the running ones.

Every database write stays on the loop thread. Workers do storage and log I/O only.

System impact

  • Off by default, so OSS deployments see no change until they set the kwarg.
  • When it is on, the outcome is unchanged. Missing outputs still fail the execution with the same message, a failed log upload is still retried, and _retry behaves the same.
  • A visit that starts a prefetch still advances last_processed_at. A row that something else moved to the front is not re-picked while its prefetch is pending.

Production test

All three orchestrator PRs ran together in production (https://github.com/Shopify/oasis-backend/pull/843). That test used a subclass with launcher and storage wrappers; this PR builds the same behavior into OrchestratorService_Sql. Before/after:

  • Completion lag (container ended → status committed): p50 128s → 1.0s, p90 182s → 5.2s.
  • Start detection (pod running → RUNNING committed): p50 81–250s → 4.8–14s.
  • Container executions never seen as RUNNING before they ended: 34–80% → 0–0.4%.

Related PRs

Testing

  • pytest tests: 602 passed.
  • New tests/test_orchestrator_completion_prefetch.py runs the orchestrator against SQLite with a fake storage provider that records the calling thread. It covers:
    • The first visit starts the prefetch and keeps the status.
    • The terminal status lands on the visit after the prefetch, and the loop thread, not the worker, moves the finished one to the front of the queue.
    • A finished prefetch never moves an ended container execution.
    • An output missing during the prefetch is read again at commit.
    • The commit reuses the ended container without refreshing it again.
    • A visit while the prefetch is pending moves the container execution back in the queue.
    • Shutdown stops new prefetches and reports one still running.
    • The loop thread does no storage or log I/O, and refreshes once per visit.
    • The outcome matches the orchestrator without the prefetch for success, missing outputs, failure and storage errors.
    • A failed log upload is retried inline.
    • Cancelling and SYSTEM_ERROR forget the prefetch.
    • With the prefetch off, the first visit commits the status.

@yuechao-qin yuechao-qin left a comment

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)

Approve with changes. The structure is right: the prefetch is off by default, every prefetched call falls back to an inline call when it is missing (_prefetched_or_call), a failed prefetch degrades to the old path, and the row is stamped last_processed_at at the top of the visit, so an execution waiting for its prefetch does not spin at the front of the queue. One change I would make before merge: a False from exists() is now remembered and reused a visit later, so a successful step whose outputs become listable a moment after the container exits can be marked FAILED for missing outputs (M1). Then two smaller decisions: a prefetch whose execution dies on the SYSTEM_ERROR path is never forgotten (m1), and the done-callback writes to the database from a worker thread with no regard for who holds the orchestrator lock (m2).

# Severity Topic
M1 🔶 major a cached exists() == False can fail an execution whose outputs land a moment later
m1 🔹 minor a prefetch for an execution that is never revisited is never forgotten
m2 🔹 minor the done-callback writes to the database from a worker thread, whoever holds the lock
N1 ▫️ nit this front-of-queue write is not status-filtered; #428's identical write is

readers = {uri: storage_provider.make_uri(uri).get_reader() for uri in output_uris}
with contextlib.suppress(Exception):
for uri, reader in readers.items():
prefetch.output_exists[uri] = _retry(reader.exists)

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)

M1 · 🔶 major — a cached exists() == False can fail an execution whose outputs land a moment later

🔶 Severity: major

Issue. The prefetch records whatever exists() returned, including False
(orchestrator_sql.py:1396), and the commit path reads that recorded answer instead of asking
again (_prefetched_or_call, :1114). A negative read taken on one visit therefore decides the
execution's fate on the next one.

visit N      container is SUCCEEDED → start prefetch
               reader.exists(uri) → False        (outputs not visible yet)
               output_exists = {uri: False}  ← recorded
               `if not all(...)`: return        orchestrator_sql.py:1397
visit N+1    prefetch is done → commit path
               _prefetched_or_call(output_exists, uri, …) → False   ← no re-read
               missing_output_names = [out]  → execution marked FAILED

Before this PR the exists() call happened inline at commit time, so a read that was merely early
was retried by the next sweep rather than remembered. The three PRs in this set make an early read
more likely, not less: #428 brings the first visit forward to the moment Kubernetes reports the
pod changed, which is the earliest point at which the container's final writes may not yet be
listable.

Fix. Only cache the affirmative answer, so a miss falls back to an inline call:

with contextlib.suppress(Exception):
    for uri, reader in readers.items():
        if _retry(reader.exists):
            prefetch.output_exists[uri] = True

_prefetched_or_call already does the right thing with a missing key, and the early return at
:1397 still works — all(prefetch.output_exists.get(uri) …) is None-false on a missing key.

Counter-argument, stated honestly. Google Cloud Storage is strongly consistent for new
objects, and the container is terminal by the time the prefetch runs, so the window needs a
container that exits before its writes are durable — a bug in the task, not in this code. The
production run saw no such failure in 733 completions. The reason to fix it anyway is that the
failure mode is silent and expensive: a successful pipeline step marked FAILED for missing
outputs, with no retry and nothing in the logs pointing at a stale read.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

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

🤖 AI-assisted: Fixed in ed84d1a. The prefetch caches only an exists() that returned True. A miss is asked again inline at commit. New test: an output missing during the prefetch is read again.

self._move_to_front_of_running_queue, container_execution.id
)
)
self._completion_prefetches[container_execution.id] = future

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)

m1 · 🔹 minor — a prefetch for an execution that is never revisited is never forgotten

🔹 Severity: minor

Issue. self._completion_prefetches[container_execution.id] = future
(orchestrator_sql.py:1330) is cleared in exactly two places: the next visit that takes the
finished prefetch (:1333), and the cancellation path (:987). Any execution that reaches a
terminal state by another route keeps its entry — and its _CompletionPrefetch, with up to 255
bytes per preloaded output — for the life of the process.

The route that matters is the refresh-error path: after
_max_container_execution_refresh_error_retries consecutive retriable failures the sweep sets
SYSTEM_ERROR and commits (:321), and the execution is never selected again. The entry stays.
Losing the orchestrator lock has the same effect for everything in flight at that moment.

Fix. Forget the prefetch wherever the execution leaves PENDING|RUNNING — the SYSTEM_ERROR
path is the one that matters today — or bound the dict, since an entry is only ever useful on the
very next visit.

Counter-argument, stated honestly. On the happy path an entry lives for one sweep, the leak
rate is one dict entry per stuck execution, and the orchestrator restarts on every deploy. This
will never be the reason the process runs out of memory. It is raised because an unbounded dict
keyed by execution id is the kind of thing that is cheap to bound now and awkward to diagnose
later.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

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

🤖 AI-assisted: Fixed in ed84d1a. _forget_completion_prefetch now runs on the SYSTEM_ERROR path as well as on cancel. It pops the future and the stored ended container, and cancels the future. Lost-lock exits the process, so nothing outlives it there. New test covers SYSTEM_ERROR.

_logger.exception("Completion prefetch failed; processing without it.")
return _CompletionPrefetch()

def _move_to_front_of_running_queue(

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)

m2 · 🔹 minor — the done-callback writes to the database from a worker thread, whoever holds the lock

🔹 Severity: minor

Issue. _move_to_front_of_running_queue runs on a completion-prefetch worker thread via
add_done_callback (orchestrator_sql.py:1335) and opens its own session to write
last_processed_at (:1341). It is the only database write in the orchestrator that happens
neither on the loop thread nor under the orchestrator lock.

In oasis every sweep is gated on _holds_orchestrator_lock (oasis-backend
orchestrator_main_oasis.py:202),
precisely because two orchestrator processes overlap during a deploy. A process that loses the lock
stops sweeping but keeps its workers, so its callbacks keep rewriting the queue order of rows the
new leader now owns.

process A (was leader)            process B (leader)
  prefetch worker ──── UPDATE last_processed_at = 1970 ──►  DB
  (no lock held)                        │
                                  B's next sweep picks that row first

Fix. Either make the callback a no-op when the process does not hold the lock, or hand the
finished prefetch back to the loop thread — a small queue the sweep drains — and do the update
there, which also puts the write back under the lock that every other write respects.

Counter-argument, stated honestly. The harm is bounded: the write only reorders the running
queue, the row is re-selected and stamped now by whoever processes it, and the status filter on
the selection query means nothing terminal can be dragged back. This ran for 11 hours in
production with the same shape and caused no observed problem. The objection is about the
invariant, not the symptom.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

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

🤖 AI-assisted: Fixed in ed84d1a. The done callback now only adds the id to a lock-protected set. The running handler drains that set on the loop thread before its next selection and does the UPDATE in its own session, so every database write stays on the loop thread.

session.execute(
sql.update(bts.ContainerExecution)
.where(bts.ContainerExecution.id == container_execution_id)
.values(last_processed_at=_FRONT_OF_RUNNING_QUEUE_LAST_PROCESSED_AT)

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)

N1 · ▫️ nit — this front-of-queue write is not status-filtered; #428's identical write is

▫️ Severity: nit — cosmetic, no behaviour change; ignore freely.

Issue. The update at orchestrator_sql.py:1354 matches on id alone, so a prefetch that
finishes after its execution was cancelled — the future is dropped but not cancelled at :987 —
writes 1970-01-01 into a terminal row's last_processed_at. The sibling write added by #428
filters on PENDING|RUNNING for exactly this reason.

Fix. Add the same status.in_((PENDING, RUNNING)) clause, so the two writes that share a
constant also share a predicate.

Counter-argument, stated honestly. Harmless today: every reader of last_processed_at filters
on those statuses too, so a terminal row at the epoch is invisible. It is a consistency nit — but
the two writes being different invites a reader to assume one of them knows something.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

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

🤖 AI-assisted: Fixed in ed84d1a together with m2. The loop-thread UPDATE filters to PENDING/RUNNING, matching #428. New test: a finished prefetch never moves an ended execution.

Also in this head, from a related Binks finding on the production test: the visits after a prefetch starts reuse the ended container that started it instead of refreshing it again.

@morgan-wowk
morgan-wowk force-pushed the orchestrator-completion-prefetch branch from b30ab49 to c86573a Compare November 5, 2026 19:09
…ff the loop thread

Committing a SUCCEEDED or FAILED container execution uploads its log and
reads every output (exists, get_info, small values) on the loop thread,
which delays every other refresh behind that I/O.

With the new internal _completion_prefetch_max_workers set, the first
visit that sees an ended container starts that I/O on a worker thread and
returns, leaving the status unchanged. When the prefetch finishes, it
records the container execution as done, and the loop thread moves every
done one that is still pending or running to the front of the running
queue before its next selection. The next visit commits it from the
prefetched results. Any call the prefetch could not make, and any output
it did not find, is made inline as before, so the outcome is the same with
the prefetch on or off. The visits after a prefetch starts reuse the
ended container that started it rather than refreshing it again, since an
ended status is final. Cancelling or a system error forgets the prefetch.

shutdown_completion_prefetches() cancels prefetches that have not started
and waits a bounded time for the running ones.

The prefetch is off by default (0 workers).
@morgan-wowk
morgan-wowk force-pushed the orchestrator-completion-prefetch branch from c86573a to ed84d1a Compare November 5, 2026 19:23
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants