Merge branch 'master' into runproc-bug-fix

Author: Brock Wade

CHANGELOG.md

@@ -1,5 +1,18 @@
 # Changelog
 
+## v2.120.0 (2022-12-07)
+
+### Features
+
+ * Add Neo image uri config for Pytorch 1.12
+ * Adding support for SageMaker Training Compiler in PyTorch estimator starting 1.12
+ * Update registries with new region account number mappings.
+ * Add DXB region to frameworks by DLC
+
+### Bug Fixes and Other Changes
+
+ * support idempotency for framework and spark processors
+
 ## v2.119.0 (2022-12-03)
 
 ### Features

VERSION

@@ -1 +1 @@
-2.119.1.dev0
+2.120.1.dev0

src/sagemaker/fw_utils.py

@@ -493,7 +493,7 @@ def framework_name_from_image(image_uri):
     # We must support both the legacy and current image name format.
     name_pattern = re.compile(
         r"""^(?:sagemaker(?:-rl)?-)?
-        (tensorflow|mxnet|chainer|pytorch|scikit-learn|xgboost
+        (tensorflow|mxnet|chainer|pytorch|pytorch-trcomp|scikit-learn|xgboost
         |huggingface-tensorflow|huggingface-pytorch
         |huggingface-tensorflow-trcomp|huggingface-pytorch-trcomp)(?:-)?
         (scriptmode|training)?

src/sagemaker/image_uri_config/neo-pytorch.json

@@ -11,7 +11,9 @@
         "1.7.0": "1.7",
         "1.7.1": "1.7",
         "1.8.0": "1.8",
-        "1.8.1": "1.8"
+        "1.8.1": "1.8",
+        "1.12.0": "1.12",
+        "1.12.1": "1.12"
     },
     "versions": {
         "1.4": {
@@ -173,6 +175,38 @@
                 "us-west-2": "301217895009"
             },
             "repository": "sagemaker-inference-pytorch"
+        },
+        "1.12": {
+            "py_versions": ["py3"],
+            "registries": {
+                "af-south-1": "774647643957",
+                "ap-east-1": "110948597952",
+                "ap-northeast-1": "941853720454",
+                "ap-northeast-2": "151534178276",
+                "ap-northeast-3": "925152966179",
+                "ap-south-1": "763008648453",
+                "ap-southeast-1": "324986816169",
+                "ap-southeast-2": "355873309152",
+                "ca-central-1": "464438896020",
+                "cn-north-1": "472730292857",
+                "cn-northwest-1": "474822919863",
+                "eu-central-1": "746233611703",
+                "eu-north-1": "601324751636",
+                "eu-south-1": "966458181534",
+                "eu-west-1": "802834080501",
+                "eu-west-2": "205493899709",
+                "eu-west-3": "254080097072",
+                "me-south-1": "836785723513",
+                "sa-east-1": "756306329178",
+                "us-east-1": "785573368785",
+                "us-east-2": "007439368137",
+                "us-gov-west-1": "263933020539",
+                "us-iso-east-1": "167761179201",
+                "us-isob-east-1": "406031935815",
+                "us-west-1": "710691900526",
+                "us-west-2": "301217895009"
+            },
+            "repository": "sagemaker-inference-pytorch"
         }
     }
 }

src/sagemaker/image_uri_config/pytorch-training-compiler.json

@@ -0,0 +1,41 @@
+{
+    "training": {
+        "processors": [
+            "gpu"
+        ],
+        "version_aliases": {
+            "1.12": "1.12.0"
+        },
+        "versions": {
+            "1.12.0": {
+                "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",
+                    "ca-central-1": "763104351884",
+                    "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-west-1": "763104351884",
+                    "us-west-2": "763104351884"
+                },
+                "repository": "pytorch-training"
+            }
+        }
+    }
+}

src/sagemaker/image_uris.py

@@ -146,7 +146,7 @@ def retrieve(
             tolerate_deprecated_model,
         )
 
-    if training_compiler_config and (framework == HUGGING_FACE_FRAMEWORK):
+    if training_compiler_config and (framework in [HUGGING_FACE_FRAMEWORK, "pytorch"]):
         final_image_scope = image_scope
         config = _config_for_framework_and_scope(
             framework + "-training-compiler", final_image_scope, accelerator_type

src/sagemaker/pytorch/__init__.py

@@ -16,3 +16,5 @@
 from sagemaker.pytorch.estimator import PyTorch  # noqa: F401
 from sagemaker.pytorch.model import PyTorchModel, PyTorchPredictor  # noqa: F401
 from sagemaker.pytorch.processing import PyTorchProcessor  # noqa: F401
+
+from sagemaker.pytorch.training_compiler.config import TrainingCompilerConfig  # noqa: F401

src/sagemaker/pytorch/estimator.py

@@ -28,6 +28,7 @@
 )
 from sagemaker.pytorch import defaults
 from sagemaker.pytorch.model import PyTorchModel
+from sagemaker.pytorch.training_compiler.config import TrainingCompilerConfig
 from sagemaker.vpc_utils import VPC_CONFIG_DEFAULT
 from sagemaker.workflow.entities import PipelineVariable
 
@@ -51,7 +52,8 @@ def __init__(
         hyperparameters: Optional[Dict[str, Union[str, PipelineVariable]]] = None,
         image_uri: Optional[Union[str, PipelineVariable]] = None,
         distribution: Optional[Dict] = None,
-        **kwargs
+        compiler_config: Optional[TrainingCompilerConfig] = None,
+        **kwargs,
     ):
         """This ``Estimator`` executes a PyTorch script in a managed PyTorch execution environment.
 
@@ -208,6 +210,31 @@ def __init__(
                     To learn more, see `Training with parameter servers
                     <https://sagemaker.readthedocs.io/en/stable/frameworks/tensorflow/using_tf.html#training-with-parameter-servers>`_.
 
+                **To enable distributed training with
+                `SageMaker Training Compiler <https://docs.aws.amazon.com/sagemaker/latest/dg/training-compiler.html>`_
+                for PyTorch:**
+
+                    .. code:: python
+
+                        {
+                            "pytorchxla": {
+                                "enabled": True
+                            }
+                        }
+
+                    To learn more, see `SageMaker Training Compiler
+                    <https://docs.aws.amazon.com/sagemaker/latest/dg/training-compiler.html>`_
+                    in the *Amazon SageMaker Developer Guide*.
+
+                    .. note::
+
+                        When you use this PyTorch XLA option for distributed training strategy,
+                        you must add the ``compiler_config`` parameter and activate SageMaker
+                        Training Compiler.
+
+                compiler_config (:class:`~sagemaker.pytorch.TrainingCompilerConfig`):
+                Configures SageMaker Training Compiler to accelerate training.
+
             **kwargs: Additional kwargs passed to the :class:`~sagemaker.estimator.Framework`
                 constructor.
 
@@ -250,6 +277,25 @@ def __init__(
 
         self.distribution = distribution or {}
 
+        if compiler_config is not None:
+            if not isinstance(compiler_config, TrainingCompilerConfig):
+                error_string = (
+                    f"Expected instance of type {TrainingCompilerConfig}"
+                    f"for argument compiler_config. "
+                    f"Instead got {type(compiler_config)}"
+                )
+                raise ValueError(error_string)
+            if compiler_config:
+                compiler_config.validate(self)
+        elif distribution is not None and "pytorchxla" in distribution:
+            raise ValueError(
+                "Distributed training through PyTorch XLA is currently only supported "
+                "when SageMaker Training Compiler is enabled. To learn more, "
+                "see Enable SageMaker Training Compiler at "
+                "https://docs.aws.amazon.com/sagemaker/latest/dg/training-compiler-enable.html."
+            )
+        self.compiler_config = compiler_config
+
     def _pytorch_distribution_configuration(self, distribution):
         """Returns a dict of distribution config for PyTorch training
 
@@ -289,6 +335,12 @@ def hyperparameters(self):
         hyperparameters.update(
             EstimatorBase._json_encode_hyperparameters(additional_hyperparameters)
         )
+        if self.compiler_config:
+            training_compiler_hyperparameters = self.compiler_config._to_hyperparameter_dict()
+            hyperparameters.update(
+                EstimatorBase._json_encode_hyperparameters(training_compiler_hyperparameters)
+            )
+
         return hyperparameters
 
     def create_model(
@@ -299,7 +351,7 @@ def create_model(
         entry_point=None,
         source_dir=None,
         dependencies=None,
-        **kwargs
+        **kwargs,
     ):
         """Create a SageMaker ``PyTorchModel`` object that can be deployed to an ``Endpoint``.
 
@@ -350,7 +402,7 @@ def create_model(
             sagemaker_session=self.sagemaker_session,
             vpc_config=self.get_vpc_config(vpc_config_override),
             dependencies=(dependencies or self.dependencies),
-            **kwargs
+            **kwargs,
         )
 
     @classmethod
@@ -371,6 +423,8 @@ def _prepare_init_params_from_job_description(cls, job_details, model_channel_na
         )
         image_uri = init_params.pop("image_uri")
         framework, py_version, tag, _ = framework_name_from_image(image_uri)
+        if framework:
+            framework = framework.split("-")[0]
 
         if tag is None:
             framework_version = None