saymrwulf/pytorch
1
0
Fork 0
mirror of https://github.com/saymrwulf/pytorch.git synced 2026-05-14 20:57:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Fix separate in process bisector cache, cleanup on exit (#143661)

Browse source 
Pull Request resolved: https://github.com/pytorch/pytorch/pull/143661
Approved by: https://github.com/ezyang
ghstack dependencies: #143657
This commit is contained in:
eellison 2024-12-20 09:52:50 -08:00 committed by PyTorch MergeBot
parent 1e246ef05b
commit d3e9133ab2
1 changed files with 19 additions and 6 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

25
torch/_inductor/compiler_bisector.py
View file

@ -1,9 +1,11 @@
import atexit
import collections
import dataclasses
import functools
import os
import shutil
import sys
import tempfile
from dataclasses import dataclass, field
from typing import Callable, Dict, List, Optional, Tuple
@ -114,9 +116,11 @@ class CompilerBisector:
bisection_enabled: bool = False
in_process_cache: Optional[str] = None
@classmethod
def get_dir(cls) -> str:
return f"{cache_dir()}/{SUBDIR_NAME}"
return f"{cache_dir() if not cls.in_process_cache else cls.in_process_cache}/{SUBDIR_NAME}"
@classmethod
def write_lines_to_file(cls, file_path: str, lines: List[str]) -> None:
@ -285,8 +289,10 @@ class CompilerBisector:
@classmethod
def delete_bisect_status(cls) -> None:
if os.path.exists(cls.get_dir()):
shutil.rmtree(cls.get_dir())
# in process_cache we have created if it exists, just the subdirectory of non created dir
dir_name = cls.in_process_cache if cls.in_process_cache else cls.get_dir()
if os.path.exists(dir_name):
shutil.rmtree(dir_name)
print("Bisection status deleted.")
else:
print("No bisection status found.")
@ -485,12 +491,19 @@ class CompilerBisector:
bisection_enabled_orig = cls.bisection_enabled
cls.delete_bisect_status()
cls.bisection_enabled = True
cls.in_process_cache = tempfile.mkdtemp()
def cleanup() -> None:
cls.bisection_enabled = bisection_enabled_orig
cls.delete_bisect_status()
cls.in_process_cache = None
cleanup_handler = atexit.register(cleanup)
# TODO - cli interface, and in-process different directories
class DisableBisect:
def __del__(self) -> None:
cls.bisection_enabled = bisection_enabled_orig
cls.delete_bisect_status()
cleanup()
atexit.unregister(cleanup_handler)
_cleanup = DisableBisect()