Add checkpoint_context.py to map flat orbax v0 flags to Orbax v1 Context.

This new module provides functions to build Orbax v1 Context, SaveDecisionPolicy, and PreservationPolicy objects based on MaxText's checkpoint configuration flags. It centralizes the logic for configuring Orbax v1, replacing various scattered v0 Orbax constructs previously used in MaxText. Unit tests are included to verify the correct mapping of flags to Orbax v1 options.

PiperOrigin-RevId: 944311354

Author: angel-core

src/dependencies/extra_deps/cuda12_overrides.txt

@@ -1,3 +1,4 @@
 datasets>=4.8.5
 fsspec==2026.2.0
 gcsfs==2026.2.0
+orbax-checkpoint>=0.12.1

src/dependencies/requirements/generated_requirements/cuda12-requirements.txt

@@ -30,7 +30,7 @@ cloudpickle>=3.1.2
 clu>=0.0.12
 colorama>=0.4.6
 contourpy>=1.3.3
-coverage>=7.14.3
+coverage>=7.15.0
 cryptography>=49.0.0
 cycler>=0.12.1
 datasets>=5.0.0
@@ -84,11 +84,11 @@ grpcio-status>=1.80.0
 gviz-api>=1.10.0
 h11>=0.16.0
 h5py>=3.14.0
-hf-xet>=1.5.1 ; platform_machine == 'AMD64' or platform_machine == 'aarch64' or platform_machine == 'amd64' or platform_machine == 'arm64' or platform_machine == 'x86_64'
+hf-xet>=1.5.1 ; platform_machine == 'AMD64' or platform_machine == 'aarch64' or platform_machine == 'x86_64'
 httpcore>=1.0.9
 httplib2>=0.32.0
 httpx>=0.28.1
-huggingface-hub>=1.21.0
+huggingface-hub>=1.22.0
 humanize>=4.16.0
 hypothesis>=6.142.1
 identify>=2.6.19
@@ -158,7 +158,7 @@ opt-einsum>=3.4.0
 optax>=0.2.8
 optree>=0.19.0
 optype>=0.17.0
-orbax-checkpoint>=0.11.39
+orbax-checkpoint>=0.12.1
 packaging>=26.1
 pandas>=3.0.3
 parameterized>=0.9.0
@@ -169,6 +169,7 @@ platformdirs>=4.10.0
 pluggy>=1.6.0
 portpicker>=1.6.0
 pre-commit>=4.6.0
+prometheus-client>=0.20.0
 promise>=2.3
 propcache>=0.5.2
 proto-plus>=1.28.0
@@ -195,7 +196,7 @@ pyproject-hooks>=1.2.0
 pytest>=8.4.2
 pytest-xdist>=3.8.0
 python-dateutil>=2.9.0.post0
-python-discovery>=1.4.2
+python-discovery>=1.4.3
 pytokens>=0.4.1
 pytype>=2024.10.11
 pyyaml>=6.0.3
@@ -229,7 +230,7 @@ tensorflow>=2.20.0
 tensorflow-datasets>=4.9.10
 tensorflow-metadata>=1.21.0
 tensorflow-text>=2.20.1
-tensorstore>=0.1.82
+tensorstore>=0.1.84
 termcolor>=3.3.0
 tiktoken>=0.13.0
 tokamax>=0.0.12
@@ -241,24 +242,24 @@ tqdm>=4.68.3
 transformer-engine>=2.16.1
 transformer-engine-cu12>=2.16.1
 transformer-engine-jax>=2.16.1
-transformers>=5.12.1
+transformers>=5.13.0
 treescope>=0.1.10
 typeguard>=2.13.3
-typer>=0.25.1
+typer>=0.26.8
 typing-extensions>=4.15.0
 typing-inspection>=0.4.2
 tzdata>=2026.2 ; sys_platform == 'emscripten' or sys_platform == 'win32'
 uritemplate>=4.2.0
 urllib3>=2.6.3
-uvicorn>=0.49.0
+uvicorn>=0.50.2
 uvloop>=0.22.1
