[fix] Fix issues for in-flight weight updates (#14064)

Co-authored-by: 赵晨阳 <zhaochen20@outlook.com>

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()

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:
+            is_paused = self.is_pause
-                result = (await self.update_weights_from_distributed_communicator(obj))[
-                    0
-                ]
-                return result.success, result.message
 
-        # This means that weight sync
+        lock_context = (
-        # cannot run while requests are in progress.
+            self.model_update_lock.writer_lock if not is_paused else nullcontext()
-        async with self.model_update_lock.writer_lock:
+        )
+        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:
+            is_paused = self.is_pause
-                result = (await self.update_weights_from_tensor_communicator(obj))[0]
-                return result.success, result.message
 
-        # This means that weight sync
+        lock_context = (
-        # cannot run while requests are in progress.
+            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_tensor_communicator(obj))[0]
+        async with lock_context:
-            success, message = result.success, result.message
+            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}."

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:
+            is_paused = self.is_pause
-                return await self._wait_for_model_update_from_disk(obj)
 
-        if True:  # Keep this redundant check to simplify some internal code sync
+        lock_context = (
-            # Hold the lock if it is not async. This means that weight sync
+            self.model_update_lock.writer_lock if not is_paused else nullcontext()
-            # cannot run while requests are in progress.
+        )
-            async with self.model_update_lock.writer_lock:
+        async with lock_context:
-                success, message, num_paused_requests = (
+            success, message, num_paused_requests = (
-                    await self._wait_for_model_update_from_disk(obj)
+                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)