pathwaysutils 0.1.7__tar.gz → 0.1.9__tar.gz

{pathwaysutils-0.1.7 → pathwaysutils-0.1.9}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pathwaysutils
-Version: 0.1.7
+Version: 0.1.9
 Summary: Package of Pathways-on-Cloud utilities.
 Keywords: 
 Author-email: Pathways-on-Cloud Utilities Developers <pathwaysutils-dev@google.com>
@@ -10,7 +10,7 @@ Classifier: License :: OSI Approved :: Apache Software License
 License-File: LICENSE
 Requires-Dist: absl-py
 Requires-Dist: fastapi
-Requires-Dist: jax>=0.8.0
+Requires-Dist: jax>=0.8.3
 Requires-Dist: orbax-checkpoint
 Requires-Dist: uvicorn
 Requires-Dist: requests

{pathwaysutils-0.1.7 → pathwaysutils-0.1.9}/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.7"
+__version__ = "v0.1.9"

{pathwaysutils-0.1.7 → pathwaysutils-0.1.9}/pathwaysutils/elastic/elastic.py

@@ -18,10 +18,13 @@ retries a function in case of `jax.errors.JaxRuntimeError` caused by slice down
 events. It also provides a utility for waiting for slices to become active.
 """
 
+from abc import ABC, abstractmethod
 import collections
-from collections.abc import Mapping, Sequence, Set
+from collections.abc import Callable, Mapping, Sequence, Set
+import itertools
 import logging
 import time
+from typing import Any
 
 import jax
 import numpy as np
@@ -41,44 +44,87 @@ _ELASTIC_DOWN_ADDITIONAL_ERROR_TYPES = frozenset([
 ])
 
 
-def _plus_one(x: jax.Array) -> jax.Array:
-  """Adds one to each element in the array.
+class SliceHealthChecker(ABC):
+  """Base class for slice health checkers.
 
-  Used to test if a slice is active.
+  Implementations must provide `dispatch` and `validate` methods.
+  `dispatch` starts the health check operations, and `validate` waits for them
+  and returns the active slice indices.
+  """
 
-  Args:
-    x: The array to add one to.
+  def __init__(self, slice_to_devices: Mapping[int, Sequence[jax.Device]]):
+    self.slice_to_devices = slice_to_devices
 
-  Returns:
-      The array with one added to each element.
-  """
-  return x + 1
+  @abstractmethod
+  def dispatch(self) -> None:
+    """Dispatches the JAX operations in the background."""
 
+  @abstractmethod
+  def validate(self) -> Set[int]:
+    """Blocks on results and returns the set of active slice indices.
 
-def _simple_execution(devices: Sequence[jax.Device]) -> jax.Array:
-  """Simple execution to test if a slice is active.
+    Returns:
+      Set of active slice indices.
+    """
 
-  This function is used to test if a slice is active. It executes a simple
-  computation on the devices and returns the result. If any of the devices are
-  not active, the returned array will fail with a JaxRuntimeError used.
 
-  Simply executing this function is not enough to determine if the slice is
-  active. We also need to check the value of the returned array.
+class DefaultSliceHealthChecker(SliceHealthChecker):
+  """Default implementation that checks each slice individually.
 
-  Args:
-    devices: The devices to execute on.
+  It executes a simple computation on each slice independently to verify
+  that the devices are active and responsive.
 
-  Returns:
-    The result of the execution.
+  The computation involves creating a zero array of the same size as the
+  number of devices in the slice, adding (_SIMPLE_EXECUTION_TEST_VALUE - 1)
+  to it, and then running a pmap to add 1 to each element. This verifies
+  that the devices are active and can perform computations.
   """
-  if not devices:
-    raise ValueError("No devices")
 
