feat: use Neo bucket in speculative decoding data source (aws#1479)

Captainia · web-flow · commit c6581ff8efc6 · 2024-06-14T12:25:03.000-07:00
* Use Neo bucket in speculative decoding data source

* address comments

* format

* address comments

* add buckets to regional config

* remove opt-in regions for neo buckets
diff --git a/src/sagemaker/jumpstart/constants.py b/src/sagemaker/jumpstart/constants.py
@@ -44,31 +44,37 @@
             region_name="us-west-2",
             content_bucket="jumpstart-cache-prod-us-west-2",
             gated_content_bucket="jumpstart-private-cache-prod-us-west-2",
+            neo_content_bucket="sagemaker-sd-models-prod-us-west-2",
         ),
         JumpStartLaunchedRegionInfo(
             region_name="us-east-1",
             content_bucket="jumpstart-cache-prod-us-east-1",
             gated_content_bucket="jumpstart-private-cache-prod-us-east-1",
+            neo_content_bucket="sagemaker-sd-models-prod-us-east-1",
         ),
         JumpStartLaunchedRegionInfo(
             region_name="us-east-2",
             content_bucket="jumpstart-cache-prod-us-east-2",
             gated_content_bucket="jumpstart-private-cache-prod-us-east-2",
+            neo_content_bucket="sagemaker-sd-models-prod-us-east-2",
         ),
         JumpStartLaunchedRegionInfo(
             region_name="eu-west-1",
             content_bucket="jumpstart-cache-prod-eu-west-1",
             gated_content_bucket="jumpstart-private-cache-prod-eu-west-1",
+            neo_content_bucket="sagemaker-sd-models-prod-eu-west-1",
         ),
         JumpStartLaunchedRegionInfo(
             region_name="eu-central-1",
             content_bucket="jumpstart-cache-prod-eu-central-1",
             gated_content_bucket="jumpstart-private-cache-prod-eu-central-1",
+            neo_content_bucket="sagemaker-sd-models-prod-eu-central-1",
         ),
         JumpStartLaunchedRegionInfo(
             region_name="eu-north-1",
             content_bucket="jumpstart-cache-prod-eu-north-1",
             gated_content_bucket="jumpstart-private-cache-prod-eu-north-1",
+            neo_content_bucket="sagemaker-sd-models-prod-eu-north-1",
         ),
         JumpStartLaunchedRegionInfo(
             region_name="me-south-1",
@@ -84,11 +90,13 @@
             region_name="ap-south-1",
             content_bucket="jumpstart-cache-prod-ap-south-1",
             gated_content_bucket="jumpstart-private-cache-prod-ap-south-1",
+            neo_content_bucket="sagemaker-sd-models-prod-ap-south-1",
         ),
         JumpStartLaunchedRegionInfo(
             region_name="eu-west-3",
             content_bucket="jumpstart-cache-prod-eu-west-3",
             gated_content_bucket="jumpstart-private-cache-prod-eu-west-3",
+            neo_content_bucket="sagemaker-sd-models-prod-eu-west-3",
         ),
         JumpStartLaunchedRegionInfo(
             region_name="af-south-1",
@@ -99,6 +107,7 @@
             region_name="sa-east-1",
             content_bucket="jumpstart-cache-prod-sa-east-1",
             gated_content_bucket="jumpstart-private-cache-prod-sa-east-1",
+            neo_content_bucket="sagemaker-sd-models-prod-sa-east-1",
         ),
         JumpStartLaunchedRegionInfo(
             region_name="ap-east-1",
@@ -109,21 +118,25 @@
             region_name="ap-northeast-2",
             content_bucket="jumpstart-cache-prod-ap-northeast-2",
             gated_content_bucket="jumpstart-private-cache-prod-ap-northeast-2",
+            neo_content_bucket="sagemaker-sd-models-prod-ap-northeast-2",
         ),
         JumpStartLaunchedRegionInfo(
             region_name="ap-northeast-3",
             content_bucket="jumpstart-cache-prod-ap-northeast-3",
             gated_content_bucket="jumpstart-private-cache-prod-ap-northeast-3",
+            neo_content_bucket="sagemaker-sd-models-prod-ap-northeast-3",
         ),
         JumpStartLaunchedRegionInfo(
             region_name="ap-southeast-3",
             content_bucket="jumpstart-cache-prod-ap-southeast-3",
             gated_content_bucket="jumpstart-private-cache-prod-ap-southeast-3",
+            neo_content_bucket="sagemaker-sd-models-prod-ap-southeast-3",
         ),
         JumpStartLaunchedRegionInfo(
             region_name="eu-west-2",
             content_bucket="jumpstart-cache-prod-eu-west-2",
             gated_content_bucket="jumpstart-private-cache-prod-eu-west-2",
+            neo_content_bucket="sagemaker-sd-models-prod-eu-west-2",
         ),
         JumpStartLaunchedRegionInfo(
             region_name="eu-south-1",
@@ -134,26 +147,31 @@
             region_name="ap-northeast-1",
             content_bucket="jumpstart-cache-prod-ap-northeast-1",
             gated_content_bucket="jumpstart-private-cache-prod-ap-northeast-1",
+            neo_content_bucket="sagemaker-sd-models-prod-ap-northeast-1",
         ),
         JumpStartLaunchedRegionInfo(
             region_name="us-west-1",
             content_bucket="jumpstart-cache-prod-us-west-1",
             gated_content_bucket="jumpstart-private-cache-prod-us-west-1",
+            neo_content_bucket="sagemaker-sd-models-prod-us-west-1",
         ),
         JumpStartLaunchedRegionInfo(
             region_name="ap-southeast-1",
             content_bucket="jumpstart-cache-prod-ap-southeast-1",
             gated_content_bucket="jumpstart-private-cache-prod-ap-southeast-1",
+            neo_content_bucket="sagemaker-sd-models-prod-ap-southeast-1",
         ),
         JumpStartLaunchedRegionInfo(
             region_name="ap-southeast-2",
             content_bucket="jumpstart-cache-prod-ap-southeast-2",
             gated_content_bucket="jumpstart-private-cache-prod-ap-southeast-2",
+            neo_content_bucket="sagemaker-sd-models-prod-ap-southeast-2",
         ),
         JumpStartLaunchedRegionInfo(
             region_name="ca-central-1",
             content_bucket="jumpstart-cache-prod-ca-central-1",
             gated_content_bucket="jumpstart-private-cache-prod-ca-central-1",
+            neo_content_bucket="sagemaker-sd-models-prod-ca-central-1",
         ),
         JumpStartLaunchedRegionInfo(
             region_name="cn-north-1",
@@ -184,6 +202,7 @@
 )
 
 JUMPSTART_DEFAULT_REGION_NAME = boto3.session.Session().region_name or "us-west-2"
