sagemaker-core 1.0.62__py3-none-any.whl → 2.1.1__py3-none-any.whl

sagemaker/core/fw_utils.py

@@ -0,0 +1,1176 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+"""Utility methods used by framework classes."""
+from __future__ import absolute_import
+
+import json
+import logging
+import os
+import re
+import shutil
+import tempfile
+import time
+from collections import namedtuple
+from typing import Dict, List, Optional, Union
+
+from packaging import version
+
+import sagemaker.core.common_utils as sagemaker_utils
+from sagemaker.core.deprecations import deprecation_warn_base, renamed_kwargs
+from sagemaker.core.instance_group import InstanceGroup
+from sagemaker.core.s3 import s3_path_join
+from sagemaker.core.session_settings import SessionSettings
+from sagemaker.core.workflow import is_pipeline_variable
+from sagemaker.core.helper.pipeline_variable import PipelineVariable
+
+logger = logging.getLogger(__name__)
+
+_TAR_SOURCE_FILENAME = "source.tar.gz"
+
+UploadedCode = namedtuple("UploadedCode", ["s3_prefix", "script_name"])
+"""sagemaker.fw_utils.UploadedCode: An object containing the S3 prefix and script name.
+
+This is for the source code used for the entry point with an ``Estimator``. It can be
+instantiated with positional or keyword arguments.
+"""
+
+PYTHON_2_DEPRECATION_WARNING = (
+    "{latest_supported_version} is the latest version of {framework} that supports "
+    "Python 2. Newer versions of {framework} will only be available for Python 3."
+    "Please set the argument \"py_version='py3'\" to use the Python 3 {framework} image."
+)
+PARAMETER_SERVER_MULTI_GPU_WARNING = (
+    "If you have selected a multi-GPU training instance type "
+    "and also enabled parameter server for distributed training, "
+    "distributed training with the default parameter server configuration will not "
+    "fully leverage all GPU cores; the parameter server will be configured to run "
+    "only one worker per host regardless of the number of GPUs."
+)
+
+DEBUGGER_UNSUPPORTED_REGIONS = (
+    "us-iso-east-1",
+    "us-isob-east-1",
+    "ap-southeast-3",
+    "ap-southeast-4",
+    "eu-south-2",
+    "me-central-1",
+    "ap-south-2",
+    "eu-central-2",
+    "us-gov-east-1",
+)
+PROFILER_UNSUPPORTED_REGIONS = (
+    "us-iso-east-1",
+    "us-isob-east-1",
+    "ap-southeast-3",
+    "ap-southeast-4",
+    "eu-south-2",
+    "me-central-1",
+    "ap-south-2",
+    "eu-central-2",
+    "us-gov-east-1",
+)
+
+SINGLE_GPU_INSTANCE_TYPES = ("ml.p2.xlarge", "ml.p3.2xlarge")
+SM_DATAPARALLEL_SUPPORTED_INSTANCE_TYPES = (
+    "ml.p3.16xlarge",
+    "ml.p3dn.24xlarge",
+    "ml.p4d.24xlarge",
+    "ml.p4de.24xlarge",
+    "local_gpu",
+)
+SM_DATAPARALLEL_SUPPORTED_FRAMEWORK_VERSIONS = {
+    # tf 2.12 should not be supported: smdataparallel excludes support for tf>=2.12.
+    "tensorflow": [
+        "2.3",
+        "2.3.1",
+        "2.3.2",
+        "2.4",
+        "2.4.1",
+        "2.4.3",
+        "2.5",
+        "2.5.0",
+        "2.5.1",
+        "2.6",
+        "2.6.0",
+        "2.6.2",
+        "2.6.3",
+        "2.7",
+        "2.7.1",
+        "2.8",
+        "2.8.0",
+        "2.9",
+        "2.9.1",
+        "2.9.2",
+        "2.10",
+        "2.10.1",
+        "2.11",
+        "2.11.0",
+    ],
+    "pytorch": [
+        "1.6",
+        "1.6.0",
+        "1.7",
+        "1.7.1",
+        "1.8",
+        "1.8.0",
+        "1.8.1",
+        "1.9",
+        "1.9.0",
+        "1.9.1",
+        "1.10",
+        "1.10.0",
+        "1.10.2",
+        "1.11",
+        "1.11.0",
+        "1.12",
+        "1.12.0",
+        "1.12.1",
+        "1.13.1",
+        "2.0.0",
+        "2.0.1",
+        "2.1.0",
+        "2.1.2",
+        "2.2.0",
+    ],
+}
+
+TORCH_DISTRIBUTED_GPU_SUPPORTED_FRAMEWORK_VERSIONS = [
+    "1.13.1",
+    "2.0.0",
+    "2.0.1",
+    "2.1.0",
+    "2.1.2",
+    "2.2.0",
+    "2.3.0",
+    "2.3.1",
+    "2.4.1",
+    "2.5.1",
+]
+
+TRAINIUM_SUPPORTED_DISTRIBUTION_STRATEGIES = ["torch_distributed"]
+TRAINIUM_SUPPORTED_TORCH_DISTRIBUTED_FRAMEWORK_VERSIONS = [
+    "1.11",
+    "1.11.0",
+    "1.12",
+    "1.12.0",
+    "1.12.1",
+    "1.13.1",
+    "2.0.0",
+]
+
+SMDISTRIBUTED_SUPPORTED_STRATEGIES = ["dataparallel", "modelparallel"]
+
+
+GRAVITON_ALLOWED_TARGET_INSTANCE_FAMILY = [
+    "m6g",
+    "m6gd",
+    "c6g",
+    "c6gd",
+    "c6gn",
+    "c7g",
+    "r6g",
+    "r6gd",
+]
+
+
+GRAVITON_ALLOWED_FRAMEWORKS = set(["tensorflow", "pytorch", "xgboost", "sklearn"])
+
+
+def validate_source_dir(script, directory):
+    """Validate that the source directory exists and it contains the user script.
+
+    Args:
+        script (str): Script filename.
+        directory (str): Directory containing the source file.
+    Raises:
+        ValueError: If ``directory`` does not exist, is not a directory, or does
+            not contain ``script``.
+    """
+    if directory:
+        if not os.path.isfile(os.path.join(directory, script)):
+            raise ValueError(
+                'No file named "{}" was found in directory "{}".'.format(script, directory)
+            )
+
+    return True
+
+
+def validate_source_code_input_against_pipeline_variables(
+    entry_point: Optional[Union[str, PipelineVariable]] = None,
+    source_dir: Optional[Union[str, PipelineVariable]] = None,
+    git_config: Optional[Dict[str, str]] = None,
+    enable_network_isolation: Union[bool, PipelineVariable] = False,
+):
+    """Validate source code input against pipeline variables.
+
+    Args:
+        entry_point (str or PipelineVariable): The path to the local Python source file that
+            should be executed as the entry point to training (default: None).
+        source_dir (str or PipelineVariable): The Path to a directory with any other
+            training source code dependencies aside from the entry point file (default: None).
+        git_config (Dict[str, str]): Git configurations used for cloning files (default: None).
+        enable_network_isolation (bool or PipelineVariable): Specifies whether container will run
+            in network isolation mode (default: False).
+    """
+    if is_pipeline_variable(enable_network_isolation) or enable_network_isolation is True:
+        if is_pipeline_variable(entry_point) or is_pipeline_variable(source_dir):
+            raise TypeError(
+                "entry_point, source_dir should not be pipeline variables "
+                "when enable_network_isolation is a pipeline variable or it is set to True."
+            )
+    if git_config:
+        if is_pipeline_variable(entry_point) or is_pipeline_variable(source_dir):
+            raise TypeError(
+                "entry_point, source_dir should not be pipeline variables when git_config is given."
+            )
+    if is_pipeline_variable(entry_point):
+        if not source_dir:
+            raise TypeError(
+                "The entry_point should not be a pipeline variable when source_dir is missing."
+            )
+        if not is_pipeline_variable(source_dir) and not source_dir.lower().startswith("s3://"):
+            raise TypeError(
+                "The entry_point should not be a pipeline variable when source_dir is a local path."
+            )
+        logger.warning(
+            "The entry_point is a pipeline variable: %s. During pipeline execution, "
+            "the interpreted value of entry_point has to be a local path in the container "
+            "pointing to a Python source file which is located at the root of source_dir.",
+            type(entry_point),
+        )
+    if is_pipeline_variable(source_dir):
+        logger.warning(
+            "The source_dir is a pipeline variable: %s. During pipeline execution, "
+            "the interpreted value of source_dir has to be an S3 URI and "
+            "must point to a file with name ``sourcedir.tar.gz``",
+            type(source_dir),
+        )
+
+
+def parse_mp_parameters(params):
+    """Parse the model parallelism parameters provided by the user.
+
+    Args:
+        params: a string representing path to an existing config, or
+                a config dict.
+
+    Returns:
+        parsed: a dict of parsed config.
+
+    Raises:
+        ValueError: if params is not a string or a dict, or
+                    the config file cannot be parsed as json.
+    """
+    parsed = None
+    if isinstance(params, dict):
+        parsed = params
+    elif os.path.exists(params):
+        try:
+            with open(params, "r") as fp:
+                parsed = json.load(fp)
+        except json.decoder.JSONDecodeError:
+            pass
+    else:
+        raise ValueError(
+            f"Expected a string path to an existing modelparallel config, or a dictionary. "
+            f"Received: {params}."
+        )
+
+    if parsed is None:
+        raise ValueError(f"Cannot parse {params} as a json file.")
+
+    return parsed
+
+
+def get_mp_parameters(distribution):
+    """Get the model parallelism parameters provided by the user.
+
+    Args:
+        distribution: distribution dictionary defined by the user.
+
+    Returns:
+        params: dictionary containing model parallelism parameters
+        used for training.
+    """
+    try:
+        mp_dict = distribution["smdistributed"]["modelparallel"]
+    except KeyError:
+        mp_dict = {}
+    if mp_dict.get("enabled", False) is True:
+        params = mp_dict.get("parameters", {})
+        params = parse_mp_parameters(params)
+        validate_mp_config(params)
+        return params
+    return None
+
+
+def validate_mp_config(config):
+    """Validate the configuration dictionary for model parallelism.
+
+    Args:
+       config (dict): Dictionary holding configuration keys and values.
+
+    Raises:
+        ValueError: If any of the keys have incorrect values.
+    """
+
+    def validate_positive(key):
+        try:
+            if not isinstance(config[key], int) or config[key] < 1:
+                raise ValueError(f"The number of {key} must be a positive integer.")
+        except KeyError:
+            pass
+
+    def validate_in(key, vals):
+        try:
+            if config[key] not in vals:
+                raise ValueError(f"{key} must be a value in: {vals}.")
+        except KeyError:
+            pass
+
+    def validate_bool(keys):
+        validate_in(keys, [True, False])
+
+    validate_in("pipeline", ["simple", "interleaved", "_only_forward"])
+    validate_in("placement_strategy", ["spread", "cluster"])
+    validate_in("optimize", ["speed", "memory"])
+
+    for key in ["microbatches", "partitions", "active_microbatches"]:
+        validate_positive(key)
+
+    for key in [
+        "auto_partition",
+        "contiguous",
+        "load_partition",
+        "horovod",
+        "ddp",
+        "deterministic_server",
+    ]:
+        validate_bool(key)
+
+    if "partition_file" in config and not isinstance(config.get("partition_file"), str):
+        raise ValueError("'partition_file' must be a str.")
+
+    if config.get("auto_partition") is False and "default_partition" not in config:
+        raise ValueError("default_partition must be supplied if auto_partition is set to False!")
+
+    if "default_partition" in config and config["default_partition"] >= config["partitions"]:
+        raise ValueError("default_partition must be less than the number of partitions!")
+
+    if "memory_weight" in config and (
+        config["memory_weight"] > 1.0 or config["memory_weight"] < 0.0
+    ):
+        raise ValueError("memory_weight must be between 0.0 and 1.0!")
+
+    if "ddp_port" in config and "ddp" not in config:
+        raise ValueError("`ddp_port` needs `ddp` to be set as well")
+
+    if "ddp_dist_backend" in config and "ddp" not in config:
+        raise ValueError("`ddp_dist_backend` needs `ddp` to be set as well")
+
+    if "ddp_port" in config:
+        if not isinstance(config["ddp_port"], int) or config["ddp_port"] < 0:
+            value = config["ddp_port"]
+            raise ValueError(f"Invalid port number {value}.")
+
+    if config.get("horovod", False) and config.get("ddp", False):
+        raise ValueError("'ddp' and 'horovod' cannot be simultaneously enabled.")
+
+
+def tar_and_upload_dir(
+    session,
+    bucket,
+    s3_key_prefix,
+    script,
+    directory=None,
+    dependencies=None,
+    kms_key=None,
+    s3_resource=None,
+    settings: Optional[SessionSettings] = None,
+) -> UploadedCode:
+    """Package source files and upload a compress tar file to S3.
+
+    The S3 location will be ``s3://<bucket>/s3_key_prefix/sourcedir.tar.gz``.
+    If directory is an S3 URI, an UploadedCode object will be returned, but
+    nothing will be uploaded to S3 (this allow reuse of code already in S3).
+    If directory is None, the script will be added to the archive at
+    ``./<basename of script>``. If directory is not None, the (recursive) contents
+    of the directory will be added to the archive. directory is treated as the base
+    path of the archive, and the script name is assumed to be a filename or relative path
+    inside the directory.
+
+    Args:
+        session (boto3.Session): Boto session used to access S3.
+        bucket (str): S3 bucket to which the compressed file is uploaded.
+        s3_key_prefix (str): Prefix for the S3 key.
+        script (str): Script filename or path.
+        directory (str): Optional. Directory containing the source file. If it
+            starts with "s3://", no action is taken.
+        dependencies (List[str]): Optional. A list of paths to directories
+            (absolute or relative) containing additional libraries that will be
+            copied into /opt/ml/lib
+        kms_key (str): Optional. KMS key ID used to upload objects to the bucket
+            (default: None).
+        s3_resource (boto3.resource("s3")): Optional. Pre-instantiated Boto3 Resource
+            for S3 connections, can be used to customize the configuration,
+            e.g. set the endpoint URL (default: None).
+        settings (sagemaker.session_settings.SessionSettings): Optional. The settings
+            of the SageMaker ``Session``, can be used to override the default encryption
+            behavior (default: None).
+    Returns:
+        sagemaker.fw_utils.UploadedCode: An object with the S3 bucket and key (S3 prefix) and
+            script name.
+    """
+    if directory and (is_pipeline_variable(directory) or directory.lower().startswith("s3://")):
+        return UploadedCode(s3_prefix=directory, script_name=script)
+
+    script_name = script if directory else os.path.basename(script)
+    dependencies = dependencies or []
+    key = "%s/sourcedir.tar.gz" % s3_key_prefix
+    if (
+        settings is not None
+        and settings.local_download_dir is not None
+        and not (
+            os.path.exists(settings.local_download_dir)
+            and os.path.isdir(settings.local_download_dir)
+        )
+    ):
+        raise ValueError(
+            "Inputted directory for storing newly generated temporary directory does "
+            f"not exist: '{settings.local_download_dir}'"
+        )
+    local_download_dir = None if settings is None else settings.local_download_dir
+    tmp = tempfile.mkdtemp(dir=local_download_dir)
+    encrypt_artifact = True if settings is None else settings.encrypt_repacked_artifacts
+
+    try:
+        source_files = _list_files_to_compress(script, directory) + dependencies
+        tar_file = sagemaker_utils.create_tar_file(
+            source_files, os.path.join(tmp, _TAR_SOURCE_FILENAME)
+        )
+
+        if kms_key:
+            extra_args = {"ServerSideEncryption": "aws:kms", "SSEKMSKeyId": kms_key}
+        elif encrypt_artifact:
+            # encrypt the tarball at rest in S3 with the default AWS managed KMS key for S3
+            # see https://docs.aws.amazon.com/AmazonS3/latest/API/API_PutObject.html#API_PutObject_RequestSyntax
+            extra_args = {"ServerSideEncryption": "aws:kms"}
+        else:
+            extra_args = None
+
+        if s3_resource is None:
+            s3_resource = session.resource("s3", region_name=session.region_name)
+        else:
+            logger.debug("Using provided s3_resource")
+
+        s3_resource.Object(bucket, key).upload_file(tar_file, ExtraArgs=extra_args)
+    finally:
+        shutil.rmtree(tmp)
+
+    return UploadedCode(s3_prefix="s3://%s/%s" % (bucket, key), script_name=script_name)
+
+
+def _list_files_to_compress(script, directory):
+    """Placeholder docstring."""
+    if directory is None:
+        return [script]
+
+    basedir = directory if directory else os.path.dirname(script)
+    return [os.path.join(basedir, name) for name in os.listdir(basedir)]
+
+
+def framework_name_from_image(image_uri):
+    # noinspection LongLine
+    """Extract the framework and Python version from the image name.
+
+    Args:
+        image_uri (str): Image URI, which should be one of the following forms:
+            legacy:
+            '<account>.dkr.ecr.<region>.amazonaws.com/sagemaker-<fw>-<py_ver>-<device>:<container_version>'
+            legacy:
+            '<account>.dkr.ecr.<region>.amazonaws.com/sagemaker-<fw>-<py_ver>-<device>:<fw_version>-<device>-<py_ver>'
+            current:
+            '<account>.dkr.ecr.<region>.amazonaws.com/sagemaker-<fw>:<fw_version>-<device>-<py_ver>'
+            current:
+            '<account>.dkr.ecr.<region>.amazonaws.com/sagemaker-rl-<fw>:<rl_toolkit><rl_version>-<device>-<py_ver>'
+            current:
+            '<account>.dkr.ecr.<region>.amazonaws.com/<fw>-<image_scope>:<fw_version>-<device>-<py_ver>'
+            current:
+            '<account>.dkr.ecr.<region>.amazonaws.com/sagemaker-xgboost:<fw_version>-<container_version>'
+
+    Returns:
+        tuple: A tuple containing:
+
+            - str: The framework name
+            - str: The Python version
+            - str: The image tag
+            - str: If the TensorFlow image is script mode
+    """
+    sagemaker_pattern = re.compile(sagemaker_utils.ECR_URI_PATTERN)
+    sagemaker_match = sagemaker_pattern.match(image_uri)
+    if sagemaker_match is None:
+        return None, None, None, None
+
+    # extract framework, python version and image tag
+    # We must support both the legacy and current image name format.
+    name_pattern = re.compile(
+        r"""^(?:sagemaker(?:-rl)?-)?
+        (tensorflow|mxnet|chainer|pytorch|pytorch-trcomp|scikit-learn|xgboost
+        |huggingface-tensorflow|huggingface-pytorch
+        |huggingface-tensorflow-trcomp|huggingface-pytorch-trcomp)(?:-)?
+        (scriptmode|training)?
+        :(.*)-(.*?)-(py2|py3\d*)(?:.*)$""",
+        re.VERBOSE,
+    )
+    name_match = name_pattern.match(sagemaker_match.group(9))
+    if name_match is not None:
+        fw, scriptmode, ver, device, py = (
+            name_match.group(1),
+            name_match.group(2),
+            name_match.group(3),
+            name_match.group(4),
+            name_match.group(5),
+        )
+        return fw, py, "{}-{}-{}".format(ver, device, py), scriptmode
+
+    legacy_name_pattern = re.compile(r"^sagemaker-(tensorflow|mxnet)-(py2|py3)-(cpu|gpu):(.*)$")
+    legacy_match = legacy_name_pattern.match(sagemaker_match.group(9))
+    if legacy_match is not None:
+        return (legacy_match.group(1), legacy_match.group(2), legacy_match.group(4), None)
+
+    # sagemaker-xgboost images are tagged with two aliases, e.g.:
+    # 1. Long tag: "315553699071.dkr.ecr.us-west-2.amazonaws.com/sagemaker-xgboost:1.5-1-cpu-py3"
+    # 2. Short tag: "315553699071.dkr.ecr.us-west-2.amazonaws.com/sagemaker-xgboost:1.5-1"
+    # Note 1: Both tags point to the same image
+    # Note 2: Both tags have full GPU capabilities, despite "cpu" delineation in the long tag
+    short_xgboost_tag_pattern = re.compile(r"^sagemaker-(xgboost):(.*)$")
+    short_xgboost_tag_match = short_xgboost_tag_pattern.match(sagemaker_match.group(9))
+    if short_xgboost_tag_match is not None:
+        return (short_xgboost_tag_match.group(1), "py3", short_xgboost_tag_match.group(2), None)
+    return None, None, None, None
+
+
+def framework_version_from_tag(image_tag):
+    """Extract the framework version from the image tag.
+
+    Args:
+        image_tag (str): Image tag, which should take the form
+            '<framework_version>-<device>-<py_version>'
+            '<xgboost_version>-<container_version>'
+
+    Returns:
+        str: The framework version.
+    """
+    tag_pattern = re.compile(r"^(.*)-(cpu|gpu)-(py2|py3\d*)$")
+    tag_match = tag_pattern.match(image_tag)
+    if tag_match is None:
+        short_xgboost_tag_pattern = re.compile(r"^(\d\.\d+\-\d)$")
+        tag_match = short_xgboost_tag_pattern.match(image_tag)
+    return None if tag_match is None else tag_match.group(1)
+
+
+def model_code_key_prefix(code_location_key_prefix, model_name, image):
+    """Returns the s3 key prefix for uploading code during model deployment.
+
+    The location returned is a potential concatenation of 2 parts
+        1. code_location_key_prefix if it exists
+        2. model_name or a name derived from the image
+    Args:
+        code_location_key_prefix (str): the s3 key prefix from code_location
+        model_name (str): the name of the model
+        image (str): the image from which a default name can be extracted
+
+    Returns:
+        str: the key prefix to be used in uploading code
+    """
+    name_from_image = f"/model_code/{int(time.time())}"
+    if not is_pipeline_variable(image):
+        name_from_image = sagemaker_utils.name_from_image(image)
+    return s3_path_join(code_location_key_prefix, model_name or name_from_image)
+
+
+def warn_if_parameter_server_with_multi_gpu(training_instance_type, distribution):
+    """Warn the user about training when it doesn't leverage all the GPU cores.
+
+    Warn the user that training will not fully leverage all the GPU
+    cores if parameter server is enabled and a multi-GPU instance is selected.
+    Distributed training with the default parameter server setup doesn't
+    support multi-GPU instances.
+
+    Args:
+        training_instance_type (str): A string representing the type of training instance selected.
+        distribution (dict): A dictionary with information to enable distributed training.
+            (Defaults to None if distributed training is not enabled.) For example:
+
+            .. code:: python
+
+                {
+                    "parameter_server": {
+                        "enabled": True
+                    }
+                }
+    """
+    if training_instance_type == "local" or distribution is None:
+        return
+    if is_pipeline_variable(training_instance_type):
+        # The training_instance_type is not available in compile time.
+        # Rather, it's given in Pipeline execution time
+        return
+
+    is_multi_gpu_instance = (
+        training_instance_type == "local_gpu"
+        or training_instance_type.split(".")[1].startswith("p")
+    ) and training_instance_type not in SINGLE_GPU_INSTANCE_TYPES
+
+    ps_enabled = "parameter_server" in distribution and distribution["parameter_server"].get(
+        "enabled", False
+    )
+
+    if is_multi_gpu_instance and ps_enabled:
+        logger.warning(PARAMETER_SERVER_MULTI_GPU_WARNING)
+
+
+def profiler_config_deprecation_warning(
+    profiler_config, image_uri, framework_name, framework_version
+):
+    """Deprecation message if framework profiling is specified TF >= 2.12 and PT >= 2.0."""
+    if profiler_config is None or profiler_config.framework_profile_params is None:
+        return
+
+    if framework_name not in ("pytorch", "tensorflow"):
+        return
+
+    if framework_version is None:
+        framework_name, _, image_tag, _ = framework_name_from_image(image_uri)
+
+        if image_tag is not None:
+            framework_version = framework_version_from_tag(image_tag)
+
+    if framework_version is not None:
+        framework_profile_thresh = (
+            version.parse("2.0") if framework_name == "pytorch" else version.parse("2.12")
+        )
+        framework_profile = version.parse(framework_version)
+        if framework_profile >= framework_profile_thresh:
+            deprecation_warn_base(
+                f"Framework profiling is deprecated from\
+                 {framework_name} version {framework_version}.\
+                 No framework metrics will be collected"
+            )
+
+
+def validate_smdistributed(
+    instance_type, framework_name, framework_version, py_version, distribution, image_uri=None
+):
+    """Check if smdistributed strategy is correctly invoked by the user.
+
+    Currently, two strategies are supported: `dataparallel` or `modelparallel`.
+    Validate if the user requested strategy is supported.
+
+    Currently, only one strategy can be specified at a time. Validate if the user has requested
+    more than one strategy simultaneously.
+
+    Validate if the smdistributed dict arg is syntactically correct.
+
+    Additionally, perform strategy-specific validations.
+
+    Args:
+        instance_type (str): A string representing the type of training instance selected.
+        framework_name (str): A string representing the name of framework selected.
+        framework_version (str): A string representing the framework version selected.
+        py_version (str): A string representing the python version selected.
+        Ex: `py38, py39, py310, py311`
+        distribution (dict): A dictionary with information to enable distributed training.
+            (Defaults to None if distributed training is not enabled.) For example:
+
+            .. code:: python
+
+                {
+                    "smdistributed": {
+                        "dataparallel": {
+                            "enabled": True
+                        }
+                    }
+                }
+        image_uri (str): A string representing a Docker image URI.
+
+    Raises:
+        ValueError: if distribution dictionary isn't correctly formatted or
+            multiple strategies are requested simultaneously or
+            an unsupported strategy is requested or
+            strategy-specific inputs are incorrect/unsupported
+    """
+    if "smdistributed" not in distribution:
+        # Distribution strategy other than smdistributed is selected
+        return
+    if is_pipeline_variable(instance_type) or is_pipeline_variable(image_uri):
+        # The instance_type is not available in compile time.
+        # Rather, it's given in Pipeline execution time
+        return
+
+    # distribution contains smdistributed
+    smdistributed = distribution["smdistributed"]
+    if not isinstance(smdistributed, dict):
+        raise ValueError("smdistributed strategy requires a dictionary")
+
+    if len(smdistributed) > 1:
+        # more than 1 smdistributed strategy requested by the user
+        err_msg = (
+            "Cannot use more than 1 smdistributed strategy. \n"
+            "Choose one of the following supported strategies:"
+            f"{SMDISTRIBUTED_SUPPORTED_STRATEGIES}"
+        )
+        raise ValueError(err_msg)
+
+    # validate if smdistributed strategy is supported
+    # currently this for loop essentially checks for only 1 key
+    for strategy in smdistributed:
+        if strategy not in SMDISTRIBUTED_SUPPORTED_STRATEGIES:
+            err_msg = (
+                f"Invalid smdistributed strategy provided: {strategy} \n"
+                f"Supported strategies: {SMDISTRIBUTED_SUPPORTED_STRATEGIES}"
+            )
+            raise ValueError(err_msg)
+
+    # smdataparallel-specific input validation
+    if "dataparallel" in smdistributed:
+        _validate_smdataparallel_args(
+            instance_type, framework_name, framework_version, py_version, distribution, image_uri
+        )
+
+
+def _validate_smdataparallel_args(
+    instance_type, framework_name, framework_version, py_version, distribution, image_uri=None
+):
+    """Check if request is using unsupported arguments.
+
+    Validate if user specifies a supported instance type, framework version, and python
+    version.
+
+    Args:
+        instance_type (str): A string representing the type of training instance selected. Ex: `ml.p3.16xlarge`
+        framework_name (str): A string representing the name of framework selected. Ex: `tensorflow`
+        framework_version (str): A string representing the framework version selected. Ex: `2.3.1`
+        py_version (str): A string representing the python version selected.
+        Ex: `py38, py39, py310, py311`
+        distribution (dict): A dictionary with information to enable distributed training.
+            (Defaults to None if distributed training is not enabled.) Ex:
+
+            .. code:: python
+
+                {
+                    "smdistributed": {
+                        "dataparallel": {
+                            "enabled": True
+                        }
+                    }
+                }
+        image_uri (str): A string representing a Docker image URI.
+
+    Raises:
+        ValueError: if
+            `py_version` is not python3 or
+            `framework_version` is not in SM_DATAPARALLEL_SUPPORTED_FRAMEWORK_VERSION
+    """
+    smdataparallel_enabled = (
+        distribution.get("smdistributed").get("dataparallel").get("enabled", False)
+    )
+
+    if not smdataparallel_enabled:
+        return
+
+    err_msg = ""
+
+    if not instance_type:
+        err_msg += "Please specify an instance_type for smdataparallel.\n"
+
+    if not image_uri:
+        # ignore framework_version & py_version if image_uri is set
+        # in case image_uri is not set, then both are mandatory
+        supported = SM_DATAPARALLEL_SUPPORTED_FRAMEWORK_VERSIONS[framework_name]
+        if framework_version not in supported:
+            err_msg += (
+                f"Provided framework_version {framework_version} is not supported by"
+                " smdataparallel.\n"
+                f"Please specify one of the supported framework versions: {supported} \n"
+            )
+
+        if "py3" not in py_version:
+            err_msg += (
+                f"Provided py_version {py_version} is not supported by smdataparallel.\n"
+                "Please specify py_version>=py3"
+            )
+
+    if err_msg:
+        raise ValueError(err_msg)
+
+
+def validate_distribution(
+    distribution: Dict,
+    instance_groups: List[InstanceGroup],
+    framework_name: str,
+    framework_version: str,
+    py_version: str,
+    image_uri: str,
+    kwargs: Dict,
+) -> Dict:
+    """Check if distribution strategy is correctly invoked by the user.
+
+    Currently, check for `dataparallel`, `modelparallel` and heterogeneous cluster set up.
+    Validate if the user requested strategy is supported.
+
+    Args:
+        distribution (dict): A dictionary with information to enable distributed training.
+            (Defaults to None if distributed training is not enabled.) For example:
+
+            .. code:: python
+
+                {
+                    "smdistributed": {
+                        "dataparallel": {
+                            "enabled": True
+                        }
+                    }
+                }
+        instance_groups ([InstanceGroup]): A list contains instance groups used for training.
+        framework_name (str): A string representing the name of framework selected.
+        framework_version (str): A string representing the framework version selected.
+        py_version (str): A string representing the python version selected.
+        Ex: `py38, py39, py310, py311`
+        image_uri (str): A string representing a Docker image URI.
+        kwargs(dict): Additional kwargs passed to this function
+
+    Returns:
+        distribution(dict): updated dictionary with validated information
+            to enable distributed training.
+
+    Raises:
+        ValueError: if distribution dictionary isn't correctly formatted or
+            multiple strategies are requested simultaneously or
+            an unsupported strategy is requested or
+            strategy-specific inputs are incorrect/unsupported or
+            heterogeneous cluster set up is incorrect
+    """
+    validated_distribution = dict(distribution)
+
+    train_instance_groups = validated_distribution.get("instance_groups", [])
+    if instance_groups is None:
+        if len(train_instance_groups) >= 1:
+            # if estimator's instance_groups is not defined but
+            # train_instance_groups are specified in distribution
+            raise ValueError("Instance groups not specified in the estimator !")
+    else:
+        if len(train_instance_groups) > len(instance_groups):
+            # if train_instance_groups in distribution are more than estimator's instance_groups
+            raise ValueError("Train instance groups oversubscribed !")
+        if len(instance_groups) == 1 and len(train_instance_groups) == 0:
+            # if just one instance_group but it is not specified in distribution, we set it for user
+            train_instance_groups = instance_groups
+        elif len(instance_groups) > 1 and len(train_instance_groups) != 1:
+            # currently we just support one train instance group
+            raise ValueError("Distribution should only contain one instance group name !")
+
+    if len(train_instance_groups) != 0:
+        # in this case, we are handling a heterogeneous cluster training job
+        instance_group_names = []
+        for train_instance_group in train_instance_groups:
+            # in future version we will support multiple train_instance_groups, so use loop here
+            if train_instance_group not in instance_groups:
+                # check if train instance groups belongs to what user defined in estimator set up
+                raise ValueError(
+                    f"Invalid training instance group {train_instance_group.instance_group_name} !"
+                )
+            instance_type = train_instance_group.instance_type
+            validate_distribution_for_instance_type(
+                instance_type=instance_type,
+                distribution=validated_distribution,
+            )
+            validate_smdistributed(
+                instance_type=instance_type,
+                framework_name=framework_name,
+                framework_version=framework_version,
+                py_version=py_version,
+                distribution=validated_distribution,
+                image_uri=image_uri,
+            )
+            if framework_name and framework_name == "pytorch":
+                # We need to validate only for PyTorch framework
+                validate_torch_distributed_distribution(
+                    instance_type=instance_type,
+                    distribution=validated_distribution,
+                    framework_version=framework_version,
+                    py_version=py_version,
+                    image_uri=image_uri,
+                    entry_point=kwargs["entry_point"],
+                )
+            warn_if_parameter_server_with_multi_gpu(
+                training_instance_type=instance_type, distribution=validated_distribution
+            )
+            # get instance group names
+            instance_group_names.append(train_instance_group.instance_group_name)
+        validated_distribution["instance_groups"] = instance_group_names
+    else:
+        # in this case, we are handling a normal training job (without heterogeneous cluster)
+        instance_type = renamed_kwargs(
+            "train_instance_type", "instance_type", kwargs.get("instance_type"), kwargs
+        )
+        validate_distribution_for_instance_type(
+            instance_type=instance_type,
+            distribution=validated_distribution,
+        )
+        validate_smdistributed(
+            instance_type=instance_type,
+            framework_name=framework_name,
+            framework_version=framework_version,
+            py_version=py_version,
+            distribution=validated_distribution,
+            image_uri=image_uri,
+        )
+        if framework_name and framework_name == "pytorch":
+            # We need to validate only for PyTorch framework
+            validate_torch_distributed_distribution(
+                instance_type=instance_type,
+                distribution=validated_distribution,
+                framework_version=framework_version,
+                py_version=py_version,
+                image_uri=image_uri,
+                entry_point=kwargs["entry_point"],
+            )
+        warn_if_parameter_server_with_multi_gpu(
+            training_instance_type=instance_type, distribution=validated_distribution
+        )
+    return validated_distribution
+
+
+def validate_distribution_for_instance_type(instance_type, distribution):
+    """Check if the provided distribution strategy is supported for the instance_type.
+
+    Args:
+        instance_type (str): A string representing the type of training instance selected.
+        distribution (dict): A dictionary with information to enable distributed training.
+    """
+    err_msg = ""
+    if isinstance(instance_type, str):
+        match = re.match(r"^ml[\._]([a-z\d]+)\.?\w*$", instance_type)
+        if match and match[1].startswith("trn"):
+            keys = list(distribution.keys())
+            if len(keys) == 0:
+                return
+            if len(keys) == 1:
+                distribution_strategy = keys[0]
+                if distribution_strategy != "torch_distributed":
+                    err_msg += (
+                        f"Provided distribution strategy {distribution_strategy} is not supported"
+                        " for Trainium instances.\n"
+                        "Please specify one of the following supported distribution strategies:"
+                        f" {TRAINIUM_SUPPORTED_DISTRIBUTION_STRATEGIES} \n"
+                    )
+            elif len(keys) > 1:
+                err_msg += (
+                    "Multiple distribution strategies are not supported for Trainium instances.\n"
+                    "Please specify one of the following supported distribution strategies:"
+                    f" {TRAINIUM_SUPPORTED_DISTRIBUTION_STRATEGIES} "
+                )
+
+    if err_msg:
+        raise ValueError(err_msg)
+
+
+def validate_torch_distributed_distribution(
+    instance_type,
+    distribution,
+    framework_version,
+    py_version,
+    image_uri,
+    entry_point,
+):
+    """Check if torch_distributed distribution strategy is correctly invoked by the user.
+
+    Args:
+        instance_type (str): A string representing the type of training instance selected.
+        distribution (dict): A dictionary with information to enable distributed training.
+            (Defaults to None if distributed training is not enabled.) For example:
+
+            .. code:: python
+
+                {
+                    "torch_distributed": {
+                        "enabled": True
+                    }
+                }
+        framework_version (str): A string representing the framework version selected.
+        py_version (str): A string representing the python version selected.
+        Ex: `py38, py39, py310, py311`
+        image_uri (str): A string representing a Docker image URI.
+        entry_point (str or PipelineVariable): The absolute or relative path to the local Python
+            source file that should be executed as the entry point to
+            training.
+
+    Raises:
+        ValueError: if
+            `py_version` is not python3 or
+            `framework_version` is not compatible with instance types
+    """
+    torch_distributed_enabled = False
+    if "torch_distributed" in distribution:
+        torch_distributed_enabled = distribution.get("torch_distributed").get("enabled", False)
+    if not torch_distributed_enabled:
+        # Distribution strategy other than torch_distributed is selected
+        return
+
+    err_msg = ""
+
+    if not image_uri:
+        # ignore framework_version and py_version if image_uri is set
+        # in case image_uri is not set, then both are mandatory
+        if "py3" not in py_version:
+            err_msg += (
+                f"Provided py_version {py_version} is not supported by torch_distributed.\n"
+                "Please specify py_version>=py3\n"
+            )
+
+        # Check instance and framework_version compatibility
+        if _is_gpu_instance(instance_type):
+            if framework_version not in TORCH_DISTRIBUTED_GPU_SUPPORTED_FRAMEWORK_VERSIONS:
+                err_msg += (
+                    f"Provided framework_version {framework_version} is not supported by"
+                    f" torch_distributed for instance {instance_type}.\n"
+                    "Please specify one of the supported framework versions:"
+                    f"{TORCH_DISTRIBUTED_GPU_SUPPORTED_FRAMEWORK_VERSIONS} \n"
+                )
+        elif _is_trainium_instance(instance_type):
+            if framework_version not in TRAINIUM_SUPPORTED_TORCH_DISTRIBUTED_FRAMEWORK_VERSIONS:
+                err_msg += (
+                    f"Provided framework_version {framework_version} is not supported by"
+                    f" torch_distributed for instance {instance_type}.\n"
+                    "Please specify one of the supported framework versions:"
+                    f"{TRAINIUM_SUPPORTED_TORCH_DISTRIBUTED_FRAMEWORK_VERSIONS} \n"
+                )
+        else:
+            err_msg += (
+                "Currently torch_distributed is supported only for GPU and Trainium instances.\n"
+            )
+
+    # Check entry point type
+    if not entry_point.endswith(".py"):
+        err_msg += (
+            "Unsupported entry point type for the distribution torch_distributed.\n"
+            "Only python programs (*.py) are supported."
+        )
+
+    if err_msg:
+        raise ValueError(err_msg)
+
+
+def _is_gpu_instance(instance_type):
+    """Returns bool indicating whether instance_type supports GPU.
+
+    Args:
+        instance_type (str): Name of the instance_type to check against.
+
+    Returns:
+        bool: Whether or not the instance_type supports GPU
+    """
+    if isinstance(instance_type, str):
+        match = re.match(r"^ml[\._]([a-z\d]+)\.?\w*$", instance_type)
+        if match:
+            if match[1].startswith("p") or match[1].startswith("g"):
+                return True
+        if instance_type == "local_gpu":
+            return True
+    return False
+
+
+def _is_trainium_instance(instance_type):
+    """Returns bool indicating whether instance_type is a Trainium instance.
+
+    Args:
+        instance_type (str): Name of the instance_type to check against.
+
+    Returns:
+        bool: Whether or not the instance_type is a Trainium instance
+    """
+    if isinstance(instance_type, str):
+        match = re.match(r"^ml[\._]([a-z\d]+)\.?\w*$", instance_type)
+        if match and match[1].startswith("trn"):
+            return True
+    return False
+
+
+def python_deprecation_warning(framework, latest_supported_version):
+    """Placeholder docstring."""
+    return PYTHON_2_DEPRECATION_WARNING.format(
+        framework=framework, latest_supported_version=latest_supported_version
+    )
+
+
+def _region_supports_debugger(region_name):
+    """Returns boolean indicating whether the region supports Amazon SageMaker Debugger.
+
+    Args:
+        region_name (str): Name of the region to check against.
+
+    Returns:
+        bool: Whether or not the region supports Amazon SageMaker Debugger.
+    """
+    return region_name.lower() not in DEBUGGER_UNSUPPORTED_REGIONS
+
+
+def _region_supports_profiler(region_name):
+    """Returns bool indicating whether region supports Amazon SageMaker Debugger profiling feature.
+
+    Args:
+        region_name (str): Name of the region to check against.
+
+    Returns:
+        bool: Whether or not the region supports Amazon SageMaker Debugger profiling feature.
+    """
+    return region_name.lower() not in PROFILER_UNSUPPORTED_REGIONS
+
+
+def _instance_type_supports_profiler(instance_type):
+    """Returns bool indicating whether instance_type supports SageMaker Debugger profiling feature.
+
+    Args:
+        instance_type (str): Name of the instance_type to check against.
+
+    Returns:
+        bool: Whether or not the region supports Amazon SageMaker Debugger profiling feature.
+    """
+    if isinstance(instance_type, str):
+        match = re.match(r"^ml[\._]([a-z\d]+)\.?\w*$", instance_type)
+        if match and match[1].startswith("trn"):
+            return True
+    return False
+
+
+def validate_version_or_image_args(framework_version, py_version, image_uri):
+    """Checks if version or image arguments are specified.
+
+    Validates framework and model arguments to enforce version or image specification.
+
+    Args:
+        framework_version (str): The version of the framework.
+        py_version (str): A string representing the python version selected.
+        Ex: `py38, py39, py310, py311`
+        image_uri (str): The URI of the image.
+
+    Raises:
+        ValueError: if `image_uri` is None and either `framework_version` or `py_version` is
+            None.
+    """
+    if (framework_version is None or py_version is None) and image_uri is None:
+        raise ValueError(
+            "framework_version or py_version was None, yet image_uri was also None. "
+            "Either specify both framework_version and py_version, or specify image_uri."
+        )