diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py index a2b4b2deb2a3..67b3fa98700f 100644 --- a/python/sglang/srt/managers/io_struct.py +++ b/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,13 +2379,19 @@ 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, ) @@ -2387,6 +2399,8 @@ 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 = ( diff --git a/python/sglang/srt/managers/scheduler_components/weight_updater.py b/python/sglang/srt/managers/scheduler_components/weight_updater.py index 2576ca613904..9de20aaa6daa 100644 --- a/python/sglang/srt/managers/scheduler_components/weight_updater.py +++ b/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,7 +440,8 @@ 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): @@ -444,7 +449,14 @@ def _stash_lora_tensors(self, lora_tensors) -> None: 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]]] diff --git a/python/sglang/srt/managers/tokenizer_control_mixin.py b/python/sglang/srt/managers/tokenizer_control_mixin.py index 18ee9949fba4..bc81fca26007 100644 --- a/python/sglang/srt/managers/tokenizer_control_mixin.py +++ b/python/sglang/srt/managers/tokenizer_control_mixin.py @@ -8,6 +8,7 @@ from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple import fastapi +import msgspec from sglang.srt.managers.communicator import FanOutCommunicator from sglang.srt.managers.io_struct import ( @@ -457,31 +458,47 @@ async def destroy_weights_update_group( return FanOutCommunicator.merge_results(results) async def _weight_update_session_call( - self: TokenizerManager, communicator, obj + self: TokenizerManager, communicator, obj, staged: bool = False ) -> Tuple[bool, str]: - """Run one weight-update session RPC under the same pause-aware locking as - update_weights_from_distributed: while the engine is paused the writer lock - is already held by whoever paused it, so taking it again would deadlock.""" + """Run one weight-update RPC with pause-aware locking (a paused engine + already holds the writer lock). A staged session takes the reader lock + and runs alongside generation: deferred names are not servable and base + tensors are rejected; other sessions keep the writer lock.""" self.auto_create_handle_loop() async with self.is_pause_cond: is_paused = self.is_pause if is_paused: results = await communicator(obj) if not is_paused: - async with self.model_update_lock.writer_lock: + lock = ( + self.model_update_lock.reader_lock + if staged + else self.model_update_lock.writer_lock + ) + async with lock: results = await communicator(obj) return FanOutCommunicator.merge_results(results) + def _staged_serving_state_error(self: TokenizerManager, obj) -> Optional[str]: + # the reader lock gave up exclusivity; serving-state mutations need the writer lock + if self._weight_update_staged_session and ( + obj.flush_cache or obj.abort_all_requests or obj.weight_version is not None + ): + return "a staged adapter session cannot change serving state" + return None + async def begin_weight_update( self: TokenizerManager, obj: BeginWeightUpdateReqInput, request: Optional[fastapi.Request] = None, ) -> Tuple[bool, str]: + staged = not obj.sync_base and bool(self._pending_lora_publications) success, message = await self._weight_update_session_call( - self.begin_weight_update_communicator, obj + self.begin_weight_update_communicator, obj, staged=staged ) if success: self._weight_update_session_open = True + self._weight_update_staged_session = staged self._weight_update_pending_version = None return success, message @@ -490,15 +507,62 @@ async def end_weight_update( obj: EndWeightUpdateReqInput, request: Optional[fastapi.Request] = None, ) -> Tuple[bool, str]: + pending = dict(self._pending_lora_publications) + # fail closed: without a full manifest a lost bucket would publish an incomplete adapter + missing = set() if obj.abort else set(pending) - set(obj.expected_lora_checksums or {}) + obj.abort = obj.abort or bool(missing) success, message = await self._weight_update_session_call( - self.end_weight_update_communicator, obj + self.end_weight_update_communicator, + obj, + staged=self._weight_update_staged_session, ) self._weight_update_session_open = False - if success: + self._weight_update_staged_session = False + if success and not obj.abort: self._update_weight_version_if_provided(self._weight_update_pending_version) self._weight_update_pending_version = None + if obj.abort or not success: + await self._discard_pending_publications() + elif pending: + await self._publish_pending_adapters() + if missing: + return False, f"deferred adapters {sorted(missing)} have no checksum manifest; session aborted" return success, message + async def _discard_pending_publications(self: TokenizerManager) -> None: + """Drop deferred publications after an aborted or failed session: the + backends must forget the zeroed identities the names never served.""" + pending, self._pending_lora_publications = self._pending_lora_publications, {} + async with self.lora_update_lock: + for ref in pending.values(): + result = _merge_lora_update_results( + await self.update_lora_adapter_communicator( + UnloadLoRAAdapterReqInput( + lora_name=ref.lora_name, lora_id=ref.lora_id + ) + ) + ) + if not result.success: + logger.warning( + "Failed to discard unpublished LoRA adapter %r: %s", + ref.lora_name, + result.error_message, + ) + + async def _publish_pending_adapters(self: TokenizerManager) -> None: + """Commit deferred publications: the names become servable atomically.""" + pending, self._pending_lora_publications = self._pending_lora_publications, {} + async with self.lora_update_lock: + for name, ref in pending.items(): + await self.lora_registry.register(ref) + self.lora_ref_cache[name] = ref + try: + await self._evict_lru_over_cap() + except ValueError: + # The publish is already committed; an unevictable over-cap + # pool is a capacity problem, not a publication failure. + logger.exception("LRU eviction after LoRA publication failed") + async def update_weights_from_distributed( self: TokenizerManager, obj: UpdateWeightsFromDistributedReqInput, @@ -509,20 +573,16 @@ async def update_weights_from_distributed( get_parallel().dp_size == 1 or get_parallel().enable_dp_attention ), "dp_size must be 1 or dp attention must be enabled for update weights from distributed" + if error := self._staged_serving_state_error(obj): + return False, error if obj.abort_all_requests: self.abort_request(abort_all=True) - # 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) - - if not is_paused: - async with self.model_update_lock.writer_lock: - results = await self.update_weights_from_distributed_communicator(obj) - - success, message = FanOutCommunicator.merge_results(results) + success, message = await self._weight_update_session_call( + self.update_weights_from_distributed_communicator, + obj, + staged=self._weight_update_staged_session, + ) if success and obj.flush_cache and self.mm_processor is not None: self.mm_processor.clear_preprocess_cache() if success and obj.weight_version is not None: @@ -569,6 +629,8 @@ async def update_weights_from_tensor( get_parallel().dp_size == 1 or get_parallel().enable_dp_attention ), "dp_size must be 1 or dp attention must be enabled for update weights from tensor" + if error := self._staged_serving_state_error(obj): + return False, error if obj.abort_all_requests: self.abort_request(abort_all=True) @@ -576,16 +638,11 @@ async def update_weights_from_tensor( obj.serialized_named_tensors ) - async with self.is_pause_cond: - is_paused = self.is_pause - if is_paused: - 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 = FanOutCommunicator.merge_results(results) + success, message = await self._weight_update_session_call( + self.update_weights_from_tensor_communicator, + obj, + staged=self._weight_update_staged_session, + ) if success and obj.flush_cache and self.mm_processor is not None: self.mm_processor.clear_preprocess_cache() if success and obj.weight_version is not None: @@ -694,35 +751,8 @@ async def load_lora_adapter( await self.lora_registry.register(new_adapter) self.lora_ref_cache[obj.lora_name] = new_adapter - if self.server_args.max_loaded_loras is not None: - while ( - self.lora_registry.num_registered_loras - > self.server_args.max_loaded_loras - ): - lru_lora_name = await self.lora_registry.lru_lora_name( - exclude_pinned=True - ) - if lru_lora_name is None: - raise ValueError( - "Didn't find any LoRA adapters when trying to evict LRU LoRA adapter. " - f"LoRA registry is: {self.lora_registry._registry}" - ) - - logger.info( - f"Unloading least recently used LoRA adapter '{lru_lora_name}' " - f"(current number of adapters: {self.lora_registry.num_registered_loras}, " - f"max allowed: {self.server_args.max_loaded_loras})" - ) - - unload_result = await self._unload_lora_adapter_locked( - UnloadLoRAAdapterReqInput(lora_name=lru_lora_name) - ) - if not unload_result.success: - raise ValueError( - f"Error while unloading LRU LoRA adapter '{lru_lora_name}': " - f"{unload_result.error_message}" - ) - del result.loaded_adapters[lru_lora_name] + for evicted in await self._evict_lru_over_cap(): + del result.loaded_adapters[evicted] return result except ValueError as e: @@ -731,6 +761,40 @@ async def load_lora_adapter( error_message=str(e), ) + async def _evict_lru_over_cap(self: TokenizerManager) -> list: + """Unload LRU adapters until the registry is back under + --max-loaded-loras; returns the evicted names. Caller must hold + lora_update_lock.""" + evicted = [] + if self.server_args.max_loaded_loras is None: + return evicted + while ( + self.lora_registry.num_registered_loras > self.server_args.max_loaded_loras + ): + lru_lora_name = await self.lora_registry.lru_lora_name(exclude_pinned=True) + if lru_lora_name is None: + raise ValueError( + "Didn't find any LoRA adapters when trying to evict LRU LoRA adapter. " + f"LoRA registry is: {self.lora_registry._registry}" + ) + + logger.info( + f"Unloading least recently used LoRA adapter '{lru_lora_name}' " + f"(current number of adapters: {self.lora_registry.num_registered_loras}, " + f"max allowed: {self.server_args.max_loaded_loras})" + ) + + unload_result = await self._unload_lora_adapter_locked( + UnloadLoRAAdapterReqInput(lora_name=lru_lora_name) + ) + if not unload_result.success: + raise ValueError( + f"Error while unloading LRU LoRA adapter '{lru_lora_name}': " + f"{unload_result.error_message}" + ) + evicted.append(lru_lora_name) + return evicted + def _validate_lora_upsert_supported(self: TokenizerManager) -> None: """Upsert resolves lora_name -> lora_id through this process's registry. @@ -770,19 +834,35 @@ async def register_lora_adapter( ) async with self.lora_update_lock: self._validate_lora_upsert_supported() - new_adapter, reused = await self.lora_registry.register_or_reuse( - LoRARef( - lora_name=obj.lora_name, - lora_path="__stream__", - pinned=obj.pinned, - reloadable=False, - ), - upsert=True, + ref = LoRARef( + lora_name=obj.lora_name, + lora_path=obj.lora_path or "__stream__", + pinned=obj.pinned, + reloadable=obj.lora_path is not None, ) + if obj.lora_name in self._pending_lora_publications: + raise ValueError( + f"LoRA adapter '{obj.lora_name}' is awaiting publication; " + "commit or abort its session before registering it again." + ) + if obj.defer_publish: + if await self.lora_registry.get_lora_id(obj.lora_name) is not None: + raise ValueError( + f"defer_publish requires a fresh adapter name, but " + f"'{obj.lora_name}' is already registered: a published " + "name has readers and must not be staged over." + ) + new_adapter, reused = ref, False + else: + new_adapter, reused = await self.lora_registry.register_or_reuse( + ref, upsert=True + ) # No path to reload a streamed adapter from: eviction would lose the - # only engine-side copy, so the cap rejects new names instead. + # only engine-side copy, so the cap rejects new non-reloadable names + # instead. Reloadable names resolve over-cap by LRU eviction. if ( not reused + and not new_adapter.reloadable and self.server_args.max_loaded_loras is not None and self.lora_registry.num_registered_loras >= self.server_args.max_loaded_loras @@ -800,11 +880,16 @@ async def register_lora_adapter( ) if result.success: - if reused: - await self.lora_registry.refresh(new_adapter) + if obj.defer_publish: + self._pending_lora_publications[obj.lora_name] = new_adapter + result = msgspec.structs.replace(result, pending=True) else: - await self.lora_registry.register(new_adapter) - self.lora_ref_cache[obj.lora_name] = new_adapter + if reused: + await self.lora_registry.refresh(new_adapter) + else: + await self.lora_registry.register(new_adapter) + self.lora_ref_cache[obj.lora_name] = new_adapter + await self._evict_lru_over_cap() return result except ValueError as e: return RegisterLoRAAdapterReqOutput( @@ -836,6 +921,11 @@ async def unload_lora_adapter( "Start unload Lora adapter. Lora name=%s", obj.lora_name, ) + if obj.lora_name in self._pending_lora_publications: + raise ValueError( + f"LoRA adapter '{obj.lora_name}' is awaiting publication; " + "abort its weight-update session instead of unloading it." + ) async with self.lora_update_lock: result = await self._unload_lora_adapter_locked(obj) diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py index 07cc6ed99aba..39c49c3edcb8 100644 --- a/python/sglang/srt/managers/tokenizer_manager.py +++ b/python/sglang/srt/managers/tokenizer_manager.py @@ -656,6 +656,9 @@ def init_weight_update(self): self.is_pause = False self.is_pause_cond = asyncio.Condition() self._weight_update_session_open = False + # A staged session (adapter-only, deferred publications pending) runs + # under the reader lock, concurrent with generation. + self._weight_update_staged_session = False self._weight_update_pending_version: Optional[str] = None def init_lora(self): @@ -676,6 +679,9 @@ def init_lora(self): if get_lora().lora_paths is not None: for lora_ref in get_lora().lora_paths: self.lora_ref_cache[lora_ref.lora_name] = lora_ref + # defer_publish registrations: zeroed identities on the backends, names + # unservable until end_weight_update commits their session + self._pending_lora_publications: Dict[str, LoRARef] = {} def init_disaggregation(self, *, start_pd_bootstrap_service: bool = True): # PD Disaggregation @@ -3450,15 +3456,30 @@ async def _resolve_lora_path(self, obj: Union[GenerateReqInput, EmbeddingReqInpu unregistered_loras = await self.lora_registry.get_unregistered_loras( unique_lora_paths ) + backfill_paths = obj.lora_backfill_paths or {} for lora_path in unregistered_loras: if lora_path is None: continue - if lora_path not in self.lora_ref_cache: + if lora_path in self._pending_lora_publications: + # a staged session is streaming this name; a disk load now would give it two identities raise ValueError( - f"Got LoRA adapter that has never been loaded: {lora_path}\n" - f"All loaded adapters: {self.lora_ref_cache.keys()}." + f"LoRA adapter '{lora_path}' is awaiting publication and " + "cannot be served or backfilled yet." ) + if lora_path not in self.lora_ref_cache: + # a request-carried path lets a fresh or restarted engine serve any published version + if lora_path in backfill_paths: + self.lora_ref_cache[lora_path] = LoRARef( + lora_name=lora_path, + lora_path=backfill_paths[lora_path], + pinned=False, + ) + else: + raise ValueError( + f"Got LoRA adapter that has never been loaded: {lora_path}\n" + f"All loaded adapters: {self.lora_ref_cache.keys()}." + ) new_lora_ref = self.lora_ref_cache[lora_path] if not new_lora_ref.reloadable: diff --git a/test/registered/unit/lora/test_lora_staged_publication.py b/test/registered/unit/lora/test_lora_staged_publication.py new file mode 100644 index 000000000000..ded5dd3b4549 --- /dev/null +++ b/test/registered/unit/lora/test_lora_staged_publication.py @@ -0,0 +1,303 @@ +"""Staged LoRA publication: fresh versions stream and commit alongside +generation, while base and fixed-name sessions keep the writer lock.""" + +import asyncio +import unittest +from types import SimpleNamespace +from unittest.mock import AsyncMock, Mock + +import torch + +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import maybe_stub_sgl_kernel + +maybe_stub_sgl_kernel() + +from sglang.srt.lora.lora_registry import LoRARef, LoRARegistry +from sglang.srt.managers.io_struct import ( + BeginWeightUpdateReqInput, + EndWeightUpdateReqInput, + GenerateReqInput, + LoRAUpdateOutput, + RegisterLoRAAdapterReqInput, + UnloadLoRAAdapterReqInput, + UpdateWeightsFromDistributedReqInput, +) +from sglang.srt.managers.scheduler_components.weight_updater import ( + SchedulerWeightUpdaterManager, +) +from sglang.srt.managers.tokenizer_manager import TokenizerManager +from sglang.srt.runtime_context import get_context + +register_cpu_ci(est_time=10, suite="base-a-test-cpu") + + +def _reply(success=True): + return SimpleNamespace(success=success, message="ok" if success else "worker failed") + + +def _manager(): + tm = TokenizerManager.__new__(TokenizerManager) + tm.server_args = SimpleNamespace(max_loaded_loras=None) + tm.elastic_worker_count = 1 + tm.init_weight_update() + tm.auto_create_handle_loop = Mock() + tm._validate_lora_upsert_supported = Mock() + tm.record_config_updates = Mock() + tm.abort_request = Mock() + tm.mm_processor = None + tm.lora_update_lock = asyncio.Lock() + tm.lora_registry = LoRARegistry( + [LoRARef(lora_name=name, reloadable=False) for name in ("A@1", "B", "C")] + ) + tm.lora_ref_cache = {} + tm._pending_lora_publications = {} + tm.update_lora_adapter_communicator = AsyncMock( + side_effect=lambda obj: [LoRAUpdateOutput(success=True)] + ) + tm.begin_weight_update_communicator = AsyncMock(return_value=[_reply()]) + tm.end_weight_update_communicator = AsyncMock(return_value=[_reply()]) + tm.update_weights_from_distributed_communicator = AsyncMock(return_value=[_reply()]) + return tm + + +def _begin(**kwargs): + return BeginWeightUpdateReqInput(**(dict(sync_base=False) | kwargs)) + + +def _bucket(**kwargs): + values = dict( + names=["A@2:model.layers.0.self_attn.q_proj.lora_A.weight"], + dtypes=["float32"], + shapes=[[1]], + flush_cache=False, + ) + return UpdateWeightsFromDistributedReqInput(**(values | kwargs)) + + +def _end(**kwargs): + values = dict(expected_lora_checksums={"A@2": {"q_proj": "hash"}}) + return EndWeightUpdateReqInput(**(values | kwargs)) + + +class TestStagedPublication(unittest.IsolatedAsyncioTestCase): + async def asyncSetUp(self): + config = get_context().override_server_args( + enable_lora=True, + pp_size=1, + dp_size=1, + enable_dp_attention=False, + checkpoint_engine_wait_weights_before_ready=False, + ) + config.install() + self.addCleanup(config.restore) + self.tm = _manager() + result = await self.tm.register_lora_adapter( + RegisterLoRAAdapterReqInput( + lora_name="A@2", config_dict={"r": 8}, defer_publish=True + ) + ) + self.assertTrue(result.success) + self.assertTrue(result.pending) + self.assertIn("A@2", self.tm._pending_lora_publications) + + async def test_deferred_name_is_unservable_until_commit(self): + tm = self.tm + old_ids = await tm.lora_registry.acquire(["A@1", "B", "C"]) + self.assertIsNone(await tm.lora_registry.get_lora_id("A@2")) + self.assertTrue((await tm.begin_weight_update(_begin()))[0]) + self.assertTrue((await tm.update_weights_from_distributed(_bucket()))[0]) + self.assertIsNone(await tm.lora_registry.get_lora_id("A@2")) + self.assertTrue((await tm.end_weight_update(_end()))[0]) + new_id = await tm.lora_registry.acquire("A@2") + self.assertNotIn(new_id, old_ids) + # Untouched adapters keep their identities through the publication. + self.assertEqual(await tm.lora_registry.get_lora_id("A@1"), old_ids[0]) + await tm.lora_registry.release(new_id) + await tm.lora_registry.release(old_ids) + + async def test_staged_session_completes_under_live_readers(self): + tm = self.tm + async with tm.model_update_lock.reader_lock: + self.assertTrue( + (await asyncio.wait_for(tm.begin_weight_update(_begin()), 1))[0] + ) + self.assertTrue(tm._weight_update_staged_session) + self.assertTrue( + ( + await asyncio.wait_for( + tm.update_weights_from_distributed(_bucket()), 1 + ) + )[0] + ) + self.assertTrue((await asyncio.wait_for(tm.end_weight_update(_end()), 1))[0]) + self.assertIsNotNone(await tm.lora_registry.get_lora_id("A@2")) + + async def test_base_and_fixed_name_sessions_still_wait_for_readers(self): + for sync_base in (True, False): + tm = _manager() # no pending publications: not a staged session + async with tm.model_update_lock.reader_lock: + task = asyncio.create_task( + tm.begin_weight_update(BeginWeightUpdateReqInput(sync_base=sync_base)) + ) + await asyncio.sleep(0) + self.assertFalse(task.done()) + tm.begin_weight_update_communicator.assert_not_awaited() + self.assertTrue((await asyncio.wait_for(task, 1))[0]) + + async def test_base_sync_with_pending_registration_keeps_writer_lock(self): + tm = self.tm + async with tm.model_update_lock.reader_lock: + task = asyncio.create_task( + tm.begin_weight_update(BeginWeightUpdateReqInput(sync_base=True)) + ) + await asyncio.sleep(0) + self.assertFalse(task.done()) + self.assertTrue((await asyncio.wait_for(task, 1))[0]) + self.assertFalse(tm._weight_update_staged_session) + + async def test_defer_publish_requires_a_fresh_name(self): + for taken in ("A@1", "A@2"): # published name; pending name + result = await self.tm.register_lora_adapter( + RegisterLoRAAdapterReqInput( + lora_name=taken, config_dict={}, defer_publish=True + ) + ) + self.assertFalse(result.success) + + async def test_pending_name_cannot_be_reregistered_or_unloaded(self): + result = await self.tm.register_lora_adapter( + RegisterLoRAAdapterReqInput(lora_name="A@2", config_dict={}) + ) + self.assertFalse(result.success) + result = await self.tm.unload_lora_adapter( + UnloadLoRAAdapterReqInput(lora_name="A@2") + ) + self.assertFalse(result.success) + + async def test_staged_bucket_cannot_change_serving_state(self): + tm = self.tm + await tm.begin_weight_update(_begin()) + for kwargs in ( + {"flush_cache": True}, + {"abort_all_requests": True}, + {"weight_version": "2"}, + ): + self.assertFalse( + (await tm.update_weights_from_distributed(_bucket(**kwargs)))[0] + ) + tm.abort_request.assert_not_called() + tm.update_weights_from_distributed_communicator.assert_not_awaited() + + async def test_any_worker_failure_discards_the_publication(self): + tm = self.tm + await tm.begin_weight_update(_begin()) + tm.end_weight_update_communicator.return_value = [_reply(), _reply(False)] + tm.update_lora_adapter_communicator.reset_mock() + self.assertFalse((await tm.end_weight_update(_end()))[0]) + self.assertIsNone(await tm.lora_registry.get_lora_id("A@2")) + self.assertEqual(tm._pending_lora_publications, {}) + # The backends held a zeroed identity for the name; it must be dropped. + discarded = tm.update_lora_adapter_communicator.await_args.args[0] + self.assertIsInstance(discarded, UnloadLoRAAdapterReqInput) + self.assertEqual(discarded.lora_name, "A@2") + + async def test_missing_manifest_fails_closed(self): + tm = self.tm + await tm.begin_weight_update(_begin()) + success, message = await tm.end_weight_update( + _end(expected_lora_checksums=None) + ) + self.assertFalse(success) + self.assertIn("A@2", message) + sent = tm.end_weight_update_communicator.await_args.args[0] + self.assertTrue(sent.abort) + self.assertIsNone(await tm.lora_registry.get_lora_id("A@2")) + + async def test_abort_never_publishes(self): + tm = self.tm + await tm.begin_weight_update(_begin()) + self.assertTrue((await tm.end_weight_update(_end(abort=True)))[0]) + self.assertIsNone(await tm.lora_registry.get_lora_id("A@2")) + self.assertEqual(tm._pending_lora_publications, {}) + + async def test_version_is_recorded_only_at_successful_end(self): + tm = _manager() + self.assertTrue( + (await tm.begin_weight_update(BeginWeightUpdateReqInput(sync_base=True)))[0] + ) + self.assertTrue( + (await tm.update_weights_from_distributed(_bucket(weight_version="2")))[0] + ) + tm.record_config_updates.assert_not_called() + self.assertTrue((await tm.end_weight_update(EndWeightUpdateReqInput()))[0]) + tm.record_config_updates.assert_called_once_with( + "tokenizer.weight_version", weight_version="2" + ) + + +class TestRequestCarriedBackfill(unittest.IsolatedAsyncioTestCase): + async def asyncSetUp(self): + config = get_context().override_server_args( + enable_lora=True, + pp_size=1, + dp_size=1, + enable_dp_attention=False, + checkpoint_engine_wait_weights_before_ready=False, + ) + config.install() + self.addCleanup(config.restore) + self.tm = _manager() + + async def load(obj): + await self.tm.lora_registry.register( + LoRARef(lora_name=obj.lora_name, lora_path=obj.lora_path) + ) + return SimpleNamespace(success=True, error_message="") + + self.tm.load_lora_adapter = AsyncMock(side_effect=load) + + async def test_backfill_path_admits_an_unknown_version(self): + obj = GenerateReqInput( + text="x", + lora_path="A@7", + lora_backfill_paths={"A@7": "/ckpt/run/sampler_weights/7"}, + ) + await self.tm._resolve_lora_path(obj) + self.assertIsNotNone(obj.lora_id) + loaded = self.tm.load_lora_adapter.await_args.args[0] + self.assertEqual( + (loaded.lora_name, loaded.lora_path), + ("A@7", "/ckpt/run/sampler_weights/7"), + ) + await self.tm.lora_registry.release(obj.lora_id) + + async def test_unknown_version_without_a_path_is_rejected(self): + with self.assertRaisesRegex(ValueError, "never been loaded"): + await self.tm._resolve_lora_path(GenerateReqInput(text="x", lora_path="A@7")) + + async def test_pending_name_is_not_servable_or_backfillable(self): + self.tm._pending_lora_publications["A@7"] = LoRARef( + lora_name="A@7", lora_path="__stream__", reloadable=False + ) + with self.assertRaisesRegex(ValueError, "awaiting publication"): + await self.tm._resolve_lora_path( + GenerateReqInput( + text="x", lora_path="A@7", lora_backfill_paths={"A@7": "/ckpt/7"} + ) + ) + + +class TestStashCopiesIPCBuckets(unittest.TestCase): + def test_stash_survives_sender_reusing_the_bucket(self): + manager = SchedulerWeightUpdaterManager.__new__(SchedulerWeightUpdaterManager) + manager.tp_worker = SimpleNamespace(ps=SimpleNamespace(tp_rank=0)) + manager._lora_stash = {} + bucket = torch.ones(4) + manager._stash_lora_tensors([("A@2:q_proj", bucket)], copy_tensors=True) + bucket.zero_() + self.assertTrue(torch.equal(manager._lora_stash["A@2"]["q_proj"], torch.ones(4))) + + +if __name__ == "__main__": + unittest.main()