|
| 1 | +"""Release LoRA GPU slots when the pinned SGLang 0.5.10 engine unloads an adapter. |
| 2 | +
|
| 3 | +Upstream issue https://github.com/sgl-project/sglang/issues/21380 describes the |
| 4 | +missing cleanup: CPU weights disappear, but the GPU pool keeps the old UID. |
| 5 | +Repeated load/unload cycles then exhaust the pool. The CUDA 13 bundle's |
| 6 | +SGLang 0.5.20 already implements this cleanup and must keep its native path. |
| 7 | +
|
| 8 | +Only the SGLang child opts into this hook. It waits for the engine's normal |
| 9 | +import rather than importing CUDA or SGLang during Python site initialization. |
| 10 | +The engine already drains requests before calling the method wrapped here. |
| 11 | +""" |
| 12 | + |
| 13 | +from __future__ import annotations |
| 14 | + |
| 15 | +import builtins |
| 16 | +import importlib.metadata |
| 17 | +import logging |
| 18 | +import sys |
| 19 | +from collections.abc import Mapping, Sequence |
| 20 | +from functools import wraps |
| 21 | +from types import ModuleType |
| 22 | +from typing import Any |
| 23 | + |
| 24 | +_AFFECTED_VERSIONS = frozenset({"0.5.10", "0.5.10.post1"}) |
| 25 | +_MANAGER_MODULE = "sglang.srt.lora.lora_manager" |
| 26 | +_POOL_MODULE = "sglang.srt.lora.mem_pool" |
| 27 | +_CLASS_MARKER = "_sie_lora_slot_release_compat" |
| 28 | +_IMPORT_MARKER = "_sie_lora_slot_release_deferred" |
| 29 | +logger = logging.getLogger(__name__) |
| 30 | + |
| 31 | + |
| 32 | +def _release_slot(pool: Any, lora_id: str, empty_slot: Any) -> None: |
| 33 | + """Forget one drained adapter, preserving the base model's ``None`` UID.""" |
| 34 | + if lora_id is None: |
| 35 | + raise ValueError("a LoRA unload must never release the base model slot") |
| 36 | + slot = pool.uid_to_buffer_id.get(lora_id) |
| 37 | + if slot is not None and pool.buffer_id_to_uid[slot] != lora_id: |
| 38 | + raise RuntimeError("LoRA GPU slot mapping is inconsistent") |
| 39 | + # The policy can remember an adapter even when it is not currently in a |
| 40 | + # GPU slot. remove() is idempotent in both pinned upstream policies. |
| 41 | + pool.eviction_policy.remove(lora_id) |
| 42 | + if slot is not None: |
| 43 | + pool.uid_to_buffer_id.pop(lora_id) |
| 44 | + pool.buffer_id_to_uid[slot] = empty_slot |
| 45 | + |
| 46 | + |
| 47 | +def _patch_manager_module(module: ModuleType) -> None: |
| 48 | + manager = getattr(module, "LoRAManager", None) |
| 49 | + if manager is None: |
| 50 | + raise RuntimeError("SGLang does not expose the expected LoRAManager") |
| 51 | + if getattr(manager, _CLASS_MARKER, False): |
| 52 | + return |
| 53 | + pool_module = sys.modules.get(_POOL_MODULE) |
| 54 | + if pool_module is None or not hasattr(pool_module, "EMPTY_SLOT"): |
| 55 | + raise RuntimeError("SGLang does not expose the expected LoRA pool sentinel") |
| 56 | + empty_slot = pool_module.EMPTY_SLOT |
| 57 | + original_unload = manager.unload_lora_adapter |
| 58 | + |
| 59 | + @wraps(original_unload) |
| 60 | + def unload_with_slot_release(self: Any, lora_ref: Any) -> Any: |
| 61 | + # Do not alter a pool for an upstream operation that did not succeed. |
| 62 | + result = original_unload(self, lora_ref) |
| 63 | + if result.success: |
| 64 | + try: |
| 65 | + _release_slot(self.memory_pool, lora_ref.lora_id, empty_slot) |
| 66 | + except Exception: |
| 67 | + logger.exception("SIE could not release an unloaded LoRA GPU slot") |
| 68 | + return self.create_lora_update_result( |
| 69 | + success=False, |
| 70 | + error_message="LoRA GPU slot cleanup failed; reload this model before loading more adapters", |
| 71 | + ) |
| 72 | + return result |
| 73 | + |
| 74 | + manager.unload_lora_adapter = unload_with_slot_release |
| 75 | + setattr(manager, _CLASS_MARKER, True) |
| 76 | + |
| 77 | + |
| 78 | +def install_lora_compat() -> bool: |
| 79 | + """Install the exact-version backport; return False for unaffected engines.""" |
| 80 | + try: |
| 81 | + version = importlib.metadata.version("sglang") |
| 82 | + except importlib.metadata.PackageNotFoundError: |
| 83 | + return False |
| 84 | + if version not in _AFFECTED_VERSIONS: |
| 85 | + return False |
| 86 | + loaded = sys.modules.get(_MANAGER_MODULE) |
| 87 | + if isinstance(loaded, ModuleType) and hasattr(loaded, "LoRAManager"): |
| 88 | + _patch_manager_module(loaded) |
| 89 | + return True |
| 90 | + |
| 91 | + current_import = builtins.__import__ |
| 92 | + if getattr(current_import, _IMPORT_MARKER, False): |
| 93 | + return True |
| 94 | + |
| 95 | + def deferred_import( |
| 96 | + name: str, |
| 97 | + globals: Mapping[str, object] | None = None, |
| 98 | + locals: Mapping[str, object] | None = None, |
| 99 | + fromlist: Sequence[str] | None = (), |
| 100 | + level: int = 0, |
| 101 | + ) -> ModuleType: |
| 102 | + module = current_import(name, globals, locals, fromlist, level) |
| 103 | + loaded_module = sys.modules.get(_MANAGER_MODULE) |
| 104 | + if not isinstance(loaded_module, ModuleType) or not hasattr(loaded_module, "LoRAManager"): |
| 105 | + return module |
| 106 | + try: |
| 107 | + _patch_manager_module(loaded_module) |
| 108 | + finally: |
| 109 | + if builtins.__import__ is deferred_import: |
| 110 | + setattr(builtins, "__import__", current_import) # noqa: B010 |
| 111 | + return module |
| 112 | + |
| 113 | + setattr(deferred_import, _IMPORT_MARKER, True) |
| 114 | + setattr(builtins, "__import__", deferred_import) # noqa: B010 |
| 115 | + return True |
0 commit comments