Remove head job sidecar mode from PathwaysJobSet.

PiperOrigin-RevId: 949622810

Author: lukebaumann

README.md

@@ -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.

pathwaysutils/experimental/gke/README.md

@@ -93,8 +93,6 @@ def __init__(
       tpu_type: str,
       topology: str,
       num_slices: int,
-      user_pod_template: Mapping[str, Any] | None = None,
-      main_container_name: str = "main",
       max_restarts: int = 0,
       max_slice_restarts: int = 0,
       termination_grace_period_seconds: int | None = None,
@@ -114,8 +112,6 @@ def __init__(
       tpu_type: TPU type (e.g., "v5e").
       topology: TPU topology (e.g., "2x2").
       num_slices: Number of slices.
-      user_pod_template: Optional user pod template for the head job.
-      main_container_name: Name of the main container in user_pod_template.
       max_restarts: Maximum number of restarts for the JobSet.
       max_slice_restarts: Maximum number of slice restarts.
       termination_grace_period_seconds: Optional termination grace period.
@@ -124,12 +120,8 @@ def __init__(
       elastic_slices: Number of elastic slices.
       labels: Optional labels for the JobSet.
       annotations: Optional annotations for the JobSet.
+      shared_pathways_service: Whether to run only RM for Shared Pathways Service.
     """
-    if shared_pathways_service and user_pod_template:
-      raise ValueError(
-          "Cannot enable shared_pathways_service when user_pod_template is"
-          " provided."
-      )
     self._shared_pathways_service = shared_pathways_service
 
     self._name = name
@@ -166,8 +158,6 @@ def __init__(
         num_slices=num_slices,
         instance_type=instance_type,
         image_tag=image_tag,
-        user_pod_template=user_pod_template,
-        main_container_name=main_container_name,
         elastic_slices=elastic_slices,
         shared_pathways_service=shared_pathways_service,
     )
@@ -185,7 +175,7 @@ def __init__(
     )
 
     self._success_policy = None
-    if user_pod_template or shared_pathways_service:
+    if shared_pathways_service:
       self._success_policy = {
           "operator": "All",
           "targetReplicatedJobs": [PATHWAYS_HEAD_JOB_NAME],
@@ -205,8 +195,6 @@ def _build_head_job_template(
       num_slices: int,
       instance_type: str,
       image_tag: str,
-      user_pod_template: Mapping[str, Any] | None,
-      main_container_name: str,
       elastic_slices: int,
       shared_pathways_service: bool,
   ) -> client.V1JobTemplateSpec:
@@ -217,9 +205,8 @@ def _build_head_job_template(
       num_slices: Number of slices.
       instance_type: TPU instance type (e.g., "tpuv5:2x2").
       image_tag: Version tag for Pathways images.
-      user_pod_template: Optional user pod template for the head job.
-      main_container_name: Name of the main container in user_pod_template.
       elastic_slices: Number of elastic slices.
+      shared_pathways_service: Whether to run only RM for Shared Pathways Service.
 
     Returns:
       The head job template.
@@ -318,74 +305,20 @@ def _build_head_job_template(
         ),
     )
 
-    api_client = client.ApiClient()
-
-    if user_pod_template:
-      user_template_obj = _deserialize_dict(
-          api_client, user_pod_template, client.V1PodTemplateSpec
-      )
-      head_pod_spec = user_template_obj.spec
-      head_pod_spec.host_network = True
-      head_pod_spec.dns_policy = "ClusterFirstWithHostNet"
-
-      rm_container.restart_policy = "Always"  # pyrefly: ignore[missing-attribute]
-      proxy_container.restart_policy = "Always"  # pyrefly: ignore[missing-attribute]
-
-      init_containers = head_pod_spec.init_containers or []
-      init_containers.extend([rm_container, proxy_container])
-      head_pod_spec.init_containers = init_containers
-
-      # Inject JAX env vars into main container.
-      jax_env = [
-          client.V1EnvVar(
-              name="PATHWAYS_HEAD",
-              value_from=client.V1EnvVarSource(
-                  field_ref=client.V1ObjectFieldSelector(
-                      field_path=(
-                          "metadata.labels['jobset.sigs.k8s.io/coordinator']"
-                      )
-                  )
-              ),
-          ),
-          client.V1EnvVar(name="JAX_PLATFORMS", value="proxy"),
-          client.V1EnvVar(name="XCLOUD_ENVIRONMENT", value="GCP"),
-          client.V1EnvVar(
-              name="JAX_BACKEND_TARGET",
-              value=f"grpc://$(PATHWAYS_HEAD):{PATHWAYS_PROXY_PORT}",
-          ),
-      ]
-      containers = head_pod_spec.containers or []
-      for c in containers:
-        if c.name == main_container_name:
-          env = c.env or []
-          env.extend(jax_env)
-          c.env = env
-          break
-      head_pod_spec.containers = containers
-
-      annotations = user_pod_template.get("metadata", {}).get("annotations", {})
-      labels = user_pod_template.get("metadata", {}).get("labels", {})
-    else:
-      # Headless mode.
-      containers = [rm_container]
-      if not shared_pathways_service:
-        containers.append(proxy_container)
-      head_pod_spec = client.V1PodSpec(
-          host_network=True,
-          dns_policy="ClusterFirstWithHostNet",
-          containers=containers,
-      )
-      annotations = {}
-      labels = {}
+    containers = [rm_container]
+    if not shared_pathways_service:
+      containers.append(proxy_container)
 
-    if not head_pod_spec.restart_policy:
-      head_pod_spec.restart_policy = "Never"
+    head_pod_spec = client.V1PodSpec(
+        host_network=True,
+        dns_policy="ClusterFirstWithHostNet",
+        containers=containers,
+        restart_policy="Never",
+    )
 
-    # Default annotations
     job_annotations = {
         "alpha.jobset.sigs.k8s.io/exclusive-topology": "kubernetes.io/hostname"
     }
-    job_annotations.update(annotations)
 
     head_job_template = client.V1JobTemplateSpec(
         metadata=client.V1ObjectMeta(annotations=job_annotations),
@@ -396,7 +329,7 @@ def _build_head_job_template(
             parallelism=1,
             template=client.V1PodTemplateSpec(
                 metadata=client.V1ObjectMeta(
-                    annotations=job_annotations, labels=labels
+                    annotations=job_annotations, labels={}
                 ),
                 spec=head_pod_spec,
             ),