sagemaker.serve.ai_inference_recommender.workload

Workload spec builder.

Classes

Workload(*, parameters, ~typing.Any], ...)

A workload specification used by benchmark and recommendation jobs.

class sagemaker.serve.ai_inference_recommender.workload.Workload(*, parameters: ~typing.Dict[str, ~typing.Any], secrets: ~typing.Dict[str, str | ~sagemaker.serve.ai_inference_recommender.secrets.Secret] = <factory>, tooling: ~typing.Dict[str, ~typing.Any] = <factory>, dataset_channels: ~typing.List[~sagemaker.serve.ai_inference_recommender.workload._DatasetChannel] = <factory>)

Bases: BaseModel

A workload specification used by benchmark and recommendation jobs.

dataset_channels: List[_DatasetChannel]

classmethod from_dataset(s3_uri: str, *, custom_dataset_type: str | None = None, tokenizer: str | None = 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: str | Secret | None = None, channel_name: str = 'dataset', **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.

Parameters:

• 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.

model_config: ClassVar[ConfigDict] = {'arbitrary_types_allowed': True}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

parameters: Dict[str, Any]

secrets: Dict[str, str | Secret]

classmethod sonnet(**kwargs: Any) → Workload

Alias for synthetic().

AIPerf seeds synthetic prompts from the Sonnet dataset by default, so Workload.sonnet(...) is the same as Workload.synthetic(...).

classmethod synthetic(*, 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: str | Secret | None = 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 Workload.from_dataset() to drive the benchmark from a real request trace instead.

Parameters:

• 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.

classmethod template(request_template: str | None = None, *, template_s3_uri: str | None = None, sagemaker_session: Any | None = None, response_field: str | None = None, tokenizer: str | None = 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: str | Secret | None = None, channel_name: str = 'template', extra_inputs: str | None = 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 template describing your endpoint’s request payload; it is rendered per request (still using synthetic prompts) and sent to your endpoint. A JMESPath 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.

Parameters:

• 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.

to_inline() → str

Serialize the workload to a JSON string.

Secret values are flattened to their ARN strings.

tooling: Dict[str, Any]