feature: Adding Training Compiler support for TensorFlow estimator

Author: Lokiiiiii

src/sagemaker/image_uris.py

@@ -134,21 +134,18 @@ def retrieve(
             tolerate_vulnerable_model,
             tolerate_deprecated_model,
         )
-    if training_compiler_config is None:
+
+    if training_compiler_config and (framework == HUGGING_FACE_FRAMEWORK):
+        config = _config_for_framework_and_scope(
+            framework + "-training-compiler", image_scope, accelerator_type
+        )
+    else:
         _framework = framework
         if framework == HUGGING_FACE_FRAMEWORK:
             inference_tool = _get_inference_tool(inference_tool, instance_type)
             if inference_tool == "neuron":
                 _framework = f"{framework}-{inference_tool}"
         config = _config_for_framework_and_scope(_framework, image_scope, accelerator_type)
-    elif framework == HUGGING_FACE_FRAMEWORK:
-        config = _config_for_framework_and_scope(
-            framework + "-training-compiler", image_scope, accelerator_type
-        )
-    else:
-        raise ValueError(
-            "Unsupported Configuration: Training Compiler is only supported with HuggingFace"
-        )
 
     original_version = version
     version = _validate_version_and_set_if_needed(version, config, framework)

src/sagemaker/tensorflow/__init__.py

@@ -16,3 +16,5 @@
 from sagemaker.tensorflow.estimator import TensorFlow  # noqa: F401 (imported but unused)
 from sagemaker.tensorflow.model import TensorFlowModel, TensorFlowPredictor  # noqa: F401
 from sagemaker.tensorflow.processing import TensorFlowProcessor  # noqa: F401
+
+from sagemaker.tensorflow.training_compiler.config import TrainingCompilerConfig  # noqa: F401

src/sagemaker/tensorflow/estimator.py

@@ -26,6 +26,7 @@
 from sagemaker.transformer import Transformer
 from sagemaker.vpc_utils import VPC_CONFIG_DEFAULT
 from sagemaker.workflow import is_pipeline_variable
+from sagemaker.tensorflow.training_compiler.config import TrainingCompilerConfig
 
 logger = logging.getLogger("sagemaker")
 
@@ -45,6 +46,7 @@ def __init__(
         model_dir=None,
         image_uri=None,
         distribution=None,
+        compiler_config=None,
         **kwargs
     ):
         """Initialize a ``TensorFlow`` estimator.
@@ -157,6 +159,8 @@ 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>`_.
+            compiler_config (:class:`~sagemaker.tensorflow.TrainingCompilerConfig`):
+                Configures SageMaker Training Compiler to accelerate training.
 
             **kwargs: Additional kwargs passed to the Framework constructor.
 
@@ -202,6 +206,16 @@ def __init__(
         self.distribution = distribution or {}
 
         self._validate_args(py_version=py_version)
+        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)
 
     def _validate_args(self, py_version):
         """Placeholder docstring"""

src/sagemaker/tensorflow/training_compiler/__init__.py

@@ -0,0 +1,105 @@
+# 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.
+"""Configuration for the SageMaker Training Compiler."""
+from __future__ import absolute_import
+import logging
+from packaging import version
+
+from sagemaker.training_compiler.config import TrainingCompilerConfig as BaseConfig
+
+logger = logging.getLogger(__name__)
+
+
+class TrainingCompilerConfig(BaseConfig):
+    """The SageMaker Training Compiler configuration class."""
+
+    SUPPORTED_INSTANCE_CLASS_PREFIXES = ["p3", "g4dn", "p4", "g5"]
+    MIN_SUPPORTED_VERSION = version.parse("2.9")
+
+    def __init__(
+        self,
+        enabled=True,
+        debug=False,
+    ):
+        """This class initializes a ``TrainingCompilerConfig`` instance.
+
+        `Amazon SageMaker Training Compiler
+        <https://docs.aws.amazon.com/sagemaker/latest/dg/training-compiler.html>`_
+        is a feature of SageMaker Training
+        and speeds up training jobs by optimizing model execution graphs.
+
+        You can compile TensorFlow models
+        by passing the object of this configuration class to the ``compiler_config``
+        parameter of the :class:`~sagemaker.tensorflow.TensorFlow`
+        estimator.
+
+        Args:
+            enabled (bool): Optional. Switch to enable SageMaker Training Compiler.
+                The default is ``True``.
+            debug (bool): Optional. Whether to dump detailed logs for debugging.
+                This comes with a potential performance slowdown.
+                The default is ``False``.
+
+        **Example**: The following code shows the basic usage of the
+        :class:`sagemaker.tensorflow.TrainingCompilerConfig()` class
+        to run a TensorFlow training job with the compiler.
+
+        .. code-block:: python
+
+            from sagemaker.tensorflow import TensorFlow, TrainingCompilerConfig
+
+            tensorflow_estimator=TensorFlow(
+                ...
+                compiler_config=TrainingCompilerConfig()
+            )
+
+        .. seealso::
+
+            For more information about how to enable SageMaker Training Compiler
+            for various training settings such as using TensorFlow-based models,
+            PyTorch-based models, and distributed training,
+            see `Enable SageMaker Training Compiler
+            <https://docs.aws.amazon.com/sagemaker/latest/dg/training-compiler-enable.html>`_
+            in the `Amazon SageMaker Training Compiler developer guide
+            <https://docs.aws.amazon.com/sagemaker/latest/dg/training-compiler.html>`_.
+
+        """
+
+        super(TrainingCompilerConfig, self).__init__(enabled=enabled, debug=debug)
+
+    @classmethod
+    def validate(
+        cls,
+        estimator,
+    ):
+        """Checks if SageMaker Training Compiler is configured correctly.
+
+        Args:
+            estimator (str): A estimator object
+                If SageMaker Training Compiler is enabled, it will validate whether
+                the estimator is configured to be compatible with Training Compiler.
+
+        Raises:
+            ValueError: Raised if the requested configuration is not compatible
+                        with SageMaker Training Compiler.
+        """
+
+        super(TrainingCompilerConfig, cls).validate(estimator)
+
+        if estimator.framework_version:
+            if version.parse(estimator.framework_version) < cls.MIN_SUPPORTED_VERSION:
+                error_helper_string = (
+                    f"SageMaker Training Compiler only supports TensorFlow version "
+                    f">= {cls.MIN_SUPPORTED_VERSION} but received {estimator.framework_version}"
+                )
+                raise ValueError(error_helper_string)