[compile] Save aot compile artifacts atomically. (#35117)
Signed-off-by: zhxchen17 <zhxchen17@fb.com>
This commit is contained in:
@@ -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(
|
||||
|
||||
Reference in New Issue
Block a user