optimize ut

intelligent-machine-learning · Oct 18, 2024 · 379c51a · 379c51a
1 parent b218326
commit 379c51a
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/dlrover/trainer/tests/torch/fsdp_ckpt_test.py b/dlrover/trainer/tests/torch/fsdp_ckpt_test.py
@@ -242,11 +242,15 @@ def test_shared_memory_writer(self):
  for _, item in files:
  write_items.append(item)
  writer = _write_state_dict_to_shm(self.shm, files, state_dict)
+ writer.reset()
+ self.assertTrue(writer.validate_checkpoint_id(None))
  self.assertTrue("dcp_metadata" in writer.metadata)
  self.assertTrue("no_shard_data" in writer.metadata)
 
  writer.shm_handler.metadata.set(writer.metadata)
  reader = SharedMemoryReader(writer.shm_handler)
+ reader.reset()
+ self.assertTrue(reader.validate_checkpoint_id(None))
  dcp_metadata = reader.read_metadata()
  self.assertTrue(_OPTIMIZER_KEY in dcp_metadata.state_dict_metadata)
  self.assertTrue(_OPTIMIZER_KEY in reader.no_shard_data)
@@ -295,6 +299,8 @@ def test_file_reader(self):
  f.write(writer.shm_handler.shared_memory.buf)
 
  reader = FileReader(tmpdir)
+ reader.reset()
+ self.assertTrue(reader.validate_checkpoint_id(None))
  metadata = reader.read_metadata()
  reader.set_up_storage_reader(metadata, True)