From ada8ce1fd0f21d5125c6ad80cd0f62dfbfa1d936 Mon Sep 17 00:00:00 2001 From: Glen Liu <62917497+glenliu21@users.noreply.github.com> Date: Thu, 20 Nov 2025 16:34:32 -0500 Subject: [PATCH] allow loras to be implicitly evicted and loaded based on max_loaded_loras (#11526) --- python/sglang/srt/lora/lora_registry.py | 76 +++++-- .../managers/tokenizer_communicator_mixin.py | 74 +++++-- .../sglang/srt/managers/tokenizer_manager.py | 55 ++++- test/srt/lora/test_lora_update.py | 197 ++++++++++++++++-- 4 files changed, 340 insertions(+), 62 deletions(-) diff --git a/python/sglang/srt/lora/lora_registry.py b/python/sglang/srt/lora/lora_registry.py index 3df0fa5be..d31c5ab93 100644 --- a/python/sglang/srt/lora/lora_registry.py +++ b/python/sglang/srt/lora/lora_registry.py @@ -14,6 +14,7 @@ import asyncio +from collections import OrderedDict from dataclasses import dataclass, field, fields from typing import Dict, List, Optional, Union from uuid import uuid4 @@ -71,8 +72,11 @@ class LoRARegistry: # Please note that the counter increment/decrement operations are not synchronized through this # lock, as they are designed to be non-blocking and can be performed concurrently. self._registry_lock = RWLock() - # A dictionary to hold LoRARef objects, mapping from LoRA name to LoRARef. - self._registry: Dict[str, LoRARef] = {} + # An ordered dictionary to hold LoRARef objects, mapping from LoRA name to LoRARef. + # The LoRARefs are stored in LRU order, such that LoRA adapters that have been + # most recently used are stored at the end. Note that lookups count for accesses. + # Ties are broken arbitrarily. + self._registry: OrderedDict[str, LoRARef] = OrderedDict() # Counters for ongoing requests, mapping from LoRA ID to ConcurrentCounter. self._counters: Dict[str, ConcurrentCounter] = {} @@ -124,29 +128,30 @@ class LoRARegistry: f"The following requested LoRA adapters are not loaded: {name}\n" f"Loaded adapters: {self._registry.keys()}." ) + self._registry.move_to_end(name) return lora_ref.lora_id - async with self._registry_lock.reader_lock: - if isinstance(lora_name, str): + 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) - return lora_id - elif isinstance(lora_name, list): + + 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 - ] - ) - return lora_ids - else: - raise TypeError( - "lora_name must be either a string or a list of strings." - ) + # 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.") async def release(self, lora_id: Union[str, List[str]]): """ @@ -186,6 +191,37 @@ class LoRARegistry: await self._counters[lora_id].wait_for_zero() del self._counters[lora_id] + async def get_unregistered_loras(self, lora_name: set[str]): + """ + Returns all LoRA adapters in lora_name that are not found in self._registry. + """ + async with self._registry_lock.writer_lock: + unregistered_loras = [] + + for name in lora_name: + if name in self._registry: + # This counts as a lookup, so we want to update the cache + self._registry.move_to_end(name) + else: + unregistered_loras.append(name) + + return unregistered_loras + + async def lru_lora_name(self, exclude_pinned=False): + """ + Returns the least recently used LoRA adapter. + 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 + def _register_adapter(self, lora_ref: LoRARef): """ Internal helper method to register a LoRA adapter. diff --git a/python/sglang/srt/managers/tokenizer_communicator_mixin.py b/python/sglang/srt/managers/tokenizer_communicator_mixin.py index 9205db0a9..107a7f6ca 100644 --- a/python/sglang/srt/managers/tokenizer_communicator_mixin.py +++ b/python/sglang/srt/managers/tokenizer_communicator_mixin.py @@ -495,6 +495,26 @@ class TokenizerCommunicatorMixin: return success, message + async def _unload_lora_adapter_locked( + self: TokenizerManager, + obj: UnloadLoRAAdapterReqInput, + ) -> UnloadLoRAAdapterReqOutput: + assert ( + self.lora_update_lock.locked() + ), "self.lora_update_lock must be locked in order for self._unload_lora_adapter_locked() to be called" + + # Unregister the LoRA adapter from the registry to stop new requests for this adapter + # from being started. + lora_id = await self.lora_registry.unregister(obj.lora_name) + obj.lora_id = lora_id + + # Initiate the actual unloading operation at the backend processes only after all + # ongoing requests using this LoRA adapter are finished. + await self.lora_registry.wait_for_unload(lora_id) + result = (await self.update_lora_adapter_communicator(obj))[0] + + return result + async def load_lora_adapter( self: TokenizerManager, obj: LoadLoRAAdapterReqInput, @@ -520,17 +540,6 @@ class TokenizerCommunicatorMixin: ) async with self.lora_update_lock: - if ( - self.server_args.max_loaded_loras is not None - and self.lora_registry.num_registered_loras - >= self.server_args.max_loaded_loras - ): - raise ValueError( - f"Cannot load LoRA adapter {obj.lora_name} at path {obj.lora_path}. " - f"Maximum number of loaded LoRA adapters is {self.server_args.max_loaded_loras}. " - "Please unload some LoRA adapters before loading new ones." - ) - # Generate new uniquely identifiable LoRARef object. new_adapter = LoRARef( lora_name=obj.lora_name, @@ -545,6 +554,37 @@ class TokenizerCommunicatorMixin: # Register the LoRA adapter only after loading is successful. if result.success: 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] return result except ValueError as e: @@ -581,17 +621,7 @@ class TokenizerCommunicatorMixin: ) async with self.lora_update_lock: - # Unregister the LoRA adapter from the registry to stop new requests for this adapter - # from being started. - lora_id = await self.lora_registry.unregister(obj.lora_name) - obj.lora_id = lora_id - - # Initiate the actual unloading operation at the backend processes only after all - # ongoing requests using this LoRA adapter are finished. - await self.lora_registry.wait_for_unload(lora_id) - result = (await self.update_lora_adapter_communicator(obj))[0] - - return result + return await self._unload_lora_adapter_locked(obj) 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 803fbdce0..7e303d8b0 100644 --- a/python/sglang/srt/managers/tokenizer_manager.py +++ b/python/sglang/srt/managers/tokenizer_manager.py @@ -42,7 +42,7 @@ from fastapi import BackgroundTasks from sglang.srt.configs.model_config import ModelConfig from sglang.srt.disaggregation.utils import DisaggregationMode from sglang.srt.environ import envs -from sglang.srt.lora.lora_registry import LoRARegistry +from sglang.srt.lora.lora_registry import LoRARef, LoRARegistry from sglang.srt.managers.async_dynamic_batch_tokenizer import AsyncDynamicbatchTokenizer from sglang.srt.managers.async_mm_data_processor import AsyncMMDataProcessor from sglang.srt.managers.disagg_service import start_disagg_service @@ -60,6 +60,7 @@ from sglang.srt.managers.io_struct import ( GenerateReqInput, GetLoadReqInput, HealthCheckOutput, + LoadLoRAAdapterReqInput, OpenSessionReqOutput, SessionParams, TokenizedEmbeddingReqInput, @@ -357,6 +358,13 @@ class TokenizerManager(TokenizerCommunicatorMixin): # Please note that, unlike `model_update_lock`, this does not block inference, allowing # LoRA updates and inference to overlap. self.lora_update_lock = asyncio.Lock() + # A cache for mapping the lora_name for LoRA adapters that have been loaded at any + # point to their latest LoRARef objects, so that they can be + # dynamically loaded if needed for inference + self.lora_ref_cache: Dict[str, LoRARef] = {} + if self.server_args.lora_paths is not None: + for lora_ref in self.server_args.lora_paths: + self.lora_ref_cache[lora_ref.lora_name] = lora_ref # Disaggregation self.disaggregation_mode = DisaggregationMode( @@ -448,6 +456,51 @@ class TokenizerManager(TokenizerCommunicatorMixin): async with self.model_update_lock.reader_lock: if self.server_args.enable_lora and obj.lora_path: + if isinstance(obj.lora_path, str): + unique_lora_paths = set([obj.lora_path]) + else: + unique_lora_paths = set(obj.lora_path) + + if ( + self.server_args.max_loaded_loras is not None + and len(unique_lora_paths) > self.server_args.max_loaded_loras + ): + raise ValueError( + f"Received request with {len(unique_lora_paths)} unique loras requested " + f"but max loaded loras is {self.server_args.max_loaded_loras}" + ) + + # Reload all existing LoRA adapters that have been dynamically unloaded + unregistered_loras = await self.lora_registry.get_unregistered_loras( + unique_lora_paths + ) + for lora_path in unregistered_loras: + if lora_path is None: + continue + + if lora_path not in self.lora_ref_cache: + raise ValueError( + f"Got LoRA adapter that has never been loaded: {lora_path}\n" + 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] + load_result = await self.load_lora_adapter( + LoadLoRAAdapterReqInput( + lora_name=new_lora_ref.lora_name, + lora_path=new_lora_ref.lora_path, + pinned=new_lora_ref.pinned, + ) + ) + if ( + not load_result.success + and "already loaded" not in load_result.error_message + ): + raise ValueError( + f"Failed to implicitly load LoRA adapter {lora_path}: {load_result.error_message}" + ) + # Look up the LoRA ID from the registry and start tracking ongoing LoRA requests. obj.lora_id = await self.lora_registry.acquire(obj.lora_path) diff --git a/test/srt/lora/test_lora_update.py b/test/srt/lora/test_lora_update.py index 09d7778b9..637f840cc 100644 --- a/test/srt/lora/test_lora_update.py +++ b/test/srt/lora/test_lora_update.py @@ -58,6 +58,9 @@ class Operation: data: Optional[Any] # If the operation is expected to fail, this is the error message to expect expected_error: Optional[str] = None + # Because the logic for implicitly evicting LoRA adapters can be complicated, we explicitly + # pass in LoRA adapters that should be implicitly evicted here + expected_implicit_evictions: Optional[set[str]] = None @dataclass @@ -142,20 +145,20 @@ BASIC_TESTS = [ data=create_batch_data( "Nutanix/Meta-Llama-3.1-8B-Instruct_lora_4_alpha_16" ), - expected_error="not loaded", ), Operation( type=OperationType.FORWARD, data=create_batch_data("pbevan11/llama-3.1-8b-ocr-correction"), - expected_error="not loaded", ), Operation( type=OperationType.LOAD, data="Nutanix/Meta-Llama-3.1-8B-Instruct_lora_4_alpha_16", + expected_error="already loaded", ), Operation( type=OperationType.LOAD, data="pbevan11/llama-3.1-8b-ocr-correction", + expected_error="already loaded", ), Operation( type=OperationType.FORWARD, @@ -174,7 +177,6 @@ BASIC_TESTS = [ Operation( type=OperationType.FORWARD, data=create_batch_data("philschmid/code-llama-3-1-8b-text-to-sql-lora"), - expected_error="not loaded", ), Operation( type=OperationType.FORWARD, @@ -198,12 +200,10 @@ BASIC_TESTS = [ data=create_batch_data( "Nutanix/Meta-Llama-3.1-8B-Instruct_lora_4_alpha_16" ), - expected_error="not loaded", ), Operation( type=OperationType.FORWARD, data=create_batch_data("pbevan11/llama-3.1-8b-ocr-correction"), - expected_error="not loaded", ), Operation( type=OperationType.FORWARD, @@ -269,7 +269,6 @@ BASIC_TESTS = [ Operation( type=OperationType.FORWARD, data=create_batch_data("philschmid/code-llama-3-1-8b-text-to-sql-lora"), - expected_error="not loaded", ), Operation( type=OperationType.FORWARD, @@ -295,12 +294,10 @@ BASIC_TESTS = [ data=create_batch_data( "Nutanix/Meta-Llama-3.1-8B-Instruct_lora_4_alpha_16" ), - expected_error="not loaded", ), Operation( type=OperationType.FORWARD, data=create_batch_data("pbevan11/llama-3.1-8b-ocr-correction"), - expected_error="not loaded", ), Operation( type=OperationType.FORWARD, @@ -309,14 +306,17 @@ BASIC_TESTS = [ Operation( type=OperationType.LOAD, data="philschmid/code-llama-3-1-8b-text-to-sql-lora", + expected_error="already loaded", ), Operation( type=OperationType.LOAD, data="Nutanix/Meta-Llama-3.1-8B-Instruct_lora_4_alpha_16", + expected_error="already loaded", ), Operation( type=OperationType.LOAD, data="pbevan11/llama-3.1-8b-ocr-correction", + expected_error="already loaded", ), Operation( type=OperationType.FORWARD, @@ -364,7 +364,7 @@ TARGET_MODULE_TESTS = [ data=create_batch_data( "Nutanix/Meta-Llama-3.1-8B-Instruct_lora_4_alpha_16" ), - expected_error="not loaded", + expected_error="never been loaded", ), Operation( type=OperationType.LOAD, @@ -404,7 +404,7 @@ TARGET_MODULE_TESTS = [ data=create_batch_data( "algoprog/fact-generation-llama-3.1-8b-instruct-lora" ), - expected_error="not loaded", + expected_error="never been loaded", ), Operation( type=OperationType.LOAD, @@ -444,7 +444,7 @@ TARGET_MODULE_TESTS = [ data=create_batch_data( "Nutanix/Meta-Llama-3.1-8B-Instruct_lora_4_alpha_16" ), - expected_error="not loaded", + expected_error="never been loaded", ), Operation( type=OperationType.LOAD, @@ -485,12 +485,12 @@ MAX_LORA_RANK_TESTS = [ Operation( type=OperationType.FORWARD, data=create_batch_data("philschmid/code-llama-3-1-8b-text-to-sql-lora"), - expected_error="not loaded", + expected_error="never been loaded", ), Operation( type=OperationType.FORWARD, data=create_batch_data("pbevan11/llama-3.1-8b-ocr-correction"), - expected_error="not loaded", + expected_error="never been loaded", ), Operation( type=OperationType.LOAD, @@ -516,7 +516,7 @@ MAX_LORA_RANK_TESTS = [ data=create_batch_data( "philschmid/code-llama-3-1-8b-text-to-sql-lora", ), - expected_error="not loaded", + expected_error="never been loaded", ), Operation( type=OperationType.FORWARD, @@ -553,7 +553,7 @@ MAX_LORA_RANK_TESTS = [ Operation( type=OperationType.FORWARD, data=create_batch_data("philschmid/code-llama-3-1-8b-text-to-sql-lora"), - expected_error="not loaded", + expected_error="never been loaded", ), Operation( type=OperationType.LOAD, @@ -580,7 +580,7 @@ MAX_LORA_RANK_TESTS = [ ] MAX_LOADED_LORAS_TESTS = [ TestCase( - description="Test max_loaded_loras limit", + description="Test max_loaded_loras limit as well as implicit eviction and reloading", base="meta-llama/Llama-3.1-8B-Instruct", max_loras_per_batch=2, max_loaded_loras=2, @@ -598,15 +598,143 @@ MAX_LOADED_LORAS_TESTS = [ Operation( type=OperationType.LOAD, data="pbevan11/llama-3.1-8b-ocr-correction", - expected_error="Maximum number of loaded LoRA adapters", + expected_implicit_evictions={ + "philschmid/code-llama-3-1-8b-text-to-sql-lora" + }, + ), + # Implicitly load "philschmid/code-llama-3-1-8b-text-to-sql-lora" + Operation( + type=OperationType.FORWARD, + data=create_batch_data( + [ + "Nutanix/Meta-Llama-3.1-8B-Instruct_lora_4_alpha_16", + "philschmid/code-llama-3-1-8b-text-to-sql-lora", + ] + ), + expected_implicit_evictions={"pbevan11/llama-3.1-8b-ocr-correction"}, + ), + Operation( + type=OperationType.UNLOAD, + data="philschmid/code-llama-3-1-8b-text-to-sql-lora", + ), + Operation( + type=OperationType.FORWARD, + data=create_batch_data( + [ + "Nutanix/Meta-Llama-3.1-8B-Instruct_lora_4_alpha_16", + ] + ), + ), + Operation( + type=OperationType.LOAD, + data="philschmid/code-llama-3-1-8b-text-to-sql-lora", + ), + # Implicitly load "pbevan11/llama-3.1-8b-ocr-correction" and make sure that "Nutanix/Meta-Llama-3.1-8B-Instruct_lora_4_alpha_16" + # isn't implicitly unloaded even though it is LRU because it is needed for this forward pass + Operation( + type=OperationType.FORWARD, + data=create_batch_data( + [ + "Nutanix/Meta-Llama-3.1-8B-Instruct_lora_4_alpha_16", + "pbevan11/llama-3.1-8b-ocr-correction", + ] + ), + expected_implicit_evictions={ + "philschmid/code-llama-3-1-8b-text-to-sql-lora" + }, ), Operation( type=OperationType.UNLOAD, data="Nutanix/Meta-Llama-3.1-8B-Instruct_lora_4_alpha_16", ), + Operation( + type=OperationType.LOAD, + data="algoprog/fact-generation-llama-3.1-8b-instruct-lora", + ), + Operation( + type=OperationType.FORWARD, + data=create_batch_data( + [ + "Nutanix/Meta-Llama-3.1-8B-Instruct_lora_4_alpha_16", + "philschmid/code-llama-3-1-8b-text-to-sql-lora", + ] + ), + expected_implicit_evictions={ + "pbevan11/llama-3.1-8b-ocr-correction", + "algoprog/fact-generation-llama-3.1-8b-instruct-lora", + }, + ), + ], + ), + TestCase( + description="Test implicit eviction and reloading with pinned LoRA adapters", + base="meta-llama/Llama-3.1-8B-Instruct", + max_loras_per_batch=2, + max_loaded_loras=2, + all_adapters=[ + "philschmid/code-llama-3-1-8b-text-to-sql-lora", + "Nutanix/Meta-Llama-3.1-8B-Instruct_lora_4_alpha_16", + "pbevan11/llama-3.1-8b-ocr-correction", + ], + initial_adapters=[ + { + "lora_name": "philschmid/code-llama-3-1-8b-text-to-sql-lora", + "lora_path": "philschmid/code-llama-3-1-8b-text-to-sql-lora", + "pinned": True, + } + ], + op_sequence=[ + Operation( + type=OperationType.LOAD, + data="Nutanix/Meta-Llama-3.1-8B-Instruct_lora_4_alpha_16", + ), Operation( type=OperationType.LOAD, data="pbevan11/llama-3.1-8b-ocr-correction", + expected_implicit_evictions={ + "Nutanix/Meta-Llama-3.1-8B-Instruct_lora_4_alpha_16" + }, + ), + # Implicitly load "Nutanix/Meta-Llama-3.1-8B-Instruct_lora_4_alpha_16" + Operation( + type=OperationType.FORWARD, + data=create_batch_data( + [ + "philschmid/code-llama-3-1-8b-text-to-sql-lora", + "Nutanix/Meta-Llama-3.1-8B-Instruct_lora_4_alpha_16", + ] + ), + expected_implicit_evictions={"pbevan11/llama-3.1-8b-ocr-correction"}, + ), + Operation( + type=OperationType.LOAD, + data={ + "lora_name": "pbevan11/llama-3.1-8b-ocr-correction", + "lora_path": "pbevan11/llama-3.1-8b-ocr-correction", + "pinned": True, + }, + expected_error="unpin some adapters", + ), + Operation( + type=OperationType.UNLOAD, + data="philschmid/code-llama-3-1-8b-text-to-sql-lora", + ), + Operation( + type=OperationType.LOAD, + data={ + "lora_name": "pbevan11/llama-3.1-8b-ocr-correction", + "lora_path": "pbevan11/llama-3.1-8b-ocr-correction", + "pinned": True, + }, + ), + Operation( + type=OperationType.FORWARD, + data=create_batch_data( + [ + "Nutanix/Meta-Llama-3.1-8B-Instruct_lora_4_alpha_16", + "pbevan11/llama-3.1-8b-ocr-correction", + ] + ), ), ], ), @@ -640,7 +768,7 @@ EVICTION_TESTS = [ "lora_path": "Nutanix/Meta-Llama-3.1-8B-Instruct_lora_4_alpha_16", "pinned": True, }, - expected_error="starvation", + expected_error="unpin some adapters", ), Operation( type=OperationType.LOAD, @@ -681,7 +809,7 @@ EVICTION_TESTS = [ "lora_path": "philschmid/code-llama-3-1-8b-text-to-sql-lora", "pinned": True, }, - expected_error="starvation", + expected_error="unpin some adapters", ), Operation( type=OperationType.LOAD, @@ -805,6 +933,7 @@ class LoRAUpdateTestSessionBase: lora_name: str, lora_path: Optional[str] = None, expected_error: Optional[str] = None, + expected_implicit_evictions: Optional[set[str]] = None, ): """ Load a LoRA adapter by name and path. @@ -823,6 +952,7 @@ class LoRAUpdateTestSessionBase: lora_paths: List[str], max_new_tokens: int = 32, expected_error: Optional[str] = None, + expected_implicit_evictions: Optional[set[str]] = None, ): """ Perform a batch forward pass with the current set of loaded LoRA adapters. @@ -869,6 +999,7 @@ class LoRAUpdateEngineTestSession(LoRAUpdateTestSessionBase): lora_path: Optional[str] = None, expected_error: Optional[str] = None, pinned: bool = False, + expected_implicit_evictions: Optional[set[str]] = None, ): """ Load a LoRA adapter by name and path. @@ -893,6 +1024,9 @@ class LoRAUpdateEngineTestSession(LoRAUpdateTestSessionBase): print(f"Received error as expected: {response.error_message}") else: self.expected_adapters.add(lora_name) + if expected_implicit_evictions is not None: + self.expected_adapters -= expected_implicit_evictions + self.testcase.assertTrue( response.success, f"Failed to load LoRA adapter {lora_name}: {response.error_message}", @@ -933,6 +1067,7 @@ class LoRAUpdateEngineTestSession(LoRAUpdateTestSessionBase): lora_paths: List[str], max_new_tokens: int = 32, expected_error: Optional[str] = None, + expected_implicit_evictions: Optional[set[str]] = None, ): """ Perform a batch forward pass with the current set of loaded LoRA adapters. @@ -964,6 +1099,13 @@ class LoRAUpdateEngineTestSession(LoRAUpdateTestSessionBase): output = response.output_strs print(f"output_strs: {output}") + self.expected_adapters.update( + [lora_path for lora_path in lora_paths if lora_path is not None] + ) + + if expected_implicit_evictions is not None: + self.expected_adapters -= expected_implicit_evictions + return output @@ -1026,6 +1168,7 @@ class LoRAUpdateServerTestSession(LoRAUpdateTestSessionBase): lora_path: Optional[str] = None, expected_error: Optional[str] = None, pinned: bool = False, + expected_implicit_evictions: Optional[set[str]] = None, ): """ Load a LoRA adapter by name and path. @@ -1051,6 +1194,9 @@ class LoRAUpdateServerTestSession(LoRAUpdateTestSessionBase): print(f"Received error as expected: {response.text}") else: self.expected_adapters.add(lora_name) + if expected_implicit_evictions is not None: + self.expected_adapters -= expected_implicit_evictions + self.testcase.assertTrue( response.ok, f"Failed to load LoRA adapter {lora_name}: {response.text}" ) @@ -1072,6 +1218,7 @@ class LoRAUpdateServerTestSession(LoRAUpdateTestSessionBase): DEFAULT_URL_FOR_TEST + "/unload_lora_adapter", json={"lora_name": lora_name}, ) + self.testcase.assertTrue( response.ok, f"Failed to unload LoRA adapter {lora_name}: {response.text}" ) @@ -1090,6 +1237,7 @@ class LoRAUpdateServerTestSession(LoRAUpdateTestSessionBase): lora_paths: List[str], max_new_tokens: int = 32, expected_error: Optional[str] = None, + expected_implicit_evictions: Optional[set[str]] = None, ): """ Perform a batch forward pass with the current set of loaded LoRA adapters. @@ -1131,6 +1279,14 @@ class LoRAUpdateServerTestSession(LoRAUpdateTestSessionBase): f"Expected {len(prompts)} outputs, but got {len(output)}", ) print(f"output_strs: {output}") + + self.expected_adapters.update( + [lora_path for lora_path in lora_paths if lora_path is not None] + ) + + if expected_implicit_evictions is not None: + self.expected_adapters -= expected_implicit_evictions + return output @@ -1192,6 +1348,7 @@ class TestLoRADynamicUpdate(CustomTestCase): op_type = op.type data = op.data expected_error = op.expected_error + expected_implicit_evictions = op.expected_implicit_evictions print("-" * 100) print( f"Running operation: {op_type} --- data: {data} --- mode: {mode} ---" @@ -1208,6 +1365,7 @@ class TestLoRADynamicUpdate(CustomTestCase): result = session.load_lora_adapter( expected_error=expected_error, + expected_implicit_evictions=expected_implicit_evictions, **adapter_info, ) elif op_type == OperationType.UNLOAD: @@ -1221,6 +1379,7 @@ class TestLoRADynamicUpdate(CustomTestCase): lora_paths=list(adapters), max_new_tokens=max_new_tokens, expected_error=expected_error, + expected_implicit_evictions=expected_implicit_evictions, ) if not expected_error: forward_outputs.append(result)