Training a machine learning model across multiple TPU devices usually requires coordination between several machines. Each machine runs its own Python process, and those processes communicate to distribute data, execute computations, and synchronize results.
This approach works, but managing multiple processes can make development and debugging difficult.
Google Cloud Pathways offers another approach. It allows developers to control distributed TPU workloads through a single Python client process. Instead of managing a separate Python interpreter on every TPU host, developers can write their JAX code in one place while Pathways coordinates execution across multiple TPU slices.
On November 9, 2026, Google announced the general availability of Cloud Pathways on Google Kubernetes Engine (GKE).
In this article, we will understand how this architecture works, connect a Python application to a Pathways cluster, and run a basic JAX computation across multiple TPU devices.
Why Running Distributed TPU Workloads Is Difficult
Consider a machine learning team training a large model across several TPU hosts.
In a traditional multi-controller JAX setup, each host runs its own Python process. Every process participates in the computation and communicates with other processes when data needs to be exchanged.
This approach is commonly called Single Program, Multiple Data (SPMD).
For example, a workload running across four TPU hosts might have the following structure:
Host 1 Host 2
+----------------+ +----------------+
| Python Process | | Python Process |
| | | |
| TPU Devices | | TPU Devices |
+-------+--------+ +-------+--------+
| |
+----------+-----------+
|
Communication
|
+----------+-----------+
| |
+-------+--------+ +-------+--------+
| Python Process | | Python Process |
| | | |
| TPU Devices | | TPU Devices |
+----------------+ +----------------+
Host 3 Host 4
Each process has a local view of its devices. The processes must coordinate with one another to execute the complete workload.
This creates several development challenges:
- The same program must be launched on multiple hosts.
- Distributed initialization must complete successfully.
- Communication failures can interrupt the workload.
- Debugging requires investigating multiple processes.
- Scaling across TPU slices introduces additional coordination.
For smaller workloads, this approach is manageable. For large machine learning systems, it can become difficult to maintain.
What Is Google Cloud Pathways?
Pathways is a distributed machine learning runtime designed to coordinate computations across large accelerator environments.
Google originally developed it for large-scale machine learning workloads, including training and serving models such as Gemini.
Cloud Pathways brings this execution model to Google Cloud.
Its main feature is a single-controller architecture. A single JAX client acts as the central point of control while Pathways manages the execution across distributed TPU resources.
The architecture looks like this:
Python Application
|
v
Single JAX Client
|
v
IFRT Proxy Client
|
v
Pathways Proxy Server
|
v
Pathways Resource Manager
|
+------+------+
| |
v v
TPU Slice 1 TPU Slice 2
| |
v v
TPU Workers TPU Workers
| |
+------+------+
|
Computation
The Python application does not need to manage every TPU host separately.
Pathways receives the computation from the JAX client, coordinates its execution, and sends compiled programs to the appropriate TPU workers.
The application interacts with the devices through a unified view.
Single Controller vs. Multi-Controller JAX
It is important to understand that Pathways does not simply combine multiple Python processes into one machine.
The execution architecture is different.
| Area |
Multi-controller JAX |
Single-controller Pathways |
| Python processes |
Multiple client processes |
One client process |
| Device visibility |
Each process has a local view |
Client has a unified device view |
| Coordination |
Processes participate in distributed execution |
Pathways coordinates execution |
| Programming model |
Primarily SPMD |
Supports more flexible execution patterns |
| Debugging |
Requires investigating multiple processes |
One central client process |
| TPU execution |
Distributed across hosts |
Distributed across Pathways workers |
Pathways still uses multiple worker processes on TPU hosts. The difference is that developers do not need to manage a separate JAX client process on every host.
This makes it easier to develop workloads that combine different computational stages, such as training, sampling, and inference.
Understand the Main Pathways Components
Pathways has several components that work together.
Pathways resource manager
The resource manager manages accelerator resources and coordinates their allocation.
It monitors worker health and handles job scheduling, pausing, and resuming.
It provides the central control plane for the distributed workload.
Pathways client
The Pathways client receives operations from the JAX application and determines how compiled programs should be placed across the available devices.
It provides the application with a unified view of the distributed TPU environment.
IFRT proxy client and server
IFRT stands for Interim Framework Runtime.
It provides the interface between the JAX application and the distributed Pathways runtime.
The proxy client sends requests to the proxy server, which forwards them to the Pathways system.
Pathways workers
Pathways workers run on the TPU hosts.
They receive compiled programs, execute computations on the TPU devices, and return results through the Pathways communication layer.
This separation allows the Python client to control the workload without running directly on every TPU host.
Prepare the Google Cloud Environment
Before running a Python workload, you need a working Pathways cluster.
Cloud Pathways runs on GKE and requires the appropriate TPU infrastructure.
The main requirements include:
- A Google Cloud project with Pathways enabled.
- A GKE cluster configured for Pathways.
- Supported Cloud TPU resources and sufficient quota.
- Google Cloud CLI.
- Kubernetes tools.
- XPK for supported cluster and workload management workflows.
- A Cloud Storage bucket for temporary files and compilation-related data.
As of November 2026, Cloud Pathways is generally available on GKE in regions supporting Cloud TPU v5e, v5p, Trillium (v6e), and Ironwood (TPU 7x). The documented versioned Pathways container images support JAX 0.10.1 and later.
The exact TPU availability depends on the selected region and capacity.
For this example, assume that the Pathways cluster and its TPU workers have already been deployed.
We will focus on connecting the Python application to the running cluster rather than creating the complete GKE infrastructure.
Connect a Python Application to Pathways
Pathways provides an IFRT proxy server through which JAX applications communicate with the distributed TPU environment.
For an interactive workload, you can expose the proxy server using Kubernetes port forwarding.
First, identify the Pathways proxy Pod:
kubectl get pods
Find the Pod running the Pathways head or proxy components.
Set up port forwarding:
PROXY_POD=$(kubectl get pods | grep pathways-head | awk '{print $1}')
PROXY_PORT=29000
kubectl port-forward ${PROXY_POD} ${PROXY_PORT}:${PROXY_PORT}
The terminal should display a message similar to:
Forwarding from 127.0.0.1:29000 -> 29000
Keep this terminal running while you connect your Python application.
Open another terminal and create a virtual environment:
python3 -m venv .venv
source .venv/bin/activate
Install JAX and the Pathways utilities:
pip install pathwaysutils jax
These packages provide the client-side integration. Use the version combination recommended for your Pathways cluster.
Run a Python Program Across TPU Devices
Create a Python file named tpu_workload.py.
The following example connects to Pathways, discovers the available TPU devices, and performs a distributed calculation.
import jax
import jax.numpy as jnp
from jax import lax
import pathwaysutils
# Initialize the Pathways client
pathwaysutils.initialize()
# Discover all available TPU devices
devices = jax.devices()
print("Available devices:", len(devices))
print("Device details:", devices)
# Get the number of available devices
num_tpus = jax.device_count()
# Create input data
x = jnp.arange(num_tpus)
# Define a distributed computation
def distributed_sum(value):
return lax.psum(1, "devices")
# Execute the computation across TPU devices
calculate = jax.pmap(
distributed_sum,
axis_name="devices"
)
result = calculate(x)
print("Distributed result:", result)
Run the program with the Pathways proxy configuration:
JAX_PLATFORMS=proxy \
JAX_BACKEND_TARGET=grpc://127.0.0.1:29000 \
python tpu_workload.py
The important parts of the program are explained below.
First, pathwaysutils.initialize() initializes the Pathways integration.
This must happen before querying devices or performing JAX computations.
Next, jax.devices() returns the devices visible through the Pathways client.
The jax.pmap function maps the computation across the selected devices. The lax.psum operation performs a sum across the mapped devices.
Since every device contributes the value 1, the result on each device is the total number of participating devices.
For example, if eight TPU devices participate, each device receives the value 8.
The example demonstrates the execution model. It is not a machine learning training loop, but the same connection can be used as the starting point for more complex workloads.
What Happens When the Workload Runs?
When you execute the Python program, several operations happen behind the scenes.
- The Python application connects to the IFRT proxy server.
- Pathways initializes the client connection.
- The client discovers the available TPU devices.
- JAX builds the computation.
- Pathways coordinates the placement and execution of compiled programs.
- TPU workers perform the computation.
- Results are returned to the Python application.
The Python application does not need to establish individual SSH connections to every TPU host or launch a separate JAX program on each one.
This is one of the main advantages of the single-controller model.
Where Pathways Helps in Real Applications
Pathways is particularly useful when a workload requires multiple computational stages or uses several TPU slices.
Large model training
A large model might not fit within the memory of a single TPU slice.
Pathways can coordinate computations across multiple slices and provide a unified programming interface.
Reinforcement learning
Reinforcement learning workloads often combine training, sampling, and evaluation.
These stages do not always require the same amount of compute.
Pathways supports more flexible execution patterns than a traditional lockstep SPMD approach.
Large-scale inference
An inference workload might distribute different model components across multiple TPU resources.
A single controller can coordinate execution while Pathways manages the underlying accelerator workers.
Interactive development
Developers can connect through a terminal or notebook, inspect available devices, execute JAX code, and debug their workload without manually launching a client process on every TPU host.
Important Considerations
Pathways simplifies distributed execution, but it does not remove the need to understand the underlying infrastructure.
Keep the following points in mind:
- A Pathways cluster must be deployed and running before the Python client connects.
- The client and Pathways components must use compatible versions.
- Network access to the IFRT proxy server must be configured correctly.
- TPU availability, quota, and topology affect the resources available to the workload.
- Large data transfers can still affect performance.
- Model sharding and data placement still require careful design.
- Monitor TPU utilization, memory usage, and execution time rather than assuming that all devices are being used efficiently.
For production training, also review Pathways' elastic training and checkpointing capabilities. These features help workloads recover from certain TPU host or slice failures, but recovery still depends on the application and its checkpointing strategy.
Summary
Google Cloud Pathways provides a single-controller approach to running distributed JAX workloads across multiple TPU devices. Instead of launching and coordinating a separate Python client process on every TPU host, developers can use one Python application to control the workload. Pathways handles device coordination, program placement, and execution across the distributed TPU infrastructure. This makes it easier to develop and debug large machine learning workloads, including model training, reinforcement learning, and inference. Developers still need to configure the TPU infrastructure, manage data placement, and monitor performance, but they can work with a much simpler programming interface.
Join the conversation! Your thoughts help the community grow.