From 31956b7a66d46661ec27a35a88a40787be985634 Mon Sep 17 00:00:00 2001 From: ZealSV Date: Mon, 10 Aug 2026 15:44:24 -0700 Subject: [PATCH] fix(serve): pre-deploy JumpStart benchmark data + public HuggingFace download helper Fixes ModelBuilder benchmark APIs for pre-deploy JumpStart models and adds a public HuggingFace-Hub download helper: - benchmark_metrics: display_benchmark_metrics() read self.benchmark_metrics, a property that existed in v2 but was dropped in the v3 port, so it raised AttributeError. Restore the property (a pandas DataFrame built from the model's published deployment-config benchmark data). - list_deployment_configs(): with no instance_type it called _get_deployment_configs() before the JumpStart metadata configs were loaded, so a pre-deploy builder (_metadata_configs=None) got []. _get_deployment_configs now calls _ensure_metadata_configs() itself, so every caller lazily loads them. - ragged benchmark data: get_metrics_from_deployment_configs built columns independently, appending the instance-rate (pricing) column only for instances that had it. That produced unequal-length columns (crashing pd.DataFrame with "All arrays must be of the same length") AND misaligned the present rate values onto the wrong rows. Assemble one record per row and pivot to columns at the end, padding unset cells with None, so columns stay equal length and every value stays on its own row. Column layout is unchanged (the rate column still comes last). - download_huggingface_model(): new public helper in serve.utils.hf_utils, exported from sagemaker.serve. Downloads a Hub snapshot and optionally uploads it to S3, so notebooks no longer hand-roll snapshot_download + S3Uploader. Supports revision / allow_patterns / ignore_patterns passthroughs; when staging to S3 without a caller-supplied local_dir it uses a temporary directory that is removed after the upload. migration.md documents both new surfaces. Unit: 11 new tests (benchmark/config, ragged-metrics alignment/order, HF helper incl. temp-dir cleanup and passthroughs). --- migration.md | 2 + .../src/sagemaker/core/jumpstart/utils.py | 45 ++++++--- .../tests/unit/test_jumpstart_utils.py | 47 +++++++++ .../src/sagemaker/serve/__init__.py | 2 + .../src/sagemaker/serve/model_builder.py | 14 +++ .../sagemaker/serve/model_builder_utils.py | 5 + .../src/sagemaker/serve/utils/hf_utils.py | 94 ++++++++++++++++++ .../unit/test_model_builder_coverage_boost.py | 81 ++++++++++++++- .../tests/unit/utils/test_hf_utils.py | 99 ++++++++++++++++++- 9 files changed, 373 insertions(+), 16 deletions(-) diff --git a/migration.md b/migration.md index a7f353e747..0ea3e9bef1 100644 --- a/migration.md +++ b/migration.md @@ -441,6 +441,8 @@ pipeline = Pipeline( | Predictor | Endpoint | Replaced with sagemaker-core | | MultiDataModel | ModelBuilder | Multi-model endpoints | | AsyncPredictor | ModelBuilder | Async inference | +| JumpStartModel.benchmark_metrics / display_benchmark_metrics() / list_deployment_configs() | ModelBuilder.benchmark_metrics / display_benchmark_metrics() / list_deployment_configs() | Same names; usable on a pre-deploy JumpStart `ModelBuilder` | +| snapshot_download + S3Uploader (hand-rolled) | `from sagemaker.serve import download_huggingface_model` | Downloads a Hub snapshot and optionally uploads it to S3 | ### Processing Features diff --git a/sagemaker-core/src/sagemaker/core/jumpstart/utils.py b/sagemaker-core/src/sagemaker/core/jumpstart/utils.py index 91d6eec1c8..92109f203e 100644 --- a/sagemaker-core/src/sagemaker/core/jumpstart/utils.py +++ b/sagemaker-core/src/sagemaker/core/jumpstart/utils.py @@ -1442,8 +1442,19 @@ def get_metrics_from_deployment_configs( if not deployment_configs: return {} - data = {"Instance Type": [], "Config Name": [], "Concurrent Users": []} - instance_rate_data = {} + # Build per-row records, then pivot to column-oriented at the end, padding + # columns a row didn't set with None. Keeps every column the same length and + # each value on its own row; appending columns independently would crash + # pd.DataFrame on ragged lengths and misalign the sparse (pricing) column. + rows: List[Dict[str, Any]] = [] + column_order: List[str] = ["Instance Type", "Config Name", "Concurrent Users"] + seen_columns = set(column_order) + + def _register_column(name: str) -> None: + if name not in seen_columns: + seen_columns.add(name) + column_order.append(name) + for index, deployment_config in enumerate(deployment_configs): benchmark_metrics = deployment_config.benchmark_metrics if not deployment_config.deployment_args or not benchmark_metrics: @@ -1465,25 +1476,31 @@ def get_metrics_from_deployment_configs( else current_instance_type ) - data["Config Name"].append(deployment_config.deployment_config_name) - data["Instance Type"].append(instance_type_to_display) - data["Concurrent Users"].append(concurrent_user) + row: Dict[str, Any] = { + "Instance Type": instance_type_to_display, + "Config Name": deployment_config.deployment_config_name, + "Concurrent Users": concurrent_user, + } + + for metric in metrics: + column_name = _normalize_benchmark_metric_column_name(metric.name, metric.unit) + _register_column(column_name) + row[column_name] = metric.value if instance_type_rate: instance_rate_column_name = ( f"{instance_type_rate.name} ({instance_type_rate.unit})" ) - instance_rate_data[instance_rate_column_name] = instance_rate_data.get( - instance_rate_column_name, [] - ) - instance_rate_data[instance_rate_column_name].append(instance_type_rate.value) + _register_column(instance_rate_column_name) + row[instance_rate_column_name] = instance_type_rate.value - for metric in metrics: - column_name = _normalize_benchmark_metric_column_name(metric.name, metric.unit) - data[column_name] = data.get(column_name, []) - data[column_name].append(metric.value) + rows.append(row) - data = {**data, **instance_rate_data} + # Pivot to column-oriented, padding unset cells with None. + data: Dict[str, List[Any]] = {column: [] for column in column_order} + for row in rows: + for column in column_order: + data[column].append(row.get(column)) return data diff --git a/sagemaker-core/tests/unit/test_jumpstart_utils.py b/sagemaker-core/tests/unit/test_jumpstart_utils.py index 8cd6edd510..08eb20ccdc 100644 --- a/sagemaker-core/tests/unit/test_jumpstart_utils.py +++ b/sagemaker-core/tests/unit/test_jumpstart_utils.py @@ -1679,6 +1679,53 @@ def test_get_metrics_from_deployment_configs_with_metrics(self): assert "Instance Type" in result assert "Config Name" in result + def test_get_metrics_ragged_pricing_stays_equal_length_and_aligned(self): + """One instance with a pricing overlay, one without: columns stay equal + length and the rate stays aligned to its instance (None for the other).""" + + def _stat(name, unit, value, concurrency): + stat = Mock() + stat.name = name + stat.unit = unit + stat.value = value + stat.concurrency = concurrency + return stat + + mock_args = Mock() + mock_args.default_instance_type = "ml.g5.xlarge" + mock_args.instance_type = "ml.g5.xlarge" + + mock_config = Mock(spec=DeploymentConfigMetadata) + mock_config.deployment_args = mock_args + mock_config.deployment_config_name = "config1" + # priced instance carries an "Instance Rate" stat; the other does not. + mock_config.benchmark_metrics = { + "ml.g5.xlarge": [ + _stat("Latency", "ms", 100, "1"), + _stat("Instance Rate", "USD/Hr", 1.23, "1"), + ], + "ml.g5.2xlarge": [ + _stat("Latency", "ms", 90, "1"), + ], + } + + result = utils.get_metrics_from_deployment_configs([mock_config]) + + # Equal-length columns -> pd.DataFrame(result) will not raise. + lengths = {column: len(values) for column, values in result.items()} + assert len(set(lengths.values())) == 1, lengths + + # Rate value stays on the priced row, None on the other (not shifted). + rate_cols = [c for c in result if "Instance Rate" in c] + assert rate_cols, result.keys() + rate_col = rate_cols[0] + by_instance = dict(zip(result["Instance Type"], result[rate_col])) + assert by_instance["ml.g5.xlarge (Default)"] == 1.23 + assert by_instance["ml.g5.2xlarge"] is None + + # Rate column stays last, after the metric columns (layout unchanged). + assert list(result).index(rate_col) == len(result) - 1 + class TestNormalizeBenchmarkMetricColumnName: """Test cases for _normalize_benchmark_metric_column_name function""" diff --git a/sagemaker-serve/src/sagemaker/serve/__init__.py b/sagemaker-serve/src/sagemaker/serve/__init__.py index 9cf115d1ab..bca977b54d 100644 --- a/sagemaker-serve/src/sagemaker/serve/__init__.py +++ b/sagemaker-serve/src/sagemaker/serve/__init__.py @@ -41,6 +41,7 @@ WorkloadValidationError, start_benchmark, ) +from sagemaker.serve.utils.hf_utils import download_huggingface_model __all__ = [ "InferenceSpec", @@ -56,4 +57,5 @@ "FeatureGatedError", "WorkloadValidationError", "start_benchmark", + "download_huggingface_model", ] diff --git a/sagemaker-serve/src/sagemaker/serve/model_builder.py b/sagemaker-serve/src/sagemaker/serve/model_builder.py index 7d37bf4d57..2e86e47ca1 100644 --- a/sagemaker-serve/src/sagemaker/serve/model_builder.py +++ b/sagemaker-serve/src/sagemaker/serve/model_builder.py @@ -4782,6 +4782,20 @@ def transformer( sagemaker_session=self.sagemaker_session, ) + @property + def benchmark_metrics(self): + """Benchmark metrics for the model's JumpStart deployment configs. + + Returns a pandas ``DataFrame`` (one row per config/instance) built from + the model's published benchmark data. Available for JumpStart models + (or HuggingFace models with a JumpStart equivalent) before deploy. + """ + import pandas as pd + + df = pd.DataFrame(self._get_deployment_configs_benchmarks_data()) + df.index = [""] * len(df) + return df + @_telemetry_emitter( feature=Feature.MODEL_CUSTOMIZATION, func_name="model_builder.display_benchmark_metrics" ) diff --git a/sagemaker-serve/src/sagemaker/serve/model_builder_utils.py b/sagemaker-serve/src/sagemaker/serve/model_builder_utils.py index f4a90a2d05..1ea1265f53 100644 --- a/sagemaker-serve/src/sagemaker/serve/model_builder_utils.py +++ b/sagemaker-serve/src/sagemaker/serve/model_builder_utils.py @@ -2709,6 +2709,11 @@ def _get_deployment_configs( selected_config_name (Optional[str]): The name of the selected deployment config. selected_instance_type (Optional[str]): The selected instance type. """ + # Lazily load the JumpStart metadata configs. Without this a pre-deploy + # builder (model set, build()/deploy() not yet called) has + # _metadata_configs=None, so both list_deployment_configs() and the + # benchmark-metrics data would come back empty. + self._ensure_metadata_configs() deployment_configs = [] if not self._metadata_configs: return deployment_configs diff --git a/sagemaker-serve/src/sagemaker/serve/utils/hf_utils.py b/sagemaker-serve/src/sagemaker/serve/utils/hf_utils.py index b980edabe4..e403255e43 100644 --- a/sagemaker-serve/src/sagemaker/serve/utils/hf_utils.py +++ b/sagemaker-serve/src/sagemaker/serve/utils/hf_utils.py @@ -14,14 +14,108 @@ from __future__ import absolute_import import json +import os +import tempfile import urllib.request from json import JSONDecodeError +from typing import Optional from urllib.error import HTTPError, URLError import logging logger = logging.getLogger(__name__) +def download_huggingface_model( + model_id: str, + *, + local_dir: Optional[str] = None, + s3_uri: Optional[str] = None, + hf_hub_token: Optional[str] = None, + revision: Optional[str] = None, + allow_patterns=None, + ignore_patterns=None, + sagemaker_session=None, +) -> str: + """Download a HuggingFace Hub model snapshot, optionally staging it to S3. + + A supported, importable helper so notebooks and scripts don't hand-roll + ``huggingface_hub.snapshot_download`` + ``S3Uploader``. Downloads the full + model snapshot from the Hub, then either leaves it on local disk or uploads + it to S3 and returns the resulting location. + + Args: + model_id: The HuggingFace Hub model id (e.g. ``"gpt2"``). + local_dir: Local directory to download into. When ``s3_uri`` is given + and ``local_dir`` is omitted, the snapshot is downloaded into a + temporary directory that is removed after the upload. Defaults to + ``None``. + s3_uri: Optional ``s3://bucket/prefix`` destination. When set, the + snapshot is uploaded there and the returned value is the S3 URI. + hf_hub_token: Optional HuggingFace Hub token for gated/private models. + revision: Optional Hub revision (branch, tag, or commit) to pin; passed + through to ``snapshot_download``. Defaults to the repo's default + branch. + allow_patterns: Optional glob(s) of files to include, passed through to + ``snapshot_download`` (e.g. ``"*.safetensors"`` to skip duplicate + ``.bin`` weights). + ignore_patterns: Optional glob(s) of files to exclude, passed through to + ``snapshot_download``. + sagemaker_session: Optional session used for the S3 upload. Defaults to + a new session built from the ambient AWS configuration. + + Returns: + The S3 URI the snapshot was uploaded to when ``s3_uri`` is given, + otherwise the local directory path the snapshot was downloaded into. + + Raises: + ImportError: If ``huggingface_hub`` is not installed. + ValueError: If neither ``local_dir`` nor ``s3_uri`` is given. + """ + if local_dir is None and s3_uri is None: + raise ValueError("Provide local_dir, s3_uri, or both.") + + try: + from huggingface_hub import snapshot_download + except ImportError as exc: + raise ImportError( + "download_huggingface_model requires huggingface_hub, which is not " + "installed. Install it with `pip install huggingface_hub`." + ) from exc + + def _download(target: str) -> None: + os.makedirs(target, exist_ok=True) + logger.info("Downloading model %s from Hugging Face Hub to %s", model_id, target) + snapshot_download( + repo_id=model_id, + local_dir=target, + token=hf_hub_token, + revision=revision, + allow_patterns=allow_patterns, + ignore_patterns=ignore_patterns, + ) + + if s3_uri is None: + _download(local_dir) + return local_dir + + from sagemaker.core.s3 import S3Uploader + + def _upload(source: str) -> str: + logger.info("Uploading model %s snapshot to %s", model_id, s3_uri) + return S3Uploader.upload( + local_path=source, + desired_s3_uri=s3_uri, + sagemaker_session=sagemaker_session, + ) + + if local_dir is not None: + _download(local_dir) + return _upload(local_dir) + with tempfile.TemporaryDirectory(prefix="hf-model-") as staging_dir: + _download(staging_dir) + return _upload(staging_dir) + + def _get_model_config_properties_from_hf(model_id: str, hf_hub_token: str = None): """Placeholder docstring""" diff --git a/sagemaker-serve/tests/unit/test_model_builder_coverage_boost.py b/sagemaker-serve/tests/unit/test_model_builder_coverage_boost.py index e67b7eaee8..d6c9f35dbf 100644 --- a/sagemaker-serve/tests/unit/test_model_builder_coverage_boost.py +++ b/sagemaker-serve/tests/unit/test_model_builder_coverage_boost.py @@ -4,7 +4,7 @@ """ import unittest -from unittest.mock import Mock, patch, MagicMock +from unittest.mock import Mock, patch, MagicMock, PropertyMock from dataclasses import dataclass import tempfile @@ -333,6 +333,85 @@ def test_list_deployment_configs_non_string_model(self): self.assertIn("only supported for JumpStart", str(context.exception)) +class TestPreDeployBenchmarkData(unittest.TestCase): + """Pre-deploy JumpStart benchmark data: list_deployment_configs() and the + benchmark_metrics property must work before build()/deploy() is called. + """ + + def _jumpstart_mb(self): + """A JumpStart-config ModelBuilder, constructed the same way as + TestFromJumpStartConfig (explicit role_arn, no live role/instance + resolution).""" + js_config = JumpStartConfig(model_id="test-model", model_version="1.0.0") + return ModelBuilder.from_jumpstart_config( + jumpstart_config=js_config, + role_arn="arn:aws:iam::123456789012:role/SageMakerRole", + ) + + def test_get_deployment_configs_ensures_metadata(self): + """_get_deployment_configs lazily loads metadata configs itself, so every + caller benefits. Pre-fix it read _metadata_configs (None pre-deploy) and + returned [] without ever loading them.""" + mb = self._jumpstart_mb() + mb._metadata_configs = None + + with patch.object(mb, "_ensure_metadata_configs") as ensure: + result = mb._get_deployment_configs(None, None) + + ensure.assert_called_once() + self.assertEqual(result, []) + + def test_list_deployment_configs_loads_metadata_when_pre_deploy(self): + """list_deployment_configs() with no instance_type returns configs for a + pre-deploy JumpStart model instead of [].""" + mb = self._jumpstart_mb() + mb.config_name = None + mb.instance_type = None + + with patch.object(mb, "_is_jumpstart_model_id", return_value=True), patch.object( + mb, "_is_model_customization", return_value=False + ), patch.object(mb, "_use_jumpstart_equivalent", return_value=False), patch.object( + mb, "_get_deployment_configs", return_value=[Mock()] + ), patch.object( + mb, "deployment_config_response_data", return_value=[{"DeploymentConfigName": "c1"}] + ): + result = mb.list_deployment_configs() + + self.assertEqual(result, [{"DeploymentConfigName": "c1"}]) + + def test_benchmark_metrics_property_returns_dataframe(self): + """The benchmark_metrics property builds a DataFrame from the config + benchmark data (it did not exist before, causing AttributeError).""" + mb = self._jumpstart_mb() + sample = { + "Instance Type": ["ml.g5.2xlarge", "ml.g5.12xlarge"], + "Latency (ms)": [100.0, 80.0], + } + + with patch.object(mb, "_get_deployment_configs_benchmarks_data", return_value=sample): + df = mb.benchmark_metrics + + self.assertEqual(list(df["Instance Type"]), ["ml.g5.2xlarge", "ml.g5.12xlarge"]) + + def test_display_benchmark_metrics_no_attribute_error(self): + """display_benchmark_metrics() reads the benchmark_metrics property and + no longer raises AttributeError for a JumpStart model. The property is + patched to a stand-in frame so the assertion is on the wiring, not on + pandas' optional markdown renderer.""" + mb = self._jumpstart_mb() + df = MagicMock() + df.to_markdown.return_value = "table" + + with patch.object(mb, "_is_jumpstart_model_id", return_value=True), patch.object( + mb, "_use_jumpstart_equivalent", return_value=False + ), patch.object( + type(mb), "benchmark_metrics", new_callable=PropertyMock, return_value=df + ): + mb.display_benchmark_metrics() + + df.to_markdown.assert_called_once() + + class TestTransformer(unittest.TestCase): """Test transformer method.""" diff --git a/sagemaker-serve/tests/unit/utils/test_hf_utils.py b/sagemaker-serve/tests/unit/utils/test_hf_utils.py index df86175382..8577f085e5 100644 --- a/sagemaker-serve/tests/unit/utils/test_hf_utils.py +++ b/sagemaker-serve/tests/unit/utils/test_hf_utils.py @@ -1,10 +1,17 @@ """Unit tests for sagemaker.serve.utils.hf_utils module.""" import unittest +import os +import shutil +import sys +import tempfile from unittest.mock import Mock, patch, mock_open import json from urllib.error import HTTPError, URLError from json import JSONDecodeError -from sagemaker.serve.utils.hf_utils import _get_model_config_properties_from_hf +from sagemaker.serve.utils.hf_utils import ( + _get_model_config_properties_from_hf, + download_huggingface_model, +) class TestGetModelConfigPropertiesFromHf(unittest.TestCase): @@ -216,5 +223,95 @@ def _urlopen_side_effect(request): self.assertEqual(result, adapter_config) +class TestDownloadHuggingfaceModel(unittest.TestCase): + """Test cases for the public download_huggingface_model helper.""" + + def setUp(self): + # download_huggingface_model imports huggingface_hub lazily; provide a + # stub module so the tests do not require the real package installed. + self._hf_stub = Mock() + self._patcher = patch.dict(sys.modules, {"huggingface_hub": self._hf_stub}) + self._patcher.start() + self.addCleanup(self._patcher.stop) + # Keep the suite side-effect free: never create real directories. + makedirs_patcher = patch("sagemaker.serve.utils.hf_utils.os.makedirs") + makedirs_patcher.start() + self.addCleanup(makedirs_patcher.stop) + self.local_dir = tempfile.mkdtemp(prefix="hf-test-") + self.addCleanup(shutil.rmtree, self.local_dir, ignore_errors=True) + + def test_requires_local_dir_or_s3_uri(self): + """Neither destination given is a caller error, before any download.""" + with self.assertRaises(ValueError) as context: + download_huggingface_model("gpt2") + self.assertIn("local_dir, s3_uri", str(context.exception)) + self._hf_stub.snapshot_download.assert_not_called() + + def test_downloads_to_local_dir_and_returns_it(self): + """With only local_dir, it downloads there and returns the path.""" + result = download_huggingface_model("gpt2", local_dir=self.local_dir) + self.assertEqual(result, self.local_dir) + self.assertEqual( + self._hf_stub.snapshot_download.call_args.kwargs["local_dir"], self.local_dir + ) + + def test_uploads_to_s3_and_returns_uri(self): + """With s3_uri, it uploads the snapshot and returns the S3 URI.""" + with patch("sagemaker.core.s3.S3Uploader") as mock_uploader: + mock_uploader.upload.return_value = "s3://bucket/prefix/gpt2" + result = download_huggingface_model( + "gpt2", local_dir=self.local_dir, s3_uri="s3://bucket/prefix" + ) + self.assertEqual(result, "s3://bucket/prefix/gpt2") + mock_uploader.upload.assert_called_once() + self.assertEqual( + mock_uploader.upload.call_args.kwargs["desired_s3_uri"], "s3://bucket/prefix" + ) + + def test_s3_only_uses_temp_dir_and_cleans_up(self): + """With s3_uri and no local_dir, the snapshot stages in a temp dir that + is removed after the upload (no snapshot left on the local volume).""" + with patch("sagemaker.core.s3.S3Uploader") as mock_uploader: + mock_uploader.upload.return_value = "s3://bucket/prefix/gpt2" + result = download_huggingface_model("gpt2", s3_uri="s3://bucket/prefix") + self.assertEqual(result, "s3://bucket/prefix/gpt2") + staging_dir = self._hf_stub.snapshot_download.call_args.kwargs["local_dir"] + self.assertFalse(os.path.exists(staging_dir)) + + def test_forwards_hf_token_and_snapshot_passthroughs(self): + """token / revision / patterns are passed through to snapshot_download.""" + download_huggingface_model( + "gpt2", + local_dir=self.local_dir, + hf_hub_token="hf_tok", + revision="v1.0", + allow_patterns="*.safetensors", + ignore_patterns="*.bin", + ) + kwargs = self._hf_stub.snapshot_download.call_args.kwargs + self.assertEqual(kwargs["token"], "hf_tok") + self.assertEqual(kwargs["revision"], "v1.0") + self.assertEqual(kwargs["allow_patterns"], "*.safetensors") + self.assertEqual(kwargs["ignore_patterns"], "*.bin") + + def test_missing_huggingface_hub_raises_import_error(self): + """The helper's own ImportError (with install guidance) is raised when + huggingface_hub cannot be imported.""" + import builtins + + real_import = builtins.__import__ + + def _raise_for_hf(name, *args, **kwargs): + if name == "huggingface_hub": + raise ImportError("No module named 'huggingface_hub'") + return real_import(name, *args, **kwargs) + + with patch("builtins.__import__", side_effect=_raise_for_hf), self.assertRaises( + ImportError + ) as context: + download_huggingface_model("gpt2", local_dir=self.local_dir) + self.assertIn("pip install huggingface_hub", str(context.exception)) + + if __name__ == "__main__": unittest.main()