feat: AMI support for BRM (#1589)

* feat: AMI support for BRM

Author: chrstfu

src/sagemaker/jumpstart/factory/model.py

@@ -662,6 +662,7 @@ def get_deploy_kwargs(
     config_name: Optional[str] = None,
     routing_config: Optional[Dict[str, Any]] = None,
     model_access_configs: Optional[Dict[str, ModelAccessConfig]] = None,
+    inference_ami_version: Optional[str] = None,
 ) -> JumpStartModelDeployKwargs:
     """Returns kwargs required to call `deploy` on `sagemaker.estimator.Model` object."""
 
@@ -699,6 +700,7 @@ def get_deploy_kwargs(
         config_name=config_name,
         routing_config=routing_config,
         model_access_configs=model_access_configs,
+        inference_ami_version=inference_ami_version,
     )
     deploy_kwargs, orig_session = _set_temp_sagemaker_session_if_not_set(kwargs=deploy_kwargs)
     deploy_kwargs.specs = verify_model_region_and_return_specs(

src/sagemaker/jumpstart/hub/interfaces.py

@@ -471,6 +471,7 @@ class HubModelDocument(HubDataHolderType):
         "hosting_use_script_uri",
         "hosting_eula_uri",
         "hosting_model_package_arn",
+        "inference_ami_version",
         "model_subscription_link",
         "inference_configs",
         "inference_config_components",
@@ -593,6 +594,8 @@ def from_json(self, json_obj: Dict[str, Any]) -> None:
         self.hosting_eula_uri: Optional[str] = json_obj.get("HostingEulaUri")
         self.hosting_model_package_arn: Optional[str] = json_obj.get("HostingModelPackageArn")
 
+        self.inference_ami_version: Optional[str] = json_obj.get("InferenceAmiVersion")
+
         self.model_subscription_link: Optional[str] = json_obj.get("ModelSubscriptionLink")
 
         self.inference_config_rankings = self._get_config_rankings(json_obj)

src/sagemaker/jumpstart/hub/parsers.py

@@ -72,7 +72,11 @@ def get_model_spec_arg_keys(
     """
     arg_keys: List[str] = []
     if arg_type == ModelSpecKwargType.DEPLOY:
-        arg_keys = ["ModelDataDownloadTimeout", "ContainerStartupHealthCheckTimeout"]
+        arg_keys = [
+            "ModelDataDownloadTimeout",
+            "ContainerStartupHealthCheckTimeout",
+            "InferenceAmiVersion",
+        ]
     elif arg_type == ModelSpecKwargType.ESTIMATOR:
         arg_keys = [
             "EncryptInterContainerTraffic",

src/sagemaker/jumpstart/model.py

@@ -666,6 +666,7 @@ def deploy(
         endpoint_type: EndpointType = EndpointType.MODEL_BASED,
         routing_config: Optional[Dict[str, Any]] = None,
         model_access_configs: Optional[Dict[str, ModelAccessConfig]] = None,
+        inference_ami_version: Optional[str] = None,
     ) -> PredictorBase:
         """Creates endpoint by calling base ``Model`` class `deploy` method.
 
@@ -808,6 +809,7 @@ def deploy(
             config_name=self.config_name,
             routing_config=routing_config,
             model_access_configs=model_access_configs,
+            inference_ami_version=inference_ami_version,
         )
         if (
             self.model_type == JumpStartModelType.PROPRIETARY

src/sagemaker/jumpstart/types.py

@@ -1389,6 +1389,7 @@ def from_json(self, json_obj: Dict[str, Any]) -> None:
         self.hosting_model_package_arns: Optional[Dict] = (
             model_package_arns if model_package_arns is not None else {}
         )
+
         self.hosting_use_script_uri: bool = json_obj.get("hosting_use_script_uri", True)
 
         self.hosting_instance_type_variants: Optional[JumpStartInstanceTypeVariants] = (
@@ -2245,6 +2246,7 @@ class JumpStartModelDeployKwargs(JumpStartKwargs):
         "routing_config",
         "specs",
         "model_access_configs",
+        "inference_ami_version",
     ]
 
     SERIALIZATION_EXCLUSION_SET = {
@@ -2298,6 +2300,7 @@ def __init__(
         config_name: Optional[str] = None,
         routing_config: Optional[Dict[str, Any]] = None,
         model_access_configs: Optional[Dict[str, CoreModelAccessConfig]] = None,
+        inference_ami_version: Optional[str] = None,
     ) -> None:
         """Instantiates JumpStartModelDeployKwargs object."""
 
@@ -2336,6 +2339,7 @@ def __init__(
         self.config_name = config_name
         self.routing_config = routing_config
         self.model_access_configs = model_access_configs
+        self.inference_ami_version = inference_ami_version
 
 
 class JumpStartEstimatorInitKwargs(JumpStartKwargs):

src/sagemaker/model.py

@@ -1383,6 +1383,7 @@ def deploy(
         inference_component_name=None,
         routing_config: Optional[Dict[str, Any]] = None,
         model_reference_arn: Optional[str] = None,
+        inference_ami_version: Optional[str] = None,
         **kwargs,
     ):
         """Deploy this ``Model`` to an ``Endpoint`` and optionally return a ``Predictor``.
@@ -1652,6 +1653,7 @@ def deploy(
                     container_startup_health_check_timeout=container_startup_health_check_timeout,
                     managed_instance_scaling=managed_instance_scaling_config,
                     routing_config=routing_config,
+                    inference_ami_version=inference_ami_version,
                 )
 
                 self.sagemaker_session.endpoint_from_production_variants(

src/sagemaker/session.py

@@ -7735,6 +7735,7 @@ def production_variant(
     container_startup_health_check_timeout=None,
     managed_instance_scaling=None,
     routing_config=None,
+    inference_ami_version=None,
 ):
     """Create a production variant description suitable for use in a ``ProductionVariant`` list.
 
@@ -7799,6 +7800,9 @@ def production_variant(
             RoutingConfig=routing_config,
         )
 
+    if inference_ami_version:
+        production_variant_configuration["InferenceAmiVersion"] = inference_ami_version
+
     return production_variant_configuration