Use a custom OS image

You can use a custom OS image for your TPU VMs to pre-load software, use a specific OS distribution, or apply custom kernel modifications. Creating a custom image involves making specific system modifications during the image creation process and configuring the image to handle boot-time tasks required for TPU functionality.

Keep the following disclaimers in mind if you use a custom OS image with TPUs:

• Google provides default TPU-optimized Ubuntu long-term support (LTS) images. The OS changes listed on this page are only validated for the Google-supported, TPU-optimized Ubuntu LTS images.
• You are responsible for extrapolating the required OS changes for any other OS distribution or custom images. Google doesn't guarantee that the modifications for Ubuntu listed on this page work with other OS distributions or another Ubuntu image with a custom kernel.
• Google doesn't build or provide testing for any OS images other than the default TPU-optimized Ubuntu LTS images. You must build and test your custom OS image.

For more information about the default TPU-optimized Ubuntu LTS images, see TPU OS images.

Prerequisites

Your base image must have the following components installed:

• Python 3
• gcloud CLI

Make modifications during image creation

Apply the following modifications while building your custom Ubuntu image.

Bind TPU devices to VFIO

To allow the guest OS to access TPU hardware, you must bind TPU devices to the vfio-pci driver.

1. Create a udev rules file named 99-tpu-vfiopci.rules in /etc/udev/rules.d/:

   # Rules for binding vfio-enabled TPU devices to vfio-pci.
   
   # v5p
   SUBSYSTEM=="pci", ACTION=="add", ATTRS{vendor}=="0x1ae0", ATTRS{device}=="0x0062", ATTRS{subsystem_vendor}=="0x1ae0", ATTRS{subsystem_device}=="0x00ad", DRIVER!="vfio-pci", TAG+="bind_to_vfio_pci"
   
   # v6e
   SUBSYSTEM=="pci", ACTION=="add", ATTRS{vendor}=="0x1ae0", ATTRS{device}=="0x006f", ATTRS{subsystem_vendor}=="0x1ae0", ATTRS{subsystem_device}=="0x00d1", DRIVER!="vfio-pci", TAG+="bind_to_vfio_pci"
   
   # TPU7x
   SUBSYSTEM=="pci", ACTION=="add", ATTRS{vendor}=="0x1ae0", ATTRS{device}=="0x0076", ATTRS{subsystem_vendor}=="0x1ae0", ATTRS{subsystem_device}=="0x00f2", DRIVER!="vfio-pci", TAG+="bind_to_vfio_pci"
   
   # Bind all 'bind_to_vfio_pci' tagged devices to vfio-pci.
   TAG=="bind_to_vfio_pci", RUN+="/lib/udev/bind_to_vfio_pci.sh $kernel"
   

2. Create a script named bind_to_vfio_pci.sh in /lib/udev/:

   #!/bin/bash
   #!/usr/bin/env bash
   
   # Run ./bind_to_vfio_pci.sh <DBDF>
   # Binds the device at <DBDF> to vfio-pci.
   # If the device is already bound to a driver, unbinds it first.
   
   # Load the vfio-pci module into the kernel. No-op if already loaded.
   modprobe vfio-pci
   
   DBDF_REGEX="^[[:xdigit:]]{4}:[[:xdigit:]]{2}:[[:xdigit:]]{2}.[[:xdigit:]]$"
   
   unset BDF
   if [[ $1