| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | """Shared fixtures and utilities for SageMaker tests.""" |
| 5 | |
| 6 | import pytest |
| 7 | import pytest_asyncio |
| 8 | |
| 9 | from tests.utils import RemoteOpenAIServer |
| 10 | |
| 11 | # Model name constants used across tests |
| 12 | MODEL_NAME_SMOLLM = "HuggingFaceTB/SmolLM2-135M-Instruct" |
| 13 | LORA_ADAPTER_NAME_SMOLLM = "jekunz/smollm-135m-lora-fineweb-faroese" |
| 14 | |
| 15 | # SageMaker header constants |
| 16 | HEADER_SAGEMAKER_CLOSED_SESSION_ID = "X-Amzn-SageMaker-Closed-Session-Id" |
| 17 | HEADER_SAGEMAKER_SESSION_ID = "X-Amzn-SageMaker-Session-Id" |
| 18 | HEADER_SAGEMAKER_NEW_SESSION_ID = "X-Amzn-SageMaker-New-Session-Id" |
| 19 | |
| 20 | |
| 21 | @pytest.fixture(scope="session") |
| 22 | def smollm2_lora_files(): |
| 23 | """Download LoRA files once per test session.""" |
| 24 | from vllm.transformers_utils.repo_utils import hf_api |
| 25 | |
| 26 | return hf_api().snapshot_download(repo_id=LORA_ADAPTER_NAME_SMOLLM) |
| 27 | |
| 28 | |
| 29 | @pytest.fixture(scope="module") |
| 30 | def basic_server_with_lora(smollm2_lora_files): |
| 31 | """Basic server fixture with standard configuration.""" |
| 32 | args = [ |
| 33 | "--dtype", |
| 34 | "bfloat16", |
| 35 | "--max-model-len", |
| 36 | "8192", |
| 37 | "--enforce-eager", |
| 38 | # lora config below |
| 39 | "--enable-lora", |
| 40 | "--max-lora-rank", |
| 41 | "256", |
| 42 | "--max-cpu-loras", |
| 43 | "2", |
| 44 | "--max-num-seqs", |
| 45 | "64", |
| 46 | ] |
| 47 | |
| 48 | envs = { |
| 49 | "VLLM_ALLOW_RUNTIME_LORA_UPDATING": "True", |
| 50 | "SAGEMAKER_ENABLE_STATEFUL_SESSIONS": "True", |
| 51 | } |
| 52 | with RemoteOpenAIServer(MODEL_NAME_SMOLLM, args, env_dict=envs) as remote_server: |
| 53 | yield remote_server |
| 54 | |
| 55 | |
| 56 | @pytest_asyncio.fixture |
| 57 | async def async_client(basic_server_with_lora: RemoteOpenAIServer): |
| 58 | """Async OpenAI client fixture for use with basic_server.""" |
| 59 | async with basic_server_with_lora.get_async_client() as async_client: |
| 60 | yield async_client |