-virtualenv>=21.5.1
+virtualenv>=21.5.2
 wadler-lindig>=0.1.7
 websockets>=16.0
 werkzeug>=3.1.8
 wheel>=0.46.3
 wrapt>=2.2.2
-xxhash>=3.8.0
+xxhash>=3.8.1
 yarl>=1.24.2
 zipp>=3.23.1
 zstandard>=0.25.0

src/maxtext/common/checkpoint_context.py

@@ -0,0 +1,171 @@
+# Copyright 2026 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.
+
+"""Maps MaxText checkpoint config onto the Orbax v1 Context + training policies.
+
+This module is the single place that translates MaxText's flat checkpoint flags
+into those objects. It builds configuration only.
+"""
+import datetime
+
+from orbax.checkpoint import pathways as ocp_pathways
+from orbax.checkpoint import v1 as ocp
+
+
+# v0 PyTreeCheckpointHandler converts `*_concurrent_gb` with GB = 10**9 bytes.
+_BYTES_PER_GB = 10**9
+
+# Matches the v0 SingleReplicaArrayHandler broadcast limit (1000 MB) that
+# MaxText used when restoring a single replica and broadcasting to the rest.
+_SINGLE_REPLICA_BROADCAST_MEMORY_LIMIT_BYTES = 1024 * 1024 * 1000
+
+
+def build_save_decision_policy(
+    *,
+    save_interval_steps: int,
+    enable_continuous_checkpointing: bool = False,
+    enable_autocheckpoint: bool = False,
+) -> ocp.training.save_decision_policies.SaveDecisionPolicy:
+  """Builds the v1 SaveDecisionPolicy.
+
+  - continuous: save as often as possible (async-friendly).
+  - autocheckpoint: save on preemption OR at the fixed interval.
+  - otherwise: save at the fixed interval.
+
+  Args:
+    save_interval_steps: Save every N steps.
+    enable_continuous_checkpointing: If true, save as often as possible.
+    enable_autocheckpoint: If true, save on preemption OR at the fixed interval.
+
+  Returns:
+    A configured ``ocp.training.save_decision_policies.SaveDecisionPolicy``.
+  """
+  policies = ocp.training.save_decision_policies
+  if enable_continuous_checkpointing:
+    return policies.ContinuousCheckpointingPolicy()
+  if enable_autocheckpoint:
+    return policies.AnySavePolicy(
+        [
+            policies.PreemptionCheckpointingPolicy(),
+            policies.FixedIntervalPolicy(save_interval_steps),
+        ]
+    )
+  return policies.FixedIntervalPolicy(interval=save_interval_steps)
+
+
+def build_preservation_policy(*, max_to_keep: int) -> ocp.training.preservation_policies.PreservationPolicy:
+  """Builds the v1 PreservationPolicy (keep the latest N checkpoints).
+
+  Args:
+    max_to_keep: The maximum number of checkpoints to keep.
+
+  Returns:
+    A configured ``ocp.training.preservation_policies.PreservationPolicy``.
+  """
+  return ocp.training.preservation_policies.LatestN(max_to_keep)
+
+
+def build_context(
+    *,
+    use_ocdbt: bool = True,
+    use_zarr3: bool = True,
+    ocdbt_target_data_file_size_bytes: int | None = None,
+    checkpoint_storage_concurrent_gb: int | None = None,
+    enable_continuous_checkpointing: bool = False,
+    todelete_full_path: str | None = None,
+    todelete_subdir: str | None = None,
+    enable_single_replica_ckpt_restoring: bool = False,
+    replica_axis_index: int = 0,
+    colocated_python_checkpointing: bool = False,
+    partial_load: bool = False,
+    checkpoint_layout: ocp.options.CheckpointLayout | None = None,
+) -> ocp.Context:
+  """Builds an Orbax v1 ``Context`` from MaxText checkpoint flags.
+
+  The returned Context is unfrozen (its options are mutable until it is entered
+  via ``with ctx:``); callers pass it to ``ocp_v1.training.Checkpointer``, which
+  applies it to every save/load.
+
+  Args:
+    use_ocdbt: Use OCDBT storage format.
+    use_zarr3: Use Zarr3 storage format.
+    ocdbt_target_data_file_size_bytes: Target OCDBT data-file size; also used as
+      the per-array ``chunk_byte_size`` (matching the v0 ``SaveArgs`` value).
+    checkpoint_storage_concurrent_gb: Concurrent IO budget in GB; applied to
+      both write and read as a byte limit (v0 used one value for both).
+    enable_continuous_checkpointing: If true, set a 60-minute async timeout.
+    todelete_full_path: GCS soft-delete path.
+    todelete_subdir: Subdirectory renaming hook for deletions.
+    enable_single_replica_ckpt_restoring: Restore on one replica and broadcast
+      to the rest (replaces the v0 ``SingleReplicaArrayHandler``).
+    replica_axis_index: Mesh axis separating replicas for load-and-broadcast.
+    colocated_python_checkpointing: Use Pathways colocated-python checkpointing.
+    partial_load: Restore only the keys present in the abstract tree (the v1
+      equivalent of v0 ``partial_restore=True``).
+    checkpoint_layout: On-disk layout (``ORBAX`` or ``SAFETENSORS``) for
+      loading.
+
+  Returns:
+    A configured, unfrozen ``ocp_v1.Context``.
+  """
+  ctx = ocp.Context()
+
+  # Array storage format + file sizing.
+  ctx.array.saving.use_ocdbt = use_ocdbt
+  ctx.array.saving.use_zarr3 = use_zarr3
+  if ocdbt_target_data_file_size_bytes is not None:
+    ctx.array.saving.ocdbt_target_data_file_size = ocdbt_target_data_file_size_bytes
+    ctx.array.saving.storage_options.chunk_byte_size = ocdbt_target_data_file_size_bytes
+
+  # Concurrent IO budget: v0 GB -> v1 bytes, applied to both directions.
+  if checkpoint_storage_concurrent_gb is not None:
+    concurrent_bytes = checkpoint_storage_concurrent_gb * _BYTES_PER_GB
+    ctx.memory.write_concurrent_bytes = concurrent_bytes
+    ctx.memory.read_concurrent_bytes = concurrent_bytes
+
+  if enable_continuous_checkpointing:
+    ctx.asynchronous.timeout_secs = int(datetime.timedelta(minutes=60).total_seconds())
+
+  if todelete_full_path is not None:
+    ctx.deletion.gcs_deletion_options.todelete_full_path = todelete_full_path
+
+  if todelete_subdir is not None:
+    raise ValueError("Renaming to subdirectory before deleting (todelete_subdir) is now unsupported by Orbax v1.")
+
+  # Single-replica restore (load on one replica, broadcast to the others).
+  if enable_single_replica_ckpt_restoring:
+    ctx.array.loading.use_load_and_broadcast = True
+    ctx.array.loading.load_and_broadcast_options.replica_axis_index = replica_axis_index
+    ctx.array.loading.load_and_broadcast_options.broadcast_memory_limit_bytes = (
+        _SINGLE_REPLICA_BROADCAST_MEMORY_LIMIT_BYTES
+    )
+
+  if colocated_python_checkpointing:
+    ctx.pathways.checkpointing_impl = ocp_pathways.CheckpointingImpl.from_options(
+        use_colocated_python=True,
+    )
+  else:
+    # v0 only used Pathways handlers when explicitly registered,
+    # and the persistence handler rejects non-NamedSharding arrays and the
+    # OCDBT/zarr3 layout MaxText writes. NO_DISPATCHER restores the standard
+    # controller-side ArrayHandler.
+    ctx.pathways.checkpointing_impl = ocp_pathways.CheckpointingImpl.NO_DISPATCHER
+
+  if partial_load:
+    ctx.pytree.loading.partial_load = True
+
+  if checkpoint_layout is not None:
+    ctx.checkpoint_layout = checkpoint_layout
+
+  return ctx