From d92d15116f984713ad1ee5cc61ba339295f91c7a Mon Sep 17 00:00:00 2001
From: Orbax Authors <no-reply@google.com>
Date: Tue, 9 Jun 2026 05:40:37 -0700
Subject: [PATCH] Internal

PiperOrigin-RevId: 929137076
---
 .../_src/checkpointers/checkpointer.py        |  4 ++-
 .../handlers/composite_checkpoint_handler.py  |  6 +---
 .../orbax/checkpoint/_src/path/atomicity.py   |  4 ++-
 .../experimental/v1/_src/saving/path_utils.py |  5 +--
 checkpoint/orbax/checkpoint/options.py        | 32 +++++++++++++++++++
 5 files changed, 40 insertions(+), 11 deletions(-)

diff --git a/checkpoint/orbax/checkpoint/_src/checkpointers/checkpointer.py b/checkpoint/orbax/checkpoint/_src/checkpointers/checkpointer.py
index b4a168585..42d4c9df9 100644
--- a/checkpoint/orbax/checkpoint/_src/checkpointers/checkpointer.py
+++ b/checkpoint/orbax/checkpoint/_src/checkpointers/checkpointer.py
@@ -14,6 +14,7 @@
 
 """Synchronous Checkpointer implementation."""
 
+import dataclasses
 import time
 from typing import Any, Iterable, Optional, Type
 
@@ -155,10 +156,11 @@ def get_temporary_path(
         self._temporary_path_class
         or atomicity_defaults.get_default_temporary_path_class(directory)
     )
+    file_options = self._file_options
     tmpdir = temporary_path_class.from_final(
         directory,
         checkpoint_metadata_store=self._metadata_store,
-        file_options=self._file_options,
+        file_options=file_options,
     )
     return tmpdir
 
diff --git a/checkpoint/orbax/checkpoint/_src/handlers/composite_checkpoint_handler.py b/checkpoint/orbax/checkpoint/_src/handlers/composite_checkpoint_handler.py
index 3c382d800..598371b5a 100644
--- a/checkpoint/orbax/checkpoint/_src/handlers/composite_checkpoint_handler.py
+++ b/checkpoint/orbax/checkpoint/_src/handlers/composite_checkpoint_handler.py
@@ -72,7 +72,6 @@
 from orbax.checkpoint._src.path import atomicity_defaults
 from orbax.checkpoint._src.path import atomicity_types
 
-
 CheckpointArgs = checkpoint_args.CheckpointArgs
 Future = future.Future
 CheckpointArgs = checkpoint_args.CheckpointArgs
@@ -1044,10 +1043,7 @@ def finalize(self, directory: epath.Path):
         # Not an error, as some items may not have been saved.
         continue
       handler.finalize(tmp_dir.get())
-      asyncio_utils.run_sync(
-          tmp_dir.finalize(
-          )
-      )
+      asyncio_utils.run_sync(tmp_dir.finalize())
 
       # Remove the temporary path once it has been finalized.
       self._current_temporary_paths.pop(item_name)
diff --git a/checkpoint/orbax/checkpoint/_src/path/atomicity.py b/checkpoint/orbax/checkpoint/_src/path/atomicity.py
index 35dc921c8..1a9d0243a 100644
--- a/checkpoint/orbax/checkpoint/_src/path/atomicity.py
+++ b/checkpoint/orbax/checkpoint/_src/path/atomicity.py
@@ -549,7 +549,9 @@ async def finalize(
     if self._snapshot is not None:
       await self._snapshot.replace_source()
     else:
-      await async_path.rename(self._tmp_path, self._final_path)  # pyrefly: ignore[bad-argument-type]
+      rename_src = self._tmp_path
+      rename_dst = self._final_path
+      await async_path.rename(rename_src, rename_dst)
 
   def __repr__(self) -> str:
     return (
diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/saving/path_utils.py b/checkpoint/orbax/checkpoint/experimental/v1/_src/saving/path_utils.py
index e26f5f5fe..36489c938 100644
--- a/checkpoint/orbax/checkpoint/experimental/v1/_src/saving/path_utils.py
+++ b/checkpoint/orbax/checkpoint/experimental/v1/_src/saving/path_utils.py
@@ -25,7 +25,6 @@
 from orbax.checkpoint.experimental.v1._src.synchronization import synchronization
 
 
-
 def get_temporary_path(
     path: path_types.Path,
     *,
@@ -42,9 +41,7 @@ def get_temporary_path(
   Returns:
     A TemporaryPath for the given path.
   """
-  temporary_path_cls = atomicity_defaults.get_default_temporary_path_class(
-      path
-  )
+  temporary_path_cls = atomicity_defaults.get_default_temporary_path_class(path)
   tmpdir = temporary_path_cls.from_final(
       path,
       # Ensure metadata store is NOT passed, to prevent separate metadata
diff --git a/checkpoint/orbax/checkpoint/options.py b/checkpoint/orbax/checkpoint/options.py
index a3f636319..b9a6fc372 100644
--- a/checkpoint/orbax/checkpoint/options.py
+++ b/checkpoint/orbax/checkpoint/options.py
@@ -74,6 +74,38 @@ class FileOptions:
   path_permission_mode: int | None = None
 
 
+# TODO: b/425293362 - Remove this once the migration to TemporaryPath is
+# complete.
+def ensure_cns_file_options(
+    file_options: FileOptions,
+) -> FileOptions:
+  """Ensures CNS file options exist, creating them if they don't.
+
+  Args:
+    file_options: Existing file options.
+
+  Returns:
+    File options with CNS file options populated.
+  """
+  if file_options.cns_file_options:
+    return file_options
+  else:
+    # For backward compatibility, file_options.cns_file_options is created
+    # from deprecated fields in file_options.
+    cns_file_options = ColossusFileOptions(
+        data_governance_annotations=file_options.data_governance_annotations,
+        cns2_storage_options=file_options.cns2_storage_options,
+    )
+    return dataclasses.replace(
+        file_options,
+        cns_file_options=cns_file_options,
+        cns2_storage_options=None,
+        data_governance_annotations=None,
+    )
+
+
+
+
 @dataclasses.dataclass
 class MemoryLimitOptions:
   """Options for configuring memory limits for save.