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()