feat(serve): make BenchmarkResult sweep-aware for concurrency search runs

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

@@ -29,6 +29,7 @@
     BenchmarkMetric,
     BenchmarkMetrics,
     BenchmarkResult,
+    BenchmarkSearchResult,
 )
 from sagemaker.serve.ai_inference_recommender.secrets import Secret
 from sagemaker.serve.ai_inference_recommender.workload import Workload
@@ -42,6 +43,7 @@
     "BenchmarkMetric",
     "BenchmarkMetrics",
     "BenchmarkResult",
+    "BenchmarkSearchResult",
     "FeatureGatedError",
     "InferenceFramework",
     "PerformanceTarget",

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

@@ -24,6 +24,10 @@
 
 PROFILE_EXPORT_FILENAME = "profile_export_aiperf.json"
 OUTPUT_ARCHIVE_FILENAME = "output.tar.gz"
+# A concurrency search / magic-list sweep writes this at the artifact root
+# instead of a single top-level profile_export_aiperf.json. Its presence in the
+# archive is how we tell a sweep run apart from a single run.
+SEARCH_HISTORY_FILENAME = "search_history.json"
 
 
 @dataclass
@@ -124,18 +128,117 @@ def from_profile_json(cls, profile: Dict[str, Any]) -> "BenchmarkMetrics":
 )
 
 
+@dataclass
+class BenchmarkSearchResult:
+    """Outcome of a concurrency search / magic-list sweep benchmark.
+
+    A search run does not produce a single ``profile_export_aiperf.json``; it
+    sweeps a dimension (typically concurrency) and records the outcome in
+    ``search_history.json``. This captures the parts callers care about: which
+    swept value won, and the raw history for anything deeper.
+
+    Attributes:
+        swept_dim: dotted path of the swept dimension, e.g.
+            ``"phases.profiling.concurrency"``.
+        winner: the largest feasible swept value (``boundary_summary.feasible_max.value``).
+            ``None`` if no feasible point was found (e.g. every level breached the SLA).
+        winner_objective: the objective value at the winning point, when reported.
+        infeasible_min: the smallest swept value that breached a constraint, if any.
+        first_breach: details of the constraint the ``infeasible_min`` level breached
+            (metric tag / stat / threshold / observed), if reported.
+        raw: the full parsed ``search_history.json`` for callers who need more.
+    """
+
+    swept_dim: Optional[str] = None
+    winner: Optional[float] = None
+    winner_objective: Optional[float] = None
+    infeasible_min: Optional[float] = None
+    first_breach: Optional[Dict[str, Any]] = None
+    raw: Dict[str, Any] = field(default_factory=dict)
+
+    @classmethod
+    def from_history_json(cls, history: Dict[str, Any]) -> "BenchmarkSearchResult":
+        # boundary_summary is present for a single-dimension search with a
+        # resolved boundary; it is null/absent for a multi-dim search or one
+        # that never ran, in which case we still return a result carrying the
+        # raw history rather than fabricating a winner.
+        boundary = history.get("boundary_summary")
+        if not isinstance(boundary, dict):
+            return cls(raw=history)
+        feasible_max = boundary.get("feasible_max")
+        if not isinstance(feasible_max, dict):
+            feasible_max = {}
+        infeasible = boundary.get("infeasible_min")
+        if not isinstance(infeasible, dict):
+            infeasible = {}
+        first_breach = infeasible.get("first_breach")
+        return cls(
+            swept_dim=boundary.get("swept_dim_path"),
+            winner=_as_float(feasible_max.get("value")),
+            winner_objective=_as_float(feasible_max.get("objective_value")),
+            infeasible_min=_as_float(infeasible.get("value")),
+            first_breach=first_breach if isinstance(first_breach, dict) else None,
+            raw=history,
+        )
+
+    def __str__(self) -> str:
+        breach = ""
+        if self.infeasible_min is not None:
+            metric = (self.first_breach or {}).get("metric_tag", "?")
+            breach = f"\n  first_breach:       {metric} at {self.infeasible_min}"
+        return (
+            f"BenchmarkSearchResult\n"
+            f"  swept_dim:          {self.swept_dim or '-'}\n"
+            f"  winner:             {_fmt_number(self.winner)}\n"
+            f"  winner_objective:   {_fmt_number(self.winner_objective)}"
+            f"{breach}\n"
+            f"  raw history available via .raw"
+        )
+
+    def __repr__(self) -> str:
+        return f"BenchmarkSearchResult(swept_dim={self.swept_dim!r}, winner={self.winner!r})"
+
+    def _repr_pretty_(self, p, cycle):
+        # Render the full summary in notebooks (Jupyter uses this hook).
+        p.text("..." if cycle else str(self))
+
+
 @dataclass
 class BenchmarkResult:
