feat(serve): add SageMaker GenAI inference benchmarking and recommendation

Adds sagemaker.serve.ai_inference_recommender, a thin ergonomic layer
over sagemaker-core's AIBenchmarkJob, AIRecommendationJob, and
AIWorkloadConfig resources.

ModelBuilder gains a new entry point and extends two existing verbs:

  # Benchmark a deployed endpoint
  job = mb.start_benchmark(endpoint=ep, workload=Workload.synthetic(...))
  result = BenchmarkResult.from_job(job)

  # Recommendation flow extends optimize() and deploy()
  mb.optimize(workload=..., performance_target="throughput",
              instance_types=["ml.g6.12xlarge"])
  endpoint = mb.deploy(role=role)              # top recommendation
  endpoint = mb.deploy(role=role, recommendation_index=2)  # alternative

print(result) and print(mb.recommendations[0]) render their data as
tables.

Public surface added under sagemaker.serve:

* Workload -- typed factory; extras pass through **params, validated
  server-side.
* BenchmarkResult / BenchmarkMetrics / BenchmarkMetric -- parses the
  AIPerf output.tar.gz from S3.
* Secret -- opt-in helper for tokens >512 chars (Secrets Manager).
* BenchmarkJob, RecommendationJob -- re-exports without the AI prefix.
* FeatureGatedError, WorkloadValidationError -- typed exceptions.

Pin-mode and workload-mode optimize() kwargs are mutually exclusive.
Recommendation deploy uses the ModelPackage path (auto-approves the
package the rec job publishes).

Includes 51 unit tests and 2 slow_test integ tests
(tests/integ/test_ai_inference_recommender_integration.py) verified
end-to-end against real AWS.

Rebased onto upstream to pick up #5860 (preserve falsy values in
sagemaker-core serialize), required so optimize_model=False reaches
the wire.

Author: ZealSV

sagemaker-serve/src/sagemaker/serve/ai_inference_recommender/_model_builder_methods.py

@@ -31,6 +31,7 @@
     AIBenchmarkOutputConfig,
     AIBenchmarkTarget,
     AICapacityReservationConfig,
+    AIDatasetConfig,
     AIModelSource,
     AIModelSourceS3,
     AIRecommendationComputeSpec,
@@ -39,6 +40,9 @@
     AIRecommendationOutputConfig,
     AIRecommendationPerformanceTarget,
     AIWorkloadConfigs,
+    AIWorkloadDataSource,
+    AIWorkloadInputDataConfig,
+    AIWorkloadS3DataSource,
     Tag,
     VpcConfig,
     WorkloadSpec,
@@ -120,6 +124,8 @@ def start_benchmark(
         network_config=network_config,
         tags=tags,
     )
+    if builder is not None:
+        builder._benchmark_job = job
     if wait:
         job.wait()
     return job
