Internal change

PiperOrigin-RevId: 459553867

Author: cpgaffney1

orbax/checkpoint/__init__.py

@@ -15,19 +15,15 @@
 """Defines exported symbols for the `orbax` package."""
 
 from .abstract_checkpoint_manager import AbstractCheckpointManager
+from .checkpoint_handler import CheckpointHandler
 from .checkpoint_manager import CheckpointManager
 from .checkpoint_manager import CheckpointManagerOptions
-from .checkpointer import Checkpointer
-from .dataset_checkpointer import DatasetCheckpointer
-from .json_checkpointer import JsonCheckpointer
+from .dataset_checkpoint_handler import DatasetCheckpointHandler
+from .json_checkpoint_handler import JsonCheckpointHandler
 from orbax.checkpoint import lazy_array
-from .pytree_checkpointer import PyTreeCheckpointer
-from .pytree_checkpointer import RestoreArgs
-from .pytree_checkpointer import SaveArgs
+from .pytree_checkpoint_handler import PyTreeCheckpointHandler
+from .pytree_checkpoint_handler import RestoreArgs
+from .pytree_checkpoint_handler import SaveArgs
 from .transform_utils import apply_transformations
 from .transform_utils import Transform
 from .utils import checkpoints_iterator
-
-# TODO(cpgaffney) Remove when handler classes are fully rolled out.
-CheckpointHandler = Checkpointer
-PyTreeCheckpointHandler = PyTreeCheckpointer

orbax/checkpoint/abstract_checkpoint_manager.py

@@ -22,20 +22,22 @@
 
 
 class AbstractCheckpointManager(abc.ABC):
-  """An interface that manages multiple Checkpointer classes.
+  """An interface that manages multiple CheckpointHandler classes.
 
   CheckpointManager coordinates save/restore operations across multiple
-  Checkpointer classes, and also provides useful methods describing checkpoint
+  CheckpointHandler classes, and also provides useful methods describing
+  checkpoint
   states.
 
   For example, CheckpointManager may be responsible for managing a parameter
   state in the form of a PyTree and a dataset iterator state in the form of
   tf.data.Iterator.
 
-  Each item should be handled by a separate Checkpointer.
+  Each item should be handled by a separate CheckpointHandler.
 
-  For instance, item "a" is handled by Checkpointer A, while item "b" is handled
-  by Checkpointer B.
+  For instance, item "a" is handled by CheckpointHandler A, while item "b" is
+  handled
+  by CheckpointHandler B.
   """
 
   @abc.abstractmethod
@@ -54,27 +56,28 @@ def save(
       ...
     }
     Each of these values is a saveable item that should be written with a
-    specific Checkpointer.
+    specific CheckpointHandler.
 
     Similarly, save_kwargs takes the form:
     {
       'params': {
-        <kwargs for PyTreeCheckpointer.save>
+        <kwargs for PyTreeCheckpointHandler.save>
       },
       'dataset': {
-        <kwargs for DatasetCheckpointer.save>
+        <kwargs for DatasetCheckpointHandler.save>
       }
       ...
     }
     The dict of kwargs for each key in save_kwargs is provided as extra
-    arguments to the save method of the corresponding Checkpointer.
+    arguments to the save method of the corresponding CheckpointHandler.
 
     Args:
       step: current step, int
       items: a savable object, or a dictionary of object name to savable object.
-      save_kwargs: save kwargs for a single Checkpointer, or a dictionary of
-        object name to kwargs needed by the Checkpointer implementation to save
-        the object.
+      save_kwargs: save kwargs for a single CheckpointHandler, or a dictionary
+        of object name to kwargs needed by the CheckpointHandler implementation
+        to save the object.
+
     Returns:
       bool indicating whether save was performed or not.
     """
@@ -97,29 +100,30 @@ def restore(
       ...
     }
     Each of these values is a restoreable item that should be read with a
