From 9be92bd5688724308173b52ad969aa5e5faaf26c Mon Sep 17 00:00:00 2001 From: Sujeeth Jinesh Date: Tue, 25 Aug 2026 20:42:21 -0700 Subject: [PATCH] No public description PiperOrigin-RevId: 970979032 --- .../colocated_controller.py | 57 +++++++------------ .../replicator_checkpoint_manager.py | 4 +- 2 files changed, 20 insertions(+), 41 deletions(-) diff --git a/checkpoint/orbax/checkpoint/experimental/emergency/multi_tier_checkpointing/colocated_controller.py b/checkpoint/orbax/checkpoint/experimental/emergency/multi_tier_checkpointing/colocated_controller.py index 184f660343..a6b7dbcafd 100644 --- a/checkpoint/orbax/checkpoint/experimental/emergency/multi_tier_checkpointing/colocated_controller.py +++ b/checkpoint/orbax/checkpoint/experimental/emergency/multi_tier_checkpointing/colocated_controller.py @@ -14,10 +14,11 @@ """Controller-side colocated orchestration for Pathways MTC.""" -from collections.abc import Mapping +from collections.abc import Mapping, Sequence import dataclasses import logging as python_logging import math +import os import threading import time from typing import Any, Callable @@ -49,7 +50,9 @@ PyTree = Any _STATE_ITEM_NAME = 'state' _DATASET_ITEM_NAME = 'dataset' -_LATEST_STEP_RETRY_TIMEOUT_SECS = 30 +_LATEST_STEP_RETRY_TIMEOUT_SECS = int( + os.environ.get('LATEST_STEP_RETRY_TIMEOUT_SECS', '180') +) # Abseil maps standard levels below DEBUG to increasing VLOG levels. _VLOG2_LEVEL = python_logging.DEBUG - 1 _RETRIABLE_COLOCATED_CALL_EXCEPTIONS = ( @@ -833,8 +836,8 @@ def _worker_dummy(self) -> jax.Array: """Returns a fresh dummy input array for worker-management RPCs.""" return dispatchers.get_dummy_input_array(self._worker_cpu_devices) - def latest_step(self) -> int | None: - """Returns the highest step present on every worker, or `None`.""" + def all_steps(self) -> Sequence[int]: + """Returns all common steps present on every worker, sorted.""" attempt = 0 deadline = time.time() + _LATEST_STEP_RETRY_TIMEOUT_SECS last_error = None @@ -842,7 +845,7 @@ def latest_step(self) -> int | None: attempt += 1 try: with TimeBlock( - f'Pathways colocated MTC latest_step attempt={attempt}', + f'Pathways colocated MTC all_steps attempt={attempt}', level=_VLOG2_LEVEL, ): result = self._worker_manager.all_steps(self._worker_dummy()) @@ -851,7 +854,7 @@ def latest_step(self) -> int | None: result, op_name='all_steps' ) if not worker_step_arrays: - return None + return [] worker_step_sets = [] for steps in worker_step_arrays: @@ -861,50 +864,28 @@ def latest_step(self) -> int | None: if int(step) != colocated_utils.NO_STEP_SENTINEL }) + if not worker_step_sets: + return [] common_steps = set.intersection(*worker_step_sets) - worker_latest_steps = [ - max(steps) if steps else None for steps in worker_step_sets - ] - if not common_steps: - logging.vlog( - 2, - 'Workers reported no common checkpoint steps: %s', - [sorted(steps) for steps in worker_step_sets], - ) - return None - latest_common_step = max(common_steps) - max_worker_step = max( - step for step in worker_latest_steps if step is not None - ) - if latest_common_step < max_worker_step: - logging.info( - 'Pathways colocated MTC latest_step selected lower common ' - 'step=%d while worker_latest_steps=%s.', - latest_common_step, - worker_latest_steps, - ) - else: - logging.vlog( - 2, - 'Pathways colocated MTC latest_step selected step=%d from ' - 'worker_latest_steps=%s.', - latest_common_step, - worker_latest_steps, - ) - return latest_common_step + return sorted(common_steps) except _RETRIABLE_COLOCATED_CALL_EXCEPTIONS as e: last_error = e logging.info( - 'latest_step transient failure on attempt=%s (%s), retrying...', + 'all_steps transient failure on attempt=%s (%s), retrying...', attempt, e, ) time.sleep(1) raise RuntimeError( - 'latest_step failed after retry budget' + 'all_steps failed after retry budget' f' ({_LATEST_STEP_RETRY_TIMEOUT_SECS}s).' ) from last_error + def latest_step(self) -> int | None: + """Returns the highest step present on every worker, or `None`.""" + steps = self.all_steps() + return max(steps) if steps else None + def should_save(self, step: int) -> bool: """Returns whether workers want to save `step`.""" if step == colocated_utils.NO_STEP_SENTINEL: diff --git a/checkpoint/orbax/checkpoint/experimental/emergency/multi_tier_checkpointing/replicator_checkpoint_manager.py b/checkpoint/orbax/checkpoint/experimental/emergency/multi_tier_checkpointing/replicator_checkpoint_manager.py index a2f3b12f64..8b947c3af4 100644 --- a/checkpoint/orbax/checkpoint/experimental/emergency/multi_tier_checkpointing/replicator_checkpoint_manager.py +++ b/checkpoint/orbax/checkpoint/experimental/emergency/multi_tier_checkpointing/replicator_checkpoint_manager.py @@ -934,9 +934,7 @@ def global_mesh(self) -> jax.sharding.Mesh: def all_steps(self, read: bool = False) -> Sequence[int]: if self._colocated_controller is not None: - raise NotImplementedError( - 'all_steps is not supported in colocated mode.' - ) + return self._colocated_controller.all_steps() return self._non_null_local_engine.all_steps(read=read) def latest_step(self) -> int | None: