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

checkpoint.py

@@ -46,26 +46,35 @@ def copy_to_shm(file: str):
         yield file
         return
 
-    tmp_dir = "/dev/shm/"
+    with tempfile.NamedTemporaryFile(dir="/dev/shm", delete=False) as tmp_file:
-    fd, tmp_path = tempfile.mkstemp(dir=tmp_dir)
+        tmp_path = tmp_file.name
-    try:
+        try:
-        shutil.copyfile(file, tmp_path)
+            shutil.copyfile(file, tmp_path)
-        yield tmp_path
+            yield tmp_path
-    finally:
+        finally:
-        os.remove(tmp_path)
+            try:
-        os.close(fd)
+                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)
+    with tempfile.NamedTemporaryFile(dir=tmp_dir, delete=False) as tmp_file:
-    try:
+        tmp_path = tmp_file.name
-        yield tmp_path
+        try:
-        shutil.copyfile(tmp_path, file)
+            yield tmp_path
-    finally:
+            shutil.copyfile(tmp_path, file)
-        os.remove(tmp_path)
+        finally:
-        os.close(fd)
+            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: