Running Dataflow on TPUs: Quickstart examples

This Colab notebook shows you how to set up two pipelines:

A pipeline that runs a trivial computation on a TPU.
A pipeline that runs inference using the Gemma-3-27b-it model on TPUs .

Both pipelines use a custom Docker image. The Dataflow jobs will launch using a Flex Template to allow the same job to be reproduced in different Colab environments.

Prerequisites

First, you need to authenticate to your Google Cloud Project. After running the cell below, you might need to click on the text prompts in the cell and enter inputs as prompted.

import sys
if 'google.colab' in sys.modules:
    from google.colab import auth
    auth.authenticate_user()
!gcloud auth login

Now, set environment variables to access pipeline resources, such as a Cloud Storage bucket or a repository to host container images in Artifact Registry.

import os
import datetime

project_id = "some-project" # @param {type:"string"}
gcs_bucket = "some-bucket" # @param {type:"string"}
ar_repository = "some-ar-repo" # @param {type:"string"}

# Use a region where you have TPU accelerator quota.
region = "some-region1" # @param {type:"string"}
!gcloud config set project {project_id}

Enable the necessary APIs if your project hasn't enabled them yet. If you have the appropriate permissions, you can enable the APIs by running the following cell.

!gcloud services enable \
    dataflow.googleapis.com \
    compute.googleapis.com \
    logging.googleapis.com \
    storage.googleapis.com \
    cloudresourcemanager.googleapis.com \
    artifactregistry.googleapis.com \
    cloudbuild.googleapis.com

Now, you'll create a Cloud Storage bucket and Artifact Registry repository if you don't already have these resources.

gcloud storage buckets describe gs://{gcs_bucket} >/dev/null 2>&1 || gcloud storage buckets create gs://{gcs_bucket} --location={region}
gcloud artifacts repositories describe {ar_repository} --location={region} >/dev/null 2>&1 || gcloud artifacts repositories create {ar_repository} --repository-format=docker --location={region}

Example 1: Minimal computation pipeline using TPU V5E

First, create a simple pipeline you can run to verify that TPUs are accessible, your custom Docker image has the necessary dependencies to interact with the TPUs and your Dataflow pipeline launch configuration is valid.

With this sample you use the PyTorch library to interact with a TPU device.

%%writefile minimal_tpu_pipeline.py
from __future__ import annotations
import torch
import torch_xla
import argparse
import logging
import apache_beam as beam
from apache_beam.options.pipeline_options import PipelineOptions


class check_tpus(beam.DoFn):
    """Validates that a TPU is accessible."""
    def setup(self):
        tpu_devices = torch_xla.xm.get_xla_supported_devices()
        if not tpu_devices:
            raise RuntimeError("No TPUs found on the worker.")
        logging.info(f"Found TPU devices: {tpu_devices}")
        tpu = torch_xla.device()
        t1 = torch.randn(3, 3, device=tpu)
        t2 = torch.randn(3, 3, device=tpu)
        result = t1 + t2
        logging.info(f"Result of a sample TPU computation: {result}")

    def process(self, element):
        yield element


def run(input_text: str, beam_args: list[str] | None = None) -> None:
    beam_options = PipelineOptions(beam_args, save_main_session=True)
    pipeline = beam.Pipeline(options=beam_options)
    (
        pipeline
        | "Create data" >> beam.Create([input_text])
        | "Check TPU availability" >> beam.ParDo(check_tpus())
        | "My transform" >> beam.LogElements(level=logging.INFO)
    )
    pipeline.run()


if __name__ == "__main__":
    logging.getLogger().setLevel(logging.INFO)

    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--input-text",
        default="Hello! This pipeline verified that TPUs are accessible.",
        help="Input text to display.",
    )
    args, beam_args = parser.parse_known_args()

    run(args.input_text, beam_args)

Create a Dockerfile for your TPU-compatible container image.

In your Dockerfile you configure the environment variables to use with a V5E 1x1 TPU device.

You must use the region where you have V5E TPU quota to run this example.

To use a different TPU, adjust the configuration according to the Dataflow documentation.

