Merge branch 'master' into master

Author: aws-patlin

.flake8

@@ -3,3 +3,4 @@ application_import_names = sagemaker, tests
 import-order-style = google
 per-file-ignores =
     tests/unit/test_tuner.py: F405
+    src/sagemaker/config/config_schema.py: E501

CHANGELOG.md

@@ -1,5 +1,16 @@
 # Changelog
 
+## v2.148.0 (2023-04-20)
+
+### Features
+
+ * [huggingface] Add `torch.distributed` support for Trainium and `torchrun`
+ * Add PyTorch 2.0 to SDK
+
+### Bug Fixes and Other Changes
+
+ * updating batch transform job in monitoring schedule
+
 ## v2.147.0 (2023-04-18)
 
 ### Features

README.rst

@@ -133,6 +133,8 @@ To run the integration tests, the following prerequisites must be met
 1. AWS account credentials are available in the environment for the boto3 client to use.
 2. The AWS account has an IAM role named :code:`SageMakerRole`.
    It should have the AmazonSageMakerFullAccess policy attached as well as a policy with `the necessary permissions to use Elastic Inference <https://docs.aws.amazon.com/sagemaker/latest/dg/ei-setup.html>`__.
+3. To run remote_function tests, dummy ecr repo should be created. It can be created by running -
+    :code:`aws ecr create-repository --repository-name remote-function-dummy-container`
 
 We recommend selectively running just those integration tests you'd like to run. You can filter by individual test function names with:
 

VERSION

@@ -1 +1 @@
-2.147.1.dev0
+2.148.1.dev0

requirements/extras/test_requirements.txt

@@ -21,3 +21,4 @@ sagemaker-experiments==0.1.35
 Jinja2==3.0.3
 pandas>=1.3.5,<1.5
 scikit-learn==1.0.2
+cloudpickle==2.2.1

setup.py

@@ -49,6 +49,7 @@ def read_requirements(filename):
 required_packages = [
     "attrs>=20.3.0,<23",
     "boto3>=1.26.28,<2.0",
+    "cloudpickle==2.2.1",
     "google-pasta",
     "numpy>=1.9.0,<2.0",
     "protobuf>=3.1,<4.0",
@@ -62,6 +63,7 @@ def read_requirements(filename):
     "PyYAML==5.4.1",
     "jsonschema",
     "platformdirs",
+    "tblib==1.7.0",
 ]
 
 # Specific use case dependencies

src/sagemaker/clarify.py

@@ -26,7 +26,7 @@
 import tempfile
 from abc import ABC, abstractmethod
 from typing import List, Union, Dict, Optional, Any
-
+from enum import Enum
 from schema import Schema, And, Use, Or, Optional as SchemaOptional, Regex
 
 from sagemaker import image_uris, s3, utils
@@ -304,6 +304,16 @@
 )
 
 
+class DatasetType(Enum):
+    """Enum to store different dataset types supported in the Analysis config file"""
+
+    TEXTCSV = "text/csv"
+    JSONLINES = "application/jsonlines"
+    JSON = "application/json"
+    PARQUET = "application/x-parquet"
+    IMAGE = "application/x-image"
+
+
 class DataConfig:
     """Config object related to configurations of the input and output dataset."""
 
