Store

A Store is initialized with a Connector instance and provides extra functionality. Similar to the Connector, the Store exposes evict, exists, get, and put operations; however, these operations act on Python objects rather than bytes. The Store will (de)serialize objects accordingly before invoking the corresponding operation on the Connector. The Store also provides caching of operations to reduce communication costs, and objects are cached after deserialization to avoid duplicate deserialization.

However, instead of the application directly invoking these aforementioned operations, the proxy method, also provided by the Store, is used. Calling Store.proxy() puts an object in the mediated channel and returns a proxy (see example below). The object is serialized before being put in the mediated channel, a factory with the key returned by the Connector and other information necessary to retrieve the object from the mediated channel is generated, and then a new proxy, initialized with the factory, is returned.

from proxystore.connectors.local import LocalConnector
from proxystore.proxy import Proxy
from proxystore.store import Store


def process(x: dict[str, str]) -> None:
    assert isinstance(x, dict)  # (1)!
    assert x['hello'] == 'world'
    # More computation using x...


with Store(LocalConnector(), name='example') as store:  # (2)!
    x = {'hello': 'world'}
    proxy = store.proxy(x)  # (3)!
    assert isinstance(proxy, Proxy)
    process(proxy)  # (4)!

1. x is resolved from the store on the first use of x and x then acts as a transparent reference to the target object (the dictionary).
2. The Connector defines the low-level communication method used by the Store. Here, the LocalConnector stores data in the processes memory, but other connectors are provided to remote storage and transfer services. Every store is registered globally within the process by its unique ID, so proxies resolved in the same process reuse the same store instance to improve performance. Stores made when resolving proxies are closed when the process exits, without clearing any objects. The optional name is used in logs.
3. Store the object and get a proxy.
4. Always succeeds regardless of if proxy is the true object or a proxy.

Asynchronous Resolving

It is common in distributed computation for inputs to functions executed remotely to not be needed immediately upon execution. Proxies created by a Store support asynchronous resolution to overlap communication and computation.

from proxystore.store.utils import resolve_async


def complex_function(large_proxied_input):
    resolve_async(large_proxied_input)

    # More computation...

    # First access to the proxy will not be as expensive because
    # of the asynchronous resolution
    compute_input(large_proxied_input)

Caching

The Store provides built in caching functionality. Caches are local to the Python process but will speed up the resolution when multiple proxies refer to the same object.

from proxystore.store import Store

# 256 MiB is the default
Store(..., cache_size_bytes=256 * 1024 * 1024)

The cache size is in bytes, and each entry counts as the size of the object's serialized data. When the cache is full, the least recently used objects are removed, and an object larger than the cache is not cached. Set cache_size_bytes=0 to turn off the cache. The cache is per process, so with many worker processes on a node, the cache can use up to cache_size_bytes in each of them.

With the default cache mode (described below), the serialized size is an estimate of the memory used by the cached object. It is close for objects like bytes and NumPy arrays, but Python objects such as dicts and lists can use several times more memory than their serialized data.

By default, the cache holds deserialized objects, so a cache hit is free, but every get() and every proxy of the same key in a process share one object. Changing that object (e.g., appending to a list) changes it for all of them. Set cache_mode='bytes' to cache the serialized data instead. Each cache hit then deserializes the data again, so each caller gets its own object, at the cost of deserializing on every hit.

from proxystore.store import Store

with Store(..., cache_mode='bytes') as store:
    key = store.put([1, 2, 3])
    a = store.get(key)
    b = store.get(key)  # (1)!
    a.append(4)
    assert b == [1, 2, 3]

1. Comes from the cache, but b is a new object.

The cache mode is part of the store's config(), so proxies resolved in other processes use the same mode. With cache_mode='bytes', callers only get separate objects if the deserializer makes new objects from the data. The default deserializer does, but a custom deserializer that returns the data itself, or views of it (e.g., numpy.frombuffer()), shares that data between callers.

Transactional Guarantees

ProxyStore is designed around optimizing the communication of ephemeral data (e.g., inputs and outputs of functions) which is typically write-once, read-many. Thus, ProxyStore does not provide update semantics on keys.

Serialization

All Store operations use ProxyStore's provided serialization utilities (proxystore.serialize) by default. However, the Store can be initialized with a custom serializer and deserializer of the form:

serializer = Callable[[Any], BytesLike]
deserializer = Callable[[BytesLike], Any]

Implementing a custom serializer may be beneficial for complex structures where pickle/cloudpickle (the default serializers used by ProxyStore) are inefficient. E.g.,

import torch
import io

from proxystore.serialize import serialize, deserialize, SerializationError
from proxystore.store import Store


def serialize_torch_model(obj: Any) -> bytes:
    if isinstance(obj, torch.nn.Module):
        buffer = io.BytesIO()
        buffer.write(b'PT\n')
        torch.save(obj, buffer)
        return buffer.getvalue()
    else:
        # Fallback for unsupported types
        return serialize(obj)


def deserialize_torch_model(raw: BytesLike) -> Any:
    try:
        return deserialize(raw)
    except SerializationError:
        buffer = io.BytesIO(raw)
        assert buffer.readline() == b'PT\n'
        return torch.load(buffer, weights_only=False)


model = torch.nn.Module()

with Store(
    ...,
    serializer=serialize_torch_model,
    deserializer=deserialize_torch_model,
) as store:
    key = store.put(model)
    model = store.get(key)

Tip

In some cases, data may already be serialized in which case an identity function can be passed as the serializer. Use an importable function rather than a lambda so the store's config can be pickled with its proxies. However, populate_target=False should also be set in this case to avoid prepopulating the proxy with the serialized target object. See the Store docstring for more information.

Every operation of a Store, and every proxy it creates, uses the serializer and deserializer of the store. To store some types in a custom format, handle those types in the serializer and fall back to serialize() for the rest, as in the example above. The serializer and deserializer are part of the config() of the store, so they must be importable wherever proxies of the store are resolved. See Issue #146 for further discussion on where custom serializers can be helpful.

Security

The default deserializer uses pickle, and unpickling data can run arbitrary code. So anyone who can write to the storage behind a store's connector (e.g., the directory of a FileConnector or the Redis server of a RedisConnector) can run code in every process that gets objects from the store or resolves its proxies. Treat write access to a connector's storage like write access to your code, and see Connector Security for what protects each connector.

Proxies and the config() of a store are also trusted. A store config names the connector, serializer, and deserializer to import and create, so only resolve proxies, and only create stores from configs, from sources you trust (see Proxy Security).