Move lora request validation to tokenizer_manager from server (#18962)

Co-authored-by: Satyam Kumar <satyamk@linkedin.com>
This commit is contained in:
satyamk7054
2026-02-19 21:03:19 +08:00
committed by GitHub
co-authored by Satyam Kumar
parent d07e8aa4a3
commit 963def7f26
7 changed files with 40 additions and 114 deletions
@@ -70,19 +70,6 @@ class OpenAIServingBase(ABC):
# Fall back to explicit lora_path
return explicit_lora_path
def _validate_lora_enabled(self, adapter_name: str) -> None:
"""Check that LoRA is enabled before attempting to use an adapter.
Raises ValueError with actionable guidance if --enable-lora flag is missing.
Adapter existence is validated later by TokenizerManager.lora_registry.
"""
if not self.tokenizer_manager.server_args.enable_lora:
raise ValueError(
f"LoRA adapter '{adapter_name}' was requested, but LoRA is not enabled. "
"Please launch the server with --enable-lora flag and preload adapters "
"using --lora-paths or /load_lora_adapter endpoint."
)
async def handle_request(
self, request: OpenAIServingRequest, raw_request: Request
) -> Union[Any, StreamingResponse, ErrorResponse]:
@@ -277,15 +277,6 @@ class OpenAIServingChat(OpenAIServingBase):
# Resolve LoRA adapter from model parameter or explicit lora_path
lora_path = self._resolve_lora_path(request.model, request.lora_path)
if lora_path:
first_adapter = (
lora_path
if isinstance(lora_path, str)
else next((a for a in lora_path if a), None)
)
if first_adapter:
self._validate_lora_enabled(first_adapter)
img_max_dynamic_patch, vid_max_dynamic_patch = _extract_max_dynamic_patch(
request
)
@@ -98,14 +98,6 @@ class OpenAIServingCompletion(OpenAIServingBase):
# Resolve LoRA adapter from model parameter or explicit lora_path
lora_path = self._resolve_lora_path(request.model, request.lora_path)
if lora_path:
first_adapter = (
lora_path
if isinstance(lora_path, str)
else next((a for a in lora_path if a), None)
)
if first_adapter:
self._validate_lora_enabled(first_adapter)
adapted_request = GenerateReqInput(
**prompt_kwargs,
@@ -128,14 +128,6 @@ class OpenAIServingEmbedding(OpenAIServingBase):
# Resolve LoRA adapter from model parameter or explicit lora_path
lora_path = self._resolve_lora_path(request.model, request.lora_path)
if lora_path:
first_adapter = (
lora_path
if isinstance(lora_path, str)
else next((a for a in lora_path if a), None)
)
if first_adapter:
self._validate_lora_enabled(first_adapter)
adapted_request = EmbeddingReqInput(
**prompt_kwargs,
@@ -511,8 +511,7 @@ class TokenizerManager(TokenizerCommunicatorMixin, TokenizerManagerMultiItemMixi
await self.is_pause_cond.wait_for(lambda: not self.is_pause)
async with self.model_update_lock.reader_lock:
if self.server_args.enable_lora and obj.lora_path:
await self._resolve_lora_path(obj)
await self._validate_and_resolve_lora(obj)
# Tokenize the request and send it to the scheduler
if obj.is_single:
@@ -2213,6 +2212,27 @@ class TokenizerManager(TokenizerCommunicatorMixin, TokenizerManagerMultiItemMixi
if len(self.model_update_tmp) == self.server_args.dp_size:
self.model_update_result.set_result(self.model_update_tmp)
async def _validate_and_resolve_lora(
self, obj: Union[GenerateReqInput, EmbeddingReqInput]
) -> None:
if not obj.lora_path:
return
if not self.server_args.enable_lora:
first_adapter = (
obj.lora_path
if isinstance(obj.lora_path, str)
else next((a for a in obj.lora_path if a), None)
)
raise ValueError(
f"LoRA adapter '{first_adapter}' was requested, but LoRA is not enabled. "
"Please launch the server with --enable-lora flag and preload adapters "
"using --lora-paths or /load_lora_adapter endpoint."
)
await self._resolve_lora_path(obj)
async def _resolve_lora_path(self, obj: Union[GenerateReqInput, EmbeddingReqInput]):
if isinstance(obj.lora_path, str):
unique_lora_paths = set([obj.lora_path])
@@ -46,6 +46,24 @@ register_cuda_ci(
class TestEmbeddingLoraSupport(unittest.TestCase):
"""Test LoRA support in embedding request structures."""
def test_engine_encode_validates_enable_lora(self):
"""Test Engine.encode() validates enable_lora before processing lora_path."""
# Use a simple non-gated model for this validation test
with SRTRunner(
MODEL_PATH,
torch_dtype=torch.float16,
model_type="embedding",
port=DEFAULT_PORT_FOR_SRT_TEST_RUNNER,
) as runner:
# Should raise ValueError because enable_lora was not set for the server
with self.assertRaises(ValueError) as context:
runner.engine.encode(prompt="Test", lora_path="fake-adapter")
error_msg = str(context.exception)
self.assertIn("not enabled", error_msg.lower())
self.assertIn("--enable-lora", error_msg)
self.assertIn("fake-adapter", error_msg)
def test_embedding_lora_fields(self):
"""Test LoRA fields exist and work correctly across all embedding structures."""
# EmbeddingReqInput: fields exist, normalization expands single to batch, indexing works
@@ -142,44 +142,6 @@ class TestResolveLoraPath(unittest.TestCase):
self.assertEqual(result, "adapter-name")
class TestValidateLoraEnabled(unittest.TestCase):
"""Test _validate_lora_enabled method."""
def test_validation_passes_when_lora_enabled(self):
"""Test validation passes when LoRA is enabled."""
tokenizer_manager = MockTokenizerManager(enable_lora=True)
serving = ConcreteServingBase(tokenizer_manager)
# Should not raise
try:
serving._validate_lora_enabled("sql-expert")
except ValueError:
self.fail("_validate_lora_enabled raised ValueError unexpectedly")
def test_validation_fails_when_lora_disabled(self):
"""Test validation fails with helpful message when LoRA is disabled."""
tokenizer_manager = MockTokenizerManager(enable_lora=False)
serving = ConcreteServingBase(tokenizer_manager)
with self.assertRaises(ValueError) as context:
serving._validate_lora_enabled("sql-expert")
error_message = str(context.exception)
self.assertIn("sql-expert", error_message)
self.assertIn("--enable-lora", error_message)
self.assertIn("not enabled", error_message)
def test_validation_error_mentions_adapter_name(self):
"""Test that error message includes the requested adapter name."""
tokenizer_manager = MockTokenizerManager(enable_lora=False)
serving = ConcreteServingBase(tokenizer_manager)
with self.assertRaises(ValueError) as context:
serving._validate_lora_enabled("my-custom-adapter")
self.assertIn("my-custom-adapter", str(context.exception))
class TestIntegrationScenarios(unittest.TestCase):
"""Integration tests for common usage scenarios."""
@@ -196,9 +158,6 @@ class TestIntegrationScenarios(unittest.TestCase):
lora_path = self.serving._resolve_lora_path(model, explicit_lora)
self.assertEqual(lora_path, "sql-expert")
# Validation should pass
self.serving._validate_lora_enabled(lora_path)
def test_backward_compatible_usage(self):
"""Test backward-compatible usage with explicit lora_path."""
model = "meta-llama/Llama-3.1-8B"
@@ -207,9 +166,6 @@ class TestIntegrationScenarios(unittest.TestCase):
lora_path = self.serving._resolve_lora_path(model, explicit_lora)
self.assertEqual(lora_path, "sql-expert")
# Validation should pass
self.serving._validate_lora_enabled(lora_path)
def test_base_model_usage(self):
"""Test using base model without any adapter."""
model = "meta-llama/Llama-3.1-8B"
@@ -228,10 +184,6 @@ class TestIntegrationScenarios(unittest.TestCase):
lora_path = self.serving._resolve_lora_path(model, explicit_lora)
self.assertEqual(lora_path, explicit_lora)
# Validate first adapter in list
if isinstance(lora_path, list) and lora_path[0]:
self.serving._validate_lora_enabled(lora_path[0])
def test_adapter_in_model_overrides_batch_list(self):
"""Test that adapter in model parameter overrides batch list."""
model = "meta-llama/Llama-3.1-8B:preferred-adapter"
@@ -240,24 +192,6 @@ class TestIntegrationScenarios(unittest.TestCase):
lora_path = self.serving._resolve_lora_path(model, explicit_lora)
self.assertEqual(lora_path, "preferred-adapter")
def test_error_when_lora_not_enabled(self):
"""Test comprehensive error flow when LoRA is not enabled."""
# Setup server without LoRA enabled
tokenizer_manager = MockTokenizerManager(enable_lora=False)
serving = ConcreteServingBase(tokenizer_manager)
# User tries to use adapter
model = "meta-llama/Llama-3.1-8B:sql-expert"
lora_path = serving._resolve_lora_path(model, None)
# Should get helpful error
with self.assertRaises(ValueError) as context:
serving._validate_lora_enabled(lora_path)
error = str(context.exception)
self.assertIn("--enable-lora", error)
self.assertIn("sql-expert", error)
class TestEdgeCases(unittest.TestCase):
"""Test edge cases and error conditions."""
@@ -318,14 +252,6 @@ class TestEdgeCases(unittest.TestCase):
result = self.serving._resolve_lora_path("model-name", "")
self.assertEqual(result, "")
def test_validation_with_empty_adapter_name(self):
"""Test validation with empty adapter name still raises error."""
tokenizer_manager = MockTokenizerManager(enable_lora=False)
serving = ConcreteServingBase(tokenizer_manager)
with self.assertRaises(ValueError):
serving._validate_lora_enabled("")
if __name__ == "__main__":
unittest.main()