proxystore.store.factory

proxystore/store/factory.py

Factory implementations.

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)

PollingStoreFactory

PollingStoreFactory(
    key: ConnectorKeyT,
    store_config: dict[str, Any],
    *,
    deserializer: DeserializerT | None = None,
    evict: bool = False,
    polling_interval: float = 1,
    polling_timeout: float | None = None
)

Bases: StoreFactory[ConnectorT, T]

Factory that polls a store until and object can be resolved.

This is an extension of the StoreFactory with the resolve() method overridden to poll the store until the target object is available.

Parameters:

• key (ConnectorKeyT) –

  Key corresponding to object in store.

• store_config (dict[str, Any]) –

  Store configuration used to reinitialize the store if needed.

• deserializer (DeserializerT | None, default: None ) –

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

• evict (bool, default: False ) –

  If True, evict the object from the store once resolve() is called.

• polling_interval (float, default: 1 ) –

  Seconds to sleep between polling the store for the object.

• polling_timeout (float | None, default: None ) –

  Optional maximum number of seconds to poll for.

Source code in proxystore/store/factory.py

def __init__(
    self,
    key: ConnectorKeyT,
    store_config: dict[str, Any],
    *,
    deserializer: DeserializerT | None = None,
    evict: bool = False,
    polling_interval: float = 1,
    polling_timeout: float | None = None,
) -> None:
    super().__init__(
        key,
        store_config,
        evict=evict,
        deserializer=deserializer,
    )
    self._polling_interval = polling_interval
    self._polling_timeout = polling_timeout

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_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)

resolve()

resolve() -> T

Get object associated with key from store.

Raises:

• ProxyResolveMissingKeyError –

  If the object associated with the key is not available after polling_timeout seconds.

Source code in proxystore/store/factory.py

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

    Raises:
        ProxyResolveMissingKeyError: If the object associated with the
            key is not available after `polling_timeout` seconds.
    """
    with Timer() as timer:
        store = self.get_store()
        time_waited = 0.0

        while True:
            obj = store.get(
                self.key,
                deserializer=self.deserializer,
                default=_MISSING_OBJECT,
            )

            # Break because we found the object or we hit the timeout
            if obj is not _MISSING_OBJECT or (
                self._polling_timeout is not None
                and time_waited >= self._polling_timeout
            ):
                break

            time.sleep(self._polling_interval)
            time_waited += self._polling_interval

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

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

    return cast(T, obj)