Addressing issues

analog/storage/default.py

@@ -87,8 +87,10 @@ def _flush_unsafe(self, buffer, push_count) -> str:
         """
         _flush_unsafe is thread unsafe flush of current buffer. No shared variable must be allowed.
         """
-        save_path = str(os.path.join(self.log_dir, f"data_{push_count}.pt"))
+        save_path = self.file_prefix + f"{push_count}.mmap"
         torch.save(buffer, save_path)
+        buffer_list = [(k, v) for k, v in buffer]
+        self.mmap_handler.write(buffer_list, save_path)
         return save_path
 
     def _flush_safe(self) -> str:
@@ -227,8 +229,8 @@ def _find_chunk_indices(self, directory):
 
     def _add_metadata_and_mmap(self, mmap_filename, chunk_index):
         # Load the memmap file
-        mmap, metadata = self.mmap_handler.read(mmap_filename)
-        self.memmaps.append(mmap)
+        with self.mmap_handler.read(mmap_filename) as (mmap, metadata):
+            self.memmaps.append(mmap)
 
         # Update the mapping from data_id to chunk
         for entry in metadata:

analog/storage/utils.py

@@ -3,7 +3,7 @@
 import json
 import numpy as np
 
-from einops import rearrange
+from contextlib import contextmanager
 
 
 def extract_arrays(obj, base_path=()):
@@ -67,6 +67,7 @@ def write(self, data_buffer, filename):
         with open(metadata_filename, "w") as f:
             json.dump(metadata, f, indent=2)
 
+    @contextmanager
     def read(self, filename):
         """
         read reads the file by chunk index, it will return the data_buffer with metadata.
@@ -84,7 +85,11 @@ def read(self, filename):
             os.path.join(self.save_path, filename), dtype=self.mmap_dtype, mode="r"
         )
         metadata = self.read_metafile(file_root + "_metadata.json")
-        return mmap, metadata
+
+        try:
+            yield mmap, metadata
+        finally:
+            del mmap
 
     def read_metafile(self, meta_filename):
         file_root, file_ext = os.path.splitext(meta_filename)

tests/storage/test_util_integration.py

@@ -57,18 +57,20 @@ def test_write_and_read(self):
 
         handler.write(data_buffer, filename)
 
-        mmap, metadata = handler.read(filename)
-        for item in metadata:
-            offset = item["offset"]
-            size = item["size"]
-            shape = tuple(item["shape"])
-            dtype = np.dtype(item["dtype"])
-            expected_data = data_buffer[item["data_id"]][1]["dummy_data"]
-            read_data = np.frombuffer(
-                mmap, dtype=dtype, count=size // dtype.itemsize, offset=offset
-            ).reshape(shape)
-            # Test if expected value and read value equals
-            self.assertTrue(np.array_equal(read_data, expected_data), "Data mismatch")
+        with handler.read(filename) as (mmap, metadata):
+            for item in metadata:
+                offset = item["offset"]
+                size = item["size"]
+                shape = tuple(item["shape"])
+                dtype = np.dtype(item["dtype"])
+                expected_data = data_buffer[item["data_id"]][1]["dummy_data"]
+                read_data = np.frombuffer(
+                    mmap, dtype=dtype, count=size // dtype.itemsize, offset=offset
+                ).reshape(shape)
+                # Test if expected value and read value equals
+                self.assertTrue(
+                    np.array_equal(read_data, expected_data), "Data mismatch"
+                )
 
     def test_read(self):
         expected_files_path = os.path.join(
@@ -82,11 +84,15 @@ def test_read(self):
         ]
 
         handler.write(data_buffer, filename)
-        mmap, metadata = handler.read(filename)
+        mmap, metadata, expected_mmap = None, None, None
+
+        with handler.read(filename) as (mm, md):
+            mmap = mm
+            metadata = md
+
+        with handler.read("expected_data.mmap") as (em, _):
+            expected_mmap = em
 
-        expected_mmap, _ = handler.read(
-            "expected_data.mmap"
-        )  # Using same metadata file.
         for item in metadata:
             offset = item["offset"]
             size = item["size"]