[Bugfix] Hold references to fire-and-forget tasks in disaggregation (#33048)

This commit is contained in:
Lin Junrong
2026-08-30 14:26:34 +08:00
committed by GitHub
parent 5ec959965b
commit 512df615de
2 changed files with 17 additions and 4 deletions
@@ -1669,6 +1669,9 @@ class CommonKVBootstrapServer(BaseKVBootstrapServer):
self.app = web.Application()
self.store = dict()
self.lock = asyncio.Lock()
# The event loop only keeps weak references to tasks, so a long-lived
# task needs a strong reference to survive garbage collection.
self._background_tasks: Set[asyncio.Task] = set()
self._setup_routes()
self.pp_size = None
self.attn_tp_size = None
@@ -1916,7 +1919,9 @@ class CommonKVBootstrapServer(BaseKVBootstrapServer):
self._loop = asyncio.new_event_loop()
asyncio.set_event_loop(self._loop)
self._loop.create_task(self._cleanup_expired_entries())
cleanup_task = self._loop.create_task(self._cleanup_expired_entries())
self._background_tasks.add(cleanup_task)
cleanup_task.add_done_callback(self._background_tasks.discard)
access_log = None
if logging.getLogger(__name__).getEffectiveLevel() <= logging.DEBUG:
@@ -404,6 +404,9 @@ class DPDispatcher:
self._pending_send_at: Dict[str, float] = {}
# Set when _result_listener gives up; makes alive_ranks report empty.
self._listener_failed = False
# The event loop only keeps weak references to tasks, so the long-lived
# loops started in `start()` need a strong reference to survive GC.
self.background_tasks: Set[asyncio.Task] = set()
# Prometheus gauge: pending requests per DP rank. Lives in the main
# process (the dispatcher), unlike the per-worker EncoderMetricsCollector.
@@ -443,9 +446,14 @@ class DPDispatcher:
def start(self) -> None:
logger.info(f"DP dispatcher started: {self.dp_size} ranks (all remote)")
asyncio.create_task(self._result_listener())
asyncio.create_task(self._worker_watchdog())
asyncio.create_task(self._cleanup_stale_mappings())
for coro in (
self._result_listener(),
self._worker_watchdog(),
self._cleanup_stale_mappings(),
):
task = asyncio.create_task(coro)
self.background_tasks.add(task)
task.add_done_callback(self.background_tasks.discard)
def _drop_pending_and_mapping(self, rank: int, req_id: str) -> None:
# dispatch / broadcast failure: no follow-up /send expected.