diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index baaf63e97..152d9717b 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -2449,8 +2449,19 @@ class Scheduler( # Process the results of the last batch tmp_batch, tmp_result = self.result_queue.popleft() self.process_batch_result(tmp_batch, tmp_result) - self.last_batch = None - self.cur_batch = None + + if self.last_batch and self.last_batch.forward_mode.is_extend(): + chunked_req_to_exclude = set() + if recv_req.mode == "in_place": + if self.chunked_req is not None: + chunked_req_to_exclude.add(self.chunked_req) + self.last_batch.filter_batch( + chunked_req_to_exclude=list(chunked_req_to_exclude) + ) + self.running_batch.merge_batch(self.last_batch) + + self.last_batch = None + self.cur_batch = None if recv_req.mode == "retract": self.running_batch.filter_batch() diff --git a/python/sglang/srt/managers/tokenizer_communicator_mixin.py b/python/sglang/srt/managers/tokenizer_communicator_mixin.py index edbc52526..a5626f0b7 100644 --- a/python/sglang/srt/managers/tokenizer_communicator_mixin.py +++ b/python/sglang/srt/managers/tokenizer_communicator_mixin.py @@ -6,6 +6,7 @@ import logging import time import uuid from collections import deque +from contextlib import nullcontext from typing import ( TYPE_CHECKING, Any, @@ -417,18 +418,15 @@ class TokenizerCommunicatorMixin: # Immediately update the weights if the engine is in paused state async with self.is_pause_cond: - if self.is_pause: - result = (await self.update_weights_from_distributed_communicator(obj))[ - 0 - ] - return result.success, result.message + is_paused = self.is_pause - # This means that weight sync - # cannot run while requests are in progress. - async with self.model_update_lock.writer_lock: + 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) - success, message = _Communicator.merge_results(results) + success, message = _Communicator.merge_results(results) if success and obj.weight_version is not None: self._update_weight_version_if_provided(obj.weight_version) message += f" Weight version updated to {obj.weight_version}." @@ -478,16 +476,15 @@ class TokenizerCommunicatorMixin: # Immediately update the weights if the engine is in paused state async with self.is_pause_cond: - if self.is_pause: - result = (await self.update_weights_from_tensor_communicator(obj))[0] - return result.success, result.message + is_paused = self.is_pause - # 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_tensor_communicator(obj))[0] - success, message = result.success, result.message + 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) + success, message = _Communicator.merge_results(results) if success and obj.weight_version is not None: self._update_weight_version_if_provided(obj.weight_version) message += f" Weight version updated to {obj.weight_version}." diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py index f339d2e9f..f563ab121 100644 --- a/python/sglang/srt/managers/tokenizer_manager.py +++ b/python/sglang/srt/managers/tokenizer_manager.py @@ -1209,16 +1209,15 @@ class TokenizerManager(TokenizerCommunicatorMixin, TokenizerManagerMultiItemMixi # Immediately update the weights if the engine is in paused state async with self.is_pause_cond: - if self.is_pause: - return await self._wait_for_model_update_from_disk(obj) + is_paused = self.is_pause - if True: # Keep this redundant check to simplify some internal code sync - # Hold the lock if it is not async. This means that weight sync - # cannot run while requests are in progress. - async with self.model_update_lock.writer_lock: - success, message, num_paused_requests = ( - await self._wait_for_model_update_from_disk(obj) - ) + lock_context = ( + self.model_update_lock.writer_lock if not is_paused else nullcontext() + ) + async with lock_context: + success, message, num_paused_requests = ( + await self._wait_for_model_update_from_disk(obj) + ) if success and obj.weight_version is not None: self._update_weight_version_if_provided(obj.weight_version)