pathwaysutils 0.1.11__tar.gz → 0.1.12__tar.gz

{pathwaysutils-0.1.11 → pathwaysutils-0.1.12}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pathwaysutils
-Version: 0.1.11
+Version: 0.1.12
 Summary: Package of Pathways-on-Cloud utilities.
 Keywords: 
 Author-email: Pathways-on-Cloud Utilities Developers <pathwaysutils-dev@google.com>
@@ -11,16 +11,31 @@ License-File: LICENSE
 Requires-Dist: absl-py
 Requires-Dist: fastapi
 Requires-Dist: jax>=0.8.3
-Requires-Dist: orbax-checkpoint
+Requires-Dist: orbax-checkpoint>=0.12.3
 Requires-Dist: uvicorn
 Requires-Dist: requests
 Requires-Dist: packaging
 Requires-Dist: pylint>=2.6.0 ; extra == "dev"
 Requires-Dist: pyink ; extra == "dev"
+Requires-Dist: kubernetes ; extra == "gke"
+Requires-Dist: pyyaml ; extra == "gke"
+Requires-Dist: portpicker ; extra == "sps"
+Requires-Dist: google-cloud-monitoring ; extra == "sps"
+Requires-Dist: portpicker ; extra == "sps-without-monitoring"
+Requires-Dist: google-cloud-monitoring ; extra == "test"
+Requires-Dist: httpx ; extra == "test"
+Requires-Dist: kubernetes ; extra == "test"
+Requires-Dist: portpicker ; extra == "test"
+Requires-Dist: pyyaml ; extra == "test"
+Requires-Dist: pytest ; extra == "test"
 Project-URL: changelog, https://github.com/AI-Hypercomputer/pathways-utils/blob/main/CHANGELOG.md
 Project-URL: homepage, https://github.com/AI-Hypercomputer/pathways-utils
 Project-URL: repository, https://github.com/AI-Hypercomputer/pathways-utils
 Provides-Extra: dev
+Provides-Extra: gke
+Provides-Extra: sps
+Provides-Extra: sps-without-monitoring
+Provides-Extra: test
 
 # Package of Pathways-on-Cloud utilities
 

{pathwaysutils-0.1.11 → pathwaysutils-0.1.12}/pathwaysutils/__init__.py

@@ -21,4 +21,4 @@ is_pathways_backend_used: Callable[[], bool] = _initialize.is_pathways_backend_u
 del _initialize
 
 # When changing this, also update the CHANGELOG.md.
-__version__ = "v0.1.11"
+__version__ = "v0.1.12"

{pathwaysutils-0.1.11 → pathwaysutils-0.1.12}/pathwaysutils/_initialize.py

@@ -18,10 +18,11 @@ import logging
 import os
 
 import jax
+from orbax.checkpoint import type_handlers
 from orbax.checkpoint._src.metadata import array_metadata_store as array_metadata_store_lib
+from orbax.checkpoint._src.serialization import cloud_pathways_array_handler
 from pathwaysutils import profiling
 from pathwaysutils import proxy_backend
-from pathwaysutils.persistence import orbax_handler
 
 
 _logger = logging.getLogger(__name__)
@@ -91,11 +92,14 @@ def initialize() -> None:
     _logger.debug("Detected Pathways-on-Cloud backend. Applying changes.")
     proxy_backend.register_backend_factory()
     profiling.monkey_patch_jax()
-    # TODO: b/365549911 - Remove when OCDBT-compatible
     if _is_persistence_enabled():
