From 66d72be1f40d87e1285c526a92d8e953e4f5a095 Mon Sep 17 00:00:00 2001 From: Andre Merzky Date: Wed, 6 May 2026 23:48:06 +0200 Subject: [PATCH 1/2] backends.dragon V3: chunk submit_tasks so dragon's bounded work-queue doesn't deadlock the asyncio event loop MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Dragon's ``Batch`` keeps an internal work-queue (``self.batch.work_q``, a stdlib ``queue.Queue`` with ``maxsize=manager_work_queue_max_batch_size`` — 256 by default) and every ``self.batch.process()/.job()/.function()`` call ends in a blocking ``self.work_q.put(task)``. Once the queue is full, ``put`` blocks the *calling thread* until dragon's dispatcher drains a slot. Previously, ``DragonExecutionBackendV3.submit_tasks`` invoked ``self.batch.{process,job,function}`` per task on the asyncio event-loop thread, in a tight ``for ... await build_task`` loop. At scale (1000+ tasks per batch, multi-node alloc where dragon's dispatcher is the slow side), the queue fills early and every subsequent put freezes the event loop for the duration of the drain. WS keepalives drop, the bridge times out the proxied request, the edge gets unregistered, and ``submit_tasks`` appears to hang indefinitely from the client's view. This change restructures the submit loop: * ``build_task``'s body -- fully synchronous already, with no awaits -- is moved to a new private ``_build_task_sync(task)`` helper. The async ``build_task`` is kept as a thin wrapper for backward compatibility with existing tests and call sites. * ``submit_tasks`` builds in chunks sized at ``self.batch.work_q.maxsize`` (with a 4096 fallback for unbounded queues) and offloads each chunk's blocking puts to a worker thread via ``asyncio.to_thread``. The event loop stays responsive for WS pings and other coroutines while the worker fills (and waits on) the dragon queue. * Tasks are now registered into ``_monitored_batches`` per chunk, not after the whole batch has been built -- so the monitor thread observes early completions instead of stalling them until all build_task calls return. Verified locally with a 500-task stress test (well above the 256-slot queue ceiling): submit returns in ~0.04s, tasks complete in ~3.3s, all 500 done. Single-node dragon's fast dispatcher means the local test cannot reproduce the original multi-node freeze, but the architectural property -- submit_tasks never blocking the event-loop thread -- holds regardless. Co-Authored-By: Claude Opus 4.7 (1M context) --- src/rhapsody/backends/execution/dragon.py | 90 +++++++++++++++++------ 1 file changed, 69 insertions(+), 21 deletions(-) diff --git a/src/rhapsody/backends/execution/dragon.py b/src/rhapsody/backends/execution/dragon.py index 8628275..e1ee4fc 100644 --- a/src/rhapsody/backends/execution/dragon.py +++ b/src/rhapsody/backends/execution/dragon.py @@ -3275,33 +3275,81 @@ async def submit_tasks(self, tasks: list[dict]) -> None: ) self._batch_monitor_thread.start() - # Build tasks - batch_tasks_data = [] - for task in tasks: - try: - batch_task = await self.build_task(task) - batch_tasks_data.append((task["uid"], batch_task)) - # This is the moment Dragon takes ownership was called - # inside build_task) and start executing it. + # Dragon's batch keeps an internal work-queue (``self.batch.work_q``, + # a stdlib ``queue.Queue`` with ``maxsize=manager_work_queue_max_batch_size``, + # 256 by default). Each ``self.batch.process/.job/.function`` call + # ends in a blocking ``self.work_q.put(task)``: when the queue is + # full, ``put`` blocks the *calling thread* until dragon's + # dispatcher drains a slot. Doing that on the asyncio event loop + # — which is what the previous serial ``for ... await build_task`` + # loop did — freezes WS keepalives and every other coroutine for + # the duration of the drain. At scale (1000+ tasks per batch on a + # multi-node alloc) that easily exceeds bridge / client timeouts. + # + # Fix: chunk submissions to the queue's maxsize and run each + # chunk's blocking puts in a worker thread via ``asyncio.to_thread``. + # The event loop is free while the worker fills (and waits on) + # the dragon queue; chunk-by-chunk registration into + # ``_monitored_batches`` lets the monitor thread observe early + # completions instead of stalling them until the whole batch is + # built. ``maxsize == 0`` (unbounded) falls back to a sane chunk + # size so we still yield periodically. + chunk = self.batch.work_q.maxsize or 4096 + + def _build_chunk(chunk_tasks): + """Synchronous per-chunk builder; runs in a worker thread. + + Returns a list of ``(task, batch_task | None, exception | None)`` + so the event-loop side can register monitored tasks and emit + RUNNING/FAILED callbacks without doing those calls itself in + a worker thread. + """ + out = [] + for task in chunk_tasks: + try: + batch_task = self._build_task_sync(task) + out.append((task, batch_task, None)) + except Exception as e: + out.append((task, None, e)) + return out + + n_built = 0 + for start in range(0, len(tasks), chunk): + chunk_results = await asyncio.to_thread( + _build_chunk, tasks[start:start + chunk]) + for task, batch_task, exc in chunk_results: + if exc is not None: + self.logger.error( + f"Failed to create task {task.get('uid')}: {exc}", + exc_info=exc) + task["exception"] = exc + self._callback_func(task, "FAILED") + continue + # Tasks are already in-flight — the Batch background thread + # auto-dispatches them the moment they were created via + # batch.function()/process()/job() inside _build_task_sync. + self._monitored_batches[batch_task.uid] = (batch_task, task["uid"]) self._callback_func(task, "RUNNING") - except Exception as e: - self.logger.error(f"Failed to create task {task.get('uid')}: {e}", exc_info=True) - task["exception"] = e - self._callback_func(task, "FAILED") - - if not batch_tasks_data: - return + n_built += 1 - # Tasks are already in-flight — the Batch background thread auto-dispatches them - # the moment they are created via batch.function()/process()/job(). - # Register each task individually for result monitoring. - for uid, batch_task in batch_tasks_data: - self._monitored_batches[batch_task.uid] = (batch_task, uid) - self.logger.info(f"Submitted {len(batch_tasks_data)} tasks (streaming, auto-dispatched)") + if n_built: + self.logger.info( + f"Submitted {n_built} tasks (streaming, auto-dispatched)") async def build_task(self, task: dict): """Translate AsyncFlow task to Dragon Batch task. + Async wrapper retained for backward compatibility with existing + callers (tests, V1/V2-style code paths). The body is fully + synchronous and lives in :meth:`_build_task_sync`; ``submit_tasks`` + offloads chunks of those sync calls to a worker thread so dragon's + blocking ``work_q.put`` does not freeze the event loop. + """ + return self._build_task_sync(task) + + def _build_task_sync(self, task: dict): + """Translate AsyncFlow task to Dragon Batch task (synchronous). + Translation Priority (in order): 1. If process_templates (list) provided → Job mode (ignore type='mpi', ignore ranks) [function/executable] 2. If process_template (single) provided → Process mode [function/executable] From 116a998db5a02aff19114a01d537ae0784d17c57 Mon Sep 17 00:00:00 2001 From: Andre Merzky Date: Tue, 30 Jun 2026 11:57:31 +0200 Subject: [PATCH 2/2] backends.dragon V3: satisfy ruff-format + docformatter on chunked submit Pure formatting (line joining + docstring reflow) on the chunking code; no logic change. Unblocks the pre-commit CI gate so the test job can run. Co-Authored-By: Claude Opus 4.8 (1M context) --- src/rhapsody/backends/execution/dragon.py | 17 +++++++---------- 1 file changed, 7 insertions(+), 10 deletions(-) diff --git a/src/rhapsody/backends/execution/dragon.py b/src/rhapsody/backends/execution/dragon.py index 3307634..adb5f79 100644 --- a/src/rhapsody/backends/execution/dragon.py +++ b/src/rhapsody/backends/execution/dragon.py @@ -3312,10 +3312,9 @@ async def submit_tasks(self, tasks: list[dict]) -> None: def _build_chunk(chunk_tasks): """Synchronous per-chunk builder; runs in a worker thread. - Returns a list of ``(task, batch_task | None, exception | None)`` - so the event-loop side can register monitored tasks and emit - RUNNING/FAILED callbacks without doing those calls itself in - a worker thread. + Returns a list of ``(task, batch_task | None, exception | None)`` so the event-loop side + can register monitored tasks and emit RUNNING/FAILED callbacks without doing those calls + itself in a worker thread. """ out = [] for task in chunk_tasks: @@ -3328,13 +3327,12 @@ def _build_chunk(chunk_tasks): n_built = 0 for start in range(0, len(tasks), chunk): - chunk_results = await asyncio.to_thread( - _build_chunk, tasks[start:start + chunk]) + chunk_results = await asyncio.to_thread(_build_chunk, tasks[start : start + chunk]) for task, batch_task, exc in chunk_results: if exc is not None: self.logger.error( - f"Failed to create task {task.get('uid')}: {exc}", - exc_info=exc) + f"Failed to create task {task.get('uid')}: {exc}", exc_info=exc + ) task["exception"] = exc self._callback_func(task, "FAILED") continue @@ -3346,8 +3344,7 @@ def _build_chunk(chunk_tasks): n_built += 1 if n_built: - self.logger.info( - f"Submitted {n_built} tasks (streaming, auto-dispatched)") + self.logger.info(f"Submitted {n_built} tasks (streaming, auto-dispatched)") async def build_task(self, task: dict): """Translate AsyncFlow task to Dragon Batch task.