-    """Parsed result of a completed benchmark job."""
+    """Parsed result of a completed benchmark job.
+
+    For a single run, ``metrics``/``profile`` carry the AIPerf profile export
+    and ``search`` is ``None``. For a concurrency search / magic-list sweep,
+    ``search`` carries the sweep outcome (the winning level); ``metrics`` is
+    empty and ``profile`` holds the raw ``search_history.json`` — a sweep has
+    no single headline profile to report.
+    """
 
     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)
+    search: Optional[BenchmarkSearchResult] = None
+
+    @property
+    def is_search(self) -> bool:
+        """True if this result came from a concurrency search / sweep run."""
+        return self.search is not None
 
     def __str__(self) -> str:
+        # A search/sweep run has no single headline profile; render the sweep
+        # outcome (winning level) instead of an (empty) metrics table.
+        if self.search is not None:
+            return (
+                f"BenchmarkResult (search)\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"  search:\n{_indent(str(self.search), '    ')}"
+            )
         # Order: well-known headline metrics first, then everything else
         # alphabetized, then HTTP-level transport metrics last (they're
         # noise for most readers, useful only for debugging).
@@ -267,10 +370,31 @@ def from_s3(
         archive_key = _find_object(s3, bucket, prefix, OUTPUT_ARCHIVE_FILENAME)
         body = s3.get_object(Bucket=bucket, Key=archive_key)["Body"].read()
 
+        # A concurrency search / magic-list sweep writes search_history.json at
+        # the artifact root and NO top-level profile_export_aiperf.json (each
+        # swept level has its own per-trial profile export nested in a subdir).
+        # Check for the search history FIRST: those per-trial exports share the
+        # profile_export_aiperf.json name, so a plain suffix match would
+        # otherwise silently return one arbitrary level's metrics as if they
+        # were the whole benchmark.
+        history_bytes = _read_member_from_tar_gz(body, SEARCH_HISTORY_FILENAME)
+        if history_bytes is not None:
+            history = json.loads(history_bytes.decode("utf-8"))
+            return cls(
+                metrics=BenchmarkMetrics.from_profile_json({}),
+                s3_output_location=s3_output_location,
+                endpoint=endpoint,
+                workload_config=workload_config,
+                tool_version=_extract_tool_version(history),
+                profile=history,
+                search=BenchmarkSearchResult.from_history_json(history),
+            )
+
         profile_bytes = _read_member_from_tar_gz(body, PROFILE_EXPORT_FILENAME)
         if profile_bytes is None:
             raise FileNotFoundError(
-                f"{PROFILE_EXPORT_FILENAME} not found in s3://{bucket}/{archive_key}"
+                f"Neither {PROFILE_EXPORT_FILENAME} nor {SEARCH_HISTORY_FILENAME} "
+                f"found in s3://{bucket}/{archive_key}"
             )
         profile = json.loads(profile_bytes.decode("utf-8"))
         return cls(

sagemaker-serve/tests/unit/test_ai_inference_recommender/test_result.py

@@ -20,6 +20,7 @@
     BenchmarkMetric,
     BenchmarkMetrics,
     BenchmarkResult,
+    BenchmarkSearchResult,
 )
 
 
@@ -286,6 +287,150 @@ def client(self, name):
             )
 
 
