pathways-cli 0.1.0__py3-none-any.whl

pathways_cli-0.1.0.dist-info/METADATA

@@ -0,0 +1,122 @@
+Metadata-Version: 2.4
+Name: pathways-cli
+Version: 0.1.0
+Summary: Pathways CLI to easily bring up pathways clusters.
+Author-email: Sam Stoelinga <sammiestoel@gmail.com>
+Requires-Python: >=3.12
+Requires-Dist: click>=8.4.1
+Requires-Dist: python-dotenv>=1.2.2
+Description-Content-Type: text/markdown
+
+# `pwy`: Standalone Pathways GKE Cluster CLI Tool
+
+`pwy` is a lightweight, standalone Python CLI utility designed to generate, apply, and manage interactive Pathways workloads on Google Kubernetes Engine (GKE) using Kubernetes JobSets.
+
+---
+
+## Features
+
+- **Automated TPU Topology Calculations**: Translates simple TPU resource types (`v6e-4`, `v6e-16`, etc.) into GKE topologies, VM counts, and instance settings.
+- **Spot VM Support**: Dynamically injects GKE node selectors and tolerations for running workloads on cost-effective Spot VMs.
+- **Colocated Python Support**: Simplifies distributed checkpointing (e.g. via Orbax) by configuring and enabling colocated host CPU sidecars and proxy endpoints automatically.
+- **Interactive & Batch Execution**: Supports spinning up pathways servers with infinite sleep drivers for interactive debugging, or executing training scripts directly.
+- **Dry-run Manifest Generation**: Preview and inspect the GKE JobSet manifest without applying it to the cluster.
+
+---
+
+## Installation
+
+This project utilizes [uv](https://github.com/astral-sh/uv) for fast, modern Python package and dependency management.
+
+To sync the environment and install `pwy`:
+
+```bash
+uv sync
+```
+
+---
+
+## Usage
+
+You can invoke `pwy` commands directly using `uv run`:
+
+### 1. Provision / Preview a Cluster (`pwy up`)
+
+Starts a Pathways JobSet or dry-runs the configuration.
+
+```bash
+uv run pwy up \
+  --tpu-type v6e-16 \
+  --gcs-scratch-location gs://my-bucket/pathways-staging \
+  --num-slices 1 \
+  --dry-run
+```
+
+#### Key Options:
+- `--tpu-type`: **(Required)** TPU type (e.g., `v6e-4`, `v6e-8`, `v6e-16`, `v6e-32`, `v6e-64`).
+- `--gcs-scratch-location`: **(Required)** GCS scratch path for pathways synchronization.
+- `--num-slices`: Number of TPU slices to run (default: `1`).
+- `--jax-client-image`: Custom client container image (default: `python:3.12-slim`).
+- `--command`: Run a custom training/eval script in the client container. If omitted, defaults to `sleep infinity` (interactive mode).
+- `--enable-spot`: Add node affinity and toleration settings for Spot VMs.
+- `--colocated-python`: Enables colocated CPU Python sidecar/init containers on GKE workers and enables external proxy routing.
+- `--dry-run`: Prints the generated YAML to stdout instead of calling `kubectl apply`.
+- `--name`: Name of the Kubernetes JobSet resource (default: `pathways-interactive`).
+- `--namespace`: Target Kubernetes namespace (default: `default`).
+
+---
+
+### 2. Teardown a Cluster (`pwy down`)
+
+Deletes the running Pathways JobSet.
+
+```bash
+uv run pwy down --name pathways-interactive --namespace default
+```
+
+---
+
+### 3. Verification Example
+
+Once the interactive cluster is running, you can verify execution by `exec`ing into the client container:
+
+1. **Find the client pod name**:
+   ```bash
+   POD_NAME=$(kubectl get pods -l jobset.sigs.k8s.io/jobset-name=pathways-interactive -o jsonpath='{.items[?(@.metadata.labels.jobset\\.sigs\\.k8s\\.io/replicatedjob-name=="pwhd")].metadata.name}')
+   ```
+
+2. **Install JAX and Pathways utils**:
+   ```bash
+   kubectl exec $POD_NAME -c client -- pip install jax pathwaysutils
+   ```
+
+3. **Run a Python snippet to initialize and list devices**:
+   ```bash
+   kubectl exec $POD_NAME -c client -- python3 -c "import pathwaysutils; pathwaysutils.initialize(); import jax; print(jax.devices())"
+   ```
+
+   The command output should print the available virtual TPU devices (e.g., coordinates and memory spaces of the allocated chips).
+
+---
+
+## TPU Type Mappings
+
+`pwy` handles all resource-limit math and topologies automatically according to the following matrix:
+
+| TPU Type | GKE Topology | VMs Per Slice | RM Instance Type |
+| :--- | :--- | :--- | :--- |
+| `v6e-4` | `2x2` | 1 | `tpuv6e:2x2` |
+| `v6e-8` | `2x4` | 2 | `tpuv6e:2x4` |
+| `v6e-16` | `4x4` | 4 | `tpuv6e:4x4` |
+| `v6e-32` | `4x8` | 8 | `tpuv6e:4x8` |
+| `v6e-64` | `8x8` | 16 | `tpuv6e:8x8` |
+
+---
+
+## Running Tests
+
+To execute the unit test suite:
+
+```bash
+uv run pytest
+```

pathways_cli-0.1.0.dist-info/RECORD

@@ -0,0 +1,9 @@
+pwy/__init__.py,sha256=FjXvQgxYlSpOMIL8zRR5_XdoZtveGIRjiMiae3yO3aE,45
+pwy/cli.py,sha256=jLyjjz68kdDiZ2Mqt63wm0Ty8XmKcc63Uatt9WiXsKE,3258
+pwy/generator.py,sha256=ZSelmMDOAcwI-ooD458ILPh42BGs3bOESmFT--Hr84k,4915
+pwy/kubernetes.py,sha256=yxxCHVaQVwvZkppgxRMSoZNRv8RC-V8lRKupmtD1GJM,646
+pwy/templates.py,sha256=g6NlFzhbBzcMudLrYmJsF0aboCUulgYAUON0GSTbygk,7750
+pathways_cli-0.1.0.dist-info/METADATA,sha256=DphQjqzCWOD1FkeWH6VqUE9JDwN-hJcT3_uA0896nBI,4288
+pathways_cli-0.1.0.dist-info/WHEEL,sha256=QccIxa26bgl1E6uMy58deGWi-0aeIkkangHcxk2kWfw,87
+pathways_cli-0.1.0.dist-info/entry_points.txt,sha256=2GvpxKRrXF7mDFvGxcsR_2ut2eZHzAlLkgAo_qi5NME,33
+pathways_cli-0.1.0.dist-info/RECORD,,

pathways_cli-0.1.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.29.0
+Root-Is-Purelib: true
+Tag: py3-none-any

pathways_cli-0.1.0.dist-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+pwy = pwy:main

pwy/__init__.py

@@ -0,0 +1,3 @@
+from pwy.cli import main
+
+__all__ = ["main"]

pwy/cli.py

@@ -0,0 +1,68 @@
+import sys
+import click
+from pwy.generator import generate_yaml
+from pwy.kubernetes import apply_manifest, delete_jobset
+
+@click.group()
+def main():
+    """pwy: Standalone Pathways GKE Cluster CLI Tool"""
+    pass
+
+@main.command()
+@click.option("--tpu-type", required=True, help="TPU type (e.g., v6e-4, v6e-8, v6e-16, v6e-32, v6e-64)")
+@click.option("--gcs-scratch-location", required=True, help="GCS scratch location (e.g., gs://bucket/staging)")
+@click.option("--num-slices", default=1, type=int, show_default=True, help="Number of TPU slices")
+@click.option("--jax-client-image", default="python:3.12-slim", show_default=True, help="Image for the JAX client container")
+@click.option("--command", default=None, help="Command to run in the JAX client container (defaults to sleep infinity)")
+@click.option("--enable-spot", is_flag=True, default=False, help="Enable spot VM scheduling")
+@click.option("--colocated-python", is_flag=True, default=False, help="Enable colocated python sidecars")
+@click.option("--dry-run", is_flag=True, default=False, help="Dry run: print generated YAML to stdout instead of applying it")
+@click.option("--name", default="pathways-interactive", show_default=True, help="Name of the JobSet resource")
+@click.option("--namespace", default="default", show_default=True, help="Kubernetes namespace")
+def up(tpu_type, gcs_scratch_location, num_slices, jax_client_image, command, enable_spot, colocated_python, dry_run, name, namespace):
+    """Starts the Pathways cluster or dry-runs the configuration."""
+    try:
+        yaml_content = generate_yaml(
+            name=name,
+            namespace=namespace,
+            tpu_type=tpu_type,
+            gcs_scratch_location=gcs_scratch_location,
+            num_slices=num_slices,
+            jax_client_image=jax_client_image,
+            command=command,
+            enable_spot=enable_spot,
+            colocated_python=colocated_python,
+        )
+    except ValueError as e:
+        click.secho(f"Error: {e}", fg="red", err=True)
+        sys.exit(1)
+
+    if dry_run:
+        click.echo(yaml_content)
+        return
+
+    click.echo(f"Applying Pathways JobSet manifest for '{name}' in namespace '{namespace}'...")
+    process = apply_manifest(yaml_content)
+    if process.returncode != 0:
+        click.secho("Failed to apply JobSet manifest.", fg="red", err=True)
+        click.echo(process.stderr.decode("utf-8"), err=True)
+        sys.exit(process.returncode)
+
+    click.secho(f"Successfully applied JobSet '{name}'!", fg="green")
+
+@main.command()
+@click.option("--name", default="pathways-interactive", show_default=True, help="Name of the JobSet resource")
+@click.option("--namespace", default="default", show_default=True, help="Kubernetes namespace")
+def down(name, namespace):
+    """Tears down the Pathways cluster JobSet resource."""
+    click.echo(f"Deleting Pathways JobSet '{name}' in namespace '{namespace}'...")
+    process = delete_jobset(name, namespace)
+    if process.returncode != 0:
+        click.secho("Failed to delete JobSet.", fg="red", err=True)
+        click.echo(process.stderr.decode("utf-8"), err=True)
+        sys.exit(process.returncode)
+
+    click.secho(f"Successfully deleted JobSet '{name}'!", fg="green")
+
+if __name__ == "__main__":
+    main()

pwy/generator.py

@@ -0,0 +1,122 @@
+from pwy.templates import YAML_TEMPLATE
+
+TPU_MAPPINGS = {
+    "v6e-4": {"topology": "2x2", "vms_per_slice": 1, "rm_type": "tpuv6e:2x2"},
+    "v6e-8": {"topology": "2x4", "vms_per_slice": 2, "rm_type": "tpuv6e:2x4"},
+    "v6e-16": {"topology": "4x4", "vms_per_slice": 4, "rm_type": "tpuv6e:4x4"},
+    "v6e-32": {"topology": "4x8", "vms_per_slice": 8, "rm_type": "tpuv6e:4x8"},
+    "v6e-64": {"topology": "8x8", "vms_per_slice": 16, "rm_type": "tpuv6e:8x8"},
+}
+
+def get_colocated_python_image(client_image: str) -> str:
+    if "/" in client_image and ":" in client_image:
+        try:
+            path, tag = client_image.rsplit(":", 1)
+            repo, _ = path.rsplit("/", 1)
+            return f"{repo}/colocated-python:{tag}"
+        except Exception:
+            pass
+    return "us-docker.pkg.dev/cloud-tpu-v2-images/pathways/colocated-python:jax-0.10.0"
+
+def generate_yaml(
+    name: str,
+    namespace: str,
+    tpu_type: str,
+    gcs_scratch_location: str,
+    num_slices: int = 1,
+    jax_client_image: str = "python:3.12-slim",
+    command: str = None,
+    enable_spot: bool = False,
+    colocated_python: bool = False,
+) -> str:
+    if tpu_type not in TPU_MAPPINGS:
+        raise ValueError(
+            f"Unsupported TPU type: {tpu_type}. Supported types: {list(TPU_MAPPINGS.keys())}"
+        )
+    
+    mapping = TPU_MAPPINGS[tpu_type]
+    gke_topology = mapping["topology"]
+    vms_per_slice = mapping["vms_per_slice"]
+    rm_instance_type = mapping["rm_type"]
+    
+    # Format client execution command
+    if not command:
+        client_command = "sleep infinity"
+    else:
+        client_command = command
+    
+    # Format Spot VM Node Selector and Tolerations
+    if enable_spot:
+        spot_toleration_head = (
+            '                - key: "cloud.google.com/gke-spot"\n'
+            '                  operator: "Equal"\n'
+            '                  value: "true"\n'
+            '                  effect: "NoSchedule"'
+        )
+        spot_node_selector_worker = '                cloud.google.com/gke-spot: "true"'
+        spot_toleration_worker = (
+            '                - key: "cloud.google.com/gke-spot"\n'
+            '                  operator: "Equal"\n'
+            '                  value: "true"\n'
+            '                  effect: "NoSchedule"'
+        )
+    else:
+        spot_toleration_head = ""
+        spot_node_selector_worker = ""
+        spot_toleration_worker = ""
+    
+    # Format colocated python options
+    if colocated_python:
+        proxy_sidecar_arg = "\n                    - --sidecar_name=external"
+        tpu_premapped_buffer_size = 34359738368  # 32 GiB
+        colocated_img = get_colocated_python_image(jax_client_image)
+        worker_init_containers = (
+            "              initContainers:\n"
+            "                - name: colocated-python\n"
+            f"                  image: {colocated_img}\n"
+            "                  imagePullPolicy: Always\n"
+            "                  restartPolicy: Always\n"
+            "                  ports:\n"
+            "                    - containerPort: 50051\n"
+            "                      protocol: TCP\n"
+            "                  env:\n"
+            "                    - name: CLOUD_PATHWAYS_SIDECAR_SHM_DIRECTORY\n"
+            "                      value: /tmp/ifrt_proxy\n"
+            "                    - name: GRPC_SERVER_ADDRESS\n"
+            "                      value: 0.0.0.0:50051\n"
+            "                  volumeMounts:\n"
+            "                    - name: shared-memory\n"
+            "                      mountPath: /tmp/ifrt_proxy"
+        )
+    else:
+        proxy_sidecar_arg = ""
+        tpu_premapped_buffer_size = 274877906944  # 256 GiB
+        worker_init_containers = ""
+    
+    # Interpolate variables in the template
+    yaml_content = YAML_TEMPLATE.format(
+        NAME=name,
+        NAMESPACE=namespace,
+        CLIENT_IMAGE=jax_client_image,
+        CLIENT_EXECUTION_COMMAND=client_command,
+        TPU_TYPE=tpu_type,
+        NUM_SLICES=num_slices,
+        RM_INSTANCE_TYPE=rm_instance_type,
+        GCS_SCRATCH_LOCATION=gcs_scratch_location,
+        GKE_TOPOLOGY=gke_topology,
+        VMS_PER_SLICE=vms_per_slice,
+        SPOT_TOLERATION_HEAD=spot_toleration_head,
+        SPOT_NODE_SELECTOR_WORKER=spot_node_selector_worker,
+        SPOT_TOLERATION_WORKER=spot_toleration_worker,
+        PROXY_SIDECAR_ARG=proxy_sidecar_arg,
+        TPU_PREMAPPED_BUFFER_SIZE=tpu_premapped_buffer_size,
+        WORKER_INIT_CONTAINERS=worker_init_containers,
+    )
+    
+    # Clean up empty lines caused by optional block placeholders
+    # (specifically ensuring there are no lines with only whitespace or empty lines where placeholders were)
+    lines = []
+    for line in yaml_content.splitlines():
+        if line.strip() or line == "":
+            lines.append(line)
+    return "\n".join(lines) + "\n"

pwy/kubernetes.py

@@ -0,0 +1,18 @@
+import subprocess
+
+def apply_manifest(yaml_content: str) -> subprocess.CompletedProcess:
+    """Applies the YAML manifest using kubectl apply -f -."""
+    process = subprocess.run(
+        ["kubectl", "apply", "-f", "-"],
+        input=yaml_content.encode("utf-8"),
+        capture_output=True,
+    )
+    return process
+
+def delete_jobset(name: str, namespace: str) -> subprocess.CompletedProcess:
+    """Deletes the JobSet using kubectl delete jobset <name> --namespace=<namespace>."""
+    process = subprocess.run(
+        ["kubectl", "delete", "jobset", name, f"--namespace={namespace}"],
+        capture_output=True,
+    )
+    return process

pwy/templates.py

@@ -0,0 +1,197 @@
+YAML_TEMPLATE = """apiVersion: jobset.x-k8s.io/v1alpha2
+kind: JobSet
+metadata:
+  name: {NAME}
+  namespace: {NAMESPACE}
+spec:
+  failurePolicy:
+    maxRestarts: 0
+    restartStrategy: BlockingRecreate
+  replicatedJobs:
+    # -------------------------------------------------------------------------
+    # 1. Pathways Head (Client Pod)
+    # -------------------------------------------------------------------------
+    - name: pwhd
+      replicas: 1
+      template:
+        spec:
+          parallelism: 1
+          completions: 1
+          backoffLimit: 32
+          template:
+            metadata:
+              annotations:
+                cluster-autoscaler.kubernetes.io/safe-to-evict: "false"
+            spec:
+              terminationGracePeriodSeconds: 60
+              restartPolicy: Never
+              hostAliases:
+                - ip: 169.254.169.254
+                  hostnames:
+                    - metadata
+                    - metadata.google.internal
+              tolerations:
+                - key: google.com/tpu
+                  operator: Equal
+                  value: "present"
+                  effect: NoSchedule
+{SPOT_TOLERATION_HEAD}
+              containers:
+                - name: client
+                  image: {CLIENT_IMAGE}
+                  command:
+                    - bash
+                    - -c
+                    - |
+                      {CLIENT_EXECUTION_COMMAND}
+                  resources:
+                    requests:
+                      cpu: "1000m"
+                      memory: "16Gi"
+                    limits:
+                      cpu: "1000m"
+                      memory: "16Gi"
+                  env:
+                    - name: TPU_TYPE
+                      value: {TPU_TYPE}
+                    - name: NUM_TPU_SLICES
+                      valueFrom:
+                        fieldRef:
+                          fieldPath: metadata.labels['jobset.sigs.k8s.io/replicatedjob-replicas']
+                    - name: JAX_BACKEND_TARGET
+                      value: grpc://localhost:29000
+                    - name: XCLOUD_ENVIRONMENT
+                      value: GCP
+                    - name: JAX_PLATFORMS
+                      value: proxy
+                    - name: ENABLE_PATHWAYS_PERSISTENCE
+                      value: "1"
+                    - name: TPU_SKIP_MDS_QUERY
+                      value: "true"
+                    - name: PYTHONUNBUFFERED
+                      value: "1"
+                    - name: TEST_UNDECLARED_OUTPUTS_DIR
+                      value: "true"
+                    - name: IFRT_PROXY_LARGE_TRANSFER_THRESHOLD
+                      value: "1"
+                    - name: IFRT_PROXY_LARGE_TRANSFER_OPTIMIZATION_DIRECTORY
+                      value: /tmp/ifrt_proxy
+                  volumeMounts:
+                    - name: shared-memory
+                      mountPath: /tmp/ifrt_proxy
+                  imagePullPolicy: Always
+              initContainers:
+                - name: pathways-proxy
+                  image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/proxy_server:jax-0.10.0
+                  restartPolicy: Always
+                  ports:
+                    - containerPort: 29000
+                  env:
+                    - name: IFRT_PROXY_USE_INSECURE_GRPC_CREDENTIALS
+                      value: "true"
+                    - name: IFRT_PROXY_LARGE_TRANSFER_OPTIMIZATION_DIRECTORY
+                      value: /tmp/ifrt_proxy
+                  args:
+                    - --resource_manager_address=localhost:29001
+                    - --server_port=29000{PROXY_SIDECAR_ARG}
+                  volumeMounts:
+                    - name: shared-memory
+                      mountPath: /tmp/ifrt_proxy
+                - name: pathways-rm
+                  image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:jax-0.10.0
+                  restartPolicy: Always
+                  env:
+                    - name: TPU_SKIP_MDS_QUERY
+                      value: "true"
+                  args:
+                    - --server_port=29001
+                    - --node_type=resource_manager
+                    - --instance_count={NUM_SLICES}
+                    - --instance_type={RM_INSTANCE_TYPE}
+                    - --gcs_scratch_location={GCS_SCRATCH_LOCATION}
+                    - --enforce_kernel_ipv6_support=false
+              volumes:
+                - name: shared-memory
+                  emptyDir:
+                    medium: Memory
+              serviceAccountName: default
+              dnsPolicy: ClusterFirstWithHostNet
+
+    # -------------------------------------------------------------------------
+    # 2. Pathways Workers (TPU Pods)
+    # -------------------------------------------------------------------------
+    - name: pwwk
+      replicas: {NUM_SLICES}
+      template:
+        metadata:
+          annotations:
+            alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool
+        spec:
+          parallelism: {VMS_PER_SLICE}
+          completions: {VMS_PER_SLICE}
+          backoffLimit: 32
+          template:
+            spec:
+              terminationGracePeriodSeconds: 60
+              hostAliases:
+                - ip: 169.254.169.254
+                  hostnames:
+                    - metadata
+                    - metadata.google.internal
+              nodeSelector:
+                cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
+                cloud.google.com/gke-tpu-topology: {GKE_TOPOLOGY}
+{SPOT_NODE_SELECTOR_WORKER}
+              tolerations:
+                - key: google.com/tpu
+                  operator: Equal
+                  value: "present"
+                  effect: NoSchedule
+{SPOT_TOLERATION_WORKER}
+{WORKER_INIT_CONTAINERS}
+              containers:
+                - name: worker
+                  image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:jax-0.10.0
+                  imagePullPolicy: Always
+                  ports:
+                    - containerPort: 8471
+                    - containerPort: 8080
+                    - containerPort: 8431
+                    - containerPort: 9000
+                    - containerPort: 29001
+                  securityContext:
+                    privileged: true
+                  resources:
+                    limits:
+                      google.com/tpu: 4
+                  env:
+                    - name: TPU_TYPE
+                      value: {TPU_TYPE}
+                    - name: NUM_TPU_SLICES
+                      valueFrom:
+                        fieldRef:
+                          fieldPath: metadata.labels['jobset.sigs.k8s.io/replicatedjob-replicas']
+                    - name: MEGASCALE_COORDINATOR_ADDRESS
+                      value: {NAME}-pwhd-0-0.{NAME}
+                    - name: MEGASCALE_NUM_SLICES
+                      valueFrom:
+                        fieldRef:
+                          fieldPath: metadata.labels['jobset.sigs.k8s.io/replicatedjob-replicas']
+                    - name: MEGASCALE_SLICE_ID
+                      valueFrom:
+                        fieldRef:
+                          fieldPath: metadata.labels['jobset.sigs.k8s.io/job-index']
+                  args:
+                    - --server_port=29001
+                    - --resource_manager_address={NAME}-pwhd-0-0.{NAME}:29001
+                    - --gcs_scratch_location={GCS_SCRATCH_LOCATION}
+                    - --tpu_pinned_host_allocation_recycle=true
+                    - --tpu_premapped_buffer_size={TPU_PREMAPPED_BUFFER_SIZE}
+                    - --enforce_kernel_ipv6_support=false
+              serviceAccountName: default
+              dnsPolicy: ClusterFirstWithHostNet
+  successPolicy:
+    operator: All
+    targetReplicatedJobs:
+      - pwhd
+"""