| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import openai # use the official async_client for correctness check |
| 4 | import pytest |
| 5 | import requests |
| 6 | |
| 7 | from tests.utils import RemoteOpenAIServer |
| 8 | |
| 9 | from .conftest import MODEL_NAME_SMOLLM |
| 10 | |
| 11 | |
| 12 | @pytest.mark.asyncio |
| 13 | async def test_sagemaker_load_adapter_happy_path( |
| 14 | async_client: openai.AsyncOpenAI, |
| 15 | basic_server_with_lora: RemoteOpenAIServer, |
| 16 | smollm2_lora_files, |
| 17 | ): |
| 18 | # The SageMaker standards library creates a POST /adapters endpoint |
| 19 | # that maps to the load_lora_adapter handler with request shape: |
| 20 | # {"lora_name": "body.name", "lora_path": "body.src"} |
| 21 | load_response = requests.post( |
| 22 | basic_server_with_lora.url_for("adapters"), |
| 23 | json={"name": "smollm2-lora-sagemaker", "src": smollm2_lora_files}, |
| 24 | ) |
| 25 | load_response.raise_for_status() |
| 26 | |
| 27 | models = await async_client.models.list() |
| 28 | models = models.data |
| 29 | dynamic_lora_model = models[-1] |
| 30 | assert dynamic_lora_model.root == smollm2_lora_files |
| 31 | assert dynamic_lora_model.parent == MODEL_NAME_SMOLLM |
| 32 | assert dynamic_lora_model.id == "smollm2-lora-sagemaker" |
| 33 | |
| 34 | |
| 35 | @pytest.mark.asyncio |
| 36 | async def test_sagemaker_unload_adapter_happy_path( |
| 37 | async_client: openai.AsyncOpenAI, |
| 38 | basic_server_with_lora: RemoteOpenAIServer, |
| 39 | smollm2_lora_files, |
| 40 | ): |
| 41 | # First, load an adapter |
| 42 | adapter_name = "smollm2-lora-sagemaker-unload" |
| 43 | load_response = requests.post( |
| 44 | basic_server_with_lora.url_for("adapters"), |
| 45 | json={"name": adapter_name, "src": smollm2_lora_files}, |
| 46 | ) |
| 47 | load_response.raise_for_status() |
| 48 | |
| 49 | # Verify it's in the models list |
| 50 | models = await async_client.models.list() |
| 51 | adapter_ids = [model.id for model in models.data] |
| 52 | assert adapter_name in adapter_ids |
| 53 | |
| 54 | # Now unload it using DELETE /adapters/{adapter_name} |
| 55 | # The SageMaker standards maps this to unload_lora_adapter with: |
| 56 | # {"lora_name": "path_params.adapter_name"} |
| 57 | unload_response = requests.delete( |
| 58 | basic_server_with_lora.url_for("adapters", adapter_name), |
| 59 | ) |
| 60 | unload_response.raise_for_status() |
| 61 | |
| 62 | # Verify it's no longer in the models list |
| 63 | models = await async_client.models.list() |
| 64 | adapter_ids = [model.id for model in models.data] |
| 65 | assert adapter_name not in adapter_ids |
| 66 | |
| 67 | |
| 68 | @pytest.mark.asyncio |
| 69 | async def test_sagemaker_load_adapter_not_found( |
| 70 | basic_server_with_lora: RemoteOpenAIServer, |
| 71 | ): |
| 72 | load_response = requests.post( |
| 73 | basic_server_with_lora.url_for("adapters"), |
| 74 | json={"name": "nonexistent-adapter", "src": "/path/does/not/exist"}, |
| 75 | ) |
| 76 | assert load_response.status_code == 404 |
| 77 | |
| 78 | |
| 79 | @pytest.mark.asyncio |
| 80 | async def test_sagemaker_load_adapter_invalid_files( |
| 81 | basic_server_with_lora: RemoteOpenAIServer, |
| 82 | tmp_path, |
| 83 | ): |
| 84 | invalid_files = tmp_path / "invalid_adapter" |
| 85 | invalid_files.mkdir() |
| 86 | (invalid_files / "adapter_config.json").write_text("not valid json") |
| 87 | |
| 88 | load_response = requests.post( |
| 89 | basic_server_with_lora.url_for("adapters"), |
| 90 | json={"name": "invalid-adapter", "src": str(invalid_files)}, |
| 91 | ) |
| 92 | assert load_response.status_code == 500 |
| 93 | |
| 94 | |
| 95 | @pytest.mark.asyncio |
| 96 | async def test_sagemaker_unload_nonexistent_adapter( |
| 97 | basic_server_with_lora: RemoteOpenAIServer, |
| 98 | ): |
| 99 | # Attempt to unload an adapter that doesn't exist |
| 100 | unload_response = requests.delete( |
| 101 | basic_server_with_lora.url_for("adapters", "nonexistent-adapter-name"), |
| 102 | ) |
| 103 | assert unload_response.status_code in (400, 404) |
| 104 | |
| 105 | |
| 106 | @pytest.mark.asyncio |
| 107 | async def test_sagemaker_invocations_with_adapter( |
| 108 | basic_server_with_lora: RemoteOpenAIServer, |
| 109 | smollm2_lora_files, |
| 110 | ): |
| 111 | # First, load an adapter via SageMaker endpoint |
| 112 | adapter_name = "smollm2-lora-invoke-test" |
| 113 | load_response = requests.post( |
| 114 | basic_server_with_lora.url_for("adapters"), |
| 115 | json={"name": adapter_name, "src": smollm2_lora_files}, |
| 116 | ) |
| 117 | load_response.raise_for_status() |
| 118 | |
| 119 | # Now test the /invocations endpoint with the adapter |
| 120 | invocation_response = requests.post( |
| 121 | basic_server_with_lora.url_for("invocations"), |
| 122 | headers={ |
| 123 | "X-Amzn-SageMaker-Adapter-Identifier": adapter_name, |
| 124 | }, |
| 125 | json={ |
| 126 | "prompt": "Hello, how are you?", |
| 127 | "max_tokens": 10, |
| 128 | }, |
| 129 | ) |
| 130 | invocation_response.raise_for_status() |
| 131 | invocation_output = invocation_response.json() |
| 132 | |
| 133 | # Verify we got a valid completion response |
| 134 | assert "choices" in invocation_output |
| 135 | assert len(invocation_output["choices"]) > 0 |
| 136 | assert "text" in invocation_output["choices"][0] |
| 137 | |
| 138 | |
| 139 | @pytest.mark.asyncio |
| 140 | async def test_sagemaker_multiple_adapters_load_unload( |
| 141 | async_client: openai.AsyncOpenAI, |
| 142 | basic_server_with_lora: RemoteOpenAIServer, |
| 143 | smollm2_lora_files, |
| 144 | ): |
| 145 | adapter_names = [f"sagemaker-adapter-{i}" for i in range(5)] |
| 146 | |
| 147 | # Load all adapters |
| 148 | for adapter_name in adapter_names: |
| 149 | load_response = requests.post( |
| 150 | basic_server_with_lora.url_for("adapters"), |
| 151 | json={"name": adapter_name, "src": smollm2_lora_files}, |
| 152 | ) |
| 153 | load_response.raise_for_status() |
| 154 | |
| 155 | # Verify all are in the models list |
| 156 | models = await async_client.models.list() |
| 157 | adapter_ids = [model.id for model in models.data] |
| 158 | for adapter_name in adapter_names: |
| 159 | assert adapter_name in adapter_ids |
| 160 | |
| 161 | # Unload all adapters |
| 162 | for adapter_name in adapter_names: |
| 163 | unload_response = requests.delete( |
| 164 | basic_server_with_lora.url_for("adapters", adapter_name), |
| 165 | ) |
| 166 | unload_response.raise_for_status() |
| 167 | |
| 168 | # Verify all are removed from models list |
| 169 | models = await async_client.models.list() |
| 170 | adapter_ids = [model.id for model in models.data] |
| 171 | for adapter_name in adapter_names: |
| 172 | assert adapter_name not in adapter_ids |