@@ -1451,7 +1461,7 @@ def _run(
                 source=self._CLARIFY_OUTPUT,
                 destination=data_config.s3_output_path,
                 output_name="analysis_result",
-                s3_upload_mode="EndOfJob",
+                s3_upload_mode=ProcessingOutputHandler.get_s3_upload_mode(analysis_config),
             )
 
             return super().run(
@@ -2171,6 +2181,33 @@ def _upload_analysis_config(analysis_config_file, s3_output_path, sagemaker_sess
     )
 
 
+class ProcessingOutputHandler:
+    """Class to handle the parameters for SagemakerProcessor.Processingoutput"""
+
+    class S3UploadMode(Enum):
+        """Enum values for different uplaod modes to s3 bucket"""
+
+        CONTINUOUS = "Continuous"
+        ENDOFJOB = "EndOfJob"
+
+    @classmethod
+    def get_s3_upload_mode(cls, analysis_config: Dict[str, Any]) -> str:
+        """Fetches s3_upload mode based on the shap_config values
+
+        Args:
+            analysis_config (dict): dict Config following the analysis_config.json format
+
+        Returns:
+            The s3_upload_mode type for the processing output.
+        """
+        dataset_type = analysis_config["dataset_type"]
+        return (
+            ProcessingOutputHandler.S3UploadMode.CONTINUOUS.value
+            if dataset_type == DatasetType.IMAGE.value
+            else ProcessingOutputHandler.S3UploadMode.ENDOFJOB.value
+        )
+
+
 def _set(value, key, dictionary):
     """Sets dictionary[key] = value if value is not None."""
     if value is not None:

src/sagemaker/config/config_schema.py

@@ -44,6 +44,17 @@
 SAGEMAKER = "SageMaker"
 PYTHON_SDK = "PythonSDK"
 MODULES = "Modules"
+REMOTE_FUNCTION = "RemoteFunction"
+DEPENDENCIES = "Dependencies"
+PRE_EXECUTION_SCRIPT = "PreExecutionScript"
+PRE_EXECUTION_COMMANDS = "PreExecutionCommands"
+ENVIRONMENT_VARIABLES = "EnvironmentVariables"
+IMAGE_URI = "ImageUri"
+INCLUDE_LOCAL_WORKDIR = "IncludeLocalWorkDir"
+INSTANCE_TYPE = "InstanceType"
+S3_KMS_KEY_ID = "S3KmsKeyId"
+S3_ROOT_URI = "S3RootUri"
+JOB_CONDA_ENV = "JobCondaEnvironment"
 OFFLINE_STORE_CONFIG = "OfflineStoreConfig"
 ONLINE_STORE_CONFIG = "OnlineStoreConfig"
 S3_STORAGE_CONFIG = "S3StorageConfig"
@@ -221,6 +232,49 @@ def _simple_path(*args: str):
     SAGEMAKER, MODEL_PACKAGE, VALIDATION_SPECIFICATION, VALIDATION_PROFILES
 )
 
+REMOTE_FUNCTION_DEPENDENCIES = _simple_path(
+    SAGEMAKER, PYTHON_SDK, MODULES, REMOTE_FUNCTION, DEPENDENCIES
+)
+REMOTE_FUNCTION_PRE_EXECUTION_COMMANDS = _simple_path(
+    SAGEMAKER, PYTHON_SDK, MODULES, REMOTE_FUNCTION, PRE_EXECUTION_COMMANDS
+)
+REMOTE_FUNCTION_PRE_EXECUTION_SCRIPT = _simple_path(
+    SAGEMAKER, PYTHON_SDK, MODULES, REMOTE_FUNCTION, PRE_EXECUTION_SCRIPT
+)
+REMOTE_FUNCTION_ENVIRONMENT_VARIABLES = _simple_path(
+    SAGEMAKER, PYTHON_SDK, MODULES, REMOTE_FUNCTION, ENVIRONMENT_VARIABLES
+)
+REMOTE_FUNCTION_IMAGE_URI = _simple_path(SAGEMAKER, PYTHON_SDK, MODULES, REMOTE_FUNCTION, IMAGE_URI)
+REMOTE_FUNCTION_INCLUDE_LOCAL_WORKDIR = _simple_path(
+    SAGEMAKER, PYTHON_SDK, MODULES, REMOTE_FUNCTION, INCLUDE_LOCAL_WORKDIR
+)
+REMOTE_FUNCTION_INSTANCE_TYPE = _simple_path(
+    SAGEMAKER, PYTHON_SDK, MODULES, REMOTE_FUNCTION, INSTANCE_TYPE
+)
+REMOTE_FUNCTION_JOB_CONDA_ENV = _simple_path(
+    SAGEMAKER, PYTHON_SDK, MODULES, REMOTE_FUNCTION, JOB_CONDA_ENV
+)
+REMOTE_FUNCTION_ROLE_ARN = _simple_path(SAGEMAKER, PYTHON_SDK, MODULES, REMOTE_FUNCTION, ROLE_ARN)
+REMOTE_FUNCTION_S3_KMS_KEY_ID = _simple_path(
+    SAGEMAKER, PYTHON_SDK, MODULES, REMOTE_FUNCTION, S3_KMS_KEY_ID
+)
+REMOTE_FUNCTION_S3_ROOT_URI = _simple_path(
+    SAGEMAKER, PYTHON_SDK, MODULES, REMOTE_FUNCTION, S3_ROOT_URI
+)
+REMOTE_FUNCTION_TAGS = _simple_path(SAGEMAKER, PYTHON_SDK, MODULES, REMOTE_FUNCTION, TAGS)
+REMOTE_FUNCTION_VOLUME_KMS_KEY_ID = _simple_path(
+    SAGEMAKER, PYTHON_SDK, MODULES, REMOTE_FUNCTION, VOLUME_KMS_KEY_ID
+)
+REMOTE_FUNCTION_VPC_CONFIG_SUBNETS = _simple_path(
+    SAGEMAKER, PYTHON_SDK, MODULES, REMOTE_FUNCTION, VPC_CONFIG, SUBNETS
+)
+REMOTE_FUNCTION_VPC_CONFIG_SECURITY_GROUP_IDS = _simple_path(
+    SAGEMAKER, PYTHON_SDK, MODULES, REMOTE_FUNCTION, VPC_CONFIG, SECURITY_GROUP_IDS
+)
+REMOTE_FUNCTION_ENABLE_INTER_CONTAINER_TRAFFIC_ENCRYPTION = _simple_path(
+    SAGEMAKER, PYTHON_SDK, MODULES, REMOTE_FUNCTION, ENABLE_INTER_CONTAINER_TRAFFIC_ENCRYPTION
+)
+
 # Paths for reference elsewhere in the SDK.
 # Names include the schema version since the paths could change with other schema versions
 MONITORING_SCHEDULE_INTER_CONTAINER_ENCRYPTION_PATH = _simple_path(
@@ -245,7 +299,6 @@ def _simple_path(*args: str):
     SAGEMAKER, TRAINING_JOB, ENABLE_INTER_CONTAINER_TRAFFIC_ENCRYPTION
 )
 