-      orbax_handler.register_pathways_handlers(
-          timeout=datetime.timedelta(hours=1),
-          array_metadata_store=array_metadata_store_lib.Store(),
+      type_handlers.register_type_handler(
+          jax.Array,
+          cloud_pathways_array_handler.CloudPathwaysArrayHandler(
+              timeout=datetime.timedelta(hours=1),
+              array_metadata_store=array_metadata_store_lib.Store(),
+          ),
+          override=True,
       )
 
     # Turn off JAX compilation cache because Pathways handles its own

{pathwaysutils-0.1.11 → pathwaysutils-0.1.12}/pathwaysutils/collect_profile.py

@@ -67,7 +67,10 @@ def main() -> None:
   args = parser.parse_args()
 
   if profiling.collect_profile(
-      args.port, args.duration_ms, args.host, args.log_dir
+      args.port,
+      args.duration_ms,
+      args.host,
+      args.log_dir,
   ):
     _logger.info("Dumped profiling information in: %s", args.log_dir)
   else:

{pathwaysutils-0.1.11 → pathwaysutils-0.1.12}/pathwaysutils/elastic/elastic.py

@@ -171,7 +171,7 @@ def get_active_slice_indices(
 
 
 def wait_for_slices(
-    slice_count: int,
+    slice_count: int | None = None,
     poll_interval: float | int = 10,
     timeout: float | int | None = None,
     slice_to_devices: Mapping[int, Sequence[jax.Device]] | None = None,
@@ -180,7 +180,8 @@ def wait_for_slices(
   """Waits until after at least `slice_count` slices become active.
 
   Args:
-    slice_count: The number of slices to wait for.
+    slice_count: The number of slices to wait for. If None, defaults to the
+      total number of slices.
     poll_interval: The minimum number of seconds to wait between availability
       checks. If the check takes longer than this, the next check will start
       immediately after the current check completes. Defaults to 10 seconds.
@@ -201,6 +202,10 @@ def wait_for_slices(
     _logger.debug("slice_to_devices is None. Getting from jax.devices().")
     slice_to_devices = get_slice_to_devices(jax.devices())
 
+  if slice_count is None:
+    _logger.debug("slice_count is None. Using len(slice_to_devices).")
+    slice_count = len(slice_to_devices)
+
   _logger.info(
       "Waiting for %s slices. Poll interval: %s, Timeout: %s",
       slice_count,

{pathwaysutils-0.1.11 → pathwaysutils-0.1.12}/pathwaysutils/elastic/manager.py

@@ -116,11 +116,12 @@ class Manager:
 
     self.all_slice_indices = frozenset(self.slice_to_devices.keys())
 
+    self._active_slice_indices = frozenset()
+    self.available_inactive_slices = frozenset()
+
     self.active_slice_indices = elastic.get_active_slice_indices(
         slice_to_devices=self.slice_to_devices
     )
-    self.inactive_slice_indices = self.all_slice_indices - self.active_slice_indices
-    self.available_inactive_slices = frozenset()
 
     self._stop_event = None
     self._monitor_thread = None
@@ -180,22 +181,26 @@ class Manager:
     except StopIteration as error:
       raise ValueError("No active slices") from error
 
+  @property
+  def active_slice_indices(self) -> frozenset[int]:
+    """The indices of active slices."""
+    return self._active_slice_indices
+
+  @active_slice_indices.setter
+  def active_slice_indices(self, value: Set[int]) -> None:
+    self._active_slice_indices = frozenset(value)
+    self.inactive_slice_indices = (
+        self.all_slice_indices - self._active_slice_indices
+    )
+    self.available_inactive_slices = frozenset(
+        self.available_inactive_slices - self._active_slice_indices
+    )
+
   @property
   def active_slice_count(self) -> int:
     """The number of active slices."""
     return len(self.active_slice_indices)
 
-  @property
-  def new_slice_event(self) -> threading.Event:
-    """Deprecated compatibility property for un-updated MaxText code.
-
-    TODO: b/527183831 - Remove this property once MaxText CL 2 is submitted.
-    """
-    event = threading.Event()
-    if self.available_inactive_slices:
-      event.set()
-    return event
-
   def scale_by_active_slices(self, x: int | float) -> int | float:
     """Scale x by the number of active slices."""
     if isinstance(x, int):
@@ -372,9 +377,6 @@ class Manager:
               poll_interval=poll_interval,
               timeout=timeout,
           )
-          self.inactive_slice_indices = (
-              self.all_slice_indices - self.active_slice_indices
-          )
           # Reset available_inactive_slices at attempt start since
           # active_slice_indices has just been updated by wait_for_slices.
           self.available_inactive_slices = frozenset()

{pathwaysutils-0.1.11 → pathwaysutils-0.1.12}/pathwaysutils/experimental/gke/jobset.py

@@ -11,17 +11,27 @@
 # limitations under the License.
 """Pathways JobSet generator and builder (with Worker Job Config)."""
 
+from __future__ import annotations
+
 import hashlib
 import json
 import logging
 import math
 import time
-from typing import Any, Mapping, Sequence
-from kubernetes import client
-from kubernetes import config as k8s_config
+from typing import TYPE_CHECKING, Any, Mapping, Sequence
 import yaml
 
-# GKE sidecar containers restartPolicy compatibility placeholder.
+try:
+  import kubernetes
+except ImportError as e:
+  raise ImportError(
+      "GKE utilities require `kubernetes`. "
+      "Please install pathwaysutils with GKE support:\n\n"
+      "    pip install 'pathwaysutils[gke]'\n"
+  ) from e
+
+from kubernetes import client
+from kubernetes import config as k8s_config
 
 _logger = logging.getLogger(__name__)
 
@@ -71,8 +81,18 @@ MACHINE_TYPE_TO_GKE_ACCELERATOR_TYPE_MAP = {
 }
 
 
+def _format_image(image: str, default_tag: str) -> str:
+  """Appends default_tag if image does not already specify a tag or digest."""
+  if "@" in image:
+    return image
+  last_slash = image.rfind("/")
+  if ":" in image[last_slash + 1:]:
+    return image
+  return f"{image}:{default_tag}"
+
+
 def _deserialize_dict(
-    api_client: client.ApiClient, data_dict: Mapping[str, Any], klass: Any
+    api_client: Any, data_dict: Mapping[str, Any], klass: Any
 ) -> Any:
   class FakeResponse:
 
@@ -93,10 +113,8 @@ class PathwaysJobSet:
       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,
+      max_slice_restarts: int = 1_000_000,
       termination_grace_period_seconds: int | None = None,
       pathways_version: str = "latest",
       jobset_api_version: str = "v1alpha2",
@@ -104,6 +122,8 @@ class PathwaysJobSet:
       labels: Mapping[str, str] | None = None,
       annotations: Mapping[str, str] | None = None,
       shared_pathways_service: bool = False,
+      pathways_rm_and_worker_image: str = DEFAULT_PATHWAYS_RM_AND_WORKER_IMAGE,
+      pathways_proxy_image: str = DEFAULT_PATHWAYS_PROXY_IMAGE,
   ):
     """Initializes the instance.
 
@@ -114,23 +134,22 @@ class PathwaysJobSet:
       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.
+      max_slice_restarts: Maximum number of slice restarts (defaults to 1_000_000 in headless and SPS mode).
       termination_grace_period_seconds: Optional termination grace period.
       pathways_version: Version tag for Pathways images.
       jobset_api_version: API version of JobSet.
       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.
+      pathways_rm_and_worker_image: Base Docker image for Resource Manager and
+        Worker containers.
+      pathways_proxy_image: Base Docker image for Proxy container.
     """
-    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._pathways_rm_and_worker_image = pathways_rm_and_worker_image
+    self._pathways_proxy_image = pathways_proxy_image
 
     self._name = name
     self._namespace = namespace
@@ -140,6 +159,11 @@ class PathwaysJobSet:
     self._labels = dict(labels) if labels else {}
     self._annotations = dict(annotations) if annotations else {}
 
+    if elastic_slices > num_slices:
+      raise ValueError(
+          f"elastic_slices ({elastic_slices}) cannot be greater than "
+          f"num_slices ({num_slices})"
+      )
     tpu_version = MACHINE_TYPE_TO_TPU_VERSION_MAP.get(tpu_type.lower())
     if not tpu_version:
       raise ValueError(f"Unsupported TPU type: {tpu_type}")
@@ -166,10 +190,10 @@ class PathwaysJobSet:
         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,
+        pathways_rm_and_worker_image=pathways_rm_and_worker_image,
+        pathways_proxy_image=pathways_proxy_image,
     )
 
     # Build worker template.
@@ -182,10 +206,11 @@ class PathwaysJobSet:
         image_tag=image_tag,
         max_slice_restarts=max_slice_restarts,
         termination_grace_period_seconds=termination_grace_period_seconds,
+        pathways_rm_and_worker_image=pathways_rm_and_worker_image,
     )
 
     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],
@@ -199,16 +224,24 @@ class PathwaysJobSet:
   def worker_job_template(self) -> client.V1JobTemplateSpec:
     return self._worker_job_template
 
+  @property
+  def pathways_rm_and_worker_image(self) -> str:
+    return self._pathways_rm_and_worker_image
+
+  @property
+  def pathways_proxy_image(self) -> str:
+    return self._pathways_proxy_image
+
   def _build_head_job_template(
       self,
       pathways_dir: str,
       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,
+      pathways_rm_and_worker_image: str,
+      pathways_proxy_image: str,
   ) -> client.V1JobTemplateSpec:
     """Builds the head job template for the JobSet.
 
@@ -217,15 +250,16 @@ class PathwaysJobSet:
       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.
+      pathways_rm_and_worker_image: Base Docker image for Resource Manager.
+      pathways_proxy_image: Base Docker image for Proxy container.
 
     Returns:
       The head job template.
     """
-    rm_image = f"{DEFAULT_PATHWAYS_RM_AND_WORKER_IMAGE}:{image_tag}"
-    proxy_image = f"{DEFAULT_PATHWAYS_PROXY_IMAGE}:{image_tag}"
+    rm_image = _format_image(pathways_rm_and_worker_image, image_tag)
+    proxy_image = _format_image(pathways_proxy_image, image_tag)
 
     rm_args = [
         f"--server_port={PATHWAYS_RM_PORT}",
@@ -318,74 +352,18 @@ class PathwaysJobSet:
         ),
     )
 
-    api_client = client.ApiClient()
+    containers = [rm_container]
+    if not shared_pathways_service:
+      containers.append(proxy_container)
 
-    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 = {}
-
-    if not head_pod_spec.restart_policy:
-      head_pod_spec.restart_policy = "Never"
+    head_pod_spec = client.V1PodSpec(
+        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 +374,7 @@ class PathwaysJobSet:
             parallelism=1,
             template=client.V1PodTemplateSpec(
                 metadata=client.V1ObjectMeta(
-                    annotations=job_annotations, labels=labels
+                    annotations=job_annotations, labels={}
                 ),
                 spec=head_pod_spec,
             ),
@@ -414,8 +392,10 @@ class PathwaysJobSet:
       image_tag: str,
       max_slice_restarts: int,
       termination_grace_period_seconds: int | None,
+      pathways_rm_and_worker_image: str,
   ) -> client.V1JobTemplateSpec:
-    worker_image = f"{DEFAULT_PATHWAYS_RM_AND_WORKER_IMAGE}:{image_tag}"
+    """Builds the worker job template for the JobSet."""
+    worker_image = _format_image(pathways_rm_and_worker_image, image_tag)
 
     args = [
         f"--resource_manager_address=$(PATHWAYS_HEAD):{PATHWAYS_RM_PORT}",
@@ -510,9 +490,7 @@ class PathwaysJobSet:
         "cloud.google.com/gke-tpu-topology": topology,
     }
 
-    backoff_limit = num_vms * 4
-    if max_slice_restarts > 0:
-      backoff_limit = num_vms * max_slice_restarts
+    backoff_limit = num_vms * max_slice_restarts
 
     worker_pod_spec = client.V1PodSpec(
         containers=[worker_container],
@@ -525,8 +503,6 @@ class PathwaysJobSet:
                 ),
             )
         ],
-        host_network=True,
-        dns_policy="ClusterFirstWithHostNet",
         restart_policy="OnFailure",
     )
     if termination_grace_period_seconds is not None:
@@ -637,14 +613,15 @@ class PathwaysJobSet:
         ports=[client.V1ContainerPort(container_port=50051)],
         volume_mounts=[
             client.V1VolumeMount(name="shared-tmp", mount_path="/tmp"),
-            client.V1VolumeMount(name=shm_volume_name, mount_path=shm_mount_path),
+            client.V1VolumeMount(
+                name=shm_volume_name, mount_path=shm_mount_path
+            ),
         ],
     )
-    colocated_container.restart_policy = "Always"  # pyrefly: ignore[missing-attribute]
 
-    init_containers = pod_spec.init_containers or []
-    init_containers.append(colocated_container)
-    pod_spec.init_containers = init_containers
+    containers = pod_spec.containers or []
+    containers.append(colocated_container)
+    pod_spec.containers = containers
 
     # Add volume mount to pathways-worker.
     for container in pod_spec.containers:
@@ -660,7 +637,7 @@ class PathwaysJobSet:
         env = container.env or []
         env.append(
             client.V1EnvVar(
-                name="cloud_pathways_sidecar_shm_directory",
+                name="CLOUD_PATHWAYS_SIDECAR_SHM_DIRECTORY",
                 value=shm_mount_path,
             )
         )
@@ -694,7 +671,9 @@ class PathwaysJobSet:
 
     for job_template in (self._head_job_template, self._worker_job_template):
       pod_spec = job_template.spec.template.spec
-      all_containers = (pod_spec.containers or []) + (pod_spec.init_containers or [])
+      all_containers = (pod_spec.containers or []) + (
+          pod_spec.init_containers or []
+      )
 
       matching = self._filter_matching_containers(containers, all_containers)
       if not matching:
@@ -824,6 +803,25 @@ class PathwaysJobSet:
     instance._head_job_template = head_job_template
     instance._worker_job_template = worker_job_template
 
+    rm_image = DEFAULT_PATHWAYS_RM_AND_WORKER_IMAGE
+    proxy_image = DEFAULT_PATHWAYS_PROXY_IMAGE
+    if (
+        head_job_template.spec
+        and head_job_template.spec.template
+        and head_job_template.spec.template.spec
+    ):
+      pod_spec = head_job_template.spec.template.spec
+      all_containers = (pod_spec.containers or []) + (
+          pod_spec.init_containers or []
+      )
+      for c in all_containers:
+        if c.name == "pathways-rm" and c.image:
+          rm_image = c.image
+        elif c.name == "pathways-proxy" and c.image:
+          proxy_image = c.image
+    instance._pathways_rm_and_worker_image = rm_image
+    instance._pathways_proxy_image = proxy_image
+
     instance._success_policy = config["spec"].get("successPolicy")
     return instance
 

{pathwaysutils-0.1.11 → pathwaysutils-0.1.12}/pathwaysutils/experimental/shared_pathways_service/README.md

@@ -16,13 +16,19 @@ You have a GKE cluster with at least 1 TPU slice (v5e, v5p or v6e).
 
 ### 2. Deploy the Pathways head pod
 
-Start the Shared Pathways Service by using [pw-service-example.yaml](yamls/pw-service-example.yaml).
-Make sure to modify the following values to deploy the Pathways pods:
+Start the Shared Pathways Service by running `deploy_pathways_service.py`.
 
-- A unique Jobset name for the head pod
-- GCS bucket path
-- TPU type and topology
-- Number of slices
+```shell
+python3 -m pathwaysutils.experimental.shared_pathways_service.deploy_pathways_service \
+--jobset_name ${SERVICE_JOBSET_NAME} \
+--gcs_bucket ${GCS_BUCKET} \
+--server_image ${SERVER_IMAGE} \
+--tpu_type ${TPU_TYPE} \
+--topology ${TOPOLOGY} \
+--num_slices ${NUM_SLICES} [--sidecar_image ${SIDECAR_IMAGE}] [--dry_run]
+```
+
+Note, `server_image` is the Pathways server image, e.g, `us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:<tag>`.
 
 ### 3. Verify that the pods created in [Step#2](#2-deploy-the-pathways-head-pod) are running
 
@@ -55,13 +61,13 @@ pathways-cluster-worker-1-0-km2rf          1/1     Running   0          3m36s
 
 ```shell
 # e.g., pathways-cluster
-$ JOBSET_NAME=<your-jobset-name>  # same as you used in [pw-service-example.yaml](#pw-service-yaml)
+$ SERVICE_JOBSET_NAME=<your-jobset-name>  # same as you used in [pw-service-example.yaml](#pw-service-yaml)
 
 # e.g., pathways-cluster-pathways-head-0-0-zzmn2
-$ HEAD_POD_NAME=$(kubectl get pods --selector=jobset.sigs.k8s.io/jobset-name=${JOBSET_NAME} -o jsonpath='{.items[?(@.status.phase=="Running")].metadata.name}' | sed 's/ /\n/g' | grep head)
+$ HEAD_POD_NAME=$(kubectl get pods --selector=jobset.sigs.k8s.io/jobset-name=${SERVICE_JOBSET_NAME} -o jsonpath='{.items[?(@.status.phase=="Running")].metadata.name}' | sed 's/ /\n/g' | grep head)
 
 # e.g., pathways-cluster-worker-0-0-bdzq4
-$ WORKER0_POD_NAME=$(kubectl get pods --selector=jobset.sigs.k8s.io/jobset-name=${JOBSET_NAME} -o jsonpath='{.items[?(@.status.phase=="Running")].metadata.name}' | sed 's/ /\n/g' | grep 'worker-0-0-')
+$ WORKER0_POD_NAME=$(kubectl get pods --selector=jobset.sigs.k8s.io/jobset-name=${SERVICE_JOBSET_NAME} -o jsonpath='{.items[?(@.status.phase=="Running")].metadata.name}' | sed 's/ /\n/g' | grep 'worker-0-0-')
 ```
 
 #### Option 3: Check project logs
@@ -80,15 +86,33 @@ I1208 20:10:18.148825       ...] argv[2]: '--resource_manager_address=pathways-c
 
 ## Instructions
 
-### 1. Clone `pathwaysutils`.
+### 1. Install `pathwaysutils` and SPS dependencies.
 
 ```shell
-git clone https://github.com/AI-Hypercomputer/pathways-utils.git
+pip install --upgrade "pathwaysutils[sps] @ git+https://github.com/AI-Hypercomputer/pathways-utils.git"
+
+# Note: google-cloud-monitoring is optional. If you don't wish to collect SPS usage metrics, install with pathways-utils[sps_without_monitoring].
 ```
 
-### 2. Use the `isc_pathways` Context Manager
+### 2. Either use `run_workload.py` or move your workload under the `isc_pathways` Context Manager
+
+Make sure to use the proxy image compatible with the server image used while deploying SPS.
+
+```shell
+python3 -m pathwaysutils.experimental.shared_pathways_service.run_workload \
+--cluster=$CLUSTER_NAME \
+--project=$PROJECT \
+--region=$REGION \
+--gcs_bucket=$GCS_BUCKET \
+--pathways_service="$SERVICE_JOBSET_NAME-pathways-head-0-0.$SERVICE_JOBSET_NAME:29001" \
+--tpu_type="$TPU_TYPE:$TOPOLOGY" \
+--tpu_count=$NUM_SLICES \
+--proxy_server_image=$PROXY_IMAGE \
+--collect_service_metrics \
+--command="python3 run_my_sample_workload.py" [--proxy_options=sidecar:true]
+```
 
-In your script,
+Alternatively, in your script,
 
 1.  Import `isc_pathways`
 2. Add `with isc_pathways.connect(...)` statement. The function takes the below values:
@@ -97,6 +121,10 @@ In your script,
     - Region
     - GCS bucket name
     - Pathways Service (See instructions to find the RM address [here](#4-find-the-pathways-service-address))
+    - Expected TPU instances
+    - Proxy server image
+    - Proxy options
+    - Whether to collect SPS usage metrics
 <a name="ml-code"></a>
 3. Write your ML code under this context manager (the `with` block) to run your JAX code on the underlying TPUs.
 
@@ -111,7 +139,8 @@ python3 pathwaysutils/experimental/shared_pathways_service/run_connect_example.p
 --gcs_bucket="gs://user-bucket" \
 --pathways_service="pathways-cluster-pathways-head-0-0.pathways-cluster:29001" \
 --tpu_type="tpuv6e:2x2" \
---tpu_count=1
+--tpu_count=1 \
+--proxy_server_image="proxy_image"
 ```
 
 The connect block will deploy a proxy pod dedicated to your client and connect