From 5f7c0e979f6b66a754c903232ef160ab6bf74a06 Mon Sep 17 00:00:00 2001 From: Luke Baumann Date: Mon, 26 Oct 2026 17:51:22 -0700 Subject: [PATCH] Import external READMEs and update PathwaysJobSet documentation. PiperOrigin-RevId: 1007692294 --- README.md | 10 +- pathwaysutils/experimental/gke/README.md | 105 ++++++++++++++++++ .../shared_pathways_service/README.md | 7 +- pathwaysutils/extras/elastic/README.md | 27 +++++ 4 files changed, 135 insertions(+), 14 deletions(-) create mode 100644 pathwaysutils/experimental/gke/README.md create mode 100644 pathwaysutils/extras/elastic/README.md diff --git a/README.md b/README.md index 9a4334a..b309e76 100644 --- a/README.md +++ b/README.md @@ -1,4 +1,4 @@ -# Package of Pathways-on-Cloud utilities +# APIs to get started with using Pathways on Cloud Pathways is a novel compute orchestration method that enables single controller runtime ML workloads. This approach unlocks a range of new capabilities for training resiliency, multihost inference and interactive development on cluster of arbitrary scale. @@ -10,12 +10,6 @@ This package handles the underlying adaptations required for the cloud environme # Further Reading [Pathways Research Paper](https://arxiv.org/abs/2203.12533) -[More resources coming Soon] - - -# How to get access to Pathways on Cloud - -Pathways on Cloud is currently in Preview. -Please contact your Google Cloud account representative to request access to Pathways. +[Get started with Pathways on Google Cloud](https://docs.cloud.google.com/ai-hypercomputer/docs/workloads/pathways-on-cloud/pathways-intro) diff --git a/pathwaysutils/experimental/gke/README.md b/pathwaysutils/experimental/gke/README.md new file mode 100644 index 0000000..416df3e --- /dev/null +++ b/pathwaysutils/experimental/gke/README.md @@ -0,0 +1,105 @@ +# PathwaysJobSet Builder + +`PathwaysJobSet` generates and deploys Kubernetes `JobSet` resources (`jobset.x-k8s.io`) for running distributed Pathways workloads on Google Kubernetes Engine (GKE) without requiring custom controllers. + +## Installation + +```bash +pip install 'pathwaysutils[gke]' +``` + +## Quickstart + +```python +from pathwaysutils.experimental.gke import jobset + +pw_jobset = jobset.PathwaysJobSet( + name="my-workload", + namespace="default", + pathways_dir="gs://my-bucket/pathways-scratch", + tpu_type="v5e", + topology="2x4", + num_slices=1, +) + +# Export to YAML or deploy directly to the active cluster +pw_jobset.export_yaml("jobset.yaml") +pw_jobset.apply(recreate=True) +``` + +## Features + +### Elasticity & Multislice +Scale across multiple TPU slices with dynamic resilience: + +```python +pw_jobset = jobset.PathwaysJobSet( + name="elastic-v6e", + namespace="default", + pathways_dir="gs://my-bucket/scratch", + tpu_type="v6e", + topology="4x4", + num_slices=4, + elastic_slices=2, # Configures Pathways Proxy for dynamic slice recovery +) +``` + +### In-Pod User Workload +Run training or inference directly in the head pod alongside the Resource Manager and Proxy sidecars: + +```python +pw_jobset.add_user_workload( + image="us-docker.pkg.dev/my-project/images/my-trainer:latest", + command="python3 train.py --batch_size=32", + env={"LOG_LEVEL": "INFO"}, +) +``` + +### Colocated Python Sidecar +Inject a sidecar container onto worker pods sharing memory with `pathways-worker` for fast tensor transfers: + +```python +pw_jobset.add_colocated_python( + image="us-docker.pkg.dev/my-project/images/worker-sidecar:latest", + shm_mount_path="/tmp/shared-memory", + shm_size_limit="16Gi", +) +``` + +### GCSFuse Cloud Storage Mounts +Mount GCS buckets to head, worker, or all containers: + +```python +pw_jobset.add_gcsfuse( + containers="all", + mount_path="/data", + bucket="my-data-bucket", + read_only=True, +) +``` + +### Container & Pod Spec Overrides +Customize arguments, environment variables, resources, or volumes on any container or pod: + +```python +pw_jobset = jobset.PathwaysJobSet( + ..., + rm_container_kwargs={"args": ["--v=1"]}, + proxy_container_kwargs={"env": {"XLA_FLAGS": "--xla_dump_to=/tmp"}}, + worker_container_kwargs={"args": ["--stderrthreshold=0"]}, + worker_pod_spec_kwargs={"priority_class_name": "high-priority"}, +) +``` + +## Serialization & Round-Tripping + +Convert between Python objects, YAML files, and dictionaries: + +```python +# Save and load JobSet configurations +pw_jobset.export_yaml("jobset.yaml") +loaded_jobset = jobset.PathwaysJobSet.import_yaml("jobset.yaml") + +# Inspect raw Kubernetes resource dictionary +config_dict = pw_jobset.to_dict() +``` diff --git a/pathwaysutils/experimental/shared_pathways_service/README.md b/pathwaysutils/experimental/shared_pathways_service/README.md index f191ee6..5e86c1b 100644 --- a/pathwaysutils/experimental/shared_pathways_service/README.md +++ b/pathwaysutils/experimental/shared_pathways_service/README.md @@ -30,11 +30,6 @@ python3 -m pathwaysutils.experimental.shared_pathways_service.deploy_pathways_se Note, `server_image` is the Pathways server image, e.g, `us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:`. -`--sidecar_image` is optional. Set it to deploy the colocated Python sidecar -next to the TPU workers, which workloads need to use -`--proxy_options=sidecar:true`. Without it, the service is deployed without the -sidecar. - ### 3. Verify that the pods created in [Step#2](#2-deploy-the-pathways-head-pod) are running Verify that the Shared Pathways Service components are started, specifically the Pathways resource manager (RM) and @@ -78,7 +73,7 @@ $ WORKER0_POD_NAME=$(kubectl get pods --selector=jobset.sigs.k8s.io/jobset-name= #### Option 3: Check project logs Find the detailed instructions -here). +here. ### 4. Find the Pathways service address diff --git a/pathwaysutils/extras/elastic/README.md b/pathwaysutils/extras/elastic/README.md new file mode 100644 index 0000000..a5d2d77 --- /dev/null +++ b/pathwaysutils/extras/elastic/README.md @@ -0,0 +1,27 @@ +# Recovery Mode Simulation & Comparison + +Evaluating distributed machine learning training on elastic TPU topologies requires understanding how varying failure scenarios and reinitialization latencies affect cumulative goodput and step progression. This folder provides an interactive browser-based simulation comparing three distinct recovery paradigms: McJAX, Snapshot Replica Resize, and Snapshot Pause Resume. + +## Live Interactive Simulation + +* [Hosted GitHub Pages Simulation](https://ai-hypercomputer.github.io/pathways-utils/pathwaysutils/extras/elastic/model_comparison.html) is a live interactive web application with simulation controls, or [HTMLPreview Instant Mirror](https://htmlpreview.github.io/?https://github.com/AI-Hypercomputer/pathways-utils/blob/main/pathwaysutils/extras/elastic/model_comparison.html). +* Local browser viewing allows running `model_comparison.html` directly on your workstation by opening the file in Chrome, Firefox, or Safari, requiring zero installation or server dependencies. + +## Recovery Paradigms Compared + +* Traditional Checkpoint Rollback (McJAX) halts cluster computation upon slice preemption, incurs a reinitialization latency, and rewinds training state to the most recent periodic checkpoint or snapshot, resulting in lost forward progress. +* Snapshot Replica Resize dynamically adjusts the mesh to continue training on surviving healthy slices with a smaller reinitialization latency than a cold start, maintaining forward progress at a degraded throughput factor (depending on your recovery implementation) with no time spent waiting for slices to rejoin. +* Snapshot Pause Resume pauses execution during slice preemption, rewinds to the last snapshot, and resumes with negligible reinitialization once replacement slices become healthy. + +## Interactive Parameter Controls + +The simulation toolbar allows tuning key workload and infrastructure characteristics in real time: + +* Cold Start Initialization Time (seconds) configures the initial cluster startup delay prior to the first training step and is reused as the reinitialization latency for McJAX. +* Scale-Down Initialization Time (seconds) sets the reconfiguration overhead when transitioning into degraded execution. +* Scale-Up Initialization Time (seconds) governs the mesh reconfiguration latency when a replacement slice joins the workload. +* Degraded Scaling Factor (multiplier) sets the relative step throughput while executing on a reduced slice topology. +* Initial Base Loss (L₀) sets the initial loss baseline at step 0 in the exponential decay function L = L₀ * exp(-α * step). +* Loss Decay Rate (α) dictates the per-step exponential decay rate, multiplied by the degraded scaling factor during degraded execution so training loss decays slower per step. +* Checkpoint Interval (steps) dictates the rollback penalty for checkpoint-based training. +* Interruption Duration for Scenario (seconds) configures hardware disruption length during automated scenario runs.