589 lines
19 KiB
Python
589 lines
19 KiB
Python
# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
from typing import Any, Dict, List, Optional, Union
|
|
|
|
from tests.utils.client import send_request
|
|
from tests.utils.constants import DefaultPort
|
|
from tests.utils.payloads import (
|
|
AnthropicMessagesPayload,
|
|
AnthropicMessagesStreamPayload,
|
|
CachedTokensChatPayload,
|
|
ChatPayload,
|
|
ChatPayloadWithLogprobs,
|
|
CompletionPayload,
|
|
CompletionPayloadWithLogprobs,
|
|
EmbeddingPayload,
|
|
LMCacheMetricsPayload,
|
|
MetricsPayload,
|
|
ResponsesPayload,
|
|
ResponsesStreamPayload,
|
|
SGLangMetricsPayload,
|
|
TRTLLMMetricsPayload,
|
|
VLLMMetricsPayload,
|
|
)
|
|
|
|
# Common default text prompt used across tests
|
|
TEXT_PROMPT = "Tell me a knock knock joke about AI."
|
|
|
|
# Longer prompt for prefix caching tests - needs to be > 64 tokens (typical block size)
|
|
# to ensure at least one full block gets cached
|
|
LONG_PROMPT_FOR_CACHING = """In the heart of Eldoria, an ancient land of boundless magic and mysterious creatures, \
|
|
lies the long-forgotten city of Aeloria. Once a beacon of knowledge and power, Aeloria was buried beneath the \
|
|
shifting sands of time, lost to the world for centuries. You are an intrepid explorer, known for your unparalleled \
|
|
curiosity and courage, who has stumbled upon an ancient map hinting at the city's location. The map suggests that \
|
|
Aeloria holds a secret so profound that it has the potential to reshape the very fabric of reality. Your journey \
|
|
will take you through treacherous deserts, enchanted forests, and across perilous mountain ranges. \
|
|
Your Task: Character Background: Develop a detailed background for your character. Describe their motivations \
|
|
for seeking out Aeloria, their skills and weaknesses, and any personal connections to the ancient city or its legends."""
|
|
|
|
|
|
def chat_payload_default(
|
|
repeat_count: int = 3,
|
|
expected_response: Optional[List[str]] = None,
|
|
expected_log: Optional[List[str]] = None,
|
|
max_tokens: int = 1000,
|
|
temperature: float = 0.0,
|
|
stream: bool = False,
|
|
) -> ChatPayload:
|
|
return ChatPayload(
|
|
body={
|
|
"messages": [
|
|
{
|
|
"role": "user",
|
|
"content": TEXT_PROMPT,
|
|
}
|
|
],
|
|
"max_tokens": max_tokens,
|
|
"temperature": temperature,
|
|
"stream": stream,
|
|
},
|
|
repeat_count=repeat_count,
|
|
expected_log=expected_log or [],
|
|
# Accept any of these keywords in the response (case-insensitive)
|
|
expected_response=expected_response
|
|
or ["AI", "knock", "joke", "think", "artificial", "intelligence"],
|
|
)
|
|
|
|
|
|
def cached_tokens_chat_payload(
|
|
repeat_count: int = 3,
|
|
expected_response: Optional[List[str]] = None,
|
|
expected_log: Optional[List[str]] = None,
|
|
max_tokens: int = 100,
|
|
temperature: float = 0.0,
|
|
min_cached_tokens: int = 64,
|
|
) -> CachedTokensChatPayload:
|
|
"""Create a chat payload that validates cached tokens in usage field.
|
|
|
|
This is useful for testing KV router cache-aware routing where repeated
|
|
identical prompts should result in cached tokens being reported.
|
|
|
|
Uses a longer prompt (~196 tokens) to ensure at least one full block (64 tokens)
|
|
gets cached. vLLM only caches complete blocks, so short prompts won't trigger
|
|
the cached_tokens field in the response.
|
|
|
|
Args:
|
|
repeat_count: Number of times to repeat the request (>1 needed to see caching)
|
|
expected_response: List of expected strings in response
|
|
expected_log: List of expected log patterns
|
|
max_tokens: Maximum tokens to generate
|
|
temperature: Sampling temperature
|
|
min_cached_tokens: Minimum cached tokens expected after first request (default: 64, one block)
|
|
|
|
Returns:
|
|
CachedTokensChatPayload configured for testing prefix caching
|
|
"""
|
|
return CachedTokensChatPayload(
|
|
body={
|
|
"messages": [
|
|
{
|
|
"role": "user",
|
|
"content": LONG_PROMPT_FOR_CACHING,
|
|
}
|
|
],
|
|
"max_tokens": max_tokens,
|
|
"temperature": temperature,
|
|
"stream": False,
|
|
},
|
|
repeat_count=repeat_count,
|
|
expected_log=expected_log or [],
|
|
expected_response=expected_response
|
|
or ["Aeloria", "Eldoria", "explorer", "ancient", "character", "background"],
|
|
min_cached_tokens=min_cached_tokens,
|
|
)
|
|
|
|
|
|
def completion_payload_default(
|
|
repeat_count: int = 3,
|
|
expected_response: Optional[List[str]] = None,
|
|
expected_log: Optional[List[str]] = None,
|
|
max_tokens: int = 1000,
|
|
temperature: float = 0.0,
|
|
stream: bool = False,
|
|
) -> CompletionPayload:
|
|
return CompletionPayload(
|
|
body={
|
|
"prompt": TEXT_PROMPT,
|
|
"max_tokens": max_tokens,
|
|
"temperature": temperature,
|
|
"stream": stream,
|
|
},
|
|
repeat_count=repeat_count,
|
|
expected_log=expected_log or [],
|
|
# Accept any of these keywords in the response (case-insensitive)
|
|
expected_response=expected_response
|
|
or ["AI", "knock", "joke", "think", "artificial", "intelligence"],
|
|
)
|
|
|
|
|
|
def multimodal_payload_default(
|
|
image_url: str = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/inpaint.png",
|
|
text: str = "Describe the image",
|
|
repeat_count: int = 1,
|
|
expected_response: Optional[List[str]] = None,
|
|
expected_log: Optional[List[str]] = None,
|
|
max_tokens: int = 160,
|
|
temperature: Optional[float] = None,
|
|
stream: bool = False,
|
|
) -> ChatPayload:
|
|
"""Create a multimodal chat payload with image and text content.
|
|
|
|
Args:
|
|
image_url: URL of the image to include in the request
|
|
text: Text prompt to accompany the image
|
|
repeat_count: Number of times to repeat the request
|
|
expected_response: List of strings expected in the response
|
|
expected_log: List of regex patterns expected in logs
|
|
max_tokens: Maximum tokens to generate
|
|
temperature: Sampling temperature (optional)
|
|
stream: Whether to stream the response
|
|
|
|
Returns:
|
|
ChatPayload configured for multimodal requests
|
|
"""
|
|
return chat_payload(
|
|
content=[
|
|
{"type": "text", "text": text},
|
|
{
|
|
"type": "image_url",
|
|
"image_url": {"url": image_url},
|
|
},
|
|
],
|
|
repeat_count=repeat_count,
|
|
expected_response=expected_response or ["image"],
|
|
expected_log=expected_log or [],
|
|
max_tokens=max_tokens,
|
|
temperature=temperature,
|
|
stream=stream,
|
|
)
|
|
|
|
|
|
def metric_payload_default(
|
|
min_num_requests: int,
|
|
repeat_count: int = 1,
|
|
expected_log: Optional[List[str]] = None,
|
|
backend: Optional[str] = None,
|
|
port: int = DefaultPort.SYSTEM1.value,
|
|
) -> MetricsPayload:
|
|
"""Create a metrics payload for the specified backend.
|
|
|
|
Args:
|
|
min_num_requests: Minimum number of requests expected in metrics
|
|
repeat_count: Number of times to repeat the request
|
|
expected_log: Expected log messages
|
|
backend: Backend type ('vllm', 'sglang', 'trtllm', 'lmcache')
|
|
port: Port to use for metrics endpoint
|
|
|
|
Returns:
|
|
Backend-specific MetricsPayload subclass based on backend parameter
|
|
"""
|
|
common_args: dict[str, Any] = {
|
|
"body": {},
|
|
"repeat_count": repeat_count,
|
|
"expected_log": expected_log or [],
|
|
"expected_response": [],
|
|
"min_num_requests": min_num_requests,
|
|
"port": port,
|
|
}
|
|
|
|
# Return backend-specific payload class
|
|
if backend == "vllm":
|
|
return VLLMMetricsPayload(**common_args)
|
|
elif backend == "sglang":
|
|
return SGLangMetricsPayload(**common_args)
|
|
elif backend == "trtllm":
|
|
return TRTLLMMetricsPayload(**common_args)
|
|
elif backend == "lmcache":
|
|
return LMCacheMetricsPayload(**common_args)
|
|
else:
|
|
# Default to base MetricsPayload for unknown backends
|
|
return MetricsPayload(**common_args)
|
|
|
|
|
|
def chat_payload(
|
|
content: Union[str, List[Dict[str, Any]]],
|
|
repeat_count: int = 1,
|
|
expected_response: Optional[List[str]] = None,
|
|
expected_log: Optional[List[str]] = None,
|
|
max_tokens: int = 300,
|
|
temperature: Optional[float] = None,
|
|
stream: bool = False,
|
|
logprobs: bool = False,
|
|
top_logprobs: Optional[int] = None,
|
|
extra_body: Optional[Dict[str, Any]] = None,
|
|
) -> ChatPayload:
|
|
body: Dict[str, Any] = {
|
|
"messages": [
|
|
{
|
|
"role": "user",
|
|
"content": content,
|
|
}
|
|
],
|
|
"max_tokens": max_tokens,
|
|
"stream": stream,
|
|
"logprobs": logprobs,
|
|
}
|
|
if temperature is not None:
|
|
body["temperature"] = temperature
|
|
if logprobs is not None:
|
|
body["logprobs"] = logprobs
|
|
if top_logprobs is not None:
|
|
body["top_logprobs"] = top_logprobs
|
|
|
|
if top_logprobs is not None:
|
|
body["top_logprobs"] = top_logprobs
|
|
|
|
if extra_body:
|
|
body.update(extra_body)
|
|
|
|
if logprobs:
|
|
return ChatPayloadWithLogprobs(
|
|
body=body,
|
|
repeat_count=repeat_count,
|
|
expected_log=expected_log or [],
|
|
expected_response=expected_response or [],
|
|
)
|
|
else:
|
|
return ChatPayload(
|
|
body=body,
|
|
repeat_count=repeat_count,
|
|
expected_log=expected_log or [],
|
|
expected_response=expected_response or [],
|
|
)
|
|
|
|
|
|
def completion_payload(
|
|
prompt: str,
|
|
repeat_count: int = 3,
|
|
expected_response: Optional[List[str]] = None,
|
|
expected_log: Optional[List[str]] = None,
|
|
max_tokens: int = 150,
|
|
temperature: float = 0.1,
|
|
stream: bool = False,
|
|
logprobs: Optional[int] = None,
|
|
) -> CompletionPayload:
|
|
body: Dict[str, Any] = {
|
|
"prompt": prompt,
|
|
"max_tokens": max_tokens,
|
|
"temperature": temperature,
|
|
"stream": stream,
|
|
}
|
|
if logprobs is not None:
|
|
body["logprobs"] = logprobs
|
|
return CompletionPayloadWithLogprobs(
|
|
body=body,
|
|
repeat_count=repeat_count,
|
|
expected_log=expected_log or [],
|
|
expected_response=expected_response or [],
|
|
)
|
|
else:
|
|
return CompletionPayload(
|
|
body=body,
|
|
repeat_count=repeat_count,
|
|
expected_log=expected_log or [],
|
|
expected_response=expected_response or [],
|
|
)
|
|
|
|
|
|
def embedding_payload_default(
|
|
repeat_count: int = 3,
|
|
expected_response: Optional[List[str]] = None,
|
|
expected_log: Optional[List[str]] = None,
|
|
) -> EmbeddingPayload:
|
|
return EmbeddingPayload(
|
|
body={
|
|
"input": ["The sky is blue.", "Machine learning is fascinating."],
|
|
},
|
|
repeat_count=repeat_count,
|
|
expected_log=expected_log or [],
|
|
expected_response=expected_response
|
|
or ["Generated 2 embeddings with dimension"],
|
|
)
|
|
|
|
|
|
def embedding_payload(
|
|
input_text: Union[str, List[str]],
|
|
repeat_count: int = 3,
|
|
expected_response: Optional[List[str]] = None,
|
|
expected_log: Optional[List[str]] = None,
|
|
) -> EmbeddingPayload:
|
|
# Normalize input to list for consistent processing
|
|
if isinstance(input_text, str):
|
|
input_list = [input_text]
|
|
expected_count = 1
|
|
else:
|
|
input_list = input_text
|
|
expected_count = len(input_text)
|
|
|
|
return EmbeddingPayload(
|
|
body={
|
|
"input": input_list,
|
|
},
|
|
repeat_count=repeat_count,
|
|
expected_log=expected_log or [],
|
|
expected_response=expected_response
|
|
or [f"Generated {expected_count} embeddings with dimension"],
|
|
)
|
|
|
|
|
|
# Build small request-based health checks for chat and completions
|
|
# these should only be used as a last resort. Generally want to use an actual health check
|
|
|
|
|
|
def make_chat_health_check(port: int, model: str):
|
|
def _check_chat_endpoint(remaining_timeout: float = 30.0) -> bool:
|
|
payload = chat_payload_default(
|
|
repeat_count=1,
|
|
expected_response=[],
|
|
max_tokens=8,
|
|
temperature=0.0,
|
|
stream=False,
|
|
).with_model(model)
|
|
payload.port = port
|
|
try:
|
|
resp = send_request(
|
|
payload.url(),
|
|
payload.body,
|
|
timeout=min(max(1.0, remaining_timeout), 5.0),
|
|
method=payload.method,
|
|
log_level=10,
|
|
)
|
|
# Validate structure only; expected_response is empty
|
|
_ = payload.response_handler(resp)
|
|
return True
|
|
except Exception:
|
|
return False
|
|
|
|
return _check_chat_endpoint
|
|
|
|
|
|
def make_completions_health_check(port: int, model: str):
|
|
def _check_completions_endpoint(remaining_timeout: float = 30.0) -> bool:
|
|
payload = completion_payload_default(
|
|
repeat_count=1,
|
|
expected_response=[],
|
|
max_tokens=8,
|
|
temperature=0.0,
|
|
stream=False,
|
|
).with_model(model)
|
|
payload.port = port
|
|
try:
|
|
resp = send_request(
|
|
payload.url(),
|
|
payload.body,
|
|
timeout=min(max(1.0, remaining_timeout), 5.0),
|
|
method=payload.method,
|
|
log_level=10,
|
|
)
|
|
out = payload.response_handler(resp)
|
|
if not out:
|
|
raise ValueError("")
|
|
return True
|
|
except Exception:
|
|
return False
|
|
|
|
return _check_completions_endpoint
|
|
|
|
|
|
def chat_payload_with_logprobs(
|
|
content: Union[str, List[Dict[str, Any]]] = TEXT_PROMPT,
|
|
repeat_count: int = 1,
|
|
expected_response: Optional[List[str]] = None,
|
|
max_tokens: int = 50,
|
|
temperature: float = 0.0,
|
|
top_logprobs: int = 3,
|
|
) -> ChatPayloadWithLogprobs:
|
|
"""
|
|
Create a chat payload that requests and validates logprobs in the response.
|
|
|
|
Args:
|
|
content: Message content (text or structured content list)
|
|
repeat_count: Number of times to repeat the request
|
|
expected_response: List of strings expected in the response text
|
|
max_tokens: Maximum tokens to generate
|
|
temperature: Sampling temperature
|
|
top_logprobs: Number of top logprobs to return per token
|
|
|
|
Returns:
|
|
ChatPayloadWithLogprobs that validates logprobs in response
|
|
"""
|
|
body: Dict[str, Any] = {
|
|
"messages": [
|
|
{
|
|
"role": "user",
|
|
"content": content,
|
|
}
|
|
],
|
|
"max_tokens": max_tokens,
|
|
"temperature": temperature,
|
|
"logprobs": True,
|
|
"top_logprobs": top_logprobs,
|
|
}
|
|
|
|
return ChatPayloadWithLogprobs(
|
|
body=body,
|
|
repeat_count=repeat_count,
|
|
expected_log=[],
|
|
expected_response=expected_response or ["AI", "knock", "joke"],
|
|
)
|
|
|
|
|
|
def completion_payload_with_logprobs(
|
|
prompt: str = TEXT_PROMPT,
|
|
repeat_count: int = 1,
|
|
expected_response: Optional[List[str]] = None,
|
|
max_tokens: int = 50,
|
|
temperature: float = 0.0,
|
|
logprobs: int = 5,
|
|
) -> CompletionPayloadWithLogprobs:
|
|
"""
|
|
Create a completion payload that requests and validates logprobs in the response.
|
|
|
|
Args:
|
|
prompt: Text prompt
|
|
repeat_count: Number of times to repeat the request
|
|
expected_response: List of strings expected in the response text
|
|
max_tokens: Maximum tokens to generate
|
|
temperature: Sampling temperature
|
|
logprobs: Number of logprobs to return per token
|
|
|
|
Returns:
|
|
CompletionPayloadWithLogprobs that validates logprobs in response
|
|
"""
|
|
body: Dict[str, Any] = {
|
|
"prompt": prompt,
|
|
"max_tokens": max_tokens,
|
|
"temperature": temperature,
|
|
"logprobs": logprobs,
|
|
}
|
|
|
|
return CompletionPayloadWithLogprobs(
|
|
body=body,
|
|
repeat_count=repeat_count,
|
|
expected_log=[],
|
|
expected_response=expected_response or ["AI", "knock", "joke"],
|
|
)
|
|
|
|
|
|
def responses_payload_default(
|
|
repeat_count: int = 1,
|
|
expected_response: Optional[List[str]] = None,
|
|
expected_log: Optional[List[str]] = None,
|
|
max_tokens: int = 200,
|
|
temperature: float = 0.0,
|
|
) -> ResponsesPayload:
|
|
"""Create a default Responses API payload (non-streaming).
|
|
|
|
For full compliance testing, use the OpenResponses bun CLI instead.
|
|
"""
|
|
return ResponsesPayload(
|
|
body={
|
|
"input": TEXT_PROMPT,
|
|
"max_output_tokens": max_tokens,
|
|
"temperature": temperature,
|
|
},
|
|
repeat_count=repeat_count,
|
|
expected_log=expected_log or [],
|
|
expected_response=expected_response
|
|
or ["AI", "knock", "joke", "think", "artificial", "intelligence"],
|
|
)
|
|
|
|
|
|
def responses_stream_payload_default(
|
|
repeat_count: int = 1,
|
|
expected_response: Optional[List[str]] = None,
|
|
expected_log: Optional[List[str]] = None,
|
|
max_tokens: int = 200,
|
|
temperature: float = 0.0,
|
|
) -> ResponsesStreamPayload:
|
|
"""Create a default Responses API streaming payload.
|
|
|
|
For full compliance testing, use the OpenResponses bun CLI instead.
|
|
"""
|
|
return ResponsesStreamPayload(
|
|
body={
|
|
"input": TEXT_PROMPT,
|
|
"stream": True,
|
|
"max_output_tokens": max_tokens,
|
|
"temperature": temperature,
|
|
},
|
|
repeat_count=repeat_count,
|
|
expected_log=expected_log or [],
|
|
expected_response=expected_response
|
|
or ["AI", "knock", "joke", "think", "artificial", "intelligence"],
|
|
)
|
|
|
|
|
|
def anthropic_messages_payload_default(
|
|
repeat_count: int = 1,
|
|
expected_response: Optional[List[str]] = None,
|
|
expected_log: Optional[List[str]] = None,
|
|
max_tokens: int = 200,
|
|
temperature: float = 0.0,
|
|
) -> AnthropicMessagesPayload:
|
|
"""Create a default Anthropic Messages API payload (non-streaming)."""
|
|
return AnthropicMessagesPayload(
|
|
body={
|
|
"max_tokens": max_tokens,
|
|
"messages": [
|
|
{
|
|
"role": "user",
|
|
"content": TEXT_PROMPT,
|
|
}
|
|
],
|
|
"temperature": temperature,
|
|
},
|
|
repeat_count=repeat_count,
|
|
expected_log=expected_log or [],
|
|
expected_response=expected_response
|
|
or ["AI", "knock", "joke", "think", "artificial", "intelligence"],
|
|
)
|
|
|
|
|
|
def anthropic_messages_stream_payload_default(
|
|
repeat_count: int = 1,
|
|
expected_response: Optional[List[str]] = None,
|
|
expected_log: Optional[List[str]] = None,
|
|
max_tokens: int = 200,
|
|
temperature: float = 0.0,
|
|
) -> AnthropicMessagesStreamPayload:
|
|
"""Create a default Anthropic Messages API streaming payload."""
|
|
return AnthropicMessagesStreamPayload(
|
|
body={
|
|
"max_tokens": max_tokens,
|
|
"messages": [
|
|
{
|
|
"role": "user",
|
|
"content": TEXT_PROMPT,
|
|
}
|
|
],
|
|
"stream": True,
|
|
"temperature": temperature,
|
|
},
|
|
repeat_count=repeat_count,
|
|
expected_log=expected_log or [],
|
|
expected_response=expected_response
|
|
or ["AI", "knock", "joke", "think", "artificial", "intelligence"],
|
|
)
|