Merge branch 'master' into master

Author: andre-marcos-perez

.gitignore

@@ -29,4 +29,6 @@ venv/
 env/
 .vscode/
 **/tmp
-.python-version
+.python-version
+**/_repack_model.py
+**/_repack_script_launcher.sh

CHANGELOG.md

@@ -1,5 +1,20 @@
 # Changelog
 
+## v2.113.0 (2022-10-21)
+
+### Features
+
+ * support torch_distributed distribution for Trainium instances
+
+### Bug Fixes and Other Changes
+
+ * bump apache-airflow from 2.4.0 to 2.4.1 in /requirements/extras
+
+### Documentation Changes
+
+ * fix kwargs and descriptions of the smdmp checkpoint function
+ * add the doc for the MonitorBatchTransformStep
+
 ## v2.112.2 (2022-10-11)
 
 ### Bug Fixes and Other Changes

VERSION

@@ -1 +1 @@
-2.112.3.dev0
+2.113.1.dev0

src/sagemaker/fw_utils.py

@@ -134,10 +134,13 @@
     "1.12.0",
 ]
 
+
 TORCH_DISTRIBUTED_SUPPORTED_FRAMEWORK_VERSIONS = ["1.11", "1.11.0"]
 
+
 TRAINIUM_SUPPORTED_DISTRIBUTION_STRATEGIES = ["torch_distributed"]
 
+
 SMDISTRIBUTED_SUPPORTED_STRATEGIES = ["dataparallel", "modelparallel"]
 
 
@@ -160,6 +163,12 @@ def validate_source_dir(script, directory):
     return True
 
 
+GRAVITON_ALLOWED_TARGET_INSTANCE_FAMILY = ["c6g", "t4g", "r6g", "m6g"]
+
+
+GRAVITON_ALLOWED_FRAMEWORKS = set(["tensorflow", "pytorch"])
+
+
 def validate_source_code_input_against_pipeline_variables(
     entry_point: Optional[Union[str, PipelineVariable]] = None,
     source_dir: Optional[Union[str, PipelineVariable]] = None,

src/sagemaker/image_uri_config/pytorch.json

@@ -654,6 +654,51 @@
             }
         }
     },
+    "inference_graviton": {
+        "processors": [
+            "cpu"
+        ],
+        "version_aliases": {
+            "1.12": "1.12.1"
+        },
+        "versions": {
+            "1.12.1": {
+                "py_versions": [
+                    "py38"
+                ],
+                "registries": {
+                    "af-south-1": "626614931356",
+                    "ap-east-1": "871362719292",
+                    "ap-northeast-1": "763104351884",
+                    "ap-northeast-2": "763104351884",
+                    "ap-northeast-3": "364406365360",
+                    "ap-south-1": "763104351884",
+                    "ap-southeast-1": "763104351884",
+                    "ap-southeast-2": "763104351884",
+                    "ap-southeast-3": "907027046896",
+                    "ca-central-1": "763104351884",
+                    "cn-north-1": "727897471807",
+                    "cn-northwest-1": "727897471807",
+                    "eu-central-1": "763104351884",
+                    "eu-north-1": "763104351884",
+                    "eu-west-1": "763104351884",
+                    "eu-west-2": "763104351884",
+                    "eu-west-3": "763104351884",
+                    "eu-south-1": "692866216735",
+                    "me-south-1": "217643126080",
+                    "sa-east-1": "763104351884",
+                    "us-east-1": "763104351884",
+                    "us-east-2": "763104351884",
+                    "us-gov-west-1": "442386744353",
+                    "us-iso-east-1": "886529160074",
+                    "us-west-1": "763104351884",
+                    "us-west-2": "763104351884"
+                },
+                "repository": "pytorch-inference-graviton",
+                "container_version": {"cpu": "ubuntu20.04"}
+            }
+        }
+    },
     "training": {
         "processors": [
             "cpu",

src/sagemaker/image_uri_config/tensorflow.json

@@ -1471,6 +1471,51 @@
             }
         }
     },
