From 1d6d032f70f1f00b14242cd2b49be7118e3b086c Mon Sep 17 00:00:00 2001 From: maocheng23 Date: Tue, 7 Apr 2026 19:17:21 -0700 Subject: [PATCH] [fix] Fix deadlock and TOCTOU race in pause-aware weight update locking Fix writer lock deadlock in update_weights_from_ipc and post_process_weights during pause_generation, and close a TOCTOU race in all four weight update paths (update_weights_from_distributed, update_weights_from_tensor, update_weights_from_ipc, post_process_weights). When paused, hold is_pause_cond during the weight update so resume_generation cannot race between the check and the update. When not paused, acquire the writer lock as before. Combines #22211 (deadlock fix) and #22304 (TOCTOU fix). Co-Authored-By: Claude Opus 4.6 (1M context) --- .../managers/tokenizer_communicator_mixin.py | 53 ++++++++++++------- 1 file changed, 33 insertions(+), 20 deletions(-) diff --git a/python/sglang/srt/managers/tokenizer_communicator_mixin.py b/python/sglang/srt/managers/tokenizer_communicator_mixin.py index 50a1c4045adb..3dce75154f58 100644 --- a/python/sglang/srt/managers/tokenizer_communicator_mixin.py +++ b/python/sglang/srt/managers/tokenizer_communicator_mixin.py @@ -6,7 +6,6 @@ import time import uuid from collections import deque -from contextlib import nullcontext from typing import ( TYPE_CHECKING, Any, @@ -522,15 +521,15 @@ async def update_weights_from_distributed( if obj.abort_all_requests: self.abort_request(abort_all=True) - # Immediately update the weights if the engine is in paused state + # 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) - lock_context = ( - self.model_update_lock.writer_lock if not is_paused else nullcontext() - ) - async with lock_context: - 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) success, message = _Communicator.merge_results(results) if success and obj.weight_version is not None: @@ -580,15 +579,14 @@ async def update_weights_from_tensor( 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 + if is_paused: + results = await self.update_weights_from_tensor_communicator(obj) - lock_context = ( - self.model_update_lock.writer_lock if not is_paused else nullcontext() - ) - async with lock_context: - 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) success, message = _Communicator.merge_results(results) if success and obj.weight_version is not None: @@ -610,10 +608,17 @@ 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") - # This means that weight sync cannot run while requests are in progress. - 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.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 except Exception as e: error_msg = f"IPC weight update failed: {str(e)}" logger.error(error_msg) @@ -632,9 +637,17 @@ async def post_process_weights( ) -> Tuple[bool, str]: """Trigger post-processing hooks for weights after loading (e.g., Marlin conversion).""" self.auto_create_handle_loop() - async with self.model_update_lock.writer_lock: - results = await self.post_process_weights_communicator(obj) - return _Communicator.merge_results(results) + + async with self.is_pause_cond: + is_paused = self.is_pause + if is_paused: + results = await self.post_process_weights_communicator(obj) + + if not is_paused: + async with self.model_update_lock.writer_lock: + results = await self.post_process_weights_communicator(obj) + + return _Communicator.merge_results(results) async def _unload_lora_adapter_locked( self: TokenizerManager,