[compile] Save aot compile artifacts atomically. (#35117)

Signed-off-by: zhxchen17 <zhxchen17@fb.com>
This commit is contained in:
Zhengxu Chen
2026-02-24 07:11:01 -05:00
committed by GitHub
parent 012dee9233
commit c77f3e1207

View File

@@ -575,7 +575,11 @@ def _support_torch_compile(
logger.info("saving AOT compiled function to %s", self._aot_compilation_path)
try:
os.makedirs(self._aot_cache_dir, exist_ok=True)
self.aot_compiled_fn.save_compiled_function(self._aot_compilation_path)
# File saving should be atomic, so we will save to a temporary location
# first. Should be upstreamed to PyTorch 2.12 as well.
tmp_file = f"{self._aot_compilation_path}.{os.getpid()}.tmp"
self.aot_compiled_fn.save_compiled_function(tmp_file)
os.replace(tmp_file, self._aot_compilation_path)
logger.info("saved AOT compiled function to %s", self._aot_compilation_path)
except Exception as e:
logger.warning(