-
 SAGEMAKER_PYTHON_SDK_CONFIG_SCHEMA = {
     "$schema": "https://json-schema.org/draft/2020-12/schema",
     TYPE: OBJECT,
@@ -377,6 +430,23 @@ def _simple_path(*args: str):
             "minItems": 0,
             "maxItems": 50,
         },
+        # Regex is taken from https://docs.aws.amazon.com/sagemaker/latest/APIReference/API_CreateTrainingJob.html#sagemaker-CreateTrainingJob-request-Environment
+        "environmentVariables": {
+            TYPE: OBJECT,
+            ADDITIONAL_PROPERTIES: False,
+            PATTERN_PROPERTIES: {
+                r"([a-zA-Z_][a-zA-Z0-9_]*){1,512}": {
+                    TYPE: "string",
+                    "pattern": r"[\S\s]*",
+                    "maxLength": 512,
+                }
+            },
+            "maxProperties": 48,
+        },
+        # Regex is taken from https://docs.aws.amazon.com/sagemaker/latest/APIReference/API_S3DataSource.html#sagemaker-Type-S3DataSource-S3Uri
+        "s3Uri": {TYPE: "string", "pattern": "^(https|s3)://([^/]+)/?(.*)$", "maxLength": 1024},
+        # Regex is taken from https://docs.aws.amazon.com/sagemaker/latest/APIReference/API_AlgorithmSpecification.html#sagemaker-Type-AlgorithmSpecification-ContainerEntrypoint
+        "preExecutionCommand": {TYPE: "string", "pattern": r".*"},
     },
     PROPERTIES: {
         SCHEMA_VERSION: {
@@ -406,6 +476,36 @@ def _simple_path(*args: str):
                             # Any SageMaker Python SDK specific configuration will be added here.
                             TYPE: OBJECT,
                             ADDITIONAL_PROPERTIES: False,
+                            PROPERTIES: {
+                                REMOTE_FUNCTION: {
+                                    TYPE: OBJECT,
+                                    ADDITIONAL_PROPERTIES: False,
+                                    PROPERTIES: {
+                                        DEPENDENCIES: {TYPE: "string"},
+                                        PRE_EXECUTION_COMMANDS: {
+                                            TYPE: "array",
+                                            "items": {"$ref": "#/definitions/preExecutionCommand"},
+                                        },
+                                        PRE_EXECUTION_SCRIPT: {TYPE: "string"},
+                                        ENABLE_INTER_CONTAINER_TRAFFIC_ENCRYPTION: {
+                                            TYPE: "boolean"
+                                        },
+                                        ENVIRONMENT_VARIABLES: {
+                                            "$ref": "#/definitions/environmentVariables"
+                                        },
+                                        IMAGE_URI: {TYPE: "string"},
+                                        INCLUDE_LOCAL_WORKDIR: {TYPE: "boolean"},
+                                        INSTANCE_TYPE: {TYPE: "string"},
+                                        JOB_CONDA_ENV: {TYPE: "string"},
+                                        ROLE_ARN: {"$ref": "#/definitions/roleArn"},
+                                        S3_KMS_KEY_ID: {"$ref": "#/definitions/kmsKeyId"},
+                                        S3_ROOT_URI: {"$ref": "#/definitions/s3Uri"},
+                                        TAGS: {"$ref": "#/definitions/tags"},
+                                        VOLUME_KMS_KEY_ID: {"$ref": "#/definitions/kmsKeyId"},
+                                        VPC_CONFIG: {"$ref": "#/definitions/vpcConfig"},
+                                    },
+                                }
+                            },
                         }
                     },
                 },

src/sagemaker/experiments/run.py

@@ -715,6 +715,14 @@ def __exit__(self, exc_type, exc_value, exc_traceback):
 
         self.close()
 
+    def __getstate__(self):
+        """Overriding this method to prevent instance of Run from being pickled.
+
+        Raise:
+            NotImplementedError: If attempting to pickle this instance.
+        """
+        raise NotImplementedError("Instance of Run type is not allowed to be pickled.")
+
 
 def load_run(
     run_name: Optional[str] = None,
@@ -787,36 +795,38 @@ def load_run(
     Returns:
         Run: The loaded Run object.
     """
-    sagemaker_session = sagemaker_session or _utils.default_session()
     environment = _RunEnvironment.load()
 
     verify_load_input_names(run_name=run_name, experiment_name=experiment_name)
 
-    if run_name or environment:
-        if run_name:
-            logger.warning(
-                "run_name is explicitly supplied in load_run, "
-                "which will be prioritized to load the Run object. "
-                "In other words, the run name in the experiment config, fetched from the "
-                "job environment or the current run context, will be ignored."
-            )
-        else:
-            exp_config = get_tc_and_exp_config_from_job_env(
-                environment=environment, sagemaker_session=sagemaker_session
-            )
-            run_name = Run._extract_run_name_from_tc_name(
-                trial_component_name=exp_config[RUN_NAME],
-                experiment_name=exp_config[EXPERIMENT_NAME],
-            )
-            experiment_name = exp_config[EXPERIMENT_NAME]
-
+    if run_name:
+        logger.warning(
+            "run_name is explicitly supplied in load_run, "
+            "which will be prioritized to load the Run object. "
+            "In other words, the run name in the experiment config, fetched from the "
+            "job environment or the current run context, will be ignored."
+        )
         run_instance = Run(
             experiment_name=experiment_name,
             run_name=run_name,
-            sagemaker_session=sagemaker_session,
+            sagemaker_session=sagemaker_session or _utils.default_session(),
         )
     elif _RunContext.get_current_run():
         run_instance = _RunContext.get_current_run()
+    elif environment:
+        exp_config = get_tc_and_exp_config_from_job_env(
+            environment=environment, sagemaker_session=sagemaker_session or _utils.default_session()
+        )
+        run_name = Run._extract_run_name_from_tc_name(
+            trial_component_name=exp_config[RUN_NAME],
+            experiment_name=exp_config[EXPERIMENT_NAME],
+        )
+        experiment_name = exp_config[EXPERIMENT_NAME]
+        run_instance = Run(
+            experiment_name=experiment_name,
+            run_name=run_name,
+            sagemaker_session=sagemaker_session or _utils.default_session(),
+        )
     else:
         raise RuntimeError(
             "Failed to load a Run object. "

src/sagemaker/image_uri_config/sagemaker-base-python.json

@@ -0,0 +1,36 @@
+{
+    "versions": {
+        "1.0": {
+            "registries": {
+                "us-east-2": "429704687514",
+                "me-south-1": "117516905037",
+                "us-west-2": "236514542706",
+                "ca-central-1": "310906938811",
+                "ap-east-1": "493642496378",
+                "us-east-1": "081325390199",
+                "ap-northeast-2": "806072073708",
+                "eu-west-2": "712779665605",
+                "ap-southeast-2": "52832661640",
+                "cn-northwest-1": "390780980154",
+                "eu-north-1": "243637512696",
+                "cn-north-1": "390048526115",
+                "ap-south-1": "394103062818",
+                "eu-west-3": "615547856133",
+                "ap-southeast-3": "276181064229",
+                "af-south-1": "559312083959",
+                "eu-west-1": "470317259841",
+                "eu-central-1": "936697816551",
+                "sa-east-1": "782484402741",
+                "ap-northeast-3": "792733760839",
+                "eu-south-1": "592751261982",
+                "ap-northeast-1": "102112518831",
+                "us-west-1": "742091327244",
+                "ap-southeast-1": "492261229750",
+                "me-central-1": "103105715889",
+                "us-gov-east-1": "107072934176",
+                "us-gov-west-1": "107173498710"
+            },
+            "repository": "sagemaker-base-python"
+        }
+    }
+}