diff --git a/vllm/compilation/decorators.py b/vllm/compilation/decorators.py index f97467ad6..68be29cca 100644 --- a/vllm/compilation/decorators.py +++ b/vllm/compilation/decorators.py @@ -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(