AMajor/grok-1
1
0
Fork 0
mirror of https://github.com/xai-org/grok-1.git synced 2024-11-24 04:29:53 +03:00
Code Issues Packages Projects Releases Wiki Activity

Improve tempfile handling in checkpoint.py

Browse Source 
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
1 changed files with 24 additions and 15 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

19
checkpoint.py
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)
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)
os.close(fd)
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)
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)
os.close(fd)
except OSError as e:
# Handle file deletion error gracefully
logger.error(f"Error deleting temporary file: {e}")
raise
def fast_unpickle(path: str) -> Any: