mirror of
https://github.com/xai-org/grok-1.git
synced 2024-11-24 04:29:53 +03:00
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:
parent
7050ed204b
commit
defd415f9a
@ -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:
|
||||
|
Loading…
Reference in New Issue
Block a user