Amazon SageMaker Distribution

Amazon SageMaker Distribution is a set of Docker images that include popular frameworks for machine learning, data science and visualization.

These images come in two variants, CPU and GPU, and include deep learning frameworks like PyTorch, TensorFlow and Keras; popular Python packages like numpy, scikit-learn and pandas; and IDEs like Jupyter Lab. The distribution contains the latest versions of all these packages such that they are mutually compatible.

This project follows semver (more on that below) and comes with a helper tool to automate new releases of the distribution.

Currently under beta

This feature is provided as a ‘Beta Service’ as defined in the AWS Service Terms. It is governed by your Agreement with AWS and the AWS Service Terms. Before using this Beta Service, please review the Betas and Previews terms found here. In particular, please note that the Beta Service is confidential and you may not discuss the features, functionality, or documentation related to the Beta Service with any parties that are not authorized by AWS. The Beta Service is subject to change and cancellation.

During this beta, we may modify some interfaces of the images, their feature-set or the list of packages that are included in it. However, the image versions will continue to be v0.x.y and such modifications, if any, will only occur as part of a minor upgrade (such as 0.4.x to 0.5.0).

Getting started

If you just want to use the images, you do not need to use this GitHub repository. Instead, you can pull pre-built and ready-to-use images from our AWS ECR Gallery repository.

Dependency versions included in a particular Amazon SageMaker Distribution version

If you want to check what packages are installed in a given version of Amazon SageMaker Distribution, you can find that in the relevant env.out file in the build_artifacts directory.

Versioning strategy

Amazon SageMaker Distribution supports semantic versioning as described on semver.org. A major version upgrade of Amazon SageMaker Distribution allows major version upgrades of all its dependencies, and similarly for minor and patch version upgrades. However, it is important to note that Amazon SageMaker Distribution’s ability to follow semver guidelines is currently dependent on how its dependencies adhere to them.

Some dependencies, such as Python, will be treated differently. Amazon SageMaker Distribution will allow a minor upgrade of Python (say, 3.10 to 3.11) only during a major upgrade (say, 4.8 to 5.0).

Image tags

Our current image tagging scheme is: <AMAZON_SAGEMAKER_DISTRIBUTION_VERSION_NUMBER>-<CPU_OR_GPU>. For example, the CPU version of Amazon SageMaker Distribution's v0.1.2 will carry the following tags:

1. 0.1.2-cpu: Once an image is tagged with such a patch version, that tag will not be assigned to any other image in future.
2. 0.1-cpu: this, and the two below, can change when new versions of Amazon SageMaker Distribution are released.
3. 0-cpu
4. latest-cpu

So, if you want to stay on the latest software as and when release by Amazon SageMaker Distribution, you can use latest-cpu and do a docker pull latest-cpu when needed. If you use, say, 0.1.2-cpu, the underlying distribution will remain the same over time.

Example use cases

Here are some examples on how you can try out one of our images.

Local environment, such as your laptop

The easiest way to get it running on your laptop is through the Docker CLI:

export ECR_IMAGE_ID='INSERT_IMAGE_YOU_WANT_TO_USE'
docker run -it \
    -p 8888:8888 \
    --user `id -u`:`id -g` \
    -v `pwd`/sample-notebooks:/home/sagemaker-user/sample-notebooks \
    $ECR_IMAGE_ID jupyter-lab --no-browser --ip=0.0.0.0

(If you have access to Nvidia GPUs, you can pass --gpus=all to the Docker command.)

In the console output, you'll then see a URL similar to http://127.0.0.1:8888/lab?token=foo. Just open that URL in your browser, create a Jupyter Lab notebook or open a terminal, and start hacking.

Note that the sample command above bind mounts a directory in pwd inside the container. That way, if you were to re-create the container (say, to use a different version or CPU/GPU variant), any files you created within that directory (such as Jupyter Lab notebooks) will persist.

Amazon SageMaker Studio

Amazon SageMaker Studio is a web-based, integrated development environment (IDE) for machine learning that lets you build, train, debug, deploy, and monitor your machine learning models.

You can use Amazon SageMaker Studio's BYOI (i.e. bring your own image) feature to try out the images in the cloud. See relevant documentation here.

Let's say you already have a Amazon SageMaker Studio Domain in your AWS account. You can then run the following steps to pull in the image into Studio. (Note that we use AWS CLI below, but all of the same operations are possible either through the AWS console or Python SDK.)

