diff --git a/python/sglang/srt/entrypoints/http_server.py b/python/sglang/srt/entrypoints/http_server.py index 5b7de28d519c..4b4a34674ef4 100644 --- a/python/sglang/srt/entrypoints/http_server.py +++ b/python/sglang/srt/entrypoints/http_server.py @@ -1558,6 +1558,7 @@ async def load_lora_adapter( @app.api_route("/load_lora_adapter_from_tensors", methods=["POST"]) +@auth_level(AuthLevel.ADMIN_OPTIONAL) async def load_lora_adapter_from_tensors( obj: Annotated[LoadLoRAAdapterFromTensorsReqInput, Body()], request: Request ): @@ -1570,6 +1571,7 @@ async def load_lora_adapter_from_tensors( @app.api_route("/load_lora_adapter_from_distributed", methods=["POST"]) +@auth_level(AuthLevel.ADMIN_OPTIONAL) async def load_lora_adapter_from_distributed( obj: Annotated[LoadLoRAAdapterFromDistributedReqInput, Body()], request: Request ): diff --git a/python/sglang/srt/lora/lora_registry.py b/python/sglang/srt/lora/lora_registry.py index 8ee933a8c402..ae6242b1ecbf 100644 --- a/python/sglang/srt/lora/lora_registry.py +++ b/python/sglang/srt/lora/lora_registry.py @@ -38,6 +38,12 @@ class LoRARef(msgspec.Struct, frozen=True, array_like=True): lora_name: Optional[str] = None lora_path: Optional[str] = None pinned: Optional[bool] = None + # False for adapters whose weights arrived over the wire (lora_path + # "__distributed__" / "__tensor__"): there is no disk artifact to reload + # from, so they must never be LRU-evicted nor implicitly reloaded. + # Trailing field with a default keeps the array_like wire format + # compatible with refs encoded before this field existed. + reloadable: bool = True def __post_init__(self): if self.lora_id is None: @@ -184,24 +190,28 @@ def _lookup(name: str) -> str: self._registry.move_to_end(name) return lora_ref.lora_id + # Lookup and increment must be one atomic admission step under the + # writer lock: with the increment outside, an unload can slip between + # them — unregister the adapter, observe the counter at zero, delete + # the counter — leaving this request to KeyError or to run on an + # adapter that is already gone. wait_for_unload never awaits inside + # this lock, so holding it across the increment cannot deadlock. if isinstance(lora_name, str): async with self._registry_lock.writer_lock: lora_id = _lookup(lora_name) - - await self._counters[lora_id].increment(notify_all=False) + await self._counters[lora_id].increment(notify_all=False) return lora_id elif isinstance(lora_name, list): async with self._registry_lock.writer_lock: lora_ids = [_lookup(name) for name in lora_name] - - # Increment the counters only after all IDs are looked up. - await asyncio.gather( - *[ - self._counters[id].increment(notify_all=False) - for id in lora_ids - if id is not None - ] - ) + # Increment the counters only after all IDs are looked up. + await asyncio.gather( + *[ + self._counters[id].increment(notify_all=False) + for id in lora_ids + if id is not None + ] + ) return lora_ids else: raise TypeError("lora_name must be either a string or a list of strings.") @@ -266,14 +276,15 @@ async def lru_lora_name(self, exclude_pinned=False): If exclude_pinned is True, then return the LRU LoRA adapter that isn't pinned. """ async with self._registry_lock.reader_lock: - if not exclude_pinned: - return next(iter(self._registry), None) - for lora_name, lora_ref in self._registry.items(): - if not lora_ref.pinned: - return lora_name - else: - return None + # Evicting a non-reloadable (wire-loaded) adapter is always + # destructive: there is no artifact to reload it from. + if not lora_ref.reloadable: + continue + if exclude_pinned and lora_ref.pinned: + continue + return lora_name + return None def _register_adapter(self, lora_ref: LoRARef): """ diff --git a/python/sglang/srt/managers/tokenizer_control_mixin.py b/python/sglang/srt/managers/tokenizer_control_mixin.py index 1d3ff8d6ddd2..9894a6605ecb 100644 --- a/python/sglang/srt/managers/tokenizer_control_mixin.py +++ b/python/sglang/srt/managers/tokenizer_control_mixin.py @@ -735,6 +735,7 @@ async def load_lora_adapter_from_tensors( lora_name=obj.lora_name, lora_path="__tensor__", pinned=obj.pinned, + reloadable=False, ) obj.lora_id = new_adapter.lora_id result = (await self.update_lora_adapter_communicator(obj))[0] @@ -811,6 +812,7 @@ async def load_lora_adapter_from_distributed( lora_name=obj.lora_name, lora_path="__distributed__", pinned=obj.pinned, + reloadable=False, ), upsert=obj.upsert, ) @@ -886,7 +888,14 @@ async def unload_lora_adapter( ) async with self.lora_update_lock: - return await self._unload_lora_adapter_locked(obj) + result = await self._unload_lora_adapter_locked(obj) + # Explicit unload is a DELETE: drop the reload-catalog entry too. + # The max_loaded_loras LRU loop calls _unload_lora_adapter_locked + # directly — an EVICT — and must keep the entry so disk-backed + # adapters can be implicitly reloaded later. + if result.success: + self.lora_ref_cache.pop(obj.lora_name, None) + return result except ValueError as e: return UnloadLoRAAdapterReqOutput(success=False, error_message=str(e)) diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py index 0b435574b208..a6b83a7b07e8 100644 --- a/python/sglang/srt/managers/tokenizer_manager.py +++ b/python/sglang/srt/managers/tokenizer_manager.py @@ -188,6 +188,13 @@ class ReqState: # request as aborted instead of sending it. abort_before_dispatch: bool = False + # The LoRA lease for this request has already been released. Every terminal + # path funnels through TokenizerManager._finalize_lora_lease, which uses + # this flag to stay idempotent: a leaked release hangs unload's + # wait_for_zero forever, and a double release drives the ConcurrentCounter + # negative, which hangs it just the same (it waits for exactly zero). + lora_lease_released: bool = False + # For streaming output last_output_offset: int = 0 @@ -1473,6 +1480,25 @@ def _coalesce_streaming_chunks( out["meta_info"] = meta_info return out + def _finalize_lora_lease(self, state: Optional[ReqState]) -> None: + """Release the request's LoRA lease exactly once, however it terminates: + normal finish, scheduler abort echo (queued / tokenizer-held / disagg), + status-code abort, or a failed dispatch. A request whose LoRA was never + acquired (pre-acquire failure leaves lora_id unset) has no lease to + release.""" + if state is None or state.lora_lease_released: + return + # lora_path / lora_id are declared on both input structs with default + # None (= base model / lease never acquired), so plain None checks are + # the contract. State-derived checks come first: a request without a + # lease has nothing to release, whatever the server config says. + if state.obj.lora_path is None or state.obj.lora_id is None: + return + if not self.enable_lora: + return + state.lora_lease_released = True + asyncio.create_task(self.lora_registry.release(state.obj.lora_id)) + async def _handle_abort_finish_reason( self, out: dict, @@ -1505,9 +1531,10 @@ async def _handle_abort_finish_reason( if state.obj.rid in self.rid_to_state: del self.rid_to_state[state.obj.rid] - # Mark ongoing LoRA request as finished. - if self.enable_lora and state.obj.lora_path: - await self.lora_registry.release(state.obj.lora_id) + # Status-code aborts also arrive through _handle_batch_output, which + # already released the lease and deleted the state — the finalizer's + # idempotency is what prevents the counter from going negative here. + self._finalize_lora_lease(state) if not is_stream: raise fastapi.HTTPException( status_code=finish_reason["status_code"], @@ -2220,8 +2247,7 @@ async def _handle_batch_output( del self.rid_to_state[rid] # Mark ongoing LoRA request as finished. - if self.enable_lora and state.obj.lora_path: - asyncio.create_task(self.lora_registry.release(state.obj.lora_id)) + self._finalize_lora_lease(state) if out_dict is not None: state.out_list.append(out_dict) @@ -2879,6 +2905,10 @@ def _handle_abort_req(self, recv_obj: AbortReq): "output_ids": output_ids, "meta_info": meta_info, } + # This abort echo is the scheduler-side terminal ACK for queued, + # tokenizer-held, and disagg-retracted requests — none of them reach + # _handle_batch_output, so this is their only lease-release point. + self._finalize_lora_lease(state) del self.rid_to_state[recv_obj.rid] state.out_list.append(out) @@ -3010,8 +3040,15 @@ async def _resolve_lora_path(self, obj: Union[GenerateReqInput, EmbeddingReqInpu f"All loaded adapters: {self.lora_ref_cache.keys()}." ) - logger.info(f"Reloading evicted adapter: {lora_path}") new_lora_ref = self.lora_ref_cache[lora_path] + if not new_lora_ref.reloadable: + raise ValueError( + f"LoRA adapter '{lora_path}' was loaded over the wire " + f"(lora_path={new_lora_ref.lora_path!r}) and has no disk " + "artifact to reload from; it must be re-pushed by the " + "trainer instead of implicitly reloaded." + ) + logger.info(f"Reloading evicted adapter: {lora_path}") load_result = await self.load_lora_adapter( LoadLoRAAdapterReqInput( lora_name=new_lora_ref.lora_name, @@ -3092,7 +3129,12 @@ def _discard_pending_req_states(self, obj): else: rids = obj.rid for rid in rids: - self.rid_to_state.pop(rid, None) + state = self.rid_to_state.pop(rid, None) + # A failed/partial dispatch may never produce a scheduler terminal + # for this rid, and once the state is dropped a late terminal has no + # release point either — so release here. Pre-acquire failures are + # covered by the finalizer's lora_id check. + self._finalize_lora_lease(state) def _should_dispatch_to_encoder( self, obj: Union[GenerateReqInput, EmbeddingReqInput] diff --git a/test/registered/unit/lora/test_lora_lease.py b/test/registered/unit/lora/test_lora_lease.py new file mode 100644 index 000000000000..8ab54f161b0b --- /dev/null +++ b/test/registered/unit/lora/test_lora_lease.py @@ -0,0 +1,233 @@ +"""Unit tests for the exactly-once LoRA lease lifecycle. + +A leaked lease leaves the adapter's in-flight counter above zero forever, so +unload's wait_for_zero never returns; a double release drives the counter +negative, which hangs it just the same (it waits for exactly zero). Covers: + + * TokenizerManager._finalize_lora_lease is the single, idempotent release + owner for every terminal path + * pre-acquire failures (lora_id never set) release nothing + * LoRARegistry.acquire: lookup + counter increment are one atomic admission + step under the writer lock, so acquire racing an unload fails cleanly + instead of touching a deleted counter + * LoRARegistry.lru_lora_name never picks non-reloadable (wire-loaded) refs + * explicit unload drops the lora_ref_cache entry (DELETE) while the locked + helper alone (the max_loaded_loras EVICT path) keeps it +""" + +import asyncio +import unittest +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock, Mock + +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase, maybe_stub_sgl_kernel + +maybe_stub_sgl_kernel() + +from sglang.srt.lora.lora_registry import LoRARef, LoRARegistry +from sglang.srt.managers.io_struct import UnloadLoRAAdapterReqInput +from sglang.srt.managers.tokenizer_manager import ReqState, TokenizerManager + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + + +def _make_tm(enable_lora: bool = True) -> TokenizerManager: + """TokenizerManager with only the fields the lease path reads.""" + tm = TokenizerManager.__new__(TokenizerManager) + tm.enable_lora = enable_lora + tm.lora_registry = MagicMock() + tm.lora_registry.release = AsyncMock() + return tm + + +def _make_state(lora_path="adapter-a", lora_id="id-1") -> ReqState: + obj = SimpleNamespace(lora_path=lora_path, lora_id=lora_id, rid="rid-1") + return ReqState([], False, asyncio.Event(), obj, Mock()) + + +class TestFinalizeLoraLease(CustomTestCase): + def _finalize(self, tm, state, times=1): + async def run(): + for _ in range(times): + tm._finalize_lora_lease(state) + # Let the create_task'd release coroutines run. + await asyncio.sleep(0) + + asyncio.run(run()) + + def test_release_exactly_once_on_double_finalize(self): + tm = _make_tm() + state = _make_state() + self._finalize(tm, state, times=3) + tm.lora_registry.release.assert_awaited_once_with("id-1") + self.assertTrue(state.lora_lease_released) + + def test_pre_acquire_failure_releases_nothing(self): + # _init_req_state runs before the LoRA acquire; a failure in between + # leaves lora_id unset and there is no lease to release. + tm = _make_tm() + state = _make_state(lora_id=None) + self._finalize(tm, state) + tm.lora_registry.release.assert_not_awaited() + self.assertFalse(state.lora_lease_released) + + def test_no_lora_path_is_noop(self): + tm = _make_tm() + state = _make_state(lora_path=None) + self._finalize(tm, state) + tm.lora_registry.release.assert_not_awaited() + + def test_lora_disabled_is_noop(self): + tm = _make_tm(enable_lora=False) + state = _make_state() + self._finalize(tm, state) + tm.lora_registry.release.assert_not_awaited() + + def test_none_state_is_noop(self): + tm = _make_tm() + + async def run(): + tm._finalize_lora_lease(None) + + asyncio.run(run()) + tm.lora_registry.release.assert_not_awaited() + + +class TestAcquireAtomicity(CustomTestCase): + def test_acquire_after_unregister_raises_value_error(self): + async def run(): + registry = LoRARegistry() + ref = LoRARef(lora_name="a", lora_path="/x") + await registry.register(ref) + await registry.unregister("a") + with self.assertRaises(ValueError): + await registry.acquire("a") + + asyncio.run(run()) + + def test_unload_waits_for_inflight_lease_then_completes(self): + async def run(): + registry = LoRARegistry() + ref = LoRARef(lora_name="a", lora_path="/x") + await registry.register(ref) + lora_id = await registry.acquire("a") + await registry.unregister("a") + + wait_task = asyncio.create_task(registry.wait_for_unload(lora_id)) + await asyncio.sleep(0.01) + self.assertFalse(wait_task.done(), "unload must wait for the lease") + + await registry.release(lora_id) + await asyncio.wait_for(wait_task, timeout=2) + + asyncio.run(run()) + + def test_batch_acquire_release_roundtrip(self): + async def run(): + registry = LoRARegistry() + for name in ("a", "b"): + await registry.register(LoRARef(lora_name=name, lora_path=f"/{name}")) + lora_ids = await registry.acquire(["a", "b", None]) + self.assertEqual(len(lora_ids), 3) + self.assertIsNone(lora_ids[2]) + await registry.release([i for i in lora_ids if i is not None]) + for name in ("a", "b"): + await registry.unregister(name) + + asyncio.run(run()) + + +class TestLruSkipsNonReloadable(CustomTestCase): + def test_lru_prefers_reloadable_even_if_older_entry_is_wire_loaded(self): + async def run(): + registry = LoRARegistry() + # Registration order = LRU order; the wire-loaded ref is oldest. + await registry.register( + LoRARef(lora_name="wire", lora_path="__distributed__", reloadable=False) + ) + await registry.register(LoRARef(lora_name="disk", lora_path="/d")) + self.assertEqual(await registry.lru_lora_name(), "disk") + + asyncio.run(run()) + + def test_all_non_reloadable_yields_no_victim(self): + async def run(): + registry = LoRARegistry() + await registry.register( + LoRARef(lora_name="w1", lora_path="__tensor__", reloadable=False) + ) + await registry.register( + LoRARef(lora_name="w2", lora_path="__distributed__", reloadable=False) + ) + self.assertIsNone(await registry.lru_lora_name()) + self.assertIsNone(await registry.lru_lora_name(exclude_pinned=True)) + + asyncio.run(run()) + + def test_exclude_pinned_still_respected(self): + async def run(): + registry = LoRARegistry() + await registry.register( + LoRARef(lora_name="pinned-disk", lora_path="/p", pinned=True) + ) + await registry.register(LoRARef(lora_name="plain-disk", lora_path="/q")) + self.assertEqual( + await registry.lru_lora_name(exclude_pinned=True), "plain-disk" + ) + self.assertEqual(await registry.lru_lora_name(), "pinned-disk") + + asyncio.run(run()) + + +class TestUnloadRefCacheCleanup(CustomTestCase): + def _make_unload_tm(self, success: bool) -> TokenizerManager: + tm = TokenizerManager.__new__(TokenizerManager) + tm.auto_create_handle_loop = Mock() + tm.server_args = MagicMock() + tm.server_args.enable_lora = True + tm.server_args.dp_size = 1 + tm.lora_update_lock = asyncio.Lock() + tm._unload_lora_adapter_locked = AsyncMock( + return_value=SimpleNamespace(success=success) + ) + tm.lora_ref_cache = {"a": LoRARef(lora_name="a", lora_path="/x")} + return tm + + def test_explicit_unload_drops_ref_cache_entry(self): + tm = self._make_unload_tm(success=True) + asyncio.run(tm.unload_lora_adapter(UnloadLoRAAdapterReqInput(lora_name="a"))) + self.assertNotIn("a", tm.lora_ref_cache) + + def test_failed_unload_keeps_ref_cache_entry(self): + tm = self._make_unload_tm(success=False) + asyncio.run(tm.unload_lora_adapter(UnloadLoRAAdapterReqInput(lora_name="a"))) + self.assertIn("a", tm.lora_ref_cache) + + def test_locked_helper_alone_keeps_ref_cache_entry(self): + # The max_loaded_loras LRU loop calls _unload_lora_adapter_locked + # directly — an EVICT — and a disk-backed adapter must keep its + # reload-catalog entry or it can never be implicitly reloaded. + tm = TokenizerManager.__new__(TokenizerManager) + tm.lora_update_lock = asyncio.Lock() + tm.lora_registry = MagicMock() + tm.lora_registry.unregister = AsyncMock(return_value="id-a") + tm.lora_registry.wait_for_unload = AsyncMock() + tm.update_lora_adapter_communicator = AsyncMock( + return_value=[SimpleNamespace(success=True)] + ) + tm.lora_ref_cache = {"a": LoRARef(lora_name="a", lora_path="/x")} + + async def run(): + async with tm.lora_update_lock: + return await tm._unload_lora_adapter_locked( + UnloadLoRAAdapterReqInput(lora_name="a") + ) + + result = asyncio.run(run()) + self.assertTrue(result.success) + self.assertIn("a", tm.lora_ref_cache) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/managers/test_abort_request_prefix.py b/test/registered/unit/managers/test_abort_request_prefix.py index 5734193e845e..36e83d046150 100644 --- a/test/registered/unit/managers/test_abort_request_prefix.py +++ b/test/registered/unit/managers/test_abort_request_prefix.py @@ -34,6 +34,7 @@ def _make_tokenizer_manager(rids=(), tokenizer_worker_num=1) -> TokenizerManager tm.server_args = MagicMock() tm.server_args.tokenizer_worker_num = tokenizer_worker_num tm.enable_metrics = False + tm.enable_lora = False tm.rid_to_state = {rid: Mock() for rid in rids} tm.send_to_scheduler = MagicMock() tm.tokenizer_ipc_name = None @@ -104,7 +105,11 @@ def test_multi_tokenizer_worker_skips_local_check(self): def _make_state(rid: str) -> ReqState: - obj = SimpleNamespace(rid=rid, stream=False, return_logprob=False) + # lora_path / lora_id model the real input-struct contract: always present, + # None = base model (the lease finalizer does plain None checks on them). + obj = SimpleNamespace( + rid=rid, stream=False, return_logprob=False, lora_path=None, lora_id=None + ) return ReqState([], False, asyncio.Event(), obj, MagicMock()) @@ -209,10 +214,15 @@ def _make_scheduler(waiting_rids=(), running_rids=(), chunked_rid=None): sched.chunked_req = FakeReq(chunked_rid) if chunked_rid is not None else None sched.waiting_queue = [FakeReq(rid) for rid in waiting_rids] sched.enable_hicache_storage = False + sched.dllm_config = None # abort_request reads it since the dLLM rework (#27877) sched.disaggregation_mode = DisaggregationMode.NULL sched.grammar_manager = MagicMock() sched.running_batch = SimpleNamespace(reqs=[FakeReq(rid) for rid in running_rids]) sched.cur_batch = None + # abort_request collects in-flight reqs from (running_batch, last_batch) on + # the pp_size == 1 path since the v0.5.16 rebase. + sched.ps = SimpleNamespace(pp_size=1) + sched.last_batch = None sched.ipc_channels = MagicMock() return sched