Auto-end PipelineRunner.run() when all root workers finish

run() now defaults to auto_end=True: the runner ends once every root
worker has finished, so single-pipeline bots end naturally when their
pipeline does and tests no longer need an explicit runner.end() call.
Multi-worker bots whose helpers run forever still trigger shutdown via
end() / cancel() from an event handler (typically on transport
disconnect). Hosts that add and remove workers across many sessions can
pass auto_end=False to keep the runner up.
This commit is contained in:
Aleix Conchillo Flaqué
2026-05-21 23:08:47 -07:00
parent afa880f523
commit e47f7d0e63
3 changed files with 52 additions and 27 deletions

View File

@@ -0,0 +1 @@
- `PipelineRunner.run()` now ends automatically once every root worker has finished, so single-pipeline bots no longer need an explicit `runner.end()` / `runner.cancel()` call. Multi-worker bots whose helpers run forever (waiting for bus messages) still trigger shutdown by calling `end()` / `cancel()` from an event handler (typically on transport disconnect). Pass `auto_end=False` to `run()` for long-lived hosts (e.g. a FastAPI server) that add and remove workers across many sessions.

View File

@@ -1 +1 @@
- Passing a worker to `PipelineRunner.run()` is deprecated. Register the worker with `PipelineRunner.add_workers()` before calling `run()` instead; `run()` now blocks until `end()` / `cancel()` is called rather than until the passed worker finishes. The `worker` argument still works but emits a `DeprecationWarning` and will be removed in a future release. - Passing a worker to `PipelineRunner.run()` is deprecated. Register the worker with `PipelineRunner.add_workers()` before calling `run()` instead. The `worker` argument still works but emits a `DeprecationWarning` and will be removed in a future release.

View File

