Run TensorFlow Training Jobs with SageMaker Training Compiler

You can use any of the SageMaker interfaces to run a training job with SageMaker Training Compiler: Amazon SageMaker Studio Classic, Amazon SageMaker notebook instances, Amazon SDK for Python (Boto3), and Amazon Command Line Interface.

Using the SageMaker Python SDK

To turn on SageMaker Training Compiler, add the compiler_config parameter to the SageMaker TensorFlow or Hugging Face estimator. Import the TrainingCompilerConfig class and pass an instance of it to the compiler_config parameter. The following code examples show the structure of the SageMaker estimator classes with SageMaker Training Compiler turned on.

For information that fits your use case, see one of the following options.

The following list is the minimal set of parameters required to run a SageMaker training job with the compiler.

disable_profiler=True,
debugger_hook_config=False

If the training job with the compiler is launched successfully, you receive the following logs during the job initialization phase:

Using the SageMaker Python SDK and Extending SageMaker Framework Deep Learning Containers

Amazon Deep Learning Containers (DLC) for TensorFlow use adapted versions of TensorFlow that include changes on top of the open-source TensorFlow framework. The SageMaker Framework Deep Learning Containers are optimized for the underlying Amazon infrastructure and Amazon SageMaker. With the advantage of using the DLCs, SageMaker Training Compiler integration adds more performance improvements over the native TensorFlow. Furthermore, you can create a custom training container by extending the DLC image.

To extend and customize the SageMaker TensorFlow DLCs for your use-case, use the following instructions.

Create a Dockerfile

Use the following Dockerfile template to extend the SageMaker TensorFlow DLC. You must use the SageMaker TensorFlow DLC image as the base image of your Docker container. To find the SageMaker TensorFlow DLC image URIs, see Supported Frameworks.

# SageMaker TensorFlow Deep Learning Container image
FROM 763104351884.dkr.ecr.<aws-region>.amazonaws.com/tensorflow-training:<image-tag>

ENV PATH="/opt/ml/code:${PATH}"

# This environment variable is used by the SageMaker container 
# to determine user code directory.
ENV SAGEMAKER_SUBMIT_DIRECTORY /opt/ml/code

# Add more code lines to customize for your use-case
...

For more information, see Step 2: Create and upload the Dockerfile and Python training scripts.

Consider the following pitfalls when extending SageMaker Framework DLCs:

For example, there’s a known issue that the tensorflow/models and tensorflow/text libraries always attempt to reinstall open source TensorFlow. If you need to install these libraries to choose a specific version for your use case, we recommend that you look into the SageMaker TensorFlow DLC Dockerfiles for v2.9 or later. The paths to the Dockerfiles are typically in the following format: tensorflow/training/docker/<tensorflow-version>/py3/<cuda-version>/Dockerfile.gpu. In the Dockerfiles, you should find the code lines to reinstall Amazon managed TensorFlow binary (specified to the TF_URL environment variable) and other dependencies in order. The reinstallation section should look like the following example:

# tf-models does not respect existing installations of TensorFlow 
# and always installs open source TensorFlow

RUN pip3 install --no-cache-dir -U \
    tf-models-official==x.y.z

RUN pip3 uninstall -y tensorflow tensorflow-gpu \
  ; pip3 install --no-cache-dir -U \
    ${TF_URL} \
    tensorflow-io==x.y.z \
    tensorflow-datasets==x.y.z

Build and push to ECR

To build and push your Docker container to Amazon ECR, follow the instructions in the following links:

Run using the SageMaker Python SDK Estimator

Use the SageMaker TensorFlow framework estimator as usual. You must specify image_uri to use the new container you hosted in Amazon ECR.

import sagemaker, boto3
from sagemaker import get_execution_role
from sagemaker.tensorflow import TensorFlow, TrainingCompilerConfig

account_id = boto3.client('sts').get_caller_identity().get('Account')
ecr_repository = 'tf-custom-container-test'
tag = ':latest'

region = boto3.session.Session().region_name

uri_suffix = 'amazonaws.com'

byoc_image_uri = '{}.dkr.ecr.{}.{}/{}'.format(
    account_id, region, uri_suffix, ecr_repository + tag
)

byoc_image_uri
# This should return something like
# 111122223333.dkr.ecr.us-east-2.amazonaws.com/tf-custom-container-test:latest

estimator = TensorFlow(
    image_uri=image_uri,
    role=get_execution_role(),
    base_job_name='tf-custom-container-test-job',
    instance_count=1,
    instance_type='ml.p3.8xlarge'
    compiler_config=TrainingCompilerConfig(),
    disable_profiler=True,
    debugger_hook_config=False
)

# Start training
estimator.fit()

Enable SageMaker Training Compiler Using the SageMaker CreateTrainingJob API Operation

SageMaker Training Compiler configuration options must be specified through the AlgorithmSpecification and HyperParameters field in the request syntax for the CreateTrainingJob API operation.

"AlgorithmSpecification": {
    "TrainingImage": "<sagemaker-training-compiler-enabled-dlc-image>"
},

"HyperParameters": {
    "sagemaker_training_compiler_enabled": "true",
    "sagemaker_training_compiler_debug_mode": "false"
}

To find a complete list of deep learning container image URIs that have SageMaker Training Compiler implemented, see Supported Frameworks.