本文属于机器翻译版本。若本译文内容与英语原文存在差异，则一律以英文原文为准。

为生成式 AI 推理建议设置工作负载配置

工作负载配置定义了 SageMaker AI 在评估您的模型或端点时使用的流量模式和基准参数。在运行推荐作业或基准测试作业之前，您需要创建工作负载配置。相同的工作负载配置可以在多个作业中重复使用。

您可以通过两种方式定义工作负载：

使用内联参数创建工作负载配置

当您没有具有代表性的数据集时，使用内联参数来指定代币分布。

Python (boto3)

import boto3
import json

client = boto3.client("sagemaker", region_name="us-west-2")

workload_spec = {
    "benchmark": {"type": "aiperf"},
    "parameters": {
        "prompt_input_tokens_mean": 550,
        "prompt_input_tokens_stddev": 150,
        "output_tokens_mean": 150,
        "output_tokens_stddev": 50,
    },
}

response = client.create_ai_workload_config(
    AIWorkloadConfigName="my-workload-config",
    AIWorkloadConfigs={
        "WorkloadSpec": {"Inline": json.dumps(workload_spec)}
    },
)
print(response["AIWorkloadConfigArn"])
        

AWS CLI

aws sagemaker create-ai-workload-config \
  --ai-workload-config-name "my-workload-config" \
  --ai-workload-configs '{"WorkloadSpec": {"Inline": "{\"benchmark\": {\"type\": \"aiperf\"}, \"parameters\": {\"prompt_input_tokens_mean\": 550, \"output_tokens_mean\": 150}}"}}' \
  --region us-west-2
        

使用数据集创建工作负载配置

如果您有真实请求的代表性数据集，请使用带InputDataConfig通道的DatasetConfig参数通过 Amazon S3 提供该数据集。

response = client.create_ai_workload_config(
    AIWorkloadConfigName="my-dataset-workload",
    DatasetConfig={
        "InputDataConfig": [
            {
                "ChannelName": "traffic",
                "DataSource": {
                    "S3DataSource": {
                        "S3Uri": "s3://DOC-EXAMPLE-BUCKET/datasets/traffic-trace/"
                    }
                }
            }
        ]
    },
    AIWorkloadConfigs={
        "WorkloadSpec": {"Inline": json.dumps(workload_spec)}
    },
)
        

默认情况下，会生成合成提示。您也可以使用公共数据集或提供来自 Amazon S3 的自定义数据集。

用于基准测试的工作负载配置

在创建用于对现有端点进行基准测试的工作负载配置时，您可以指定其他参数，例如分词器、并发度、请求计数和请求速率。

workload_spec = {
    "benchmark": {"type": "aiperf"},
    "parameters": {
        "tokenizer": "TinyLlama/TinyLlama-1.1B-Chat-v1.0",
        "concurrency": 1,
        "request_count": 10,
        "streaming": True,
        "prompt_input_tokens_mean": 550,
        "prompt_input_tokens_stddev": 150,
        "output_tokens_mean": 50,
        "output_tokens_stddev": 10,
        "request_rate": 1.0,
        "benchmark_duration": 60,
    },
    "tooling": {"api_standard": "openai", "version": "0.6.0"},
}
        

管理工作负载配置

使用以下操作来管理您的工作负载配置。

# List workload configurations
response = client.list_ai_workload_configs(MaxResults=10)
for config in response["AIWorkloadConfigs"]:
    print(f"{config['AIWorkloadConfigName']} - {config['AIWorkloadConfigArn']}")

# Describe a workload configuration
response = client.describe_ai_workload_config(
    AIWorkloadConfigName="my-workload-config"
)

# Delete a workload configuration
client.delete_ai_workload_config(
    AIWorkloadConfigName="my-workload-config"
)
        

工作负载规格参考

本节提供了您在创建工作负载配置时在WorkloadSpec.Inline字段中传递的工作负载规范 JSON 文档的完整架构。

语法

以下是带有常用参数的工作负载规范的代表性示例。所有参数都记录在下面的参考表中。

{
    "benchmark": {
        "type": "aiperf"
    },
    "parameters": {
        "prompt_input_tokens_mean": 550,
        "prompt_input_tokens_stddev": 150.0,
        "output_tokens_mean": 150,
        "output_tokens_stddev": 50.0,
        "concurrency": 10,
        "request_count": 100,
        "request_rate": 5.0,
        "benchmark_duration": 120,
        "streaming": true,
        "tokenizer": "meta-llama/Llama-3.2-1B"
    },
    "secrets": {
        "hf_token": "arn:aws:secretsmanager:us-west-2:111122223333:secret:my-hf-token-AbCdEf"
    },
    "tooling": {
        "api_standard": "openai"
    }
}
            

工作负载规格密钥

工作负载规范包含以下顶级密钥。未知密钥被拒绝。

基准测试

必需的映射。确定要使用的基准测试工具。

参数

可选的映射。传递给 aiPerf 引擎的基准参数。未知的参数名称会被拒绝。除非另有说明，否则所有参数均为可选参数。

代币分配

流量整形

一般性问题

热身

预热参数控制一个可选的预热阶段，该阶段在测得的基准测试之前运行。这会启动模型服务器的缓存和 JIT 编译。

数据集

图像输入

用于对接受图像输入的多模态模型进行基准测试的参数。

视频输入

用于对接受视频输入的多模态模型进行基准测试的参数。

密钥

可选的映射。使用本节通过 AWS 引用 Secrets Manager 机密而不是将它们作为纯文本包含在中来传递敏感值（例如 Hugging Huging Face 访问令牌）。parameters

每个密钥都是一个参数名称，其值是 Secrets Manager 中 AWS 密钥的 ARN。在作业启动时，该服务会解析密钥并将该值注入基准测试环境。

"secrets": {
    "hf_token": "arn:aws:secretsmanager:us-west-2:111122223333:secret:my-hf-token-AbCdEf"
}
                

工具

可选的映射。指定 API 标准和工具版本。