Move Orbax path package under _src.

PiperOrigin-RevId: 693839844

checkpoint/CHANGELOG.md

@@ -14,6 +14,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Move `type_handlers` to `_src/serialization`
 - Add notes to Barrier error `XlaRuntimeError(DEADLINE_EXCEEDED)` with
  actionable info.
+- Make `NameFormat.find_all` impls concurrent.
+- Move `path` package under `_src` package.
 
 ## [0.8.0] - 2024-10-29
 
@@ -37,7 +39,6 @@ exported in the same way.
 - De-duplicate `get_ts_context` usages and move to ts_utils.
 - Move `logging` to `_src`.
 - Move `metadata` to `_src`.
-- Make `NameFormat.find_all` impls concurrent.
 
 ## [0.7.0] - 2024-10-07
 

checkpoint/orbax/__init__.py

@@ -37,7 +37,7 @@
 from orbax.checkpoint import version
 # TODO(cpgaffney): Import the public multihost API.
 from orbax.checkpoint._src.multihost import multihost
-from orbax.checkpoint.path import step
+from orbax.checkpoint._src.path import step
 
 from orbax.checkpoint.future import Future
 

checkpoint/orbax/checkpoint/__init__.py

@@ -37,7 +37,7 @@
 from orbax.checkpoint import version
 # TODO(cpgaffney): Import the public multihost API.
 from orbax.checkpoint._src.multihost import multihost
-from orbax.checkpoint.path import step
+from orbax.checkpoint._src.path import step
 
 from orbax.checkpoint.future import Future
 

checkpoint/orbax/checkpoint/_src/handlers/base_pytree_checkpoint_handler.py

@@ -43,11 +43,11 @@
 from orbax.checkpoint._src.handlers import async_checkpoint_handler
 from orbax.checkpoint._src.metadata import tree as tree_metadata
 from orbax.checkpoint._src.multihost import multihost
+from orbax.checkpoint._src.path import format_utils
 from orbax.checkpoint._src.serialization import serialization
 from orbax.checkpoint._src.serialization import tensorstore_utils as ts_utils
 from orbax.checkpoint._src.serialization import type_handlers
 from orbax.checkpoint._src.tree import utils as tree_utils
-from orbax.checkpoint.path import format_utils
 import tensorstore as ts
 
 

checkpoint/orbax/checkpoint/_src/handlers/composite_checkpoint_handler.py

@@ -63,8 +63,8 @@
 from orbax.checkpoint._src.handlers import checkpoint_handler
 from orbax.checkpoint._src.handlers import handler_registration
 from orbax.checkpoint._src.handlers import proto_checkpoint_handler
-from orbax.checkpoint.path import atomicity
-from orbax.checkpoint.path import utils as path_utils
+from orbax.checkpoint._src.path import atomicity
+from orbax.checkpoint._src.path import locking
 
 CheckpointArgs = checkpoint_args.CheckpointArgs
 Future = future.Future
@@ -721,7 +721,7 @@ def _existing_items(self, directory: epath.Path) -> List[str]:
     return [
         p.name
         for p in directory.iterdir()
-        if p.is_dir() and p != path_utils.lockdir(directory)
+        if p.is_dir() and p != locking.lockdir(directory)
     ]
 
   def restore(
@@ -883,7 +883,9 @@ def close(self):
 )
 class CompositeArgs(Composite, CheckpointArgs):
   """Args for wrapping multiple checkpoint items together."""
+
   ...
 
+
 # Returned object of CompositeCheckpointHandler is an alias of CompositeArgs.
 CompositeResults = CompositeArgs

checkpoint/orbax/checkpoint/_src/handlers/composite_checkpoint_handler_test.py

@@ -27,7 +27,7 @@
 from orbax.checkpoint._src.handlers import standard_checkpoint_handler
 from orbax.checkpoint._src.metadata import value as value_metadata
 from orbax.checkpoint._src.multihost import multihost
-from orbax.checkpoint.path import step
+from orbax.checkpoint._src.path import step
 
 CompositeArgs = composite_checkpoint_handler.CompositeArgs
 JsonCheckpointHandler = json_checkpoint_handler.JsonCheckpointHandler

