Merge branch 'master' into neo-input-shape-derivation-comment

Author: HappyAmazonian

src/sagemaker/estimator.py

@@ -3082,6 +3082,7 @@ def __init__(
             hyperparameters=hyperparameters,
             instance_groups=instance_groups,
             training_repository_access_mode=training_repository_access_mode,
+            enable_infra_check=enable_infra_check,
             training_repository_credentials_provider_arn=training_repository_credentials_provider_arn,  # noqa: E501 # pylint: disable=line-too-long
             container_entry_point=container_entry_point,
             container_arguments=container_arguments,

src/sagemaker/jumpstart/estimator.py

@@ -102,6 +102,7 @@ def __init__(
         instance_groups: Optional[List[InstanceGroup]] = None,
         training_repository_access_mode: Optional[Union[str, PipelineVariable]] = None,
         training_repository_credentials_provider_arn: Optional[Union[str, PipelineVariable]] = None,
+        enable_infra_check: Optional[Union[bool, PipelineVariable]] = None,
         container_entry_point: Optional[List[str]] = None,
         container_arguments: Optional[List[str]] = None,
         disable_output_compression: Optional[bool] = None,
@@ -485,6 +486,8 @@ def __init__(
                 private Docker registry where your training image is hosted (Default: None).
                 When it's set to None, SageMaker will not do authentication before pulling the image
                 in the private Docker registry. (Default: None).
+            enable_infra_check (Optional[Union[bool, PipelineVariable]]):
+                Specifies whether it is running Sagemaker built-in infra check jobs.
             container_entry_point (Optional[List[str]]): The entrypoint script for a Docker
                 container used to run a training job. This script takes precedence over
                 the default train processing instructions.
@@ -565,6 +568,7 @@ def _is_valid_model_id_hook():
             container_entry_point=container_entry_point,
             container_arguments=container_arguments,
             disable_output_compression=disable_output_compression,
+            enable_infra_check=enable_infra_check,
         )
 
         self.model_id = estimator_init_kwargs.model_id

src/sagemaker/jumpstart/factory/estimator.py

@@ -123,6 +123,7 @@ def get_init_kwargs(
     container_entry_point: Optional[List[str]] = None,
     container_arguments: Optional[List[str]] = None,
     disable_output_compression: Optional[bool] = None,
+    enable_infra_check: Optional[Union[bool, PipelineVariable]] = None,
 ) -> JumpStartEstimatorInitKwargs:
     """Returns kwargs required to instantiate `sagemaker.estimator.Estimator` object."""
 
@@ -178,6 +179,7 @@ def get_init_kwargs(
         container_entry_point=container_entry_point,
         container_arguments=container_arguments,
         disable_output_compression=disable_output_compression,
+        enable_infra_check=enable_infra_check,
     )
 
     estimator_init_kwargs = _add_model_version_to_kwargs(estimator_init_kwargs)

src/sagemaker/jumpstart/types.py

@@ -18,6 +18,7 @@
 from sagemaker.utils import get_instance_type_family
 
 from sagemaker.session import Session
+from sagemaker.workflow.entities import PipelineVariable
 
 
 class JumpStartDataHolderType:
@@ -1011,6 +1012,7 @@ class JumpStartEstimatorInitKwargs(JumpStartKwargs):
         "container_entry_point",
         "container_arguments",
         "disable_output_compression",
+        "enable_infra_check",
     ]
 
     SERIALIZATION_EXCLUSION_SET = {
@@ -1074,6 +1076,7 @@ def __init__(
         container_entry_point: Optional[List[str]] = None,
         container_arguments: Optional[List[str]] = None,
         disable_output_compression: Optional[bool] = None,
+        enable_infra_check: Optional[Union[bool, PipelineVariable]] = None,
     ) -> None:
         """Instantiates JumpStartEstimatorInitKwargs object."""
 
@@ -1130,6 +1133,7 @@ def __init__(
         self.container_entry_point = container_entry_point
         self.container_arguments = container_arguments
         self.disable_output_compression = disable_output_compression
+        self.enable_infra_check = enable_infra_check
 
 
 class JumpStartEstimatorFitKwargs(JumpStartKwargs):

src/sagemaker/jumpstart/utils.py

@@ -14,7 +14,7 @@
 from __future__ import absolute_import
 import logging
 import os
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, Union
 from urllib.parse import urlparse
 from packaging.version import Version
 import sagemaker
@@ -277,7 +277,7 @@ def get_jumpstart_base_name_if_jumpstart_model(
 
 def add_jumpstart_tags(
     tags: Optional[List[Dict[str, str]]] = None,
-    inference_model_uri: Optional[str] = None,
+    inference_model_uri: Optional[Union[str, dict]] = None,
     inference_script_uri: Optional[str] = None,
     training_model_uri: Optional[str] = None,
     training_script_uri: Optional[str] = None,
@@ -289,7 +289,7 @@ def add_jumpstart_tags(
     Args:
         tags (Optional[List[Dict[str,str]]): Current tags for JumpStart inference
             or training job. (Default: None).
-        inference_model_uri (Optional[str]): S3 URI for inference model artifact.
+        inference_model_uri (Optional[Union[dict, str]]): S3 URI for inference model artifact.
             (Default: None).
         inference_script_uri (Optional[str]): S3 URI for inference script tarball.
             (Default: None).
@@ -302,6 +302,10 @@ def add_jumpstart_tags(
         "The URI (%s) is a pipeline variable which is only interpreted at execution time. "
         "As a result, the JumpStart resources will not be tagged."
     )
+
+    if isinstance(inference_model_uri, dict):
+        inference_model_uri = inference_model_uri.get("S3DataSource", {}).get("S3Uri", None)
+
     if inference_model_uri:
         if is_pipeline_variable(inference_model_uri):
             logging.warning(warn_msg, "inference_model_uri")

src/sagemaker/model.py

@@ -1347,7 +1347,9 @@ def deploy(
 
         tags = add_jumpstart_tags(
             tags=tags,
-            inference_model_uri=self.model_data if isinstance(self.model_data, str) else None,
+            inference_model_uri=self.model_data
+            if isinstance(self.model_data, (str, dict))
+            else None,
             inference_script_uri=self.source_dir,
         )
 

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

@@ -677,7 +677,7 @@ def test_jumpstart_estimator_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(["kwargs", "enable_infra_check"])
+        init_args_to_skip: Set[str] = set(["kwargs"])
         fit_args_to_skip: Set[str] = set()
         deploy_args_to_skip: Set[str] = set(["kwargs"])
 

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

@@ -331,7 +331,7 @@ def test_jumpstart_model_kwargs_match_parent_class(self):
 
         """If you add arguments to <Model constructor>, this test will fail.
         Please add the new argument to the skip set below,
-        and cut a ticket sev-3 to JumpStart team: AWS > SageMaker > JumpStart"""
+        and reach out to JumpStart team."""
 
         init_args_to_skip: Set[str] = set()
         deploy_args_to_skip: Set[str] = set(["kwargs"])

tests/unit/sagemaker/jumpstart/test_utils.py

@@ -216,6 +216,34 @@ def test_add_jumpstart_tags_inference():
         inference_script_uri=inference_script_uri,
     ) == [{"Key": JumpStartTag.INFERENCE_MODEL_URI.value, "Value": inference_model_uri}]
 
+    tags = []
+    inference_model_uri = {"S3DataSource": {"S3Uri": random_jumpstart_s3_uri("random_key")}}
+    inference_script_uri = "dfsdfs"
+    assert utils.add_jumpstart_tags(
+        tags=tags,
+        inference_model_uri=inference_model_uri,
+        inference_script_uri=inference_script_uri,
+    ) == [
+        {
+            "Key": JumpStartTag.INFERENCE_MODEL_URI.value,
+            "Value": inference_model_uri["S3DataSource"]["S3Uri"],
+        }
+    ]
+
+    tags = []
+    inference_model_uri = {"S3DataSource": {"S3Uri": random_jumpstart_s3_uri("random_key/prefix/")}}
+    inference_script_uri = "dfsdfs"
+    assert utils.add_jumpstart_tags(
+        tags=tags,
+        inference_model_uri=inference_model_uri,
+        inference_script_uri=inference_script_uri,
+    ) == [
+        {
+            "Key": JumpStartTag.INFERENCE_MODEL_URI.value,
+            "Value": inference_model_uri["S3DataSource"]["S3Uri"],
+        }
+    ]
+
     tags = [{"Key": "some", "Value": "tag"}]
     inference_model_uri = random_jumpstart_s3_uri("random_key")
     inference_script_uri = "dfsdfs"