A simple shared memory contextmanager based on https://docs.python.org/3/library/multiprocessing.shared_memory.html#multiprocessing.shared_memory.SharedMemory
| 271 | |
| 272 | |
| 273 | class SharedMemoryManager: |
| 274 | """ |
| 275 | A simple shared memory contextmanager based on |
| 276 | https://docs.python.org/3/library/multiprocessing.shared_memory.html#multiprocessing.shared_memory.SharedMemory |
| 277 | """ |
| 278 | def __init__(self, create=False) -> None: |
| 279 | self._shms: list[SharedMemory] = [] |
| 280 | self.__create = create |
| 281 | |
| 282 | def SharedMemory(self, *, name=None, create=False, size=0, track=True): |
| 283 | shm = SharedMemory(name=name, create=create, size=size, track=track) |
| 284 | shm._create = create |
| 285 | # Essential to keep refs on Windows |
| 286 | # https://stackoverflow.com/questions/74193377/filenotfounderror-when-passing-a-shared-memory-to-a-new-process#comment130999060_74194875 # noqa: E501 |
| 287 | self._shms.append(shm) |
| 288 | return shm |
| 289 | |
| 290 | def __enter__(self): |
| 291 | return self |
| 292 | |
| 293 | def __exit__(self, *args, **kwargs): |
| 294 | for shm in self._shms: |
| 295 | try: |
| 296 | shm.close() |
| 297 | if shm._create: |
| 298 | shm.unlink() |
| 299 | except Exception: |
| 300 | warnings.warn(f'Failed to unlink shared memory {shm.name!r}', |
| 301 | category=ResourceWarning, stacklevel=2) |
| 302 | raise |
| 303 | |
| 304 | def arr2shm(self, vals): |
| 305 | """Array to shared memory. Returns (shm_name, shape, dtype) used for restore.""" |
| 306 | assert vals.ndim == 1, (vals.ndim, vals.shape, vals) |
| 307 | shm = self.SharedMemory(size=vals.nbytes, create=True) |
| 308 | # np.array can't handle pandas' tz-aware datetimes |
| 309 | # https://github.com/numpy/numpy/issues/18279 |
| 310 | buf = np.ndarray(vals.shape, dtype=vals.dtype.base, buffer=shm.buf) |
| 311 | has_tz = getattr(vals.dtype, 'tz', None) |
| 312 | buf[:] = vals.tz_localize(None) if has_tz else vals # Copy into shared memory |
| 313 | return shm.name, vals.shape, vals.dtype |
| 314 | |
| 315 | def df2shm(self, df): |
| 316 | return tuple(( |
| 317 | (column, *self.arr2shm(values)) |
| 318 | for column, values in chain([(self._DF_INDEX_COL, df.index)], df.items()) |
| 319 | )) |
| 320 | |
| 321 | @staticmethod |
| 322 | def shm2s(shm, shape, dtype) -> pd.Series: |
| 323 | arr = np.ndarray(shape, dtype=dtype.base, buffer=shm.buf) |
| 324 | arr.setflags(write=False) |
| 325 | return pd.Series(arr, dtype=dtype) |
| 326 | |
| 327 | _DF_INDEX_COL = '__bt_index' |
| 328 | |
| 329 | @staticmethod |
| 330 | def shm2df(data_shm): |
no outgoing calls