Improve tempfile handling in checkpoint.py

Using more context managers, as well as dynamic temp dir creation, improve temp file handling, error handling, and logging
This commit is contained in:
Aareon Sullivan 2024-03-22 19:22:04 -05:00 committed by GitHub
parent 7050ed204b
commit defd415f9a
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -46,26 +46,35 @@ def copy_to_shm(file: str):
yield file
return
tmp_dir = "/dev/shm/"
fd, tmp_path = tempfile.mkstemp(dir=tmp_dir)
try:
shutil.copyfile(file, tmp_path)
yield tmp_path
finally:
os.remove(tmp_path)
os.close(fd)
with tempfile.NamedTemporaryFile(dir="/dev/shm", delete=False) as tmp_file:
tmp_path = tmp_file.name
try:
shutil.copyfile(file, tmp_path)
yield tmp_path
finally:
try:
os.remove(tmp_path)
except OSError as e:
# Handle file deletion error gracefully
logger.error(f"Error deleting temporary file: {e}")
raise
@contextlib.contextmanager
def copy_from_shm(file: str):
tmp_dir = "/dev/shm/"
fd, tmp_path = tempfile.mkstemp(dir=tmp_dir)
try:
yield tmp_path
shutil.copyfile(tmp_path, file)
finally:
os.remove(tmp_path)
os.close(fd)
with tempfile.NamedTemporaryFile(dir=tmp_dir, delete=False) as tmp_file:
tmp_path = tmp_file.name
try:
yield tmp_path
shutil.copyfile(tmp_path, file)
finally:
try:
os.remove(tmp_path)
except OSError as e:
# Handle file deletion error gracefully
logger.error(f"Error deleting temporary file: {e}")
raise
def fast_unpickle(path: str) -> Any: