Spaces:
Running
Running
| """AdapterPool — Sprint 16 runtime 核心：共享 base model 上的 LoRA 熱插拔。 | |
| 【這個模組在做什麼？】 | |
| Gemma 4 E4B 整個 ML service 只載入一份（共享 base model），但每位使用者可能 | |
| 有自己的 stylist LoRA adapter。peft 的 `set_adapter()` 會直接改動模型狀態 — | |
| 不同併發請求如果沒序列化，後一個會看到前一個的 adapter，造成跨使用者資料污染。 | |
| AdapterPool 用 asyncio.Lock 把「load → set → yield → reset」整段序列化， | |
| 並搭配 LRU 管理已載入的 adapter 數量（避免 HF Spaces Free tier VRAM 爆掉）。 | |
| 【關鍵設計決策】 | |
| 1. **default adapter 永遠 pin 在 pool** | |
| 主流程（stylist 給新手 + 非訂閱者）都走 default；如果被驅逐會造成所有 | |
| 無 adapter 請求每次都要 reload，CPU 100% 拉爆。 | |
| 2. **asyncio.Lock 序列化 set_adapter** | |
| peft 的 set_adapter 是 mutation，不 serialize 會 race。這個 lock 粒度 | |
| 粗一點（包整個 generate 區塊）比讓 peft 內部 lock 再 fine-grain 更安全。 | |
| trade-off：單 ML service process 併發容量 = 1（在 set_adapter 階段）。 | |
| 對 Wardrobe OS 的流量規模（< 10 req/s）完全夠用。 | |
| 3. **capacity 包含 default** | |
| 更精確地說：capacity 是 non-default adapter 數量的上限。default 不算進 | |
| LRU queue，永遠 pinned。所以 capacity=4 實際 loaded 總數最多是 5。 | |
| 4. **HF Spaces Free tier 預設 capacity=4** | |
| review v8 P1-7 的選項 (c)：免費方案 + capacity=4 + cold-start penalty | |
| acceptable。MVP 階段 marketplace subscribers 現實上 < 10 個，capacity=4 | |
| 夠用；> 100 時再升級到 (a) 或 (b)。 | |
| 【併發模型】 | |
| asyncio.Lock → FastAPI async 請求自然串起來。不是 threading.Lock（ML service | |
| 是 async，沒 GIL 問題；threading.Lock 會整個 event loop 掛掉）。 | |
| 【對應測試】 | |
| apps/ml-service/tests/test_adapter_pool.py（用 FakePeftModel mock，不載實模型） | |
| """ | |
| from __future__ import annotations | |
| import asyncio | |
| import logging | |
| from collections import OrderedDict | |
| from contextlib import asynccontextmanager | |
| from typing import Any, AsyncIterator | |
| logger = logging.getLogger(__name__) | |
| class AdapterPool: | |
| """LRU pool for hot-swapping PEFT LoRA adapters on a shared base model.""" | |
| def __init__( | |
| self, | |
| model: Any, | |
| *, | |
| default_name: str, | |
| capacity: int = 4, | |
| ) -> None: | |
| """初始化 pool。 | |
| Args: | |
| model: peft.PeftModel 或相容物件（需要 load_adapter / set_adapter / | |
| delete_adapter 三個方法，以及 loaded_adapters 的存在性檢查能力）。 | |
| Production 會是實際的 PeftModel；測試會是 FakePeftModel。 | |
| default_name: 已經 pin 在 base model 上的 default adapter name。 | |
| AdapterPool 假設呼叫者在 __init__ 之前已經把 default 載入好。 | |
| capacity: non-default adapter 同時 loaded 的上限。default 不算進去。 | |
| """ | |
| if capacity < 1: | |
| raise ValueError(f"capacity must be >= 1, got {capacity}") | |
| self._model = model | |
| self._default = default_name | |
| self._capacity = capacity | |
| # LRU queue 只裝 non-default adapter。最近使用的放在尾端，最舊的在頭端。 | |
| # 用 OrderedDict 是因為 move_to_end 是 O(1)，加上 popitem(last=False) 也是 O(1)。 | |
| # value 欄不重要（我們只關心 key 順序），存 hf_repo_id 方便 debug。 | |
| self._lru: OrderedDict[str, str] = OrderedDict() | |
| # asyncio.Lock — 整個 use() context 內 set_adapter + yield 都 hold 著 | |
| self._lock = asyncio.Lock() | |
| # ------------------------------------------------------------------ | |
| # 公開 API | |
| # ------------------------------------------------------------------ | |
| def is_loaded(self, name: str) -> bool: | |
| """檢查 adapter 是否在 pool 內（default 永遠是 True）。""" | |
| if name == self._default: | |
| return True | |
| return name in self._lru | |
| async def use(self, name: str, hf_repo_id: str) -> AsyncIterator[Any]: | |
| """切到指定 adapter 執行 generation，離開時切回 default。 | |
| 【使用方式】 | |
| async with pool.use("user-uuid", "hf://fengchen31/wardrobe-os-stylist-<uuid>") as model: | |
| output = model.generate(..., use_cache=False) # ← Sprint 16 Task 2.0 incident | |
| 【契約】 | |
| - 進入時：adapter 被載入（如果還沒）並 set 為 active | |
| - 離開時：無論是否有 exception，adapter 都切回 default | |
| - 整段序列化：同時呼叫 use() 的請求會 FIFO 排隊 | |
| Args: | |
| name: adapter name（通常是使用者 UUID，或 "default"）。 | |
| hf_repo_id: HF Hub 路徑，load_adapter 時用。name=="default" 時會被忽略。 | |
| Yields: | |
| 底層的 model 物件（方便 caller 直接拿來 generate）。 | |
| """ | |
| async with self._lock: | |
| # 載入 + 切換 | |
| await self._ensure_loaded(name, hf_repo_id) | |
| # peft 的 set_adapter 是同步 torch mutation — 放到 worker thread 以免 | |
| # 凍結 event loop（Gemma 11 capability 共用同一個 loop） | |
| await asyncio.to_thread(self._model.set_adapter, name) | |
| logger.debug("adapter_pool: set active → %s", name) | |
| try: | |
| yield self._model | |
| finally: | |
| # 一律切回 default — 防跨使用者污染，即使 caller 拋 exception | |
| await asyncio.to_thread(self._model.set_adapter, self._default) | |
| logger.debug("adapter_pool: reset → %s", self._default) | |
| # ------------------------------------------------------------------ | |
| # 內部 | |
| # ------------------------------------------------------------------ | |
| async def _ensure_loaded(self, name: str, hf_repo_id: str) -> None: | |
| """確保 adapter 已在 pool 內；該驅逐就驅逐，該載入就載入。""" | |
| if name == self._default: | |
| # default 永遠 loaded，no-op | |
| return | |
| if name in self._lru: | |
| # hit — 把它推到 LRU 最新端 | |
| self._lru.move_to_end(name, last=True) | |
| return | |
| # miss — 需要載入，先看是否要驅逐最舊的 | |
| if len(self._lru) >= self._capacity: | |
| oldest_name, oldest_repo = self._lru.popitem(last=False) | |
| logger.info( | |
| "adapter_pool: evict %s (hf=%s) to make room for %s", | |
| oldest_name, | |
| oldest_repo, | |
| name, | |
| ) | |
| try: | |
| await asyncio.to_thread(self._model.delete_adapter, oldest_name) | |
| except Exception as exc: | |
| # peft 的 delete_adapter 理論上不會失敗，但保險起見不讓它炸掉 | |
| # 整個 pool。記 warning 繼續 — LRU tracking 已經移除了。 | |
| logger.warning( | |
| "adapter_pool: delete_adapter(%s) failed: %s — LRU still evicted", | |
| oldest_name, | |
| exc, | |
| ) | |
| # 真正 load — HF Hub download + torch weight load，30s+ on cold miss | |
| # 必須 to_thread，否則整個 event loop 被這段凍住 | |
| logger.info("adapter_pool: load_adapter %s from %s", name, hf_repo_id) | |
| await asyncio.to_thread(self._model.load_adapter, hf_repo_id, adapter_name=name) | |
| self._lru[name] = hf_repo_id | |