+    "inference_graviton": {
+        "processors": [
+            "cpu"
+        ],
+        "version_aliases": {
+            "2.9": "2.9.1"
+        },
+        "versions": {
+            "2.9.1": {
+                "py_versions": [
+                    "py38"
+                ],
+                "registries": {
+                    "af-south-1": "626614931356",
+                    "ap-east-1": "871362719292",
+                    "ap-northeast-1": "763104351884",
+                    "ap-northeast-2": "763104351884",
+                    "ap-northeast-3": "364406365360",
+                    "ap-south-1": "763104351884",
+                    "ap-southeast-1": "763104351884",
+                    "ap-southeast-2": "763104351884",
+                    "ap-southeast-3": "907027046896",
+                    "ca-central-1": "763104351884",
+                    "cn-north-1": "727897471807",
+                    "cn-northwest-1": "727897471807",
+                    "eu-central-1": "763104351884",
+                    "eu-north-1": "763104351884",
+                    "eu-west-1": "763104351884",
+                    "eu-west-2": "763104351884",
+                    "eu-west-3": "763104351884",
+                    "eu-south-1": "692866216735",
+                    "me-south-1": "217643126080",
+                    "sa-east-1": "763104351884",
+                    "us-east-1": "763104351884",
+                    "us-east-2": "763104351884",
+                    "us-gov-west-1": "442386744353",
+                    "us-iso-east-1": "886529160074",
+                    "us-west-1": "763104351884",
+                    "us-west-2": "763104351884"
+                },
+                "repository": "tensorflow-inference-graviton",
+                "container_version": {"cpu": "ubuntu20.04"}
+            }
+        }
+    },
     "training": {
         "processors": [
             "cpu",

src/sagemaker/image_uris.py

@@ -25,6 +25,7 @@
 from sagemaker.jumpstart import artifacts
 from sagemaker.workflow import is_pipeline_variable
 from sagemaker.workflow.utilities import override_pipeline_parameter_var
+from sagemaker.fw_utils import GRAVITON_ALLOWED_TARGET_INSTANCE_FAMILY, GRAVITON_ALLOWED_FRAMEWORKS
 
 logger = logging.getLogger(__name__)
 
@@ -151,6 +152,7 @@ def retrieve(
             inference_tool = _get_inference_tool(inference_tool, instance_type)
             if inference_tool == "neuron":
                 _framework = f"{framework}-{inference_tool}"
+        image_scope = _get_image_scope_for_instance_type(_framework, instance_type, image_scope)
         config = _config_for_framework_and_scope(_framework, image_scope, accelerator_type)
 
     original_version = version
@@ -216,6 +218,9 @@ def retrieve(
     else:
         tag_prefix = version_config.get("tag_prefix", version)
 
+    if repo == f"{framework}-inference-graviton":
+        container_version = f"{container_version}-sagemaker"
+
     tag = _format_tag(tag_prefix, processor, py_version, container_version, inference_tool)
 
     if instance_type is not None and _should_auto_select_container_version(
@@ -287,6 +292,15 @@ def config_for_framework(framework):
         return json.load(f)
 
 
+def _get_image_scope_for_instance_type(framework, instance_type, image_scope):
+    """Extract the image scope from instance type."""
+    if framework in GRAVITON_ALLOWED_FRAMEWORKS and isinstance(instance_type, str):
+        match = re.match(r"^ml[\._]([a-z\d]+)\.?\w*$", instance_type)
+        if match and match[1] in GRAVITON_ALLOWED_TARGET_INSTANCE_FAMILY:
+            return "inference_graviton"
+    return image_scope
+
+
 def _get_inference_tool(inference_tool, instance_type):
     """Extract the inference tool name from instance type."""
     if not inference_tool and instance_type:

src/sagemaker/inputs.py

@@ -176,6 +176,7 @@ class TransformInput(object):
     output_filter: str = attr.ib(default=None)
     join_source: str = attr.ib(default=None)
     model_client_config: dict = attr.ib(default=None)
+    batch_data_capture_config: dict = attr.ib(default=None)
 
 
 class FileSystemInput(object):
@@ -232,3 +233,45 @@ def __init__(
 
         if content_type:
             self.config["ContentType"] = content_type
+
+
+class BatchDataCaptureConfig(object):
+    """Configuration object passed in when create a batch transform job.
+
+    Specifies configuration related to batch transform job data capture for use with
+    Amazon SageMaker Model Monitoring
+    """
+
+    def __init__(
+        self,
+        destination_s3_uri: str,
+        kms_key_id: str = None,
+        generate_inference_id: bool = None,
+    ):
+        """Create new BatchDataCaptureConfig
+
+        Args:
+            destination_s3_uri (str): S3 Location to store the captured data
+            kms_key_id (str): The KMS key to use when writing to S3.
+                KmsKeyId can be an ID of a KMS key, ARN of a KMS key, alias of a KMS key,
+                or alias of a KMS key. The KmsKeyId is applied to all outputs.
+                (default: None)
+            generate_inference_id (bool): Flag to generate an inference id
+                (default: None)
+        """
+        self.destination_s3_uri = destination_s3_uri
+        self.kms_key_id = kms_key_id
+        self.generate_inference_id = generate_inference_id
+
+    def _to_request_dict(self):
+        """Generates a request dictionary using the parameters provided to the class."""
+        batch_data_capture_config = {
+            "DestinationS3Uri": self.destination_s3_uri,
+        }
+
+        if self.kms_key_id is not None:
+            batch_data_capture_config["KmsKeyId"] = self.kms_key_id
+        if self.generate_inference_id is not None:
+            batch_data_capture_config["GenerateInferenceId"] = self.generate_inference_id
+
+        return batch_data_capture_config

src/sagemaker/job.py

@@ -68,7 +68,7 @@ def _load_config(inputs, estimator, expand_role=True, validate_uri=True):
         input_config = _Job._format_inputs_to_input_config(inputs, validate_uri)
         role = (
             estimator.sagemaker_session.expand_role(estimator.role)
-            if expand_role
+            if (expand_role and not is_pipeline_variable(estimator.role))
             else estimator.role
         )
         output_config = _Job._prepare_output_config(estimator.output_path, estimator.output_kms_key)

src/sagemaker/model_monitor/__init__.py

@@ -23,6 +23,7 @@
 from sagemaker.model_monitor.model_monitoring import BaseliningJob  # noqa: F401
 from sagemaker.model_monitor.model_monitoring import MonitoringExecution  # noqa: F401
 from sagemaker.model_monitor.model_monitoring import EndpointInput  # noqa: F401
+from sagemaker.model_monitor.model_monitoring import BatchTransformInput  # noqa: F401
 from sagemaker.model_monitor.model_monitoring import MonitoringOutput  # noqa: F401
 from sagemaker.model_monitor.model_monitoring import ModelQualityMonitor  # noqa: F401
 
@@ -42,5 +43,6 @@
 
 from sagemaker.model_monitor.data_capture_config import DataCaptureConfig  # noqa: F401
 from sagemaker.model_monitor.dataset_format import DatasetFormat  # noqa: F401
+from sagemaker.model_monitor.dataset_format import MonitoringDatasetFormat  # noqa: F401
 
 from sagemaker.network import NetworkConfig  # noqa: F401