fix(dragon V3): chunk submit_tasks via asyncio.to_thread so dragon's bounded work-queue doesn't freeze the event loop

[andre-merzky]

Summary

Restructures DragonExecutionBackendV3.submit_tasks so that dragon's
bounded internal work-queue can never block the asyncio event loop.
Submission is chunked at self.batch.work_q.maxsize (default 256, 4096
fallback for unbounded), and each chunk's blocking self.batch.{process, job,function}(...) calls run on a worker thread via
asyncio.to_thread. Tasks register into _monitored_batches chunk-by-
chunk, not en bloc, so the monitor thread observes early completions
instead of stalling them until the whole batch is built.

Background

Dragon's Batch keeps an internal queue.Queue(maxsize=manager_work_queue_max_batch_size)
— 256 by default. Every self.batch.process()/.job()/.function() call
ends in a blocking self.work_q.put(task): when the queue is full, the
calling thread blocks until dragon's dispatcher drains a slot.

Previously submit_tasks invoked those builders per-task on the asyncio
event-loop thread in a serial for ... await build_task loop. At scale
(1000+ tasks per batch on a multi-node alloc, where dragon's dispatcher is
the slow side) the queue fills early and every subsequent put freezes the
loop for the duration of the drain. The user-visible failure mode:

• WS keepalives drop on the bridge↔edge channel.
• The bridge times out the proxied request, the edge gets unregistered.
• submit_tasks appears to hang indefinitely from the client's view.

What changes

• build_task's body — fully synchronous already, no awaits — is
  moved to a new private _build_task_sync(task) helper. The async
  build_task is kept as a thin one-liner wrapper for back-compat with
  existing tests / call sites.
• submit_tasks chunks the input at work_q.maxsize (or 4096 if
  unbounded) and runs each chunk's blocking puts in 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 registered into _monitored_batches per chunk, not after
  the whole batch is built.

History note

This commit was previously merged to main (df26eef) and then immediately
reverted (bf89444) because the merge had been done by force-push instead
of via a proper PR review — this PR is the proper-review re-apply, with no
content changes vs. df26eef. The branch was rebased / cherry-picked onto
the current main tip; the diff against main is purely the chunking change.

Test plan

• Local 500-task stress test (well above the 256-slot ceiling): submit
  returns in ~0.04s, all 500 tasks complete in ~3.3s. Single-node
  dispatcher is too fast to reproduce the original multi-node freeze,
  but the architectural property — submit_tasks never blocking the
  event-loop thread — holds regardless.
• Multi-node soak test on Perlmutter once fix(dragon V3): guard results_ddict subscript with __contains__ to unwedge multi-node monitor loop #52 lands (the
  __contains__ guard for results_ddict is a prerequisite for any
  multi-node V3 run to actually deliver completions).

[gemini-code-assist]

Code Review

This pull request refactors the submit_tasks method to process tasks in chunks and offload them to worker threads using asyncio.to_thread, preventing the asyncio event loop from blocking when the internal Dragon work queue is full. A new synchronous helper, _build_task_sync, was introduced to facilitate this. Feedback was provided regarding the build_task method, which remains async but currently calls the synchronous logic directly; it should also use asyncio.to_thread to maintain its non-blocking contract.

[gemini-code-assist]

The build_task method is still marked as async, but it now directly calls the synchronous and potentially blocking _build_task_sync. If a caller (such as a test or a V1/V2-style code path) uses await backend.build_task(task), it will now block the asyncio event loop during the self.batch.put call inside Dragon's internal logic. To maintain the non-blocking contract of the async method, you should offload this call to a worker thread as well.

Suggested change

	return self._build_task_sync(task)
	return await asyncio.to_thread(self._build_task_sync, task)

[andre-merzky]

/gemini review

[gemini-code-assist]

Code Review

This pull request refactors task submission in the Dragon backend to prevent blocking the asyncio event loop by chunking task submissions and offloading synchronous task building to a worker thread via asyncio.to_thread. The reviewer identified a critical thread-safety issue where concurrent access to self._task_registry and self._monitored_batches from the worker thread and the event loop/monitor threads could cause race conditions or runtime errors, and provided a detailed thread-safe proxying solution.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

Thread-Safety Issue with self._task_registry and self._monitored_batches

Running _build_task_sync in a worker thread via asyncio.to_thread introduces a critical thread-safety issue:

