Image
class. The options specified in the Image
class will influence how the image is built.
Exploring the Beam Image Class
Every application that runs on Beam instantiates theImage
class. This class provides a variety of methods for customizing the container image used to run your application.
It exposes options for:
- Installing a specific version of Python
- Adding custom shell commands that run during the build process
- Adding custom Python packages to install in the container
- Choosing a custom base image to build on top of
- Using a custom Dockerfile to build your own base image
- Setting up a custom conda environment using micromamba
The default Beam image uses
ubuntu:22.04
as its base and installs Python
3.10.Adding Python Packages
The most common way to customize your image is to add the Python packages required by your application. This is done by calling theadd_python_packages
method on the Image
object with a list of package names.
Pinning the version of the package is recommended. This ensures that when you
re-deploy your application, you won’t accidentally pick up a new version that
breaks your application.
Importing requirements.txt
If you already have a requirements.txt
file, you can also use that directly using the Image
constructor’s python_packages
parameter:
Adding Shell Commands
Sometimes, it is necessary to run additional shell commands while building your image. This can be achieved by calling theadd_commands
method on the Image
object with a list of commands.
For instance, you might need to install libjpeg-dev
when using the Pillow
library. In the example below, we’ll install libjpeg-dev
and then install Pillow
.
Customizing the Base Image
Some applications and libraries require specific dependencies that are not available in the default Beam image. In these cases, you can use a custom base image. Some of the most common custom base images are the CUDA development images from NVIDIA (e.g.nvcr.io/nvidia/cuda:12.4.1-cudnn-devel-ubuntu22.04
). These images come with additional libraries, debugging tools, and nvcc
installed.
The image below will use a custom CUDA image as the base.
CUDA Drivers & NVIDIA Kernel Drivers
When choosing a custom base image, it is important to understand the difference between the NVIDIA Kernel Driver and the CUDA Runtime & Libraries.Component | Location | Role |
---|---|---|
NVIDIA Kernel Driver | Host Machine | Low-level GPU management, talks directly to hardware. |
CUDA Runtime & Libraries | Container | Provides high-level APIs and libraries for applications. |
For example, using a CUDA 12.2 image on a host with a CUDA 12.4 driver will
work. However, using a CUDA 12.8 image on a host with a CUDA 12.4 driver will
not work.
GPU | Driver Version | CUDA Version |
---|---|---|
T4 | 550.127.05 | 12.4 |
A10G | 550.90.12 | 12.4 |
RTX4090 | 550.127.05 | 12.4 |
L40s | 550.127.05 | 12.4 |
A100-40 | 550.127.05 | 12.4 |
H100 | 550.127.05 | 12.4 |
Using a Specific Python Version
To install a specific version of Python, you can use thepython_version
parameter:
python_version
is specified and the CUDA image has no Python version installed.
python_version
is specified.
python_version
in your Image
constructor. This function will use the PyTorch image as the base and will use the Python version that already exists in the PyTorch image.
Building on GPU
By default, Beam builds your images on CPU-only machines. However, sometimes you might need the build to occur on a machine with a GPU. For instance, some libraries might compile CUDA kernels during installation. In these cases, you can use thebuild_with_gpu()
command to run your build on the GPU of your choice.
Building with Environment Variables
Often, shell commands require certain environment variables to be set. You can set these using thewith_envs
command:
Injecting Secrets
Sometimes, you might not want the environment variables to be set in plain text. In these cases, you can leverage Beam secrets and thewith_secrets
command:
You can create secrets like this, using the CLI:
beam secret create HF_TOKEN <your_token>
.Using a Dockerfile
You also have the option to build your own custom base image using a Dockerfile. Thefrom_dockerfile()
command accepts a path to a valid Dockerfile as well as an optional path to a context directory:
COPY
and ADD
, meaning all relative paths are relative to this directory.
The image built from your Dockerfile will be used as the base image for a Beam application.
Ports will not be exposed in the runtime environment, and the entrypoint
will be overridden.
Conda Environments
Beam supports using Anaconda environments via micromamba. To get started, you can chain themicromamba
method to your Image
definition and then specify packages and channels via the add_micromamba_packages
method.
pip
to install additional packages in the conda
environment and you can run shell commands too.
If you need to run a shell command inside the conda environment, you should
prepend the command with
micromamba run -n beta9
as shown above.