This Dockerfile creates an image that serves both as a custom worker image for your Beam pipeline and also as a launcher image for your Flex template.

%%writefile Dockerfile

FROM python:3.11-slim

COPY minimal_tpu_pipeline.py minimal_tpu_pipeline.py

# Copy the Apache Beam worker dependencies from the Beam Python 3.10 SDK image.
COPY --from=apache/beam_python3.10_sdk:2.67.0 /opt/apache/beam /opt/apache/beam

# Copy Template Launcher dependencies
COPY --from=gcr.io/dataflow-templates-base/python310-template-launcher-base /opt/google/dataflow/python_template_launcher /opt/google/dataflow/python_template_launcher

# Install TPU software and Apache Beam SDK
RUN pip install --no-cache-dir torch~=2.8.0 torch_xla[tpu]~=2.8.0 apache-beam[gcp]==2.67.0 -f https://storage.googleapis.com/libtpu-releases/index.html

# Configuration for v5e 1x1 accelerator type.
ENV TPU_CHIPS_PER_HOST_BOUNDS=1,1,1
ENV TPU_ACCELERATOR_TYPE=v5litepod-1
ENV TPU_SKIP_MDS_QUERY=1
ENV TPU_HOST_BOUNDS=1,1,1
ENV TPU_WORKER_HOSTNAMES=localhost
ENV TPU_WORKER_ID=0

ENV FLEX_TEMPLATE_PYTHON_PY_FILE=minimal_tpu_pipeline.py

# Set the entrypoint to Apache Beam SDK worker launcher.
ENTRYPOINT [ "/opt/apache/beam/boot"]

Push your Docker image to Artifact Registry.

Finally, build your Docker image, and push it in Artifact Registry. This process should take about 15 minutes or so.

container_tag = "20250801"
container_image = ''.join([
    region, "-docker.pkg.dev/",
    project_id, "/",
    ar_repository, "/",
    "tpu-minimal-example", ":", container_tag
])

!gcloud builds submit --tag {container_image}

Build the Dataflow Flex Template.

To create a reproducible environment for launching the pipeline, build a Flex Template.

First, create a metadata.json file to change the default Dataflow worker disk size when launching the template.

%%writefile metadata.json
{
    "name": "Minimal TPU Example on Dataflow",
    "description": "A Flex template launching a Dataflow Job doing a TPU computation ",
    "parameters": [
      {
        "name": "disk_size_gb",
        "label": "disk_size_gb",
        "helpText": "disk_size_gb for worker",
        "isOptional": true
      }
    ]
}

Run the following cell to build the Flex Template and save it Cloud Storage.

!gcloud dataflow flex-template build gs://{gcs_bucket}/minimal_tpu_pipeline.json \
  --image {container_image} \
  --sdk-language "PYTHON" \
  --metadata-file metadata.json \
  --project {project_id}

Submit your pipeline to Dataflow.

Since you launch the pipeline as a Flex Template, make the following adjustments to the command line:

Use --parameters option to specify the container image and disk size.
Use --additional-experiments option to specify the necessary Dataflow service options.
To avoid using more than one process on a TPU simultaneously, limit process-level parallelism with the no_use_multiple_sdk_containers experiment.

!gcloud dataflow flex-template run "minimal-tpu-example-`date +%Y%m%d-%H%M%S`" \
  --template-file-gcs-location gs://{gcs_bucket}/minimal_tpu_pipeline.json \
  --region {region} \
  --project {project_id} \
  --temp-location gs://{gcs_bucket}/tmp \
  --parameters sdk_container_image={container_image} \
  --worker-machine-type "ct5lp-hightpu-1t" \
  --parameters disk_size_gb=50 \
  --additional-experiments "worker_accelerator=type:tpu-v5-lite-podslice;topology:1x1" \
  --additional-experiments "no_use_multiple_sdk_containers"

Once the job is launched, use the following link to monitor its status: https://console.cloud.google.com/dataflow/jobs/

Sample worker logs for the Check TPU availability step look like the following:

Found TPU devices: ['xla:0']
Result of a sample TPU computation: tensor([[ 0.3355, -1.4628, -3.2610], [-1.4656, 0.3196, -2.8766], [ 0.8667, -1.5060, 0.7125]], device='xla:0')

