pathwaysutils 0.1.9__tar.gz → 0.1.10__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pathwaysutils
-Version: 0.1.9
+Version: 0.1.10
 Summary: Package of Pathways-on-Cloud utilities.
 Keywords: 
 Author-email: Pathways-on-Cloud Utilities Developers <pathwaysutils-dev@google.com>

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

{pathwaysutils-0.1.9 → pathwaysutils-0.1.10}/pathwaysutils/elastic/manager.py

@@ -22,7 +22,7 @@ from collections.abc import Callable, Mapping, Sequence, Set
 import functools
 import logging
 import threading
-from typing import Any, TypeVar
+from typing import Any, TypeAlias, TypeVar
 import warnings
 
 import jax
@@ -48,6 +48,7 @@ class ScaleUpSignalError(Exception):
 
 
 _F = TypeVar("_F", bound=Callable[..., Any])
+RetryPolicy: TypeAlias = Callable[[int, Exception], bool]
 
 
 def _elastic_event_cleanup() -> None:
@@ -66,6 +67,19 @@ def _elastic_event_cleanup() -> None:
     array.delete()
 
 
+class ElasticRetryLimit:
+  """A retry callback that limits the number of attempts."""
+
+  def __init__(self, max_attempts: int):
+    if max_attempts <= 0:
+      raise ValueError("max_attempts must be positive.")
+    self.max_attempts = max_attempts
+
+  def __call__(self, attempt: int, error: Exception) -> bool:
+    del error  # Unused
+    return attempt < self.max_attempts
+
+
 class Manager:
   """Utility class for elastic training.
 
@@ -191,12 +205,14 @@ class Manager:
 
   def elastic_retry(
       self,
-      max_retries: int,
+      # TODO: b/523384843 - Remove max_retries parameter.
+      max_retries: int | None = None,
       minimum_slice_count: int | None = None,
       poll_interval: float | int = 10,
       timeout: float | None = None,
       pre_callback: Callable[..., Any] | None = None,
       on_elastic_event_callback: Callable[..., Any] | None = None,
+      retry_policy: RetryPolicy | None = None,
   ) -> Callable[[_F], _F]:
     """Retries a function with elasticity fault tolerance.
 
@@ -224,6 +240,8 @@ class Manager:
 
     Args:
       max_retries: The maximum number of times to retry the function.
+        Deprecated: Use `retry_policy` instead.
+        TODO: b/523384843 - Remove max_retries parameter.
       minimum_slice_count: The minimum number of slices required to run the
         function. If None, defaults to the total number of slices.
       poll_interval: The number of seconds to wait between activity checks.
@@ -233,6 +251,11 @@ class Manager:
       pre_callback: A callback to call before the function is attempted.
       on_elastic_event_callback: A callback to call after an elastic failure
         occurs.
+      retry_policy: A policy (callable) to determine if a retry should be
+        attempted. It accepts the attempt number (1-indexed) and the exception
+        that triggered the retry. If it returns False, no more retries are
+        attempted. If neither `retry_policy` nor `max_retries` is specified, it
+        defaults to unlimited retries.
 
     Returns:
       A decorator that retries the wrapped function.
@@ -248,17 +271,24 @@ class Manager:
         else minimum_slice_count
     )
 
-    if max_retries <= 0:
-      raise ValueError("max_retries must be positive.")
+    if max_retries is not None and retry_policy is not None:
+      raise ValueError("Cannot specify both max_retries and retry_policy.")
+
+    if retry_policy is None:
+      if max_retries is None:
+        # Default to unlimited retries if neither parameter is supplied.
+        retry_policy = lambda attempt, error: True
+      else:
+        if max_retries <= 0:
+          raise ValueError("max_retries must be positive.")
+        retry_policy = ElasticRetryLimit(max_retries)
 
     def decorator(func: _F) -> _F:
       @functools.wraps(func)
       def wrapper(*args: Any, **kwargs: Any) -> Any:
 
