Chapter 2. Managing custom training images To run distributed training jobs, you can use one of the base training images that are provided with OpenShift AI, or you can create your own custom training images. You can optionally push your custom training images to the integrated OpenShift image registry, to make your images available to other users. 2.1. About base training images The base training images for distributed workloads are optimized with the tools and libraries that you need to run distributed training jobs. You can use the provided base images, or you can create custom images that are specific to your needs. For information about Red Hat support of training images and packages, see Red Hat OpenShift AI: Supported Configurations . The following table lists the training images that are installed with Red Hat OpenShift AI by default. These images are AMD64 images, which might not work on other architectures. Table 2.1. Default training base images Image type Description Ray CUDA If you are working with compute-intensive models and you want to accelerate the training job with NVIDIA GPU support, you can use the Ray Compute Unified Device Architecture (CUDA) base image to gain access to the NVIDIA CUDA Toolkit. Using this toolkit, you can accelerate your work by using libraries and tools that are optimized for NVIDIA GPUs. Ray ROCm If you are working with compute-intensive models and you want to accelerate the training job with AMD GPU support, you can use the Ray ROCm base image to gain access to the AMD ROCm software stack. Using this software stack, you can accelerate your work by using libraries and tools that are optimized for AMD GPUs. KFTO CUDA If you are working with compute-intensive models and you want to accelerate the training job with NVIDIA GPU support, you can use the Kubeflow Training Operator (KFTO) CUDA base image to gain access to the NVIDIA CUDA Toolkit. Using this toolkit, you can accelerate your work by using libraries and tools that are optimized for NVIDIA GPUs. KFTO ROCm If you are working with compute-intensive models and you want to accelerate the training job with AMD GPU support, you can use the KFTO ROCm base image to gain access to the AMD ROCm software stack. Using this software stack, you can accelerate your work by using libraries and tools that are optimized for AMD GPUs. If the preinstalled packages that are provided in these images are not sufficient for your use case, you have the following options: Install additional libraries after launching a default image. This option is good if you want to add libraries on an ad hoc basis as you run training jobs. However, it can be challenging to manage the dependencies of installed libraries. Create a custom image that includes the additional libraries or packages. For more information, see Creating a custom training image . 2.2. Creating a custom training image You can create a custom training image by adding packages to a base training image. Prerequisites You can access the training image that you have chosen to use as the base for your custom image. Select the image based on the image type (for example, Ray or KFTO), the accelerator framework (for example, CUDA for NVIDIA GPUs, or ROCm for AMD GPUs), and the Python version (for example, 3.9 or 3.11). The following table shows some example base training images: Table 2.2. Example base training images Image type Accelerator framework Python version Example base training image Preinstalled packages Ray CUDA 3.9 ray:2.35.0-py39-cu121 Ray 2.35.0, Python 3.9, CUDA 12.1 Ray CUDA 3.11 ray:2.35.0-py311-cu121 Ray 2.35.0, Python 3.11, CUDA 12.1 Ray ROCm 3.11 ray:2.35.0-py311-rocm62 Ray 2.35.0, Python 3.11, ROCm 6.2 KFTO CUDA 3.11 training:py311-cuda121-torch241 Python 3.11, CUDA 12.1, PyTorch 2.4.1 KFTO ROCm 3.11 training:py311-rocm62-torch241 Python 3.11, ROCm 6.2, PyTorch 2.4.1 For a complete list of the OpenShift AI base training images and their preinstalled packages, see Supported Configurations . You have Podman installed in your local environment, and you can access a container registry. For more information about Podman and container registries, see Building, running, and managing containers . Procedure In a terminal window, create a directory for your work, and change to that directory. Set the IMG environment variable to the name of your custom image. In the example commands in this section, my_training_image is the name of the custom image. Create a file named Dockerfile with the following content: Use the FROM instruction to specify the location of a suitable base training image. In the following command, replace _<base-training-image>_ with the name of your chosen base training image: Examples: FROM quay.io/modh/ray:2.35.0-py311-cu121 FROM quay.io/modh/training:py311-rocm62-torch241 Use the RUN instruction to install additional packages. You can also add comments to the Dockerfile by prefixing each comment line with a number sign ( # ). The following example shows how to install a specific version of the Python PyTorch package: # Install PyTorch RUN python3 -m pip install torch==2.4.0 Build the image file. Use the -t option with the podman build command to create an image tag that specifies the custom image name and version, to make it easier to reference and manage the image: Example: podman build -t USD{IMG}:0.0.1 -f Dockerfile The build output indicates when the build process is complete. Display a list of your images: If your new image was created successfully, it is included in the list of images. Push the image to your container registry: podman push USD{IMG}:0.0.1 Optional: Make your new image available to other users, as described in Pushing an image to the integrated OpenShift image registry . 2.3. Pushing an image to the integrated OpenShift image registry To make an image available to other users in your OpenShift cluster, you can push the image to the integrated OpenShift image registry , a built-in container image registry. For more information about the integrated OpenShift image registry, see Integrated OpenShift image registry . Prerequisites Your cluster administrator has exposed the integrated image registry, as described in Exposing the registry . You have Podman installed in your local environment. For more information about Podman and container registries, see Building, running, and managing containers . Procedure In a terminal window, log in to the OpenShift CLI as shown in the following example: Set the IMG environment variable to the name of your image. In the example commands in this section, my_training_image is the name of the image. Log in to the integrated image registry: Tag the image for the integrated image registry: Push the image to the integrated image registry: Retrieve the image repository location for the tag that you want: Any user can now use your image by specifying this retrieved image location value in the image parameter of a Ray cluster or training job.

Commands:
export IMG=my_training_image
FROM quay.io/modh/ <base-training-image>
FROM quay.io/modh/ray:2.35.0-py311-cu121
FROM quay.io/modh/training:py311-rocm62-torch241
Install PyTorch RUN python3 -m pip install torch==2.4.0
build -t <custom-image-name> :_<version>_ -f Dockerfile
build -t USD{IMG}:0.0.1 -f Dockerfile
images
push USD{IMG}:0.0.1
oc login <openshift_cluster_url> -u <admin_username> -p <password>
export IMG=my_training_image
login -u USD(oc whoami) -p USD(oc whoami -t) USD(oc registry info)
tag USD{IMG} USD(oc registry info)/USD(oc project -q)/USD{IMG}
push USD(oc registry info)/USD(oc project -q)/USD{IMG}
get is USD{IMG} -o jsonpath='{.status.tags[?(@.tag==" <TAG> ")].items[0].dockerImageReference}'