@@ -256,11 +262,25 @@ def _ensure_workload_config(
         return workload
 
     config_name = name or f"sm-wl-{int(time.time())}-{uuid.uuid4().hex[:8]}"
+    dataset_config = None
+    if workload.dataset_channels:
+        dataset_config = AIDatasetConfig(
+            input_data_config=[
+                AIWorkloadInputDataConfig(
+                    channel_name=channel.channel_name,
+                    data_source=AIWorkloadDataSource(
+                        s3_data_source=AIWorkloadS3DataSource(s3_uri=channel.s3_uri),
+                    ),
+                )
+                for channel in workload.dataset_channels
+            ],
+        )
     AIWorkloadConfig.create(
         ai_workload_config_name=config_name,
         ai_workload_configs=AIWorkloadConfigs(
             workload_spec=WorkloadSpec(inline=workload.to_inline()),
         ),
+        dataset_config=dataset_config,
         tags=tags,
     )
     return config_name

sagemaker-serve/src/sagemaker/serve/ai_inference_recommender/result.py

@@ -115,6 +115,9 @@ class BenchmarkResult:
 
     metrics: BenchmarkMetrics
     s3_output_location: str
+    endpoint: Optional[str] = None
+    workload_config: Optional[str] = None
+    tool_version: Optional[str] = None
     profile: Dict[str, Any] = field(default_factory=dict)
 
     def __repr__(self) -> str:
@@ -134,6 +137,9 @@ def __repr__(self) -> str:
         table = _format_metrics_table(ordered)
         return (
             f"BenchmarkResult\n"
+            f"  endpoint:           {self.endpoint or '-'}\n"
+            f"  workload_config:    {self.workload_config or '-'}\n"
+            f"  tool_version:       {self.tool_version or '-'}\n"
             f"  s3_output_location: {self.s3_output_location}\n"
             f"  metrics:\n{_indent(table, '    ')}\n"
             f"  raw profile available via .profile"
@@ -148,6 +154,10 @@ def from_job(
     ) -> "BenchmarkResult":
         """Download and parse the benchmark output for a completed ``AIBenchmarkJob``.
 
+        Populates ``endpoint``, ``workload_config``, and ``tool_version`` from
+        the job's ``BenchmarkTarget`` and ``WorkloadConfigIdentifier`` plus the
+        AIPerf profile metadata so the parsed result is self-describing.
+
         Args:
             job: An ``AIBenchmarkJob`` (or ``BenchmarkJob`` re-export) that has
                 reached a terminal state.
@@ -183,21 +193,35 @@ def from_job(
                 f"AIBenchmarkJob {job.get_name()} has no S3OutputLocation "
                 f"(status={status}). {hint}"
             )
-        return cls.from_s3(job.output_config.s3_output_location, session=session)
+        workload_config = getattr(job, "ai_workload_config_identifier", None)
+        return cls.from_s3(
+            job.output_config.s3_output_location,
+            session=session,
+            endpoint=_extract_endpoint(job),
+            # Normalize falsy sentinels (e.g. unset optional fields) to None
+            # so the result renders cleanly when fields are missing.
+            workload_config=workload_config or None,
+        )
 
     @classmethod
     def from_s3(
         cls,
         s3_output_location: str,
         *,
         session: Optional[boto3.session.Session] = None,
+        endpoint: Optional[str] = None,
+        workload_config: Optional[str] = None,
     ) -> "BenchmarkResult":
         """Download and parse the benchmark output artifact from S3.
 
         Args:
             s3_output_location: ``s3://bucket/prefix/`` location written by
                 the benchmark job.
             session: Optional boto3 session. Defaults to the ambient session.
+            endpoint: Optional endpoint identifier to attach to the result.
+                Threaded through by :meth:`from_job`.
+            workload_config: Optional workload-config identifier to attach.
+                Threaded through by :meth:`from_job`.
 
         Returns:
             A parsed ``BenchmarkResult``.
@@ -216,10 +240,39 @@ def from_s3(
         return cls(
             metrics=BenchmarkMetrics.from_profile_json(profile),
             s3_output_location=s3_output_location,
+            endpoint=endpoint,
+            workload_config=workload_config,
+            tool_version=_extract_tool_version(profile),
             profile=profile,
         )
 
 
+def _extract_endpoint(job) -> Optional[str]:
+    target = getattr(job, "benchmark_target", None) or None
+    endpoint = (getattr(target, "endpoint", None) or None) if target else None
+    identifier = getattr(endpoint, "identifier", None) if endpoint else None
+    return identifier or None
+
+
+def _extract_tool_version(profile: Dict[str, Any]) -> Optional[str]:
+    """Best-effort lookup of the AIPerf tool version from the profile JSON.
+
+    AIPerf has no single canonical key; we check a few plausible top-level
+    locations and return the first string we find.
+    """
+    for key in ("aiperf_version", "tool_version", "version"):
+        value = profile.get(key)
+        if isinstance(value, str):
+            return value
+    meta = profile.get("metadata") or profile.get("meta") or {}
+    if isinstance(meta, dict):
+        for key in ("aiperf_version", "tool_version", "version"):
+            value = meta.get(key)
+            if isinstance(value, str):
+                return value
+    return None
+
+
 def _parse_s3_uri(uri: str) -> tuple:
     parsed = urlparse(uri)
     if parsed.scheme != "s3":

sagemaker-serve/src/sagemaker/serve/ai_inference_recommender/workload.py

@@ -14,13 +14,27 @@
 from __future__ import absolute_import
 
 import json
-from typing import Any, Dict, Optional, Union
+from typing import Any, Dict, List, Optional, Union
 
 from pydantic import BaseModel, ConfigDict, Field
 
 from sagemaker.serve.ai_inference_recommender.secrets import Secret
 
 
+class _DatasetChannel(BaseModel):
+    """Internal record of an S3 channel that should be mounted at job runtime.
+
+    Used by ``Workload.from_dataset`` to capture the S3 URI and channel name
+    so :class:`AIWorkloadConfig` can be created with the matching
+    ``DatasetConfig``. AIPerf reads the dataset from the container path
+    declared in ``parameters.input_file``; this record exists only to plumb
+    the S3-side configuration through to the service.
+    """
+
+    channel_name: str
+    s3_uri: str
+
+
 class Workload(BaseModel):
     """A workload specification used by benchmark and recommendation jobs."""
 
@@ -29,6 +43,7 @@ class Workload(BaseModel):
     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)
 
     @classmethod
     def synthetic(
@@ -47,6 +62,11 @@ def synthetic(
     ) -> "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.
@@ -77,6 +97,88 @@ def synthetic(
             secrets["hf_token"] = hf_token
         return cls(parameters=parameters, secrets=secrets)
 
+    @classmethod
+    def from_dataset(
+        cls,
+        *,
+        s3_uri: str,
+        channel_name: str,
+        input_file: str,
+        custom_dataset_type: Optional[str] = None,
+        tokenizer: Optional[str] = None,
+        concurrency: int = 1,
+        request_count: int = 100,
+        streaming: bool = True,
+        hf_token: Optional[Union[str, Secret]] = None,
+        **params: Any,
+    ) -> "Workload":
+        """Build a workload that drives traffic from an S3-hosted dataset.
+
+        The S3 location is mounted into the AIPerf container at
+        ``/opt/ml/input/data/{channel_name}/`` by the SageMaker AI inference
+        recommender service via the workload config's ``DatasetConfig``.
+        AIPerf reads the file declared by ``input_file`` from inside that
+        mount.
+
+        Args:
+            s3_uri: ``s3://bucket/prefix/`` (or single-object) URI containing
+                the dataset. Mounted under the container at
+                ``/opt/ml/input/data/{channel_name}/``.
+            channel_name: Logical channel name. Used as the directory under
+                ``/opt/ml/input/data/`` where the S3 contents are mounted.
+            input_file: Container-internal path to the file AIPerf should
+                read, e.g. ``/opt/ml/input/data/traffic/requests.jsonl``.
+                Must match ``s3_uri`` + ``channel_name`` for the mount to
+                contain the file.
+            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.
+            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.
+
+        Returns:
+            A ``Workload`` whose inline payload references ``input_file`` and
+            whose ``dataset_channels`` carry the S3 URI for the service to
+            mount.
+        """
+        if not s3_uri.startswith("s3://"):
+            raise ValueError(
+                f"s3_uri must start with 's3://'; got {s3_uri!r}."
+            )
+        expected_prefix = f"/opt/ml/input/data/{channel_name}/"
+        if not input_file.startswith(expected_prefix):
+            raise ValueError(
+                f"input_file must live under {expected_prefix!r} so the "
+                f"mounted channel contains it; got {input_file!r}."
+            )
+        parameters: Dict[str, Any] = {
+            "input_file": input_file,
+            "concurrency": concurrency,
+            "request_count": request_count,
+            "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),
+            ],
+        )
+
     def to_inline(self) -> str:
         """Serialize the workload to a JSON string.