feat: Model class to support AdditionalModelDataSources (aws#1469)

* Add support for AdditionalModelDataSources

* Resolve PR comments

* Resolve PR comments

* Resolve PR comments

* fix unit tests

* Resolve PR comments

---------

Co-authored-by: Jonathan Makunga <[email protected]>

Author: makungaj1

src/sagemaker/jumpstart/model.py

@@ -369,6 +369,14 @@ def _validate_model_id_and_type():
             model_type=self.model_type,
         )
 
+        self.additional_model_data_sources = (
+            self._metadata_configs.get(self.config_name).resolved_config.get(
+                "hosting_additional_data_sources"
+            )
+            if self._metadata_configs.get(self.config_name)
+            else None
+        )
+
     def log_subscription_warning(self) -> None:
         """Log message prompting the customer to subscribe to the proprietary model."""
         subscription_link = verify_model_region_and_return_specs(

src/sagemaker/model.py

@@ -160,6 +160,7 @@ def __init__(
         dependencies: Optional[List[str]] = None,
         git_config: Optional[Dict[str, str]] = None,
         resources: Optional[ResourceRequirements] = None,
+        additional_model_data_sources: Optional[Dict[str, Any]] = None,
     ):
         """Initialize an SageMaker ``Model``.
 
@@ -323,9 +324,12 @@ def __init__(
                 for a model to be deployed to an endpoint. Only
                 EndpointType.INFERENCE_COMPONENT_BASED supports this feature.
                 (Default: None).
+            additional_model_data_sources (Optional[Dict[str, Any]]): Additional location
+                of SageMaker model data (default: None).
 
         """
         self.model_data = model_data
+        self.additional_model_data_sources = additional_model_data_sources
         self.image_uri = image_uri
         self.predictor_cls = predictor_cls
         self.name = name
@@ -671,6 +675,7 @@ def prepare_container_def(
             accept_eula=(
                 accept_eula if accept_eula is not None else getattr(self, "accept_eula", None)
             ),
+            additional_model_data_sources=self.additional_model_data_sources,
         )
 
     def is_repack(self) -> bool:

src/sagemaker/session.py

@@ -7137,6 +7137,7 @@ def container_def(
     container_mode=None,
     image_config=None,
     accept_eula=None,
+    additional_model_data_sources=None,
 ):
     """Create a definition for executing a container as part of a SageMaker model.
 
@@ -7159,6 +7160,8 @@ def container_def(
             The `accept_eula` value must be explicitly defined as `True` in order to
             accept the end-user license agreement (EULA) that some
             models require. (Default: None).
+        additional_model_data_sources (PipelineVariable or dict): Additional location
+                of SageMaker model data (default: None).
 
     Returns:
         dict[str, str]: A complete container definition object usable with the CreateModel API if
@@ -7168,6 +7171,9 @@ def container_def(
         env = {}
     c_def = {"Image": image_uri, "Environment": env}
 
+    if additional_model_data_sources:
+        c_def["AdditionalModelDataSources"] = additional_model_data_sources
+
     if isinstance(model_data_url, str) and (
         not (model_data_url.startswith("s3://") and model_data_url.endswith("tar.gz"))
         or accept_eula is None

tests/unit/sagemaker/jumpstart/estimator/test_estimator.py

@@ -1168,6 +1168,7 @@ def test_jumpstart_estimator_kwargs_match_parent_class(self):
             "inference_config_name"
         } == model_class_init_args - {
             "model_data",
+            "additional_model_data_sources",
             "self",
             "name",
             "resources",

tests/unit/sagemaker/jumpstart/model/test_model.py

@@ -759,7 +759,7 @@ def test_jumpstart_model_kwargs_match_parent_class(self):
         Please add the new argument to the skip set below,
         and reach out to JumpStart team."""
 
-        init_args_to_skip: Set[str] = set([])
+        init_args_to_skip: Set[str] = set(["additional_model_data_sources"])
         deploy_args_to_skip: Set[str] = set(["kwargs"])
 
         parent_class_init = Model.__init__