checkpoint/orbax/checkpoint/path/async_utils.py → checkpoint/orbax/checkpoint/_src/path/async_utils.py

@@ -18,7 +18,7 @@
 
 from etils import epath
 from orbax.checkpoint._src import asyncio_utils
-from orbax.checkpoint.path import step as step_lib
+from orbax.checkpoint._src.path import step as step_lib
 
 
 # TODO(b/360190539): This functionality should be provided by either an external

checkpoint/orbax/checkpoint/path/atomicity.py → checkpoint/orbax/checkpoint/_src/path/atomicity.py

@@ -64,9 +64,9 @@
 from orbax.checkpoint._src.metadata import checkpoint as checkpoint_metadata
 from orbax.checkpoint._src.multihost import counters
 from orbax.checkpoint._src.multihost import multihost
+from orbax.checkpoint._src.path import async_utils
+from orbax.checkpoint._src.path import step as step_lib
 from orbax.checkpoint._src.path import utils
-from orbax.checkpoint.path import async_utils
-from orbax.checkpoint.path import step as step_lib
 
 
 TMP_DIR_SUFFIX = step_lib.TMP_DIR_SUFFIX

checkpoint/orbax/checkpoint/path/atomicity_test.py → checkpoint/orbax/checkpoint/_src/path/atomicity_test.py

@@ -18,8 +18,8 @@
 from etils import epath
 from orbax.checkpoint import test_utils
 from orbax.checkpoint._src.multihost import multihost
-from orbax.checkpoint.path import atomicity
-from orbax.checkpoint.path import step as step_lib
+from orbax.checkpoint._src.path import atomicity
+from orbax.checkpoint._src.path import step as step_lib
 
 AtomicRenameTemporaryPath = atomicity.AtomicRenameTemporaryPath
 CommitFileTemporaryPath = atomicity.CommitFileTemporaryPath

checkpoint/orbax/checkpoint/path/deleter.py → checkpoint/orbax/checkpoint/_src/path/deleter.py

@@ -22,7 +22,7 @@
 from etils import epath
 import jax
 from orbax.checkpoint import utils
