@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / bca74e32b7
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 13, 13:15 UTC

[Frontend] Add sagemaker_standards dynamic lora adapter and stateful session management decorators to vLLM OpenAI API server (#27892)

Signed-off-by: Zuyi Zhao <zhaozuy@amazon.com> Signed-off-by: Shen Teng <sheteng@amazon.com> Co-authored-by: Shen Teng <sheteng@amazon.com> Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com>

Zuyi Zhao committed Nov 10, 2025 at 20:57 UTC bca74e32b7ef03515cda508ba88151e2e547bdc9
11 files changed +1613 -83
requirements/common.txt
+1
@@ -49,3 +49,4 @@ cbor2 # Required for cross-language serialization of hashable objects
49 setproctitle # Used to set process names for better debugging and monitoring
50 openai-harmony >= 0.0.3 # Required for gpt-oss
51 anthropic == 0.71.0
52 +model-hosting-container-standards < 1.0.0
\ No newline at end of file
tests/entrypoints/sagemaker/__init__.py
tests/entrypoints/sagemaker/conftest.py new
+58
@@ -0,0 +1,58 @@
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 ...utils import RemoteOpenAIServer
10 +
11 +# Model name constants used across tests
12 +MODEL_NAME_ZEPHYR = "HuggingFaceH4/zephyr-7b-beta"
13 +MODEL_NAME_SMOLLM = "HuggingFaceTB/SmolLM2-135M-Instruct"
14 +LORA_ADAPTER_NAME_SMOLLM = "jekunz/smollm-135m-lora-fineweb-faroese"
15 +
16 +# SageMaker header constants
17 +HEADER_SAGEMAKER_CLOSED_SESSION_ID = "X-Amzn-SageMaker-Closed-Session-Id"
18 +HEADER_SAGEMAKER_SESSION_ID = "X-Amzn-SageMaker-Session-Id"
19 +HEADER_SAGEMAKER_NEW_SESSION_ID = "X-Amzn-SageMaker-New-Session-Id"
20 +
21 +
22 +@pytest.fixture(scope="session")
23 +def smollm2_lora_files():
24 + """Download LoRA files once per test session."""
25 + from huggingface_hub import snapshot_download
26 +
27 + return snapshot_download(repo_id=LORA_ADAPTER_NAME_SMOLLM)
28 +
29 +
30 +@pytest.fixture(scope="module")
31 +def basic_server_with_lora(smollm2_lora_files):
32 + """Basic server fixture with standard configuration."""
33 + args = [
34 + "--dtype",
35 + "bfloat16",
36 + "--max-model-len",
37 + "8192",
38 + "--enforce-eager",
39 + # lora config below
40 + "--enable-lora",
41 + "--max-lora-rank",
42 + "256",
43 + "--max-cpu-loras",
44 + "2",
45 + "--max-num-seqs",
46 + "64",
47 + ]
48 +
49 + envs = {"VLLM_ALLOW_RUNTIME_LORA_UPDATING": "True"}
50 + with RemoteOpenAIServer(MODEL_NAME_SMOLLM, args, env_dict=envs) as remote_server:
51 + yield remote_server
52 +
53 +
54 +@pytest_asyncio.fixture
55 +async def async_client(basic_server_with_lora: RemoteOpenAIServer):
56 + """Async OpenAI client fixture for use with basic_server."""
57 + async with basic_server_with_lora.get_async_client() as async_client:
58 + yield async_client
tests/entrypoints/sagemaker/test_sagemaker_handler_overrides.py new
+734
@@ -0,0 +1,734 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +"""Integration tests for handler override functionality.
5 +
6 +Tests real customer usage scenarios:
7 +- Using @custom_ping_handler and @custom_invocation_handler decorators
8 + to override handlers
9 +- Setting environment variables for handler specifications
10 +- Writing customer scripts with custom_sagemaker_ping_handler() and
11 + custom_sagemaker_invocation_handler() functions
12 +- Priority: env vars > decorators > customer script files > framework
13 + defaults
14 +
15 +Note: These tests focus on validating server responses rather than directly calling
16 +get_ping_handler() and get_invoke_handler() to ensure full integration testing.
17 +"""
18 +
19 +import os
20 +import tempfile
21 +
22 +import pytest
23 +import requests
24 +
25 +from ...utils import RemoteOpenAIServer
26 +from .conftest import (
27 + MODEL_NAME_SMOLLM,
28 +)
29 +
30 +
31 +class TestHandlerOverrideIntegration:
32 + """Integration tests simulating real customer usage scenarios.
33 +
34 + Each test simulates a fresh server startup where customers:
35 + - Use @custom_ping_handler and @custom_invocation_handler decorators
36 + - Set environment variables (CUSTOM_FASTAPI_PING_HANDLER, etc.)
37 + - Write customer scripts with custom_sagemaker_ping_handler() and
38 + custom_sagemaker_invocation_handler() functions
39 + """
40 +
41 + def setup_method(self):
42 + """Setup for each test - simulate fresh server startup."""
43 + self._clear_caches()
44 + self._clear_env_vars()
45 +
46 + def teardown_method(self):
47 + """Cleanup after each test."""
48 + self._clear_env_vars()
49 +
50 + def _clear_caches(self):
51 + """Clear handler registry and function loader cache."""
52 + try:
53 + from model_hosting_container_standards.common.handler import (
54 + handler_registry,
55 + )
56 + from model_hosting_container_standards.sagemaker.sagemaker_loader import (
57 + SageMakerFunctionLoader,
58 + )
59 +
60 + handler_registry.clear()
61 + SageMakerFunctionLoader._default_function_loader = None
62 + except ImportError:
63 + pytest.skip("model-hosting-container-standards not available")
64 +
65 + def _clear_env_vars(self):
66 + """Clear SageMaker environment variables."""
67 + try:
68 + from model_hosting_container_standards.common.fastapi.config import (
69 + FastAPIEnvVars,
70 + )
71 + from model_hosting_container_standards.sagemaker.config import (
72 + SageMakerEnvVars,
73 + )
74 +
75 + # Clear SageMaker env vars
76 + for var in [
77 + SageMakerEnvVars.SAGEMAKER_MODEL_PATH,
78 + SageMakerEnvVars.CUSTOM_SCRIPT_FILENAME,
79 + ]:
80 + os.environ.pop(var, None)
81 +
82 + # Clear FastAPI env vars
83 + for var in [
84 + FastAPIEnvVars.CUSTOM_FASTAPI_PING_HANDLER,
85 + FastAPIEnvVars.CUSTOM_FASTAPI_INVOCATION_HANDLER,
86 + ]:
87 + os.environ.pop(var, None)
88 + except ImportError:
89 + pass
90 +
91 + @pytest.mark.asyncio
92 + async def test_customer_script_functions_auto_loaded(self):
93 + """Test customer scenario: script functions automatically override
94 + framework defaults."""
95 + try:
96 + from model_hosting_container_standards.sagemaker.config import (
97 + SageMakerEnvVars,
98 + )
99 + except ImportError:
100 + pytest.skip("model-hosting-container-standards not available")
101 +
102 + # Customer writes a script file with ping() and invoke() functions
103 + with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as f:
104 + f.write(
105 + """
106 +from fastapi import Request
107 +
108 +async def custom_sagemaker_ping_handler():
109 + return {
110 + "status": "healthy",
111 + "source": "customer_override",
112 + "message": "Custom ping from customer script"
113 + }
114 +
115 +async def custom_sagemaker_invocation_handler(request: Request):
116 + return {
117 + "predictions": ["Custom response from customer script"],
118 + "source": "customer_override"
119 + }
120 +"""
121 + )
122 + script_path = f.name
123 +
124 + try:
125 + script_dir = os.path.dirname(script_path)
126 + script_name = os.path.basename(script_path)
127 +
128 + # Customer sets SageMaker environment variables to point to their script
129 + env_vars = {
130 + SageMakerEnvVars.SAGEMAKER_MODEL_PATH: script_dir,
131 + SageMakerEnvVars.CUSTOM_SCRIPT_FILENAME: script_name,
132 + }
133 +
134 + args = [
135 + "--dtype",
136 + "bfloat16",
137 + "--max-model-len",
138 + "2048",
139 + "--enforce-eager",
140 + "--max-num-seqs",
141 + "32",
142 + ]
143 +
144 + with RemoteOpenAIServer(
145 + MODEL_NAME_SMOLLM, args, env_dict=env_vars
146 + ) as server:
147 + # Customer tests their server and sees their overrides work
148 + # automatically
149 + ping_response = requests.get(server.url_for("ping"))
150 + assert ping_response.status_code == 200
151 + ping_data = ping_response.json()
152 +
153 + invoke_response = requests.post(
154 + server.url_for("invocations"),
155 + json={
156 + "model": MODEL_NAME_SMOLLM,
157 + "messages": [{"role": "user", "content": "Hello"}],
158 + "max_tokens": 5,
159 + },
160 + )
161 + assert invoke_response.status_code == 200
162 + invoke_data = invoke_response.json()
163 +
164 + # Customer sees their functions are used
165 + assert ping_data["source"] == "customer_override"
166 + assert ping_data["message"] == "Custom ping from customer script"
167 + assert invoke_data["source"] == "customer_override"
168 + assert invoke_data["predictions"] == [
169 + "Custom response from customer script"
170 + ]
171 +
172 + finally:
173 + os.unlink(script_path)
174 +
175 + @pytest.mark.asyncio
176 + async def test_customer_decorator_usage(self):
177 + """Test customer scenario: using @custom_ping_handler and
178 + @custom_invocation_handler decorators."""
179 + try:
180 + from model_hosting_container_standards.sagemaker.config import (
181 + SageMakerEnvVars,
182 + )
183 + except ImportError:
184 + pytest.skip("model-hosting-container-standards not available")
185 +
186 + # Customer writes a script file with decorators
187 + with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as f:
188 + f.write(
189 + """
190 +import model_hosting_container_standards.sagemaker as sagemaker_standards
191 +from fastapi import Request
192 +
193 +@sagemaker_standards.custom_ping_handler
194 +async def my_ping():
195 + return {
196 + "type": "ping",
197 + "source": "customer_decorator"
198 + }
199 +
200 +@sagemaker_standards.custom_invocation_handler
201 +async def my_invoke(request: Request):
202 + return {
203 + "type": "invoke",
204 + "source": "customer_decorator"
205 + }
206 +"""
207 + )
208 + script_path = f.name
209 +
210 + try:
211 + script_dir = os.path.dirname(script_path)
212 + script_name = os.path.basename(script_path)
213 +
214 + env_vars = {
215 + SageMakerEnvVars.SAGEMAKER_MODEL_PATH: script_dir,
216 + SageMakerEnvVars.CUSTOM_SCRIPT_FILENAME: script_name,
217 + }
218 +
219 + args = [
220 + "--dtype",
221 + "bfloat16",
222 + "--max-model-len",
223 + "2048",
224 + "--enforce-eager",
225 + "--max-num-seqs",
226 + "32",
227 + ]
228 +
229 + with RemoteOpenAIServer(
230 + MODEL_NAME_SMOLLM, args, env_dict=env_vars
231 + ) as server:
232 + ping_response = requests.get(server.url_for("ping"))
233 + assert ping_response.status_code == 200
234 + ping_data = ping_response.json()
235 +
236 + invoke_response = requests.post(
237 + server.url_for("invocations"),
238 + json={
239 + "model": MODEL_NAME_SMOLLM,
240 + "messages": [{"role": "user", "content": "Hello"}],
241 + "max_tokens": 5,
242 + },
243 + )
244 + assert invoke_response.status_code == 200
245 + invoke_data = invoke_response.json()
246 +
247 + # Customer sees their handlers are used by the server
248 + assert ping_data["source"] == "customer_decorator"
249 + assert invoke_data["source"] == "customer_decorator"
250 +
251 + finally:
252 + os.unlink(script_path)
253 +
254 + @pytest.mark.asyncio
255 + async def test_handler_priority_order(self):
256 + """Test priority: @custom_ping_handler/@custom_invocation_handler
257 + decorators vs script functions."""
258 + try:
259 + from model_hosting_container_standards.sagemaker.config import (
260 + SageMakerEnvVars,
261 + )
262 + except ImportError:
263 + pytest.skip("model-hosting-container-standards not available")
264 +
265 + # Customer writes a script with both decorator and regular functions
266 + with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as f:
267 + f.write(
268 + """
269 +import model_hosting_container_standards.sagemaker as sagemaker_standards
270 +from fastapi import Request
271 +
272 +# Customer uses @custom_ping_handler decorator (higher priority than script functions)
273 +@sagemaker_standards.custom_ping_handler
274 +async def decorated_ping():
275 + return {
276 + "status": "healthy",
277 + "source": "ping_decorator_in_script",
278 + "priority": "decorator"
279 + }
280 +
281 +# Customer also has a regular function (lower priority than
282 +# @custom_ping_handler decorator)
283 +async def custom_sagemaker_ping_handler():
284 + return {
285 + "status": "healthy",
286 + "source": "script_function",
287 + "priority": "function"
288 + }
289 +
290 +# Customer has a regular invoke function
291 +async def custom_sagemaker_invocation_handler(request: Request):
292 + return {
293 + "predictions": ["Script function response"],
294 + "source": "script_invoke_function",
295 + "priority": "function"
296 + }
297 +"""
298 + )
299 + script_path = f.name
300 +
301 + try:
302 + script_dir = os.path.dirname(script_path)
303 + script_name = os.path.basename(script_path)
304 +
305 + env_vars = {
306 + SageMakerEnvVars.SAGEMAKER_MODEL_PATH: script_dir,
307 + SageMakerEnvVars.CUSTOM_SCRIPT_FILENAME: script_name,
308 + }
309 +
310 + args = [
311 + "--dtype",
312 + "bfloat16",
313 + "--max-model-len",
314 + "2048",
315 + "--enforce-eager",
316 + "--max-num-seqs",
317 + "32",
318 + ]
319 +
320 + with RemoteOpenAIServer(
321 + MODEL_NAME_SMOLLM, args, env_dict=env_vars
322 + ) as server:
323 + ping_response = requests.get(server.url_for("ping"))
324 + assert ping_response.status_code == 200
325 + ping_data = ping_response.json()
326 +
327 + invoke_response = requests.post(
328 + server.url_for("invocations"),
329 + json={
330 + "model": MODEL_NAME_SMOLLM,
331 + "messages": [{"role": "user", "content": "Hello"}],
332 + "max_tokens": 5,
333 + },
334 + )
335 + assert invoke_response.status_code == 200
336 + invoke_data = invoke_response.json()
337 +
338 + # @custom_ping_handler decorator has higher priority than
339 + # script function
340 + assert ping_data["source"] == "ping_decorator_in_script"
341 + assert ping_data["priority"] == "decorator"
342 +
343 + # Script function is used for invoke
344 + assert invoke_data["source"] == "script_invoke_function"
345 + assert invoke_data["priority"] == "function"
346 +
347 + finally:
348 + os.unlink(script_path)
349 +
350 + @pytest.mark.asyncio
351 + async def test_environment_variable_script_loading(self):
352 + """Test that environment variables correctly specify script location
353 + and loading."""
354 + try:
355 + from model_hosting_container_standards.sagemaker.config import (
356 + SageMakerEnvVars,
357 + )
358 + except ImportError:
359 + pytest.skip("model-hosting-container-standards not available")
360 +
361 + # Customer writes a script in a specific directory
362 + with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as f:
363 + f.write(
364 + """
365 +from fastapi import Request
366 +
367 +async def custom_sagemaker_ping_handler():
368 + return {
369 + "status": "healthy",
370 + "source": "env_loaded_script",
371 + "method": "environment_variable_loading"
372 + }
373 +
374 +async def custom_sagemaker_invocation_handler(request: Request):
375 + return {
376 + "predictions": ["Loaded via environment variables"],
377 + "source": "env_loaded_script",
378 + "method": "environment_variable_loading"
379 + }
380 +"""
381 + )
382 + script_path = f.name
383 +
384 + try:
385 + script_dir = os.path.dirname(script_path)
386 + script_name = os.path.basename(script_path)
387 +
388 + # Test environment variable script loading
389 + env_vars = {
390 + SageMakerEnvVars.SAGEMAKER_MODEL_PATH: script_dir,
391 + SageMakerEnvVars.CUSTOM_SCRIPT_FILENAME: script_name,
392 + }
393 +
394 + args = [
395 + "--dtype",
396 + "bfloat16",
397 + "--max-model-len",
398 + "2048",
399 + "--enforce-eager",
400 + "--max-num-seqs",
401 + "32",
402 + ]
403 +
404 + with RemoteOpenAIServer(
405 + MODEL_NAME_SMOLLM, args, env_dict=env_vars
406 + ) as server:
407 + ping_response = requests.get(server.url_for("ping"))
408 + assert ping_response.status_code == 200
409 + ping_data = ping_response.json()
410 +
411 + invoke_response = requests.post(
412 + server.url_for("invocations"),
413 + json={
414 + "model": MODEL_NAME_SMOLLM,
415 + "messages": [{"role": "user", "content": "Hello"}],
416 + "max_tokens": 5,
417 + },
418 + )
419 + assert invoke_response.status_code == 200
420 + invoke_data = invoke_response.json()
421 +
422 + # Verify that the script was loaded via environment variables
423 + assert ping_data["source"] == "env_loaded_script"
424 + assert ping_data["method"] == "environment_variable_loading"
425 + assert invoke_data["source"] == "env_loaded_script"
426 + assert invoke_data["method"] == "environment_variable_loading"
427 +
428 + finally:
429 + os.unlink(script_path)
430 +
431 + @pytest.mark.asyncio
432 + async def test_framework_default_handlers(self):
433 + """Test that framework default handlers work when no customer
434 + overrides exist."""
435 + args = [
436 + "--dtype",
437 + "bfloat16",
438 + "--max-model-len",
439 + "2048",
440 + "--enforce-eager",
441 + "--max-num-seqs",
442 + "32",
443 + ]
444 +
445 + # Explicitly pass empty env_dict to ensure no SageMaker env vars are set
446 + # This prevents pollution from previous tests
447 + try:
448 + from model_hosting_container_standards.common.fastapi.config import (
449 + FastAPIEnvVars,
450 + )
451 + from model_hosting_container_standards.sagemaker.config import (
452 + SageMakerEnvVars,
453 + )
454 +
455 + env_dict = {
456 + SageMakerEnvVars.SAGEMAKER_MODEL_PATH: "",
457 + SageMakerEnvVars.CUSTOM_SCRIPT_FILENAME: "",
458 + FastAPIEnvVars.CUSTOM_FASTAPI_PING_HANDLER: "",
459 + FastAPIEnvVars.CUSTOM_FASTAPI_INVOCATION_HANDLER: "",
460 + }
461 + except ImportError:
462 + env_dict = {}
463 +
464 + with RemoteOpenAIServer(MODEL_NAME_SMOLLM, args, env_dict=env_dict) as server:
465 + # Test that default ping works
466 + ping_response = requests.get(server.url_for("ping"))
467 + assert ping_response.status_code == 200
468 +
469 + # Test that default invocations work
470 + invoke_response = requests.post(
471 + server.url_for("invocations"),
472 + json={
473 + "model": MODEL_NAME_SMOLLM,
474 + "messages": [{"role": "user", "content": "Hello"}],
475 + "max_tokens": 5,
476 + },
477 + )
478 + assert invoke_response.status_code == 200
479 +
480 + @pytest.mark.asyncio
481 + async def test_handler_env_var_override(self):
482 + """Test CUSTOM_FASTAPI_PING_HANDLER and CUSTOM_FASTAPI_INVOCATION_HANDLER
483 + environment variable overrides."""
484 + try:
485 + from model_hosting_container_standards.common.fastapi.config import (
486 + FastAPIEnvVars,
487 + )
488 + from model_hosting_container_standards.sagemaker.config import (
489 + SageMakerEnvVars,
490 + )
491 + except ImportError:
492 + pytest.skip("model-hosting-container-standards not available")
493 +
494 + # Create a script with both env var handlers and script functions
495 + with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as f:
496 + f.write(
497 + """
498 +from fastapi import Request, Response
499 +import json
500 +
501 +async def env_var_ping_handler(raw_request: Request) -> Response:
502 + return Response(
503 + content=json.dumps({
504 + "status": "healthy",
505 + "source": "env_var_ping",
506 + "method": "environment_variable"
507 + }),
508 + media_type="application/json"
509 + )
510 +
511 +async def env_var_invoke_handler(raw_request: Request) -> Response:
512 + return Response(
513 + content=json.dumps({
514 + "predictions": ["Environment variable response"],
515 + "source": "env_var_invoke",
516 + "method": "environment_variable"
517 + }),
518 + media_type="application/json"
519 + )
520 +
521 +async def custom_sagemaker_ping_handler():
522 + return {
523 + "status": "healthy",
524 + "source": "script_ping",
525 + "method": "script_function"
526 + }
527 +
528 +async def custom_sagemaker_invocation_handler(request: Request):
529 + return {
530 + "predictions": ["Script function response"],
531 + "source": "script_invoke",
532 + "method": "script_function"
533 + }
534 +"""
535 + )
536 + script_path = f.name
537 +
538 + try:
539 + script_dir = os.path.dirname(script_path)
540 + script_name = os.path.basename(script_path)
541 +
542 + # Set environment variables to override both handlers
543 + env_vars = {
544 + SageMakerEnvVars.SAGEMAKER_MODEL_PATH: script_dir,
545 + SageMakerEnvVars.CUSTOM_SCRIPT_FILENAME: script_name,
546 + FastAPIEnvVars.CUSTOM_FASTAPI_PING_HANDLER: (
547 + f"{script_name}:env_var_ping_handler"
548 + ),
549 + FastAPIEnvVars.CUSTOM_FASTAPI_INVOCATION_HANDLER: (
550 + f"{script_name}:env_var_invoke_handler"
551 + ),
552 + }
553 +
554 + args = [
555 + "--dtype",
556 + "bfloat16",
557 + "--max-model-len",
558 + "2048",
559 + "--enforce-eager",
560 + "--max-num-seqs",
561 + "32",
562 + ]
563 +
564 + with RemoteOpenAIServer(
565 + MODEL_NAME_SMOLLM, args, env_dict=env_vars
566 + ) as server:
567 + # Test ping handler override
568 + ping_response = requests.get(server.url_for("ping"))
569 + assert ping_response.status_code == 200
570 + ping_data = ping_response.json()
571 +
572 + # Environment variable should override script function
573 + assert ping_data["method"] == "environment_variable"
574 + assert ping_data["source"] == "env_var_ping"
575 +
576 + # Test invocation handler override
577 + invoke_response = requests.post(
578 + server.url_for("invocations"),
579 + json={
580 + "model": MODEL_NAME_SMOLLM,
581 + "messages": [{"role": "user", "content": "Hello"}],
582 + "max_tokens": 5,
583 + },
584 + )
585 + assert invoke_response.status_code == 200
586 + invoke_data = invoke_response.json()
587 +
588 + # Environment variable should override script function
589 + assert invoke_data["method"] == "environment_variable"
590 + assert invoke_data["source"] == "env_var_invoke"
591 +
592 + finally:
593 + os.unlink(script_path)
594 +
595 + @pytest.mark.asyncio
596 + async def test_env_var_priority_over_decorator_and_script(self):
597 + """Test that environment variables have highest priority over decorators
598 + and script functions for both ping and invocation handlers."""
599 + try:
600 + from model_hosting_container_standards.common.fastapi.config import (
601 + FastAPIEnvVars,
602 + )
603 + from model_hosting_container_standards.sagemaker.config import (
604 + SageMakerEnvVars,
605 + )
606 + except ImportError:
607 + pytest.skip("model-hosting-container-standards not available")
608 +
609 + # Create a script with all three handler types for both ping and invocation
610 + with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as f:
611 + f.write(
612 + """
613 +import model_hosting_container_standards.sagemaker as sagemaker_standards
614 +from fastapi import Request, Response
615 +import json
616 +
617 +# Environment variable handlers (highest priority)
618 +async def env_priority_ping(raw_request: Request) -> Response:
619 + return Response(
620 + content=json.dumps({
621 + "status": "healthy",
622 + "source": "env_var",
623 + "priority": "environment_variable"
624 + }),
625 + media_type="application/json"
626 + )
627 +
628 +async def env_priority_invoke(raw_request: Request) -> Response:
629 + return Response(
630 + content=json.dumps({
631 + "predictions": ["Environment variable response"],
632 + "source": "env_var",
633 + "priority": "environment_variable"
634 + }),
635 + media_type="application/json"
636 + )
637 +
638 +# Decorator handlers (medium priority)
639 +@sagemaker_standards.custom_ping_handler
640 +async def decorator_ping(raw_request: Request) -> Response:
641 + return Response(
642 + content=json.dumps({
643 + "status": "healthy",
644 + "source": "decorator",
645 + "priority": "decorator"
646 + }),
647 + media_type="application/json"
648 + )
649 +
650 +@sagemaker_standards.custom_invocation_handler
651 +async def decorator_invoke(raw_request: Request) -> Response:
652 + return Response(
653 + content=json.dumps({
654 + "predictions": ["Decorator response"],
655 + "source": "decorator",
656 + "priority": "decorator"
657 + }),
658 + media_type="application/json"
659 + )
660 +
661 +# Script functions (lowest priority)
662 +async def custom_sagemaker_ping_handler():
663 + return {
664 + "status": "healthy",
665 + "source": "script",
666 + "priority": "script_function"
667 + }
668 +
669 +async def custom_sagemaker_invocation_handler(request: Request):
670 + return {
671 + "predictions": ["Script function response"],
672 + "source": "script",
673 + "priority": "script_function"
674 + }
675 +"""
676 + )
677 + script_path = f.name
678 +
679 + try:
680 + script_dir = os.path.dirname(script_path)
681 + script_name = os.path.basename(script_path)
682 +
683 + # Set environment variables to specify highest priority handlers
684 + env_vars = {
685 + SageMakerEnvVars.SAGEMAKER_MODEL_PATH: script_dir,
686 + SageMakerEnvVars.CUSTOM_SCRIPT_FILENAME: script_name,
687 + FastAPIEnvVars.CUSTOM_FASTAPI_PING_HANDLER: (
688 + f"{script_name}:env_priority_ping"
689 + ),
690 + FastAPIEnvVars.CUSTOM_FASTAPI_INVOCATION_HANDLER: (
691 + f"{script_name}:env_priority_invoke"
692 + ),
693 + }
694 +
695 + args = [
696 + "--dtype",
697 + "bfloat16",
698 + "--max-model-len",
699 + "2048",
700 + "--enforce-eager",
701 + "--max-num-seqs",
702 + "32",
703 + ]
704 +
705 + with RemoteOpenAIServer(
706 + MODEL_NAME_SMOLLM, args, env_dict=env_vars
707 + ) as server:
708 + # Test ping handler priority
709 + ping_response = requests.get(server.url_for("ping"))
710 + assert ping_response.status_code == 200
711 + ping_data = ping_response.json()
712 +
713 + # Environment variable has highest priority and should be used
714 + assert ping_data["priority"] == "environment_variable"
715 + assert ping_data["source"] == "env_var"
716 +
717 + # Test invocation handler priority
718 + invoke_response = requests.post(
719 + server.url_for("invocations"),
720 + json={
721 + "model": MODEL_NAME_SMOLLM,
722 + "messages": [{"role": "user", "content": "Hello"}],
723 + "max_tokens": 5,
724 + },
725 + )
726 + assert invoke_response.status_code == 200
727 + invoke_data = invoke_response.json()
728 +
729 + # Environment variable has highest priority and should be used
730 + assert invoke_data["priority"] == "environment_variable"
731 + assert invoke_data["source"] == "env_var"
732 +
733 + finally:
734 + os.unlink(script_path)
tests/entrypoints/sagemaker/test_sagemaker_lora_adapters.py new
+171
@@ -0,0 +1,171 @@
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 ...utils import RemoteOpenAIServer
8 +from .conftest import MODEL_NAME_SMOLLM
9 +
10 +
11 +@pytest.mark.asyncio
12 +async def test_sagemaker_load_adapter_happy_path(
13 + async_client: openai.AsyncOpenAI,
14 + basic_server_with_lora: RemoteOpenAIServer,
15 + smollm2_lora_files,
16 +):
17 + # The SageMaker standards library creates a POST /adapters endpoint
18 + # that maps to the load_lora_adapter handler with request shape:
19 + # {"lora_name": "body.name", "lora_path": "body.src"}
20 + load_response = requests.post(
21 + basic_server_with_lora.url_for("adapters"),
22 + json={"name": "smollm2-lora-sagemaker", "src": smollm2_lora_files},
23 + )
24 + load_response.raise_for_status()
25 +
26 + models = await async_client.models.list()
27 + models = models.data
28 + dynamic_lora_model = models[-1]
29 + assert dynamic_lora_model.root == smollm2_lora_files
30 + assert dynamic_lora_model.parent == MODEL_NAME_SMOLLM
31 + assert dynamic_lora_model.id == "smollm2-lora-sagemaker"
32 +
33 +
34 +@pytest.mark.asyncio
35 +async def test_sagemaker_unload_adapter_happy_path(
36 + async_client: openai.AsyncOpenAI,
37 + basic_server_with_lora: RemoteOpenAIServer,
38 + smollm2_lora_files,
39 +):
40 + # First, load an adapter
41 + adapter_name = "smollm2-lora-sagemaker-unload"
42 + load_response = requests.post(
43 + basic_server_with_lora.url_for("adapters"),
44 + json={"name": adapter_name, "src": smollm2_lora_files},
45 + )
46 + load_response.raise_for_status()
47 +
48 + # Verify it's in the models list
49 + models = await async_client.models.list()
50 + adapter_ids = [model.id for model in models.data]
51 + assert adapter_name in adapter_ids
52 +
53 + # Now unload it using DELETE /adapters/{adapter_name}
54 + # The SageMaker standards maps this to unload_lora_adapter with:
55 + # {"lora_name": "path_params.adapter_name"}
56 + unload_response = requests.delete(
57 + basic_server_with_lora.url_for("adapters", adapter_name),
58 + )
59 + unload_response.raise_for_status()
60 +
61 + # Verify it's no longer in the models list
62 + models = await async_client.models.list()
63 + adapter_ids = [model.id for model in models.data]
64 + assert adapter_name not in adapter_ids
65 +
66 +
67 +@pytest.mark.asyncio
68 +async def test_sagemaker_load_adapter_not_found(
69 + basic_server_with_lora: RemoteOpenAIServer,
70 +):
71 + load_response = requests.post(
72 + basic_server_with_lora.url_for("adapters"),
73 + json={"name": "nonexistent-adapter", "src": "/path/does/not/exist"},
74 + )
75 + assert load_response.status_code == 404
76 +
77 +
78 +@pytest.mark.asyncio
79 +async def test_sagemaker_load_adapter_invalid_files(
80 + basic_server_with_lora: RemoteOpenAIServer,
81 + tmp_path,
82 +):
83 + invalid_files = tmp_path / "invalid_adapter"
84 + invalid_files.mkdir()
85 + (invalid_files / "adapter_config.json").write_text("not valid json")
86 +
87 + load_response = requests.post(
88 + basic_server_with_lora.url_for("adapters"),
89 + json={"name": "invalid-adapter", "src": str(invalid_files)},
90 + )
91 + assert load_response.status_code == 400
92 +
93 +
94 +@pytest.mark.asyncio
95 +async def test_sagemaker_unload_nonexistent_adapter(
96 + basic_server_with_lora: RemoteOpenAIServer,
97 +):
98 + # Attempt to unload an adapter that doesn't exist
99 + unload_response = requests.delete(
100 + basic_server_with_lora.url_for("adapters", "nonexistent-adapter-name"),
101 + )
102 + assert unload_response.status_code in (400, 404)
103 +
104 +
105 +@pytest.mark.asyncio
106 +async def test_sagemaker_invocations_with_adapter(
107 + basic_server_with_lora: RemoteOpenAIServer,
108 + smollm2_lora_files,
109 +):
110 + # First, load an adapter via SageMaker endpoint
111 + adapter_name = "smollm2-lora-invoke-test"
112 + load_response = requests.post(
113 + basic_server_with_lora.url_for("adapters"),
114 + json={"name": adapter_name, "src": smollm2_lora_files},
115 + )
116 + load_response.raise_for_status()
117 +
118 + # Now test the /invocations endpoint with the adapter
119 + invocation_response = requests.post(
120 + basic_server_with_lora.url_for("invocations"),
121 + headers={
122 + "X-Amzn-SageMaker-Adapter-Identifier": adapter_name,
123 + },
124 + json={
125 + "prompt": "Hello, how are you?",
126 + "max_tokens": 10,
127 + },
128 + )
129 + invocation_response.raise_for_status()
130 + invocation_output = invocation_response.json()
131 +
132 + # Verify we got a valid completion response
133 + assert "choices" in invocation_output
134 + assert len(invocation_output["choices"]) > 0
135 + assert "text" in invocation_output["choices"][0]
136 +
137 +
138 +@pytest.mark.asyncio
139 +async def test_sagemaker_multiple_adapters_load_unload(
140 + async_client: openai.AsyncOpenAI,
141 + basic_server_with_lora: RemoteOpenAIServer,
142 + smollm2_lora_files,
143 +):
144 + adapter_names = [f"sagemaker-adapter-{i}" for i in range(5)]
145 +
146 + # Load all adapters
147 + for adapter_name in adapter_names:
148 + load_response = requests.post(
149 + basic_server_with_lora.url_for("adapters"),
150 + json={"name": adapter_name, "src": smollm2_lora_files},
151 + )
152 + load_response.raise_for_status()
153 +
154 + # Verify all are in the models list
155 + models = await async_client.models.list()
156 + adapter_ids = [model.id for model in models.data]
157 + for adapter_name in adapter_names:
158 + assert adapter_name in adapter_ids
159 +
160 + # Unload all adapters
161 + for adapter_name in adapter_names:
162 + unload_response = requests.delete(
163 + basic_server_with_lora.url_for("adapters", adapter_name),
164 + )
165 + unload_response.raise_for_status()
166 +
167 + # Verify all are removed from models list
168 + models = await async_client.models.list()
169 + adapter_ids = [model.id for model in models.data]
170 + for adapter_name in adapter_names:
171 + assert adapter_name not in adapter_ids
tests/entrypoints/sagemaker/test_sagemaker_middleware_integration.py new
+346
@@ -0,0 +1,346 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +"""Integration test for middleware loader functionality.
5 +
6 +Tests that customer middlewares get called correctly with a vLLM server.
7 +"""
8 +
9 +import os
10 +import tempfile
11 +
12 +import pytest
13 +import requests
14 +
15 +from ...utils import RemoteOpenAIServer
16 +from .conftest import (
17 + MODEL_NAME_SMOLLM,
18 +)
19 +
20 +
21 +class TestMiddlewareIntegration:
22 + """Integration test for middleware with vLLM server."""
23 +
24 + def setup_method(self):
25 + """Setup for each test - simulate fresh server startup."""
26 + self._clear_caches()
27 +
28 + def _clear_caches(self):
29 + """Clear middleware registry and function loader cache."""
30 + try:
31 + from model_hosting_container_standards.common.fastapi.middleware import (
32 + middleware_registry,
33 + )
34 + from model_hosting_container_standards.common.fastapi.middleware.source.decorator_loader import ( # noqa: E501
35 + decorator_loader,
36 + )
37 + from model_hosting_container_standards.sagemaker.sagemaker_loader import (
38 + SageMakerFunctionLoader,
39 + )
40 +
41 + middleware_registry.clear_middlewares()
42 + decorator_loader.clear()
43 + SageMakerFunctionLoader._default_function_loader = None
44 + except ImportError:
45 + pytest.skip("model-hosting-container-standards not available")
46 +
47 + @pytest.mark.asyncio
48 + async def test_customer_middleware_with_vllm_server(self):
49 + """Test that customer middlewares work with actual vLLM server.
50 +
51 + Tests decorator-based middlewares (@custom_middleware, @input_formatter,
52 + @output_formatter)
53 + on multiple endpoints (chat/completions, invocations).
54 + """
55 + try:
56 + from model_hosting_container_standards.sagemaker.config import (
57 + SageMakerEnvVars,
58 + )
59 + except ImportError:
60 + pytest.skip("model-hosting-container-standards not available")
61 +
62 + # Customer writes a middleware script with multiple decorators
63 + with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as f:
64 + f.write(
65 + """
66 +from model_hosting_container_standards.common.fastapi.middleware import (
67 + custom_middleware, input_formatter, output_formatter
68 +)
69 +
70 +# Global flag to track if input formatter was called
71 +_input_formatter_called = False
72 +
73 +@input_formatter
74 +async def customer_input_formatter(request):
75 + # Process input - mark that input formatter was called
76 + global _input_formatter_called
77 + _input_formatter_called = True
78 + return request
79 +
80 +@custom_middleware("throttle")
81 +async def customer_throttle_middleware(request, call_next):
82 + response = await call_next(request)
83 + response.headers["X-Customer-Throttle"] = "applied"
84 + order = response.headers.get("X-Middleware-Order", "")
85 + response.headers["X-Middleware-Order"] = order + "throttle,"
86 + return response
87 +
88 +@output_formatter
89 +async def customer_output_formatter(response):
90 + global _input_formatter_called
91 + response.headers["X-Customer-Processed"] = "true"
92 + # Since input_formatter and output_formatter are combined into
93 + # pre_post_process middleware,
94 + # if output_formatter is called, input_formatter should have been called too
95 + if _input_formatter_called:
96 + response.headers["X-Input-Formatter-Called"] = "true"
97 + order = response.headers.get("X-Middleware-Order", "")
98 + response.headers["X-Middleware-Order"] = order + "output_formatter,"
99 + return response
100 +"""
101 + )
102 + script_path = f.name
103 +
104 + try:
105 + script_dir = os.path.dirname(script_path)
106 + script_name = os.path.basename(script_path)
107 +
108 + # Set environment variables to point to customer script
109 + env_vars = {
110 + SageMakerEnvVars.SAGEMAKER_MODEL_PATH: script_dir,
111 + SageMakerEnvVars.CUSTOM_SCRIPT_FILENAME: script_name,
112 + }
113 +
114 + args = [
115 + "--dtype",
116 + "bfloat16",
117 + "--max-model-len",
118 + "2048",
119 + "--enforce-eager",
120 + "--max-num-seqs",
121 + "32",
122 + ]
123 +
124 + with RemoteOpenAIServer(
125 + MODEL_NAME_SMOLLM, args, env_dict=env_vars
126 + ) as server:
127 + # Test 1: Middlewares applied to chat/completions endpoint
128 + chat_response = requests.post(
129 + server.url_for("v1/chat/completions"),
130 + json={
131 + "model": MODEL_NAME_SMOLLM,
132 + "messages": [{"role": "user", "content": "Hello"}],
133 + "max_tokens": 5,
134 + "temperature": 0.0,
135 + },
136 + )
137 +
138 + assert chat_response.status_code == 200
139 +
140 + # Verify all middlewares were executed
141 + assert "X-Customer-Throttle" in chat_response.headers
142 + assert chat_response.headers["X-Customer-Throttle"] == "applied"
143 + assert "X-Customer-Processed" in chat_response.headers
144 + assert chat_response.headers["X-Customer-Processed"] == "true"
145 +
146 + # Verify input formatter was called
147 + assert "X-Input-Formatter-Called" in chat_response.headers
148 + assert chat_response.headers["X-Input-Formatter-Called"] == "true"
149 +
150 + # Verify middleware execution order
151 + execution_order = chat_response.headers.get(
152 + "X-Middleware-Order", ""
153 + ).rstrip(",")
154 + order_parts = execution_order.split(",") if execution_order else []
155 + assert "throttle" in order_parts
156 + assert "output_formatter" in order_parts
157 +
158 + # Test 2: Middlewares applied to invocations endpoint
159 + invocations_response = requests.post(
160 + server.url_for("invocations"),
161 + json={
162 + "model": MODEL_NAME_SMOLLM,
163 + "messages": [{"role": "user", "content": "Hello"}],
164 + "max_tokens": 5,
165 + "temperature": 0.0,
166 + },
167 + )
168 +
169 + assert invocations_response.status_code == 200
170 +
171 + # Verify all middlewares were executed
172 + assert "X-Customer-Throttle" in invocations_response.headers
173 + assert invocations_response.headers["X-Customer-Throttle"] == "applied"
174 + assert "X-Customer-Processed" in invocations_response.headers
175 + assert invocations_response.headers["X-Customer-Processed"] == "true"
176 +
177 + # Verify input formatter was called
178 + assert "X-Input-Formatter-Called" in invocations_response.headers
179 + assert (
180 + invocations_response.headers["X-Input-Formatter-Called"] == "true"
181 + )
182 +
183 + finally:
184 + os.unlink(script_path)
185 +
186 + @pytest.mark.asyncio
187 + async def test_middleware_with_ping_endpoint(self):
188 + """Test that middlewares work with SageMaker ping endpoint."""
189 + try:
190 + from model_hosting_container_standards.sagemaker.config import (
191 + SageMakerEnvVars,
192 + )
193 + except ImportError:
194 + pytest.skip("model-hosting-container-standards not available")
195 +
196 + # Customer writes a middleware script
197 + with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as f:
198 + f.write(
199 + """
200 +from model_hosting_container_standards.common.fastapi.middleware import (
201 + custom_middleware
202 +)
203 +
204 +@custom_middleware("pre_post_process")
205 +async def ping_tracking_middleware(request, call_next):
206 + response = await call_next(request)
207 + if request.url.path == "/ping":
208 + response.headers["X-Ping-Tracked"] = "true"
209 + return response
210 +"""
211 + )
212 + script_path = f.name
213 +
214 + try:
215 + script_dir = os.path.dirname(script_path)
216 + script_name = os.path.basename(script_path)
217 +
218 + env_vars = {
219 + SageMakerEnvVars.SAGEMAKER_MODEL_PATH: script_dir,
220 + SageMakerEnvVars.CUSTOM_SCRIPT_FILENAME: script_name,
221 + }
222 +
223 + args = [
224 + "--dtype",
225 + "bfloat16",
226 + "--max-model-len",
227 + "2048",
228 + "--enforce-eager",
229 + "--max-num-seqs",
230 + "32",
231 + ]
232 +
233 + with RemoteOpenAIServer(
234 + MODEL_NAME_SMOLLM, args, env_dict=env_vars
235 + ) as server:
236 + # Test ping endpoint with middleware
237 + response = requests.get(server.url_for("ping"))
238 +
239 + assert response.status_code == 200
240 + assert "X-Ping-Tracked" in response.headers
241 + assert response.headers["X-Ping-Tracked"] == "true"
242 +
243 + finally:
244 + os.unlink(script_path)
245 +
246 + @pytest.mark.asyncio
247 + async def test_middleware_env_var_override(self):
248 + """Test middleware environment variable overrides."""
249 + try:
250 + from model_hosting_container_standards.common.fastapi.config import (
251 + FastAPIEnvVars,
252 + )
253 + from model_hosting_container_standards.sagemaker.config import (
254 + SageMakerEnvVars,
255 + )
256 + except ImportError:
257 + pytest.skip("model-hosting-container-standards not available")
258 +
259 + # Create a script with middleware functions specified via env vars
260 + with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as f:
261 + f.write(
262 + """
263 +from fastapi import Request
264 +
265 +# Global flag to track if pre_process was called
266 +_pre_process_called = False
267 +
268 +async def env_throttle_middleware(request, call_next):
269 + response = await call_next(request)
270 + response.headers["X-Env-Throttle"] = "applied"
271 + return response
272 +
273 +async def env_pre_process(request: Request) -> Request:
274 + # Mark that pre_process was called
275 + global _pre_process_called
276 + _pre_process_called = True
277 + return request
278 +
279 +async def env_post_process(response):
280 + global _pre_process_called
281 + if hasattr(response, 'headers'):
282 + response.headers["X-Env-Post-Process"] = "applied"
283 + # Since pre_process and post_process are combined into
284 + # pre_post_process middleware,
285 + # if post_process is called, pre_process should have been called too
286 + if _pre_process_called:
287 + response.headers["X-Pre-Process-Called"] = "true"
288 + return response
289 +"""
290 + )
291 + script_path = f.name
292 +
293 + try:
294 + script_dir = os.path.dirname(script_path)
295 + script_name = os.path.basename(script_path)
296 +
297 + # Set environment variables for middleware
298 + # Use script_name with .py extension as per plugin example
299 + env_vars = {
300 + SageMakerEnvVars.SAGEMAKER_MODEL_PATH: script_dir,
301 + SageMakerEnvVars.CUSTOM_SCRIPT_FILENAME: script_name,
302 + FastAPIEnvVars.CUSTOM_FASTAPI_MIDDLEWARE_THROTTLE: (
303 + f"{script_name}:env_throttle_middleware"
304 + ),
305 + FastAPIEnvVars.CUSTOM_PRE_PROCESS: f"{script_name}:env_pre_process",
306 + FastAPIEnvVars.CUSTOM_POST_PROCESS: f"{script_name}:env_post_process",
307 + }
308 +
309 + args = [
310 + "--dtype",
311 + "bfloat16",
312 + "--max-model-len",
313 + "2048",
314 + "--enforce-eager",
315 + "--max-num-seqs",
316 + "32",
317 + ]
318 +
319 + with RemoteOpenAIServer(
320 + MODEL_NAME_SMOLLM, args, env_dict=env_vars
321 + ) as server:
322 + response = requests.get(server.url_for("ping"))
323 + assert response.status_code == 200
324 +
325 + # Check if environment variable middleware was applied
326 + headers = response.headers
327 +
328 + # Verify that env var middlewares were applied
329 + assert "X-Env-Throttle" in headers, (
330 + "Throttle middleware should be applied via env var"
331 + )
332 + assert headers["X-Env-Throttle"] == "applied"
333 +
334 + assert "X-Env-Post-Process" in headers, (
335 + "Post-process middleware should be applied via env var"
336 + )
337 + assert headers["X-Env-Post-Process"] == "applied"
338 +
339 + # Verify that pre_process was called
340 + assert "X-Pre-Process-Called" in headers, (
341 + "Pre-process should be called via env var"
342 + )
343 + assert headers["X-Pre-Process-Called"] == "true"
344 +
345 + finally:
346 + os.unlink(script_path)
tests/entrypoints/sagemaker/test_sagemaker_stateful_sessions.py new
+153
@@ -0,0 +1,153 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +
5 +import openai # use the official client for correctness check
6 +import pytest
7 +import requests
8 +
9 +from ...utils import RemoteOpenAIServer
10 +from .conftest import (
11 + HEADER_SAGEMAKER_CLOSED_SESSION_ID,
12 + HEADER_SAGEMAKER_NEW_SESSION_ID,
13 + HEADER_SAGEMAKER_SESSION_ID,
14 + MODEL_NAME_SMOLLM,
15 +)
16 +
17 +CLOSE_BADREQUEST_CASES = [
18 + (
19 + "nonexistent_session_id",
20 + {"session_id": "nonexistent-session-id"},
21 + {},
22 + "session not found",
23 + ),
24 + ("malformed_close_request", {}, {"extra-field": "extra-field-data"}, None),
25 +]
26 +
27 +
28 +@pytest.mark.asyncio
29 +async def test_create_session_badrequest(basic_server_with_lora: RemoteOpenAIServer):
30 + bad_response = requests.post(
31 + basic_server_with_lora.url_for("invocations"),
32 + json={"requestType": "NEW_SESSION", "extra-field": "extra-field-data"},
33 + )
34 +
35 + assert bad_response.status_code == 400
36 +
37 +
38 +@pytest.mark.asyncio
39 +@pytest.mark.parametrize(
40 + "test_name,session_id_change,request_body_change,expected_error",
41 + CLOSE_BADREQUEST_CASES,
42 +)
43 +async def test_close_session_badrequest(
44 + basic_server_with_lora: RemoteOpenAIServer,
45 + test_name: str,
46 + session_id_change: dict[str, str],
47 + request_body_change: dict[str, str],
48 + expected_error: str | None,
49 +):
50 + # first attempt to create a session
51 + url = basic_server_with_lora.url_for("invocations")
52 + create_response = requests.post(url, json={"requestType": "NEW_SESSION"})
53 + create_response.raise_for_status()
54 + valid_session_id, expiration = create_response.headers.get(
55 + HEADER_SAGEMAKER_NEW_SESSION_ID, ""
56 + ).split(";")
57 + assert valid_session_id
58 +
59 + close_request_json = {"requestType": "CLOSE"}
60 + if request_body_change:
61 + close_request_json.update(request_body_change)
62 + bad_session_id = session_id_change.get("session_id")
63 + bad_close_response = requests.post(
64 + url,
65 + headers={HEADER_SAGEMAKER_SESSION_ID: bad_session_id or valid_session_id},
66 + json=close_request_json,
67 + )
68 +
69 + # clean up created session, should succeed
70 + clean_up_response = requests.post(
71 + url,
72 + headers={HEADER_SAGEMAKER_SESSION_ID: valid_session_id},
73 + json={"requestType": "CLOSE"},
74 + )
75 + clean_up_response.raise_for_status()
76 +
77 + assert bad_close_response.status_code == 400
78 + if expected_error:
79 + assert expected_error in bad_close_response.json()["error"]["message"]
80 +
81 +
82 +@pytest.mark.asyncio
83 +async def test_close_session_invalidrequest(
84 + basic_server_with_lora: RemoteOpenAIServer, async_client: openai.AsyncOpenAI
85 +):
86 + # first attempt to create a session
87 + url = basic_server_with_lora.url_for("invocations")
88 + create_response = requests.post(url, json={"requestType": "NEW_SESSION"})
89 + create_response.raise_for_status()
90 + valid_session_id, expiration = create_response.headers.get(
91 + HEADER_SAGEMAKER_NEW_SESSION_ID, ""
92 + ).split(";")
93 + assert valid_session_id
94 +
95 + close_request_json = {"requestType": "CLOSE"}
96 + invalid_close_response = requests.post(
97 + url,
98 + # no headers to specify session_id
99 + json=close_request_json,
100 + )
101 +
102 + # clean up created session, should succeed
103 + clean_up_response = requests.post(
104 + url,
105 + headers={HEADER_SAGEMAKER_SESSION_ID: valid_session_id},
106 + json={"requestType": "CLOSE"},
107 + )
108 + clean_up_response.raise_for_status()
109 +
110 + assert invalid_close_response.status_code == 424
111 + assert "invalid session_id" in invalid_close_response.json()["error"]["message"]
112 +
113 +
114 +@pytest.mark.asyncio
115 +async def test_session(basic_server_with_lora: RemoteOpenAIServer):
116 + # first attempt to create a session
117 + url = basic_server_with_lora.url_for("invocations")
118 + create_response = requests.post(url, json={"requestType": "NEW_SESSION"})
119 + create_response.raise_for_status()
120 + valid_session_id, expiration = create_response.headers.get(
121 + HEADER_SAGEMAKER_NEW_SESSION_ID, ""
122 + ).split(";")
123 + assert valid_session_id
124 +
125 + # test invocation with session id
126 +
127 + request_args = {
128 + "model": MODEL_NAME_SMOLLM,
129 + "prompt": "what is 1+1?",
130 + "max_completion_tokens": 5,
131 + "temperature": 0.0,
132 + "logprobs": False,
133 + }
134 +
135 + invocation_response = requests.post(
136 + basic_server_with_lora.url_for("invocations"),
137 + headers={HEADER_SAGEMAKER_SESSION_ID: valid_session_id},
138 + json=request_args,
139 + )
140 + invocation_response.raise_for_status()
141 +
142 + # close created session, should succeed
143 + close_response = requests.post(
144 + url,
145 + headers={HEADER_SAGEMAKER_SESSION_ID: valid_session_id},
146 + json={"requestType": "CLOSE"},
147 + )
148 + close_response.raise_for_status()
149 +
150 + assert (
151 + close_response.headers.get(HEADER_SAGEMAKER_CLOSED_SESSION_ID)
152 + == valid_session_id
153 + )
vllm/entrypoints/dynamic_lora.py new
+57
@@ -0,0 +1,57 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import model_hosting_container_standards.sagemaker as sagemaker_standards
4 +from fastapi import APIRouter, Depends, Request
5 +from fastapi.responses import JSONResponse, Response
6 +
7 +from vllm.entrypoints.openai.api_server import models, validate_json_request
8 +from vllm.entrypoints.openai.protocol import (
9 + ErrorResponse,
10 + LoadLoRAAdapterRequest,
11 + UnloadLoRAAdapterRequest,
12 +)
13 +from vllm.entrypoints.openai.serving_models import OpenAIServingModels
14 +from vllm.logger import init_logger
15 +
16 +logger = init_logger(__name__)
17 +
18 +
19 +def register_dynamic_lora_routes(router: APIRouter):
20 + @sagemaker_standards.register_load_adapter_handler(
21 + request_shape={
22 + "lora_name": "body.name",
23 + "lora_path": "body.src",
24 + },
25 + )
26 + @router.post("/v1/load_lora_adapter", dependencies=[Depends(validate_json_request)])
27 + async def load_lora_adapter(request: LoadLoRAAdapterRequest, raw_request: Request):
28 + handler: OpenAIServingModels = models(raw_request)
29 + response = await handler.load_lora_adapter(request)
30 + if isinstance(response, ErrorResponse):
31 + return JSONResponse(
32 + content=response.model_dump(), status_code=response.error.code
33 + )
34 +
35 + return Response(status_code=200, content=response)
36 +
37 + @sagemaker_standards.register_unload_adapter_handler(
38 + request_shape={
39 + "lora_name": "path_params.adapter_name",
40 + }
41 + )
42 + @router.post(
43 + "/v1/unload_lora_adapter", dependencies=[Depends(validate_json_request)]
44 + )
45 + async def unload_lora_adapter(
46 + request: UnloadLoRAAdapterRequest, raw_request: Request
47 + ):
48 + handler: OpenAIServingModels = models(raw_request)
49 + response = await handler.unload_lora_adapter(request)
50 + if isinstance(response, ErrorResponse):
51 + return JSONResponse(
52 + content=response.model_dump(), status_code=response.error.code
53 + )
54 +
55 + return Response(status_code=200, content=response)
56 +
57 + return router
vllm/entrypoints/openai/api_server.py
+17 -83
@@ -19,6 +19,7 @@ from contextlib import asynccontextmanager
19 from http import HTTPStatus
20 from typing import Annotated, Any, Literal
21
22 +import model_hosting_container_standards.sagemaker as sagemaker_standards
23 import prometheus_client
24 import pydantic
25 import regex as re
@@ -65,7 +66,6 @@ from vllm.entrypoints.openai.protocol import (
66 ErrorInfo,
67 ErrorResponse,
68 IOProcessorResponse,
68 - LoadLoRAAdapterRequest,
69 PoolingBytesResponse,
70 PoolingRequest,
71 PoolingResponse,
@@ -82,7 +82,6 @@ from vllm.entrypoints.openai.protocol import (
82 TranscriptionResponse,
83 TranslationRequest,
84 TranslationResponse,
85 - UnloadLoRAAdapterRequest,
85 )
86 from vllm.entrypoints.openai.serving_chat import OpenAIServingChat
87 from vllm.entrypoints.openai.serving_classification import ServingClassification
@@ -387,13 +386,6 @@ async def get_server_load_metrics(request: Request):
386 return JSONResponse(content={"server_load": request.app.state.server_load_metrics})
387
388
390 -@router.get("/ping", response_class=Response)
391 -@router.post("/ping", response_class=Response)
392 -async def ping(raw_request: Request) -> Response:
393 - """Ping check. Endpoint required for SageMaker"""
394 - return await health(raw_request)
395 -
396 -
389 @router.post(
390 "/tokenize",
391 dependencies=[Depends(validate_json_request)],
@@ -1236,47 +1228,6 @@ INVOCATION_VALIDATORS = [
1228 ]
1229
1230
1239 -@router.post(
1240 - "/invocations",
1241 - dependencies=[Depends(validate_json_request)],
1242 - responses={
1243 - HTTPStatus.BAD_REQUEST.value: {"model": ErrorResponse},
1244 - HTTPStatus.UNSUPPORTED_MEDIA_TYPE.value: {"model": ErrorResponse},
1245 - HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse},
1246 - },
1247 -)
1248 -async def invocations(raw_request: Request):
1249 - """For SageMaker, routes requests based on the request type."""
1250 - try:
1251 - body = await raw_request.json()
1252 - except json.JSONDecodeError as e:
1253 - raise HTTPException(
1254 - status_code=HTTPStatus.BAD_REQUEST.value, detail=f"JSON decode error: {e}"
1255 - ) from e
1256 -
1257 - valid_endpoints = [
1258 - (validator, endpoint)
1259 - for validator, (get_handler, endpoint) in INVOCATION_VALIDATORS
1260 - if get_handler(raw_request) is not None
1261 - ]
1262 -
1263 - for request_validator, endpoint in valid_endpoints:
1264 - try:
1265 - request = request_validator.validate_python(body)
1266 - except pydantic.ValidationError:
1267 - continue
1268 -
1269 - return await endpoint(request, raw_request)
1270 -
1271 - type_names = [
1272 - t.__name__ if isinstance(t := validator._type, type) else str(t)
1273 - for validator, _ in valid_endpoints
1274 - ]
1275 - msg = f"Cannot find suitable handler for request. Expected one of: {type_names}"
1276 - res = base(raw_request).create_error_response(message=msg)
1277 - return JSONResponse(content=res.model_dump(), status_code=res.error.code)
1278 -
1279 -
1231 if envs.VLLM_TORCH_PROFILER_DIR:
1232 logger.warning_once(
1233 "Torch Profiler is enabled in the API server. This should ONLY be "
@@ -1304,39 +1255,6 @@ if envs.VLLM_TORCH_PROFILER_DIR or envs.VLLM_TORCH_CUDA_PROFILE:
1255 return Response(status_code=200)
1256
1257
1307 -if envs.VLLM_ALLOW_RUNTIME_LORA_UPDATING:
1308 - logger.warning(
1309 - "LoRA dynamic loading & unloading is enabled in the API server. "
1310 - "This should ONLY be used for local development!"
1311 - )
1312 -
1313 - @router.post("/v1/load_lora_adapter", dependencies=[Depends(validate_json_request)])
1314 - async def load_lora_adapter(request: LoadLoRAAdapterRequest, raw_request: Request):
1315 - handler = models(raw_request)
1316 - response = await handler.load_lora_adapter(request)
1317 - if isinstance(response, ErrorResponse):
1318 - return JSONResponse(
1319 - content=response.model_dump(), status_code=response.error.code
1320 - )
1321 -
1322 - return Response(status_code=200, content=response)
1323 -
1324 - @router.post(
1325 - "/v1/unload_lora_adapter", dependencies=[Depends(validate_json_request)]
1326 - )
1327 - async def unload_lora_adapter(
1328 - request: UnloadLoRAAdapterRequest, raw_request: Request
1329 - ):
1330 - handler = models(raw_request)
1331 - response = await handler.unload_lora_adapter(request)
1332 - if isinstance(response, ErrorResponse):
1333 - return JSONResponse(
1334 - content=response.model_dump(), status_code=response.error.code
1335 - )
1336 -
1337 - return Response(status_code=200, content=response)
1338 -
1339 -
1258 def load_log_config(log_config_file: str | None) -> dict | None:
1259 if not log_config_file:
1260 return None
@@ -1606,6 +1524,20 @@ def build_app(args: Namespace) -> FastAPI:
1524 )
1525 else:
1526 app = FastAPI(lifespan=lifespan)
1527 +
1528 + if envs.VLLM_ALLOW_RUNTIME_LORA_UPDATING:
1529 + logger.warning(
1530 + "LoRA dynamic loading & unloading is enabled in the API server. "
1531 + "This should ONLY be used for local development!"
1532 + )
1533 + from vllm.entrypoints.dynamic_lora import register_dynamic_lora_routes
1534 +
1535 + register_dynamic_lora_routes(router)
1536 +
1537 + from vllm.entrypoints.sagemaker.routes import register_sagemaker_routes
1538 +
1539 + register_sagemaker_routes(router)
1540 +
1541 app.include_router(router)
1542 app.root_path = args.root_path
1543
@@ -1696,6 +1628,8 @@ def build_app(args: Namespace) -> FastAPI:
1628 f"Invalid middleware {middleware}. Must be a function or a class."
1629 )
1630
1631 + app = sagemaker_standards.bootstrap(app)
1632 +
1633 return app
1634
1635
vllm/entrypoints/sagemaker/__init__.py new
+4
@@ -0,0 +1,4 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +"""SageMaker-specific integration for vLLM."""
vllm/entrypoints/sagemaker/routes.py new
+72
@@ -0,0 +1,72 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import json
4 +from http import HTTPStatus
5 +
6 +import model_hosting_container_standards.sagemaker as sagemaker_standards
7 +import pydantic
8 +from fastapi import APIRouter, Depends, HTTPException, Request
9 +from fastapi.responses import JSONResponse, Response
10 +
11 +from vllm.entrypoints.openai.api_server import (
12 + INVOCATION_VALIDATORS,
13 + base,
14 + health,
15 + validate_json_request,
16 +)
17 +from vllm.entrypoints.openai.protocol import ErrorResponse
18 +
19 +
20 +def register_sagemaker_routes(router: APIRouter):
21 + @router.post("/ping", response_class=Response)
22 + @router.get("/ping", response_class=Response)
23 + @sagemaker_standards.register_ping_handler
24 + async def ping(raw_request: Request) -> Response:
25 + """Ping check. Endpoint required for SageMaker"""
26 + return await health(raw_request)
27 +
28 + @router.post(
29 + "/invocations",
30 + dependencies=[Depends(validate_json_request)],
31 + responses={
32 + HTTPStatus.BAD_REQUEST.value: {"model": ErrorResponse},
33 + HTTPStatus.UNSUPPORTED_MEDIA_TYPE.value: {"model": ErrorResponse},
34 + HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse},
35 + },
36 + )
37 + @sagemaker_standards.register_invocation_handler
38 + @sagemaker_standards.stateful_session_manager()
39 + @sagemaker_standards.inject_adapter_id(adapter_path="model")
40 + async def invocations(raw_request: Request):
41 + """For SageMaker, routes requests based on the request type."""
42 + try:
43 + body = await raw_request.json()
44 + except json.JSONDecodeError as e:
45 + raise HTTPException(
46 + status_code=HTTPStatus.BAD_REQUEST.value,
47 + detail=f"JSON decode error: {e}",
48 + ) from e
49 +
50 + valid_endpoints = [
51 + (validator, endpoint)
52 + for validator, (get_handler, endpoint) in INVOCATION_VALIDATORS
53 + if get_handler(raw_request) is not None
54 + ]
55 +
56 + for request_validator, endpoint in valid_endpoints:
57 + try:
58 + request = request_validator.validate_python(body)
59 + except pydantic.ValidationError:
60 + continue
61 +
62 + return await endpoint(request, raw_request)
63 +
64 + type_names = [
65 + t.__name__ if isinstance(t := validator._type, type) else str(t)
66 + for validator, _ in valid_endpoints
67 + ]
68 + msg = f"Cannot find suitable handler for request. Expected one of: {type_names}"
69 + res = base(raw_request).create_error_response(message=msg)
70 + return JSONResponse(content=res.model_dump(), status_code=res.error.code)
71 +
72 + return router