@@ -29,12 +29,14 @@ For multi-worker setups, register every worker the same way:
await runner.add_workers(CodeWorker("code_worker", ...), worker) await runner.add_workers(CodeWorker("code_worker", ...), worker)
await runner.run() await runner.run()
``run()`` blocks until :meth:`PipelineRunner.end` / By default, ``run()`` ends once every root worker has finished — so a
:meth:`PipelineRunner.cancel` is called (or an incoming ``BusEndMessage`` / single-pipeline bot naturally ends when its pipeline does. Multi-worker
``BusCancelMessage`` triggers the same path). Added workers finishing on bots whose helpers run forever (e.g. waiting for bus messages) end by
their own does **not** unblock it — use ``end()`` / ``cancel()`` from an calling :meth:`PipelineRunner.end` / :meth:`PipelineRunner.cancel` from
event handler (e.g. when the transport disconnects) to shut the runner an event handler (typically on transport disconnect). For long-lived
down. hosts that add and remove workers over many sessions (e.g. a FastAPI
server), pass ``auto_end=False`` to ``run()`` so the runner does not
exit when no workers are left.
""" """
import asyncio import asyncio
@@ -87,11 +89,12 @@ class PipelineRunner(BaseObject, BusSubscriber):
- :meth:`add_workers(*workers)` — register one or more workers on the - :meth:`add_workers(*workers)` — register one or more workers on the
runner's bus and start them in the background. Workers run runner's bus and start them in the background. Workers run
concurrently and are cancelled when :meth:`end` / :meth:`cancel` concurrently and remaining workers are cancelled when the runner
is called. ends.
- :meth:`run` — block until :meth:`end` / :meth:`cancel` is called - :meth:`run` — block until the runner ends. By default
(or until an incoming ``BusEndMessage`` / ``BusCancelMessage`` (``auto_end=True``) the runner ends once every root worker has
triggers the same path). finished; pass ``auto_end=False`` to keep the runner up until
:meth:`end` / :meth:`cancel` is called.
Event handlers available: Event handlers available:
@@ -132,6 +135,7 @@ class PipelineRunner(BaseObject, BusSubscriber):
self._entries: dict[str, _WorkerEntry] = {} self._entries: dict[str, _WorkerEntry] = {}
self._known_runners: set[str] = set() self._known_runners: set[str] = set()
self._running: bool = False self._running: bool = False
self._auto_end: bool = True
self._shutdown_event = asyncio.Event() self._shutdown_event = asyncio.Event()
self._sig_task: asyncio.Task | None = None self._sig_task: asyncio.Task | None = None
@@ -184,15 +188,23 @@ class PipelineRunner(BaseObject, BusSubscriber):
if self._running: if self._running:
await self._start_worker(entry) await self._start_worker(entry)
async def run(self, worker: PipelineWorker | None = None) -> None: async def run(
self,
worker: PipelineWorker | None = None,
*,
auto_end: bool = True,
) -> None:
"""Run all added workers until the runner is stopped. """Run all added workers until the runner is stopped.
Blocks until :meth:`end` or :meth:`cancel` is called (or until an By default (``auto_end=True``), the runner ends once every root
incoming ``BusEndMessage`` / ``BusCancelMessage`` triggers the worker has finished — so a single-pipeline bot naturally ends
same path). Added workers finishing on their own does **not** when its pipeline does. Multi-worker bots whose helpers run
unblock the runner — call ``end()`` / ``cancel()`` from an forever (e.g. waiting for bus messages) end by calling
event handler (e.g. when the transport disconnects) to shut the :meth:`end` / :meth:`cancel` from an event handler (typically on
runner down. transport disconnect). For long-lived hosts that add and remove
workers over many sessions (e.g. a FastAPI server), pass
``auto_end=False`` so the runner does not exit when no workers
are left.
Args: Args:
worker: Optional pipeline worker to run. worker: Optional pipeline worker to run.
@@ -201,6 +213,10 @@ class PipelineRunner(BaseObject, BusSubscriber):
Register the worker with :meth:`add_workers` before Register the worker with :meth:`add_workers` before
calling ``run()`` instead. Passing ``worker`` here calling ``run()`` instead. Passing ``worker`` here
will be removed in a future release. will be removed in a future release.
auto_end: When ``True`` (the default), the runner ends once
every root worker has finished. When ``False``, the
runner blocks until :meth:`end` or :meth:`cancel` is
called.
""" """
if worker is not None: if worker is not None:
warnings.warn( warnings.warn(
@@ -211,6 +227,7 @@ class PipelineRunner(BaseObject, BusSubscriber):
) )
logger.debug(f"PipelineRunner '{self}': started running") logger.debug(f"PipelineRunner '{self}': started running")
self._auto_end = auto_end
self._shutdown_event.clear() self._shutdown_event.clear()
# Treat the main worker as any other added worker: ``add_workers`` attaches # Treat the main worker as any other added worker: ``add_workers`` attaches
@@ -222,15 +239,10 @@ class PipelineRunner(BaseObject, BusSubscriber):
await self._setup_session() await self._setup_session()
await self._call_event_handler("on_ready") await self._call_event_handler("on_ready")
# Wait for the main worker's background runner worker to finish # Wait for shutdown. With ``auto_end=True``, ``_run_worker`` sets
# (or for an explicit shutdown when there's no main worker). # ``_shutdown_event`` as soon as any root worker finishes.
try: try:
if worker is not None: await self._shutdown_event.wait()
runner_task = self._entries[worker.name].runner_task
if runner_task is not None:
await runner_task
else:
await self._shutdown_event.wait()
except asyncio.CancelledError: except asyncio.CancelledError:
pass pass
@@ -383,6 +395,18 @@ class PipelineRunner(BaseObject, BusSubscriber):
await worker.run(params) await worker.run(params)
except asyncio.CancelledError: except asyncio.CancelledError:
pass pass
finally:
# End the runner once every root worker has finished. The
# current worker's task is still "running" (we're inside its
# body), so exclude it from the check.
if self._auto_end and worker.parent is None:
others_running = any(
e.runner_task is not None and not e.runner_task.done()
for e in self._entries.values()
if e.worker.parent is None and e.worker is not worker
)
if not others_running:
self._shutdown_event.set()
async def _on_local_worker_ready(self, data: WorkerReadyData) -> None: async def _on_local_worker_ready(self, data: WorkerReadyData) -> None:
"""Called when a local added worker registers as ready.""" """Called when a local added worker registers as ready."""