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.
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.
If you want to check what packages are installed in a given version of Amazon SageMaker Distribution, you can find that in the relevant RELEASE.md file in the build_artifacts directory.
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).
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:
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.0.1-cpu
: this, and the two below, can change when new versions of Amazon SageMaker Distribution are released.0-cpu
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.
If you want to generate/view the staleness report for each of the individual packages in a given SageMaker distribution image version, then run the following command:
VERSION=<Insert SageMaker Distribution version in semver format here. example: 0.4.2>
python ./src/main.py generate-staleness-report --target-patch-version $VERSION
Here are some examples on how you can try out one of our images.
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 is a web-based, integrated development environment (IDE) for machine learning that lets you build, train, debug, deploy, and monitor your machine learning models.
To use the sagemaker-distribution image in SageMaker Studio, select SageMaker Distribution v{Major_version} {CPU/GPU}
using the SageMaker Studio Launcher.
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.
If you'd like to create a new Docker image on top of what we offer, we recommend you use micromamba install ...
instead of pip install ...
.
For example:
FROM public.ecr.aws/sagemaker/sagemaker-distribution:latest-cpu
USER $ROOT
RUN apt-get install -y vim
USER $MAMBA_USER
RUN micromamba install sagemaker-inference --freeze-installed --yes --channel conda-forge --name base
See CONTRIBUTING for more information.
This project is licensed under the Apache-2.0 License.