-  test_input = np.zeros(len(devices), dtype=float) + (
-      _SIMPLE_EXECUTION_TEST_VALUE - 1
-  )
+  def _plus_one(self, x: jax.Array) -> jax.Array:
+    return x + 1
+
+  def _simple_execution(self, devices: Sequence[jax.Device]) -> jax.Array:
+    if not devices:
+      raise ValueError("No devices")
+
+    test_input = np.zeros(len(devices), dtype=float) + (
+        _SIMPLE_EXECUTION_TEST_VALUE - 1
+    )
+
+    return jax.pmap(self._plus_one, devices=devices)(test_input)
+
+  def dispatch(self) -> None:
+    self.results = {
+        slice_index: self._simple_execution(devices)
+        for slice_index, devices in self.slice_to_devices.items()
+    }
+
+  def validate(self) -> Set[int]:
+    active_slice_indices = set()
+    for slice_index, x in self.results.items():
+      expected = (
+          np.zeros(len(self.slice_to_devices[slice_index]), dtype=float)
+          + _SIMPLE_EXECUTION_TEST_VALUE
+      )
+      try:
+        jax.block_until_ready(x)
+        if np.allclose(x, expected):
+          active_slice_indices.add(slice_index)
+        else:
+          msg = (
+              f"Error with _simple_execution for slice_index={slice_index}. "
+              f"Expected: {expected!r}, Actual: {x!r}"
+          )
+          _logger.error(msg)
+          raise ValueError(msg)
+      except jax.errors.JaxRuntimeError as error:
+        _logger.debug(
+            "Caught JaxRuntimeError for slice_index=%s: %s", slice_index, error
+        )
+        if not is_error_due_to_slice_down(error):
+          raise
+    return active_slice_indices
 
