pikers
/
piker
3
0
Fork 0
Code Issues 2 Pull Requests 4 Packages Projects Releases Wiki Activity

Store shm array size in token schema, use for loading

size_in_shm_token
Tyler Goodlet 2022-08-28 20:28:42 -04:00
parent 77a687bced
commit 7b653fe4f4
1 changed files with 15 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

20
piker/data/_sharedmem.py
View File

@ -118,6 +118,7 @@ class _Token(Struct, frozen=True):
shm_first_index_name: str shm_first_index_name: str
shm_last_index_name: str shm_last_index_name: str
dtype_descr: tuple dtype_descr: tuple
size: int
@property @property
def dtype(self) -> np.dtype: def dtype(self) -> np.dtype:
@ -158,6 +159,7 @@ def get_shm_token(key: str) -> _Token:
def _make_token( def _make_token(
key: str, key: str,
size: int,
dtype: Optional[np.dtype] = None, dtype: Optional[np.dtype] = None,
) -> _Token: ) -> _Token:
''' '''
@ -170,7 +172,8 @@ def _make_token(
shm_name=key, shm_name=key,
shm_first_index_name=key + "_first", shm_first_index_name=key + "_first",
shm_last_index_name=key + "_last", shm_last_index_name=key + "_last",
dtype_descr=tuple(np.dtype(dtype).descr) dtype_descr=tuple(np.dtype(dtype).descr),
size=size,
) )
@ -222,6 +225,7 @@ class ShmArray:
shm_first_index_name=self._first._shm.name, shm_first_index_name=self._first._shm.name,
shm_last_index_name=self._last._shm.name, shm_last_index_name=self._last._shm.name,
dtype_descr=tuple(self._array.dtype.descr), dtype_descr=tuple(self._array.dtype.descr),
size=self._len,
) )
@property @property
@ -467,7 +471,8 @@ def open_shm_array(
token = _make_token( token = _make_token(
key=key, key=key,
dtype=dtype size=size,
dtype=dtype,
) )
# create single entry arrays for storing an first and last indices # create single entry arrays for storing an first and last indices
@ -527,7 +532,7 @@ def open_shm_array(
def attach_shm_array( def attach_shm_array(
token: tuple[str, str, tuple[str, str]], token: tuple[str, str, tuple[str, str]],
size: int = _default_size, # size: int = _default_size,
readonly: bool = True, readonly: bool = True,
) -> ShmArray: ) -> ShmArray:
@ -566,7 +571,7 @@ def attach_shm_array(
raise _err raise _err
shmarr = np.ndarray( shmarr = np.ndarray(
(size,), (token.size,),
dtype=token.dtype, dtype=token.dtype,
buffer=shm.buf buffer=shm.buf
) )
@ -634,6 +639,7 @@ def maybe_open_shm_array(
use ``attach_shm_array``. use ``attach_shm_array``.
''' '''
size = kwargs.pop('size', _default_size)
try: try:
# see if we already know this key # see if we already know this key
token = _known_tokens[key] token = _known_tokens[key]
@ -641,7 +647,11 @@ def maybe_open_shm_array(
except KeyError: except KeyError:
log.warning(f"Could not find {key} in shms cache") log.warning(f"Could not find {key} in shms cache")
if dtype: if dtype:
token = _make_token(key, dtype) token = _make_token(
key,
size=size,
dtype=dtype,
)
try: try:
return attach_shm_array(token=token, **kwargs), False return attach_shm_array(token=token, **kwargs), False
except FileNotFoundError: except FileNotFoundError: