[Bugfix] Hold references to fire-and-forget tasks in disaggregation (#33048)
This commit is contained in:
@@ -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.
|
||||
|
||||
Reference in New Issue
Block a user