| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import gc |
| 4 | import os |
| 5 | import tempfile |
| 6 | |
| 7 | import openai |
| 8 | import pytest |
| 9 | import pytest_asyncio |
| 10 | import torch.cuda |
| 11 | |
| 12 | from tests.utils import RemoteOpenAIServer |
| 13 | from vllm.engine.arg_utils import EngineArgs |
| 14 | from vllm.model_executor.model_loader.tensorizer import ( |
| 15 | TensorizerConfig, |
| 16 | tensorize_lora_adapter, |
| 17 | tensorize_vllm_model, |
| 18 | ) |
| 19 | from vllm.platforms import current_platform |
| 20 | |
| 21 | MODEL_NAME = "unsloth/llama-3.2-1b-Instruct" |
| 22 | LORA_PATH = "davzoku/finqa_adapter_1b" |
| 23 | |
| 24 | |
| 25 | def _cleanup(): |
| 26 | gc.collect() |
| 27 | torch.accelerator.empty_cache() |
| 28 | |
| 29 | |
| 30 | @pytest.fixture(autouse=True) |
| 31 | def cleanup(): |
| 32 | _cleanup() |
| 33 | |
| 34 | |
| 35 | @pytest.fixture(scope="module") |
| 36 | def tmp_dir(): |
| 37 | with tempfile.TemporaryDirectory() as path: |
| 38 | yield path |
| 39 | |
| 40 | |
| 41 | @pytest.fixture(scope="module") |
| 42 | def model_uri(tmp_dir): |
| 43 | yield f"{tmp_dir}/model.tensors" |
| 44 | |
| 45 | |
| 46 | @pytest.fixture(scope="module") |
| 47 | def tensorize_model_and_lora(tmp_dir, model_uri): |
| 48 | tensorizer_config = TensorizerConfig(tensorizer_uri=model_uri, lora_dir=tmp_dir) |
| 49 | args = EngineArgs(model=MODEL_NAME) |
| 50 | |
| 51 | tensorize_lora_adapter(LORA_PATH, tensorizer_config) |
| 52 | tensorize_vllm_model(args, tensorizer_config) |
| 53 | |
| 54 | # Manually invoke a _cleanup() here, as the cleanup() |
| 55 | # fixture won't be guaranteed to be called after this |
| 56 | # when this fixture is used for a test |
| 57 | _cleanup() |
| 58 | yield |
| 59 | |
| 60 | |
| 61 | @pytest.fixture(scope="module") |
| 62 | def server(model_uri, tensorize_model_and_lora): |
| 63 | # In this case, model_uri is a directory with a model.tensors |
| 64 | # file and all necessary model artifacts, particularly a |
| 65 | # HF `config.json` file. In this case, Tensorizer can infer the |
| 66 | # `TensorizerConfig` so --model-loader-extra-config can be completely |
| 67 | # omitted. |
| 68 | |
| 69 | ## Start OpenAI API server |
| 70 | args = [ |
| 71 | "--load-format", |
| 72 | "tensorizer", |
| 73 | "--served-model-name", |
| 74 | MODEL_NAME, |
| 75 | "--enable-lora", |
| 76 | ] |
| 77 | if current_platform.is_rocm(): |
| 78 | args += ["--attention-backend", "TRITON_ATTN"] |
| 79 | |
| 80 | model_dir = os.path.dirname(model_uri) |
| 81 | with RemoteOpenAIServer(model_dir, args) as remote_server: |
| 82 | yield remote_server |
| 83 | |
| 84 | |
| 85 | @pytest_asyncio.fixture |
| 86 | async def client(server): |
| 87 | async with server.get_async_client() as async_client: |
| 88 | yield async_client |
| 89 | |
| 90 | |
| 91 | @pytest.mark.asyncio |
| 92 | @pytest.mark.parametrize("model_name", [MODEL_NAME]) |
| 93 | async def test_single_completion(client: openai.AsyncOpenAI, model_name: str): |
| 94 | _cleanup() |
| 95 | completion = await client.completions.create( |
| 96 | model=model_name, prompt="Hello, my name is", max_tokens=5, temperature=0.0 |
| 97 | ) |
| 98 | |
| 99 | assert completion.id is not None |
| 100 | assert completion.choices is not None and len(completion.choices) == 1 |
| 101 | assert completion.model == MODEL_NAME |
| 102 | assert len(completion.choices) == 1 |
| 103 | assert len(completion.choices[0].text) >= 5 |
| 104 | assert completion.choices[0].finish_reason == "length" |
| 105 | assert completion.usage == openai.types.CompletionUsage( |
| 106 | completion_tokens=5, prompt_tokens=6, total_tokens=11 |
| 107 | ) |