1. self._task_registry: Written to from the worker thread inside _build_task_sync while concurrently read/popped from the asyncio event loop thread (e.g., in _deliver_batch and cancel_task). This can lead to RuntimeError or lost task registrations.
2. self._monitored_batches: Written to/popped from the event loop thread while concurrently iterated/popped from the background monitor thread (_monitor_loop).

To resolve this safely without modifying code outside the diff hunks, we can proxy both dictionaries with thread-aware/thread-safe wrappers at the start of submit_tasks:

• self._task_registry is proxied to redirect worker thread writes to thread-local storage, which are then safely applied to the main dictionary on the event loop thread after asyncio.to_thread completes.
• self._monitored_batches is proxied with a thread-safe wrapper using a threading.Lock to synchronize accesses between the event loop thread and the monitor thread.

        # 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

        # Thread-safety: self._task_registry is modified in _build_task_sync.
        # Since _build_task_sync is now run in a worker thread via asyncio.to_thread,
        # we must prevent concurrent writes to self._task_registry from the worker thread.
        # We proxy self._task_registry with a thread-aware dict that redirects worker thread
        # writes to thread-local storage, then apply them on the event loop thread.
        local_storage = threading.local()
        main_thread_id = threading.get_ident()
        original_registry = self._task_registry

        if not getattr(self, "_task_registry_proxied", False):
            class ThreadAwareDict(dict):
                def __setitem__(self, key, value):
                    if threading.get_ident() == main_thread_id:
                        original_registry[key] = value
                    else:
                        if not hasattr(local_storage, "tasks"):
                            local_storage.tasks = {}
                        local_storage.tasks[key] = value

                def __getitem__(self, key):
                    if threading.get_ident() == main_thread_id:
                        return original_registry[key]
                    return getattr(local_storage, "tasks", {}).get(key) or original_registry[key]

                def pop(self, key, default=None):
                    if threading.get_ident() == main_thread_id:
                        return original_registry.pop(key, default)
                    return getattr(local_storage, "tasks", {}).pop(key, default)

                def __contains__(self, key):
                    if threading.get_ident() == main_thread_id:
                        return key in original_registry
                    return key in getattr(local_storage, "tasks", {}) or key in original_registry

                def clear(self):
                    original_registry.clear()

            self._task_registry = ThreadAwareDict()
            self._task_registry_proxied = True

        # Thread-safety: self._monitored_batches is accessed concurrently by the event loop
        # thread and the background monitor thread. We proxy it with a thread-safe dict.
        if not getattr(self, "_monitored_batches_proxied", False):
            original_monitored = self._monitored_batches
            monitored_lock = threading.Lock()
            class ThreadSafeDict(dict):
                def __setitem__(self, key, value):
                    with monitored_lock:
                        original_monitored[key] = value
                def __getitem__(self, key):
                    with monitored_lock:
                        return original_monitored[key]
                def pop(self, key, default=None):
                    with monitored_lock:
                        return original_monitored.pop(key, default)
                def keys(self):
                    with monitored_lock:
                        return list(original_monitored.keys())
                def get(self, key, default=None):
                    with monitored_lock:
                        return original_monitored.get(key, default)
                def __contains__(self, key):
                    with monitored_lock:
                        return key in original_monitored
                def __bool__(self):
                    with monitored_lock:
                        return bool(original_monitored)
                def __len__(self):
                    with monitored_lock:
                        return len(original_monitored)
            self._monitored_batches = ThreadSafeDict()
            self._monitored_batches_proxied = True

        def _build_chunk(chunk_tasks):
            """Synchronous per-chunk builder; runs in a worker thread.

            Returns a list of ``(task, batch_task | None, task_info | 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)
                    task_info = getattr(local_storage, "tasks", {}).pop(task["uid"], None)
                    out.append((task, batch_task, task_info, None))
                except Exception as e:
                    out.append((task, None, 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, task_info, 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
                
                # Safely register the task info on the event loop thread
                if task_info:
                    original_registry[task["uid"]] = task_info

                # 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")
                n_built += 1

        if n_built:
            self.logger.info(
                f"Submitted {n_built} tasks (streaming, auto-dispatched)")

[andre-merzky]

On the flagged _task_registry thread-safety: under CPython's GIL the dict ops here are atomic, nothing iterates the registry, and per-uid there's a build→register→complete→pop happens-before — so there's no observable race today. Deferring the hardening to a dedicated free-threaded (no-GIL) review pass rather than adding the proxy machinery now.