# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License"). You
# may not use this file except in compliance with the License. A copy of
# the License is located at
#
# http://aws.amazon.com/apache2.0/
#
# or in the "license" file accompanying this file. This file is
# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
# ANY KIND, either express or implied. See the License for the specific
# language governing permissions and limitations under the License.
"""Workload spec builder."""
from __future__ import absolute_import
import json
import os
import re
import uuid
from typing import Any, Dict, List, Optional, Union
from urllib.parse import urlparse
from pydantic import BaseModel, ConfigDict, Field, field_validator
from sagemaker.serve.ai_inference_recommender.secrets import Secret
# Default input-data channel names; the channel is mounted at
# {_CONTAINER_INPUT_DATA_DIR}/{channel_name}/ inside the benchmark container.
_DEFAULT_CHANNEL_NAME = "dataset"
_TEMPLATE_CHANNEL_NAME = "template"
_CONTAINER_INPUT_DATA_DIR = "/opt/ml/input/data"
# A Secrets Manager secret ARN, e.g.
# arn:aws:secretsmanager:us-east-1:123456789012:secret:my-secret-AbCdEf
_SECRET_ARN_PATTERN = re.compile(r"^arn:aws[a-z\-]*:secretsmanager:[^:]+:\d+:secret:.+")
class _DatasetChannel(BaseModel):
"""An S3 channel to mount into the benchmark container at job runtime."""
channel_name: str
s3_uri: str
[docs]
class Workload(BaseModel):
"""A workload specification used by benchmark and recommendation jobs."""
model_config = ConfigDict(arbitrary_types_allowed=True)
parameters: Dict[str, Any]
secrets: Dict[str, Union[str, Secret]] = Field(default_factory=dict)
tooling: Dict[str, Any] = Field(default_factory=lambda: {"api_standard": "openai"})
dataset_channels: List[_DatasetChannel] = Field(default_factory=list)
@field_validator("secrets")
@classmethod
def _reject_plaintext_secrets(
cls, value: Dict[str, Union[str, Secret]]
) -> Dict[str, Union[str, Secret]]:
"""Reject plaintext secret values so they are never persisted server-side.
A workload's secrets are serialized into the ``AIWorkloadConfig`` the
service stores; a raw string that is not a Secrets Manager ARN would be
saved in cleartext. Require a ``Secret`` or an ARN string, and point
callers at :meth:`Secret.from_string` to store a plaintext value safely.
"""
for key, secret in value.items():
if isinstance(secret, str) and not _SECRET_ARN_PATTERN.match(secret):
raise ValueError(
f"secret {key!r} looks like a plaintext value, not a Secrets "
"Manager ARN. Passing a plaintext secret would persist it in "
"cleartext. Store it first with Secret.from_string(<value>) "
"and pass the returned Secret (or its ARN)."
)
return value
[docs]
@classmethod
def synthetic(
cls,
*,
tokenizer: str,
concurrency: int = 1,
request_count: int = 100,
prompt_input_tokens_mean: int = 256,
prompt_input_tokens_stddev: float = 0.0,
output_tokens_mean: int = 256,
output_tokens_stddev: float = 0.0,
streaming: bool = True,
hf_token: Optional[Union[str, Secret]] = None,
**params: Any,
) -> "Workload":
"""Build a workload that uses synthetic prompts.
Synthetic prompts are generated by AIPerf from the Sonnet dataset,
producing realistic token distributions. Use
:meth:`Workload.from_dataset` to drive the benchmark from a real
request trace instead.
Args:
tokenizer: HuggingFace tokenizer id (e.g. ``meta-llama/Llama-3.2-1B``).
concurrency: Number of in-flight requests.
request_count: Total number of requests to issue.
prompt_input_tokens_mean: Mean input prompt length in tokens.
prompt_input_tokens_stddev: Standard deviation of input token count.
output_tokens_mean: Mean output response length in tokens.
output_tokens_stddev: Standard deviation of output token count.
streaming: Whether to use streaming chat completions.
hf_token: HuggingFace access token for gated tokenizers. Accepts a
``Secret`` or a Secrets Manager ARN string.
**params: Additional parameters merged into the workload's
``parameters`` map.
"""
parameters: Dict[str, Any] = {
"tokenizer": tokenizer,
"concurrency": concurrency,
"request_count": request_count,
"prompt_input_tokens_mean": prompt_input_tokens_mean,
"prompt_input_tokens_stddev": prompt_input_tokens_stddev,
"output_tokens_mean": output_tokens_mean,
"output_tokens_stddev": output_tokens_stddev,
"streaming": streaming,
**params,
}
secrets: Dict[str, Union[str, Secret]] = {}
if hf_token is not None:
secrets["hf_token"] = hf_token
return cls(parameters=parameters, secrets=secrets)
[docs]
@classmethod
def sonnet(cls, **kwargs: Any) -> "Workload":
"""Alias for :meth:`synthetic`.
AIPerf seeds synthetic prompts from the Sonnet dataset by default,
so ``Workload.sonnet(...)`` is the same as ``Workload.synthetic(...)``.
"""
return cls.synthetic(**kwargs)
[docs]
@classmethod
def from_dataset(
cls,
s3_uri: str,
*,
custom_dataset_type: Optional[str] = None,
tokenizer: Optional[str] = None,
concurrency: int = 1,
request_count: int = 100,
prompt_input_tokens_mean: int = 256,
prompt_input_tokens_stddev: float = 0.0,
output_tokens_mean: int = 256,
output_tokens_stddev: float = 0.0,
streaming: bool = True,
hf_token: Optional[Union[str, Secret]] = None,
channel_name: str = _DEFAULT_CHANNEL_NAME,
**params: Any,
) -> "Workload":
"""Build a workload that drives traffic from an S3-hosted dataset.
The benchmark replays requests from the dataset at the given S3 prefix.
Args:
s3_uri: ``s3://bucket/prefix/`` URI containing the dataset.
custom_dataset_type: Optional AIPerf custom-dataset format
(e.g. ``"openai-chat"``).
tokenizer: Optional HuggingFace tokenizer id; required for some
AIPerf metrics that compute per-token statistics.
concurrency: Number of in-flight requests.
request_count: Total number of requests to issue.
prompt_input_tokens_mean: Mean input prompt length in tokens.
prompt_input_tokens_stddev: Standard deviation of input token count.
output_tokens_mean: Mean output response length in tokens.
output_tokens_stddev: Standard deviation of output token count.
streaming: Whether to use streaming chat completions.
hf_token: HuggingFace access token for gated tokenizers. Accepts
a ``Secret`` or a Secrets Manager ARN string.
channel_name: Name of the input data channel the dataset is
mounted under. Defaults to ``"dataset"``.
**params: Additional AIPerf parameters for the workload.
Returns:
A ``Workload`` configured to drive traffic from the dataset.
"""
if not s3_uri.startswith("s3://"):
raise ValueError(f"s3_uri must start with 's3://'; got {s3_uri!r}.")
parameters: Dict[str, Any] = {
"concurrency": concurrency,
"request_count": request_count,
"prompt_input_tokens_mean": prompt_input_tokens_mean,
"prompt_input_tokens_stddev": prompt_input_tokens_stddev,
"output_tokens_mean": output_tokens_mean,
"output_tokens_stddev": output_tokens_stddev,
"streaming": streaming,
**params,
}
if custom_dataset_type is not None:
parameters["custom_dataset_type"] = custom_dataset_type
if tokenizer is not None:
parameters["tokenizer"] = tokenizer
secrets: Dict[str, Union[str, Secret]] = {}
if hf_token is not None:
secrets["hf_token"] = hf_token
return cls(
parameters=parameters,
secrets=secrets,
dataset_channels=[
_DatasetChannel(channel_name=channel_name, s3_uri=s3_uri),
],
)
[docs]
@classmethod
def template(
cls,
request_template: Optional[str] = None,
*,
template_s3_uri: Optional[str] = None,
sagemaker_session: Optional[Any] = None,
response_field: Optional[str] = None,
tokenizer: Optional[str] = None,
concurrency: int = 1,
request_count: int = 100,
prompt_input_tokens_mean: int = 256,
prompt_input_tokens_stddev: float = 0.0,
output_tokens_mean: int = 256,
output_tokens_stddev: float = 0.0,
streaming: bool = True,
hf_token: Optional[Union[str, Secret]] = None,
channel_name: str = _TEMPLATE_CHANNEL_NAME,
extra_inputs: Optional[str] = None,
**params: Any,
) -> "Workload":
"""Build a workload that benchmarks a custom-format (non-OpenAI) endpoint.
Use this for endpoints that don't speak the OpenAI chat-completions
format (e.g. DJL custom handlers, TensorRT-LLM native format). You
provide a `Jinja2 <https://jinja.palletsprojects.com>`__ template
describing your endpoint's request payload; it is rendered per request
(still using synthetic prompts) and sent to your endpoint. A
`JMESPath <https://jmespath.org>`__ ``response_field`` extracts the
generated text from the response.
Provide the template in one of two ways: pass ``request_template`` as
a local file path or an inline Jinja2 string and it is uploaded to the
session default bucket for you, or pass ``template_s3_uri`` if the
template is already in S3. Exactly one of the two is required.
Args:
request_template: The Jinja2 template for your endpoint's request
payload, given as either a local file path or an inline string.
It is uploaded to the session default bucket for you; mutually
exclusive with ``template_s3_uri``.
template_s3_uri: ``s3://bucket/key`` URI of an already-uploaded
Jinja2 template file (a single object, not a prefix). Use this
only when the template is already in S3; otherwise pass
``request_template``.
sagemaker_session: Session used to upload a local/inline
``request_template``. Defaults to a new ``Session``.
response_field: Optional JMESPath query that extracts the response
text from your endpoint's output (e.g. ``"generated_text"`` or
``"choices[0].message.content"``). Omit to let AIPerf
auto-detect the response format.
tokenizer: Optional HuggingFace tokenizer id; required for AIPerf
metrics that compute per-token statistics.
concurrency: Number of in-flight requests.
request_count: Total number of requests to issue.
prompt_input_tokens_mean: Mean input prompt length in tokens.
prompt_input_tokens_stddev: Standard deviation of input token count.
output_tokens_mean: Mean output response length in tokens.
output_tokens_stddev: Standard deviation of output token count.
streaming: Whether the endpoint streams its response.
hf_token: HuggingFace access token for gated tokenizers. Accepts a
``Secret`` or a Secrets Manager ARN string.
channel_name: Name of the input data channel the template is
mounted under. Defaults to ``"template"``.
extra_inputs: Optional space-separated ``key:value`` pairs for
advanced AIPerf options beyond ``response_field``
(e.g. ``"ignore_eos:true"``).
**params: Additional AIPerf parameters for the workload.
Returns:
A ``Workload`` configured for template mode.
"""
if (request_template is None) == (template_s3_uri is None):
raise ValueError(
"Provide exactly one of 'request_template' (a local path or "
"inline Jinja2 string, uploaded for you) or 'template_s3_uri' "
"(an already-uploaded s3:// URI)."
)
if request_template is not None:
template_s3_uri = cls._upload_template(
request_template, channel_name=channel_name, sagemaker_session=sagemaker_session
)
if not template_s3_uri.startswith("s3://"):
raise ValueError(f"template_s3_uri must start with 's3://'; got {template_s3_uri!r}.")
template_filename = urlparse(template_s3_uri).path.rsplit("/", 1)[-1]
if not template_filename:
raise ValueError(
"template_s3_uri must point at a template file, not a prefix; "
f"got {template_s3_uri!r}."
)
template_local_path = f"{_CONTAINER_INPUT_DATA_DIR}/{channel_name}/{template_filename}"
extra_inputs_pairs = [f"payload_template:{template_local_path}"]
if response_field is not None:
extra_inputs_pairs.append(f"response_field:{response_field}")
if extra_inputs:
extra_inputs_pairs.append(extra_inputs)
parameters: Dict[str, Any] = {
"concurrency": concurrency,
"request_count": request_count,
"prompt_input_tokens_mean": prompt_input_tokens_mean,
"prompt_input_tokens_stddev": prompt_input_tokens_stddev,
"output_tokens_mean": output_tokens_mean,
"output_tokens_stddev": output_tokens_stddev,
"streaming": streaming,
"extra_inputs": " ".join(extra_inputs_pairs),
**params,
}
if tokenizer is not None:
parameters["tokenizer"] = tokenizer
secrets: Dict[str, Union[str, Secret]] = {}
if hf_token is not None:
secrets["hf_token"] = hf_token
return cls(
parameters=parameters,
secrets=secrets,
dataset_channels=[
_DatasetChannel(channel_name=channel_name, s3_uri=template_s3_uri),
],
)
@staticmethod
def _upload_template(
request_template: str, *, channel_name: str, sagemaker_session: Optional[Any] = None
) -> str:
"""Upload a local-path or inline Jinja2 template to S3 and return its URI.
``request_template`` is treated as a local file path when it points at
an existing file (absolute, relative, or ``~``-prefixed); otherwise it
is treated as an inline template string. The object is uploaded to the
session default bucket so the caller never has to stage it by hand.
"""
from sagemaker.core.helper.session_helper import Session
from sagemaker.core.s3 import S3Uploader
session = sagemaker_session or Session()
bucket = session.default_bucket()
key_prefix = f"ai-inference-recommender/templates/{uuid.uuid4().hex}"
# Expand ~ so home-relative paths are recognized as files rather than
# being mistaken for inline template content.
local_path = os.path.expanduser(request_template)
if os.path.isfile(local_path):
# upload() appends the local filename to the prefix and returns the
# full S3 URI of the uploaded object.
return S3Uploader.upload(
local_path=local_path,
desired_s3_uri=f"s3://{bucket}/{key_prefix}",
sagemaker_session=session,
)
# Inline template string.
return S3Uploader.upload_string_as_file_body(
body=request_template,
desired_s3_uri=f"s3://{bucket}/{key_prefix}/endpoint_template.jinja",
sagemaker_session=session,
)
[docs]
def to_inline(self) -> str:
"""Serialize the workload to a JSON string.
``Secret`` values are flattened to their ARN strings.
"""
flat_secrets = {k: (v.arn if isinstance(v, Secret) else v) for k, v in self.secrets.items()}
payload: Dict[str, Any] = {
"benchmark": {"type": "aiperf"},
"parameters": self.parameters,
"tooling": dict(self.tooling),
}
if flat_secrets:
payload["secrets"] = flat_secrets
return json.dumps(payload)