proxystore.store¶

proxystore/store/__init__.py

The ProxyStore Store interface.

Store ¶

Store(
    name: str,
    connector: ConnectorT,
    *,
    serializer: SerializerT | None = None,
    deserializer: DeserializerT | None = None,
    cache_size: int = 16,
    metrics: bool = False
)

Bases: Generic[ConnectorT]

Key-value store interface for proxies.

Tip

A Store instance can be used as a context manager which will automatically call close() on exit.

with Store('my-store', connector=...) as store:
    key = store.put('value')
    store.get(key)

Parameters:

name (str) –

Name of the store instance.
connector (ConnectorT) –

Connector instance to use for object storage.
serializer (SerializerT | None, default: None ) –

Optional callable which serializes the object. If None, the default serializer (serialize()) will be used.
deserializer (DeserializerT | None, default: None ) –

Optional callable used by the factory to deserialize the byte string. If None, the default deserializer (deserialize()) will be used.
cache_size (int, default: 16 ) –

Size of LRU cache (in # of objects). If 0, the cache is disabled. The cache is local to the Python process.
metrics (bool, default: False ) –

Enable recording operation metrics.

Raises:

ValueError –

If cache_size is less than zero.

Source code in proxystore/store/base.py

def __init__(
    self,
    name: str,
    connector: ConnectorT,
    *,
    serializer: SerializerT | None = None,
    deserializer: DeserializerT | None = None,
    cache_size: int = 16,
    metrics: bool = False,
) -> None:
    if cache_size < 0:
        raise ValueError(
            f'Cache size cannot be negative. Got {cache_size}.',
        )

    self.connector = connector
    self.cache: LRUCache[ConnectorKeyT, Any] = LRUCache(cache_size)
    self._name = name
    self._metrics = StoreMetrics() if metrics else None
    self._cache_size = cache_size
    self._serializer = serializer
    self._deserializer = deserializer

    logger.info(f'Initialized {self}')

name `property` ¶

name: str

Name of this Store instance.

metrics `property` ¶

metrics: StoreMetrics | None

Optional metrics for this instance.

serializer `property` ¶

serializer: SerializerT

Serializer for this instance.

deserializer `property` ¶

deserializer: DeserializerT

Deserializer for this instance.

close() ¶

close(*args: Any, **kwargs: Any) -> None

Close the connector associated with the store.

Warning

This method should only be called at the end of the program when the store will no longer be used, for example once all proxies have been resolved.

Parameters:

args (Any, default: () ) –

Positional arguments to pass to Connector.close().
kwargs (Any, default: {} ) –

Keyword arguments to pass to Connector.close().

Source code in proxystore/store/base.py

def close(self, *args: Any, **kwargs: Any) -> None:
    """Close the connector associated with the store.

    Warning:
        This method should only be called at the end of the program
        when the store will no longer be used, for example once all
        proxies have been resolved.

    Args:
        args: Positional arguments to pass to
            [`Connector.close()`][proxystore.connectors.protocols.Connector.close].
        kwargs: Keyword arguments to pass to
            [`Connector.close()`][proxystore.connectors.protocols.Connector.close].
    """
    self.connector.close(*args, **kwargs)

config() ¶

config() -> dict[str, Any]

Get the store configuration.

Example

>>> store = Store(...)
>>> config = store.config()
>>> store = Store.from_config(config)

Returns:

dict[str, Any] –

Store configuration.

Source code in proxystore/store/base.py

def config(self) -> dict[str, Any]:
    """Get the store configuration.

    Example:
        ```python
        >>> store = Store(...)
        >>> config = store.config()
        >>> store = Store.from_config(config)
        ```

    Returns:
        Store configuration.
    """
    return {
        'name': self.name,
        'connector_type': get_class_path(type(self.connector)),
        'connector_config': self.connector.config(),
        'serializer': self._serializer,
        'deserializer': self._deserializer,
        'cache_size': self._cache_size,
        'metrics': self.metrics is not None,
    }

from_config() `classmethod` ¶

from_config(config: dict[str, Any]) -> Store[Any]

Create a new store instance from a configuration.

Parameters:

config (dict[str, Any]) –

Configuration returned by .config().

Returns:

Store[Any] –

Store instance.

Source code in proxystore/store/base.py

@classmethod
def from_config(cls, config: dict[str, Any]) -> Store[Any]:
    """Create a new store instance from a configuration.

    Args:
        config: Configuration returned by `#!python .config()`.

    Returns:
        Store instance.
    """
    config = config.copy()  # Avoid messing with callers version
    connector_type = config.pop('connector_type')
    connector_config = config.pop('connector_config')
    connector = import_class(connector_type)
    config['connector'] = connector.from_config(connector_config)
    return cls(**config)

future() ¶

future(
    *,
    evict: bool = False,
    serializer: SerializerT | None = None,
    deserializer: DeserializerT | None = None,
    polling_interval: float = 1,
    polling_timeout: float | None = None
) -> ProxyFuture[T]

Create a future to an object.

Example

from proxystore.connectors.file import FileConnector
from proxystore.store import Store
from proxystore.store.future import ProxyFuture

def remote_foo(future: ProxyFuture) -> None:
    # Computation that generates a result value needed by
    # the remote_bar function.
    future.set_result(...)

def remote_bar(data: Any) -> None:
    # Function uses data, which is a proxy, as normal, blocking
    # until the remote_foo function has called set_result.
    ...

with Store('future-example', FileConnector(...)) as store:
    future = store.future()

    # The invoke_remove function invokes a provided function
    # on a remote process. For example, this could be a serverless
    # function execution.
    foo_result_future = invoke_remote(remote_foo, future)
    bar_result_future = invoke_remote(remote_bar, future.proxy())

    foo_result_future.result()
    bar_result_future.result()

Warning

This method only works if the connector is of type DeferrableConnector.

Warning

This method and the ProxyFuture.proxy() are experimental features and may change in future releases.

Parameters:

evict (bool, default: False ) –

If a proxy returned by ProxyFuture.proxy() should evict the object once resolved.
serializer (SerializerT | None, default: None ) –

Optionally override the default serializer for the store instance.
deserializer (DeserializerT | None, default: None ) –

Optionally override the default deserializer for the store instance.
polling_interval (float, default: 1 ) –

Seconds to sleep between polling the store for the object when getting the result of the future.
polling_timeout (float | None, default: None ) –

Optional maximum number of seconds to poll for when getting the result of the future.

Returns:

ProxyFuture[T] –

Future which can be used to get the result object at a later time or create a proxy which will resolve to the result of the future.

Raises:

NotImplementedError –

If the connector is not of type DeferrableConnector.

Source code in proxystore/store/base.py

def future(
    self,
    *,
    evict: bool = False,
    serializer: SerializerT | None = None,
    deserializer: DeserializerT | None = None,
    polling_interval: float = 1,
    polling_timeout: float | None = None,
) -> ProxyFuture[T]:
    """Create a future to an object.

    Example:
        ```python
        from proxystore.connectors.file import FileConnector
        from proxystore.store import Store
        from proxystore.store.future import ProxyFuture

        def remote_foo(future: ProxyFuture) -> None:
            # Computation that generates a result value needed by
            # the remote_bar function.
            future.set_result(...)

        def remote_bar(data: Any) -> None:
            # Function uses data, which is a proxy, as normal, blocking
            # until the remote_foo function has called set_result.
            ...

        with Store('future-example', FileConnector(...)) as store:
            future = store.future()

            # The invoke_remove function invokes a provided function
            # on a remote process. For example, this could be a serverless
            # function execution.
            foo_result_future = invoke_remote(remote_foo, future)
            bar_result_future = invoke_remote(remote_bar, future.proxy())

            foo_result_future.result()
            bar_result_future.result()
        ```

    Warning:
        This method only works if the `connector` is of type
        [`DeferrableConnector`][proxystore.connectors.protocols.DeferrableConnector].

    Warning:
        This method and the
        [`ProxyFuture.proxy()`][proxystore.store.future.ProxyFuture.proxy]
        are experimental features and may change in future releases.

    Args:
        evict: If a proxy returned by
            [`ProxyFuture.proxy()`][proxystore.store.future.ProxyFuture.proxy]
            should evict the object once resolved.
        serializer: Optionally override the default serializer for the
            store instance.
        deserializer: Optionally override the default deserializer for the
            store instance.
        polling_interval: Seconds to sleep between polling the store for
            the object when getting the result of the future.
        polling_timeout: Optional maximum number of seconds to poll for
            when getting the result of the future.

    Returns:
        Future which can be used to get the result object at a later time \
        or create a proxy which will resolve to the result of the future.

    Raises:
        NotImplementedError: If the `connector` is not of type
            [`DeferrableConnector`][proxystore.connectors.protocols.DeferrableConnector].
    """
    if not isinstance(self.connector, DeferrableConnector):
        raise NotImplementedError(
            'The provided connector is type '
            f'{type(self.connector).__name__} which does not implement '
            f'the {DeferrableConnector.__name__} necessary to use the '
            f'{ProxyFuture.__name__} interface.',
        )

    warnings.warn(
        'The Store.future() and ProxyFuture interfaces are experimental '
        'and may change in future releases.',
        category=ExperimentalWarning,
        stacklevel=2,
    )

    key = self.connector.new_key()
    factory: PollingStoreFactory[ConnectorT, T] = PollingStoreFactory(
        key,
        store_config=self.config(),
        deserializer=deserializer,
        evict=evict,
        polling_interval=polling_interval,
        polling_timeout=polling_timeout,
    )
    return ProxyFuture(factory, serializer=serializer)

evict() ¶

evict(key: ConnectorKeyT) -> None

Evict the object associated with the key.

Parameters:

key (ConnectorKeyT) –

Key associated with object to evict.

Source code in proxystore/store/base.py

def evict(self, key: ConnectorKeyT) -> None:
    """Evict the object associated with the key.

    Args:
        key: Key associated with object to evict.
    """
    with Timer() as timer:
        with Timer() as connector_timer:
            self.connector.evict(key)

        if self.metrics is not None:
            ctime = connector_timer.elapsed_ns
            self.metrics.add_time('store.evict.connector', key, ctime)

        self.cache.evict(key)

    if self.metrics is not None:
        self.metrics.add_time('store.evict', key, timer.elapsed_ns)

    logger.debug(
        f'Store(name="{self.name}"): EVICT {key} in '
        f'{timer.elapsed_ms:.3f} ms',
    )

exists() ¶

exists(key: ConnectorKeyT) -> bool

Check if an object associated with the key exists.

Parameters:

key (ConnectorKeyT) –

Key potentially associated with stored object.

Returns:

bool –

If an object associated with the key exists.

Source code in proxystore/store/base.py

def exists(self, key: ConnectorKeyT) -> bool:
    """Check if an object associated with the key exists.

    Args:
        key: Key potentially associated with stored object.

    Returns:
        If an object associated with the key exists.
    """
    with Timer() as timer:
        res = self.cache.exists(key)
        if not res:
            with Timer() as connector_timer:
                res = self.connector.exists(key)

            if self.metrics is not None:
                ctime = connector_timer.elapsed_ns
                self.metrics.add_time('store.exists.connector', key, ctime)

    if self.metrics is not None:
        self.metrics.add_time('store.exists', key, timer.elapsed_ns)

    logger.debug(
        f'Store(name="{self.name}"): EXISTS {key} in '
        f'{timer.elapsed_ms:.3f} ms',
    )
    return res

get() ¶

get(
    key: ConnectorKeyT,
    *,
    deserializer: DeserializerT | None = None,
    default: object | None = None
) -> Any | None

Get the object associated with the key.

Parameters:

key (ConnectorKeyT) –

Key associated with the object to retrieve.
deserializer (DeserializerT | None, default: None ) –

Optionally override the default deserializer for the store instance.
default (object | None, default: None ) –

An optional value to be returned if an object associated with the key does not exist.

Returns:

Any | None –

Object or None if the object does not exist.

Source code in proxystore/store/base.py

def get(
    self,
    key: ConnectorKeyT,
    *,
    deserializer: DeserializerT | None = None,
    default: object | None = None,
) -> Any | None:
    """Get the object associated with the key.

    Args:
        key: Key associated with the object to retrieve.
        deserializer: Optionally override the default deserializer for the
            store instance.
        default: An optional value to be returned if an object
            associated with the key does not exist.

    Returns:
        Object or `None` if the object does not exist.
    """
    timer = Timer()
    timer.start()

    if self.is_cached(key):
        value = self.cache.get(key)

        timer.stop()
        if self.metrics is not None:
            self.metrics.add_counter('store.get.cache_hits', key, 1)
            self.metrics.add_time('store.get', key, timer.elapsed_ns)

        logger.debug(
            f'Store(name="{self.name}"): GET {key} in '
            f'{timer.elapsed_ms:.3f} ms (cached=True)',
        )
        return value

    with Timer() as connector_timer:
        value = self.connector.get(key)

    if self.metrics is not None:
        ctime = connector_timer.elapsed_ns
        self.metrics.add_counter('store.get.cache_misses', key, 1)
        self.metrics.add_time('store.get.connector', key, ctime)

    if value is not None:
        with Timer() as deserializer_timer:
            if deserializer is not None:
                result = deserializer(value)
            else:
                result = self.deserializer(value)

        if self.metrics is not None:
            dtime = deserializer_timer.elapsed_ns
            obj_size = len(value)
            self.metrics.add_time('store.get.deserialize', key, dtime)
            self.metrics.add_attribute(
                'store.get.object_size',
                key,
                obj_size,
            )

        self.cache.set(key, result)
    else:
        result = default

    timer.stop()
    if self.metrics is not None:
        self.metrics.add_time('store.get', key, timer.elapsed_ns)

    logger.debug(
        f'Store(name="{self.name}"): GET {key} in '
        f'{timer.elapsed_ms:.3f} ms (cached=False)',
    )
    return result

is_cached() ¶

is_cached(key: ConnectorKeyT) -> bool

Check if an object associated with the key is cached locally.

Parameters:

key (ConnectorKeyT) –

Key potentially associated with stored object.

Returns:

bool –

If the object is cached.

Source code in proxystore/store/base.py

def is_cached(self, key: ConnectorKeyT) -> bool:
    """Check if an object associated with the key is cached locally.

    Args:
        key: Key potentially associated with stored object.

    Returns:
        If the object is cached.
    """
    return self.cache.exists(key)

proxy() ¶

proxy(
    obj: T | NonProxiableT,
    *,
    evict: bool = False,
    serializer: SerializerT | None = None,
    deserializer: DeserializerT | None = None,
    skip_nonproxiable: bool = False,
    **kwargs: Any
) -> Proxy[T] | NonProxiableT

Create a proxy that will resolve to an object in the store.

Parameters:

obj (T | NonProxiableT) –

The object to place in store and return a proxy for.
evict (bool, default: False ) –

If the proxy should evict the object once resolved.
serializer (SerializerT | None, default: None ) –

Optionally override the default serializer for the store instance.
deserializer (DeserializerT | None, default: None ) –

Optionally override the default deserializer for the store instance.
skip_nonproxiable (bool, default: False ) –

Return non-proxiable types (e.g., built-in constants like bool or None) rather than raising a NonProxiableTypeError.
kwargs (Any, default: {} ) –

Additional keyword arguments to pass to Connector.put().

Returns:

Proxy[T] | NonProxiableT –

A proxy of the object unless obj is a non-proxiable type skip_nonproxiable is True in which case obj is returned directly.

Raises:

NonProxiableTypeError –

If obj is a non-proxiable type. This behavior can be overridden by setting skip_nonproxiable=True.

Source code in proxystore/store/base.py

def proxy(
    self,
    obj: T | NonProxiableT,
    *,
    evict: bool = False,
    serializer: SerializerT | None = None,
    deserializer: DeserializerT | None = None,
    skip_nonproxiable: bool = False,
    **kwargs: Any,
) -> Proxy[T] | NonProxiableT:
    """Create a proxy that will resolve to an object in the store.

    Args:
        obj: The object to place in store and return a proxy for.
        evict: If the proxy should evict the object once resolved.
        serializer: Optionally override the default serializer for the
            store instance.
        deserializer: Optionally override the default deserializer for the
            store instance.
        skip_nonproxiable: Return non-proxiable types (e.g., built-in
            constants like `bool` or `None`) rather than raising a
            [`NonProxiableTypeError`][proxystore.store.exceptions.NonProxiableTypeError].
        kwargs: Additional keyword arguments to pass to
            [`Connector.put()`][proxystore.connectors.protocols.Connector.put].

    Returns:
        A proxy of the object unless `obj` is a non-proxiable type \
        `#!python skip_nonproxiable is True` in which case `obj` is \
        returned directly.

    Raises:
        NonProxiableTypeError: If `obj` is a non-proxiable type. This
            behavior can be overridden by setting
            `#!python skip_nonproxiable=True`.
    """
    if isinstance(obj, _NON_PROXIABLE_TYPES):
        if skip_nonproxiable:
            # MyPy raises the following error which is not correct:
            #     Incompatible return value type (got "Optional[bool]",
            #     expected "Optional[Proxy[T]]")  [return-value]
            return obj  # type: ignore[return-value]
        else:
            raise NonProxiableTypeError(
                f'Object of {type(obj)} is not proxiable.',
            )

    with Timer() as timer:
        key = self.put(obj, serializer=serializer, **kwargs)
        factory: StoreFactory[ConnectorT, T] = StoreFactory(
            key,
            store_config=self.config(),
            deserializer=deserializer,
            evict=evict,
        )
        proxy = Proxy(factory)

    if self.metrics is not None:
        self.metrics.add_time('store.proxy', key, timer.elapsed_ns)

    logger.debug(
        f'Store(name="{self.name}"): PROXY {key} in '
        f'{timer.elapsed_ms:.3f} ms',
    )
    return proxy

proxy_batch() ¶

proxy_batch(
    objs: Sequence[T | NonProxiableT],
    *,
    evict: bool = False,
    serializer: SerializerT | None = None,
    deserializer: DeserializerT | None = None,
    skip_nonproxiable: bool = False,
    **kwargs: Any
) -> list[Proxy[T] | NonProxiableT]

Create proxies that will resolve to an object in the store.

Parameters:

objs (Sequence[T | NonProxiableT]) –

The objects to place in store and return a proxies for.
evict (bool, default: False ) –

If a proxy should evict its object once resolved.
serializer (SerializerT | None, default: None ) –

Optionally override the default serializer for the store instance.
deserializer (DeserializerT | None, default: None ) –

Optionally override the default deserializer for the store instance.
skip_nonproxiable (bool, default: False ) –

Return non-proxiable types (e.g., built-in constants like bool or None) rather than raising a NonProxiableTypeError.
kwargs (Any, default: {} ) –

Additional keyword arguments to pass to Connector.put_batch().

Returns:

list[Proxy[T] | NonProxiableT] –

A list of proxies of each object or the object itself if said object is not proxiable and skip_nonproxiable is True.

Raises:

NonProxiableTypeError –

If obj is a non-proxiable type. This behavior can be overridden by setting skip_nonproxiable=True.

Source code in proxystore/store/base.py

def proxy_batch(  # type: ignore[misc]
    self,
    objs: Sequence[T | NonProxiableT],
    *,
    evict: bool = False,
    serializer: SerializerT | None = None,
    deserializer: DeserializerT | None = None,
    skip_nonproxiable: bool = False,
    **kwargs: Any,
) -> list[Proxy[T] | NonProxiableT]:
    """Create proxies that will resolve to an object in the store.

    Args:
        objs: The objects to place in store and return a proxies for.
        evict: If a proxy should evict its object once resolved.
        serializer: Optionally override the default serializer for the
            store instance.
        deserializer: Optionally override the default deserializer for the
            store instance.
        skip_nonproxiable: Return non-proxiable types (e.g., built-in
            constants like `bool` or `None`) rather than raising a
            [`NonProxiableTypeError`][proxystore.store.exceptions.NonProxiableTypeError].
        kwargs: Additional keyword arguments to pass to
            [`Connector.put_batch()`][proxystore.connectors.protocols.Connector.put_batch].

    Returns:
        A list of proxies of each object or the object itself if said \
        object is not proxiable and `#!python skip_nonproxiable is True`.

    Raises:
        NonProxiableTypeError: If `obj` is a non-proxiable type. This
            behavior can be overridden by setting
            `#!python skip_nonproxiable=True`.
    """
    with Timer() as timer:
        # Find if there are non-proxiable types and if that's okay
        non_proxiable: list[tuple[int, Any]] = []
        for i, obj in enumerate(objs):
            if isinstance(obj, _NON_PROXIABLE_TYPES):
                non_proxiable.append((i, obj))

        if len(non_proxiable) > 0 and not skip_nonproxiable:
            raise NonProxiableTypeError(
                f'Input sequence contains {len(non_proxiable)} '
                'objects that are not proxiable.',
            )

        # Pop non-proxiable types so we can batch proxy the proxiable ones
        non_proxiable_indicies = [i for i, _ in non_proxiable]
        proxiable_objs = [
            obj
            for i, obj in enumerate(objs)
            if i not in non_proxiable_indicies
        ]

        keys = self.put_batch(
            proxiable_objs,
            serializer=serializer,
            **kwargs,
        )
        proxies: list[Proxy[T]] = [
            self.proxy_from_key(
                key,
                evict=evict,
                deserializer=deserializer,
            )
            for key in keys
        ]

        # Put non-proxiable objects back in their original positions.
        # The indices of non_proxiable must still be sorted
        for original_index, original_object in non_proxiable:
            proxies.insert(original_index, original_object)

    if self.metrics is not None:
        self.metrics.add_time('store.proxy_batch', keys, timer.elapsed_ns)

    logger.debug(
        f'Store(name="{self.name}"): PROXY_BATCH ({len(proxies)} items) '
        f'in {timer.elapsed_ms:.3f} ms',
    )
    return cast(List[Union[Proxy[T], NonProxiableT]], proxies)

proxy_from_key() ¶

proxy_from_key(
    key: ConnectorKeyT,
    *,
    evict: bool = False,
    deserializer: DeserializerT | None = None
) -> Proxy[T]

Create a proxy that will resolve to an object already in the store.

Parameters:

key (ConnectorKeyT) –

The key associated with an object already in the store.
evict (bool, default: False ) –

If the proxy should evict the object once resolved.
deserializer (DeserializerT | None, default: None ) –

Optionally override the default deserializer for the store instance.

Returns:

Proxy[T] –

A proxy of the object.

Source code in proxystore/store/base.py

def proxy_from_key(
    self,
    key: ConnectorKeyT,
    *,
    evict: bool = False,
    deserializer: DeserializerT | None = None,
) -> Proxy[T]:
    """Create a proxy that will resolve to an object already in the store.

    Args:
        key: The key associated with an object already in the store.
        evict: If the proxy should evict the object once resolved.
        deserializer: Optionally override the default deserializer for the
            store instance.

    Returns:
        A proxy of the object.
    """
    factory: StoreFactory[ConnectorT, T] = StoreFactory(
        key,
        store_config=self.config(),
        deserializer=deserializer,
        evict=evict,
    )
    logger.debug(f'Store(name="{self.name}"): PROXY_FROM_KEY {key}')
    return Proxy(factory)

locked_proxy() ¶

locked_proxy(
    obj: T | NonProxiableT,
    *,
    evict: bool = False,
    serializer: SerializerT | None = None,
    deserializer: DeserializerT | None = None,
    skip_nonproxiable: bool = True,
    **kwargs: Any
) -> ProxyLocker[T] | NonProxiableT

Proxy an object and return ProxyLocker.

Parameters:

obj (T | NonProxiableT) –

The object to place in store and return a proxy for.
evict (bool, default: False ) –

If the proxy should evict the object once resolved.
serializer (SerializerT | None, default: None ) –

Optionally override the default serializer for the store instance.
deserializer (DeserializerT | None, default: None ) –

Optionally override the default deserializer for the store instance.
skip_nonproxiable (bool, default: True ) –

Return non-proxiable types (e.g., built-in constants like bool or None) rather than raising a NonProxiableTypeError.
kwargs (Any, default: {} ) –

Additional keyword arguments to pass to Connector.put().

Returns:

ProxyLocker[T] | NonProxiableT –

A proxy wrapped in a ProxyLocker unless obj is a non-proxiable type skip_nonproxiable is True in which case obj is returned directly.

Raises:

NonProxiableTypeError –

If obj is a non-proxiable type. This behavior can be overridden by setting skip_nonproxiable=True.

Source code in proxystore/store/base.py

def locked_proxy(
    self,
    obj: T | NonProxiableT,
    *,
    evict: bool = False,
    serializer: SerializerT | None = None,
    deserializer: DeserializerT | None = None,
    skip_nonproxiable: bool = True,
    **kwargs: Any,
) -> ProxyLocker[T] | NonProxiableT:
    """Proxy an object and return [`ProxyLocker`][proxystore.proxy.ProxyLocker].

    Args:
        obj: The object to place in store and return a proxy for.
        evict: If the proxy should evict the object once resolved.
        serializer: Optionally override the default serializer for the
            store instance.
        deserializer: Optionally override the default deserializer for the
            store instance.
        skip_nonproxiable: Return non-proxiable types (e.g., built-in
            constants like `bool` or `None`) rather than raising a
            [`NonProxiableTypeError`][proxystore.store.exceptions.NonProxiableTypeError].
        kwargs: Additional keyword arguments to pass to
            [`Connector.put()`][proxystore.connectors.protocols.Connector.put].

    Returns:
        A proxy wrapped in a \
        [`ProxyLocker`][proxystore.proxy.ProxyLocker] unless `obj` is a \
        non-proxiable type `#!python skip_nonproxiable is True` in which \
        case `obj` is returned directly.

    Raises:
        NonProxiableTypeError: If `obj` is a non-proxiable type. This
            behavior can be overridden by setting
            `#!python skip_nonproxiable=True`.
    """
    possible_proxy = self.proxy(
        obj,
        evict=evict,
        serializer=serializer,
        deserializer=deserializer,
        skip_nonproxiable=skip_nonproxiable,
        **kwargs,
    )

    if isinstance(possible_proxy, Proxy):
        return ProxyLocker(possible_proxy)
    return possible_proxy

put() ¶

put(
    obj: Any,
    *,
    serializer: SerializerT | None = None,
    **kwargs: Any
) -> ConnectorKeyT

Put an object in the store.

Parameters:

obj (Any) –

Object to put in the store.
serializer (SerializerT | None, default: None ) –

Optionally override the default serializer for the store instance.
kwargs (Any, default: {} ) –

Additional keyword arguments to pass to Connector.put().

Returns:

ConnectorKeyT –

A key which can be used to retrieve the object.

Raises:

TypeError –

If the output of serializer is not bytes.

Source code in proxystore/store/base.py

def put(
    self,
    obj: Any,
    *,
    serializer: SerializerT | None = None,
    **kwargs: Any,
) -> ConnectorKeyT:
    """Put an object in the store.

    Args:
        obj: Object to put in the store.
        serializer: Optionally override the default serializer for the
            store instance.
        kwargs: Additional keyword arguments to pass to
            [`Connector.put()`][proxystore.connectors.protocols.Connector.put].

    Returns:
        A key which can be used to retrieve the object.

    Raises:
        TypeError: If the output of `serializer` is not bytes.
    """
    timer = Timer()
    timer.start()

    with Timer() as serialize_timer:
        if serializer is not None:
            obj = serializer(obj)
        else:
            obj = self.serializer(obj)

    if not isinstance(obj, bytes):
        raise TypeError('Serializer must produce bytes.')

    with Timer() as connector_timer:
        key = self.connector.put(obj, **kwargs)

    timer.stop()
    if self.metrics is not None:
        ctime = connector_timer.elapsed_ns
        stime = serialize_timer.elapsed_ns
        self.metrics.add_attribute('store.put.object_size', key, len(obj))
        self.metrics.add_time('store.put.serialize', key, stime)
        self.metrics.add_time('store.put.connector', key, ctime)
        self.metrics.add_time('store.put', key, timer.elapsed_ns)

    logger.debug(
        f'Store(name="{self.name}"): PUT {key} in '
        f'{timer.elapsed_ms:.3f} ms',
    )
    return key

put_batch() ¶

put_batch(
    objs: Sequence[Any],
    *,
    serializer: SerializerT | None = None,
    **kwargs: Any
) -> list[ConnectorKeyT]

Put multiple objects in the store.

Parameters:

objs (Sequence[Any]) –

Sequence of objects to put in the store.
serializer (SerializerT | None, default: None ) –

Optionally override the default serializer for the store instance.
kwargs (Any, default: {} ) –

Additional keyword arguments to pass to Connector.put_batch().

Returns:

list[ConnectorKeyT] –

A list of keys which can be used to retrieve the objects.

Raises:

TypeError –

If the output of serializer is not bytes.

Source code in proxystore/store/base.py

def put_batch(
    self,
    objs: Sequence[Any],
    *,
    serializer: SerializerT | None = None,
    **kwargs: Any,
) -> list[ConnectorKeyT]:
    """Put multiple objects in the store.

    Args:
        objs: Sequence of objects to put in the store.
        serializer: Optionally override the default serializer for the
            store instance.
        kwargs: Additional keyword arguments to pass to
            [`Connector.put_batch()`][proxystore.connectors.protocols.Connector.put_batch].

    Returns:
        A list of keys which can be used to retrieve the objects.

    Raises:
        TypeError: If the output of `serializer` is not bytes.
    """
    timer = Timer()
    timer.start()

    def _serialize(obj: Any) -> bytes:
        if serializer is not None:
            obj = serializer(obj)
        else:
            obj = self.serializer(obj)

        if not isinstance(obj, bytes):
            raise TypeError('Serializer must produce bytes.')

        return obj

    with Timer() as serialize_timer:
        _objs = list(map(_serialize, objs))

    with Timer() as connector_timer:
        keys = self.connector.put_batch(_objs, **kwargs)

    timer.stop()
    if self.metrics is not None:
        ctime = connector_timer.elapsed_ns
        stime = serialize_timer.elapsed_ns
        sizes = sum(len(obj) for obj in _objs)
        self.metrics.add_attribute(
            'store.put_batch.object_sizes',
            keys,
            sizes,
        )
        self.metrics.add_time('store.put_batch.serialize', keys, stime)
        self.metrics.add_time('store.put_batch.connector', keys, ctime)
        self.metrics.add_time('store.put_batch', keys, timer.elapsed_ns)

    logger.debug(
        f'Store(name="{self.name}"): PUT_BATCH ({len(keys)} items) in '
        f'{timer.elapsed_ms:.3f} ms',
    )
    return keys

StoreFactory ¶

StoreFactory(
    key: ConnectorKeyT,
    store_config: dict[str, Any],
    *,
    evict: bool = False,
    deserializer: DeserializerT | None = None
)

Bases: Generic[ConnectorT, T]

Factory that resolves an object from a store.

Adds support for asynchronously retrieving objects from a Store instance.

The factory takes the store_config parameter that is used to reinitialize the store if the factory is sent to a remote process where the store has not already been initialized.

Parameters:

key (ConnectorKeyT) –

Key corresponding to object in store.
store_config (dict[str, Any]) –

Store configuration used to reinitialize the store if needed.
evict (bool, default: False ) –

If True, evict the object from the store once resolve() is called.
deserializer (DeserializerT | None, default: None ) –

Optional callable used to deserialize the byte string. If None, the default deserializer (deserialize()) will be used.

Source code in proxystore/store/factory.py

def __init__(
    self,
    key: ConnectorKeyT,
    store_config: dict[str, Any],
    *,
    evict: bool = False,
    deserializer: DeserializerT | None = None,
) -> None:
    self.key = key
    self.store_config = store_config
    self.evict = evict
    self.deserializer = deserializer

    # The following are not included when a factory is serialized
    # because they are specific to that instance of the factory
    self._obj_future: Future[T] | None = None

get_store() ¶

get_store() -> Store[ConnectorT]

Get store and reinitialize if necessary.

Raises:

ValueError –

If the type of the returned store does not match the expected store type passed to the factory constructor.

Source code in proxystore/store/factory.py

def get_store(self) -> Store[ConnectorT]:
    """Get store and reinitialize if necessary.

    Raises:
        ValueError: If the type of the returned store does not match the
            expected store type passed to the factory constructor.
    """
    with _factory_get_store_lock:
        store = proxystore.store.get_store(self.store_config['name'])
        if store is None:
            store = proxystore.store.Store.from_config(self.store_config)
            proxystore.store.register_store(store)
        return store

resolve() ¶

resolve() -> T

Get object associated with key from store.

Raises:

ProxyResolveMissingKeyError –

If the key associated with this factory does not exist in the store.

Source code in proxystore/store/factory.py

def resolve(self) -> T:
    """Get object associated with key from store.

    Raises:
        ProxyResolveMissingKeyError: If the key associated with this
            factory does not exist in the store.
    """
    with Timer() as timer:
        store = self.get_store()
        obj = store.get(
            self.key,
            deserializer=self.deserializer,
            default=_MISSING_OBJECT,
        )

        if obj is _MISSING_OBJECT:
            raise ProxyResolveMissingKeyError(
                self.key,
                type(store),
                store.name,
            )

        if self.evict:
            store.evict(self.key)

    if store.metrics is not None:
        total_time = timer.elapsed_ns
        store.metrics.add_time('factory.resolve', self.key, total_time)

    return cast(T, obj)

resolve_async() ¶

resolve_async() -> None

Asynchronously get object associated with key from store.

Source code in proxystore/store/factory.py

def resolve_async(self) -> None:
    """Asynchronously get object associated with key from store."""
    logger.debug(f'Starting asynchronous resolve of {self.key}')
    self._obj_future = _default_pool.submit(self.resolve)

get_store() ¶

get_store(val: str | Proxy[T]) -> Store[Any] | None

Get the backend store with name.

Parameters:

val (str | Proxy[T]) –

name of the store to get or a Proxy instance.

Returns:

Store[Any] | None –

Store if a store matching the name or belonging to the proxy exists. If the store does not exist, returns None.

Raises:

ProxyStoreFactoryError –

If the value is a proxy but does not contain a factory of type StoreFactory.

Source code in proxystore/store/__init__.py

def get_store(val: str | Proxy[T]) -> Store[Any] | None:
    """Get the backend store with name.

    Args:
        val: name of the store to get or a [`Proxy`][proxystore.proxy.Proxy]
            instance.

    Returns:
        [`Store`][proxystore.store.base.Store] if a store matching the \
        name or belonging to the proxy exists. If the store does not exist, \
        returns `None`.

    Raises:
        ProxyStoreFactoryError: If the value is a proxy but does not contain a
            factory of type
            [`StoreFactory`][proxystore.store.base.StoreFactory].
    """
    if isinstance(val, Proxy):
        # If the object is a proxy, get the factory that will access the store
        factory = val.__factory__
        if isinstance(factory, StoreFactory):
            return factory.get_store()
        else:
            raise ProxyStoreFactoryError(
                'The proxy must contain a factory with type '
                f'{type(StoreFactory).__name__}. {type(factory).__name__} '
                'is not supported.',
            )
    else:
        name = val

    if name in _stores:
        return _stores[name]
    return None

register_store() ¶

register_store(
    store: Store[Any], exist_ok: bool = False
) -> None

Register the store instance to the global registry.

Note

Global means globally accessible within the Python process.

Tip

Use the store_registration context manager to automatically register and unregister as store.

Parameters:

store (Store[Any]) –

Store instance to register.
exist_ok (bool, default: False ) –

If a store with the same name exists, overwrite it.

Raises:

StoreExistsError –

If a store with the same name is already registered and exist_ok is false.

Source code in proxystore/store/__init__.py

def register_store(store: Store[Any], exist_ok: bool = False) -> None:
    """Register the store instance to the global registry.

    Note:
        Global means globally accessible within the Python process.

    Tip:
        Use the [`store_registration`][proxystore.store.store_registration]
        context manager to automatically register and unregister as store.

    Args:
        store: Store instance to register.
        exist_ok: If a store with the same name exists, overwrite it.

    Raises:
        StoreExistsError: If a store with the same name is already registered
            and `exist_ok` is false.
    """
    if store.name in _stores and not exist_ok:
        raise StoreExistsError(f'A store named {store.name} already exists.')

    _stores[store.name] = store
    logger.info(f'Registered a store named {store.name}')

store_registration() ¶

store_registration(
    *stores: Store[Any], exist_ok: bool = False
) -> Generator[None, None, None]

Context manager that registers and unregisters a set of stores.

Example

from proxystore.connectors.local import LocalConnector
from proxystore.store import Store
from proxystore.store import store_registration

with Store('store', LocalConnector()) as store:
    with store_registration(store):
        ...

stores = [
    Store('store1', LocalConnector()),
    Store('store2', LocalConnector()),
]
with store_registration(*stores):
    ...

Parameters:

stores (Store[Any], default: () ) –

Set of Store instances to register then unregister when the context manager is exited.
exist_ok (bool, default: False ) –

If a store with the same name exists, overwrite it.

Raises:

StoreExistsError –

If a store with the same name is already registered and exist_ok is false.

Source code in proxystore/store/__init__.py

@contextlib.contextmanager
def store_registration(
    *stores: Store[Any],
    exist_ok: bool = False,
) -> Generator[None, None, None]:
    """Context manager that registers and unregisters a set of stores.

    Example:
        ```python
        from proxystore.connectors.local import LocalConnector
        from proxystore.store import Store
        from proxystore.store import store_registration

        with Store('store', LocalConnector()) as store:
            with store_registration(store):
                ...

        stores = [
            Store('store1', LocalConnector()),
            Store('store2', LocalConnector()),
        ]
        with store_registration(*stores):
            ...
        ```

    Args:
        stores: Set of [`Store`][proxystore.store.base.Store] instances to
            register then unregister when the context manager is exited.
        exist_ok: If a store with the same name exists, overwrite it.

    Raises:
        StoreExistsError: If a store with the same name is already registered
            and `exist_ok` is false.
    """
    for store in stores:
        register_store(store, exist_ok=exist_ok)

    yield

    for store in stores:
        unregister_store(store)

unregister_store() ¶

unregister_store(name_or_store: str | Store[Any]) -> None

Unregisters the store instance from the global registry.

Note

This function is a no-op if no store matching the name exists (i.e., no exception will be raised).

Parameters:

name_or_store (str | Store[Any]) –

Name of the store to unregister or a store itself.

Source code in proxystore/store/__init__.py

def unregister_store(name_or_store: str | Store[Any]) -> None:
    """Unregisters the store instance from the global registry.

    Note:
        This function is a no-op if no store matching the name
        exists (i.e., no exception will be raised).

    Args:
        name_or_store: Name of the store to unregister or a store itself.
    """
    name = (
        name_or_store if isinstance(name_or_store, str) else name_or_store.name
    )
    if name in _stores:
        del _stores[name]
        logger.debug(f'Unregistered a store named {name}')

proxystore.store¶

Store ¶

name property ¶

metrics property ¶

serializer property ¶

deserializer property ¶

close() ¶

config() ¶

from_config() classmethod ¶

future() ¶

evict() ¶

exists() ¶

get() ¶

is_cached() ¶

proxy() ¶

proxy_batch() ¶

proxy_from_key() ¶

locked_proxy() ¶

put() ¶

put_batch() ¶

StoreFactory ¶

get_store() ¶

resolve() ¶

resolve_async() ¶

get_store() ¶

register_store() ¶

store_registration() ¶

unregister_store() ¶

name `property` ¶

metrics `property` ¶

serializer `property` ¶

deserializer `property` ¶

from_config() `classmethod` ¶