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

sagemaker/core/training_compiler_config.py

@@ -0,0 +1,197 @@
+# 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 sagemaker.core.workflow import is_pipeline_variable
+
+logger = logging.getLogger(__name__)
+
+
+class TrainingCompilerConfig(object):
+    """The SageMaker Training Compiler configuration class."""
+
+    DEBUG_PATH = "/opt/ml/output/data/compiler/"
+    SUPPORTED_INSTANCE_CLASS_PREFIXES = ["p3", "p3dn", "g4dn", "p4d", "g5"]
+
+    HP_ENABLE_COMPILER = "sagemaker_training_compiler_enabled"
+    HP_ENABLE_DEBUG = "sagemaker_training_compiler_debug_mode"
+
+    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 Hugging Face models
+        by passing the object of this configuration class to the ``compiler_config``
+        parameter of the :class:`~sagemaker.huggingface.HuggingFace`
+        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.huggingface.TrainingCompilerConfig()` class
+        to run a HuggingFace training job with the compiler.
+
+        .. code-block:: python
+
+            from sagemaker.huggingface import HuggingFace, TrainingCompilerConfig
+
+            huggingface_estimator=HuggingFace(
+                ...
+                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>`_.
+
+        """
+
+        self.enabled = enabled
+        self.debug = debug
+
+        self.disclaimers_and_warnings()
+
+    def __nonzero__(self):
+        """Evaluates to 0 if SM Training Compiler is disabled."""
+        return self.enabled
+
+    def disclaimers_and_warnings(self):
+        """Disclaimers and warnings.
+
+        Logs disclaimers and warnings about the
+        requested configuration of SageMaker Training Compiler.
+
+        """
+
+        if self.enabled and self.debug:
+            logger.warning(
+                "Debugging is enabled."
+                "This will dump detailed logs from compilation to %s"
+                "This might impair training performance.",
+                self.DEBUG_PATH,
+            )
+
+    def _to_hyperparameter_dict(self):
+        """Converts configuration object into hyperparameters.
+
+        Returns:
+            dict: A portion of the hyperparameters passed to the training job as a dictionary.
+
+        """
+
+        compiler_config_hyperparameters = {
+            self.HP_ENABLE_COMPILER: self.enabled,
+            self.HP_ENABLE_DEBUG: self.debug,
+        }
+
+        return compiler_config_hyperparameters
+
+    @classmethod
+    def validate(cls, estimator):
+        """Checks if SageMaker Training Compiler is configured correctly.
+
+        Args:
+            estimator (:class:`sagemaker.estimator.Estimator`): An estimator object.
+                When SageMaker Training Compiler is enabled, it validates if
+                the estimator is configured to be compatible with Training Compiler.
+
+
+        Raises:
+            ValueError: Raised if the requested configuration is not compatible
+                        with SageMaker Training Compiler.
+        """
+        if is_pipeline_variable(estimator.instance_type):
+            warn_msg = (
+                "Estimator instance_type is a PipelineVariable (%s), "
+                "which has to be interpreted as one of the "
+                "%s classes in execution time."
+            )
+            logger.warning(
+                warn_msg,
+                type(estimator.instance_type),
+                str(cls.SUPPORTED_INSTANCE_CLASS_PREFIXES).replace(",", ""),
+            )
+        elif estimator.instance_type:
+            if "local" not in estimator.instance_type:
+                requested_instance_class = estimator.instance_type.split(".")[
+                    1
+                ]  # Expecting ml.class.size
+                if not any(
+                    [requested_instance_class == i for i in cls.SUPPORTED_INSTANCE_CLASS_PREFIXES]
+                ):
+                    error_helper_string = (
+                        "Unsupported Instance class {}."
+                        "SageMaker Training Compiler only supports {}"
+                    )
+                    error_helper_string = error_helper_string.format(
+                        requested_instance_class, cls.SUPPORTED_INSTANCE_CLASS_PREFIXES
+                    )
+                    raise ValueError(error_helper_string)
+            elif estimator.instance_type == "local":
+                error_helper_string = (
+                    "SageMaker Training Compiler doesn't support local mode."
+                    "It only supports the following GPU instances: {}"
+                )
+                error_helper_string = error_helper_string.format(
+                    cls.SUPPORTED_INSTANCE_CLASS_PREFIXES
+                )
+                raise ValueError(error_helper_string)
+
+        if estimator.distribution and "smdistributed" in estimator.distribution:
+            raise ValueError(
+                "SageMaker distributed training configuration is currently not compatible with "
+                "SageMaker Training Compiler."
+            )
+
+        if estimator.debugger_hook_config or (not estimator.disable_profiler):
+            helper_string = (
+                "Using Debugger and/or Profiler with SageMaker Training Compiler "
+                "might add recompilation overhead and degrade"
+                "performance. Found debugger_hook_config={} "
+                "disable_profiler={}. Please set "
+                "debugger_hook_config=None and disable_profiler=True for optimal "
+                "performance. For more information, see Training Compiler "
+                "Performance Considerations "
+                "(https://docs.aws.amazon.com/sagemaker/latest/dg/training-compiler-tips-pitfalls.html"
+                "#training-compiler-tips-pitfalls-considerations)."
+            )
+            helper_string = helper_string.format(
+                estimator.debugger_hook_config, estimator.disable_profiler
+            )
+            logger.warning(helper_string)
+
+        if estimator.instance_groups:
+            raise ValueError(
+                "SageMaker Training Compiler currently only supports homogeneous clusters of "
+                "the following GPU instance families: {}. Please use the 'instance_type' "
+                "and 'instance_count' parameters instead of 'instance_groups'".format(
+                    cls.SUPPORTED_INSTANCE_CLASS_PREFIXES
+                )
+            )