Run TensorFlow Training Jobs with SageMaker Training Compiler

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

Using the SageMaker Python SDK

To turn on SageMaker Training Compiler, add the compiler_config parameter to the SageMaker AI 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 AI 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 AI Python SDK and Extending SageMaker AI Framework Deep Learning Containers

AWS Deep Learning Containers (DLC) for TensorFlow use adapted versions of TensorFlow that include changes on top of the open-source TensorFlow framework. The SageMaker AI Framework Deep Learning Containers are optimized for the underlying AWS infrastructure and Amazon SageMaker AI. 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 AI TensorFlow DLCs for your use-case, use the following instructions.

Create a Dockerfile

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

# SageMaker AI 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 AI 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 AI 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 AI 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 AWS 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 AI 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 AI 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.