[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.app = web.Application()
|
||||||
self.store = dict()
|
self.store = dict()
|
||||||
self.lock = asyncio.Lock()
|
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._setup_routes()
|
||||||
self.pp_size = None
|
self.pp_size = None
|
||||||
self.attn_tp_size = None
|
self.attn_tp_size = None
|
||||||
@@ -1916,7 +1919,9 @@ class CommonKVBootstrapServer(BaseKVBootstrapServer):
|
|||||||
self._loop = asyncio.new_event_loop()
|
self._loop = asyncio.new_event_loop()
|
||||||
asyncio.set_event_loop(self._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
|
access_log = None
|
||||||
if logging.getLogger(__name__).getEffectiveLevel() <= logging.DEBUG:
|
if logging.getLogger(__name__).getEffectiveLevel() <= logging.DEBUG:
|
||||||
|
|||||||
@@ -404,6 +404,9 @@ class DPDispatcher:
|
|||||||
self._pending_send_at: Dict[str, float] = {}
|
self._pending_send_at: Dict[str, float] = {}
|
||||||
# Set when _result_listener gives up; makes alive_ranks report empty.
|
# Set when _result_listener gives up; makes alive_ranks report empty.
|
||||||
self._listener_failed = False
|
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
|
# Prometheus gauge: pending requests per DP rank. Lives in the main
|
||||||
# process (the dispatcher), unlike the per-worker EncoderMetricsCollector.
|
# process (the dispatcher), unlike the per-worker EncoderMetricsCollector.
|
||||||
@@ -443,9 +446,14 @@ class DPDispatcher:
|
|||||||
|
|
||||||
def start(self) -> None:
|
def start(self) -> None:
|
||||||
logger.info(f"DP dispatcher started: {self.dp_size} ranks (all remote)")
|
logger.info(f"DP dispatcher started: {self.dp_size} ranks (all remote)")
|
||||||
asyncio.create_task(self._result_listener())
|
for coro in (
|
||||||
asyncio.create_task(self._worker_watchdog())
|
self._result_listener(),
|
||||||
asyncio.create_task(self._cleanup_stale_mappings())
|
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:
|
def _drop_pending_and_mapping(self, rank: int, req_id: str) -> None:
|
||||||
# dispatch / broadcast failure: no follow-up /send expected.
|
# dispatch / broadcast failure: no follow-up /send expected.
|
||||||
|
|||||||
Reference in New Issue
Block a user