trainer/examples/jax/tpu_smoke_test.ipynb

Distributed JAX TPU Smoke Test with Kubeflow Trainer

This notebook demonstrates how to run a distributed JAX verification job on physical Cloud TPUs using the Kubeflow Trainer Python SDK.

It initializes a distributed JAX environment, verifies hardware discovery across all physical TPU chips, and executes a basic multi-host matrix multiplication.

Install the Kubeflow SDK

You need to install the Kubeflow SDK to interact with Kubeflow Trainer.

In [ ]:

# !pip install -U kubeflow

Define JAX Distributed Verification Function

This function will run on each TPU worker. It initializes JAX's distributed coordination service, prints diagnostic information about local and global TPU devices, and runs a distributed matrix multiplication.

In [ ]:

def jax_tpu_smoke_test():
    import jax
    import jax.numpy as jnp

    # Initialize distributed JAX coordination
    jax.distributed.initialize()

    # Log device count
    print("==========================")
    print(f"Total global TPU device count: {jax.device_count()}")
    print(f"Local TPU device count: {jax.local_device_count()}")
    print(f"Global TPU devices details: {jax.devices()}")
    print("==========================")

    # Execute distributed matrix multiplication to verify JAX TPU backend
    x = jnp.ones((1000, 1000))
    y = (x @ x).block_until_ready()
    
    print(f"JAX Matrix Multiplication verification successfully completed! Result shape: {y.shape}")

Initialize Kubeflow Trainer Client

We initialize the TrainerClient which communicates with the Kubeflow Trainer API on the Kubernetes cluster.

In [ ]:

from kubeflow.trainer import TrainerClient, CustomTrainer

client = TrainerClient()

Configure GKE TPU Scheduling Patches

Since physical TPU hardware requires GKE-specific node selectors, topologies, and tolerations, we configure a RuntimePatch programmatically in the Python SDK. This will inject the GKE scheduling parameters directly into our training Pod specifications.

In [ ]:

from kubeflow.trainer.options import RuntimePatch
from kubeflow.trainer.options.kubernetes import (
    TrainingRuntimeSpecPatch,
    JobSetTemplatePatch,
    JobSetSpecPatch,
    ReplicatedJobPatch,
    JobTemplatePatch,
    JobSpecPatch,
    PodTemplatePatch,
    PodSpecPatch,
)

# We define a GKE scheduling patch using typed SDK dataclasses
tpu_patch = RuntimePatch(
    training_runtime_spec=TrainingRuntimeSpecPatch(
        template=JobSetTemplatePatch(
            spec=JobSetSpecPatch(
                replicated_jobs=[
                    ReplicatedJobPatch(
                        name="node",
                        template=JobTemplatePatch(
                            spec=JobSpecPatch(
                                template=PodTemplatePatch(
                                    spec=PodSpecPatch(
                                        node_selector={
                                            "cloud.google.com/gke-tpu-accelerator": "tpu-v5-lite-podslice",
                                            "cloud.google.com/gke-tpu-topology": "2x2",
                                        },
                                        tolerations=[
                                            {
                                                "key": "google.com/tpu",
                                                "operator": "Exists",
                                                "effect": "NoSchedule",
                                            }
                                        ],
                                    )
                                )
                            )
                        )
                    )
                ]
            )
        )
    )
)

Submit the Distributed JAX TPU TrainJob

We submit our JAX TPU job by referencing the jax-distributed ClusterTrainingRuntime and applying our custom TPU scheduling runtime patch.

Note that we configure:

1. resources_per_node requesting "google.com/tpu": 4 to specify the physical TPU chips.
2. JAX platform coordination environment variables: JAX_PLATFORMS="tpu,cpu" and ENABLE_PJRT_COMPATIBILITY="true".
3. A JAX compatible TPU image: us-docker.pkg.dev/cloud-tpu-images/jax-ai-image/tpu:latest.

In [ ]:

num_tpu_per_node = 4
num_nodes = 1

job_name = client.train(
    trainer=CustomTrainer(
        func=jax_tpu_smoke_test,
        image="us-docker.pkg.dev/cloud-tpu-images/jax-ai-image/tpu:latest",
        num_nodes=num_nodes,
        resources_per_node={
            "google.com/tpu": num_tpu_per_node
        },
        env={
            "JAX_PLATFORMS": "tpu,cpu",
            "ENABLE_PJRT_COMPATIBILITY": "true"
        }
    ),
    runtime="jax-distributed",
    options=[tpu_patch]
)

print(f"Training job '{job_name}' submitted successfully for JAX TPU execution!")

Monitor TrainJob Steps & Devices

Wait for the TrainJob to transition into Running status, and check details about scheduled worker Pods, devices, and replica counts.

In [ ]:

# Wait for the job to reach 'Running' status
client.wait_for_job_status(name=job_name, status={"Running"})

# Display step properties and active Pod metadata
for step in client.get_job(name=job_name).steps:
    print(f"Step Name: {step.name}, Pod Name: {step.pod_name}, Status: {step.status}, Device Requested: {step.device}, Device Count: {step.device_count}")

Fetch JAX TPU Coordinator Logs

Retrieve and print logs from all active distributed JAX hosts. You should see JAX successfully initialize PJRT on the physical Cloud TPU backend and display details of the global devices.

In [ ]:

import time

# Give JAX initialization some time to write stdout
time.sleep(10)

# Fetch and print logs from each worker replica
logs = client.get_job_logs(name=job_name)
for worker_name, log_text in logs.items():
    print(f"\n=== Distributed JAX Logs for Worker Replica: {worker_name} ===")
    print(log_text)
    print("=======================================================")

Verify TrainJob Successful Completion

We block and wait for the job to finish with a success state.

In [ ]:

# Wait for job to reach 'Succeeded' status
client.wait_for_job_status(name=job_name, status={"Succeeded"}, timeout=120)
print(f"Distributed JAX TPU TrainJob '{job_name}' finished and verified successfully!")

Clean Up Resources

Once training or verification is finished, clean up the Kubernetes resources by deleting the TrainJob object.

In [ ]:

# client.delete_job(job_name)