Unverified Commit c77f3e12 authored by Zhengxu Chen's avatar Zhengxu Chen Committed by GitHub
Browse files

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


Signed-off-by: default avatarzhxchen17 <zhxchen17@fb.com>
parent 012dee92
......@@ -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(
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment