Skip to content

Commit 4cbb4b5

Browse files
lukebaumanncopybara-github
authored andcommitted
pathwaysutils: Update README.md with reference implementation purpose.
PiperOrigin-RevId: 949623042
1 parent 1089127 commit 4cbb4b5

9 files changed

Lines changed: 344 additions & 465 deletions

File tree

README.md

Lines changed: 0 additions & 21 deletions
This file was deleted.
Lines changed: 223 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,223 @@
1+
# PathwaysJobSet Reference Implementation
2+
3+
`PathwaysJobSet` is a client-side Python library designed to generate and
4+
deploy Kubernetes `JobSet` resources (`kind: JobSet`) for running Pathways
5+
workloads on Google Kubernetes Engine (GKE).
6+
7+
It serves as a direct replacement for the custom `PathwaysJob` CRD (`kind:
8+
PathwaysJob`), allowing you to deploy Pathways workloads using standard,
9+
community-supported `JobSet` resources without needing a custom controller
10+
installed on your GKE cluster.
11+
12+
---
13+
14+
## Purpose & Overview
15+
16+
This directory serves as the official **reference implementation** for
17+
launching Pathways workloads via Kubernetes `JobSet`:
18+
19+
- **Official Reference Standard:** Other Google teams and external customers
20+
can base their JobSet specifications on this reference implementation and
21+
expect their workloads to work reliably with Pathways.
22+
- **Internal Integration Testing:** Our team uses this reference
23+
implementation as the foundation for all of our automated integration
24+
tests.
25+
- **Minimal Required Specification:** This library defines the baseline,
26+
minimal required Kubernetes specification to run Pathways on Cloud
27+
workloads.
28+
29+
### Production Workloads & Cluster Toolkit
30+
While this reference implementation focuses on the minimal required
31+
specification for Pathways, full-featured production deployments often
32+
require additional Google Cloud Platform (GCP) features and infrastructure
33+
integrations.
34+
35+
- **Bridging the Gap with Cluster Toolkit:** Solutions such as
36+
[Cluster Toolkit](https://cloud.google.com/cluster-toolkit/docs/overview)
37+
bridge the gap between this minimal required implementation and
38+
full-featured, enterprise-grade production workloads.
39+
- **Guidance for Custom Workloads:** Advanced users and customers who need
40+
custom JobSets or desire fine-grained control over their workload
41+
definitions can refer to this `pathwaysutils` JobSet generator to see our
42+
team's recommended best practices for Pathways JobSets, as well as reference
43+
implementations for integrating with other GCP services (such as Cloud
44+
Storage FUSE).
45+
46+
---
47+
48+
## Key Features
49+
50+
- **Client-Side Generation:** Generates standard `JobSet` YAML manifests
51+
that can be inspected, version-controlled, or applied manually.
52+
- **Standard Headless Execution:** Runs the Pathways Resource Manager (RM)
53+
and Proxy as standalone containers in the head job.
54+
- **GCSFuse Integration:** Easily mount Cloud Storage buckets to head
55+
and/or worker pods using GCSFuse.
56+
- **Colocated Python Support:** Inject a colocated Python sidecar container
57+
and shared memory volume into worker pods for multi-agent or hybrid
58+
workloads.
59+
- **Elasticity Support:** Configure elastic slices for dynamic scaling.
60+
61+
---
62+
63+
## Basic Usage
64+
65+
```python
66+
from pathwaysutils.experimental.gke import jobset
67+
68+
# 1. Initialize the builder
69+
pw_jobset = jobset.PathwaysJobSet(
70+
name="my-pathways-workload",
71+
namespace="default",
72+
pathways_dir="gs://my-bucket/pathways-scratch",
73+
tpu_type="v5e",
74+
topology="4x8",
75+
num_slices=1,
76+
)
77+
78+
# 2. Export to a standard JobSet YAML file
79+
pw_jobset.export_yaml("jobset.yaml")
80+
81+
# 3. Deploy directly to a GKE cluster (requires kubernetes configured)
82+
pw_jobset.apply(
83+
project_id="my-gcp-project",
84+
region="us-central1",
85+
cluster_id="my-gke-cluster",
86+
)
87+
```
88+
89+
---
90+
91+
## Examples
92+
93+
### 1. Single-Slice v5e-32 (Non-elastic)
94+
A standard single-slice workload on TPU v5e-32 (32 chips, `4x8` topology, 8
95+
VMs).
96+
97+
```python
98+
from pathwaysutils.experimental.gke import jobset
99+
100+
js = jobset.PathwaysJobSet(
101+
name="v5e-32-headless",
102+
namespace="default",
103+
pathways_dir="gs://my-bucket/scratch",
104+
tpu_type="v5e",
105+
topology="4x8",
106+
num_slices=1,
107+
)
108+
js.export_yaml("v5e_32_headless.yaml")
109+
```
110+
111+
### 2. Multislice v5p-4x4x4 (2 Slices)
112+
A multislice workload using 2 slices of TPU v5p-4x4x4 (64 chips per slice).
113+
114+
```python
115+
from pathwaysutils.experimental.gke import jobset
116+
117+
js = jobset.PathwaysJobSet(
118+
name="v5p-multislice",
119+
namespace="default",
120+
pathways_dir="gs://my-bucket/scratch",
121+
tpu_type="v5p",
122+
topology="4x4x4",
123+
num_slices=2, # 2 slices
124+
)
125+
js.export_yaml("v5p_multislice.yaml")
126+
```
127+
128+
### 3. Multislice v6e with Elasticity
129+
A multislice workload on TPU v6e-16 (topology `4x4`, 16 chips, 4 VMs per
130+
slice) with 2 active slices initially, and elasticity configured for up to 4
131+
slices.
132+
133+
```python
134+
from pathwaysutils.experimental.gke import jobset
135+
136+
js = jobset.PathwaysJobSet(
137+
name="v6e-elastic",
138+
namespace="default",
139+
pathways_dir="gs://my-bucket/scratch",
140+
tpu_type="v6e",
141+
topology="4x4",
142+
num_slices=2, # Initial/Active slices
143+
elastic_slices=4, # Enable elasticity (informs proxy of max slices)
144+
)
145+
js.export_yaml("v6e_elastic.yaml")
146+
```
147+
148+
### 4. Advanced Features: GCSFuse and Colocated Python
149+
This example demonstrates chaining advanced features:
150+
151+
- Mounting a GCS bucket to all containers (head and workers) using GCSFuse.
152+
- Adding a colocated Python sidecar to the worker pods.
153+
154+
```python
155+
from pathwaysutils.experimental.gke import jobset
156+
157+
js = (
158+
jobset.PathwaysJobSet(
159+
name="advanced-workload",
160+
namespace="default",
161+
pathways_dir="gs://my-bucket/scratch",
162+
tpu_type="v5e",
163+
topology="4x8",
164+
num_slices=1,
165+
)
166+
.add_colocated_python() # Injects colocated python sidecar to workers
167+
.add_gcsfuse(
168+
containers="all",
169+
mount_path="/data",
170+
bucket="my-data-bucket",
171+
read_only=True,
172+
)
173+
)
174+
js.export_yaml("advanced_workload.yaml")
175+
```
176+
177+
---
178+
179+
## API Reference: `PathwaysJobSet`
180+
181+
### Constructor `__init__`
182+
183+
```python
184+
def __init__(
185+
self,
186+
name: str,
187+
namespace: str,
188+
pathways_dir: str,
189+
tpu_type: str,
190+
topology: str,
191+
num_slices: int,
192+
max_restarts: int = 0,
193+
max_slice_restarts: int = 0,
194+
termination_grace_period_seconds: int | None = None,
195+
pathways_version: str = "latest",
196+
jobset_api_version: str = "v1alpha2",
197+
elastic_slices: int = 0,
198+
labels: Mapping[str, str] | None = None,
199+
annotations: Mapping[str, str] | None = None,
200+
shared_pathways_service: bool = False,
201+
)
202+
```
203+
204+
- `tpu_type`: Supported values include `"v5e"`, `"v5p"`, `"v6e"`, `"v4"`.
205+
- `elastic_slices`: If `> 0`, enables elasticity by passing
206+
`--num_elastic_slices` to the Pathways Proxy.
207+
- `shared_pathways_service`: If `True`, configures the head job to run only
208+
the Resource Manager (`pathways-rm`).
209+
210+
### Methods
211+
212+
- **`add_colocated_python()`**: Injects a colocated Python container and
213+
shared memory volume (`/tmp/shared-memory`) into worker pods.
214+
- **`add_gcsfuse(containers, mount_path, bucket, read_only=False)`**: Mounts
215+
a GCS bucket using GCSFuse CSI driver. `containers` can be `"head"`,
216+
`"worker"`, or `"all"`.
217+
- **`to_dict()`**: Returns the compiled K8s JobSet resource as a Python
218+
dictionary.
219+
- **`export_yaml(filepath)`**: Serializes and writes the JobSet to a YAML
220+
file.
221+
- **`apply(project_id, region, cluster_id)`**: Deploys the JobSet to the
222+
specified GKE cluster. Supports delete-and-recreate lifecycle if the JobSet
223+
already exists.

0 commit comments

Comments
 (0)