Stop writing msgpack file for new checkpoints and update empty nodes …

…handling so that it no longer depends on this file.

PiperOrigin-RevId: 649179323

flax/training/checkpoints.py

@@ -31,6 +31,7 @@
 )
 from collections.abc import Callable, Iterable
 
+from etils import epath
 import jax
 import orbax.checkpoint as ocp
 from absl import logging
@@ -76,7 +77,7 @@
 
 # Orbax main checkpoint file name.
 ORBAX_CKPT_FILENAME = 'checkpoint'
-ORBAX_MANIFEST_OCDBT = 'manifest.ocdbt'
+ORBAX_METADATA_FILENAME = '_METADATA'
 
 PyTree = Any
 
@@ -123,7 +124,8 @@ def _safe_remove(path: str):
 
 def _is_orbax_checkpoint(path: str) -> bool:
   return io.exists(os.path.join(path, ORBAX_CKPT_FILENAME)) or io.exists(
-    os.path.join(path, ORBAX_MANIFEST_OCDBT)
+      os.path.join(path, ORBAX_METADATA_FILENAME)
+      or ocp.type_handlers.is_ocdbt_checkpoint(epath.Path(path))
   )