| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import asyncio |
| 4 | import time |
| 5 | |
| 6 | import pytest |
| 7 | |
| 8 | from vllm.engine.arg_utils import AsyncEngineArgs |
| 9 | from vllm.entrypoints.launchers.api_server.entry import ( |
| 10 | build_async_engine_client_from_engine_args, |
| 11 | ) |
| 12 | from vllm.inputs import TextPrompt |
| 13 | from vllm.lora.request import LoRARequest |
| 14 | from vllm.sampling_params import SamplingParams |
| 15 | from vllm.utils.async_utils import merge_async_iterators |
| 16 | |
| 17 | MODEL_PATH = "zai-org/chatglm3-6b" |
| 18 | LORA_RANK = 64 |
| 19 | DEFAULT_MAX_LORAS = 4 * 3 |
| 20 | |
| 21 | |
| 22 | def get_lora_requests(lora_path) -> list[LoRARequest]: |
| 23 | lora_requests: list[LoRARequest] = [ |
| 24 | LoRARequest(lora_name=f"{i}", lora_int_id=i, lora_path=lora_path) |
| 25 | for i in range(1, DEFAULT_MAX_LORAS + 1) |
| 26 | ] |
| 27 | return lora_requests |
| 28 | |
| 29 | |
| 30 | async def requests_processing_time(llm, lora_requests: list[LoRARequest]) -> float: |
| 31 | sampling_params = SamplingParams( |
| 32 | n=1, temperature=0.0, top_p=1.0, ignore_eos=True, max_tokens=1 |
| 33 | ) |
| 34 | |
| 35 | generators = [] |
| 36 | start = time.perf_counter() |
| 37 | |
| 38 | for lora_request in lora_requests: |
| 39 | lora_int_id = lora_request.lora_int_id |
| 40 | generator = llm.generate( |
| 41 | prompt=TextPrompt(prompt=f"hello {lora_int_id}", multi_modal_data=None), # type: ignore |
| 42 | sampling_params=sampling_params, |
| 43 | lora_request=lora_request, |
| 44 | request_id=f"test{lora_int_id}", |
| 45 | ) |
| 46 | generators.append(generator) |
| 47 | |
| 48 | all_gens = merge_async_iterators(*generators) |
| 49 | async for i, res in all_gens: |
| 50 | pass |
| 51 | |
| 52 | end = time.perf_counter() |
| 53 | return end - start |
| 54 | |
| 55 | |
| 56 | @pytest.mark.asyncio |
| 57 | async def test_add_lora(chatglm3_lora_files): |
| 58 | """The add_lora function is used to preload some LoRA adapters into the |
| 59 | engine in anticipation of future requests using these adapters. To test |
| 60 | this functionality, we use the async engine to process some requests - We |
| 61 | do it twice, once with add_lora() preloading and once without. |
| 62 | |
| 63 | We measure the request processing time in both cases and expect the time |
| 64 | to be lesser in the case with add_lora() calls. |
| 65 | """ |
| 66 | lora_requests: list[LoRARequest] = get_lora_requests(chatglm3_lora_files) |
| 67 | |
| 68 | max_loras = len(set([lr.lora_int_id for lr in lora_requests])) |
| 69 | # Create engine in eager-mode. Due to high max_loras, the CI can |
| 70 | # OOM during cuda-graph capture. |
| 71 | engine_args = AsyncEngineArgs( |
| 72 | model=MODEL_PATH, |
| 73 | enable_lora=True, |
| 74 | max_loras=max_loras, |
| 75 | max_lora_rank=LORA_RANK, |
| 76 | max_model_len=128, |
| 77 | gpu_memory_utilization=0.8, # avoid OOM |
| 78 | trust_remote_code=True, |
| 79 | enforce_eager=True, |
| 80 | ) |
| 81 | |
| 82 | # split lora_requests into 3 parts |
| 83 | part_size = len(lora_requests) // 3 |
| 84 | dummy_run_requests = lora_requests[:part_size] |
| 85 | warmup_run_requests = lora_requests[part_size : part_size * 2] |
| 86 | cold_run_requests = lora_requests[part_size * 2 :] |
| 87 | |
| 88 | async with build_async_engine_client_from_engine_args(engine_args) as llm: |
| 89 | # Dummy run - So any 1-time functionality like triton kernel compilation |
| 90 | # is complete here. |
| 91 | await requests_processing_time(llm, dummy_run_requests) |
| 92 | |
| 93 | # Run with warmup |
| 94 | add_lora_tasks = [llm.add_lora(lr) for lr in warmup_run_requests] |
| 95 | add_lora_results = await asyncio.gather(*add_lora_tasks) |
| 96 | |
| 97 | # Test that all all_lora calls are successful. |
| 98 | assert all(add_lora_results) |
| 99 | |
| 100 | time_with_add_lora = await requests_processing_time(llm, warmup_run_requests) |
| 101 | |
| 102 | # Run without any warmup |
| 103 | time_cold_start = await requests_processing_time(llm, cold_run_requests) |
| 104 | |
| 105 | print(f"time hot-start {time_with_add_lora} vs time cold-start {time_cold_start} ") |
| 106 | |
| 107 | assert time_with_add_lora < time_cold_start, ( |
| 108 | f"time_with_add_lora={time_with_add_lora}, " |
| 109 | f"time_cold_start={time_cold_start}" |
| 110 | "The engine request processing time with LoRA pre-loading " |
| 111 | "must be less than the version that does on-demand LoRA loading." |
| 112 | ) |