Add merge / finalize step when using OCDBT driver. Files will be firs…

…t written to per-process subdirectories, which are later copied by reference to the main directory before the checkpoint is finalized.

PiperOrigin-RevId: 572392273

tests/checkpoints_test.py

@@ -195,9 +195,14 @@ def test_overwrite_checkpoints(self, use_orbax):
 
     os.chdir(os.path.dirname(tmp_dir))
     rel_tmp_dir = './' + os.path.basename(tmp_dir)
-    checkpoints.save_checkpoint(rel_tmp_dir, test_object, 3, keep=1)
-    new_object = checkpoints.restore_checkpoint(rel_tmp_dir, test_object0)
-    check_eq(new_object, test_object)
+    if use_orbax:
+      with self.assertRaises(ValueError):
+        checkpoints.save_checkpoint(rel_tmp_dir, test_object, 3, keep=1)
+    else:
+      checkpoints.save_checkpoint(rel_tmp_dir, test_object, 3, keep=1)
+      new_object = checkpoints.restore_checkpoint(rel_tmp_dir, test_object0)
+      check_eq(new_object, test_object)
+
     non_norm_dir_path = tmp_dir + '//'
     checkpoints.save_checkpoint(non_norm_dir_path, test_object, 4, keep=1)
     new_object = checkpoints.restore_checkpoint(non_norm_dir_path, test_object0)