Upgrade orbax checkpointer to 0.11.15 (#1255)

* Support Orbax 0.11.14 and newer

* upgrade orbax to 0.11.14

* move orbax dependency to core group

* Revert "move orbax dependency to core group"

This reverts commit ef2566a.

* upgrade orbax to 0.11.15

Author: samos123

axlearn/common/checkpointer_orbax_emergency.py

@@ -751,7 +751,7 @@ def save(
         # including step time in total blocking time.
         start_t = time.perf_counter()
         self._get_tensor_manager(state_with_tensors).save(
-            step=step, args=ocp.args.PyTreeSave(item=state_with_tensors)
+            step=step, args=ocp.args.Composite(state=ocp.args.PyTreeSave(item=state_with_tensors))
         )
         time_diff = time.perf_counter() - start_t
         if self._composite_save_policy(step=step, evaler_summaries=self._eval_summaries):
@@ -808,7 +808,9 @@ def restore(
 
         restored_state_with_tensors = tensor_manager.restore(
             step=step,
-            args=ocp.args.PyTreeRestore(item=self._get_abstract_state(state_with_tensors)),
+            args=ocp.args.Composite(
+                state=ocp.args.PyTreeRestore(item=self._get_abstract_state(state_with_tensors))
+            ),
         )
         # Merge non-tensor and tensor states by replacing leaves of the non-tensor Pytree with the
         # not-None leaves of the tensor Pytree.
@@ -826,7 +828,8 @@ def wait_until_finished(self):
         self._non_tensor_manager.wait_until_finished()
         self._tensor_manager.wait_until_finished()
 
-    def stop(self):
+    def stop(self, *, has_exception: bool = False):
         """See `BaseCheckpointer.stop` for details."""
-        self._non_tensor_manager.stop()
-        self._tensor_manager.close()
+        self._non_tensor_manager.stop(has_exception=has_exception)
+        if self._tensor_manager:
+            self._tensor_manager.close()

pyproject.toml

@@ -154,7 +154,7 @@ mmau = [
 # Orbax checkpointing.
 orbax = [
     "humanize==4.10.0",
-    "orbax-checkpoint==0.11.1",
+    "orbax-checkpoint==0.11.15",
 ]
 # Audio dependencies.
 audio = [