Cloud Pathways for JAX: Migration Checks Before Your First GKE Run

Cloud Pathways became generally available on Google Kubernetes Engine on November 9, bringing Google's distributed machine-learning runtime to JAX teams managing multiple TPU hosts and slices. Its single-controller model simplifies coordination, but moving an existing multi-controller workload requires checking assumptions about devices, input pipelines and recovery.

That is the useful migration question: which responsibilities move into the runtime, and which remain in the training code? A job that starts successfully can still load data inefficiently or recover incorrectly when a slice disappears.

What general availability covers

Google's GA announcement covers GKE in regions supporting TPU v5e, v5p, v6e and v7x, with versioned containers supporting JAX 0.10.1 or later. Pathways on Slurm clusters through Cluster Director remains a preview. Do not carry the GKE availability label across to every deployment route.

The runtime coordinates a fleet through a resource manager, a proxy and accelerator workers. For the Cluster Toolkit route, submitting with the --pathways flag generates the JobSet configuration for those roles. The toolkit removes manifest work, but you still need the cluster, cloud permissions, capacity and a suitable storage bucket.

Check process indexing before the first migration

Google's JAX porting guide states that jax.process_index() is always zero under Pathways, while jax.devices() and jax.local_devices() return the TPU devices across the whole job.

Review any code that uses process index to select a data shard, initialize state or decide which process writes a checkpoint. A branch written for one Python process per host can have a different meaning when one controller sees the whole allocation.

For example, suppose a data loader previously selected a shard using the host's process index. Simply retaining that code does not prove the migrated workload is consuming the intended global dataset. Validate sample identities and ordering on a small run before measuring throughput. Correct distribution is part of the migration, not a later performance improvement.

Match containers and put compiler settings in the right place

Pathways releases are coupled to JAX versions. Pin the corresponding jax-version container tags and align the workload and any colocated sidecar environments. A floating latest tag makes it harder to reproduce a failed run or distinguish a model regression from an environment change.

There is also an easy configuration trap: XLA flags belong on the pathways-proxy container. Cluster Toolkit exposes --pathways-proxy-args for this. Setting a flag only in the client shell is not sufficient evidence that the remote compiler used it.

The porting guide describes a persistent compilation cache using the Cloud Storage location configured through --gcs_scratch_location. If that value is missing or empty, Pathways falls back to an in-memory cache that does not survive redeployment. When startup repeatedly recompiles, inspect that configuration and its storage permissions before changing the model.

Keep the input pipeline from starving the TPUs

A centralized client can become a data bottleneck if it must prepare and dispatch every batch. Google's guide describes asynchronous read-ahead as suitable for smaller deployments and recommends colocated Python for moving input processing onto accelerator hosts at larger sizes.

For a multiprocess loader in the GKE head container, set an explicit worker count within the container's CPU allocation. The guide warns that os.cpu_count() reports the host VM's CPU count, which can oversubscribe the container while it shares the node with Pathways services. It also calls for a memory-backed /dev/shm volume so prefetched batches do not hit the default container shared-memory limit.

Start by measuring input wait time alongside device utilization. More accelerator capacity will not fix a loader that cannot feed the existing allocation, and it can make the cost of that bottleneck larger.

Elastic recovery requires training-code changes

Google's resilient training guide distinguishes planned interruption handling from elastic recovery after a hardware failure. Some workloads can suspend and resume transparently; others rely on stateful components that are not restorable through the same path.

Elastic training keeps the client and healthy slices alive, but recovery logic still has to rebuild the device mesh and restore the training state. The guide uses pathwaysutils with Orbax snapshots and distinguishes waiting for the original allocation from continuing on fewer slices. Resizing can require recompilation, so measure recovery time rather than assuming it disappears.

Use the active-device set for recovery, not the original device list. The guide notes that jax.devices() retains the initial allocation for the client's lifetime. Its elastic manager exposes active_devices for building a mesh from healthy resources.

Dataset position needs the same care. A loader whose state is tied to the original host count may not restore correctly onto a smaller topology. The documentation uses an elastic iterator with a global dataset position to address that problem.

A small acceptance test before a large run

For a first migration, keep the model and dataset small enough to inspect. Compare a short baseline run with the migrated job, verify dataset progression and checkpoint restoration, and separate compilation time from steady-state training. Then test an approved interruption scenario with the recovery mode you actually intend to use.

When debugging, collect the workers' logs as well as the controller's. Google's troubleshooting guide uses --main-only=false with gcluster job logs to include worker pods. A quiet client log does not establish that every worker is healthy.

These are documentation-based checks, not results from an independently run cluster. Pathways can make a large TPU allocation easier to program, but the migration is complete only when the job still consumes the right data and produces recoverable state after its ideal execution path breaks.

I’m Paras. I work at the intersection of AI, engineering and distribution.

I help technical teams explain infrastructure changes and developer workflows without turning them into marketing copy.

Book a 30-minute call.