-from orbax.checkpoint.path import step as step_lib
+from orbax.checkpoint._src.path import step as step_lib
 
 _THREADED_DELETE_DURATION = (
     '/jax/orbax/checkpoint_manager/threaded_checkpoint_deleter/duration'

checkpoint/orbax/checkpoint/path/deleter_test.py → checkpoint/orbax/checkpoint/_src/path/deleter_test.py

@@ -17,8 +17,8 @@
 from absl.testing import absltest
 from absl.testing import parameterized
 from etils import epath
-from orbax.checkpoint.path import deleter as deleter_lib
-from orbax.checkpoint.path import step as step_lib
+from orbax.checkpoint._src.path import deleter as deleter_lib
+from orbax.checkpoint._src.path import step as step_lib
 
 
 class CheckpointDeleterTest(parameterized.TestCase):

checkpoint/orbax/checkpoint/path/format_utils_test.py → checkpoint/orbax/checkpoint/_src/path/format_utils_test.py

@@ -22,7 +22,7 @@
 from orbax.checkpoint._src.handlers import pytree_checkpoint_handler
 from orbax.checkpoint._src.handlers import standard_checkpoint_handler
 from orbax.checkpoint._src.metadata import checkpoint as checkpoint_metadata
-from orbax.checkpoint.path import format_utils
+from orbax.checkpoint._src.path import format_utils
 
 
 

checkpoint/orbax/checkpoint/path/utils.py → checkpoint/orbax/checkpoint/_src/path/locking.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Path utility functions for Orbax."""
+"""Manages locking of checkpoint step dirs."""
 
 # TODO(b/337137764): Add unit tests.
 # TODO(b/337137764): If needed, export the functions from
@@ -24,8 +24,8 @@
 
 from etils import epath
 from orbax.checkpoint._src import asyncio_utils
-from orbax.checkpoint.path import async_utils
-from orbax.checkpoint.path import step as step_lib
+from orbax.checkpoint._src.path import async_utils
+from orbax.checkpoint._src.path import step as step_lib
 
 
 _LOCK_ITEM_NAME = 'LOCKED'

checkpoint/orbax/checkpoint/path/step_test.py → checkpoint/orbax/checkpoint/_src/path/step_test.py

@@ -21,8 +21,8 @@
 from etils import epath
 from orbax.checkpoint import test_utils
 from orbax.checkpoint._src.metadata import checkpoint
-from orbax.checkpoint.path import atomicity
-from orbax.checkpoint.path import step as step_lib
+from orbax.checkpoint._src.path import atomicity
+from orbax.checkpoint._src.path import step as step_lib
 
 
 class StandardNameFormatTest(parameterized.TestCase):

checkpoint/orbax/checkpoint/_src/path/utils.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Utils for snapshotting."""
+"""Utils for path constructs."""
 
 import os
 import time

checkpoint/orbax/checkpoint/_src/serialization/type_handlers.py

@@ -42,10 +42,10 @@
 from orbax.checkpoint._src.metadata import value as value_metadata
 from orbax.checkpoint._src.multihost import multihost
 from orbax.checkpoint._src.multihost import multislice
+from orbax.checkpoint._src.path import async_utils
+from orbax.checkpoint._src.path import format_utils
 from orbax.checkpoint._src.serialization import serialization
 from orbax.checkpoint._src.serialization import tensorstore_utils as ts_utils
-from orbax.checkpoint.path import async_utils
-from orbax.checkpoint.path import format_utils
 import tensorstore as ts
 
 

checkpoint/orbax/checkpoint/async_checkpointer.py

@@ -31,8 +31,8 @@
 from orbax.checkpoint._src.metadata import checkpoint
 from orbax.checkpoint._src.multihost import counters
 from orbax.checkpoint._src.multihost import multihost
-from orbax.checkpoint.path import async_utils
-from orbax.checkpoint.path import atomicity
+from orbax.checkpoint._src.path import async_utils
+from orbax.checkpoint._src.path import atomicity
 
 
 BarrierSyncFn = multihost.BarrierSyncFn

checkpoint/orbax/checkpoint/checkpoint_manager.py

@@ -44,13 +44,13 @@
 from orbax.checkpoint._src.handlers import proto_checkpoint_handler
 from orbax.checkpoint._src.metadata import checkpoint
 from orbax.checkpoint._src.multihost import multihost
+from orbax.checkpoint._src.path import atomicity
+from orbax.checkpoint._src.path import deleter
+from orbax.checkpoint._src.path import step as step_lib
 from orbax.checkpoint._src.path import utils as path_utils
 from orbax.checkpoint.logging import abstract_logger
 from orbax.checkpoint.logging import standard_logger
 from orbax.checkpoint.logging import step_statistics
-from orbax.checkpoint.path import atomicity
-from orbax.checkpoint.path import deleter
-from orbax.checkpoint.path import step as step_lib
 from typing_extensions import Self  # for Python version < 3.11
 
 

checkpoint/orbax/checkpoint/checkpoint_utils.py

@@ -25,9 +25,9 @@
 from orbax.checkpoint import utils
 from orbax.checkpoint._src.metadata import value as value_metadata
 from orbax.checkpoint._src.multihost import multihost
+from orbax.checkpoint._src.path import step as step_lib
 from orbax.checkpoint._src.path.snapshot import snapshot as snapshot_lib
 from orbax.checkpoint._src.serialization import type_handlers
-from orbax.checkpoint.path import step as step_lib
 
 
 PyTree = Any

checkpoint/orbax/checkpoint/checkpoint_utils_test.py

@@ -27,7 +27,7 @@
 from orbax.checkpoint import utils
 from orbax.checkpoint._src.handlers import pytree_checkpoint_handler
 from orbax.checkpoint._src.metadata import value as value_metadata
-from orbax.checkpoint.path import step as step_lib
+from orbax.checkpoint._src.path import step as step_lib
 
 
 RestoreArgs = pytree_checkpoint_handler.RestoreArgs

checkpoint/orbax/checkpoint/checkpointer.py

@@ -30,7 +30,7 @@
 from orbax.checkpoint._src.handlers import composite_checkpoint_handler
 from orbax.checkpoint._src.metadata import checkpoint
 from orbax.checkpoint._src.multihost import multihost
-from orbax.checkpoint.path import atomicity
+from orbax.checkpoint._src.path import atomicity
 from typing_extensions import Self  # for Python version < 3.11
 
 

checkpoint/orbax/checkpoint/experimental/emergency/checkpoint_manager.py

@@ -48,12 +48,12 @@
 from orbax.checkpoint._src.multihost import counters
 from orbax.checkpoint._src.multihost import multihost
 from orbax.checkpoint._src.multihost import multislice
+from orbax.checkpoint._src.path import step as step_lib
 from orbax.checkpoint._src.serialization import type_handlers
 from orbax.checkpoint.experimental.emergency import multihost as emergency_multihost
 from orbax.checkpoint.logging import abstract_logger
 from orbax.checkpoint.logging import standard_logger
 from orbax.checkpoint.logging import step_statistics
-from orbax.checkpoint.path import step as step_lib
 from typing_extensions import Self  # for Python version < 3.11
 
 

checkpoint/orbax/checkpoint/path.py

@@ -0,0 +1,23 @@
+# Copyright 2024 The Orbax Authors.
+#
+# 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
+#
+#     http://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.
+
+"""Defines exported symbols from orbax.checkpoint.path package."""
+
+# pylint: disable=unused-import
+
+from orbax.checkpoint._src.path import async_utils
+from orbax.checkpoint._src.path import atomicity
+from orbax.checkpoint._src.path import deleter
+from orbax.checkpoint._src.path import format_utils
+from orbax.checkpoint._src.path import step

checkpoint/orbax/checkpoint/standard_checkpointer.py

@@ -20,7 +20,7 @@
 from orbax.checkpoint import options as options_lib
 from orbax.checkpoint._src.handlers import standard_checkpoint_handler
 from orbax.checkpoint._src.metadata import checkpoint
-from orbax.checkpoint.path import atomicity
+from orbax.checkpoint._src.path import atomicity
 
 
 StandardCheckpointHandler = (

checkpoint/orbax/checkpoint/test_utils.py

@@ -42,12 +42,12 @@
 from orbax.checkpoint._src.multihost import counters
 from orbax.checkpoint._src.multihost import multihost
 from orbax.checkpoint._src.multihost import multislice
+from orbax.checkpoint._src.path import atomicity
+from orbax.checkpoint._src.path import step as step_lib
 from orbax.checkpoint._src.serialization import serialization
 from orbax.checkpoint._src.serialization import tensorstore_utils as ts_utils
 from orbax.checkpoint._src.serialization import type_handlers
 from orbax.checkpoint._src.tree import utils as tree_utils
-from orbax.checkpoint.path import atomicity
-from orbax.checkpoint.path import step as step_lib
 
 
 class MuNu(NamedTuple):

checkpoint/orbax/checkpoint/utils.py

@@ -27,10 +27,10 @@
 import jax
 import numpy as np
 from orbax.checkpoint._src.multihost import multihost
+from orbax.checkpoint._src.path import async_utils
+from orbax.checkpoint._src.path import locking
+from orbax.checkpoint._src.path import step as step_lib
 from orbax.checkpoint._src.tree import utils as tree_utils
-from orbax.checkpoint.path import async_utils
-from orbax.checkpoint.path import step as step_lib
-from orbax.checkpoint.path import utils as path_utils
 
 
 TMP_DIR_SUFFIX = step_lib.TMP_DIR_SUFFIX
@@ -54,8 +54,8 @@
 async_makedirs = async_utils.async_makedirs
 async_write_bytes = async_utils.async_write_bytes
 async_exists = async_utils.async_exists
-lockdir = path_utils.lockdir
-is_locked = path_utils.is_locked
+lockdir = locking.lockdir
+is_locked = locking.is_locked
 
 
 is_gcs_path = step_lib.is_gcs_path
@@ -115,9 +115,7 @@ def name_from_leaf_placeholder(placeholder: str) -> str:
 
 def all_leaves_are_placeholders(tree: PyTree) -> bool:
   """Determines if all leaves in `tree` are placeholders."""
-  return all(
-      leaf_is_placeholder(leaf) for leaf in jax.tree.leaves(tree)
-  )
+  return all(leaf_is_placeholder(leaf) for leaf in jax.tree.leaves(tree))
 
 
 def pytree_structure(directory: epath.PathLike) -> PyTree: