Load skims into shared memory to be accessed by later models

SANDAG · Apr 8, 2024 · 422ca6c · 422ca6c
1 parent 2c11873
commit 422ca6c
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 15 deletions.
diff --git a/activitysim/core/mem.py b/activitysim/core/mem.py
@@ -296,6 +296,9 @@ def shared_memory_size(data_buffers=None):
 
             shared_size += Dataset.shm.preload_shared_memory_size(data_buffer[11:])
             continue
+        if isinstance(data_buffer, multiprocessing.shared_memory.SharedMemory):
+            shared_size += data_buffer.size
+            continue
         try:
             obj = data_buffer.get_obj()
         except Exception:

diff --git a/activitysim/core/skim_dict_factory.py b/activitysim/core/skim_dict_factory.py
@@ -409,19 +409,26 @@ def allocate_skim_buffer(self, skim_info, shared=False):
             f"total size: {util.INT(csz)} ({util.GB(csz)})"
         )
 
-        if shared:
-            if dtype_name == "float64":
-                typecode = "d"
-            elif dtype_name == "float32":
-                typecode = "f"
-            else:
-                raise RuntimeError(
-                    "allocate_skim_buffer unrecognized dtype %s" % dtype_name
-                )
-
-            buffer = multiprocessing.RawArray(typecode, buffer_size)
-        else:
-            buffer = np.zeros(buffer_size, dtype=dtype)
+        # if shared:
+            # if dtype_name == "float64":
+            #     typecode = "d"
+            # elif dtype_name == "float32":
+            #     typecode = "f"
+            # else:
+            #     raise RuntimeError(
+            #         "allocate_skim_buffer unrecognized dtype %s" % dtype_name
+            #     )
+
+            # buffer = multiprocessing.RawArray(typecode, buffer_size)
+        shared_mem_name = f"skim_shared_memory__{skim_info.skim_tag}"
+        try:
+            buffer = multiprocessing.shared_memory.SharedMemory(name=shared_mem_name)
+            logger.info(f"skim buffer already allocated in shared memory: {shared_mem_name}, size: {buffer.size}")
+        except FileNotFoundError:
+            buffer = multiprocessing.shared_memory.SharedMemory(create=True, size=csz, name=shared_mem_name)
+            logger.info(f"allocating skim buffer in shared memory: {shared_mem_name}, size: {buffer.size}")
+        # else:
+        #     buffer = np.zeros(buffer_size, dtype=dtype)
 
         return buffer
 
@@ -440,8 +447,9 @@ def _skim_data_from_buffer(self, skim_info, skim_buffer):
         """
 
         dtype = np.dtype(skim_info.dtype_name)
-        assert len(skim_buffer) == util.iprod(skim_info.skim_data_shape)
-        skim_data = np.frombuffer(skim_buffer, dtype=dtype).reshape(
+        # assert len(skim_buffer) == util.iprod(skim_info.skim_data_shape)
+        assert skim_buffer.size >= util.iprod(skim_info.skim_data_shape) * dtype.itemsize
+        skim_data = np.frombuffer(skim_buffer.buf, dtype=dtype, count=util.iprod(skim_info.skim_data_shape)).reshape(
             skim_info.skim_data_shape
         )
         return skim_data
@@ -462,6 +470,9 @@ def load_skims_to_buffer(self, skim_info, skim_buffer):
         skim_data = self._skim_data_from_buffer(skim_info, skim_buffer)
         assert skim_data.shape == skim_info.skim_data_shape
 
+        if skim_data.any():
+            return
+
         if read_cache:
             # returns None if cache file not found
             cache_data = self._open_existing_readonly_memmap_skim_cache(skim_info)