forked from goodboy/tractor
Add `ShmList` wrapping the stdlib's `ShareableList`
First attempt at getting `multiprocessing.shared_memory.ShareableList` working; we wrap the stdlib type with a readonly attr and a `.key` for cross-actor lookup. Also, rename all `numpy` specific routines to have a `ndarray` suffix in the func names.macos_in_ci
parent
901353c213
commit
0e4b37d122
tractor
206
tractor/_shm.py
206
tractor/_shm.py
|
@ -28,6 +28,7 @@ import time
|
||||||
from typing import Optional
|
from typing import Optional
|
||||||
from multiprocessing.shared_memory import (
|
from multiprocessing.shared_memory import (
|
||||||
SharedMemory,
|
SharedMemory,
|
||||||
|
ShareableList,
|
||||||
_USE_POSIX,
|
_USE_POSIX,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
@ -87,10 +88,11 @@ disable_mantracker()
|
||||||
|
|
||||||
|
|
||||||
class SharedInt:
|
class SharedInt:
|
||||||
"""Wrapper around a single entry shared memory array which
|
'''
|
||||||
|
Wrapper around a single entry shared memory array which
|
||||||
holds an ``int`` value used as an index counter.
|
holds an ``int`` value used as an index counter.
|
||||||
|
|
||||||
"""
|
'''
|
||||||
def __init__(
|
def __init__(
|
||||||
self,
|
self,
|
||||||
shm: SharedMemory,
|
shm: SharedMemory,
|
||||||
|
@ -117,10 +119,13 @@ class SharedInt:
|
||||||
log.warning(f'Shm for {name} already unlinked?')
|
log.warning(f'Shm for {name} already unlinked?')
|
||||||
|
|
||||||
|
|
||||||
class _Token(Struct, frozen=True):
|
class _NpToken(Struct, frozen=True):
|
||||||
'''
|
'''
|
||||||
Internal represenation of a shared memory "token"
|
Internal represenation of a shared memory ``numpy`` array "token"
|
||||||
which can be used to key a system wide post shm entry.
|
which can be used to key and load a system (OS) wide shm entry
|
||||||
|
and correctly read the array by type signature.
|
||||||
|
|
||||||
|
This type is msg safe.
|
||||||
|
|
||||||
'''
|
'''
|
||||||
shm_name: str # this servers as a "key" value
|
shm_name: str # this servers as a "key" value
|
||||||
|
@ -137,18 +142,18 @@ class _Token(Struct, frozen=True):
|
||||||
return self.to_dict()
|
return self.to_dict()
|
||||||
|
|
||||||
@classmethod
|
@classmethod
|
||||||
def from_msg(cls, msg: dict) -> _Token:
|
def from_msg(cls, msg: dict) -> _NpToken:
|
||||||
if isinstance(msg, _Token):
|
if isinstance(msg, _NpToken):
|
||||||
return msg
|
return msg
|
||||||
|
|
||||||
# TODO: native struct decoding
|
# TODO: native struct decoding
|
||||||
# return _token_dec.decode(msg)
|
# return _token_dec.decode(msg)
|
||||||
|
|
||||||
msg['dtype_descr'] = tuple(map(tuple, msg['dtype_descr']))
|
msg['dtype_descr'] = tuple(map(tuple, msg['dtype_descr']))
|
||||||
return _Token(**msg)
|
return _NpToken(**msg)
|
||||||
|
|
||||||
|
|
||||||
# _token_dec = msgspec.msgpack.Decoder(_Token)
|
# _token_dec = msgspec.msgpack.Decoder(_NpToken)
|
||||||
|
|
||||||
# TODO: this api?
|
# TODO: this api?
|
||||||
# _known_tokens = tractor.ActorVar('_shm_tokens', {})
|
# _known_tokens = tractor.ActorVar('_shm_tokens', {})
|
||||||
|
@ -159,10 +164,14 @@ class _Token(Struct, frozen=True):
|
||||||
_known_tokens = {}
|
_known_tokens = {}
|
||||||
|
|
||||||
|
|
||||||
def get_shm_token(key: str) -> _Token:
|
def get_shm_token(key: str) -> _NpToken | str:
|
||||||
"""Convenience func to check if a token
|
'''
|
||||||
|
Convenience func to check if a token
|
||||||
for the provided key is known by this process.
|
for the provided key is known by this process.
|
||||||
"""
|
|
||||||
|
Returns either the ``numpy`` token or a string for a shared list.
|
||||||
|
|
||||||
|
'''
|
||||||
return _known_tokens.get(key)
|
return _known_tokens.get(key)
|
||||||
|
|
||||||
|
|
||||||
|
@ -171,13 +180,13 @@ def _make_token(
|
||||||
size: int,
|
size: int,
|
||||||
dtype: np.dtype,
|
dtype: np.dtype,
|
||||||
|
|
||||||
) -> _Token:
|
) -> _NpToken:
|
||||||
'''
|
'''
|
||||||
Create a serializable token that can be used
|
Create a serializable token that can be used
|
||||||
to access a shared array.
|
to access a shared array.
|
||||||
|
|
||||||
'''
|
'''
|
||||||
return _Token(
|
return _NpToken(
|
||||||
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",
|
||||||
|
@ -188,7 +197,7 @@ def _make_token(
|
||||||
|
|
||||||
class ShmArray:
|
class ShmArray:
|
||||||
'''
|
'''
|
||||||
A shared memory ``numpy`` (compatible) array API.
|
A shared memory ``numpy.ndarray`` API.
|
||||||
|
|
||||||
An underlying shared memory buffer is allocated based on
|
An underlying shared memory buffer is allocated based on
|
||||||
a user specified ``numpy.ndarray``. This fixed size array
|
a user specified ``numpy.ndarray``. This fixed size array
|
||||||
|
@ -228,8 +237,8 @@ class ShmArray:
|
||||||
# TODO: ringbuf api?
|
# TODO: ringbuf api?
|
||||||
|
|
||||||
@property
|
@property
|
||||||
def _token(self) -> _Token:
|
def _token(self) -> _NpToken:
|
||||||
return _Token(
|
return _NpToken(
|
||||||
shm_name=self._shm.name,
|
shm_name=self._shm.name,
|
||||||
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,
|
||||||
|
@ -446,15 +455,17 @@ class ShmArray:
|
||||||
...
|
...
|
||||||
|
|
||||||
|
|
||||||
def open_shm_array(
|
def open_shm_ndarray(
|
||||||
|
|
||||||
key: Optional[str] = None,
|
key: Optional[str] = None,
|
||||||
size: int = _default_size, # see above
|
size: int = int(2 ** 10),
|
||||||
dtype: Optional[np.dtype] = None,
|
dtype: np.dtype | None = None,
|
||||||
|
append_start_index: int = 0,
|
||||||
readonly: bool = False,
|
readonly: bool = False,
|
||||||
|
|
||||||
) -> ShmArray:
|
) -> ShmArray:
|
||||||
'''Open a memory shared ``numpy`` using the standard library.
|
'''
|
||||||
|
Open a memory shared ``numpy`` using the standard library.
|
||||||
|
|
||||||
This call unlinks (aka permanently destroys) the buffer on teardown
|
This call unlinks (aka permanently destroys) the buffer on teardown
|
||||||
and thus should be used from the parent-most accessor (process).
|
and thus should be used from the parent-most accessor (process).
|
||||||
|
@ -501,10 +512,10 @@ def open_shm_array(
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
|
|
||||||
# start the "real-time" updated section after 3-days worth of 1s
|
# Start the "real-time" append-updated (or "pushed-to") section
|
||||||
# sampled OHLC. this allows appending up to a days worth from
|
# after some start index: ``append_start_index``. This allows appending
|
||||||
# tick/quote feeds before having to flush to a (tsdb) storage
|
# from a start point in the array which isn't the 0 index and looks
|
||||||
# backend, and looks something like,
|
# something like,
|
||||||
# -------------------------
|
# -------------------------
|
||||||
# | | i
|
# | | i
|
||||||
# _________________________
|
# _________________________
|
||||||
|
@ -518,7 +529,7 @@ def open_shm_array(
|
||||||
# this sets the index to 3/4 of the length of the buffer
|
# this sets the index to 3/4 of the length of the buffer
|
||||||
# leaving a "days worth of second samples" for the real-time
|
# leaving a "days worth of second samples" for the real-time
|
||||||
# section.
|
# section.
|
||||||
last.value = first.value = _rt_buffer_start
|
last.value = first.value = append_start_index
|
||||||
|
|
||||||
shmarr = ShmArray(
|
shmarr = ShmArray(
|
||||||
array,
|
array,
|
||||||
|
@ -540,7 +551,7 @@ def open_shm_array(
|
||||||
return shmarr
|
return shmarr
|
||||||
|
|
||||||
|
|
||||||
def attach_shm_array(
|
def attach_shm_ndarray(
|
||||||
token: tuple[str, str, tuple[str, str]],
|
token: tuple[str, str, tuple[str, str]],
|
||||||
readonly: bool = True,
|
readonly: bool = True,
|
||||||
|
|
||||||
|
@ -553,11 +564,11 @@ def attach_shm_array(
|
||||||
access are constructed.
|
access are constructed.
|
||||||
|
|
||||||
'''
|
'''
|
||||||
token = _Token.from_msg(token)
|
token = _NpToken.from_msg(token)
|
||||||
key = token.shm_name
|
key = token.shm_name
|
||||||
|
|
||||||
if key in _known_tokens:
|
if key in _known_tokens:
|
||||||
assert _Token.from_msg(_known_tokens[key]) == token, "WTF"
|
assert _NpToken.from_msg(_known_tokens[key]) == token, "WTF"
|
||||||
|
|
||||||
# XXX: ugh, looks like due to the ``shm_open()`` C api we can't
|
# XXX: ugh, looks like due to the ``shm_open()`` C api we can't
|
||||||
# actually place files in a subdir, see discussion here:
|
# actually place files in a subdir, see discussion here:
|
||||||
|
@ -625,10 +636,14 @@ def attach_shm_array(
|
||||||
return sha
|
return sha
|
||||||
|
|
||||||
|
|
||||||
def maybe_open_shm_array(
|
def maybe_open_shm_ndarray(
|
||||||
key: str,
|
key: str, # unique identifier for segment
|
||||||
dtype: Optional[np.dtype] = None,
|
|
||||||
**kwargs,
|
# from ``open_shm_array()``
|
||||||
|
size: int = int(2 ** 10), # array length in index terms
|
||||||
|
dtype: np.dtype | None = None,
|
||||||
|
append_start_index: int = 0,
|
||||||
|
readonly: bool = True,
|
||||||
|
|
||||||
) -> tuple[ShmArray, bool]:
|
) -> tuple[ShmArray, bool]:
|
||||||
'''
|
'''
|
||||||
|
@ -641,18 +656,23 @@ def maybe_open_shm_array(
|
||||||
running in **this** process. Systems where multiple actors may seek
|
running in **this** process. Systems where multiple actors may seek
|
||||||
to access a common block can use this function to attempt to acquire
|
to access a common block can use this function to attempt to acquire
|
||||||
a token as discovered by the actors who have previously stored
|
a token as discovered by the actors who have previously stored
|
||||||
a "key" -> ``_Token`` map in an actor local (aka python global)
|
a "key" -> ``_NpToken`` map in an actor local (aka python global)
|
||||||
variable.
|
variable.
|
||||||
|
|
||||||
If you know the explicit ``_Token`` for your memory segment instead
|
If you know the explicit ``_NpToken`` for your memory segment instead
|
||||||
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]
|
||||||
return attach_shm_array(token=token, **kwargs), False
|
return (
|
||||||
|
attach_shm_ndarray(
|
||||||
|
token=token,
|
||||||
|
readonly=readonly,
|
||||||
|
),
|
||||||
|
False, # not newly opened
|
||||||
|
)
|
||||||
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:
|
||||||
|
@ -661,8 +681,16 @@ def maybe_open_shm_array(
|
||||||
size=size,
|
size=size,
|
||||||
dtype=dtype,
|
dtype=dtype,
|
||||||
)
|
)
|
||||||
|
else:
|
||||||
|
|
||||||
try:
|
try:
|
||||||
return attach_shm_array(token=token, **kwargs), False
|
return (
|
||||||
|
attach_shm_ndarray(
|
||||||
|
token=token,
|
||||||
|
readonly=readonly,
|
||||||
|
),
|
||||||
|
False,
|
||||||
|
)
|
||||||
except FileNotFoundError:
|
except FileNotFoundError:
|
||||||
log.warning(f"Could not attach to shm with token {token}")
|
log.warning(f"Could not attach to shm with token {token}")
|
||||||
|
|
||||||
|
@ -671,36 +699,84 @@ def maybe_open_shm_array(
|
||||||
# Attempt to open a block and expect
|
# Attempt to open a block and expect
|
||||||
# to fail if a block has been allocated
|
# to fail if a block has been allocated
|
||||||
# on the OS by someone else.
|
# on the OS by someone else.
|
||||||
return open_shm_array(key=key, dtype=dtype, **kwargs), True
|
return (
|
||||||
|
open_shm_ndarray(
|
||||||
|
key=key,
|
||||||
|
size=size,
|
||||||
|
dtype=dtype,
|
||||||
|
append_start_index=append_start_index,
|
||||||
|
readonly=readonly,
|
||||||
|
),
|
||||||
|
True,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
def try_read(
|
class ShmList(ShareableList):
|
||||||
array: np.ndarray
|
|
||||||
|
|
||||||
) -> Optional[np.ndarray]:
|
|
||||||
'''
|
'''
|
||||||
Try to read the last row from a shared mem array or ``None``
|
Carbon copy of ``.shared_memory.ShareableList`` but add a
|
||||||
if the array read returns a zero-length array result.
|
readonly state instance var.
|
||||||
|
|
||||||
Can be used to check for backfilling race conditions where an array
|
|
||||||
is currently being (re-)written by a writer actor but the reader is
|
|
||||||
unaware and reads during the window where the first and last indexes
|
|
||||||
are being updated.
|
|
||||||
|
|
||||||
'''
|
'''
|
||||||
try:
|
def __init__(
|
||||||
return array[-1]
|
self,
|
||||||
except IndexError:
|
sequence: list | None = None,
|
||||||
# XXX: race condition with backfilling shm.
|
*,
|
||||||
#
|
name: str | None = None,
|
||||||
# the underlying issue is that a backfill (aka prepend) and subsequent
|
readonly: bool = True
|
||||||
# shm array first/last index update could result in an empty array
|
|
||||||
# read here since the indices may be updated in such a way that
|
|
||||||
# a read delivers an empty array (though it seems like we
|
|
||||||
# *should* be able to prevent that?). also, as and alt and
|
|
||||||
# something we need anyway, maybe there should be some kind of
|
|
||||||
# signal that a prepend is taking place and this consumer can
|
|
||||||
# respond (eg. redrawing graphics) accordingly.
|
|
||||||
|
|
||||||
# the array read was emtpy
|
) -> None:
|
||||||
return None
|
self._readonly = readonly
|
||||||
|
self._key = name
|
||||||
|
return super().__init__(
|
||||||
|
sequence=sequence,
|
||||||
|
name=name,
|
||||||
|
)
|
||||||
|
|
||||||
|
@property
|
||||||
|
def key(self) -> str:
|
||||||
|
return self._key
|
||||||
|
|
||||||
|
def __setitem__(
|
||||||
|
self,
|
||||||
|
position,
|
||||||
|
value,
|
||||||
|
|
||||||
|
) -> None:
|
||||||
|
|
||||||
|
# mimick ``numpy`` error
|
||||||
|
if self._readonly:
|
||||||
|
raise ValueError('assignment destination is read-only')
|
||||||
|
|
||||||
|
return super().__setitem__(position, value)
|
||||||
|
|
||||||
|
|
||||||
|
def open_shm_list(
|
||||||
|
key: str,
|
||||||
|
sequence: list | None = None,
|
||||||
|
size: int = int(2 ** 10),
|
||||||
|
dtype: np.dtype | None = None,
|
||||||
|
readonly: bool = True,
|
||||||
|
|
||||||
|
) -> ShmList:
|
||||||
|
|
||||||
|
if sequence is None:
|
||||||
|
sequence = list(map(float, range(size)))
|
||||||
|
|
||||||
|
shml = ShmList(
|
||||||
|
sequence=sequence,
|
||||||
|
name=key,
|
||||||
|
readonly=readonly,
|
||||||
|
)
|
||||||
|
|
||||||
|
# "close" attached shm on actor teardown
|
||||||
|
tractor.current_actor().lifetime_stack.callback(shml.shm.close)
|
||||||
|
tractor.current_actor().lifetime_stack.callback(shml.shm.unlink)
|
||||||
|
|
||||||
|
return shml
|
||||||
|
|
||||||
|
|
||||||
|
def attach_shm_list(
|
||||||
|
key: str,
|
||||||
|
) -> ShmList:
|
||||||
|
|
||||||
|
return ShmList(name=key)
|
||||||
|
|
Loading…
Reference in New Issue