class_docstr: Handle training of custom HuggingFace code.
class_docstr_tok: ["Handle","training","of","custom","HuggingFace","code","."]
docstr: This estimator runs a Hugging Face training script in a SageMaker training environment.
  
  The estimator initiates the SageMaker-managed Hugging Face environment
  by using the pre-built Hugging Face Docker container and runs
  the Hugging Face training script that user provides through
  the ``entry_point`` argument.
  
  After configuring the estimator class, use the class method
  :meth:`~sagemaker.amazon.estimator.Framework.fit()` to start a training job.
  
  Args:
      py_version (str): Python version you want to use for executing your model training
          code. Defaults to ``None``. Required unless ``image_uri`` is provided.  If
          using PyTorch, the current supported version is ``py36``. If using TensorFlow,
          the current supported version is ``py37``.
      entry_point (str or PipelineVariable): Path (absolute or relative) to the Python source
          file which should be executed as the entry point to training.
          If ``source_dir`` is specified, then ``entry_point``
          must point to a file located at the root of ``source_dir``.
      transformers_version (str): Transformers version you want to use for
          executing your model training code. Defaults to ``None``. Required unless
          ``image_uri`` is provided. The current supported version is ``4.6.1``.
      tensorflow_version (str): TensorFlow version you want to use for
          executing your model training code. Defaults to ``None``. Required unless
          ``pytorch_version`` is provided. The current supported version is ``2.4.1``.
      pytorch_version (str): PyTorch version you want to use for
          executing your model training code. Defaults to ``None``. Required unless
          ``tensorflow_version`` is provided. The current supported versions are ``1.7.1`` and ``1.6.0``.
      source_dir (str or PipelineVariable): Path (absolute, relative or an S3 URI) to a
          directory with any other training source code dependencies aside from the entry
          point file (default: None). If ``source_dir`` is an S3 URI, it must
          point to a tar.gz file. Structure within this directory are preserved
          when training on Amazon SageMaker.
      hyperparameters (dict[str, str] or dict[str, PipelineVariable]): Hyperparameters
          that will be used for training (default: None). The hyperparameters are made
          accessible as a dict[str, str] to the training code on
          SageMaker. For convenience, this accepts other types for keys
          and values, but ``str()`` will be called to convert them before
          training.
      image_uri (str or PipelineVariable): If specified, the estimator will use this image
          for training and hosting, instead of selecting the appropriate
          SageMaker official image based on framework_version and
          py_version. It can be an ECR url or dockerhub image and tag.
          Examples:
              * ``123412341234.dkr.ecr.us-west-2.amazonaws.com/my-custom-image:1.0``
              * ``custom-image:latest``
  
          If ``framework_version`` or ``py_version`` are ``None``, then
          ``image_uri`` is required. If also ``None``, then a ``ValueError``
          will be raised.
      distribution (dict): A dictionary with information on how to run distributed training
          (default: None).  Currently, the following are supported:
          distributed training with parameter servers, SageMaker Distributed (SMD) Data
          and Model Parallelism, and MPI. SMD Model Parallelism can only be used with MPI.
          To enable parameter server use the following setup:
  
          .. code:: python
  
              {
                  "parameter_server": {
                      "enabled": True
                  }
              }
  
          To enable MPI:
  
          .. code:: python
  
              {
                  "mpi": {
                      "enabled": True
                  }
              }
  
          To enable SMDistributed Data Parallel or Model Parallel:
  
          .. code:: python
  
              {
                  "smdistributed": {
                      "dataparallel": {
                          "enabled": True
                      },
                      "modelparallel": {
                          "enabled": True,
                          "parameters": {}
                      }
                  }
              }
  
          **To enable PyTorch DDP:**
  
              .. code:: python
  
                  {
                      "pytorchddp": {
                          "enabled": True
                      }
                  }
  
              To learn more, see `Distributed PyTorch Training
              <https://sagemaker.readthedocs.io/en/stable/frameworks/pytorch/using_pytorch.html#distributed-pytorch-training>`_.
  
          **To enable Torch Distributed:**
  
              This is available for general distributed training on
              GPU instances from PyTorch v1.13.1 and later.
  
              .. code:: python
  
                  {
                      "torch_distributed": {
                          "enabled": True
                      }
                  }
  
              This option also supports distributed training on Trn1.
              To learn more, see `Distributed PyTorch Training on Trainium
              <https://sagemaker.readthedocs.io/en/stable/frameworks/pytorch/using_pytorch.html#distributed-pytorch-training-on-trainium>`_.
  
          To enable distributed training with
          `SageMaker Training Compiler <https://docs.aws.amazon.com/sagemaker/latest/dg/training-compiler.html>`_
          for Hugging Face Transformers with 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.huggingface.TrainingCompilerConfig`):
          Configures SageMaker Training Compiler to accelerate training.
  
      **kwargs: Additional kwargs passed to the :class:`~sagemaker.estimator.Framework`
          constructor.
  
  .. tip::
  
      You can find additional parameters for initializing this class at
      :class:`~sagemaker.estimator.Framework` and
      :class:`~sagemaker.estimator.EstimatorBase`.
docstr_tok: ["This","estimator","runs","a","Hugging","Face","training","script","in","a","SageMaker","training","environment",".","The","estimator","initiates","the","SageMaker-managed","Hugging","Face","environment","by","using","the","pre-built","Hugging","Face","Docker","container","and","runs","the","Hugging","Face","training","script","that","user","provides","through","the","``","entry_point","``","argument",".","After","configuring","the","estimator","class",",","use","the","class","method",":","meth",":","`","~sagemaker.amazon.estimator.Framework.fit","(",")","`","to","start","a","training","job",".","Args",":","py_version","(","str",")",":","Python","version","you","want","to","use","for","executing","your","model","training","code",".","Defaults","to","``","None","``",".","Required","unless","``","image_uri","``","is","provided",".","If","using","PyTorch",",","the","current","supported","version","is","``","py36","``",".","If","using","TensorFlow",",","the","current","supported","version","is","``","py37","``",".","entry_point","(","str","or","PipelineVariable",")",":","Path","(","absolute","or","relative",")","to","the","Python","source","file","which","should","be","executed","as","the","entry","point","to","training",".","If","``","source_dir","``","is","specified",",","then","``","entry_point","``","must","point","to","a","file","located","at","the","root","of","``","source_dir","``",".","transformers_version","(","str",")",":","Transformers","version","you","want","to","use","for","executing","your","model","training","code",".","Defaults","to","``","None","``",".","Required","unless","``","image_uri","``","is","provided",".","The","current","supported","version","is","``","4.6.1","``",".","tensorflow_version","(","str",")",":","TensorFlow","version","you","want","to","use","for","executing","your","model","training","code",".","Defaults","to","``","None","``",".","Required","unless","``","pytorch_version","``","is","provided",".","The","current","supported","version","is","``","2.4.1","``",".","pytorch_version","(","str",")",":","PyTorch","version","you","want","to","use","for","executing","your","model","training","code",".","Defaults","to","``","None","``",".","Required","unless","``","tensorflow_version","``","is","provided",".","The","current","supported","versions","are","``","1.7.1","``","and","``","1.6.0","``",".","source_dir","(","str","or","PipelineVariable",")",":","Path","(","absolute",",","relative","or","an","S3","URI",")","to","a","directory","with","any","other","training","source","code","dependencies","aside","from","the","entry","point","file","(","default",":","None",")",".","If","``","source_dir","``","is","an","S3","URI",",","it","must","point","to","a","tar.gz","file",".","Structure","within","this","directory","are","preserved","when","training","on","Amazon","SageMaker",".","hyperparameters","(","dict","[","str",",","str","]","or","dict","[","str",",","PipelineVariable","]",")",":","Hyperparameters","that","will","be","used","for","training","(","default",":","None",")",".","The","hyperparameters","are","made","accessible","as","a","dict","[","str",",","str","]","to","the","training","code","on","SageMaker",".","For","convenience",",","this","accepts","other","types","for","keys","and","values",",","but","``","str","(",")","``","will","be","called","to","convert","them","before","training",".","image_uri","(","str","or","PipelineVariable",")",":","If","specified",",","the","estimator","will","use","this","image","for","training","and","hosting",",","instead","of","selecting","the","appropriate","SageMaker","official","image","based","on","framework_version","and","py_version",".","It","can","be","an","ECR","url","or","dockerhub","image","and","tag",".","Examples",":","*","``","123412341234.dkr.ecr.us-west-2.amazonaws.com\/my-custom-image:1.0","``","*","``","custom-image",":","latest","``","If","``","framework_version","``","or","``","py_version","``","are","``","None","``",",","then","``","image_uri","``","is","required",".","If","also","``","None","``",",","then","a","``","ValueError","``","will","be","raised",".","distribution","(","dict",")",":","A","dictionary","with","information","on","how","to","run","distributed","training","(","default",":","None",")",".","Currently",",","the","following","are","supported",":","distributed","training","with","parameter","servers",",","SageMaker","Distributed","(","SMD",")","Data","and","Model","Parallelism",",","and","MPI",".","SMD","Model","Parallelism","can","only","be","used","with","MPI",".","To","enable","parameter","server","use","the","following","setup",":","..","code",":",":","python","{","``","parameter_server","''",":","{","``","enabled","''",":","True","}","}","To","enable","MPI",":","..","code",":",":","python","{","``","mpi","''",":","{","``","enabled","''",":","True","}","}","To","enable","SMDistributed","Data","Parallel","or","Model","Parallel",":","..","code",":",":","python","{","``","smdistributed","''",":","{","``","dataparallel","''",":","{","``","enabled","''",":","True","}",",","``","modelparallel","''",":","{","``","enabled","''",":","True",",","``","parameters","''",":","{","}","}","}","}","*","*","To","enable","PyTorch","DDP",":","*","*","..","code",":",":","python","{","``","pytorchddp","''",":","{","``","enabled","''",":","True","}","}","To","learn","more",",","see","`","Distributed","PyTorch","Training","<","https",":","\/\/sagemaker.readthedocs.io\/en\/stable\/frameworks\/pytorch\/using_pytorch.html","#","distributed-pytorch-training",">","`","_",".","*","*","To","enable","Torch","Distributed",":","*","*","This","is","available","for","general","distributed","training","on","GPU","instances","from","PyTorch","v1.13.1","and","later",".","..","code",":",":","python","{","``","torch_distributed","''",":","{","``","enabled","''",":","True","}","}","This","option","also","supports","distributed","training","on","Trn1",".","To","learn","more",",","see","`","Distributed","PyTorch","Training","on","Trainium","<","https",":","\/\/sagemaker.readthedocs.io\/en\/stable\/frameworks\/pytorch\/using_pytorch.html","#","distributed-pytorch-training-on-trainium",">","`","_",".","To","enable","distributed","training","with","`","SageMaker","Training","Compiler","<","https",":","\/\/docs.aws.amazon.com\/sagemaker\/latest\/dg\/training-compiler.html",">","`","_","for","Hugging","Face","Transformers","with","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.huggingface.TrainingCompilerConfig","`",")",":","Configures","SageMaker","Training","Compiler","to","accelerate","training",".","*","*","kwargs",":","Additional","kwargs","passed","to","the",":","class",":","`","~sagemaker.estimator.Framework","`","constructor",".","..","tip",":",":","You","can","find","additional","parameters","for","initializing","this","class","at",":","class",":","`","~sagemaker.estimator.Framework","`","and",":","class",":","`","~sagemaker.estimator.EstimatorBase","`","."]
code: def __init__(
      self,
      py_version: str,
      entry_point: Union[str, PipelineVariable],
      transformers_version: Optional[str] = None,
      tensorflow_version: Optional[str] = None,
      pytorch_version: Optional[str] = None,
      source_dir: Optional[Union[str, PipelineVariable]] = None,
      hyperparameters: Optional[
          Dict[str, Union[str, PipelineVariable]]
      ] = None,
      image_uri: Optional[Union[str, PipelineVariable]] = None,
      distribution: Optional[Dict] = None,
      compiler_config: Optional[TrainingCompilerConfig] = None,
      **kwargs,
  ):
      """This estimator runs a Hugging Face training script in a SageMaker training environment.
  
      The estimator initiates the SageMaker-managed Hugging Face environment
      by using the pre-built Hugging Face Docker container and runs
      the Hugging Face training script that user provides through
      the ``entry_point`` argument.
  
      After configuring the estimator class, use the class method
      :meth:`~sagemaker.amazon.estimator.Framework.fit()` to start a training job.
  
      Args:
          py_version (str): Python version you want to use for executing your model training
              code. Defaults to ``None``. Required unless ``image_uri`` is provided.  If
              using PyTorch, the current supported version is ``py36``. If using TensorFlow,
              the current supported version is ``py37``.
          entry_point (str or PipelineVariable): Path (absolute or relative) to the Python source
              file which should be executed as the entry point to training.
              If ``source_dir`` is specified, then ``entry_point``
              must point to a file located at the root of ``source_dir``.
          transformers_version (str): Transformers version you want to use for
              executing your model training code. Defaults to ``None``. Required unless
              ``image_uri`` is provided. The current supported version is ``4.6.1``.
          tensorflow_version (str): TensorFlow version you want to use for
              executing your model training code. Defaults to ``None``. Required unless
              ``pytorch_version`` is provided. The current supported version is ``2.4.1``.
          pytorch_version (str): PyTorch version you want to use for
              executing your model training code. Defaults to ``None``. Required unless
              ``tensorflow_version`` is provided. The current supported versions are ``1.7.1`` and ``1.6.0``.
          source_dir (str or PipelineVariable): Path (absolute, relative or an S3 URI) to a
              directory with any other training source code dependencies aside from the entry
              point file (default: None). If ``source_dir`` is an S3 URI, it must
              point to a tar.gz file. Structure within this directory are preserved
              when training on Amazon SageMaker.
          hyperparameters (dict[str, str] or dict[str, PipelineVariable]): Hyperparameters
              that will be used for training (default: None). The hyperparameters are made
              accessible as a dict[str, str] to the training code on
              SageMaker. For convenience, this accepts other types for keys
              and values, but ``str()`` will be called to convert them before
              training.
          image_uri (str or PipelineVariable): If specified, the estimator will use this image
              for training and hosting, instead of selecting the appropriate
              SageMaker official image based on framework_version and
              py_version. It can be an ECR url or dockerhub image and tag.
              Examples:
                  * ``123412341234.dkr.ecr.us-west-2.amazonaws.com/my-custom-image:1.0``
                  * ``custom-image:latest``
  
              If ``framework_version`` or ``py_version`` are ``None``, then
              ``image_uri`` is required. If also ``None``, then a ``ValueError``
              will be raised.
          distribution (dict): A dictionary with information on how to run distributed training
              (default: None).  Currently, the following are supported:
              distributed training with parameter servers, SageMaker Distributed (SMD) Data
              and Model Parallelism, and MPI. SMD Model Parallelism can only be used with MPI.
              To enable parameter server use the following setup:
  
              .. code:: python
  
                  {
                      "parameter_server": {
                          "enabled": True
                      }
                  }
  
              To enable MPI:
  
              .. code:: python
  
                  {
                      "mpi": {
                          "enabled": True
                      }
                  }
  
              To enable SMDistributed Data Parallel or Model Parallel:
  
              .. code:: python
  
                  {
                      "smdistributed": {
                          "dataparallel": {
                              "enabled": True
                          },
                          "modelparallel": {
                              "enabled": True,
                              "parameters": {}
                          }
                      }
                  }
  
              **To enable PyTorch DDP:**
  
                  .. code:: python
  
                      {
                          "pytorchddp": {
                              "enabled": True
                          }
                      }
  
                  To learn more, see `Distributed PyTorch Training
                  <https://sagemaker.readthedocs.io/en/stable/frameworks/pytorch/using_pytorch.html#distributed-pytorch-training>`_.
  
              **To enable Torch Distributed:**
  
                  This is available for general distributed training on
                  GPU instances from PyTorch v1.13.1 and later.
  
                  .. code:: python
  
                      {
                          "torch_distributed": {
                              "enabled": True
                          }
                      }
  
                  This option also supports distributed training on Trn1.
                  To learn more, see `Distributed PyTorch Training on Trainium
                  <https://sagemaker.readthedocs.io/en/stable/frameworks/pytorch/using_pytorch.html#distributed-pytorch-training-on-trainium>`_.
  
              To enable distributed training with
              `SageMaker Training Compiler <https://docs.aws.amazon.com/sagemaker/latest/dg/training-compiler.html>`_
              for Hugging Face Transformers with 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.huggingface.TrainingCompilerConfig`):
              Configures SageMaker Training Compiler to accelerate training.
  
          **kwargs: Additional kwargs passed to the :class:`~sagemaker.estimator.Framework`
              constructor.
  
      .. tip::
  
          You can find additional parameters for initializing this class at
          :class:`~sagemaker.estimator.Framework` and
          :class:`~sagemaker.estimator.EstimatorBase`.
      """
      self.framework_version = transformers_version
      self.py_version = py_version
      self.tensorflow_version = tensorflow_version
      self.pytorch_version = pytorch_version
  
      self._validate_args(image_uri=image_uri)
  
      if "enable_sagemaker_metrics" not in kwargs:
          kwargs["enable_sagemaker_metrics"] = True
  
      kwargs["py_version"] = self.py_version
  
      super(HuggingFace, self).__init__(
          entry_point,
          source_dir,
          hyperparameters,
          image_uri=image_uri,
          **kwargs,
      )
  
      if "entry_point" not in kwargs:
          kwargs["entry_point"] = entry_point
  
      self.base_framework_name = (
          "tensorflow" if tensorflow_version is not None else "pytorch"
      )
      self.base_framework_version = (
          tensorflow_version
          if tensorflow_version is not None
          else pytorch_version
      )
  
      if distribution is not None:
          distribution = validate_distribution(
              distribution,
              self.instance_groups,
              self.base_framework_name,
              self.base_framework_version,
              py_version,
              image_uri,
              kwargs,
          )
  
      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
  
file_path: sagemaker-python-sdk/src/sagemaker/huggingface/estimator.py
imports: from __future__ import absolute_import
  import logging
  import re
  from typing import Optional, Union, Dict
  from sagemaker.estimator import Framework, EstimatorBase
  from sagemaker.fw_utils import framework_name_from_image, validate_distribution
  from sagemaker.huggingface.model import HuggingFaceModel
  from sagemaker.vpc_utils import VPC_CONFIG_DEFAULT
  from sagemaker.huggingface.training_compiler.config import TrainingCompilerConfig
  from sagemaker.workflow.entities import PipelineVariable
label_desc: Use image node_id 1 to create a new HuggingFace object from inherited base classes: Framework with example: obj = HuggingFace(py_version, entry_point, transformers_version, tensorflow_version, pytorch_version, source_dir, hyperparameters, image_uri, distribution, compiler_config)