From c0087535d7f2e5bfcbf2a7be6825b9f5055a54c6 Mon Sep 17 00:00:00 2001 From: Colin Gaffney Date: Sat, 28 Oct 2023 09:44:11 -0700 Subject: [PATCH] Add merge / finalize step when using OCDBT driver. Files will be first written to per-process subdirectories, which are later copied by reference to the main directory before the checkpoint is finalized. PiperOrigin-RevId: 577475560 --- tests/checkpoints_test.py | 5 ----- 1 file changed, 5 deletions(-) diff --git a/tests/checkpoints_test.py b/tests/checkpoints_test.py index b49d715561..013e28729c 100644 --- a/tests/checkpoints_test.py +++ b/tests/checkpoints_test.py @@ -193,11 +193,6 @@ def test_overwrite_checkpoints(self, use_orbax): new_object = checkpoints.restore_checkpoint(tmp_dir, test_object0) check_eq(new_object, test_object) - 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) 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)