Fix Orbax error (#548)

This is addressing the following exception that can sometimes happen during checkpointing:
```
FileExistsError: [Errno 17] File exists: '..../10000.orbax-checkpoint-tmp-4'
```

Author: uzhilinsky

scripts/train_test.py

@@ -16,7 +16,7 @@ def test_train(tmp_path: pathlib.Path, config_name: str):
     config = dataclasses.replace(
         _config._CONFIGS_DICT[config_name],  # noqa: SLF001
         batch_size=2,
-        checkpoint_base_dir=tmp_path / "checkpoint",
+        checkpoint_base_dir=str(tmp_path / "checkpoint"),
         exp_name="test",
         overwrite=False,
         resume=False,

src/openpi/training/checkpoints.py

@@ -1,3 +1,6 @@
+from __future__ import annotations
+
+import asyncio
 import concurrent.futures as futures
 import dataclasses
 import logging
@@ -6,6 +9,7 @@
 from etils import epath
 import jax
 import orbax.checkpoint as ocp
+import orbax.checkpoint.future as future
 
 from openpi.shared import array_typing as at
 import openpi.shared.normalize as _normalize
@@ -117,18 +121,12 @@ def __call__(self, directory: epath.Path) -> None: ...
 class CallbackHandler(ocp.AsyncCheckpointHandler):
     """A CheckpointHandler for calling an arbitrary function asynchronously. Only for saving, not for restoring."""
 
-    def __init__(self):
-        self._executor = futures.ThreadPoolExecutor(max_workers=1)
-
-    def close(self):
-        self._executor.shutdown()
-
-    def save(self, directory: epath.Path, args: "CallbackSave"):
+    def save(self, directory: epath.Path, args: CallbackSave):
         if jax.process_index() == 0:
             args.callback(directory)
 
-    async def async_save(self, directory: epath.Path, args: "CallbackSave") -> list[futures.Future]:
-        return [self._executor.submit(self.save, directory, args)]
+    async def async_save(self, directory: epath.Path, args: CallbackSave) -> list[futures.Future]:
+        return [future.CommitFutureAwaitingContractedSignals(asyncio.to_thread(self.save, directory, args))]
 
     def restore(self, *args, **kwargs):
         raise NotImplementedError("CallbackHandler does not support restore")