[LoRA] Publish fresh streamed versions alongside generation

python/sglang/srt/managers/io_struct.py

@@ -265,6 +265,8 @@ class GenerateReqInput:
     lora_path: Optional[Union[List[Optional[str]], str]] = None
     # The uid of LoRA adaptors, should be initialized by tokenizer manager
     lora_id: Optional[Union[List[Optional[str]], str]] = None
+    # {lora_name: disk path}: backfill source when this engine does not hold the adapter
+    lora_backfill_paths: Optional[Dict[str, str]] = None
 
     # Custom logit processor for advanced sampling control. Must be a serialized instance
     # of `CustomLogitProcessor` in python/sglang/srt/sampling/custom_logit_processor.py
@@ -1133,6 +1135,8 @@ class EmbeddingReqInput:
     lora_path: Optional[Union[List[Optional[str]], str]] = None
     # The uid of LoRA adaptors, should be initialized by tokenizer manager
     lora_id: Optional[Union[List[Optional[str]], str]] = None
+    # {lora_name: disk path}: backfill source when this engine does not hold the adapter
+    lora_backfill_paths: Optional[Dict[str, str]] = None
     # Resolved embedding overrides with positions (set by tokenizer manager or score mixin).
     # Runtime type: Optional[Union[PositionalEmbeds, List[Optional[PositionalEmbeds]]]]
     positional_embed_overrides: Any = None
@@ -2049,6 +2053,8 @@ class EndWeightUpdateReqInput(BaseReq, kw_only=True):
     # {lora_name: {hf_key: sha256}}; when set, each stashed adapter is verified
     # (set equality + per-tensor checksum) before it is applied.
     expected_lora_checksums: Optional[Dict[str, Dict[str, str]]] = None
+    # discard the streamed LoRA stash and deferred publications; in-place base writes are not rolled back
+    abort: bool = False
 
 
 class EndWeightUpdateReqOutput(BaseReq, kw_only=True):
@@ -2373,20 +2379,28 @@ class RegisterLoRAAdapterReqInput(BaseReq, kw_only=True):
     # and pinning every slot would trip the anti-starvation check.
     pinned: bool = False
     lora_id: Optional[str] = None
+    # Disk artifact holding the same adapter (PEFT dir). With a path the
+    # adapter is reloadable: it may be LRU-evicted and refilled from disk.
+    lora_path: Optional[str] = None
+    # keep the name unservable until end_weight_update commits its session; fresh names only
+    defer_publish: bool = False
 
     def to_ref(self) -> LoRARef:
         return LoRARef(
             lora_id=self.lora_id,
             lora_name=self.lora_name,
-            lora_path="__stream__",
+            lora_path=self.lora_path or "__stream__",
             pinned=self.pinned,
+            reloadable=self.lora_path is not None,
         )
 
 
 class LoRAUpdateOutput(BaseReq, kw_only=True):
     success: bool
     error_message: Optional[str] = None
     loaded_adapters: Optional[Dict[str, Union[str, LoRARef]]] = None
+    # acks a defer_publish registration; the trainer fails closed without it
+    pending: bool = False
 
 
 LoadLoRAAdapterReqOutput = UnloadLoRAAdapterReqOutput = RegisterLoRAAdapterReqOutput = (

python/sglang/srt/managers/scheduler_components/weight_updater.py

@@ -333,7 +333,7 @@ def update_weights_from_tensor(self, recv_req: UpdateWeightsFromTensorReqInput):
             else:
                 base_tensors, lora_tensors = _split_lora_named_tensors(named_tensors)
                 named_tensors = base_tensors
-            self._stash_lora_tensors(lora_tensors)
+            self._stash_lora_tensors(lora_tensors, copy_tensors=True)
             success, message = True, "Success"
             if base_tensors:
                 assert (
@@ -423,9 +423,13 @@ def end_weight_update(self, recv_req: EndWeightUpdateReqInput):
             run_post_load = not self._weight_update_loaded
             for _, runner in self.get_model_runners(self._weight_update_selector):
                 runner.end_weight_update(run_post_load=run_post_load)
-        success, message = self._apply_lora_stash(recv_req.expected_lora_checksums)
+        if recv_req.abort:
+            self._lora_stash = {}
+            success, message = True, "Aborted: streamed adapters discarded"
+        else:
+            success, message = self._apply_lora_stash(recv_req.expected_lora_checksums)
         self._weight_update_in_progress = False
-        if success:
+        if success and not recv_req.abort:
             self.record_weight_version_after_update(self._weight_update_pending_version)
         self._weight_update_pending_version = None
         torch.distributed.barrier(group=self.tp_cpu_group)
@@ -436,15 +440,23 @@ def forget_lora_adapter(self, lora_name: str) -> None:
         is a new adapter identity and may stream a different tensor set."""
         self._lora_applied_names.pop(lora_name, None)
 
-    def _stash_lora_tensors(self, lora_tensors) -> None:
+    def _stash_lora_tensors(self, lora_tensors, *, copy_tensors: bool = False) -> None:
+        copied_devices = set()
         for prefixed_name, tensor in lora_tensors:
             lora_name, hf_key = prefixed_name.split(":", 1)
             if isinstance(tensor, LocalSerializedTensor):
                 tensor = tensor.get(self.tp_worker.ps.tp_rank)
             assert isinstance(
                 tensor, torch.Tensor
             ), f"streamed LoRA tensor {prefixed_name!r} must arrive as a plain tensor"
+            if copy_tensors:
+                # the stash outlives this RPC, but an IPC sender may reuse the bucket once it replies
+                tensor = tensor.clone()
+                if tensor.is_cuda:
+                    copied_devices.add(tensor.device)
             self._lora_stash.setdefault(lora_name, {})[hf_key] = tensor
+        for device in copied_devices:
+            torch.cuda.current_stream(device).synchronize()
 
     def _apply_lora_stash(
         self, expected_checksums: Optional[Dict[str, Dict[str, str]]]