pathways-cli 0.1.0__tar.gz

pathways_cli-0.1.0/.gitignore

@@ -0,0 +1,16 @@
+# Python-generated files
+__pycache__/
+*.py[oc]
+build/
+dist/
+wheels/
+*.egg-info
+
+# Virtual environments
+.venv
+
+# Local env files
+.env
+
+# Test caches
+.pytest_cache/

pathways_cli-0.1.0/.python-version

@@ -0,0 +1 @@
+3.12

pathways_cli-0.1.0/GEMINI.md

@@ -0,0 +1,52 @@
+# Gemini Agent & Developer Guide: `pathways-cli`
+
+This document details the codebase design, architecture, key lessons, and integration verification guides for the `pwy` CLI tool, serving as a developer-facing companion to the user-facing `README.md`.
+
+---
+
+## 1. Codebase Architecture
+
+The project is structured around a standard PEP 621 package layout:
+
+```
+/Users/stoelinga/workspace/pathways-cli/
+├── pyproject.toml         # Package configurations, CLI scripts, and Pytest options
+├── .gitignore             # Excludes local environments, caches, and secrets
+├── README.md              # User documentation and example verification steps
+├── GEMINI.md              # Codebase design and developer/agent context (this file)
+├── src/
+│   └── pwy/
+│       ├── __init__.py    # Exposes cli entry points
+│       ├── cli.py         # click CLI definition: up, down commands
+│       ├── generator.py   # Topology math, spot VM toggles, colocated python configurations
+│       ├── templates.py   # Complete GKE JobSet multi-line YAML manifest template
+│       └── kubernetes.py  # Wrapper invoking kubectl subprocesses
+└── tests/
+    ├── __init__.py
+    ├── test_cli.py        # CLI option validations & mocks
+    ├── test_generator.py  # Mappings and string-formatting unit tests
+    └── test_e2e.py        # Real GKE cluster integration execution verifying JAX setup
+```
+
+---
+
+## 2. Testing Workflows
+
+Verify changes using one of the three testing scopes:
+
+### 1. Unit Tests (Mocked)
+Tests calculations and YAML generation without cluster access.
+```bash
+uv run pytest tests/test_generator.py tests/test_cli.py
+```
+
+### 2. End-to-End Integration Tests (Active Cluster)
+Runs actual deployments on a running TPU nodepool, installs JAX, executes verification scripts, and tears the setup down.
+1. Configure your GCS path in a local `.env` file:
+   ```env
+   PWY_E2E_GCS_SCRATCH_LOCATION=gs://my-staging-bucket/pathways
+   ```
+2. Run pytest targeting the `e2e` mark:
+   ```bash
+   uv run pytest tests/test_e2e.py -m e2e -s
+   ```

pathways_cli-0.1.0/PKG-INFO

@@ -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/README.md

@@ -0,0 +1,112 @@
+# `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/plan.md