+NEO_DEFAULT_REGION_NAME = boto3.session.Session().region_name or "us-west-2"
 
 JUMPSTART_DEFAULT_MANIFEST_FILE_S3_KEY = "models_manifest.json"
 JUMPSTART_DEFAULT_PROPRIETARY_MANIFEST_KEY = "proprietary-sdk-manifest.json"
@@ -201,6 +220,7 @@
     "AWS_JUMPSTART_MANIFEST_LOCAL_ROOT_DIR_OVERRIDE"
 )
 ENV_VARIABLE_JUMPSTART_SPECS_LOCAL_ROOT_DIR_OVERRIDE = "AWS_JUMPSTART_SPECS_LOCAL_ROOT_DIR_OVERRIDE"
+ENV_VARIABLE_NEO_CONTENT_BUCKET_OVERRIDE = "AWS_NEO_CONTENT_BUCKET_OVERRIDE"
 
 JUMPSTART_RESOURCE_BASE_NAME = "sagemaker-jumpstart"
 
diff --git a/src/sagemaker/jumpstart/factory/model.py b/src/sagemaker/jumpstart/factory/model.py
@@ -46,6 +46,7 @@
 from sagemaker.jumpstart.utils import (
     add_jumpstart_model_info_tags,
     get_default_jumpstart_session_with_user_agent_suffix,
+    get_neo_content_bucket,
     update_dict_if_key_not_present,
     resolve_model_sagemaker_config_field,
     verify_model_region_and_return_specs,
@@ -631,14 +632,16 @@ def _add_additional_model_data_sources_to_kwargs(
         model_type=kwargs.model_type,
         config_name=kwargs.config_name,
     )
-
-    additional_data_sources = specs.get_additional_s3_data_sources()
+    # Append speculative decoding data source from metadata
+    speculative_decoding_data_sources = specs.get_speculative_decoding_s3_data_sources()
+    for data_source in speculative_decoding_data_sources:
+        data_source.s3_data_source.set_bucket(get_neo_content_bucket())
     api_shape_additional_model_data_sources = (
         [
             camel_case_to_pascal_case(data_source.to_json())
-            for data_source in additional_data_sources
+            for data_source in speculative_decoding_data_sources
         ]
-        if specs.get_additional_s3_data_sources()
+        if specs.get_speculative_decoding_s3_data_sources()
         else None
     )
 
diff --git a/src/sagemaker/jumpstart/types.py b/src/sagemaker/jumpstart/types.py
@@ -15,7 +15,13 @@
 from copy import deepcopy
 from enum import Enum
 from typing import Any, Dict, List, Optional, Set, Union
-from sagemaker.utils import get_instance_type_family, format_tags, Tags, deep_override_dict
+from sagemaker.utils import (
+    S3_PREFIX,
+    get_instance_type_family,
+    format_tags,
+    Tags,
+    deep_override_dict,
+)
 from sagemaker.model_metrics import ModelMetrics
 from sagemaker.metadata_properties import MetadataProperties
 from sagemaker.drift_check_baselines import DriftCheckBaselines
@@ -116,10 +122,14 @@ class JumpStartS3FileType(str, Enum):
 class JumpStartLaunchedRegionInfo(JumpStartDataHolderType):
     """Data class for launched region info."""
 
-    __slots__ = ["content_bucket", "region_name", "gated_content_bucket"]
+    __slots__ = ["content_bucket", "region_name", "gated_content_bucket", "neo_content_bucket"]
 
     def __init__(
-        self, content_bucket: str, region_name: str, gated_content_bucket: Optional[str] = None
+        self,
+        content_bucket: str,
+        region_name: str,
+        gated_content_bucket: Optional[str] = None,
+        neo_content_bucket: Optional[str] = None,
     ):
         """Instantiates JumpStartLaunchedRegionInfo object.
 
@@ -128,10 +138,13 @@ def __init__(
             region_name (str): Name of JumpStart launched region.
             gated_content_bucket (Optional[str[]): Name of JumpStart gated s3 content bucket
                 optionally associated with region.
+            neo_content_bucket (Optional[str]): Name of Neo service s3 content bucket
+                optionally associated with region.
         """
         self.content_bucket = content_bucket
         self.gated_content_bucket = gated_content_bucket
         self.region_name = region_name
+        self.neo_content_bucket = neo_content_bucket
 
 
 class JumpStartModelHeader(JumpStartDataHolderType):
@@ -848,6 +861,21 @@ def to_json(self) -> Dict[str, Any]:
                     json_obj[att] = cur_val
         return json_obj
 
+    def set_bucket(self, bucket: str) -> None:
+        """Sets bucket name from S3 URI."""
+
+        if self.s3_uri.startswith(S3_PREFIX):
+            s3_path = self.s3_uri[len(S3_PREFIX) :]
+            old_bucket = s3_path.split("/")[0]
+            key = s3_path[len(old_bucket) :]
+            self.s3_uri = f"{S3_PREFIX}{bucket}{key}"  # pylint: disable=W0201
+            return
+
+        if not bucket.endswith("/"):
+            bucket += "/"
+
+        self.s3_uri = f"{S3_PREFIX}{bucket}{self.s3_uri}"  # pylint: disable=W0201
+
 
 class AdditionalModelDataSource(JumpStartDataHolderType):
     """Data class of additional model data source mirrors CreateModel API."""
@@ -1638,8 +1666,10 @@ def supports_incremental_training(self) -> bool:
         """Returns True if the model supports incremental training."""
         return self.incremental_training_supported
 
-    def get_speculative_decoding_s3_data_sources(self) -> List[JumpStartAdditionalDataSources]:
+    def get_speculative_decoding_s3_data_sources(self) -> List[JumpStartModelDataSource]:
         """Returns data sources for speculative decoding."""
+        if not self.hosting_additional_data_sources:
+            return []
         return self.hosting_additional_data_sources.speculative_decoding or []
 
     def get_additional_s3_data_sources(self) -> List[JumpStartAdditionalDataSources]:
diff --git a/src/sagemaker/jumpstart/utils.py b/src/sagemaker/jumpstart/utils.py
@@ -156,7 +156,7 @@ def get_jumpstart_content_bucket(
         except KeyError:
             formatted_launched_regions_str = get_jumpstart_launched_regions_message()
             raise ValueError(
-                f"Unable to get content bucket for JumpStart in {region} region. "
+                f"Unable to get content bucket for Neo in {region} region. "
                 f"{formatted_launched_regions_str}"
             )
 
@@ -170,6 +170,34 @@ def get_jumpstart_content_bucket(
     return bucket_to_return
 
 
+def get_neo_content_bucket(
+    region: str = constants.NEO_DEFAULT_REGION_NAME,
+) -> str:
+    """Returns the regionalized S3 bucket name for Neo service.
+
+    Raises:
+        ValueError: If Neo is not launched in ``region``.
+    """
+
+    bucket_to_return: Optional[str] = None
+    if (
+        constants.ENV_VARIABLE_NEO_CONTENT_BUCKET_OVERRIDE in os.environ
+        and len(os.environ[constants.ENV_VARIABLE_NEO_CONTENT_BUCKET_OVERRIDE]) > 0
+    ):
+        bucket_to_return = os.environ[constants.ENV_VARIABLE_NEO_CONTENT_BUCKET_OVERRIDE]
+        info_log = f"Using Neo bucket override: '{bucket_to_return}'"
+        constants.JUMPSTART_LOGGER.info(info_log)
+    else:
+        try:
+            bucket_to_return = constants.JUMPSTART_REGION_NAME_TO_LAUNCHED_REGION_DICT[
+                region
+            ].neo_content_bucket
+        except KeyError:
+            raise ValueError(f"Unable to get content bucket for Neo in {region} region.")
+
+    return bucket_to_return
+
+
 def get_formatted_manifest(
     manifest: List[Dict],
 ) -> Dict[JumpStartVersionedModelId, JumpStartModelHeader]:
diff --git a/tests/unit/sagemaker/jumpstart/constants.py b/tests/unit/sagemaker/jumpstart/constants.py
@@ -7806,6 +7806,7 @@
             },
         },
         "gpu-accelerated": {
+            "supported_inference_instance_types": ["ml.p2.xlarge", "ml.p3.2xlarge"],
             "hosting_instance_type_variants": {
                 "regional_aliases": {
                     "us-west-2": {
diff --git a/tests/unit/sagemaker/jumpstart/model/test_model.py b/tests/unit/sagemaker/jumpstart/model/test_model.py
@@ -1687,6 +1687,72 @@ def test_model_set_deployment_config(
             endpoint_logging=False,
         )
 
+    @mock.patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor._get_manifest")
+    @mock.patch(
+        "sagemaker.jumpstart.factory.model.get_default_jumpstart_session_with_user_agent_suffix"
+    )
+    @mock.patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs")
+    @mock.patch("sagemaker.jumpstart.model.Model.deploy")
+    @mock.patch("sagemaker.jumpstart.model.Model.__init__")
+    @mock.patch("sagemaker.jumpstart.factory.model.JUMPSTART_DEFAULT_REGION_NAME", region)
+    def test_model_deployment_config_additional_model_data_source(
+        self,
+        mock_model_init: mock.Mock,
+        mock_model_deploy: mock.Mock,
+        mock_get_model_specs: mock.Mock,
+        mock_session: mock.Mock,
+        mock_get_manifest: mock.Mock,
+    ):
+        mock_session.return_value = sagemaker_session
+        mock_get_model_specs.side_effect = get_prototype_spec_with_configs
+        mock_get_manifest.side_effect = (
+            lambda region, model_type, *args, **kwargs: get_prototype_manifest(region, model_type)
+        )
+        mock_model_deploy.return_value = default_predictor
+
+        model_id, _ = "pytorch-eqa-bert-base-cased", "*"
+
+        model = JumpStartModel(model_id=model_id, config_name="gpu-accelerated")
+
+        mock_model_init.assert_called_once_with(
+            image_uri="763104351884.dkr.ecr.us-west-2.amazonaws.com/"
+            "pytorch-hosting-neuronx:1.13.1-neuronx-py310-sdk2.14.1-ubuntu20.04",
+            model_data="s3://jumpstart-cache-prod-us-west-2/pytorch-infer/"
+            "infer-pytorch-eqa-bert-base-cased.tar.gz",
+            source_dir="s3://jumpstart-cache-prod-us-west-2/source-directory-tarballs/"
+            "pytorch/inference/eqa/v1.0.0/sourcedir.tar.gz",
+            entry_point="inference.py",
+            predictor_cls=Predictor,
+            role=execution_role,
+            sagemaker_session=sagemaker_session,
+            enable_network_isolation=False,
+            additional_model_data_sources=[
+                {
+                    "ChannelName": "draft_model_name",
+                    "S3DataSource": {
+                        "CompressionType": "None",
+                        "S3DataType": "S3Prefix",
+                        "S3Uri": "s3://sagemaker-sd-models-prod-us-west-2/key/to/draft/model/artifact/",
+                        "ModelAccessConfig": {"AcceptEula": False},
+                    },
+                }
+            ],
+        )
+
+        model.deploy()
+
+        mock_model_deploy.assert_called_once_with(
+            initial_instance_count=1,
+            instance_type="ml.p2.xlarge",
+            tags=[
+                {"Key": JumpStartTag.MODEL_ID, "Value": "pytorch-eqa-bert-base-cased"},
+                {"Key": JumpStartTag.MODEL_VERSION, "Value": "1.0.0"},
+                {"Key": JumpStartTag.INFERENCE_CONFIG_NAME, "Value": "gpu-accelerated"},
+            ],
+            wait=True,
+            endpoint_logging=False,
+        )
+
     @mock.patch(
         "sagemaker.jumpstart.model.get_jumpstart_configs", side_effect=lambda *args, **kwargs: {}
     )
diff --git a/tests/unit/sagemaker/jumpstart/test_types.py b/tests/unit/sagemaker/jumpstart/test_types.py
diff --git a/tests/unit/sagemaker/jumpstart/test_utils.py b/tests/unit/sagemaker/jumpstart/test_utils.py
