Internal

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
 

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)

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 (

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

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.