mirror of
https://github.com/hanzo-ml/trainer.git
synced 2026-09-24 00:31:34 +00:00
feat(examples): add interactive JAX distributed TPU smoke test notebook (#3531)
Signed-off-by: Richa Banker <[email protected]>
This commit is contained in:
Richa Banker
@@ -0,0 +1,315 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "48dd25df",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# Distributed JAX TPU Smoke Test with Kubeflow Trainer\n",
|
||||
"\n",
|
||||
"This notebook demonstrates how to run a distributed JAX verification job on physical Cloud TPUs using the Kubeflow Trainer Python SDK.\n",
|
||||
"\n",
|
||||
"It initializes a distributed JAX environment, verifies hardware discovery across all physical TPU chips, and executes a basic multi-host matrix multiplication."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "b4af3557",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Install the Kubeflow SDK\n",
|
||||
"\n",
|
||||
"You need to install the Kubeflow SDK to interact with Kubeflow Trainer."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "a0434dd0",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"# !pip install -U kubeflow"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "c7b57327",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Define JAX Distributed Verification Function\n",
|
||||
"\n",
|
||||
"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."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "4a05ab4c",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"def jax_tpu_smoke_test():\n",
|
||||
" import jax\n",
|
||||
" import jax.numpy as jnp\n",
|
||||
"\n",
|
||||
" # Initialize distributed JAX coordination\n",
|
||||
" jax.distributed.initialize()\n",
|
||||
"\n",
|
||||
" # Log device count\n",
|
||||
" print(\"==========================\")\n",
|
||||
" print(f\"Total global TPU device count: {jax.device_count()}\")\n",
|
||||
" print(f\"Local TPU device count: {jax.local_device_count()}\")\n",
|
||||
" print(f\"Global TPU devices details: {jax.devices()}\")\n",
|
||||
" print(\"==========================\")\n",
|
||||
"\n",
|
||||
" # Execute distributed matrix multiplication to verify JAX TPU backend\n",
|
||||
" x = jnp.ones((1000, 1000))\n",
|
||||
" y = (x @ x).block_until_ready()\n",
|
||||
" \n",
|
||||
" print(f\"JAX Matrix Multiplication verification successfully completed! Result shape: {y.shape}\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "2481ec3e",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Initialize Kubeflow Trainer Client\n",
|
||||
"\n",
|
||||
"We initialize the `TrainerClient` which communicates with the Kubeflow Trainer API on the Kubernetes cluster."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "591048ad",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"from kubeflow.trainer import TrainerClient, CustomTrainer\n",
|
||||
"\n",
|
||||
"client = TrainerClient()"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "183e82e4",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Configure GKE TPU Scheduling Patches\n",
|
||||
"\n",
|
||||
"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."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "cd87905d",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"from kubeflow.trainer.options import RuntimePatch\n",
|
||||
"from kubeflow.trainer.options.kubernetes import (\n",
|
||||
" TrainingRuntimeSpecPatch,\n",
|
||||
" JobSetTemplatePatch,\n",
|
||||
" JobSetSpecPatch,\n",
|
||||
" ReplicatedJobPatch,\n",
|
||||
" JobTemplatePatch,\n",
|
||||
" JobSpecPatch,\n",
|
||||
" PodTemplatePatch,\n",
|
||||
" PodSpecPatch,\n",
|
||||
")\n",
|
||||
"\n",
|
||||
"# We define a GKE scheduling patch using typed SDK dataclasses\n",
|
||||
"tpu_patch = RuntimePatch(\n",
|
||||
" training_runtime_spec=TrainingRuntimeSpecPatch(\n",
|
||||
" template=JobSetTemplatePatch(\n",
|
||||
" spec=JobSetSpecPatch(\n",
|
||||
" replicated_jobs=[\n",
|
||||
" ReplicatedJobPatch(\n",
|
||||
" name=\"node\",\n",
|
||||
" template=JobTemplatePatch(\n",
|
||||
" spec=JobSpecPatch(\n",
|
||||
" template=PodTemplatePatch(\n",
|
||||
" spec=PodSpecPatch(\n",
|
||||
" node_selector={\n",
|
||||
" \"cloud.google.com/gke-tpu-accelerator\": \"tpu-v5-lite-podslice\",\n",
|
||||
" \"cloud.google.com/gke-tpu-topology\": \"2x2\",\n",
|
||||
" },\n",
|
||||
" tolerations=[\n",
|
||||
" {\n",
|
||||
" \"key\": \"google.com/tpu\",\n",
|
||||
" \"operator\": \"Exists\",\n",
|
||||
" \"effect\": \"NoSchedule\",\n",
|
||||
" }\n",
|
||||
" ],\n",
|
||||
" )\n",
|
||||
" )\n",
|
||||
" )\n",
|
||||
" )\n",
|
||||
" )\n",
|
||||
" ]\n",
|
||||
" )\n",
|
||||
" )\n",
|
||||
" )\n",
|
||||
")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "56ba99df",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Submit the Distributed JAX TPU TrainJob\n",
|
||||
"\n",
|
||||
"We submit our JAX TPU job by referencing the `jax-distributed` ClusterTrainingRuntime and applying our custom TPU scheduling runtime patch.\n",
|
||||
"\n",
|
||||
"Note that we configure:\n",
|
||||
"1. `resources_per_node` requesting `\"google.com/tpu\": 4` to specify the physical TPU chips.\n",
|
||||
"2. JAX platform coordination environment variables: `JAX_PLATFORMS=\"tpu,cpu\"` and `ENABLE_PJRT_COMPATIBILITY=\"true\"`.\n",
|
||||
"3. A JAX compatible TPU image: `us-docker.pkg.dev/cloud-tpu-images/jax-ai-image/tpu:latest`."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "d5243de1",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"num_tpu_per_node = 4\n",
|
||||
"num_nodes = 1\n",
|
||||
"\n",
|
||||
"job_name = client.train(\n",
|
||||
" trainer=CustomTrainer(\n",
|
||||
" func=jax_tpu_smoke_test,\n",
|
||||
" image=\"us-docker.pkg.dev/cloud-tpu-images/jax-ai-image/tpu:latest\",\n",
|
||||
" num_nodes=num_nodes,\n",
|
||||
" resources_per_node={\n",
|
||||
" \"google.com/tpu\": num_tpu_per_node\n",
|
||||
" },\n",
|
||||
" env={\n",
|
||||
" \"JAX_PLATFORMS\": \"tpu,cpu\",\n",
|
||||
" \"ENABLE_PJRT_COMPATIBILITY\": \"true\"\n",
|
||||
" }\n",
|
||||
" ),\n",
|
||||
" runtime=\"jax-distributed\",\n",
|
||||
" options=[tpu_patch]\n",
|
||||
")\n",
|
||||
"\n",
|
||||
"print(f\"Training job '{job_name}' submitted successfully for JAX TPU execution!\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "2a037675",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Monitor TrainJob Steps & Devices\n",
|
||||
"\n",
|
||||
"Wait for the TrainJob to transition into `Running` status, and check details about scheduled worker Pods, devices, and replica counts."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "70deefeb",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"# Wait for the job to reach 'Running' status\n",
|
||||
"client.wait_for_job_status(name=job_name, status={\"Running\"})\n",
|
||||
"\n",
|
||||
"# Display step properties and active Pod metadata\n",
|
||||
"for step in client.get_job(name=job_name).steps:\n",
|
||||
" print(f\"Step Name: {step.name}, Pod Name: {step.pod_name}, Status: {step.status}, Device Requested: {step.device}, Device Count: {step.device_count}\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "41677279",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Fetch JAX TPU Coordinator Logs\n",
|
||||
"\n",
|
||||
"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."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "b410f857",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"import time\n",
|
||||
"\n",
|
||||
"# Give JAX initialization some time to write stdout\n",
|
||||
"time.sleep(10)\n",
|
||||
"\n",
|
||||
"# Fetch and print logs from each worker replica\n",
|
||||
"logs = client.get_job_logs(name=job_name)\n",
|
||||
"for worker_name, log_text in logs.items():\n",
|
||||
" print(f\"\\n=== Distributed JAX Logs for Worker Replica: {worker_name} ===\")\n",
|
||||
" print(log_text)\n",
|
||||
" print(\"=======================================================\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "4162b7ac",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Verify TrainJob Successful Completion\n",
|
||||
"\n",
|
||||
"We block and wait for the job to finish with a success state."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "26cab010",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"# Wait for job to reach 'Succeeded' status\n",
|
||||
"client.wait_for_job_status(name=job_name, status={\"Succeeded\"}, timeout=120)\n",
|
||||
"print(f\"Distributed JAX TPU TrainJob '{job_name}' finished and verified successfully!\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "5b926dfd",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Clean Up Resources\n",
|
||||
"\n",
|
||||
"Once training or verification is finished, clean up the Kubernetes resources by deleting the TrainJob object."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "13759389",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"# client.delete_job(job_name)"
|
||||
]
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"kernelspec": {
|
||||
"display_name": "Python 3",
|
||||
"language": "python",
|
||||
"name": "python3"
|
||||
},
|
||||
"language_info": {
|
||||
"name": "python",
|
||||
"version": "3.11.15"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 5
|
||||
}
|
||||
Reference in New Issue
Block a user