pathwaysutils 0.1.12__tar.gz → 0.1.14__tar.gz

pathwaysutils-0.1.14/PKG-INFO

@@ -0,0 +1,90 @@
+Metadata-Version: 2.4
+Name: pathwaysutils
+Version: 0.1.14
+Summary: Package of Pathways-on-Cloud utilities.
+Keywords: 
+Author-email: Pathways-on-Cloud Utilities Developers <pathwaysutils-dev@google.com>
+Requires-Python: >=3.11
+Description-Content-Type: text/markdown
+Classifier: License :: OSI Approved :: Apache Software License
+License-File: LICENSE
+Requires-Dist: absl-py
+Requires-Dist: fastapi
+Requires-Dist: jax>=0.10.0
+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
+
+# APIs to get started with using Pathways on Cloud
+
+Pathways is a compute orchestration runtime for single-controller ML workloads
+across multiple Cloud TPU slices. A single Python controller orchestrates
+computation across hundreds or thousands of TPU chips, unlocking slice-level
+training resiliency, fast cross-slice data movement, and interactive
+multi-slice development.
+
+This package bundles the Python client integrations, resharding and elasticity
+helpers, and GKE deployment builders for running JAX workloads on Pathways on
+Cloud.
+
+## Installation
+
+```bash
+pip install pathwaysutils
+# Optional GKE JobSet generator and kubectl/Kubernetes client dependencies:
+pip install 'pathwaysutils[gke]'
+# Optional experimental Shared Pathways Service (SPS) dependencies:
+pip install 'pathwaysutils[sps]'
+```
+
+## Core Modules
+
+- `pathwaysutils.initialize()` registers the Pathways IFRT proxy backend,
+  configures the compilation cache and XProf trace hooks, and registers Orbax
+  checkpoint handlers.
+- `pathwaysutils.reshard.reshard` reshards arrays and PyTrees across device
+  meshes directly on Pathways workers via the IFRT `RemapArray` API without
+  routing tensors through controller memory.
+- `pathwaysutils.elastic.manager.Manager` monitors slice health in the
+  background and provides `@elastic_retry` and `ScaleUpSignalError` for
+  step-level fault recovery and dynamic slice scaling.
+- [`PathwaysJobSet`][gke-readme] (`pathwaysutils.experimental.gke.jobset`)
+  builds, customizes, exports, and deploys Kubernetes `JobSet` manifests for
+  Pathways workloads on GKE.
+- [Shared Pathways Service][sps-readme]
+  (`pathwaysutils.experimental.shared_pathways_service`) deploys a persistent
+  multi-tenant Pathways service on GKE for interactive JAX sessions.
+
+## Further Reading
+
+- [Get started with Pathways on Google Cloud][cloud-docs]
+- [Pathways Research Paper](https://arxiv.org/abs/2203.12533)
+
+[gke-readme]: https://github.com/AI-Hypercomputer/pathways-utils/tree/main/pathwaysutils/experimental/gke
+[sps-readme]: https://github.com/AI-Hypercomputer/pathways-utils/tree/main/pathwaysutils/experimental/shared_pathways_service
+[cloud-docs]: https://docs.cloud.google.com/ai-hypercomputer/docs/workloads/pathways-on-cloud/pathways-intro
+
+
+

pathwaysutils-0.1.14/README.md

@@ -0,0 +1,50 @@
+# APIs to get started with using Pathways on Cloud
+
+Pathways is a compute orchestration runtime for single-controller ML workloads
+across multiple Cloud TPU slices. A single Python controller orchestrates
+computation across hundreds or thousands of TPU chips, unlocking slice-level
+training resiliency, fast cross-slice data movement, and interactive
+multi-slice development.
+
+This package bundles the Python client integrations, resharding and elasticity
+helpers, and GKE deployment builders for running JAX workloads on Pathways on
+Cloud.
+
+## Installation
+
+```bash
+pip install pathwaysutils
+# Optional GKE JobSet generator and kubectl/Kubernetes client dependencies:
+pip install 'pathwaysutils[gke]'
+# Optional experimental Shared Pathways Service (SPS) dependencies:
+pip install 'pathwaysutils[sps]'
+```
+
+## Core Modules
+
+- `pathwaysutils.initialize()` registers the Pathways IFRT proxy backend,
+  configures the compilation cache and XProf trace hooks, and registers Orbax
+  checkpoint handlers.
+- `pathwaysutils.reshard.reshard` reshards arrays and PyTrees across device
+  meshes directly on Pathways workers via the IFRT `RemapArray` API without
+  routing tensors through controller memory.
+- `pathwaysutils.elastic.manager.Manager` monitors slice health in the
+  background and provides `@elastic_retry` and `ScaleUpSignalError` for
+  step-level fault recovery and dynamic slice scaling.
+- [`PathwaysJobSet`][gke-readme] (`pathwaysutils.experimental.gke.jobset`)
+  builds, customizes, exports, and deploys Kubernetes `JobSet` manifests for
+  Pathways workloads on GKE.
+- [Shared Pathways Service][sps-readme]
+  (`pathwaysutils.experimental.shared_pathways_service`) deploys a persistent
+  multi-tenant Pathways service on GKE for interactive JAX sessions.
+
+## Further Reading
+
+- [Get started with Pathways on Google Cloud][cloud-docs]
+- [Pathways Research Paper](https://arxiv.org/abs/2203.12533)
+
+[gke-readme]: https://github.com/AI-Hypercomputer/pathways-utils/tree/main/pathwaysutils/experimental/gke
+[sps-readme]: https://github.com/AI-Hypercomputer/pathways-utils/tree/main/pathwaysutils/experimental/shared_pathways_service
+[cloud-docs]: https://docs.cloud.google.com/ai-hypercomputer/docs/workloads/pathways-on-cloud/pathways-intro
+
+

{pathwaysutils-0.1.12 → pathwaysutils-0.1.14}/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.12"
+__version__ = "v0.1.14"

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

@@ -40,7 +40,9 @@ def is_pathways_backend_used() -> bool:
   is not necessarily which backend is ultimately selected or there may be more
   than one platform specified and another may have higher priority.
   """
-  return jax.config.jax_platforms and "proxy" in jax.config.jax_platforms
+  return bool(
+      jax.config.jax_platforms and "proxy" in jax.config.jax_platforms
+  )
 
 
 def _is_persistence_enabled() -> bool:

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

@@ -198,6 +198,10 @@ def wait_for_slices(
     TimeoutError: If the timeout is reached before the slices become
       active.
   """
+  if poll_interval < 0:
+    raise ValueError(f"poll_interval must be non-negative, got {poll_interval}")
+  if timeout is not None and timeout < 0:
+    raise ValueError(f"timeout must be non-negative, got {timeout}")
   if slice_to_devices is None:
     _logger.debug("slice_to_devices is None. Getting from jax.devices().")
     slice_to_devices = get_slice_to_devices(jax.devices())

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

@@ -52,7 +52,7 @@ RetryPolicy: TypeAlias = Callable[[int, Exception], bool]
 
 
 def _elastic_event_cleanup() -> None:
-  """Cleans up JAX profiles, caches, and live arrays."""
+  """Stops any ongoing JAX profiler trace."""
   try:
     _logger.info("Cleaning up any ongoing traces")
     jax.profiler.stop_trace()
@@ -62,10 +62,6 @@ def _elastic_event_cleanup() -> None:
     _logger.exception("Error cleaning up ongoing traces")
     raise
 
-  jax.clear_caches()
-  for array in jax.live_arrays():
-    array.delete()
-
 
 class ElasticRetryLimit:
   """A retry callback that limits the number of attempts."""
@@ -162,8 +158,9 @@ class Manager:
           pass
         else:
           raise
-      self._monitor_thread = None
-      self._stop_event = None
+      finally:
+        self._monitor_thread = None
+        self._stop_event = None
 
   @functools.cached_property
   def total_slice_count(self) -> int:
@@ -201,6 +198,14 @@ class Manager:
     """The number of active slices."""
     return len(self.active_slice_indices)
 
+  @property
+  def active_devices(self) -> list[jax.Device]:
+    """The devices belonging to the active slices, ordered by slice index."""
+    devices = []
+    for slice_index in sorted(self.active_slice_indices):
+      devices.extend(self.slice_to_devices[slice_index])
+    return devices
+
   def scale_by_active_slices(self, x: int | float) -> int | float:
     """Scale x by the number of active slices."""
     if isinstance(x, int):
@@ -218,20 +223,6 @@ class Manager:
     else:
       raise ValueError(f"Unsupported type: {type(x)=}")
 
-  def _cleanup_on_retry(self):
-    """Cleans up JAX caches and traces on retry."""
-    try:
-      _logger.debug("Cleaning up any ongoing traces")
-      jax.profiler.stop_trace()
-    except (RuntimeError, ValueError):
-      _logger.debug("No ongoing traces to clean up")
-    except Exception:  # pylint: disable=broad-exception-caught
-      _logger.exception("Error cleaning up ongoing traces")
-
-    jax.clear_caches()
-    for array in jax.live_arrays():
-      array.delete()
-
   def _check_inactive_slices(self) -> None:
     """Checks inactive slices and updates available_inactive_slices."""
     if not self.inactive_slice_indices:
@@ -299,8 +290,7 @@ class Manager:
 
     This decorator wraps a function to automatically retry execution in case of
     `jax.errors.JaxRuntimeError` caused by slice down events. It waits for
-    `minimum_slice_count` active slices before each attempt and cleans up JAX
-    caches on failure.
+    `minimum_slice_count` active slices before each attempt.
 
     If `minimum_slice_count` is not met, the function will wait until at least
     `minimum_slice_count` slices are active before execution. If
@@ -313,6 +303,12 @@ class Manager:
     boundaries) and raise a `ScaleUpSignalError` to gracefully interrupt the
     current execution and trigger a retry with the expanded hardware.
 
+    On an elastic event the arrays of the failed attempt are not deleted
+    explicitly; JAX frees them once the attempt has unwound and nothing
+    references them anymore. The wrapped function and the callbacks must not
+    keep references to them, or the memory stays allocated during the next
+    attempt.
+
     Often, the function will dispatch JAX operations and wait for them to
     complete while creating a log message. If using Python logging, it is
     recommended to set `logging.raiseExceptions=True` to ensure that the

{pathwaysutils-0.1.12 → pathwaysutils-0.1.14}/pathwaysutils/experimental/concatenate_by_mesh_axis.py

@@ -68,10 +68,21 @@ def concatenate_by_mesh_axis(
   if not flattened_arrays[0]:
     return array_trees[0]
 
-  # Extract the shared mesh from each PyTree (from an arbitrary array in each).
+  # Extract the shared mesh from each PyTree and verify all leaves in a PyTree
+  # share that mesh.
   meshes_to_concatenate = [
       _get_named_sharding(array).mesh for array in input_flat_arrays[0]
   ]
+  for tree_arrays, expected_mesh in zip(
+      flattened_arrays, meshes_to_concatenate
+  ):
+    for array in tree_arrays:
+      sharding = _get_named_sharding(array)
+      if sharding.mesh != expected_mesh:
+        raise ValueError(
+            "All arrays in each PyTree must have the same mesh. Got"
+            f" {sharding.mesh} vs. {expected_mesh}."
+        )
 
   # Validate that the meshes are compatible.
   reference_mesh = meshes_to_concatenate[0]
@@ -117,6 +128,19 @@ def concatenate_by_mesh_axis(
   ) -> jax.sharding.NamedSharding:
     reference_sharding = _get_named_sharding(arrays[0])
     reference_spec = reference_sharding.spec
+    for array in arrays[1:]:
+      sharding = _get_named_sharding(array)
+      if sharding.spec != reference_spec:
+        raise ValueError(
+            "Corresponding arrays across PyTrees must have the same"
+            f" PartitionSpec. Got {sharding.spec} vs. {reference_spec}."
+        )
+      if sharding.memory_kind != reference_sharding.memory_kind:
+        raise ValueError(
+            "Corresponding arrays across PyTrees must have the same"
+            f" memory_kind. Got {sharding.memory_kind} vs."
+            f" {reference_sharding.memory_kind}."
+        )
     return jax.sharding.NamedSharding(
         concatenated_mesh,
         reference_spec,

pathwaysutils-0.1.14/pathwaysutils/experimental/gke/README.md

@@ -0,0 +1,142 @@
+# PathwaysJobSet Builder
+
+`PathwaysJobSet` generates and deploys Kubernetes `JobSet` resources (`jobset.x-k8s.io`) for running distributed Pathways workloads on Google Kubernetes Engine (GKE) without requiring custom controllers.
+
+## Installation
+
+```bash
+pip install 'pathwaysutils[gke]'
+```
+
+## Quickstart
+
+```python
+from pathwaysutils.experimental.gke import jobset
+
+pw_jobset = jobset.PathwaysJobSet(
+    name="my-workload",
+    namespace="default",
+    pathways_dir="gs://my-bucket/pathways-scratch",
+    tpu_type="v5e",
+    topology="2x4",
+    num_slices=1,
+)
+
+# Export to YAML or deploy directly to the active cluster
+pw_jobset.export_yaml("jobset.yaml")
+pw_jobset.apply(recreate=True)
+```
+
+## Features
+
+### Elasticity & Multislice
+Scale across multiple TPU slices with dynamic resilience:
+
+```python
+pw_jobset = jobset.PathwaysJobSet(
+    name="elastic-v6e",
+    namespace="default",
+    pathways_dir="gs://my-bucket/scratch",
+    tpu_type="v6e",
+    topology="4x4",
+    num_slices=4,
+    elastic_slices=2,  # Configures Pathways Proxy for dynamic slice recovery
+)
+```
+
+### Heterogeneous Multi-Topology Slices
+Combine multiple TPU slice topologies as separate worker `replicatedJobs`
+within a single `JobSet`. Call `add_worker_topology` before worker mutators
+(`add_colocated_python`, `add_gcsfuse`, or `with_custom_libtpu`) so they apply
+across all worker templates:
+
+```python
+pw_jobset = (
+    jobset.PathwaysJobSet(
+        name="heterogeneous-v5e",
+        namespace="default",
+        pathways_dir="gs://my-bucket/scratch",
+        tpu_type="v5e",
+        topology="4x8",
+        num_slices=1,
+    )
+    .add_worker_topology(
+        tpu_type="ct5lp-hightpu-8t",
+        topology="2x4",
+        num_slices=1,
+    )
+)
+```
+
+### In-Pod User Workload
+Run training or inference directly in the head pod alongside the Resource Manager and Proxy sidecars:
+
+```python
+pw_jobset.add_user_workload(
+    image="us-docker.pkg.dev/my-project/images/my-trainer:latest",
+    command="python3 train.py --batch_size=32",
+    env={"LOG_LEVEL": "INFO"},
+)
+```
+
+### Colocated Python Sidecar
+Inject a sidecar container onto worker pods sharing memory with
+`pathways-worker` for zero-copy tensor transfers:
+
+```python
+pw_jobset.add_colocated_python(
+    image="us-docker.pkg.dev/my-project/images/worker-sidecar:latest",
+    shm_mount_path="/tmp/shared-memory",
+    shm_size_limit="16Gi",
+)
+```
+
+### GCSFuse Cloud Storage Mounts
+Mount GCS buckets to head, worker, or all containers:
+
+```python
+pw_jobset.add_gcsfuse(
+    containers="all",
+    mount_path="/data",
+    bucket="my-data-bucket",
+    read_only=True,
+)
+```
+
+### Custom `libtpu.so` Runtime Override
+Set `TPU_LIBRARY_PATH` across the Pathways containers (`pathways-rm`,
+`pathways-proxy`, and `pathways-worker` when `containers="all"`) without
+injecting it into user or sidecar containers:
+
+```python
+pw_jobset.with_custom_libtpu(
+    libtpu_path="/data/libtpu/libtpu.so",
+    containers="all",
+)
+```
+
+### Container & Pod Spec Overrides
+Customize arguments, environment variables, resources, or volumes on any container or pod:
+
+```python
+pw_jobset = jobset.PathwaysJobSet(
+    ...,
+    rm_container_kwargs={"args": ["--v=1"]},
+    proxy_container_kwargs={"env": {"XLA_FLAGS": "--xla_dump_to=/tmp"}},
+    worker_container_kwargs={"args": ["--stderrthreshold=0"]},
+    worker_pod_spec_kwargs={"priority_class_name": "high-priority"},
+)
+```
+
+## Serialization & Round-Tripping
+
+Convert between Python objects, YAML files, and dictionaries:
+
+```python
+# Save and load JobSet configurations
+pw_jobset.export_yaml("jobset.yaml")
+loaded_jobset = jobset.PathwaysJobSet.import_yaml("jobset.yaml")
+
+# Inspect raw Kubernetes resource dictionary
+config_dict = pw_jobset.to_dict()
+```