+# A concurrency search writes search_history.json at the artifact root. Schema
+# per aiperf v0.11.0 (docs/sweeping/search-recipes.md): boundary_summary carries
+# the winning (feasible_max) and first-breaching (infeasible_min) swept values.
+SAMPLE_SEARCH_HISTORY = {
+    "aiperf_version": "0.11.0",
+    "boundary_summary": {
+        "swept_dim_path": "phases.profiling.concurrency",
+        "feasible_max": {"value": 256, "iteration_idx": 0, "objective_value": 4172.3},
+        "infeasible_min": {
+            "value": 512,
+            "first_breach": {
+                "metric_tag": "time_to_first_token",
+                "stat": "p95",
+                "op": "lt",
+                "threshold": 200.0,
+                "observed": 213.4,
+            },
+        },
+    },
+}
+
+
+def _make_search_archive(history=None, include_per_trial_profile=True) -> bytes:
+    """Build an output.tar.gz for a sweep run.
+
+    Mirrors what the benchmark container ships: search_history.json at the root
+    and (optionally) a per-trial profile_export_aiperf.json nested in a swept
+    level's subdir — the exact shape that would trip a naive suffix match.
+    """
+    history_bytes = json.dumps(history or SAMPLE_SEARCH_HISTORY).encode("utf-8")
+    buf = io.BytesIO()
+    with tarfile.open(fileobj=buf, mode="w:gz") as tar:
+        info = tarfile.TarInfo(name="search_history.json")
+        info.size = len(history_bytes)
+        tar.addfile(info, io.BytesIO(history_bytes))
+        if include_per_trial_profile:
+            # A per-level export shares the single-run profile filename; it must
+            # NOT be mistaken for the headline profile.
+            trial = json.dumps({"request_throughput": {"avg": 999.0}}).encode("utf-8")
+            tinfo = tarfile.TarInfo(
+                name="search_iter_0000/profile_runs/run_0000/profile_export_aiperf.json"
+            )
+            tinfo.size = len(trial)
+            tar.addfile(tinfo, io.BytesIO(trial))
+    return buf.getvalue()
+
+
+def _parse_archive_from_s3(archive_bytes: bytes) -> BenchmarkResult:
+    s3_client = boto3.session.Session(region_name="us-east-1").client("s3")
+    with Stubber(s3_client) as stub:
+        stub.add_response(
+            "list_objects_v2",
+            {
+                "Contents": [{"Key": "p/output.tar.gz", "Size": len(archive_bytes)}],
+                "KeyCount": 1,
+            },
+            expected_params={"Bucket": "b", "Prefix": "p/"},
+        )
+        stub.add_response(
+            "get_object",
+            {"Body": _StreamingBody(archive_bytes)},
+            expected_params={"Bucket": "b", "Key": "p/output.tar.gz"},
+        )
+
+        class _SessionStub:
+            def client(self, name):
+                return s3_client
+
+        return BenchmarkResult.from_s3("s3://b/p/", session=_SessionStub())
+
+
+class TestBenchmarkSearchResultFromHistory:
+    def test_winner_and_boundary_parsed(self):
+        result = BenchmarkSearchResult.from_history_json(SAMPLE_SEARCH_HISTORY)
+        assert result.swept_dim == "phases.profiling.concurrency"
+        assert result.winner == 256.0
+        assert result.winner_objective == 4172.3
+        assert result.infeasible_min == 512.0
+        assert result.first_breach["metric_tag"] == "time_to_first_token"
+        assert result.raw == SAMPLE_SEARCH_HISTORY
+
+    def test_no_infeasible_min_when_no_sla_breach(self):
+        # No SLA filter → nothing breaches → infeasible_min is null, winner is
+        # the highest swept value.
+        history = {
+            "boundary_summary": {
+                "swept_dim_path": "phases.profiling.concurrency",
+                "feasible_max": {"value": 1024},
+                "infeasible_min": None,
+            }
+        }
+        result = BenchmarkSearchResult.from_history_json(history)
+        assert result.winner == 1024.0
+        assert result.infeasible_min is None
+        assert result.first_breach is None
+
+    def test_null_boundary_summary_yields_no_winner(self):
+        # Multi-dim search (or one that never resolved a boundary) → null block.
+        # We keep the raw history but report no winner rather than fabricating one.
+        history = {"boundary_summary": None, "iterations": []}
+        result = BenchmarkSearchResult.from_history_json(history)
+        assert result.winner is None
+        assert result.swept_dim is None
+        assert result.raw == history
+
+    def test_missing_boundary_summary_key(self):
+        result = BenchmarkSearchResult.from_history_json({"iterations": []})
+        assert result.winner is None
+        assert result.raw == {"iterations": []}
+
+
+class TestBenchmarkResultFromS3Search:
+    def test_search_run_parsed_from_history_not_per_trial_profile(self):
+        result = _parse_archive_from_s3(_make_search_archive())
+        # The result must reflect the search, NOT the stray per-trial profile
+        # (request_throughput avg 999.0) that shares the single-run filename.
+        assert result.is_search is True
+        assert result.search is not None
+        assert result.search.winner == 256.0
+        assert result.metrics.all_metrics == {}
+        assert result.metrics.request_throughput is None
+        assert result.profile == SAMPLE_SEARCH_HISTORY
+        assert result.tool_version == "0.11.0"
+
+    def test_search_history_preferred_even_without_per_trial_profile(self):
+        result = _parse_archive_from_s3(_make_search_archive(include_per_trial_profile=False))
+        assert result.is_search is True
+        assert result.search.swept_dim == "phases.profiling.concurrency"
+
+    def test_single_run_still_not_flagged_as_search(self):
+        # Regression guard: the single-run path must be untouched.
+        result = _parse_archive_from_s3(_make_output_archive())
+        assert result.is_search is False
+        assert result.search is None
+        assert result.metrics.request_throughput.avg == 12.5
+
+    def test_str_renders_search_summary(self):
+        result = _parse_archive_from_s3(_make_search_archive())
+        text = str(result)
+        assert "BenchmarkResult (search)" in text
+        assert "phases.profiling.concurrency" in text
+        assert "256" in text
+
+
 class TestFindObjectPagination:
     """_find_object must paginate so keys beyond the first page are found."""