@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 114 lines 4.36 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Script to test add_lora, remove_lora, pin_lora, list_loras functions."""
4
5 import pytest
6
7 from vllm.engine.arg_utils import AsyncEngineArgs, EngineArgs
8 from vllm.entrypoints.launchers.api_server.entry import (
9 build_async_engine_client_from_engine_args,
10 )
11 from vllm.lora.request import LoRARequest
12 from vllm.v1.engine.llm_engine import LLMEngine
13
14 MODEL_PATH = "Qwen/Qwen3-0.6B"
15 LORA_MODULE_PATH = "charent/self_cognition_Alice"
16 LORA_RANK = 8
17
18
19 def make_lora_request(lora_id: int):
20 return LoRARequest(
21 lora_name=f"{lora_id}", lora_int_id=lora_id, lora_path=LORA_MODULE_PATH
22 )
23
24
25 def test_lora_functions_sync():
26 max_loras = 4
27 # Create engine in eager-mode. Due to high max_loras, the CI can
28 # OOM during cuda-graph capture.
29 engine_args = EngineArgs(
30 model=MODEL_PATH,
31 enable_lora=True,
32 max_loras=max_loras,
33 max_lora_rank=LORA_RANK,
34 max_model_len=128,
35 gpu_memory_utilization=0.8,
36 enforce_eager=True,
37 )
38
39 llm = LLMEngine.from_engine_args(engine_args)
40
41 def run_check(fn, args, expected: list):
42 fn(args)
43 assert set(llm.list_loras()) == set(expected)
44
45 run_check(llm.add_lora, make_lora_request(1), [1])
46 run_check(llm.add_lora, make_lora_request(2), [1, 2])
47
48 # Pin LoRA 1 and test that it is never removed on subsequent adds.
49 run_check(llm.pin_lora, 1, [1, 2])
50 run_check(llm.add_lora, make_lora_request(3), [1, 2, 3])
51 run_check(llm.add_lora, make_lora_request(4), [1, 2, 3, 4])
52 run_check(llm.add_lora, make_lora_request(5), [1, 5, 3, 4])
53 run_check(llm.add_lora, make_lora_request(6), [1, 5, 6, 4])
54 run_check(llm.add_lora, make_lora_request(7), [1, 5, 6, 7])
55 run_check(llm.add_lora, make_lora_request(8), [1, 8, 6, 7])
56 run_check(llm.add_lora, make_lora_request(9), [1, 8, 9, 7])
57 run_check(llm.add_lora, make_lora_request(10), [1, 8, 9, 10])
58
59 # Remove LoRA 1 and continue adding.
60 run_check(llm.remove_lora, 1, [8, 9, 10])
61 run_check(llm.add_lora, make_lora_request(11), [8, 9, 10, 11])
62 run_check(llm.add_lora, make_lora_request(12), [12, 9, 10, 11])
63 run_check(llm.add_lora, make_lora_request(13), [12, 13, 10, 11])
64
65 # Remove all LoRAs.
66 run_check(llm.remove_lora, 13, [12, 10, 11])
67 run_check(llm.remove_lora, 12, [10, 11])
68 run_check(llm.remove_lora, 11, [10])
69 run_check(llm.remove_lora, 10, [])
70
71
72 @pytest.mark.asyncio
73 async def test_lora_functions_async():
74 max_loras = 4
75 engine_args = AsyncEngineArgs(
76 model=MODEL_PATH,
77 enable_lora=True,
78 max_loras=max_loras,
79 max_lora_rank=LORA_RANK,
80 max_model_len=128,
81 gpu_memory_utilization=0.8,
82 enforce_eager=True,
83 )
84
85 async def run_check(fn, args, expected: list):
86 await fn(args)
87 assert set(await llm.list_loras()) == set(expected)
88
89 async with build_async_engine_client_from_engine_args(engine_args) as llm:
90 await run_check(llm.add_lora, make_lora_request(1), [1])
91 await run_check(llm.add_lora, make_lora_request(2), [1, 2])
92
93 # Pin LoRA 1 and test that it is never removed on subsequent adds.
94 await run_check(llm.pin_lora, 1, [1, 2])
95 await run_check(llm.add_lora, make_lora_request(3), [1, 2, 3])
96 await run_check(llm.add_lora, make_lora_request(4), [1, 2, 3, 4])
97 await run_check(llm.add_lora, make_lora_request(5), [1, 5, 3, 4])
98 await run_check(llm.add_lora, make_lora_request(6), [1, 5, 6, 4])
99 await run_check(llm.add_lora, make_lora_request(7), [1, 5, 6, 7])
100 await run_check(llm.add_lora, make_lora_request(8), [1, 8, 6, 7])
101 await run_check(llm.add_lora, make_lora_request(9), [1, 8, 9, 7])
102 await run_check(llm.add_lora, make_lora_request(10), [1, 8, 9, 10])
103
104 # Remove LoRA 1 and continue adding.
105 await run_check(llm.remove_lora, 1, [8, 9, 10])
106 await run_check(llm.add_lora, make_lora_request(11), [8, 9, 10, 11])
107 await run_check(llm.add_lora, make_lora_request(12), [12, 9, 10, 11])
108 await run_check(llm.add_lora, make_lora_request(13), [12, 13, 10, 11])
109
110 # Remove all LoRAs
111 await run_check(llm.remove_lora, 13, [12, 10, 11])
112 await run_check(llm.remove_lora, 12, [10, 11])
113 await run_check(llm.remove_lora, 11, [10])
114 await run_check(llm.remove_lora, 10, [])