Use correct bucket for SM/JS draft models and minor formatting/validation updates.

Author: Joseph Zhang

src/sagemaker/jumpstart/factory/model.py

@@ -54,11 +54,11 @@
     add_hub_content_arn_tags,
     add_jumpstart_model_info_tags,
     get_default_jumpstart_session_with_user_agent_suffix,
-    get_neo_content_bucket,
     get_top_ranked_config_name,
     update_dict_if_key_not_present,
     resolve_model_sagemaker_config_field,
     verify_model_region_and_return_specs,
+    get_draft_model_content_bucket,
 )
 
 from sagemaker.jumpstart.factory.utils import (
@@ -76,7 +76,6 @@
     name_from_base,
     format_tags,
     Tags,
-    get_domain_for_region,
 )
 from sagemaker.workflow.entities import PipelineVariable
 from sagemaker.compute_resource_requirements.resource_requirements import ResourceRequirements
@@ -572,7 +571,9 @@ def _add_additional_model_data_sources_to_kwargs(
     # 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(region=kwargs.region))
+        data_source.s3_data_source.set_bucket(
+            get_draft_model_content_bucket(provider=data_source.provider, region=kwargs.region)
+        )
     api_shape_additional_model_data_sources = (
         [
             camel_case_to_pascal_case(data_source.to_json())

src/sagemaker/jumpstart/model.py

@@ -458,9 +458,7 @@ def retrieve_example_payload(self) -> JumpStartSerializablePayload:
             sagemaker_session=self.sagemaker_session,
         )
 
-    def set_deployment_config(
-        self, config_name: str, instance_type: str
-    ) -> None:
+    def set_deployment_config(self, config_name: str, instance_type: str) -> None:
         """Sets the deployment config to apply to the model.
 
         Args:
@@ -479,7 +477,7 @@ def set_deployment_config(
             instance_type=instance_type,
             config_name=config_name,
             sagemaker_session=self.sagemaker_session,
-            role=self.role
+            role=self.role,
         )
 
     @property
@@ -766,11 +764,11 @@ def deploy(
                 (Default: EndpointType.MODEL_BASED).
             routing_config (Optional[Dict]): Settings the control how the endpoint routes
                 incoming traffic to the instances that the endpoint hosts.
-            model_access_configs (Optional[Dict[str, ModelAccessConfig]]): For models that require ModelAccessConfig,
-                provide a `{ "model_id", ModelAccessConfig(accept_eula=True) }` to indicate whether model terms
-                of use have been accepted. The `accept_eula` value must be explicitly defined as `True` in order to
-                accept the end-user license agreement (EULA) that some.
-                (Default: None)
+            model_access_configs (Optional[Dict[str, ModelAccessConfig]]): For models that require
+                ModelAccessConfig, provide a `{ "model_id", ModelAccessConfig(accept_eula=True) }`
+                to indicate whether model terms of use have been accepted. The `accept_eula` value
+                must be explicitly defined as `True` in order to accept the end-user license
+                agreement (EULA) that some. (Default: None)
 
         Raises:
             MarketplaceModelSubscriptionError: If the caller is not subscribed to the model.

src/sagemaker/jumpstart/types.py

@@ -1082,7 +1082,7 @@ def set_bucket(self, bucket: str) -> None:
 class AdditionalModelDataSource(JumpStartDataHolderType):
     """Data class of additional model data source mirrors CreateModel API."""
 
-    SERIALIZATION_EXCLUSION_SET: Set[str] = set()
+    SERIALIZATION_EXCLUSION_SET = {"provider"}
 
     __slots__ = ["channel_name", "s3_data_source", "hosting_eula_key"]
 
@@ -1103,6 +1103,7 @@ def from_json(self, json_obj: Dict[str, Any]) -> None:
         self.channel_name: str = json_obj["channel_name"]
         self.s3_data_source: S3DataSource = S3DataSource(json_obj["s3_data_source"])
         self.hosting_eula_key: str = json_obj.get("hosting_eula_key")
+        self.provider: Dict = json_obj.get("provider", {})
 
     def to_json(self, exclude_keys=True) -> Dict[str, Any]:
         """Returns json representation of AdditionalModelDataSource object."""
@@ -1121,7 +1122,9 @@ def to_json(self, exclude_keys=True) -> Dict[str, Any]:
 class JumpStartModelDataSource(AdditionalModelDataSource):
     """Data class JumpStart additional model data source."""
 
-    SERIALIZATION_EXCLUSION_SET = {"artifact_version"}
+    SERIALIZATION_EXCLUSION_SET = {
+        "artifact_version"
+    } | AdditionalModelDataSource.SERIALIZATION_EXCLUSION_SET
 
     __slots__ = list(SERIALIZATION_EXCLUSION_SET) + AdditionalModelDataSource.__slots__
 
@@ -2241,7 +2244,7 @@ class JumpStartModelDeployKwargs(JumpStartKwargs):
         "config_name",
         "routing_config",
         "specs",
-        "model_access_configs"
+        "model_access_configs",
     ]
 
     SERIALIZATION_EXCLUSION_SET = {
@@ -2255,7 +2258,7 @@ class JumpStartModelDeployKwargs(JumpStartKwargs):
         "sagemaker_session",
         "training_instance_type",
         "config_name",
-        "model_access_configs"
+        "model_access_configs",
     }
 
     def __init__(
@@ -2294,7 +2297,7 @@ def __init__(
         endpoint_type: Optional[EndpointType] = None,
         config_name: Optional[str] = None,
         routing_config: Optional[Dict[str, Any]] = None,
-        model_access_configs: Optional[Dict[str, CoreModelAccessConfig]] = None
+        model_access_configs: Optional[Dict[str, CoreModelAccessConfig]] = None,
     ) -> None:
         """Instantiates JumpStartModelDeployKwargs object."""
 

src/sagemaker/jumpstart/utils.py

@@ -564,6 +564,7 @@ def get_eula_message(model_specs: JumpStartModelSpecs, region: str) -> str:
 
 
 def format_eula_message_from_specs(model_id: str, region: str, hosting_eula_key: str):
+    """Returns a formatted EULA message."""
     return (
         f"Model '{model_id}' requires accepting end-user license agreement (EULA). "
         f"See https://{get_jumpstart_content_bucket(region=region)}.s3.{region}."
@@ -1552,21 +1553,25 @@ def _add_model_access_configs_to_model_data_sources(
         hosting_eula_key = model_data_source.get("HostingEulaKey")
         if hosting_eula_key:
             if not model_access_configs or not model_access_configs.get(model_id):
-                eula_message_template = "{model_source}{base_eula_message}{model_access_configs_message}"
+                eula_message_template = (
+                    "{model_source}{base_eula_message}{model_access_configs_message}"
+                )
                 model_access_config_entry = (
-                    "\"{model_id}\":ModelAccessConfig(accept_eula=True)".format(model_id=model_id)
+                    '"{model_id}":ModelAccessConfig(accept_eula=True)'.format(model_id=model_id)
                 )
-                raise ValueError(eula_message_template.format(
-                    model_source="Draft " if model_data_source.get("ChannelName") else "",
-                    base_eula_message=format_eula_message_from_specs(
-                        model_id=model_id, region=region, hosting_eula_key=hosting_eula_key
-                    ),
-                    model_access_configs_message=(
-                        " Please add a ModelAccessConfig entry:"
-                        f" {model_access_config_entry} "
-                        "to model_access_configs to acknowledge the EULA."
+                raise ValueError(
+                    eula_message_template.format(
+                        model_source="Draft " if model_data_source.get("ChannelName") else "",
+                        base_eula_message=format_eula_message_from_specs(
+                            model_id=model_id, region=region, hosting_eula_key=hosting_eula_key
+                        ),
+                        model_access_configs_message=(
+                            " Please add a ModelAccessConfig entry:"
+                            f" {model_access_config_entry} "
+                            "to model_access_configs to acknowledge the EULA."
+                        ),
                     )
-                ))
+                )
             acked_model_data_source = model_data_source.copy()
             acked_model_data_source.pop("HostingEulaKey")
             acked_model_data_source["S3DataSource"]["ModelAccessConfig"] = (
@@ -1576,3 +1581,17 @@ def _add_model_access_configs_to_model_data_sources(
         else:
             acked_model_data_sources.append(model_data_source)
     return acked_model_data_sources
+
+
+def get_draft_model_content_bucket(provider: Dict, region: str) -> str:
+    """Returns the correct content bucket for a 1p draft model."""
+    neo_bucket = get_neo_content_bucket(region=region)
+    if not provider:
+        return neo_bucket
+    provider_name = provider.get("name", "")
+    if provider_name == "JumpStart":
+        classification = provider.get("classification", "ungated")
+        if classification == "gated":
+            return get_jumpstart_gated_content_bucket(region=region)
+        return get_jumpstart_content_bucket(region=region)
+    return neo_bucket

src/sagemaker/serve/builder/jumpstart_builder.py

@@ -504,9 +504,7 @@ def tune_for_tgi_jumpstart(self, max_tuning_duration: int = 1800):
             sharded_supported=sharded_supported, max_tuning_duration=max_tuning_duration
         )
 
-    def set_deployment_config(
-        self, config_name: str, instance_type: str
-    ) -> None:
+    def set_deployment_config(self, config_name: str, instance_type: str) -> None:
         """Sets the deployment config to apply to the model.
 
         Args:

src/sagemaker/serve/builder/model_builder.py

@@ -1344,7 +1344,7 @@ def _optimize_for_hf(
             Optional[Dict[str, Any]]: Model optimization job input arguments.
         """
         if speculative_decoding_config:
-            if speculative_decoding_config.get("ModelProvider", "") == "JumpStart":
+            if speculative_decoding_config.get("ModelProvider", "").lower() == "jumpstart":
                 _jumpstart_speculative_decoding(
                     model=self.pysdk_model,
                     speculative_decoding_config=speculative_decoding_config,