Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
188 changes: 188 additions & 0 deletions pathwaysutils/experimental/gke/README.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,188 @@
# PathwaysJobSet Builder

`PathwaysJobSet` is a client-side Python library designed to generate and
deploy Kubernetes `JobSet` resources (`kind: JobSet`) for running Pathways
workloads on GKE.

It serves as a direct replacement for the custom `PathwaysJob` CRD (`kind:
PathwaysJob`), allowing you to deploy Pathways workloads using standard,
community-supported `JobSet` resources without needing a custom controller
installed on your GKE cluster.

---

## Key Features

- **Client-Side Generation:** Generates standard `JobSet` YAMLs that can be
inspected, version-controlled, or applied manually.
- **Standard Headless Execution:** Runs the Pathways Resource Manager (RM)
and Proxy as standalone containers in the head job.
- **GCSFuse Integration:** Easily mount GCS buckets to head and/or worker
pods.
- **Colocated Python Support:** Inject a colocated Python sidecar container
into worker pods for multi-agent or hybrid workloads.
- **Elasticity Support:** Configure elastic slices for dynamic scaling.

---

## Basic Usage

```python
from pathwaysutils.experimental.gke import jobset

# 1. Initialize the builder
pw_jobset = jobset.PathwaysJobSet(
name="my-pathways-workload",
namespace="default",
pathways_dir="gs://my-bucket/pathways-scratch",
tpu_type="v5e",
topology="4x8",
num_slices=1,
)

# 2. Export to a standard JobSet YAML file
pw_jobset.export_yaml("jobset.yaml")

# 3. Deploy directly to a GKE cluster (requires kubernetes configured)
pw_jobset.apply(
project_id="my-gcp-project",
region="us-central1",
cluster_id="my-gke-cluster",
)
```

---

## Examples

### 1. Single-Slice v5e-32 (Non-elastic)
A standard single-slice workload on TPU v5e-32 (32 chips, `4x8` topology, 8
VMs).

```python
from pathwaysutils.experimental.gke import jobset

js = jobset.PathwaysJobSet(
name="v5e-32-headless",
namespace="default",
pathways_dir="gs://my-bucket/scratch",
tpu_type="v5e",
topology="4x8",
num_slices=1,
)
js.export_yaml("v5e_32_headless.yaml")
```

### 2. Multislice v5p-4x4x4 (2 Slices)
A multislice workload using 2 slices of TPU v5p-4x4x4 (64 chips per slice).

```python
from pathwaysutils.experimental.gke import jobset

js = jobset.PathwaysJobSet(
name="v5p-multislice",
namespace="default",
pathways_dir="gs://my-bucket/scratch",
tpu_type="v5p",
topology="4x4x4",
num_slices=2, # 2 slices
)
js.export_yaml("v5p_multislice.yaml")
```

### 3. Multislice v6e with Elasticity
A multislice workload on TPU v6e-16 (topology `4x4`, 16 chips, 4 VMs per
slice) with 2 active slices initially, and elasticity configured for up to 4
slices.

```python
from pathwaysutils.experimental.gke import jobset

js = jobset.PathwaysJobSet(
name="v6e-elastic",
namespace="default",
pathways_dir="gs://my-bucket/scratch",
tpu_type="v6e",
topology="4x4",
num_slices=2, # Initial/Active slices
elastic_slices=4, # Enable elasticity (informs proxy of max slices)
)
js.export_yaml("v6e_elastic.yaml")
```

### 4. Advanced Features: GCSFuse and Colocated Python
This example demonstrates chaining advanced features:

- Mounting a GCS bucket to all containers (head and workers) using GCSFuse.
- Adding a colocated Python sidecar to the worker pods.

```python
from pathwaysutils.experimental.gke import jobset

js = (
jobset.PathwaysJobSet(
name="advanced-workload",
namespace="default",
pathways_dir="gs://my-bucket/scratch",
tpu_type="v5e",
topology="4x8",
num_slices=1,
)
.add_colocated_python() # Injects colocated python sidecar to workers
.add_gcsfuse(
containers="all",
mount_path="/data",
bucket="my-data-bucket",
read_only=True,
)
)
js.export_yaml("advanced_workload.yaml")
```

---

## API Reference: `PathwaysJobSet`

### Constructor `__init__`

```python
def __init__(
self,
name: str,
namespace: str,
pathways_dir: str,
tpu_type: str,
topology: str,
num_slices: int,
max_restarts: int = 0,
max_slice_restarts: int = 0,
termination_grace_period_seconds: int | None = None,
pathways_version: str = "latest",
jobset_api_version: str = "v1alpha2",
elastic_slices: int = 0,
labels: Mapping[str, str] | None = None,
annotations: Mapping[str, str] | None = None,
shared_pathways_service: bool = False,
)
```

- `tpu_type`: Supported values include `"v5e"`, `"v5p"`, `"v6e"`, `"v4"`.
- `elastic_slices`: If `> 0`, enables elasticity by passing
`--num_elastic_slices` to the Pathways Proxy.
- `shared_pathways_service`: If `True`, configures the head job to run only
the Resource Manager (`pathways-rm`).

### Methods

- **`add_colocated_python()`**: Injects a colocated Python container and
shared memory volume (`/tmp/shared-memory`) into worker pods.
- **`add_gcsfuse(containers, mount_path, bucket, read_only=False)`**: Mounts
a GCS bucket using GCSFuse CSI driver. `containers` can be `"head"`,
`"worker"`, or `"all"`.
- **`to_dict()`**: Returns the compiled K8s JobSet resource as a Python
dictionary.
- **`export_yaml(filepath)`**: Serializes and writes the JobSet to a YAML
file.
- **`apply(project_id, region, cluster_id)`**: Deploys the JobSet to the
specified GKE cluster. Supports delete-and-recreate lifecycle if the JobSet
already exists.
Loading
Loading