minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[tiny] Fix TOCTOU race in pause-aware weight update locking (#22304)

Browse Source 
Co-authored-by: maocheng23 <maocheng@berkeley.edu>
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
Liangsheng Yin
2026-04-07 18:54:28 -07:00
committed by GitHub
co-authored by maocheng23 Claude Opus 4.6
parent eca62ab8f4
commit 1c5c6dad5e
1 changed files with 18 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/managers/tokenizer_communicator_mixin.py
+18 -23
View File
@@ -6,7 +6,6 @@ import logging
import time import time
import uuid import uuid
from collections import deque from collections import deque
from contextlib import nullcontext
from typing import ( from typing import (
TYPE_CHECKING, TYPE_CHECKING,
Any, Any,
@@ -624,15 +623,15 @@ class TokenizerCommunicatorMixin:
if obj.abort_all_requests: if obj.abort_all_requests:
self.abort_request(abort_all=True) 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: async with self.is_pause_cond:
is_paused = self.is_pause is_paused = self.is_pause
if is_paused:
results = await self.update_weights_from_distributed_communicator(obj)
lock_context = ( if not is_paused:
self.model_update_lock.writer_lock if not is_paused else nullcontext() async with self.model_update_lock.writer_lock:
) results = await self.update_weights_from_distributed_communicator(obj)
async with lock_context:
results = await self.update_weights_from_distributed_communicator(obj)
success, message = _Communicator.merge_results(results) success, message = _Communicator.merge_results(results)
if success and obj.weight_version is not None: if success and obj.weight_version is not None:
@@ -682,15 +681,14 @@ class TokenizerCommunicatorMixin:
if obj.abort_all_requests: if obj.abort_all_requests:
self.abort_request(abort_all=True) self.abort_request(abort_all=True)
# Immediately update the weights if the engine is in paused state
async with self.is_pause_cond: async with self.is_pause_cond:
is_paused = self.is_pause is_paused = self.is_pause
if is_paused:
results = await self.update_weights_from_tensor_communicator(obj)
lock_context = ( if not is_paused:
self.model_update_lock.writer_lock if not is_paused else nullcontext() async with self.model_update_lock.writer_lock:
) results = await self.update_weights_from_tensor_communicator(obj)
async with lock_context:
results = await self.update_weights_from_tensor_communicator(obj)
success, message = _Communicator.merge_results(results) success, message = _Communicator.merge_results(results)
if success and obj.weight_version is not None: if success and obj.weight_version is not None:
@@ -713,19 +711,16 @@ class TokenizerCommunicatorMixin:
), "dp_size must be 1 or dp attention must be enabled for update weights from IPC" ), "dp_size must be 1 or dp attention must be enabled for update weights from IPC"
logger.info("Starting IPC weight update") logger.info("Starting IPC weight update")
# Skip the writer lock when paused: readers are blocked on
# is_pause_cond so no concurrent inference can race, and
# waiting for the writer lock would deadlock because existing
# readers are stuck waiting on the paused scheduler.
async with self.is_pause_cond: async with self.is_pause_cond:
is_paused = self.is_pause is_paused = self.is_pause
if is_paused:
result = (await self.update_weights_from_ipc_communicator(obj))[0]
success, message = result.success, result.message
lock_context = ( if not is_paused:
self.model_update_lock.writer_lock if not is_paused else nullcontext() async with self.model_update_lock.writer_lock:
) result = (await self.update_weights_from_ipc_communicator(obj))[0]
async with lock_context: success, message = result.success, result.message
result = (await self.update_weights_from_ipc_communicator(obj))[0]
success, message = result.success, result.message
except Exception as e: except Exception as e:
error_msg = f"IPC weight update failed: {str(e)}" error_msg = f"IPC weight update failed: {str(e)}"
logger.error(error_msg) logger.error(error_msg)