-    specific Checkpointer. Implementations should support items=None, and the
+    specific CheckpointHandler. Implementations should support items=None, and
+    the
     ability to restore an item which is not provided in this dict.
 
     Similarly, restore_kwargs takes the form:
     {
       'params': {
-        <kwargs for PyTreeCheckpointer.restore>
+        <kwargs for PyTreeCheckpointHandler.restore>
       },
       'dataset': {
-        <kwargs for DatasetCheckpointer.restore>
+        <kwargs for DatasetCheckpointHandler.restore>
       }
       ...
     }
     The dict of kwargs for each key in restore_kwargs is provided as extra
-    arguments to the restore method of the corresponding Checkpointer.
+    arguments to the restore method of the corresponding CheckpointHandler.
 
     Args:
       step: current step, int
       items: a restoreable object, or a dictionary of object name to restoreable
         object.
-      restore_kwargs: restore kwargs for a single Checkpointer, or a dictionary
-        of object name to kwargs needed by the Checkpointer implementation to
-        restore the object.
+      restore_kwargs: restore kwargs for a single CheckpointHandler, or a
+        dictionary of object name to kwargs needed by the CheckpointHandler
+        implementation to restore the object.
 
     Returns:
       A dictionary mapping name to restored object, or a single restored object.
@@ -128,13 +132,15 @@ def restore(
 
   @abc.abstractmethod
   def structure(self) -> Union[Any, Mapping[str, Any]]:
-    """For all Checkpointers, returns the saved structure.
+    """For all CheckpointHandlers, returns the saved structure.
 
-    Calls the `structure` method for each Checkpointer and returns a mapping of
+    Calls the `structure` method for each CheckpointHandler and returns a
+    mapping of
     each item name to the restored structure. If the manager only manages a
     single item, a single structure will be returned instead.
 
-    Note that any items for which the corresponding Checkpointer does not have
+    Note that any items for which the corresponding CheckpointHandler does not
+    have
     an implemented `structure` method, these items will simply not be contained
     in the result.
 

orbax/checkpoint/async_checkpoint_handler.py

@@ -0,0 +1,39 @@
+# Copyright 2022 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.
+
+"""AsyncCheckpointHandler interface."""
+
+import abc
+import asyncio
+from typing import Any, Optional, Sequence
+
+
+class AsyncCheckpointHandler(abc.ABC):
+  """An interface providing async methods that can be used with CheckpointHandler."""
+
+  @abc.abstractmethod
+  async def async_save(self, directory: str, item: Any, *args,
+                       **kwargs) -> Optional[Sequence[asyncio.Future]]:
+    """Constructs a save operation.
+
+    Synchronously awaits a copy of the item, before returning commit futures
+    necessary to save the item.
+
+    Args:
+      directory: the directory to save to.
+      item: the item to be saved.
+      *args: additional arguments for save.
+      **kwargs: additional arguments for save.
+    """
+    pass

orbax/checkpoint/checkpoint_handler.py

@@ -0,0 +1,69 @@
+# Copyright 2022 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.
+
+"""CheckpointHandler interface."""
+
+import abc
+from typing import Any, Optional
+
+
+class CheckpointHandler(abc.ABC):
+  """An interface providing save/restore methods used on a savable item.
+
+  Item may be a PyTree, Dataset, or any other supported object.
+  """
+
+  @abc.abstractmethod
+  def save(self, directory: str, item: Any, *args, **kwargs):
+    """Saves the provided item synchronously.
+
+    Args:
+      directory: the directory to save to.
+      item: the item to be saved.
+      *args: additional arguments for save.
+      **kwargs: additional arguments for save.
+    """
+    pass
+
+  @abc.abstractmethod
+  def restore(self,
+              directory: str,
+              item: Optional[Any] = None,
+              **kwargs) -> Any:
+    """Restores the provided item synchronously.
+
+    Args:
+      directory: the directory to restore from.
+      item: an item with the same structure as that to be restored.
+      **kwargs: additional arguments for restore.
+
+    Returns:
+      The restored item.
+    """
+    pass
+
+  @abc.abstractmethod
+  def structure(self, directory: str) -> Any:
+    """Returns the structure of the item.
+
+    This should represent the checkpointed item format without needing to fully
+    restore the entire item and all its values.
+
+    Args:
+      directory: the directory where the checkpoint is located.
+
+    Returns:
+      item structure
+    """
+    pass