# Step 0: Define the following variables specific to your setup.
export ECR_GALLERY_IMAGE_ID='sagemaker-distribution:latest-cpu'
export SAGEMAKER_IMAGE_NAME='...'
export SAGEMAKER_STUDIO_DOMAIN_ID='d-foo'
export SAGEMAKER_STUDIO_IAM_ROLE_ARN='...'

# Step 1.0: Copy pre-built image from ECR Gallery to your private ECR repository.
# Note: Currently SageMaker Studio doesn't support loading images from ECR Gallery. This could change in the future and at that point, this step might not be necessary.

# Prerequisite: AWS Account should already have a private ECR repository
# https://docs.aws.amazon.com/AmazonECR/latest/userguide/repository-create.html

# Pull the Docker image from the public ECR Gallery
docker pull public.ecr.aws/sagemaker/$ECR_GALLERY_IMAGE_ID

# Tag the image for your private ECR Repository.
# Define the following variables specific to your setup.
export ECR_PRIVATE_REPOSITORY_NAME='...'
export ECR_IMAGE_TAG='...'
export AWS_ACCOUNT_ID='...'
export AWS_ECR_REPOSITORY_REGION='...'
export ECR_IMAGE_URI=$AWS_ACCOUNT_ID.dkr.ecr.$AWS_ECR_REPOSITORY_REGION.amazonaws.com/$ECR_PRIVATE_REPOSITORY_NAME:$ECR_IMAGE_TAG
docker tag public.ecr.aws/sagemaker/$ECR_GALLERY_IMAGE_ID $ECR_IMAGE_URI

# Push the image to your private repository
docker push $ECR_IMAGE_URI

# Step 1.1: Create an Amazon SageMaker Image.
aws sagemaker create-image \
    --image-name $SAGEMAKER_IMAGE_NAME \
    --role-arn $SAGEMAKER_STUDIO_IAM_ROLE_ARN

# Step 2: Create an Amazon SageMaker Image Version.
aws sagemaker create-image-version \
    --image-name $SAGEMAKER_IMAGE_NAME \
    --base-image $ECR_IMAGE_URI

# Optional, step 2.1: Verify that the Image and Image Version were successfully created.
# You may have to specify a different `--version-number` if you created multiple Image Versions.
aws sagemaker describe-image-version \
    --image-name $SAGEMAKER_IMAGE_NAME \
    --version-number 1

# Step 3.1: Create a file that will contain App Image Config parameters.
cat > /tmp/app-config.json << EOF
{
   "AppImageConfigName": "app-image-config-$SAGEMAKER_IMAGE_NAME",
   "KernelGatewayImageConfig": { 
      "FileSystemConfig": { 
         "DefaultGid": 100,
         "DefaultUid": 1000,
         "MountPath": "/home/sagemaker-user"
      },
      "KernelSpecs": [ 
         { 
            "DisplayName": "Python 3 (ipykernel)",
            "Name": "python3"
         }
      ]
   }
}
EOF

# Step 3.2: Create an Amazon SageMaker App Image Config.
aws sagemaker create-app-image-config \
    --cli-input-json file:///tmp/app-config.json

# Step 4.1: Create a file that will contain default parameters for User Settings.
cat > /tmp/default-user-settings.json << EOF
{
    "DefaultUserSettings": {
        "KernelGatewayAppSettings": {
            "CustomImages": [
                {
                    "ImageName": "$SAGEMAKER_IMAGE_NAME",
                    "AppImageConfigName": "app-image-config-$SAGEMAKER_IMAGE_NAME",
                    "ImageVersionNumber": 1
                }
            ]
        }
    }
}
EOF

# Step 4.2: Update Amazon SageMaker Domain with the new default User Settings.
aws sagemaker update-domain \
    --domain-id $SAGEMAKER_STUDIO_DOMAIN_ID \
    --cli-input-json file:///tmp/default-user-settings.json

Now, for your User Profile, log-out and log back in to start using the custom image.

"I want to directly use the Conda environment, not via a Docker image"

Amazon SageMaker Distribution supports full reproducibility of Conda environments, so you don't necessarily need to use Docker. Just find the version number you want to use in the build_artifacts directory, open one of cpu.env.out or gpu.env.out and follow the instructions in the first 2 lines.

Security

See CONTRIBUTING for more information.

License

This project is licensed under the Apache-2.0 License.