From 8a63dfcf28d72c15d3844a84d36bc5a824c13209 Mon Sep 17 00:00:00 2001 From: Lawrence Wu Date: Wed, 18 Mar 2026 18:27:10 -0700 Subject: [PATCH 1/4] fix: /pause_generation with --tokenizer-worker-num > 1 /pause_generation sets `is_pause` only on one worker. /continue_generation unsets `is_pause` on a potentially different worker. e.g. with 8 workers, with 7/8 probability we will leave 1/8 of workers paused, thus 1/8 of all requests will hang. With PD disaggregation, with 15/64 probability a request will get routed to one worker which is paused and one worker which is not paused, and these requests will wait forever for KV transfer. We fix this by setting the pause state in shared memory and polling it to update the local `is_pause` / `is_pause_cond`. Note that the consistency of `is_pause` will be reduced to eventual consistency in the multi-worker case. e.g. a pause immediately followed by a weight update may result in the weight update not seeing `is_pause` set immediately. Thus, we change the logic to continue checking for the pause flag while waiting for the writer lock. NOTE: We do not fix the following existing issues: * /continue_generation is not safe during a weight update which started while paused * weight update with multi-worker without pausing is not safe, i.e. `self.model_update_lock.writer_lock` is local, but the lock needs to be acquired across all workers --- python/sglang/srt/entrypoints/http_server.py | 7 +- .../srt/managers/multi_tokenizer_mixin.py | 25 +++++ .../srt/managers/tokenizer_control_mixin.py | 35 ++---- .../sglang/srt/managers/tokenizer_manager.py | 101 ++++++++++++++++-- python/sglang/srt/utils/aio_rwlock.py | 1 + .../tokenizer/test_multi_tokenizer.py | 44 ++++++++ 6 files changed, 172 insertions(+), 41 deletions(-) diff --git a/python/sglang/srt/entrypoints/http_server.py b/python/sglang/srt/entrypoints/http_server.py index aba143c83c16..de55dbf91803 100644 --- a/python/sglang/srt/entrypoints/http_server.py +++ b/python/sglang/srt/entrypoints/http_server.py @@ -146,6 +146,7 @@ MultiTokenizerRouter, TokenizerWorker, get_main_process_id, + create_shared_pause_flag, monkey_patch_uvicorn_multiprocessing, read_from_shared_memory, write_data_for_multi_tokenizer, @@ -602,7 +603,6 @@ async def model_info(): "has_audio_understanding": model_config.is_audio_understandable_model, "model_type": getattr(model_config.hf_config, "model_type", None), "architectures": getattr(model_config.hf_config, "architectures", None), - "weight_version": _global_state.tokenizer_manager.server_args.weight_version, # "hf_config": model_config.hf_config.to_dict(), } return result @@ -1946,7 +1946,7 @@ def _execute_server_warmup(server_args: ServerArgs): _global_state.tokenizer_manager.server_status = ServerStatus.Up else: - logger.info(f"Start of pd disaggregation warmup ...") + logger.info("Start of pd disaggregation warmup ...") json_data = { "sampling_params": { "temperature": 0.0, @@ -2188,6 +2188,7 @@ def _setup_and_run_http_server( multi_tokenizer_args_shm = write_data_for_multi_tokenizer( port_args, server_args, scheduler_infos[0] ) + pause_flag_shm = create_shared_pause_flag(os.getpid()) try: # Update logging configs @@ -2289,6 +2290,8 @@ async def _run_with_ssl_refresh(): if server_args.tokenizer_worker_num > 1: if multi_tokenizer_args_shm is not None: multi_tokenizer_args_shm.unlink() + pause_flag_shm.close() + pause_flag_shm.unlink() if _global_state is not None: _global_state.tokenizer_manager.socket_mapping.clear_all_sockets() diff --git a/python/sglang/srt/managers/multi_tokenizer_mixin.py b/python/sglang/srt/managers/multi_tokenizer_mixin.py index 9757e4b5809f..c398ecee8aee 100644 --- a/python/sglang/srt/managers/multi_tokenizer_mixin.py +++ b/python/sglang/srt/managers/multi_tokenizer_mixin.py @@ -404,6 +404,9 @@ def __init__( self.send_to_scheduler, 2 ) + main_pid = get_main_process_id() + self._init_shared_pause(get_shared_pause_shm_name(main_pid)) + def _attach_multi_http_worker_info(self, req: Union[BaseReq, BaseBatchReq]): if isinstance(req, BaseReq): @@ -477,6 +480,28 @@ def read_from_shared_memory(name: str) -> Any: raise FileNotFoundError(f"Shared memory {name} not found") +SHARED_PAUSE_SHM_PREFIX = "sglang_pause_" + + +def get_shared_pause_shm_name(pid: int) -> str: + return f"{SHARED_PAUSE_SHM_PREFIX}{pid}" + + +def create_shared_pause_flag(pid: int) -> shared_memory.SharedMemory: + """Returns a SharedMemory object, which the caller must close and unlink on shutdown.""" + name = get_shared_pause_shm_name(pid) + try: + shm = shared_memory.SharedMemory(name=name) + if shm.size < 1: + shm.close() + shm.unlink() + shm = shared_memory.SharedMemory(create=True, size=1, name=name) + except FileNotFoundError: + shm = shared_memory.SharedMemory(create=True, size=1, name=name) + shm.buf[0] = 0 + return shm + + def write_data_for_multi_tokenizer( port_args: PortArgs, server_args: ServerArgs, scheduler_info: Dict ): diff --git a/python/sglang/srt/managers/tokenizer_control_mixin.py b/python/sglang/srt/managers/tokenizer_control_mixin.py index b18c3cd340b9..1dbb8921ef26 100644 --- a/python/sglang/srt/managers/tokenizer_control_mixin.py +++ b/python/sglang/srt/managers/tokenizer_control_mixin.py @@ -425,15 +425,8 @@ async def update_weights_from_distributed( if obj.abort_all_requests: self.abort_request(abort_all=True) - # Hold is_pause_cond while updating to prevent unpause from racing. - async with self.is_pause_cond: - is_paused = self.is_pause - if is_paused: - results = await self.update_weights_from_distributed_communicator(obj) - - if not is_paused: - async with self.model_update_lock.writer_lock: - results = await self.update_weights_from_distributed_communicator(obj) + async with self._wait_for_pause_or_lock(): + results = await self.update_weights_from_distributed_communicator(obj) success, message = FanOutCommunicator.merge_results(results) if success and obj.weight_version is not None: @@ -483,14 +476,8 @@ async def update_weights_from_tensor( if obj.abort_all_requests: self.abort_request(abort_all=True) - async with self.is_pause_cond: - is_paused = self.is_pause - if is_paused: - results = await self.update_weights_from_tensor_communicator(obj) - - if not is_paused: - async with self.model_update_lock.writer_lock: - results = await self.update_weights_from_tensor_communicator(obj) + async with self._wait_for_pause_or_lock(): + results = await self.update_weights_from_tensor_communicator(obj) success, message = FanOutCommunicator.merge_results(results) if success and obj.weight_version is not None: @@ -512,17 +499,9 @@ async def update_weights_from_ipc( self.server_args.dp_size == 1 or self.server_args.enable_dp_attention ), "dp_size must be 1 or dp attention must be enabled for update weights from IPC" logger.info("Starting IPC weight update") - - async with self.is_pause_cond: - is_paused = self.is_pause - if is_paused: - result = (await self.update_weights_from_ipc_communicator(obj))[0] - success, message = result.success, result.message - - if not is_paused: - async with self.model_update_lock.writer_lock: - result = (await self.update_weights_from_ipc_communicator(obj))[0] - success, message = result.success, result.message + async with self._wait_for_pause_or_lock(): + result = (await self.update_weights_from_ipc_communicator(obj))[0] + success, message = result.success, result.message except Exception as e: error_msg = f"IPC weight update failed: {str(e)}" logger.error(error_msg) diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py index 8a49717c58bb..daaf818bb687 100644 --- a/python/sglang/srt/managers/tokenizer_manager.py +++ b/python/sglang/srt/managers/tokenizer_manager.py @@ -25,10 +25,11 @@ import sys import threading from collections import deque -from contextlib import nullcontext +from contextlib import asynccontextmanager, nullcontext from datetime import datetime from enum import Enum from http import HTTPStatus +from multiprocessing import shared_memory from typing import Any, Awaitable, Dict, List, Optional, Tuple, Union import fastapi @@ -416,6 +417,86 @@ def init_weight_update(self): ) self.is_pause = False self.is_pause_cond = asyncio.Condition() + self._is_pause_shm = None + + def _init_shared_pause(self, shm_name: str): + """Attach to shared pause flag for multi-worker coordination. + + In multi-tokenizer mode, each worker process has its own is_pause flag + and asyncio.Condition, which cannot be shared across processes. This + method connects to a shared memory byte so that pause/continue from + any worker is visible to all workers. A background task polls the + shared flag and updates the local is_pause + is_pause_cond. + """ + self._is_pause_shm = shared_memory.SharedMemory(name=shm_name) + self.is_pause = bool(self._is_pause_shm.buf[0]) + self.asyncio_tasks.add(asyncio.create_task(self._poll_shared_pause())) + + async def _poll_shared_pause(self): + """Background task that polls shared memory and syncs local is_pause.""" + while True: + await asyncio.sleep(0.1) + shm_paused = bool(self._is_pause_shm.buf[0]) + if shm_paused != self.is_pause: + async with self.is_pause_cond: + self.is_pause = shm_paused + self.is_pause_cond.notify_all() + + @asynccontextmanager + async def _wait_for_pause_or_lock(self): + """ + Acquire writer lock OR skip it if the engine becomes paused. + + NOTE: + * /continue_generation is not safe while we are within this context + * if using --tokenizer-worker-num > 1, entering this context without + pausing is not safe, i.e. `self.model_update_lock.writer_lock` is + local and does not guarantee lock acquisition across all workers + """ + if self._is_pause_shm is not None: + # self.is_pause is eventually consistent, so check self._is_pause_shm first + is_paused = bool(self._is_pause_shm.buf[0]) + else: + async with self.is_pause_cond: + is_paused = self.is_pause + if is_paused: + yield + return + + lock_task = asyncio.create_task(self.model_update_lock.acquire_writer()) + pause_task = asyncio.create_task(self._wait_until_paused()) + lock_acquired = False + try: + _, pending = await asyncio.wait( + [lock_task, pause_task], return_when=asyncio.FIRST_COMPLETED + ) + for t in pending: + t.cancel() + for t in pending: + try: + await t + except asyncio.CancelledError: + pass + if lock_task.done() and not lock_task.cancelled(): + await lock_task + lock_acquired = True + if pause_task.done() and not pause_task.cancelled(): + await pause_task + yield + finally: + for t in (lock_task, pause_task): + if not t.done(): + t.cancel() + try: + await t + except asyncio.CancelledError: + pass + if lock_acquired: + await self.model_update_lock.release_writer() + + async def _wait_until_paused(self): + async with self.is_pause_cond: + await self.is_pause_cond.wait_for(lambda: self.is_pause) def init_lora(self): # LoRA @@ -1470,8 +1551,11 @@ def abort_request(self, rid: str = "", abort_all: bool = False): async def pause_generation(self, obj: PauseGenerationReqInput): async with self.is_pause_cond: self.is_pause = True + if self._is_pause_shm is not None: + self._is_pause_shm.buf[0] = 1 + self.is_pause_cond.notify_all() if obj.mode != "abort": - await self.send_to_scheduler.send_pyobj(obj) + self.send_to_scheduler.send_pyobj(obj) else: # we are using the model_update_lock to check if there is still on-going requests. while True: @@ -1485,7 +1569,9 @@ async def pause_generation(self, obj: PauseGenerationReqInput): async def continue_generation(self, obj: ContinueGenerationReqInput): async with self.is_pause_cond: self.is_pause = False - await self.send_to_scheduler.send_pyobj(obj) + if self._is_pause_shm is not None: + self._is_pause_shm.buf[0] = 0 + self.send_to_scheduler.send_pyobj(obj) self.is_pause_cond.notify_all() async def update_weights_from_disk( @@ -1503,14 +1589,7 @@ async def update_weights_from_disk( if obj.abort_all_requests: self.abort_request(abort_all=True) - # Immediately update the weights if the engine is in paused state - async with self.is_pause_cond: - is_paused = self.is_pause - - lock_context = ( - self.model_update_lock.writer_lock if not is_paused else nullcontext() - ) - async with lock_context: + async with self._wait_for_pause_or_lock(): success, message, num_paused_requests = ( await self._wait_for_model_update_from_disk(obj) ) diff --git a/python/sglang/srt/utils/aio_rwlock.py b/python/sglang/srt/utils/aio_rwlock.py index 79dd4f242396..0904f2d531ab 100644 --- a/python/sglang/srt/utils/aio_rwlock.py +++ b/python/sglang/srt/utils/aio_rwlock.py @@ -68,6 +68,7 @@ async def acquire_writer(self): finally: # Decrement waiting writers only after we've acquired the writer lock self._waiting_writers -= 1 + self._cond.notify_all() async def release_writer(self): async with self._lock: diff --git a/test/registered/tokenizer/test_multi_tokenizer.py b/test/registered/tokenizer/test_multi_tokenizer.py index 3fb76315998f..c3f44c0b08c3 100644 --- a/test/registered/tokenizer/test_multi_tokenizer.py +++ b/test/registered/tokenizer/test_multi_tokenizer.py @@ -1,5 +1,9 @@ +import concurrent.futures +import time import unittest +import requests + from sglang.srt.utils import kill_process_tree from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci from sglang.test.kits.eval_accuracy_kit import MMLUMixin @@ -73,6 +77,46 @@ def test_multi_tokenizer_ttft(self): self.assertLess(res["median_ttft_ms"], 86) self.assertLess(res["median_itl_ms"], 10) + def test_pause_continue_generation(self): + """Test that pause/continue works across all tokenizer workers.""" + + def generate(timeout=30): + return requests.post( + self.base_url + "/generate", + json={ + "text": "The capital of France is", + "sampling_params": {"temperature": 0, "max_new_tokens": 8}, + }, + timeout=timeout, + ) + + requests.post( + self.base_url + "/pause_generation", + json={"mode": "in_place"}, + timeout=30, + ).raise_for_status() + + num_requests = 16 + with concurrent.futures.ThreadPoolExecutor(max_workers=num_requests) as pool: + futures = [pool.submit(generate, timeout=60) for _ in range(num_requests)] + + time.sleep(2) + + done = [f for f in futures if f.done()] + self.assertEqual( + len(done), + 0, + f"{len(done)}/{num_requests} requests completed while paused", + ) + + requests.post( + self.base_url + "/continue_generation", json={} + ).raise_for_status() + + for f in concurrent.futures.as_completed(futures, timeout=60): + resp = f.result() + self.assertEqual(resp.status_code, 200) + if __name__ == "__main__": unittest.main() From 010f9000d0ad54f1922f024edb6ed5a5d4806c23 Mon Sep 17 00:00:00 2001 From: Lawrence Wu Date: Thu, 16 Apr 2026 23:02:36 +0000 Subject: [PATCH 2/4] update docstring --- python/sglang/srt/managers/tokenizer_manager.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py index daaf818bb687..588c5c84f5fc 100644 --- a/python/sglang/srt/managers/tokenizer_manager.py +++ b/python/sglang/srt/managers/tokenizer_manager.py @@ -424,9 +424,11 @@ def _init_shared_pause(self, shm_name: str): In multi-tokenizer mode, each worker process has its own is_pause flag and asyncio.Condition, which cannot be shared across processes. This - method connects to a shared memory byte so that pause/continue from - any worker is visible to all workers. A background task polls the - shared flag and updates the local is_pause + is_pause_cond. + method connects to a shared memory byte so that pause/continue from any + worker is visible to all workers. A background task polls the shared + flag and updates the local is_pause + is_pause_cond. Note that the + consistency of is_pause is reduced to eventual consistency in this + case. """ self._is_pause_shm = shared_memory.SharedMemory(name=shm_name) self.is_pause = bool(self._is_pause_shm.buf[0]) From 37a3c7dc21ce8c573e519812a0740427a5495f81 Mon Sep 17 00:00:00 2001 From: Lawrence Wu Date: Fri, 17 Apr 2026 17:25:42 +0000 Subject: [PATCH 3/4] restore toctou fix --- .../srt/managers/tokenizer_control_mixin.py | 6 +- .../sglang/srt/managers/tokenizer_manager.py | 130 ++++++++------- python/sglang/srt/utils/aio_rwlock.py | 157 +++++++++++++----- 3 files changed, 186 insertions(+), 107 deletions(-) diff --git a/python/sglang/srt/managers/tokenizer_control_mixin.py b/python/sglang/srt/managers/tokenizer_control_mixin.py index 1dbb8921ef26..82446766bedb 100644 --- a/python/sglang/srt/managers/tokenizer_control_mixin.py +++ b/python/sglang/srt/managers/tokenizer_control_mixin.py @@ -425,7 +425,7 @@ async def update_weights_from_distributed( if obj.abort_all_requests: self.abort_request(abort_all=True) - async with self._wait_for_pause_or_lock(): + async with self._ensure_paused_or_model_locked(): results = await self.update_weights_from_distributed_communicator(obj) success, message = FanOutCommunicator.merge_results(results) @@ -476,7 +476,7 @@ async def update_weights_from_tensor( if obj.abort_all_requests: self.abort_request(abort_all=True) - async with self._wait_for_pause_or_lock(): + async with self._ensure_paused_or_model_locked(): results = await self.update_weights_from_tensor_communicator(obj) success, message = FanOutCommunicator.merge_results(results) @@ -499,7 +499,7 @@ async def update_weights_from_ipc( self.server_args.dp_size == 1 or self.server_args.enable_dp_attention ), "dp_size must be 1 or dp attention must be enabled for update weights from IPC" logger.info("Starting IPC weight update") - async with self._wait_for_pause_or_lock(): + async with self._ensure_paused_or_model_locked(): result = (await self.update_weights_from_ipc_communicator(obj))[0] success, message = result.success, result.message except Exception as e: diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py index 588c5c84f5fc..3e8cc9a06068 100644 --- a/python/sglang/srt/managers/tokenizer_manager.py +++ b/python/sglang/srt/managers/tokenizer_manager.py @@ -107,7 +107,7 @@ get_or_create_event_loop, kill_process_tree, ) -from sglang.srt.utils.aio_rwlock import RWLock +from sglang.srt.utils.aio_rwlock import RWCondition, RWLock from sglang.srt.utils.hf_transformers_utils import ( get_processor, get_tokenizer, @@ -416,15 +416,15 @@ def init_weight_update(self): None ) self.is_pause = False - self.is_pause_cond = asyncio.Condition() + self.is_pause_cond = RWCondition() self._is_pause_shm = None def _init_shared_pause(self, shm_name: str): """Attach to shared pause flag for multi-worker coordination. In multi-tokenizer mode, each worker process has its own is_pause flag - and asyncio.Condition, which cannot be shared across processes. This - method connects to a shared memory byte so that pause/continue from any + and RWCondition, which cannot be shared across processes. This method + connects to a shared memory byte so that pause/continue from any worker is visible to all workers. A background task polls the shared flag and updates the local is_pause + is_pause_cond. Note that the consistency of is_pause is reduced to eventual consistency in this @@ -438,67 +438,77 @@ async def _poll_shared_pause(self): """Background task that polls shared memory and syncs local is_pause.""" while True: await asyncio.sleep(0.1) - shm_paused = bool(self._is_pause_shm.buf[0]) - if shm_paused != self.is_pause: - async with self.is_pause_cond: - self.is_pause = shm_paused - self.is_pause_cond.notify_all() + if bool(self._is_pause_shm.buf[0]) != self.is_pause: + async with self.is_pause_cond.writer_lock: + shm_paused = bool(self._is_pause_shm.buf[0]) + if shm_paused != self.is_pause: + self.is_pause = shm_paused + self.is_pause_cond.notify_all() @asynccontextmanager - async def _wait_for_pause_or_lock(self): + async def _ensure_paused_or_model_locked(self): """ - Acquire writer lock OR skip it if the engine becomes paused. + Context manager: within the context, either the engine is paused OR the + model_update_lock writer is held. - NOTE: - * /continue_generation is not safe while we are within this context - * if using --tokenizer-worker-num > 1, entering this context without - pausing is not safe, i.e. `self.model_update_lock.writer_lock` is - local and does not guarantee lock acquisition across all workers - """ - if self._is_pause_shm is not None: - # self.is_pause is eventually consistent, so check self._is_pause_shm first - is_paused = bool(self._is_pause_shm.buf[0]) - else: - async with self.is_pause_cond: - is_paused = self.is_pause - if is_paused: - yield - return + Acquires is_pause_cond as a reader, which blocks writers + (continue_generation, pause_generation) but not other readers + (send_request). This prevents continue_generation from unpausing + while a weight update is in-flight (TOCTOU fix from #22304). + + NOTE: with --tokenizer-worker-num > 1, is_pause_cond is process-local + so it cannot block continue_generation in another worker process. - lock_task = asyncio.create_task(self.model_update_lock.acquire_writer()) - pause_task = asyncio.create_task(self._wait_until_paused()) - lock_acquired = False + NOTE: with --tokenizer-worker-num > 1, model_update_lock is + process-local so it cannot block model read/update in another worker + process. + """ + await self.is_pause_cond.acquire_reader() + pause_reader_held = True try: - _, pending = await asyncio.wait( - [lock_task, pause_task], return_when=asyncio.FIRST_COMPLETED - ) - for t in pending: - t.cancel() - for t in pending: - try: - await t - except asyncio.CancelledError: - pass - if lock_task.done() and not lock_task.cancelled(): - await lock_task - lock_acquired = True - if pause_task.done() and not pause_task.cancelled(): - await pause_task - yield - finally: - for t in (lock_task, pause_task): - if not t.done(): + if self.is_pause: + yield + return + + async def wait_until_paused(): + await self.is_pause_cond.wait_for_reader(lambda: self.is_pause) + + lock_task = asyncio.create_task(self.model_update_lock.acquire_writer()) + pause_task = asyncio.create_task(wait_until_paused()) + lock_acquired = False + try: + _, pending = await asyncio.wait( + [lock_task, pause_task], return_when=asyncio.FIRST_COMPLETED + ) + for t in pending: t.cancel() + for t in pending: try: await t except asyncio.CancelledError: pass - if lock_acquired: - await self.model_update_lock.release_writer() - - async def _wait_until_paused(self): - async with self.is_pause_cond: - await self.is_pause_cond.wait_for(lambda: self.is_pause) + if lock_task.done() and not lock_task.cancelled(): + await lock_task + lock_acquired = True + if pause_task.done() and not pause_task.cancelled(): + await pause_task + if lock_acquired: + await self.is_pause_cond.release_reader() + pause_reader_held = False + yield + finally: + for t in (lock_task, pause_task): + if not t.done(): + t.cancel() + try: + await t + except asyncio.CancelledError: + pass + if lock_acquired: + await self.model_update_lock.release_writer() + finally: + if pause_reader_held: + await self.is_pause_cond.release_reader() def init_lora(self): # LoRA @@ -626,8 +636,8 @@ async def generate_request( # Log the request self.request_logger.log_received_request(obj, self.tokenizer, request) - async with self.is_pause_cond: - await self.is_pause_cond.wait_for(lambda: not self.is_pause) + async with self.is_pause_cond.reader_lock: + await self.is_pause_cond.wait_for_reader(lambda: not self.is_pause) async with self.model_update_lock.reader_lock: await self._validate_and_resolve_lora(obj) @@ -1551,7 +1561,7 @@ def abort_request(self, rid: str = "", abort_all: bool = False): ) async def pause_generation(self, obj: PauseGenerationReqInput): - async with self.is_pause_cond: + async with self.is_pause_cond.writer_lock: self.is_pause = True if self._is_pause_shm is not None: self._is_pause_shm.buf[0] = 1 @@ -1563,13 +1573,13 @@ async def pause_generation(self, obj: PauseGenerationReqInput): 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() + is_locked = self.model_update_lock.is_locked() if not is_locked: break await asyncio.sleep(1.0) async def continue_generation(self, obj: ContinueGenerationReqInput): - async with self.is_pause_cond: + async with self.is_pause_cond.writer_lock: self.is_pause = False if self._is_pause_shm is not None: self._is_pause_shm.buf[0] = 0 @@ -1591,7 +1601,7 @@ async def update_weights_from_disk( if obj.abort_all_requests: self.abort_request(abort_all=True) - async with self._wait_for_pause_or_lock(): + async with self._ensure_paused_or_model_locked(): success, message, num_paused_requests = ( await self._wait_for_model_update_from_disk(obj) ) diff --git a/python/sglang/srt/utils/aio_rwlock.py b/python/sglang/srt/utils/aio_rwlock.py index 0904f2d531ab..e159d090b830 100644 --- a/python/sglang/srt/utils/aio_rwlock.py +++ b/python/sglang/srt/utils/aio_rwlock.py @@ -1,84 +1,153 @@ import asyncio +import collections class RWLock: - def __init__(self): - # Protects internal state - self._lock = asyncio.Lock() - - # Condition variable used to wait for state changes - self._cond = asyncio.Condition(self._lock) + """asyncio readers-writer lock with writer preference. Non-reetrant.""" - # Number of readers currently holding the lock + def __init__(self): + self._cond = asyncio.Condition() self._readers = 0 - - # Whether a writer is currently holding the lock self._writer_active = False - - # How many writers are queued waiting for a turn self._waiting_writers = 0 @property def reader_lock(self): - """ - A context manager for acquiring a shared (reader) lock. - - Example: - async with rwlock.reader_lock: - # read-only access - """ return _ReaderLock(self) @property def writer_lock(self): - """ - A context manager for acquiring an exclusive (writer) lock. - - Example: - async with rwlock.writer_lock: - # exclusive access - """ return _WriterLock(self) async def acquire_reader(self): - async with self._lock: - # Wait until there is no active writer or waiting writer - # to ensure fairness. - while self._writer_active or self._waiting_writers > 0: - await self._cond.wait() + async with self._cond: + await self._cond.wait_for( + lambda: not self._writer_active and self._waiting_writers == 0 + ) self._readers += 1 async def release_reader(self): - async with self._lock: + async with self._cond: self._readers -= 1 - # If this was the last reader, wake up anyone waiting - # (potentially a writer or new readers). if self._readers == 0: self._cond.notify_all() async def acquire_writer(self): - async with self._lock: - # Increment the count of writers waiting + async with self._cond: self._waiting_writers += 1 try: - # Wait while either a writer is active or readers are present - while self._writer_active or self._readers > 0: - await self._cond.wait() + await self._cond.wait_for( + lambda: not self._writer_active and self._readers == 0 + ) self._writer_active = True finally: - # Decrement waiting writers only after we've acquired the writer lock self._waiting_writers -= 1 self._cond.notify_all() async def release_writer(self): - async with self._lock: + async with self._cond: self._writer_active = False - # 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 + def is_locked(self): + return self._writer_active or self._readers > 0 + + +class RWCondition: + """Condition variable built on an RWLock. Copies the CPython implementation: + https://github.com/python/cpython/blob/446edda20919447fdc8b5a43f2f2ae686df82e6a/Lib/asyncio/locks.py#L219 + """ + + def __init__(self, rwlock: RWLock | None = None): + self._rwlock = rwlock if rwlock is not None else RWLock() + self._waiters = collections.deque() + + @property + def reader_lock(self): + return _ReaderLock(self._rwlock) + + @property + def writer_lock(self): + return _WriterLock(self._rwlock) + + async def acquire_reader(self): + await self._rwlock.acquire_reader() + + async def release_reader(self): + await self._rwlock.release_reader() + + async def acquire_writer(self): + await self._rwlock.acquire_writer() + + async def release_writer(self): + await self._rwlock.release_writer() + + def is_locked(self): + return self._rwlock.is_locked() + + def _notify(self, n): + idx = 0 + for fut in self._waiters: + if idx >= n: + break + if not fut.done(): + idx += 1 + fut.set_result(False) + + def notify_all(self): + self._notify(len(self._waiters)) + + async def _wait(self, release_fn, acquire_fn): + fut = asyncio.get_running_loop().create_future() + await release_fn() + try: + try: + self._waiters.append(fut) + try: + await fut + return True + finally: + self._waiters.remove(fut) + finally: + err = None + while True: + try: + await acquire_fn() + break + except asyncio.CancelledError as e: + err = e + if err is not None: + try: + raise err + finally: + err = None + except BaseException: + self._notify(1) + raise + + async def wait_reader(self): + return await self._wait( + self._rwlock.release_reader, self._rwlock.acquire_reader + ) + + async def wait_writer(self): + return await self._wait( + self._rwlock.release_writer, self._rwlock.acquire_writer + ) + + async def wait_for_reader(self, predicate): + result = predicate() + while not result: + await self.wait_reader() + result = predicate() + return result + + async def wait_for_writer(self, predicate): + result = predicate() + while not result: + await self.wait_writer() + result = predicate() + return result class _ReaderLock: From fa11d6f00ca80f8673eff940b1af06666752bebc Mon Sep 17 00:00:00 2001 From: Lawrence Wu Date: Fri, 17 Apr 2026 18:45:09 +0000 Subject: [PATCH 4/4] fix lint --- python/sglang/srt/entrypoints/http_server.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/sglang/srt/entrypoints/http_server.py b/python/sglang/srt/entrypoints/http_server.py index de55dbf91803..d51f467800de 100644 --- a/python/sglang/srt/entrypoints/http_server.py +++ b/python/sglang/srt/entrypoints/http_server.py @@ -145,8 +145,8 @@ from sglang.srt.managers.multi_tokenizer_mixin import ( MultiTokenizerRouter, TokenizerWorker, - get_main_process_id, create_shared_pause_flag, + get_main_process_id, monkey_patch_uvicorn_multiprocessing, read_from_shared_memory, write_data_for_multi_tokenizer,