From 2a91cb268938f0807d02f34a7f06a59d0b7cc1b0 Mon Sep 17 00:00:00 2001 From: Hao Zhang Date: Fri, 6 Mar 2026 02:34:50 +0800 Subject: [PATCH] feat(utility): ensure atomic checkpoint updates in RuntimeContext.save --- qmp/utility/context.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/qmp/utility/context.py b/qmp/utility/context.py index 5ab948e..9b90aea 100644 --- a/qmp/utility/context.py +++ b/qmp/utility/context.py @@ -112,11 +112,13 @@ def save(self, data: typing.Any, step: int) -> None: data_path = self.folder() / "data.pth" local_data_path = self.folder() / f"data.{step}.pth" torch.save(data, local_data_path) - data_path.unlink(missing_ok=True) if step % self.checkpoint_interval == 0: - data_path.symlink_to(f"data.{step}.pth") + temp_data_path = self.folder() / "data.pth.tmp" + temp_data_path.unlink(missing_ok=True) + temp_data_path.symlink_to(f"data.{step}.pth") + temp_data_path.replace(data_path) else: - local_data_path.rename(data_path) + local_data_path.replace(data_path) if self.max_relative_step is not None: self.max_absolute_step = step + self.max_relative_step - 1 self.max_relative_step = None