SageMaker TensorFlow Containers

SageMaker TensorFlow Containers is an open source library for making the TensorFlow framework run on Amazon SageMaker.

This repository also contains Dockerfiles which install this library, TensorFlow, and dependencies for building SageMaker TensorFlow images.

For information on running TensorFlow jobs on SageMaker: Python SDK.

For notebook examples: SageMaker Notebook Examples.

Table of Contents

1. Getting Started
2. Building your Image
3. Running the tests
4. The sagemaker_tensorflow_containers Python package

Getting Started

Prerequisites

Make sure you have installed all of the following prerequisites on your development machine:

• Docker

For Testing on GPU

• Nvidia-Docker

Recommended

• A Python environment management tool. (e.g. PyEnv, VirtualEnv)

Building your Image

Amazon SageMaker utilizes Docker containers to run all training jobs & inference endpoints.

The Docker images are built from the Dockerfiles specified in Docker/.

The Docker files are grouped based on TensorFlow version and separated based on Python version and processor type.

The Docker images, used to run training & inference jobs, are built from both corresponding “base” and “final” Dockerfiles.

Base Images

The "base" Dockerfile encompass the installation of the framework and all of the dependencies needed.

Tagging scheme is based on <tensorflow_version>-<processor>-<python_version>. (e.g. 1.4 .1-cpu-py2)

All “final” Dockerfiles build images using base images that use the tagging scheme above.

If you want to build your "base" Docker image, then use:

# All build instructions assume you're building from the same directory as the Dockerfile.

# CPU
docker build -t tensorflow-base:<tensorflow_version>-cpu-<python_version> -f Dockerfile.cpu .

# GPU
docker build -t tensorflow-base:<tensorflow_version>-gpu-<python_version> -f Dockerfile.gpu .

# Example

# CPU
docker build -t tensorflow-base:1.4.1-cpu-py2 -f Dockerfile.cpu .

# GPU
docker build -t tensorflow-base:1.4.1-gpu-py2 -f Dockerfile.gpu .

Final Images

The "final" Dockerfiles encompass the installation of the SageMaker specific support code.

All “final” Dockerfiles use base images for building.

These “base” images are specified with the naming convention of tensorflow-base:<tensorflow_version>-<processor>-<python_version>.

Before building “final” images:

Build your “base” image. Make sure it is named and tagged in accordance with your “final” Dockerfile.

# Create the SageMaker TensorFlow Container Python package.
cd sagemaker-tensorflow-containers
python setup.py sdist

#. Copy your Python package to “final” Dockerfile directory that you are building.
cp dist/sagemaker_tensorflow_container-<package_version>.tar.gz docker/<tensorflow_version>/final

If you want to build "final" Docker images, then use:

# All build instructions assumes you're building from the same directory as the Dockerfile.

# CPU
docker build -t <image_name>:<tag> -f Dockerfile.cpu .

# GPU
docker build -t <image_name>:<tag> -f Dockerfile.gpu .

# Example

# CPU
docker build -t preprod-tensorflow:1.4.1-cpu-py2 -f Dockerfile.cpu .

# GPU
docker build -t preprod-tensorflow:1.4.1-gpu-py2 -f Dockerfile.gpu .

# For building images of TensorFlow versions 1.6 and above
docker build -t preprod-tensorflow:1.6.0-cpu-py2 --build-arg py_version=2
--build-arg framework_installable=tensorflow-1.6.0-cp27-cp27mu-manylinux1_x86_64.whl -f Dockerfile.cpu .

Running the tests

Running the tests requires installation of the SageMaker TensorFlow Container code and its test dependencies.

git clone https://github.com/aws/sagemaker-tensorflow-containers.git
cd sagemaker-tensorflow-containers
pip install -e .[test]

Tests are defined in test/ and include unit, integration and functional tests.

Unit Tests

If you want to run unit tests, then use:

# All test instructions should be run from the top level directory

pytest test/unit

Integration Tests

Running integration tests require Docker and AWS credentials, as the integration tests make calls to a couple AWS services. The integration and functional tests require configurations specified within their respective conftest.py.

Integration tests on GPU require Nvidia-Docker.

Before running integration tests:

1. Build your Docker image.
2. Pass in the correct pytest arguments to run tests against your Docker image.

If you want to run local integration tests, then use:

# Required arguments for integration tests are found in test/integ/conftest.py

pytest test/integ --docker-base-name <your_docker_image> \
                  --tag <your_docker_image_tag> \
                  --framework-version <tensorflow_version> \
                  --processor <cpu_or_gpu>

# Example
pytest test/integ --docker-base-name preprod-tensorflow \
                  --tag 1.0 \
                  --framework-version 1.4.1 \
                  --processor cpu

Functional Tests

Functional tests require your Docker image to be within an Amazon ECR repository.

The Docker-base-name is your ECR repository namespace.

The instance-type is your specified Amazon SageMaker Instance Type that the functional test will run on.

Before running functional tests:

1. Build your Docker image.
2. Push the image to your ECR repository.
3. Pass in the correct pytest arguments to run tests on SageMaker against the image within your ECR repository.

If you want to run a functional end to end test on Amazon SageMaker, then use:

# Required arguments for integration tests are found in test/functional/conftest.py

pytest test/functional --aws-id <your_aws_id> \
                       --docker-base-name <your_docker_image> \
                       --instance-type <amazon_sagemaker_instance_type> \
                       --tag <your_docker_image_tag> \

# Example
pytest test/functional --aws-id 12345678910 \
                       --docker-base-name preprod-tensorflow \
                       --instance-type ml.m4.xlarge \
                       --tag 1.0

The sagemaker_tensorflow_containers Python package

Building

This is a standard Python package with a Python C extension module, built using CMake, stored in the "cpp" subdirectory. This module has a series of tests which are run during a setuptools build of hte package. The tests may be explicitly run by invoking the executables built in each test directory.

Contributing

Please read CONTRIBUTING.md for details on our code of conduct, and the process for submitting pull requests to us.

License

SageMaker TensorFlow Containers is licensed under the Apache 2.0 License. It is copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved. The license is available at: http://aws.amazon.com/apache2.0/