diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py index 81b4fe0bf..3375ab60e 100644 --- a/python/sglang/srt/managers/tokenizer_manager.py +++ b/python/sglang/srt/managers/tokenizer_manager.py @@ -1190,7 +1190,14 @@ class TokenizerManager(TokenizerCommunicatorMixin): async def pause_generation(self): async with self.is_pause_cond: self.is_pause = True - self.abort_request(abort_all=True) + # we are using the model_update_lock to check if there is still on-going requests. + while True: + # TODO: maybe make it async instead of fire-and-forget + self.abort_request(abort_all=True) + is_locked = await self.model_update_lock.is_locked() + if not is_locked: + break + await asyncio.sleep(1.0) async def continue_generation(self): async with self.is_pause_cond: diff --git a/python/sglang/srt/utils/aio_rwlock.py b/python/sglang/srt/utils/aio_rwlock.py index deda1fe79..79dd4f242 100644 --- a/python/sglang/srt/utils/aio_rwlock.py +++ b/python/sglang/srt/utils/aio_rwlock.py @@ -75,6 +75,10 @@ class RWLock: # Wake up anyone waiting (readers or writers) self._cond.notify_all() + async def is_locked(self): + async with self._lock: + return self._writer_active or self._readers > 0 + class _ReaderLock: def __init__(self, rwlock: RWLock):