@@ -0,0 +1,352 @@
+# Implementation Plan: Standalone Pathways Cluster CLI Tool (`pwy`)
+
+This document outlines the final design and detailed implementation specifications for a standalone Python CLI tool to generate, apply, and manage interactive Pathways GKE cluster manifests.
+
+---
+
+## 1. CLI Commands & Arguments
+
+The CLI binary/entry point will be named `pwy`.
+
+### Commands
+
+#### 1. `pwy up`
+Starts the cluster or dry-runs the configuration.
+
+```bash
+pwy up \
+  --tpu-type=v6e-4 \
+  --gcs-scratch-location=gs://my-bucket/staging \
+  [--num-slices=1] \
+  [--jax-client-image=python:3.12-slim] \
+  [--command="python my_script.py"] \
+  [--enable-spot] \
+  [--colocated-python] \
+  [--dry-run] \
+  [--name=pathways-interactive] \
+  [--namespace=default]
+```
+
+* **Behavior**:
+  * Calculates cluster configuration based on `--tpu-type` and `--num-slices`.
+  * Generates the JobSet YAML.
+  * If `--dry-run` is set: Prints the generated YAML to stdout and exits.
+  * Otherwise: Pipes the YAML directly to `kubectl apply -f -`.
+
+#### 2. `pwy down`
+Tears down the cluster.
+
+```bash
+pwy down [--name=pathways-interactive] [--namespace=default]
+```
+* **Behavior**: Runs `kubectl delete jobset <name> --namespace=<namespace>`.
+
+---
+
+## 2. Automated TPU Mappings & Configuration Math
+
+The tool automatically maps the user-provided `--tpu-type` to GKE resources, topologies, and arguments.
+
+### Mappings Database
+
+| TPU Type | GKE Topology | VMs Per Slice (`vms_per_slice`) | RM Instance Type (`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` |
+
+### Configuration Math
+
+For any given run with `tpu-type` and `num-slices`:
+
+* **Pathways Head (`pwhd` replicated job)**:
+  * `replicas` is always `1`.
+* **Pathways Worker (`pwwk` replicated job)**:
+  * `replicas` (number of slices) = `--num-slices`
+  * `spec.parallelism` (VMs per slice) = `vms_per_slice`
+  * `spec.completions` (VMs per slice) = `vms_per_slice`
+  * `resources.limits["google.com/tpu"]` = `4` (always 4 for v6e)
+  * `nodeSelector["cloud.google.com/gke-tpu-topology"]` = `GKE Topology`
+  * `nodeSelector["cloud.google.com/gke-tpu-accelerator"]` = `tpu-v6e-slice`
+* **Resource Manager container (`pathways-rm` sidecar)**:
+  * `--instance_count` = `--num-slices`
+  * `--instance_type` = `rm_instance_type`
+
+---
+
+## 3. Client Command & Image Execution Logic
+
+The client container `command` field is generated dynamically based on the `--jax-client-image`, `--command`, and `--colocated-python` flags:
+
+1. **JAX Client Image**: 
+   * Uses `--jax-client-image` (defaulting to `python:3.12-slim`).
+2. **Command Executed**:
+   * **If `--command` is NOT provided**:
+     ```bash
+     bash -c "sleep infinity"
+     ```
+   * **If `--command` IS provided** (e.g. `--command="python training.py"`):
+     The tool boots up the environment, initializes pathways, and then executes the custom command directly:
+     ```bash
+     bash -c "python training.py"
+     ```
+
+---
+
+## 4. Full Manifest Reference Template
+
+Below is the complete reference JobSet YAML template populated with default/baseline variables for a `v6e-4` single-slice run. The CLI generator code should output a structure identical to this:
+
+```yaml
+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 only added if --enable-spot is True
+                {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.9.2
+                  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
+                  volumeMounts:
+                    - name: shared-memory
+                      mountPath: /tmp/ifrt_proxy
+                - name: pathways-rm
+                  image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:jax-0.9.2
+                  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}
+              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:
+                # Spot toleration only added if --enable-spot is True
+                {SPOT_TOLERATION_WORKER}
+              containers:
+                - name: worker
+                  image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:jax-0.9.2
+                  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=274877906944
+              serviceAccountName: default
+              dnsPolicy: ClusterFirstWithHostNet
+  successPolicy:
+    operator: All
+    targetReplicatedJobs:
+      - pwhd
+```
+
+---
+
+## 5. Repository Layout & Target Files
+
+A new standalone directory structure will be created under the workspace (mocking a new repository context):
+
+```
+/Users/stoelinga/workspace/pathways-cli/
+├── pyproject.toml
+├── README.md
+├── pwy/
+│   ├── __init__.py
+│   ├── cli.py             # Entry point (Click CLI commands: up, down)
+│   ├── generator.py       # Topology mapping and dictionary interpolation
+│   ├── templates.py       # Text template holding the YAML manifest structure
+│   └── kubernetes.py      # Subprocess module executing "kubectl apply -f" or "kubectl delete"
+└── tests/
+    ├── __init__.py
+    ├── test_generator.py
+    └── test_cli.py
+```
+
+### File Implementation Details
+
+#### `pwy/generator.py`
+Contains the lookup dictionaries and mapping functions:
+```python
+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 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,
+) -> str:
+    # 1. Look up TPU type mappings
+    # 2. Format client container commands
+    # 3. Handle spot nodeSelector and tolerations formatting
+    # 4. Interpolate templates.YAML_TEMPLATE with final string variables
+    ...
+```
+
+#### `pwy/cli.py`
+Handles options parsing and commands:
+* Imports `generate_yaml`.
+* Runs `kubectl apply` or `kubectl delete` using Python's `subprocess.run(..., input=yaml_content.encode())`.
+

pathways_cli-0.1.0/pyproject.toml

@@ -0,0 +1,33 @@
+[project]
+name = "pathways-cli"
+version = "0.1.0"
+description = "Pathways CLI to easily bring up pathways clusters."
+readme = "README.md"
+authors = [
+    { name = "Sam Stoelinga", email = "sammiestoel@gmail.com" }
+]
+requires-python = ">=3.12"
+dependencies = [
+    "click>=8.4.1",
+    "python-dotenv>=1.2.2",
+]
+
+[project.scripts]
+pwy = "pwy:main"
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[tool.hatch.build.targets.wheel]
+packages = ["src/pwy"]
+
+[dependency-groups]
+dev = [
+    "pytest>=9.0.3",
+]
+
+[tool.pytest.ini_options]
+markers = [
+    "e2e: end-to-end integration tests",
+]

pathways_cli-0.1.0/src/pwy/__init__.py

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