-        def attempt_execution(retry_index: int) -> Any:
-          _logger.info(
-              "Elastic attempt %d out of %d", retry_index + 1, max_retries
-          )
+        def attempt_execution(attempt: int) -> Any:
+          _logger.info("Elastic attempt %d", attempt)
           self.active_slice_indices = elastic.wait_for_slices(
               slice_count=target_slice_count,
               slice_to_devices=self.slice_to_devices,
@@ -289,34 +319,49 @@ class Manager:
               if monitor_thread is not None:
                 monitor_thread.join()
 
-        for retry_index in range(max_retries):
+        attempt = 1
+        while True:
           try:
-            return attempt_execution(retry_index)
-          except ScaleUpSignalError:
-            _logger.info("Scale up requested. Retrying.")
+            return attempt_execution(attempt)
+          except ScaleUpSignalError as error:
+            _logger.info("Scale up requested.")
             _elastic_event_cleanup()
 
             if on_elastic_event_callback is not None:
               on_elastic_event_callback()
+
+            if not retry_policy(attempt, error):
+              _logger.info(
+                  "Retry policy rejected retry after ScaleUpSignalError."
+              )
+              raise ElasticRuntimeError(
+                  f"Elastic attempt {attempt} failed."
+              ) from error
+
+            _logger.info("Retrying.")
           except jax.errors.JaxRuntimeError as error:
             if not elastic.is_error_due_to_slice_down(error):
               raise
 
             if self.new_slice_event.is_set():
-              _logger.info(
-                  "Slice down event and new slice available detected. Retrying."
-              )
+              _logger.info("Slice down event and new slice available detected.")
             else:
-              _logger.info("Slice down event detected. Retrying.")
+              _logger.info("Slice down event detected.")
 
             _elastic_event_cleanup()
 
             if on_elastic_event_callback is not None:
               on_elastic_event_callback()
-        else:
-          raise ElasticRuntimeError(
-              f"Elastic attempt {max_retries} out of {max_retries} failed."
-          )
+
+            if not retry_policy(attempt, error):
+              _logger.info("Retry policy rejected retry after JaxRuntimeError.")
+              raise ElasticRuntimeError(
+                  f"Elastic attempt {attempt} failed."
+              ) from error
+
+            _logger.info("Retrying.")
+
+          attempt += 1
 
       return wrapper
 

{pathwaysutils-0.1.9 → pathwaysutils-0.1.10}/pathwaysutils/experimental/gke/jobset.py

@@ -9,10 +9,11 @@
 # 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.
-"""Pathways JobSet generator and builder (Head Job Config)."""
+"""Pathways JobSet generator and builder (with Worker Job Config)."""
 
 import json
 import logging
+import math
 from typing import Any, Mapping
 from kubernetes import client
 
@@ -33,6 +34,7 @@ DEFAULT_PATHWAYS_PROXY_IMAGE = (
 
 PATHWAYS_PROXY_PORT = 29000
 PATHWAYS_RM_PORT = 29001
+PATHWAYS_WORKER_PORT = 29005
 
 MACHINE_TYPE_TO_TPU_VERSION_MAP = {
     "tpu7x-standard-4t": "tpu7x",
@@ -77,7 +79,7 @@ def _deserialize_dict(
 
 
 class PathwaysJobSet:
-  """Generates JobSet configuration for Pathways (with Head Job Config)."""
+  """JobSet configuration generator for Pathways."""
 
   def __init__(
       self,
@@ -90,6 +92,8 @@ class PathwaysJobSet:
       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,
       pathways_version: str = "latest",
       jobset_api_version: str = "v1alpha2",
       elastic_slices: int = 0,
@@ -108,6 +112,8 @@ class PathwaysJobSet:
       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.
       pathways_version: Version tag for Pathways images.
       jobset_api_version: API version of JobSet.
       elastic_slices: Number of elastic slices.
@@ -126,6 +132,19 @@ class PathwaysJobSet:
     if not tpu_version:
       raise ValueError(f"Unsupported TPU type: {tpu_type}")
 
+    gke_accel_type = MACHINE_TYPE_TO_GKE_ACCELERATOR_TYPE_MAP.get(
+        tpu_type.lower()
+    )
+
+    # Calculate VMs.
+    dims = [int(x) for x in topology.split("x")]
+    total_chips = math.prod(dims)
+    chips_per_vm = 8 if tpu_type.lower().endswith("8t") else 4
+    if total_chips < chips_per_vm:
+      num_vms = 1
+    else:
+      num_vms = total_chips // chips_per_vm
+
     instance_type = f"{tpu_version}:{topology}"
     image_tag = pathways_version
 
@@ -140,8 +159,17 @@ class PathwaysJobSet:
         elastic_slices=elastic_slices,
     )
 
-    # Build minimal worker template (placeholder)
-    self._worker_job_template = self._build_minimal_job_template("worker")
+    # Build worker template.
+    self._worker_job_template = self._build_worker_job_template(
+        pathways_dir=pathways_dir,
+        num_vms=num_vms,
+        chips_per_vm=chips_per_vm,
+        gke_accel_type=gke_accel_type,
+        topology=topology,
+        image_tag=image_tag,
+        max_slice_restarts=max_slice_restarts,
+        termination_grace_period_seconds=termination_grace_period_seconds,
+    )
 
     self._success_policy = None
     if user_pod_template:
@@ -150,20 +178,6 @@ class PathwaysJobSet:
           "targetReplicatedJobs": [PATHWAYS_HEAD_JOB_NAME],
       }
 
-  def _build_minimal_job_template(self, role: str) -> client.V1JobTemplateSpec:
-    """Builds a minimal job template for a given role."""
-    pod_spec = client.V1PodSpec(
-        containers=[
-            client.V1Container(name=f"placeholder-{role}", image="ubuntu")
-        ]
-    )
-    job_spec = client.V1JobSpec(
-        template=client.V1PodTemplateSpec(
-            metadata=client.V1ObjectMeta(labels={"role": role}), spec=pod_spec
-        )
-    )
-    return client.V1JobTemplateSpec(spec=job_spec)
-
   def _build_head_job_template(
       self,
       pathways_dir: str,
@@ -365,6 +379,157 @@ class PathwaysJobSet:
     )
     return head_job_template
 
+  def _build_worker_job_template(
+      self,
+      pathways_dir: str,
+      num_vms: int,
+      chips_per_vm: int,
+      gke_accel_type: str,
+      topology: str,
+      image_tag: str,
+      max_slice_restarts: int,
+      termination_grace_period_seconds: int | None,
+  ) -> client.V1JobTemplateSpec:
+    worker_image = f"{DEFAULT_PATHWAYS_RM_AND_WORKER_IMAGE}:{image_tag}"
+
+    args = [
+        f"--resource_manager_address=$(PATHWAYS_HEAD):{PATHWAYS_RM_PORT}",
+        f"--server_port={PATHWAYS_WORKER_PORT}",
+        f"--gcs_scratch_location={pathways_dir}",
+    ]
+    worker_env = [
+        client.V1EnvVar(name="TPU_MIN_LOG_LEVEL", value="0"),
+        client.V1EnvVar(name="TF_CPP_MIN_LOG_LEVEL", value="0"),
+        client.V1EnvVar(name="XCLOUD_ENVIRONMENT", value="GCP"),
+        client.V1EnvVar(name="MEGASCALE_GRPC_ENABLE_XOR_TRACER", value="false"),
+        client.V1EnvVar(
+            name="MEGASCALE_NUM_SLICES",
+            value_from=client.V1EnvVarSource(
+                field_ref=client.V1ObjectFieldSelector(
+                    field_path="metadata.labels['jobset.sigs.k8s.io/replicatedjob-replicas']"
+                )
+            ),
+        ),
+        client.V1EnvVar(
+            name="JOBSET_NAME",
+            value_from=client.V1EnvVarSource(
+                field_ref=client.V1ObjectFieldSelector(
+                    field_path=(
+                        "metadata.annotations['jobset.sigs.k8s.io/jobset-name']"
+                    )
+                )
+            ),
+        ),
+        client.V1EnvVar(
+            name="REPLICATED_JOB_NAME",
+            value_from=client.V1EnvVarSource(
+                field_ref=client.V1ObjectFieldSelector(
+                    field_path="metadata.annotations['jobset.sigs.k8s.io/replicatedjob-name']"
+                )
+            ),
+        ),
+        client.V1EnvVar(
+            name="MEGASCALE_SLICE_ID",
+            value_from=client.V1EnvVarSource(
+                field_ref=client.V1ObjectFieldSelector(
+                    field_path="metadata.labels['jobset.sigs.k8s.io/job-index']"
+                )
+            ),
+        ),
+        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="MEGASCALE_COORDINATOR_ADDRESS",
+            value_from=client.V1EnvVarSource(
+                field_ref=client.V1ObjectFieldSelector(
+                    field_path=(
+                        "metadata.labels['jobset.sigs.k8s.io/coordinator']"
+                    )
+                )
+            ),
+        ),
+    ]
+
+    worker_container = client.V1Container(
+        name="pathways-worker",
+        image=worker_image,
+        image_pull_policy="Always",
+        args=args,
+        env=worker_env,
+        ports=[
+            client.V1ContainerPort(
+                container_port=PATHWAYS_WORKER_PORT, protocol="TCP"
+            ),
+            client.V1ContainerPort(container_port=29006, protocol="TCP"),
+            client.V1ContainerPort(container_port=8471, protocol="TCP"),
+            client.V1ContainerPort(container_port=8080, protocol="TCP"),
+        ],
+        volume_mounts=[
+            client.V1VolumeMount(name="shared-tmp", mount_path="/tmp")
+        ],
+        resources=client.V1ResourceRequirements(
+            limits={"google.com/tpu": str(chips_per_vm)}
+        ),
+    )
+
+    node_selector = {
+        "cloud.google.com/gke-tpu-accelerator": gke_accel_type,
+        "cloud.google.com/gke-tpu-topology": topology,
+    }
+
+    backoff_limit = num_vms * 4
+    if max_slice_restarts > 0:
+      backoff_limit = num_vms * max_slice_restarts
+
+    worker_pod_spec = client.V1PodSpec(
+        containers=[worker_container],
+        node_selector=node_selector,
+        volumes=[
+            client.V1Volume(
+                name="shared-tmp",
+                host_path=client.V1HostPathVolumeSource(
+                    path="/tmp", type="DirectoryOrCreate"
+                ),
+            )
+        ],
+        host_network=True,
+        dns_policy="ClusterFirstWithHostNet",
+        restart_policy="OnFailure",
+    )
+    if termination_grace_period_seconds is not None:
+      worker_pod_spec.termination_grace_period_seconds = (
+          termination_grace_period_seconds
+      )
+
+    worker_job_template = client.V1JobTemplateSpec(
+        metadata=client.V1ObjectMeta(),
+        spec=client.V1JobSpec(
+            backoff_limit=backoff_limit,
+            completion_mode="Indexed",
+            completions=num_vms,
+            parallelism=num_vms,
+            template=client.V1PodTemplateSpec(
+                metadata=client.V1ObjectMeta(
+                    annotations={
+                        "alpha.jobset.sigs.k8s.io/exclusive-topology": (
+                            "cloud.google.com/gke-nodepool"
+                        )
+                    }
+                ),
+                spec=worker_pod_spec,
+            ),
+        ),
+    )
+    return worker_job_template
+
   def _compile_config(self) -> dict[str, Any]:
     """Compiles the JobSet configuration into a dictionary."""
     with client.ApiClient() as api_client:

pathwaysutils-0.1.10/pathwaysutils/experimental/reshard.py

@@ -0,0 +1,218 @@
+# Copyright 2025 Google LLC
+#
+# 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.
+"""Experimental resharding API for elastic device sets."""
+
+import base64
+from collections.abc import Callable, Mapping, Sequence
+import json
+import logging
+from typing import Any
+
+import jax
+from pathwaysutils import lru_cache
+from pathwaysutils import plugin_executable
+from pathwaysutils import reshard as pw_reshard
+
+# Redirects for generic intermediate sharding helpers and exceptions
+NoIntermediateShardingError = pw_reshard.NoIntermediateShardingError
+NoIntermediateShardingNeededError = pw_reshard.NoIntermediateShardingNeededError
+find_intermediate_sharding = pw_reshard.find_intermediate_sharding
+reshard_with_intermediate_sharding_generic = (
+    pw_reshard.reshard_with_intermediate_sharding_generic
+)
+
+_logger = logging.getLogger(__name__)
+
+
+def _identity(x: Any) -> Any:
+  """A helper function that returns its input."""
+  return x
+
+
+class ReshardingPlanWrapper:
+  """Wrapper around PluginProgram(reshard_request)."""
+
+  _plugin_executable: plugin_executable.PluginExecutable
+  _avals: Sequence[jax.core.ShapedArray]
+  _out_shardings: Sequence[jax.sharding.Sharding]
+
+  def __init__(
+      self,
+      avals: Sequence[jax.core.ShapedArray],
+      source_shardings: Sequence[jax.sharding.Sharding],
+      destination_shardings: Sequence[jax.sharding.Sharding],
+      donate: bool,
+  ):
+
+    def ifrt_hlo_sharding(
+        aval: jax.core.ShapedArray, sharding: jax.sharding.Sharding
+    ) -> Mapping[str, Any]:
+      result = {
+          "devices": {
+              "device_ids": [
+                  device.id for device in sharding._addressable_device_assignment  # pylint: disable=protected-access
+              ]
+          },
+          "xla_hlo_sharding": (
+              base64.b64encode(
+                  sharding._to_xla_hlo_sharding(aval.ndim)  # pylint: disable=protected-access
+                  .to_proto()
+                  .SerializeToString()
+              ).decode("utf-8")
+          ),
+      }
+      if sharding.memory_kind is not None:
+        result["memory_kind"] = sharding.memory_kind
+      return result
+
+    request = {
+        "reshardRequest": {
+            "donateInput": donate,
+            "inSharding": [
+                ifrt_hlo_sharding(aval, old_sharding)
+                for aval, old_sharding in zip(avals, source_shardings)
+            ],
+            "outSharding": [
+                ifrt_hlo_sharding(aval, new_sharding)
+                for aval, new_sharding in zip(avals, destination_shardings)
+            ],
+        }
+    }
+
+    self._plugin_executable = plugin_executable.PluginExecutable(
+        json.dumps(request)
+    )
+    self._avals = avals
+    self._out_shardings = destination_shardings
+
+  def execute(self, inp_arrays: tuple[jax.Array, ...]) -> Sequence[jax.Array]:
+    out_arrays, fut = self._plugin_executable.call(
+        inp_arrays, self._out_shardings, self._avals
+    )
+    fut.result()
+    return out_arrays
+
+
+def _get_resharding_plan(
+    avals: tuple[jax.core.ShapedArray, ...],
+    old_shardings: tuple[jax.sharding.Sharding, ...],
+    new_shardings: tuple[jax.sharding.Sharding, ...],
+    donate: bool,
+) -> ReshardingPlanWrapper:
+  """Returns a resharding plan for the given sharding task."""
+  return ReshardingPlanWrapper(avals, old_shardings, new_shardings, donate)
+
+
+_get_resharding_plan_cached = lru_cache.lru_cache()(_get_resharding_plan)
+
+
+def _sidechannel_jax_array_reshard(
+    array_info: Mapping[str, Any], *, donate: bool, cache_resharding_plans: bool
+) -> Sequence[jax.Array]:
+  get_resharding_plan_func = (
+      _get_resharding_plan_cached
+      if cache_resharding_plans
+      else _get_resharding_plan
+  )
+  return get_resharding_plan_func(
+      tuple(arr.aval for arr in array_info["arrays"]),
+      tuple(arr.sharding for arr in array_info["arrays"]),
+      tuple(array_info["dst_shardings"]),
+      donate,
+  ).execute(tuple(array_info["arrays"]))
+
+
+def reshard(
+    x: Any,
+    sharding: jax.sharding.Sharding | Any,
+    *,
+    donate: bool = False,
+    may_alias: bool | None = None,
+    cache_resharding_plans: bool = False,
+) -> Any:
+  """Reshards `x` to `sharding`.
+
+  This function is an alternative to `pathwaysutils.reshard` (which is the IFRT
+  reshard) that uses the sidechannel MPMD resharding API for the final reshard.
+
+  Args:
+    x: An array, scalar, or (nested) standard Python container thereof.
+    sharding: A `Sharding` or a (nested) `Sharding` in standard Python container
+      (must be a tree prefix of `x`), representing the device(s) and sharding to
+      which `x` should be sharded to. The result will be committed to the
+      device(s) of the sharding.
+    donate: If `True`, donate all input arrays, which may reduce the amount of
+      memory needed for resharding. Buffers donated to resharding should not be
+      reused.
+    may_alias: If `True`, may alias the input array with the output array. May
+      reduce the amount of memory needed for resharding. Not used at the moment.
+    cache_resharding_plans: If `True`, uses a resharding plan cache to avoid
+      recreating plans for the same resharding operation.
+
+  Returns:
+    A copy of `x` whose sharding is `sharding`.
+  """
+  return pw_reshard.reshard_generic(
+      x,
+      sharding,
+      donate=donate,
+      may_alias=may_alias,
+      jax_array_reshard_fn=_sidechannel_jax_array_reshard,
+      cache_resharding_plans=cache_resharding_plans,
+  )
+
+
+def sidechannel_reshard_with_intermediate_sharding(
+    x: Any,
+    in_sharding: jax.sharding.Sharding,
+    out_sharding: jax.sharding.Sharding,
+    *,
+    donate: bool = False,
+    may_alias: bool | None = None,
+    cache_resharding_plans: bool = False,
+) -> Any:
+  """Reshards `x` to `out_sharding`, using an intermediate sharding if possible.
+
+  This function is an alternative to `reshard` that may be faster and sometimes
+  essential for certain sharding combinations by using an intermediate sharding
+  to avoid expensive all-gathers. If no beneficial intermediate sharding is
+  found, it falls back to standard resharding. See `find_intermediate_sharding`
+  for more details on when an intermediate sharding is used.
+
+  Uses the sidechannel resharding API for the final reshard.
+
+  Args:
+    x: An array, scalar, or (nested) standard Python container thereof.
+    in_sharding: The source sharding of `x`.
+    out_sharding: The target sharding for `x`.
+    donate: If `True`, donate all input arrays, which may reduce the amount of
+      memory needed for resharding. Buffers donated to resharding should not be
+      reused.
+    may_alias: If `True`, may alias the input array with the output array. May
+      reduce the amount of memory needed for resharding. Not used at the moment.
+    cache_resharding_plans: If `True`, uses a resharding plan cache to avoid
+      recreating plans for the same resharding operation.
+
+  Returns:
+    A copy of `x` whose sharding is `out_sharding`.
+  """
+  return reshard_with_intermediate_sharding_generic(
+      x,
+      in_sharding,
+      out_sharding,
+      donate=donate,
+      may_alias=may_alias,
+      jax_array_reshard_fn=_sidechannel_jax_array_reshard,
+      cache_resharding_plans=cache_resharding_plans,
+  )