wehub-resource-sync
630 lines
22 KiB
Python
630 lines
22 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
"""Integration tests for handler override functionality.
|
|
|
|
Tests real customer usage scenarios:
|
|
- Using @custom_ping_handler and @custom_invocation_handler decorators
|
|
to override handlers
|
|
- Setting environment variables for handler specifications
|
|
- Writing customer scripts with custom_sagemaker_ping_handler() and
|
|
custom_sagemaker_invocation_handler() functions
|
|
- Priority: env vars > decorators > customer script files > framework
|
|
defaults
|
|
|
|
The handler-override scenarios exercise the real vLLM SageMaker router and
|
|
bootstrap path via an in-process FastAPI ``TestClient`` instead of launching a
|
|
model server. These scenarios fully replace the ``/ping`` and ``/invocations``
|
|
endpoints with customer handlers, so no inference engine is required to
|
|
validate override behavior. Avoiding the model server also keeps the tests
|
|
fast and deterministic rather than depending on the FastAPI version resolved
|
|
into the test environment at runtime.
|
|
"""
|
|
|
|
import os
|
|
|
|
import pytest
|
|
import requests
|
|
from fastapi import FastAPI
|
|
from fastapi.testclient import TestClient
|
|
|
|
from tests.utils import RemoteOpenAIServer
|
|
|
|
from .conftest import (
|
|
MODEL_NAME_SMOLLM,
|
|
)
|
|
|
|
|
|
def _build_sagemaker_test_client() -> TestClient:
|
|
"""Build a TestClient over the real SageMaker router and bootstrap path.
|
|
|
|
``attach_router`` is called with empty supported tasks because the override
|
|
tests replace the endpoints with customer handlers, so no framework
|
|
invocation handler (and therefore no engine) is exercised.
|
|
"""
|
|
from vllm.entrypoints.serve.sagemaker.api_router import (
|
|
attach_router,
|
|
sagemaker_standards_bootstrap,
|
|
)
|
|
|
|
app = FastAPI()
|
|
attach_router(app, ())
|
|
return TestClient(sagemaker_standards_bootstrap(app))
|
|
|
|
|
|
class TestHandlerOverrideIntegration:
|
|
"""Integration tests simulating real customer usage scenarios.
|
|
|
|
Each test simulates a fresh server startup where customers:
|
|
- Use @custom_ping_handler and @custom_invocation_handler decorators
|
|
- Set environment variables (CUSTOM_FASTAPI_PING_HANDLER, etc.)
|
|
- Write customer scripts with custom_sagemaker_ping_handler() and
|
|
custom_sagemaker_invocation_handler() functions
|
|
"""
|
|
|
|
def setup_method(self):
|
|
"""Setup for each test - simulate fresh server startup."""
|
|
self._clear_caches()
|
|
self._clear_env_vars()
|
|
|
|
def teardown_method(self):
|
|
"""Cleanup after each test."""
|
|
self._clear_env_vars()
|
|
|
|
def _clear_caches(self):
|
|
"""Clear handler registry and function loader cache."""
|
|
try:
|
|
from model_hosting_container_standards.common.handler import (
|
|
handler_registry,
|
|
)
|
|
from model_hosting_container_standards.sagemaker.sagemaker_loader import (
|
|
SageMakerFunctionLoader,
|
|
)
|
|
|
|
handler_registry.clear()
|
|
SageMakerFunctionLoader._default_function_loader = None
|
|
except ImportError:
|
|
pytest.skip("model-hosting-container-standards not available")
|
|
|
|
def _clear_env_vars(self):
|
|
"""Clear SageMaker environment variables."""
|
|
try:
|
|
from model_hosting_container_standards.common.fastapi.config import (
|
|
FastAPIEnvVars,
|
|
)
|
|
from model_hosting_container_standards.sagemaker.config import (
|
|
SageMakerEnvVars,
|
|
)
|
|
|
|
# Clear SageMaker env vars
|
|
for var in [
|
|
SageMakerEnvVars.SAGEMAKER_MODEL_PATH,
|
|
SageMakerEnvVars.CUSTOM_SCRIPT_FILENAME,
|
|
]:
|
|
os.environ.pop(var, None)
|
|
|
|
# Clear FastAPI env vars
|
|
for var in [
|
|
FastAPIEnvVars.CUSTOM_FASTAPI_PING_HANDLER,
|
|
FastAPIEnvVars.CUSTOM_FASTAPI_INVOCATION_HANDLER,
|
|
]:
|
|
os.environ.pop(var, None)
|
|
except ImportError:
|
|
pass
|
|
|
|
def test_customer_script_functions_auto_loaded(self, monkeypatch, tmp_path):
|
|
"""Test customer scenario: script functions automatically override
|
|
framework defaults."""
|
|
try:
|
|
from model_hosting_container_standards.sagemaker.config import (
|
|
SageMakerEnvVars,
|
|
)
|
|
except ImportError:
|
|
pytest.skip("model-hosting-container-standards not available")
|
|
|
|
# Customer writes a script file with ping() and invoke() functions
|
|
script_path = tmp_path / "model.py"
|
|
script_path.write_text(
|
|
"""
|
|
from fastapi import Request
|
|
|
|
async def custom_sagemaker_ping_handler():
|
|
return {
|
|
"status": "healthy",
|
|
"source": "customer_override",
|
|
"message": "Custom ping from customer script"
|
|
}
|
|
|
|
async def custom_sagemaker_invocation_handler(request: Request):
|
|
return {
|
|
"predictions": ["Custom response from customer script"],
|
|
"source": "customer_override"
|
|
}
|
|
"""
|
|
)
|
|
|
|
# Customer sets SageMaker environment variables to point to their script
|
|
monkeypatch.setenv(SageMakerEnvVars.SAGEMAKER_MODEL_PATH, str(tmp_path))
|
|
monkeypatch.setenv(SageMakerEnvVars.CUSTOM_SCRIPT_FILENAME, script_path.name)
|
|
|
|
with _build_sagemaker_test_client() as client:
|
|
# Customer tests their server and sees their overrides work
|
|
# automatically
|
|
ping_response = client.get("/ping")
|
|
assert ping_response.status_code == 200
|
|
ping_data = ping_response.json()
|
|
|
|
invoke_response = client.post(
|
|
"/invocations",
|
|
json={
|
|
"model": MODEL_NAME_SMOLLM,
|
|
"messages": [{"role": "user", "content": "Hello"}],
|
|
"max_tokens": 5,
|
|
},
|
|
)
|
|
assert invoke_response.status_code == 200
|
|
invoke_data = invoke_response.json()
|
|
|
|
# Customer sees their functions are used
|
|
assert ping_data["source"] == "customer_override"
|
|
assert ping_data["message"] == "Custom ping from customer script"
|
|
assert invoke_data["source"] == "customer_override"
|
|
assert invoke_data["predictions"] == [
|
|
"Custom response from customer script"
|
|
]
|
|
|
|
def test_customer_decorator_usage(self, monkeypatch, tmp_path):
|
|
"""Test customer scenario: using @custom_ping_handler and
|
|
@custom_invocation_handler decorators."""
|
|
try:
|
|
from model_hosting_container_standards.sagemaker.config import (
|
|
SageMakerEnvVars,
|
|
)
|
|
except ImportError:
|
|
pytest.skip("model-hosting-container-standards not available")
|
|
|
|
# Customer writes a script file with decorators
|
|
script_path = tmp_path / "model.py"
|
|
script_path.write_text(
|
|
"""
|
|
import model_hosting_container_standards.sagemaker as sagemaker_standards
|
|
from fastapi import Request
|
|
|
|
@sagemaker_standards.custom_ping_handler
|
|
async def my_ping():
|
|
return {
|
|
"type": "ping",
|
|
"source": "customer_decorator"
|
|
}
|
|
|
|
@sagemaker_standards.custom_invocation_handler
|
|
async def my_invoke(request: Request):
|
|
return {
|
|
"type": "invoke",
|
|
"source": "customer_decorator"
|
|
}
|
|
"""
|
|
)
|
|
|
|
monkeypatch.setenv(SageMakerEnvVars.SAGEMAKER_MODEL_PATH, str(tmp_path))
|
|
monkeypatch.setenv(SageMakerEnvVars.CUSTOM_SCRIPT_FILENAME, script_path.name)
|
|
|
|
with _build_sagemaker_test_client() as client:
|
|
ping_response = client.get("/ping")
|
|
assert ping_response.status_code == 200
|
|
ping_data = ping_response.json()
|
|
|
|
invoke_response = client.post(
|
|
"/invocations",
|
|
json={
|
|
"model": MODEL_NAME_SMOLLM,
|
|
"messages": [{"role": "user", "content": "Hello"}],
|
|
"max_tokens": 5,
|
|
},
|
|
)
|
|
assert invoke_response.status_code == 200
|
|
invoke_data = invoke_response.json()
|
|
|
|
# Customer sees their handlers are used by the server
|
|
assert ping_data["source"] == "customer_decorator"
|
|
assert invoke_data["source"] == "customer_decorator"
|
|
|
|
def test_handler_priority_order(self, monkeypatch, tmp_path):
|
|
"""Test priority: @custom_ping_handler/@custom_invocation_handler
|
|
decorators vs script functions."""
|
|
try:
|
|
from model_hosting_container_standards.sagemaker.config import (
|
|
SageMakerEnvVars,
|
|
)
|
|
except ImportError:
|
|
pytest.skip("model-hosting-container-standards not available")
|
|
|
|
# Customer writes a script with both decorator and regular functions
|
|
script_path = tmp_path / "model.py"
|
|
script_path.write_text(
|
|
"""
|
|
import model_hosting_container_standards.sagemaker as sagemaker_standards
|
|
from fastapi import Request
|
|
|
|
# Customer uses @custom_ping_handler decorator (higher priority than script functions)
|
|
@sagemaker_standards.custom_ping_handler
|
|
async def decorated_ping():
|
|
return {
|
|
"status": "healthy",
|
|
"source": "ping_decorator_in_script",
|
|
"priority": "decorator"
|
|
}
|
|
|
|
# Customer also has a regular function (lower priority than
|
|
# @custom_ping_handler decorator)
|
|
async def custom_sagemaker_ping_handler():
|
|
return {
|
|
"status": "healthy",
|
|
"source": "script_function",
|
|
"priority": "function"
|
|
}
|
|
|
|
# Customer has a regular invoke function
|
|
async def custom_sagemaker_invocation_handler(request: Request):
|
|
return {
|
|
"predictions": ["Script function response"],
|
|
"source": "script_invoke_function",
|
|
"priority": "function"
|
|
}
|
|
"""
|
|
)
|
|
|
|
monkeypatch.setenv(SageMakerEnvVars.SAGEMAKER_MODEL_PATH, str(tmp_path))
|
|
monkeypatch.setenv(SageMakerEnvVars.CUSTOM_SCRIPT_FILENAME, script_path.name)
|
|
|
|
with _build_sagemaker_test_client() as client:
|
|
ping_response = client.get("/ping")
|
|
assert ping_response.status_code == 200
|
|
ping_data = ping_response.json()
|
|
|
|
invoke_response = client.post(
|
|
"/invocations",
|
|
json={
|
|
"model": MODEL_NAME_SMOLLM,
|
|
"messages": [{"role": "user", "content": "Hello"}],
|
|
"max_tokens": 5,
|
|
},
|
|
)
|
|
assert invoke_response.status_code == 200
|
|
invoke_data = invoke_response.json()
|
|
|
|
# @custom_ping_handler decorator has higher priority than
|
|
# script function
|
|
assert ping_data["source"] == "ping_decorator_in_script"
|
|
assert ping_data["priority"] == "decorator"
|
|
|
|
# Script function is used for invoke
|
|
assert invoke_data["source"] == "script_invoke_function"
|
|
assert invoke_data["priority"] == "function"
|
|
|
|
def test_environment_variable_script_loading(self, monkeypatch, tmp_path):
|
|
"""Test that environment variables correctly specify script location
|
|
and loading."""
|
|
try:
|
|
from model_hosting_container_standards.sagemaker.config import (
|
|
SageMakerEnvVars,
|
|
)
|
|
except ImportError:
|
|
pytest.skip("model-hosting-container-standards not available")
|
|
|
|
# Customer writes a script in a specific directory
|
|
script_path = tmp_path / "model.py"
|
|
script_path.write_text(
|
|
"""
|
|
from fastapi import Request
|
|
|
|
async def custom_sagemaker_ping_handler():
|
|
return {
|
|
"status": "healthy",
|
|
"source": "env_loaded_script",
|
|
"method": "environment_variable_loading"
|
|
}
|
|
|
|
async def custom_sagemaker_invocation_handler(request: Request):
|
|
return {
|
|
"predictions": ["Loaded via environment variables"],
|
|
"source": "env_loaded_script",
|
|
"method": "environment_variable_loading"
|
|
}
|
|
"""
|
|
)
|
|
|
|
# Test environment variable script loading
|
|
monkeypatch.setenv(SageMakerEnvVars.SAGEMAKER_MODEL_PATH, str(tmp_path))
|
|
monkeypatch.setenv(SageMakerEnvVars.CUSTOM_SCRIPT_FILENAME, script_path.name)
|
|
|
|
with _build_sagemaker_test_client() as client:
|
|
ping_response = client.get("/ping")
|
|
assert ping_response.status_code == 200
|
|
ping_data = ping_response.json()
|
|
|
|
invoke_response = client.post(
|
|
"/invocations",
|
|
json={
|
|
"model": MODEL_NAME_SMOLLM,
|
|
"messages": [{"role": "user", "content": "Hello"}],
|
|
"max_tokens": 5,
|
|
},
|
|
)
|
|
assert invoke_response.status_code == 200
|
|
invoke_data = invoke_response.json()
|
|
|
|
# Verify that the script was loaded via environment variables
|
|
assert ping_data["source"] == "env_loaded_script"
|
|
assert ping_data["method"] == "environment_variable_loading"
|
|
assert invoke_data["source"] == "env_loaded_script"
|
|
assert invoke_data["method"] == "environment_variable_loading"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_framework_default_handlers(self):
|
|
"""Test that framework default handlers work when no customer
|
|
overrides exist.
|
|
|
|
This scenario exercises the real inference path (default
|
|
``/invocations``), so it keeps using a live model server rather than
|
|
the in-process TestClient.
|
|
"""
|
|
args = [
|
|
"--dtype",
|
|
"bfloat16",
|
|
"--max-model-len",
|
|
"2048",
|
|
"--enforce-eager",
|
|
"--max-num-seqs",
|
|
"32",
|
|
]
|
|
|
|
# Explicitly pass empty env_dict to ensure no SageMaker env vars are set
|
|
# This prevents pollution from previous tests
|
|
try:
|
|
from model_hosting_container_standards.common.fastapi.config import (
|
|
FastAPIEnvVars,
|
|
)
|
|
from model_hosting_container_standards.sagemaker.config import (
|
|
SageMakerEnvVars,
|
|
)
|
|
|
|
env_dict = {
|
|
SageMakerEnvVars.SAGEMAKER_MODEL_PATH: "",
|
|
SageMakerEnvVars.CUSTOM_SCRIPT_FILENAME: "",
|
|
FastAPIEnvVars.CUSTOM_FASTAPI_PING_HANDLER: "",
|
|
FastAPIEnvVars.CUSTOM_FASTAPI_INVOCATION_HANDLER: "",
|
|
}
|
|
except ImportError:
|
|
env_dict = {}
|
|
|
|
with RemoteOpenAIServer(MODEL_NAME_SMOLLM, args, env_dict=env_dict) as server:
|
|
# Test that default ping works
|
|
ping_response = requests.get(server.url_for("ping"))
|
|
assert ping_response.status_code == 200
|
|
|
|
# Test that default invocations work
|
|
invoke_response = requests.post(
|
|
server.url_for("invocations"),
|
|
json={
|
|
"model": MODEL_NAME_SMOLLM,
|
|
"messages": [{"role": "user", "content": "Hello"}],
|
|
"max_tokens": 5,
|
|
},
|
|
)
|
|
assert invoke_response.status_code == 200
|
|
|
|
def test_handler_env_var_override(self, monkeypatch, tmp_path):
|
|
"""Test CUSTOM_FASTAPI_PING_HANDLER and CUSTOM_FASTAPI_INVOCATION_HANDLER
|
|
environment variable overrides."""
|
|
try:
|
|
from model_hosting_container_standards.common.fastapi.config import (
|
|
FastAPIEnvVars,
|
|
)
|
|
from model_hosting_container_standards.sagemaker.config import (
|
|
SageMakerEnvVars,
|
|
)
|
|
except ImportError:
|
|
pytest.skip("model-hosting-container-standards not available")
|
|
|
|
# Create a script with both env var handlers and script functions
|
|
script_path = tmp_path / "model.py"
|
|
script_path.write_text(
|
|
"""
|
|
from fastapi import Request, Response
|
|
import json
|
|
|
|
async def env_var_ping_handler(raw_request: Request) -> Response:
|
|
return Response(
|
|
content=json.dumps({
|
|
"status": "healthy",
|
|
"source": "env_var_ping",
|
|
"method": "environment_variable"
|
|
}),
|
|
media_type="application/json"
|
|
)
|
|
|
|
async def env_var_invoke_handler(raw_request: Request) -> Response:
|
|
return Response(
|
|
content=json.dumps({
|
|
"predictions": ["Environment variable response"],
|
|
"source": "env_var_invoke",
|
|
"method": "environment_variable"
|
|
}),
|
|
media_type="application/json"
|
|
)
|
|
|
|
async def custom_sagemaker_ping_handler():
|
|
return {
|
|
"status": "healthy",
|
|
"source": "script_ping",
|
|
"method": "script_function"
|
|
}
|
|
|
|
async def custom_sagemaker_invocation_handler(request: Request):
|
|
return {
|
|
"predictions": ["Script function response"],
|
|
"source": "script_invoke",
|
|
"method": "script_function"
|
|
}
|
|
"""
|
|
)
|
|
|
|
# Set environment variables to override both handlers
|
|
monkeypatch.setenv(SageMakerEnvVars.SAGEMAKER_MODEL_PATH, str(tmp_path))
|
|
monkeypatch.setenv(SageMakerEnvVars.CUSTOM_SCRIPT_FILENAME, script_path.name)
|
|
monkeypatch.setenv(
|
|
FastAPIEnvVars.CUSTOM_FASTAPI_PING_HANDLER,
|
|
f"{script_path.name}:env_var_ping_handler",
|
|
)
|
|
monkeypatch.setenv(
|
|
FastAPIEnvVars.CUSTOM_FASTAPI_INVOCATION_HANDLER,
|
|
f"{script_path.name}:env_var_invoke_handler",
|
|
)
|
|
|
|
with _build_sagemaker_test_client() as client:
|
|
# Test ping handler override
|
|
ping_response = client.get("/ping")
|
|
assert ping_response.status_code == 200
|
|
ping_data = ping_response.json()
|
|
|
|
# Environment variable should override script function
|
|
assert ping_data["method"] == "environment_variable"
|
|
assert ping_data["source"] == "env_var_ping"
|
|
|
|
# Test invocation handler override
|
|
invoke_response = client.post(
|
|
"/invocations",
|
|
json={
|
|
"model": MODEL_NAME_SMOLLM,
|
|
"messages": [{"role": "user", "content": "Hello"}],
|
|
"max_tokens": 5,
|
|
},
|
|
)
|
|
assert invoke_response.status_code == 200
|
|
invoke_data = invoke_response.json()
|
|
|
|
# Environment variable should override script function
|
|
assert invoke_data["method"] == "environment_variable"
|
|
assert invoke_data["source"] == "env_var_invoke"
|
|
|
|
def test_env_var_priority_over_decorator_and_script(self, monkeypatch, tmp_path):
|
|
"""Test that environment variables have highest priority over decorators
|
|
and script functions for both ping and invocation handlers."""
|
|
try:
|
|
from model_hosting_container_standards.common.fastapi.config import (
|
|
FastAPIEnvVars,
|
|
)
|
|
from model_hosting_container_standards.sagemaker.config import (
|
|
SageMakerEnvVars,
|
|
)
|
|
except ImportError:
|
|
pytest.skip("model-hosting-container-standards not available")
|
|
|
|
# Create a script with all three handler types for both ping and invocation
|
|
script_path = tmp_path / "model.py"
|
|
script_path.write_text(
|
|
"""
|
|
import model_hosting_container_standards.sagemaker as sagemaker_standards
|
|
from fastapi import Request, Response
|
|
import json
|
|
|
|
# Environment variable handlers (highest priority)
|
|
async def env_priority_ping(raw_request: Request) -> Response:
|
|
return Response(
|
|
content=json.dumps({
|
|
"status": "healthy",
|
|
"source": "env_var",
|
|
"priority": "environment_variable"
|
|
}),
|
|
media_type="application/json"
|
|
)
|
|
|
|
async def env_priority_invoke(raw_request: Request) -> Response:
|
|
return Response(
|
|
content=json.dumps({
|
|
"predictions": ["Environment variable response"],
|
|
"source": "env_var",
|
|
"priority": "environment_variable"
|
|
}),
|
|
media_type="application/json"
|
|
)
|
|
|
|
# Decorator handlers (medium priority)
|
|
@sagemaker_standards.custom_ping_handler
|
|
async def decorator_ping(raw_request: Request) -> Response:
|
|
return Response(
|
|
content=json.dumps({
|
|
"status": "healthy",
|
|
"source": "decorator",
|
|
"priority": "decorator"
|
|
}),
|
|
media_type="application/json"
|
|
)
|
|
|
|
@sagemaker_standards.custom_invocation_handler
|
|
async def decorator_invoke(raw_request: Request) -> Response:
|
|
return Response(
|
|
content=json.dumps({
|
|
"predictions": ["Decorator response"],
|
|
"source": "decorator",
|
|
"priority": "decorator"
|
|
}),
|
|
media_type="application/json"
|
|
)
|
|
|
|
# Script functions (lowest priority)
|
|
async def custom_sagemaker_ping_handler():
|
|
return {
|
|
"status": "healthy",
|
|
"source": "script",
|
|
"priority": "script_function"
|
|
}
|
|
|
|
async def custom_sagemaker_invocation_handler(request: Request):
|
|
return {
|
|
"predictions": ["Script function response"],
|
|
"source": "script",
|
|
"priority": "script_function"
|
|
}
|
|
"""
|
|
)
|
|
|
|
# Set environment variables to specify highest priority handlers
|
|
monkeypatch.setenv(SageMakerEnvVars.SAGEMAKER_MODEL_PATH, str(tmp_path))
|
|
monkeypatch.setenv(SageMakerEnvVars.CUSTOM_SCRIPT_FILENAME, script_path.name)
|
|
monkeypatch.setenv(
|
|
FastAPIEnvVars.CUSTOM_FASTAPI_PING_HANDLER,
|
|
f"{script_path.name}:env_priority_ping",
|
|
)
|
|
monkeypatch.setenv(
|
|
FastAPIEnvVars.CUSTOM_FASTAPI_INVOCATION_HANDLER,
|
|
f"{script_path.name}:env_priority_invoke",
|
|
)
|
|
|
|
with _build_sagemaker_test_client() as client:
|
|
# Test ping handler priority
|
|
ping_response = client.get("/ping")
|
|
assert ping_response.status_code == 200
|
|
ping_data = ping_response.json()
|
|
|
|
# Environment variable has highest priority and should be used
|
|
assert ping_data["priority"] == "environment_variable"
|
|
assert ping_data["source"] == "env_var"
|
|
|
|
# Test invocation handler priority
|
|
invoke_response = client.post(
|
|
"/invocations",
|
|
json={
|
|
"model": MODEL_NAME_SMOLLM,
|
|
"messages": [{"role": "user", "content": "Hello"}],
|
|
"max_tokens": 5,
|
|
},
|
|
)
|
|
assert invoke_response.status_code == 200
|
|
invoke_data = invoke_response.json()
|
|
|
|
# Environment variable has highest priority and should be used
|
|
assert invoke_data["priority"] == "environment_variable"
|
|
assert invoke_data["source"] == "env_var"
|