-  return jax.pmap(_plus_one, devices=devices)(test_input)
 
 
 def get_slice_to_devices(
@@ -94,6 +140,7 @@ def get_slice_to_devices(
 @timing.timeit
 def get_active_slice_indices(
     slice_to_devices: Mapping[int, Sequence[jax.Device]] | None = None,
+    checker: SliceHealthChecker | None = None,
 ) -> Set[int]:
   """Returns the set of active slices indices.
 
@@ -101,6 +148,8 @@ def get_active_slice_indices(
     slice_to_devices: A mapping from slice index to devices. If None,
       `get_slice_to_devices(jax.devices())` is used to gather all available
       devices and group them by slice.
+    checker: A SliceHealthChecker instance. If None, DefaultSliceHealthChecker
+      is used.
 
   Returns:
     A set of integers representing the indices of the active slices.
@@ -109,55 +158,16 @@ def get_active_slice_indices(
     _logger.debug("slice_to_devices is None. Getting from jax.devices().")
     slice_to_devices = get_slice_to_devices(tuple(jax.devices()))
 
+  if checker is None:
+    checker = DefaultSliceHealthChecker(slice_to_devices)
+
   _logger.debug(
       "Getting active slice indices for slices: %s",
       sorted(list(slice_to_devices.keys())),
   )
 
-  active_slice_indices = set()
-
-  results = {
-      slice_index: _simple_execution(devices)
-      for slice_index, devices in slice_to_devices.items()
-  }
-
-  for slice_index, x in results.items():
-    _logger.debug("Checking slice_index=%s", slice_index)
-    expected = (
-        np.zeros(len(slice_to_devices[slice_index]), dtype=float)
-        + _SIMPLE_EXECUTION_TEST_VALUE
-    )
-    try:
-      with timing.Timer(f"Checking {slice_index=}"):
-        _logger.debug("Blocking until ready for slice_index=%s", slice_index)
-        jax.block_until_ready(x)
-        _logger.debug("Execution finished for slice_index=%s", slice_index)
-        if np.allclose(x, expected):
-          active_slice_indices.add(slice_index)
-          _logger.debug("slice_index=%s active", slice_index)
-        else:
-          _logger.error(
-              "Error with _simple_execution for slice_index=%s. "
-              "This should never happen. Expected: %r, Actual: %r",
-              slice_index,
-              expected,
-              x,
-          )
-          raise ValueError(
-              f"Error with _simple_execution for slice_index={slice_index}."
-          )
-    except jax.errors.JaxRuntimeError as error:
-      _logger.debug(
-          "Caught JaxRuntimeError for slice_index=%s: %s", slice_index, error
-      )
-      if not is_error_due_to_slice_down(error):
-        _logger.info("Re-raising error for slice_index=%s", slice_index)
-        raise
-      _logger.debug("slice_index=%s bad", slice_index)
-
-  _logger.debug("active_slice_indices=%s", active_slice_indices)
-
-  return active_slice_indices
+  checker.dispatch()
+  return checker.validate()
 
 
 def wait_for_slices(
@@ -165,6 +175,7 @@ def wait_for_slices(
     poll_interval: float | int = 10,
     timeout: float | int | None = None,
     slice_to_devices: Mapping[int, Sequence[jax.Device]] | None = None,
+    checker: SliceHealthChecker | None = None,
 ) -> Set[int]:
   """Waits until after at least `slice_count` slices become active.
 
@@ -177,6 +188,7 @@ def wait_for_slices(
       timeout.
     slice_to_devices: A mapping from slice index to devices. If None,
       `get_slice_to_devices(jax.devices())` is used.
+    checker: A SliceHealthChecker instance to use for checking health.
 
   Returns:
     The active slice indices
@@ -201,7 +213,7 @@ def wait_for_slices(
     check_start_time = time.time()
 
     _logger.debug("Checking active slices...")
-    active_slice_indices = get_active_slice_indices(slice_to_devices)
+    active_slice_indices = get_active_slice_indices(slice_to_devices, checker)
     if len(active_slice_indices) >= slice_count:
       _logger.info(
           "Sufficient slices active: %s >= %s. Active indices: %s",

pathwaysutils-0.1.9/pathwaysutils/experimental/concatenate_by_mesh_axis.py

@@ -0,0 +1,152 @@
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Pathwaysutils JAX concatenate_by_mesh_axis."""
+
+from collections.abc import Sequence
+import itertools
+from typing import Any
+import jax
+import numpy as np
+from pathwaysutils import jax as pw_jax
+
+
+def concatenate_by_mesh_axis(
+    array_trees: Sequence[Any],
+    mesh_axis: str,
+) -> Any:
+  """Concatenates meshes by an axis. Returns arrays on the concatenated mesh.
+
+  Note: This API always donates the given arrays.
+
+  Args:
+    array_trees: Sequence of PyTrees of JAX arrays with `NamedSharding`. All
+      PyTrees in the sequence have the same structure. All arrays in each PyTree
+      are sharded/replicated on the same mesh. The input arrays are always
+      donated.
+    mesh_axis: Mesh axis to concatenate.
+
+  Returns:
+    A PyTree with the same structure as `array_trees[i]`. It has arrays with
+    their shards concatenated to match a concatenated mesh.
+  """
+  if not array_trees:
+    return array_trees
+
+  def _get_named_sharding(array: jax.Array) -> jax.sharding.NamedSharding:
+    if not isinstance(array, jax.Array):
+      raise ValueError(f"Elements must be jax.Array. Got {type(array)}")
+    sharding = array.sharding
+    if not isinstance(sharding, jax.sharding.NamedSharding):
+      raise ValueError(f"Expected NamedSharding. Got {type(sharding)}")
+    return sharding
+
+  flats_and_defs = [jax.tree_util.tree_flatten(at) for at in array_trees]
+  flattened_arrays = [fd[0] for fd in flats_and_defs]
+  treedefs = [fd[1] for fd in flats_and_defs]
+  input_treedef = treedefs[0]
+  for td in treedefs[1:]:
+    if td != input_treedef:
+      raise ValueError(
+          "All array trees must have the same treedef. Got"
+          f" {td} vs. {input_treedef}"
+      )
+
+  # Convert to have the output array structure in the outer list, and each entry
+  # be a list of arrays from each shard for the concatenated output array.
+  input_flat_arrays = list(zip(*flattened_arrays))
+
+  if not flattened_arrays[0]:
+    return array_trees[0]
+
+  # Extract the shared mesh from each PyTree (from an arbitrary array in each).
+  meshes_to_concatenate = [
+      _get_named_sharding(array).mesh for array in input_flat_arrays[0]
+  ]
+
+  # Validate that the meshes are compatible.
+  reference_mesh = meshes_to_concatenate[0]
+  if mesh_axis not in reference_mesh.axis_names:
+    raise ValueError(
+        f"mesh_axis '{mesh_axis}' not found in mesh axis names:"
+        f" {reference_mesh.axis_names}"
+    )
+  mesh_axis_idx = reference_mesh.axis_names.index(mesh_axis)
+  for mesh in meshes_to_concatenate:
+    if mesh.axis_names != reference_mesh.axis_names:
+      raise ValueError(
+          "Meshes must have the same axis names. Got"
+          f" {mesh} vs. {reference_mesh}."
+      )
+    if (
+        mesh.axis_sizes[:mesh_axis_idx]
+        != reference_mesh.axis_sizes[:mesh_axis_idx]
+        or mesh.axis_sizes[mesh_axis_idx + 1 :]
+        != reference_mesh.axis_sizes[mesh_axis_idx + 1 :]
+    ):
+      raise ValueError(
+          "Arrays must have the same mesh axis sizes for all axes except"
+          f" {mesh_axis}. Got {mesh} vs. {reference_mesh}."
+      )
+
+  # Construct list of the mesh axis section boundaries.
+  devices = [mesh.devices for mesh in meshes_to_concatenate]
+  mesh_axis_sections = list(
+      itertools.accumulate(
+          mesh.axis_sizes[mesh_axis_idx] for mesh in meshes_to_concatenate
+      )
+  )
+
+  concatenated_mesh = jax.sharding.Mesh(
+      np.concatenate(devices, mesh_axis_idx),
+      axis_names=reference_mesh.axis_names,
+      axis_types=reference_mesh.axis_types,
+  )
+
+  def _get_output_sharding(
+      arrays: Sequence[jax.Array],
+  ) -> jax.sharding.NamedSharding:
+    reference_sharding = _get_named_sharding(arrays[0])
+    reference_spec = reference_sharding.spec
+    return jax.sharding.NamedSharding(
+        concatenated_mesh,
+        reference_spec,
+        memory_kind=reference_sharding.memory_kind,
+    )
+
+  def _sharded_dim_idx_for_sharding(
+      sharding: jax.sharding.NamedSharding,
+  ) -> int:
+    sharded_dim_idx = -1
+    for dim_idx, dim_spec in enumerate(sharding.spec):
+      flat_dim_spec, _ = jax.tree_util.tree_flatten(dim_spec)
+      if mesh_axis in flat_dim_spec:
+        sharded_dim_idx = dim_idx
+        break
+    return sharded_dim_idx
+
+  out_shardings = [_get_output_sharding(arrays) for arrays in input_flat_arrays]
+  sharded_dim_idxs = [
+      _sharded_dim_idx_for_sharding(sharding) for sharding in out_shardings
+  ]
+
+  flat_output_arrays = pw_jax.concatenate_by_mesh_axis(
+      arrays=input_flat_arrays,
+      sharded_dim_idxs=sharded_dim_idxs,
+      mesh_axis_sizes=concatenated_mesh.axis_sizes,
+      mesh_axis_idx=mesh_axis_idx,
+      mesh_axis_sections=mesh_axis_sections,
+      out_shardings=out_shardings,
+      donate=True,
+  )
+
+  return jax.tree_util.tree_unflatten(input_treedef, flat_output_arrays)

pathwaysutils-0.1.9/pathwaysutils/experimental/gke/__init__.py

@@ -0,0 +1,11 @@
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.