Example 2: Inference Pipeline with Gemma 3 27B using TPU V6E

This example shows you how to perform inference on a TPU using Gemma 3 27b model.

To fit this model in TPU memory, you need four V6E TPU chips connected in 2x2 topology.

You must use the region where you have V6E TPU quota to run this example.

The example uses Apache Beam RunInference APIs with the VLLM Completions model handler.

The model is downloaded from HuggingFace at runtime, and running the example requires a HuggingFace access token.

First, create a pipeline file.

%%writefile gemma_tpu_pipeline.py
from __future__ import annotations
import argparse
import logging
import apache_beam as beam
from apache_beam.ml.inference.base import RunInference
from apache_beam.options.pipeline_options import PipelineOptions
from apache_beam.ml.inference.vllm_inference import VLLMCompletionsModelHandler


def run(input_text: str, beam_args: list[str] | None = None) -> None:
    beam_options = PipelineOptions(beam_args, save_main_session=True)
    pipeline = beam.Pipeline(options=beam_options)
    (
        pipeline
        | "Create data" >> beam.Create([input_text])
        | "Run Inference" >> RunInference(
            model_handler=VLLMCompletionsModelHandler(
                'google/gemma-3-27b-it',
                {
                    'max-model-len': '4096',
                    'no-enable-prefix-caching': None,
                    'disable-log-requests': None,
                    'tensor-parallel-size': '4',
                    'limit-mm-per-prompt': '{"image": 0}'
                })
            )
        | "Log Output" >> beam.LogElements(level=logging.INFO)
    )
    pipeline.run()


if __name__ == "__main__":
    logging.getLogger().setLevel(logging.INFO)
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--input-text",
        default="What are TPUs?",
        help="Input text query.",
    )
    args, beam_args = parser.parse_known_args()
    run(args.input_text, beam_args)

Create a new Dockerfile for this pipeline with additional dependencies.

Note that this sample uses a different TPU device than the example 1, so the environment variables are different.

You must use your own HuggingFace Token in the Dockerfile. For instructions on creating a token, see User access tokens.

%%writefile Dockerfile
# Use the official vLLM TPU base image, which has TPU dependencies.
# To use the latest version, use: vllm/vllm-tpu:nightly
FROM vllm/vllm-tpu:5964069367a7d54c3816ce3faba79e02110cde17

# Copy your pipeline file.
COPY gemma_tpu_pipeline.py gemma_tpu_pipeline.py

# You can use a more recent version of Apache Beam
COPY --from=apache/beam_python3.12_sdk:2.67.0 /opt/apache/beam /opt/apache/beam
RUN pip install --no-cache-dir apache-beam[gcp]==2.67.0

# Copy Template Launcher dependencies
COPY --from=gcr.io/dataflow-templates-base/python310-template-launcher-base /opt/google/dataflow/python_template_launcher /opt/google/dataflow/python_template_launcher

# Replace the Hugginface token here.
RUN python -c 'from huggingface_hub import HfFolder; HfFolder.save_token("YOUR HUGGINGFACE TOKEN")'

# TPU environment variables.
ENV TPU_SKIP_MDS_QUERY=1

# Configuration for v6e 2x2 accelerator type.
ENV TPU_HOST_BOUNDS=1,1,1
ENV TPU_CHIPS_PER_HOST_BOUNDS=2,2,1
ENV TPU_ACCELERATOR_TYPE=v6e-4
ENV VLLM_USE_V1=1

ENV FLEX_TEMPLATE_PYTHON_PY_FILE=gemma_tpu_pipeline.py

# Set the entrypoint to Apache Beam SDK worker launcher.
ENTRYPOINT [ "/opt/apache/beam/boot"]

Run the following cell to build the Docker image and push it to Artifact Registry. This process should take 15 min or so.

container_tag = "20250801"
container_image = ''.join([
    region, "-docker.pkg.dev/",
    project_id, "/",
    ar_repository, "/",
    "tpu-run-inference-example", ":", container_tag
])
!gcloud builds submit --tag {container_image}

Build the Flex Template for this pipeline.