| 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 | The handler-override scenarios exercise the real vLLM SageMaker router and |
| 16 | bootstrap path via an in-process FastAPI ``TestClient`` instead of launching a |
| 17 | model server. These scenarios fully replace the ``/ping`` and ``/invocations`` |
| 18 | endpoints with customer handlers, so no inference engine is required to |
| 19 | validate override behavior. Avoiding the model server also keeps the tests |
| 20 | fast and deterministic rather than depending on the FastAPI version resolved |
| 21 | into the test environment at runtime. |
| 22 | """ |
| 23 | |
| 24 | import os |
| 25 | |
| 26 | import pytest |
| 27 | import requests |
| 28 | from fastapi import FastAPI |
| 29 | from fastapi.testclient import TestClient |
| 30 | |
| 31 | from tests.utils import RemoteOpenAIServer |
| 32 | |
| 33 | from .conftest import ( |
| 34 | MODEL_NAME_SMOLLM, |
| 35 | ) |
| 36 | |
| 37 | |
| 38 | def _build_sagemaker_test_client() -> TestClient: |
| 39 | """Build a TestClient over the real SageMaker router and bootstrap path. |
| 40 | |
| 41 | ``attach_router`` is called with empty supported tasks because the override |
| 42 | tests replace the endpoints with customer handlers, so no framework |
| 43 | invocation handler (and therefore no engine) is exercised. |
| 44 | """ |
| 45 | from vllm.entrypoints.serve.sagemaker.api_router import ( |
| 46 | attach_router, |
| 47 | sagemaker_standards_bootstrap, |
| 48 | ) |
| 49 | |
| 50 | app = FastAPI() |
| 51 | attach_router(app, ()) |
| 52 | return TestClient(sagemaker_standards_bootstrap(app)) |
| 53 | |
| 54 | |
| 55 | class TestHandlerOverrideIntegration: |
| 56 | """Integration tests simulating real customer usage scenarios. |
| 57 | |
| 58 | Each test simulates a fresh server startup where customers: |
| 59 | - Use @custom_ping_handler and @custom_invocation_handler decorators |
| 60 | - Set environment variables (CUSTOM_FASTAPI_PING_HANDLER, etc.) |
| 61 | - Write customer scripts with custom_sagemaker_ping_handler() and |
| 62 | custom_sagemaker_invocation_handler() functions |
| 63 | """ |
| 64 | |
| 65 | def setup_method(self): |
| 66 | """Setup for each test - simulate fresh server startup.""" |
| 67 | self._clear_caches() |
| 68 | self._clear_env_vars() |
| 69 | |
| 70 | def teardown_method(self): |
| 71 | """Cleanup after each test.""" |
| 72 | self._clear_env_vars() |
| 73 | |
| 74 | def _clear_caches(self): |
| 75 | """Clear handler registry and function loader cache.""" |
| 76 | try: |
| 77 | from model_hosting_container_standards.common.handler import ( |
| 78 | handler_registry, |
| 79 | ) |
| 80 | from model_hosting_container_standards.sagemaker.sagemaker_loader import ( |
| 81 | SageMakerFunctionLoader, |
| 82 | ) |
| 83 | |
| 84 | handler_registry.clear() |
| 85 | SageMakerFunctionLoader._default_function_loader = None |
| 86 | except ImportError: |
| 87 | pytest.skip("model-hosting-container-standards not available") |
| 88 | |
| 89 | def _clear_env_vars(self): |
| 90 | """Clear SageMaker environment variables.""" |
| 91 | try: |
| 92 | from model_hosting_container_standards.common.fastapi.config import ( |
| 93 | FastAPIEnvVars, |
| 94 | ) |
| 95 | from model_hosting_container_standards.sagemaker.config import ( |
| 96 | SageMakerEnvVars, |
| 97 | ) |
| 98 | |
| 99 | # Clear SageMaker env vars |
| 100 | for var in [ |
| 101 | SageMakerEnvVars.SAGEMAKER_MODEL_PATH, |
| 102 | SageMakerEnvVars.CUSTOM_SCRIPT_FILENAME, |
| 103 | ]: |
| 104 | os.environ.pop(var, None) |
| 105 | |
| 106 | # Clear FastAPI env vars |
| 107 | for var in [ |
| 108 | FastAPIEnvVars.CUSTOM_FASTAPI_PING_HANDLER, |
| 109 | FastAPIEnvVars.CUSTOM_FASTAPI_INVOCATION_HANDLER, |
| 110 | ]: |
| 111 | os.environ.pop(var, None) |
| 112 | except ImportError: |
| 113 | pass |
| 114 | |
| 115 | def test_customer_script_functions_auto_loaded(self, monkeypatch, tmp_path): |
| 116 | """Test customer scenario: script functions automatically override |
| 117 | framework defaults.""" |
| 118 | try: |
| 119 | from model_hosting_container_standards.sagemaker.config import ( |
| 120 | SageMakerEnvVars, |
| 121 | ) |
| 122 | except ImportError: |
| 123 | pytest.skip("model-hosting-container-standards not available") |
| 124 | |
| 125 | # Customer writes a script file with ping() and invoke() functions |
| 126 | script_path = tmp_path / "model.py" |
| 127 | script_path.write_text( |
| 128 | """ |
| 129 | from fastapi import Request |
| 130 | |
| 131 | async def custom_sagemaker_ping_handler(): |
| 132 | return { |
| 133 | "status": "healthy", |
| 134 | "source": "customer_override", |
| 135 | "message": "Custom ping from customer script" |
| 136 | } |
| 137 | |
| 138 | async def custom_sagemaker_invocation_handler(request: Request): |
| 139 | return { |
| 140 | "predictions": ["Custom response from customer script"], |
| 141 | "source": "customer_override" |
| 142 | } |
| 143 | """ |
| 144 | ) |
| 145 | |
| 146 | # Customer sets SageMaker environment variables to point to their script |
| 147 | monkeypatch.setenv(SageMakerEnvVars.SAGEMAKER_MODEL_PATH, str(tmp_path)) |
| 148 | monkeypatch.setenv(SageMakerEnvVars.CUSTOM_SCRIPT_FILENAME, script_path.name) |
| 149 | |
| 150 | with _build_sagemaker_test_client() as client: |
| 151 | # Customer tests their server and sees their overrides work |
| 152 | # automatically |
| 153 | ping_response = client.get("/ping") |
| 154 | assert ping_response.status_code == 200 |
| 155 | ping_data = ping_response.json() |
| 156 | |
| 157 | invoke_response = client.post( |
| 158 | "/invocations", |
| 159 | json={ |
| 160 | "model": MODEL_NAME_SMOLLM, |
| 161 | "messages": [{"role": "user", "content": "Hello"}], |
| 162 | "max_tokens": 5, |
| 163 | }, |
| 164 | ) |
| 165 | assert invoke_response.status_code == 200 |
| 166 | invoke_data = invoke_response.json() |
| 167 | |
| 168 | # Customer sees their functions are used |
| 169 | assert ping_data["source"] == "customer_override" |
| 170 | assert ping_data["message"] == "Custom ping from customer script" |
| 171 | assert invoke_data["source"] == "customer_override" |
| 172 | assert invoke_data["predictions"] == [ |
| 173 | "Custom response from customer script" |
| 174 | ] |
| 175 | |
| 176 | def test_customer_decorator_usage(self, monkeypatch, tmp_path): |
| 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 | script_path = tmp_path / "model.py" |
| 188 | script_path.write_text( |
| 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 | |
| 209 | monkeypatch.setenv(SageMakerEnvVars.SAGEMAKER_MODEL_PATH, str(tmp_path)) |
| 210 | monkeypatch.setenv(SageMakerEnvVars.CUSTOM_SCRIPT_FILENAME, script_path.name) |
| 211 | |
| 212 | with _build_sagemaker_test_client() as client: |
| 213 | ping_response = client.get("/ping") |
| 214 | assert ping_response.status_code == 200 |
| 215 | ping_data = ping_response.json() |
| 216 | |
| 217 | invoke_response = client.post( |
| 218 | "/invocations", |
| 219 | json={ |
| 220 | "model": MODEL_NAME_SMOLLM, |
| 221 | "messages": [{"role": "user", "content": "Hello"}], |
| 222 | "max_tokens": 5, |
| 223 | }, |
| 224 | ) |
| 225 | assert invoke_response.status_code == 200 |
| 226 | invoke_data = invoke_response.json() |
| 227 | |
| 228 | # Customer sees their handlers are used by the server |
| 229 | assert ping_data["source"] == "customer_decorator" |
| 230 | assert invoke_data["source"] == "customer_decorator" |
| 231 | |
| 232 | def test_handler_priority_order(self, monkeypatch, tmp_path): |
| 233 | """Test priority: @custom_ping_handler/@custom_invocation_handler |
| 234 | decorators vs script functions.""" |
| 235 | try: |
| 236 | from model_hosting_container_standards.sagemaker.config import ( |
| 237 | SageMakerEnvVars, |
| 238 | ) |
| 239 | except ImportError: |
| 240 | pytest.skip("model-hosting-container-standards not available") |
| 241 | |
| 242 | # Customer writes a script with both decorator and regular functions |
| 243 | script_path = tmp_path / "model.py" |
| 244 | script_path.write_text( |
| 245 | """ |
| 246 | import model_hosting_container_standards.sagemaker as sagemaker_standards |
| 247 | from fastapi import Request |
| 248 | |
| 249 | # Customer uses @custom_ping_handler decorator (higher priority than script functions) |
| 250 | @sagemaker_standards.custom_ping_handler |
| 251 | async def decorated_ping(): |
| 252 | return { |
| 253 | "status": "healthy", |
| 254 | "source": "ping_decorator_in_script", |
| 255 | "priority": "decorator" |
| 256 | } |
| 257 | |
| 258 | # Customer also has a regular function (lower priority than |
| 259 | # @custom_ping_handler decorator) |
| 260 | async def custom_sagemaker_ping_handler(): |
| 261 | return { |
| 262 | "status": "healthy", |
| 263 | "source": "script_function", |
| 264 | "priority": "function" |
| 265 | } |
| 266 | |
| 267 | # Customer has a regular invoke function |
| 268 | async def custom_sagemaker_invocation_handler(request: Request): |
| 269 | return { |
| 270 | "predictions": ["Script function response"], |
| 271 | "source": "script_invoke_function", |
| 272 | "priority": "function" |
| 273 | } |
| 274 | """ |
| 275 | ) |
| 276 | |
| 277 | monkeypatch.setenv(SageMakerEnvVars.SAGEMAKER_MODEL_PATH, str(tmp_path)) |
| 278 | monkeypatch.setenv(SageMakerEnvVars.CUSTOM_SCRIPT_FILENAME, script_path.name) |
| 279 | |
| 280 | with _build_sagemaker_test_client() as client: |
| 281 | ping_response = client.get("/ping") |
| 282 | assert ping_response.status_code == 200 |
| 283 | ping_data = ping_response.json() |
| 284 | |
| 285 | invoke_response = client.post( |
| 286 | "/invocations", |
| 287 | json={ |
| 288 | "model": MODEL_NAME_SMOLLM, |
| 289 | "messages": [{"role": "user", "content": "Hello"}], |
| 290 | "max_tokens": 5, |
| 291 | }, |
| 292 | ) |
| 293 | assert invoke_response.status_code == 200 |
| 294 | invoke_data = invoke_response.json() |
| 295 | |
| 296 | # @custom_ping_handler decorator has higher priority than |
| 297 | # script function |
| 298 | assert ping_data["source"] == "ping_decorator_in_script" |
| 299 | assert ping_data["priority"] == "decorator" |
| 300 | |
| 301 | # Script function is used for invoke |
| 302 | assert invoke_data["source"] == "script_invoke_function" |
| 303 | assert invoke_data["priority"] == "function" |
| 304 | |
| 305 | def test_environment_variable_script_loading(self, monkeypatch, tmp_path): |
| 306 | """Test that environment variables correctly specify script location |
| 307 | and loading.""" |
| 308 | try: |
| 309 | from model_hosting_container_standards.sagemaker.config import ( |
| 310 | SageMakerEnvVars, |
| 311 | ) |
| 312 | except ImportError: |
| 313 | pytest.skip("model-hosting-container-standards not available") |
| 314 | |
| 315 | # Customer writes a script in a specific directory |
| 316 | script_path = tmp_path / "model.py" |
| 317 | script_path.write_text( |
| 318 | """ |
| 319 | from fastapi import Request |
| 320 | |
| 321 | async def custom_sagemaker_ping_handler(): |
| 322 | return { |
| 323 | "status": "healthy", |
| 324 | "source": "env_loaded_script", |
| 325 | "method": "environment_variable_loading" |
| 326 | } |
| 327 | |
| 328 | async def custom_sagemaker_invocation_handler(request: Request): |
| 329 | return { |
| 330 | "predictions": ["Loaded via environment variables"], |
| 331 | "source": "env_loaded_script", |
| 332 | "method": "environment_variable_loading" |
| 333 | } |
| 334 | """ |
| 335 | ) |
| 336 | |
| 337 | # Test environment variable script loading |
| 338 | monkeypatch.setenv(SageMakerEnvVars.SAGEMAKER_MODEL_PATH, str(tmp_path)) |
| 339 | monkeypatch.setenv(SageMakerEnvVars.CUSTOM_SCRIPT_FILENAME, script_path.name) |
| 340 | |
| 341 | with _build_sagemaker_test_client() as client: |
| 342 | ping_response = client.get("/ping") |
| 343 | assert ping_response.status_code == 200 |
| 344 | ping_data = ping_response.json() |
| 345 | |
| 346 | invoke_response = client.post( |
| 347 | "/invocations", |
| 348 | json={ |
| 349 | "model": MODEL_NAME_SMOLLM, |
| 350 | "messages": [{"role": "user", "content": "Hello"}], |
| 351 | "max_tokens": 5, |
| 352 | }, |
| 353 | ) |
| 354 | assert invoke_response.status_code == 200 |
| 355 | invoke_data = invoke_response.json() |
| 356 | |
| 357 | # Verify that the script was loaded via environment variables |
| 358 | assert ping_data["source"] == "env_loaded_script" |
| 359 | assert ping_data["method"] == "environment_variable_loading" |
| 360 | assert invoke_data["source"] == "env_loaded_script" |
| 361 | assert invoke_data["method"] == "environment_variable_loading" |
| 362 | |
| 363 | @pytest.mark.asyncio |
| 364 | async def test_framework_default_handlers(self): |
| 365 | """Test that framework default handlers work when no customer |
| 366 | overrides exist. |
| 367 | |
| 368 | This scenario exercises the real inference path (default |
| 369 | ``/invocations``), so it keeps using a live model server rather than |
| 370 | the in-process TestClient. |
| 371 | """ |
| 372 | args = [ |
| 373 | "--dtype", |
| 374 | "bfloat16", |
| 375 | "--max-model-len", |
| 376 | "2048", |
| 377 | "--enforce-eager", |
| 378 | "--max-num-seqs", |
| 379 | "32", |
| 380 | ] |
| 381 | |
| 382 | # Explicitly pass empty env_dict to ensure no SageMaker env vars are set |
| 383 | # This prevents pollution from previous tests |
| 384 | try: |
| 385 | from model_hosting_container_standards.common.fastapi.config import ( |
| 386 | FastAPIEnvVars, |
| 387 | ) |
| 388 | from model_hosting_container_standards.sagemaker.config import ( |
| 389 | SageMakerEnvVars, |
| 390 | ) |
| 391 | |
| 392 | env_dict = { |
| 393 | SageMakerEnvVars.SAGEMAKER_MODEL_PATH: "", |
| 394 | SageMakerEnvVars.CUSTOM_SCRIPT_FILENAME: "", |
| 395 | FastAPIEnvVars.CUSTOM_FASTAPI_PING_HANDLER: "", |
| 396 | FastAPIEnvVars.CUSTOM_FASTAPI_INVOCATION_HANDLER: "", |
| 397 | } |
| 398 | except ImportError: |
| 399 | env_dict = {} |
| 400 | |
| 401 | with RemoteOpenAIServer(MODEL_NAME_SMOLLM, args, env_dict=env_dict) as server: |
| 402 | # Test that default ping works |
| 403 | ping_response = requests.get(server.url_for("ping")) |
| 404 | assert ping_response.status_code == 200 |
| 405 | |
| 406 | # Test that default invocations work |
| 407 | invoke_response = requests.post( |
| 408 | server.url_for("invocations"), |
| 409 | json={ |
| 410 | "model": MODEL_NAME_SMOLLM, |
| 411 | "messages": [{"role": "user", "content": "Hello"}], |
| 412 | "max_tokens": 5, |
| 413 | }, |
| 414 | ) |
| 415 | assert invoke_response.status_code == 200 |
| 416 | |
| 417 | def test_handler_env_var_override(self, monkeypatch, tmp_path): |
| 418 | """Test CUSTOM_FASTAPI_PING_HANDLER and CUSTOM_FASTAPI_INVOCATION_HANDLER |
| 419 | environment variable overrides.""" |
| 420 | try: |
| 421 | from model_hosting_container_standards.common.fastapi.config import ( |
| 422 | FastAPIEnvVars, |
| 423 | ) |
| 424 | from model_hosting_container_standards.sagemaker.config import ( |
| 425 | SageMakerEnvVars, |
| 426 | ) |
| 427 | except ImportError: |
| 428 | pytest.skip("model-hosting-container-standards not available") |
| 429 | |
| 430 | # Create a script with both env var handlers and script functions |
| 431 | script_path = tmp_path / "model.py" |
| 432 | script_path.write_text( |
| 433 | """ |
| 434 | from fastapi import Request, Response |
| 435 | import json |
| 436 | |
| 437 | async def env_var_ping_handler(raw_request: Request) -> Response: |
| 438 | return Response( |
| 439 | content=json.dumps({ |
| 440 | "status": "healthy", |
| 441 | "source": "env_var_ping", |
| 442 | "method": "environment_variable" |
| 443 | }), |
| 444 | media_type="application/json" |
| 445 | ) |
| 446 | |
| 447 | async def env_var_invoke_handler(raw_request: Request) -> Response: |
| 448 | return Response( |
| 449 | content=json.dumps({ |
| 450 | "predictions": ["Environment variable response"], |
| 451 | "source": "env_var_invoke", |
| 452 | "method": "environment_variable" |
| 453 | }), |
| 454 | media_type="application/json" |
| 455 | ) |
| 456 | |
| 457 | async def custom_sagemaker_ping_handler(): |
| 458 | return { |
| 459 | "status": "healthy", |
| 460 | "source": "script_ping", |
| 461 | "method": "script_function" |
| 462 | } |
| 463 | |
| 464 | async def custom_sagemaker_invocation_handler(request: Request): |
| 465 | return { |
| 466 | "predictions": ["Script function response"], |
| 467 | "source": "script_invoke", |
| 468 | "method": "script_function" |
| 469 | } |
| 470 | """ |
| 471 | ) |
| 472 | |
| 473 | # Set environment variables to override both handlers |
| 474 | monkeypatch.setenv(SageMakerEnvVars.SAGEMAKER_MODEL_PATH, str(tmp_path)) |
| 475 | monkeypatch.setenv(SageMakerEnvVars.CUSTOM_SCRIPT_FILENAME, script_path.name) |
| 476 | monkeypatch.setenv( |
| 477 | FastAPIEnvVars.CUSTOM_FASTAPI_PING_HANDLER, |
| 478 | f"{script_path.name}:env_var_ping_handler", |
| 479 | ) |
| 480 | monkeypatch.setenv( |
| 481 | FastAPIEnvVars.CUSTOM_FASTAPI_INVOCATION_HANDLER, |
| 482 | f"{script_path.name}:env_var_invoke_handler", |
| 483 | ) |
| 484 | |
| 485 | with _build_sagemaker_test_client() as client: |
| 486 | # Test ping handler override |
| 487 | ping_response = client.get("/ping") |
| 488 | assert ping_response.status_code == 200 |
| 489 | ping_data = ping_response.json() |
| 490 | |
| 491 | # Environment variable should override script function |
| 492 | assert ping_data["method"] == "environment_variable" |
| 493 | assert ping_data["source"] == "env_var_ping" |
| 494 | |
| 495 | # Test invocation handler override |
| 496 | invoke_response = client.post( |
| 497 | "/invocations", |
| 498 | json={ |
| 499 | "model": MODEL_NAME_SMOLLM, |
| 500 | "messages": [{"role": "user", "content": "Hello"}], |
| 501 | "max_tokens": 5, |
| 502 | }, |
| 503 | ) |
| 504 | assert invoke_response.status_code == 200 |
| 505 | invoke_data = invoke_response.json() |
| 506 | |
| 507 | # Environment variable should override script function |
| 508 | assert invoke_data["method"] == "environment_variable" |
| 509 | assert invoke_data["source"] == "env_var_invoke" |
| 510 | |
| 511 | def test_env_var_priority_over_decorator_and_script(self, monkeypatch, tmp_path): |
| 512 | """Test that environment variables have highest priority over decorators |
| 513 | and script functions for both ping and invocation handlers.""" |
| 514 | try: |
| 515 | from model_hosting_container_standards.common.fastapi.config import ( |
| 516 | FastAPIEnvVars, |
| 517 | ) |
| 518 | from model_hosting_container_standards.sagemaker.config import ( |
| 519 | SageMakerEnvVars, |
| 520 | ) |
| 521 | except ImportError: |
| 522 | pytest.skip("model-hosting-container-standards not available") |
| 523 | |
| 524 | # Create a script with all three handler types for both ping and invocation |
| 525 | script_path = tmp_path / "model.py" |
| 526 | script_path.write_text( |
| 527 | """ |
| 528 | import model_hosting_container_standards.sagemaker as sagemaker_standards |
| 529 | from fastapi import Request, Response |
| 530 | import json |
| 531 | |
| 532 | # Environment variable handlers (highest priority) |
| 533 | async def env_priority_ping(raw_request: Request) -> Response: |
| 534 | return Response( |
| 535 | content=json.dumps({ |
| 536 | "status": "healthy", |
| 537 | "source": "env_var", |
| 538 | "priority": "environment_variable" |
| 539 | }), |
| 540 | media_type="application/json" |
| 541 | ) |
| 542 | |
| 543 | async def env_priority_invoke(raw_request: Request) -> Response: |
| 544 | return Response( |
| 545 | content=json.dumps({ |
| 546 | "predictions": ["Environment variable response"], |
| 547 | "source": "env_var", |
| 548 | "priority": "environment_variable" |
| 549 | }), |
| 550 | media_type="application/json" |
| 551 | ) |
| 552 | |
| 553 | # Decorator handlers (medium priority) |
| 554 | @sagemaker_standards.custom_ping_handler |
| 555 | async def decorator_ping(raw_request: Request) -> Response: |
| 556 | return Response( |
| 557 | content=json.dumps({ |
| 558 | "status": "healthy", |
| 559 | "source": "decorator", |
| 560 | "priority": "decorator" |
| 561 | }), |
| 562 | media_type="application/json" |
| 563 | ) |
| 564 | |
| 565 | @sagemaker_standards.custom_invocation_handler |
| 566 | async def decorator_invoke(raw_request: Request) -> Response: |
| 567 | return Response( |
| 568 | content=json.dumps({ |
| 569 | "predictions": ["Decorator response"], |
| 570 | "source": "decorator", |
| 571 | "priority": "decorator" |
| 572 | }), |
| 573 | media_type="application/json" |
| 574 | ) |
| 575 | |
| 576 | # Script functions (lowest priority) |
| 577 | async def custom_sagemaker_ping_handler(): |
| 578 | return { |
| 579 | "status": "healthy", |
| 580 | "source": "script", |
| 581 | "priority": "script_function" |
| 582 | } |
| 583 | |
| 584 | async def custom_sagemaker_invocation_handler(request: Request): |
| 585 | return { |
| 586 | "predictions": ["Script function response"], |
| 587 | "source": "script", |
| 588 | "priority": "script_function" |
| 589 | } |
| 590 | """ |
| 591 | ) |
| 592 | |
| 593 | # Set environment variables to specify highest priority handlers |
| 594 | monkeypatch.setenv(SageMakerEnvVars.SAGEMAKER_MODEL_PATH, str(tmp_path)) |
| 595 | monkeypatch.setenv(SageMakerEnvVars.CUSTOM_SCRIPT_FILENAME, script_path.name) |
| 596 | monkeypatch.setenv( |
| 597 | FastAPIEnvVars.CUSTOM_FASTAPI_PING_HANDLER, |
| 598 | f"{script_path.name}:env_priority_ping", |
| 599 | ) |
| 600 | monkeypatch.setenv( |
| 601 | FastAPIEnvVars.CUSTOM_FASTAPI_INVOCATION_HANDLER, |
| 602 | f"{script_path.name}:env_priority_invoke", |
| 603 | ) |
| 604 | |
| 605 | with _build_sagemaker_test_client() as client: |
| 606 | # Test ping handler priority |
| 607 | ping_response = client.get("/ping") |
| 608 | assert ping_response.status_code == 200 |
| 609 | ping_data = ping_response.json() |
| 610 | |
| 611 | # Environment variable has highest priority and should be used |
| 612 | assert ping_data["priority"] == "environment_variable" |
| 613 | assert ping_data["source"] == "env_var" |
| 614 | |
| 615 | # Test invocation handler priority |
| 616 | invoke_response = client.post( |
| 617 | "/invocations", |
| 618 | json={ |
| 619 | "model": MODEL_NAME_SMOLLM, |
| 620 | "messages": [{"role": "user", "content": "Hello"}], |
| 621 | "max_tokens": 5, |
| 622 | }, |
| 623 | ) |
| 624 | assert invoke_response.status_code == 200 |
| 625 | invoke_data = invoke_response.json() |
| 626 | |
| 627 | # Environment variable has highest priority and should be used |
| 628 | assert invoke_data["priority"] == "environment_variable" |
| 629 | assert invoke_data["source"] == "env_var" |