From 065d642b2d6a5c974c2c47a7f314e6b981814a6f Mon Sep 17 00:00:00 2001 From: smouaa Date: Thu, 23 Jul 2026 23:41:32 +0000 Subject: [PATCH] feat(serve): make BenchmarkResult sweep-aware for concurrency search runs A concurrency search / magic-list sweep (--search-recipe or --concurrency 1,10,100) does not produce a single top-level profile_export_aiperf.json. The benchmark container instead ships search_history.json at the artifact root plus per-trial profile exports nested in subdirs. The existing reader matched profile_export_aiperf.json by filename suffix, so for a sweep it would either fail (no such root file) or silently return one arbitrary level's per-trial metrics as if they were the whole benchmark. BenchmarkResult.from_s3 now checks for search_history.json first and, when present, parses the sweep outcome into a new BenchmarkSearchResult (winning swept value from boundary_summary.feasible_max, the first constraint breach, and the raw history). The single-run path is unchanged. Adds BenchmarkSearchResult (exported), a BenchmarkResult.search field + is_search property, and unit tests covering the sweep layout, the per-trial-name collision, null boundary summaries, and the single-run regression guard. --- .../ai_inference_recommender/__init__.py | 2 + .../serve/ai_inference_recommender/result.py | 128 +++++++++++++++- .../test_result.py | 145 ++++++++++++++++++ 3 files changed, 273 insertions(+), 2 deletions(-) diff --git a/sagemaker-serve/src/sagemaker/serve/ai_inference_recommender/__init__.py b/sagemaker-serve/src/sagemaker/serve/ai_inference_recommender/__init__.py index db1a67aeeb..cad6f696d0 100644 --- a/sagemaker-serve/src/sagemaker/serve/ai_inference_recommender/__init__.py +++ b/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", diff --git a/sagemaker-serve/src/sagemaker/serve/ai_inference_recommender/result.py b/sagemaker-serve/src/sagemaker/serve/ai_inference_recommender/result.py index f1b33a600a..132313e691 100644 --- a/sagemaker-serve/src/sagemaker/serve/ai_inference_recommender/result.py +++ b/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,9 +128,91 @@ 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 @@ -134,8 +220,25 @@ class BenchmarkResult: 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( diff --git a/sagemaker-serve/tests/unit/test_ai_inference_recommender/test_result.py b/sagemaker-serve/tests/unit/test_ai_inference_recommender/test_result.py index 7addf83194..1904845a59 100644 --- a/sagemaker-serve/tests/unit/test_ai_inference_recommender/test_result.py +++ b/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."""