proxystore.connectors.endpoint¶

proxystore/connectors/endpoint.py

Endpoint connector implementation.

EndpointConnectorError ¶

Bases: Exception

Exception raised when the endpoint connector fails.

This is raised when the connector cannot connect to any of its endpoints or when a request fails. The error that caused the failure (e.g., an endpoint was not running or a peer endpoint was not allowed) is the __cause__ of the exception. The types of the causes are internal to proxystore.endpoint so only the message should be relied on.

EndpointKey ¶

Bases: NamedTuple

Key to object in an Endpoint.

Attributes:

object_id (str) –

Unique object ID.
endpoint_id (str) –

ID of the endpoint where object is stored.

EndpointConnector ¶

EndpointConnector(
    endpoints: Sequence[str] | None = None,
    proxystore_dir: str | PathLike[str] | None = None,
    reconnect_timeout: float = 5,
    request_timeout: float | None = REQUEST_TIMEOUT,
)

Connector to ProxyStore Endpoints.

Warning

Specifying a custom proxystore_dir can cause problems if the proxystore_dir is not the same on all systems that a proxy created by this store could end up on. It is recommended to leave the proxystore_dir unspecified so the correct default directory will be used.

Parameters:

endpoints (Sequence[str] | None, default: None ) –

IDs of the endpoints which the connector may use as its home endpoint, the local endpoint which all requests are sent to. The connector uses the first of these endpoints found in proxystore_dir which it can connect to. This is not an allowlist: requests for objects stored on other endpoints are forwarded by the home endpoint, and which peers an endpoint communicates with is controlled by the allowlist of peers of each endpoint. The list is part of the connector's configuration so a proxy resolved on another system will use whichever of these endpoints is on that system. If None, the connector uses the only endpoint running in proxystore_dir.
proxystore_dir (str | PathLike[str] | None, default: None ) –

Optionally specify the proxystore home directory. Defaults to home_dir().
reconnect_timeout (float, default: 5 ) –

Seconds to keep trying to reconnect to the endpoint if it is unavailable (e.g., because it is restarting) before a request fails.
request_timeout (float | None, default: REQUEST_TIMEOUT ) –

Seconds a request to the local endpoint can go without sending or receiving any data, including while the endpoint handles the request, before it fails. Requests for objects on peer endpoints have no timeout in the client because the endpoint only responds once the peer does. If None, requests have no timeout.

Raises:

ValueError –

If endpoints is an empty list or contains an invalid endpoint ID, or if request_timeout is not positive.
EndpointConnectorError –

If unable to connect to any of the endpoints provided, or if endpoints is None and there is not exactly one endpoint running in proxystore_dir.

Source code in proxystore/connectors/endpoint.py

def __init__(
    self,
    endpoints: Sequence[str] | None = None,
    proxystore_dir: str | os.PathLike[str] | None = None,
    reconnect_timeout: float = 5,
    request_timeout: float | None = REQUEST_TIMEOUT,
) -> None:
    if endpoints is not None and len(endpoints) == 0:
        raise ValueError('At least one endpoint must be specified.')
    self.endpoints: list[EndpointId] | None = (
        None
        if endpoints is None
        else [EndpointId.from_str(e) for e in endpoints]
    )
    # Stored as a str so the config is serializable (e.g., to TOML).
    self.proxystore_dir = (
        None if proxystore_dir is None else os.fspath(proxystore_dir)
    )
    self.reconnect_timeout = reconnect_timeout
    self.request_timeout = check_request_timeout(request_timeout)

    home = (
        home_dir() if self.proxystore_dir is None else self.proxystore_dir
    )
    if self.endpoints is None:
        candidates = _find_running(home)
    else:
        candidates = [
            (endpoint_dir, endpoint)
            for endpoint_dir, endpoint in EndpointDir.find_all(home)
            if endpoint.id in self.endpoints
        ]

    # Use the first locally accessible endpoint as our home endpoint
    failures: list[str] = []
    found: tuple[EndpointId, EndpointDir, EndpointClient] | None = None
    for endpoint_dir, endpoint in candidates:
        endpoint_id = endpoint.id
        logger.debug('Attempting connection to %s', endpoint_id)
        try:
            client = _connect(endpoint_dir, endpoint_id, request_timeout)
        except EndpointError as e:
            logger.debug('Connection to %s failed: %r', endpoint_id, e)
            failures.append(f'{endpoint.name} ({endpoint_id}): {e}')
            continue

        logger.debug(
            'Connection to %s successful, using as local endpoint',
            endpoint_id,
        )
        found = (endpoint_id, endpoint_dir, client)
        break

    if found is None:
        if len(failures) == 0:
            raise EndpointConnectorError(
                'Failed to find an endpoint configuration in '
                f'{home} matching one of the provided endpoint IDs.',
            )
        reasons = '\n'.join(f'  - {failure}' for failure in failures)
        raise EndpointConnectorError(
            'Failed to connect to any of the endpoints matching the '
            f'provided endpoint IDs:\n{reasons}',
        )
    endpoint_id, endpoint_dir, client = found
    self.endpoint_id: EndpointId = endpoint_id
    self.endpoint_dir = endpoint_dir

    self._pool = _ConnectionPool(
        lambda: _connect(endpoint_dir, endpoint_id, request_timeout),
        reconnect_timeout=reconnect_timeout,
    )
    self._pool.add(client)

close ¶

close(*, clear: bool | None = None) -> bool

Close the connector and clean up.

Parameters:

clear (bool | None, default: None ) –

Ignored because objects are owned by the endpoints.

Returns:

bool –

False because objects are never removed.

Source code in proxystore/connectors/endpoint.py

def close(self, *, clear: bool | None = None) -> bool:
    """Close the connector and clean up.

    Args:
        clear: Ignored because objects are owned by the endpoints.

    Returns:
        `False` because objects are never removed.
    """
    self._pool.close()
    return False

config ¶

config() -> dict[str, Any]

Get the connector configuration.

The configuration contains all the information needed to reconstruct the connector object.

Source code in proxystore/connectors/endpoint.py

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

    The configuration contains all the information needed to reconstruct
    the connector object.
    """
    return {
        'endpoints': (
            None if self.endpoints is None else list(self.endpoints)
        ),
        'proxystore_dir': self.proxystore_dir,
        'reconnect_timeout': self.reconnect_timeout,
        'request_timeout': self.request_timeout,
    }

from_config `classmethod` ¶

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

Create a new connector instance from a configuration.

Parameters:

config (dict[str, Any]) –

Configuration returned by .config().

Source code in proxystore/connectors/endpoint.py

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

    Args:
        config: Configuration returned by `#!python .config()`.
    """
    return cls(**init_kwargs(cls, config))

evict ¶

evict(key: EndpointKey) -> None

Evict the object associated with the key.

Parameters:

key (EndpointKey) –

Key associated with object to evict.

Source code in proxystore/connectors/endpoint.py

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

    Args:
        key: Key associated with object to evict.
    """
    self._request(
        'Evict',
        lambda client: client.evict(key.object_id, key.endpoint_id),
    )

exists ¶

exists(key: EndpointKey) -> bool

Check if an object associated with the key exists.

Parameters:

key (EndpointKey) –

Key potentially associated with stored object.

Returns:

bool –

If an object associated with the key exists.

Source code in proxystore/connectors/endpoint.py

def exists(self, key: EndpointKey) -> 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.
    """
    return self._request(
        'Exists',
        lambda client: client.exists(key.object_id, key.endpoint_id),
    )

get ¶

get(key: EndpointKey) -> BytesLike | None

Get the serialized object associated with the key.

Parameters:

key (EndpointKey) –

Key associated with the object to retrieve.

Returns:

BytesLike | None –

Serialized object or None if the object does not exist.

Source code in proxystore/connectors/endpoint.py

def get(self, key: EndpointKey) -> BytesLike | None:
    """Get the serialized object associated with the key.

    Args:
        key: Key associated with the object to retrieve.

    Returns:
        Serialized object or `None` if the object does not exist.
    """
    return self._request(
        'Get',
        lambda client: client.get(key.object_id, key.endpoint_id),
    )

get_batch ¶

get_batch(
    keys: Sequence[EndpointKey],
) -> list[BytesLike | None]

Get a batch of serialized objects associated with the keys.

Parameters:

keys (Sequence[EndpointKey]) –

Sequence of keys associated with objects to retrieve.

Note

Objects on peer endpoints are fetched at the same time (up to eight at a time, each with its own connection to the local endpoint) because each request waits for the peer to respond. Objects on the local endpoint are fetched one at a time while the peer requests run, because the local endpoint is not faster when sent requests at the same time.

Returns:

list[BytesLike | None] –

List with same order as keys with the serialized objects or None if the corresponding key does not have an associated object.

Source code in proxystore/connectors/endpoint.py

def get_batch(self, keys: Sequence[EndpointKey]) -> list[BytesLike | None]:
    """Get a batch of serialized objects associated with the keys.

    Args:
        keys: Sequence of keys associated with objects to retrieve.

    Note:
        Objects on peer endpoints are fetched at the same time (up to
        eight at a time, each with its own connection to the local
        endpoint) because each request waits for the peer to respond.
        Objects on the local endpoint are fetched one at a time while
        the peer requests run, because the local endpoint is not faster
        when sent requests at the same time.

    Returns:
        List with same order as `keys` with the serialized objects or \
        `None` if the corresponding key does not have an associated object.
    """
    peer = [
        i
        for i, key in enumerate(keys)
        if key.endpoint_id != self.endpoint_id
    ]
    if len(peer) == 0:
        return [self.get(key) for key in keys]

    executor = self._pool.executor()
    futures = {i: executor.submit(self.get, keys[i]) for i in peer}
    try:
        objs = [
            None if i in futures else self.get(key)
            for i, key in enumerate(keys)
        ]
        for i, future in futures.items():
            objs[i] = future.result()
    finally:
        # Requests which have not started are not needed if one failed.
        for future in futures.values():
            future.cancel()
    return objs

new_key ¶

new_key(obj: BytesLike | None = None) -> EndpointKey

Create a new key.

Warning

The returned key will be associated with this instance's local endpoint. I.e., when set() is called on this key, the connector must be connected to the same local endpoint.

Parameters:

obj (BytesLike | None, default: None ) –

Optional object which the key will be associated with. Ignored in this implementation.

Returns:

EndpointKey –

Key which can be used to retrieve an object once set() has been called on the key.

Source code in proxystore/connectors/endpoint.py

def new_key(self, obj: BytesLike | None = None) -> EndpointKey:
    """Create a new key.

    Warning:
        The returned key will be associated with this instance's local
        endpoint. I.e., when
        [`set()`][proxystore.connectors.endpoint.EndpointConnector.set]
        is called on this key, the connector must be connected to the same
        local endpoint.

    Args:
        obj: Optional object which the key will be associated with.
            Ignored in this implementation.

    Returns:
        Key which can be used to retrieve an object once \
        [`set()`][proxystore.connectors.endpoint.EndpointConnector.set] \
        has been called on the key.
    """
    return EndpointKey(
        object_id=str(uuid.uuid4()),
        endpoint_id=self.endpoint_id,
    )

put ¶

put(obj: BytesLike) -> EndpointKey

Put a serialized object in the store.

Parameters:

obj (BytesLike) –

Serialized object to put in the store.

Returns:

EndpointKey –

Key which can be used to retrieve the object.

Source code in proxystore/connectors/endpoint.py

def put(self, obj: BytesLike) -> EndpointKey:
    """Put a serialized object in the store.

    Args:
        obj: Serialized object to put in the store.

    Returns:
        Key which can be used to retrieve the object.
    """
    key = EndpointKey(
        object_id=str(uuid.uuid4()),
        endpoint_id=self.endpoint_id,
    )
    self.set(key, obj)
    return key

put_batch ¶

put_batch(objs: Sequence[BytesLike]) -> list[EndpointKey]

Put a batch of serialized objects in the store.

Parameters:

objs (Sequence[BytesLike]) –

Sequence of serialized objects to put in the store.

Returns:

list[EndpointKey] –

List of keys with the same order as objs which can be used to retrieve the objects.

Source code in proxystore/connectors/endpoint.py

def put_batch(self, objs: Sequence[BytesLike]) -> list[EndpointKey]:
    """Put a batch of serialized objects in the store.

    Args:
        objs: Sequence of serialized objects to put in the store.

    Returns:
        List of keys with the same order as `objs` which can be used to \
        retrieve the objects.
    """
    return [self.put(obj) for obj in objs]

set ¶

set(key: EndpointKey, obj: BytesLike) -> None

Set the object associated with a key.

Note

The Connector provides write-once, read-many semantics. Thus, set() should only be called once per key, otherwise unexpected behavior can occur.

Parameters:

key (EndpointKey) –

Key that the object will be associated with.
obj (BytesLike) –

Object to associate with the key.

Source code in proxystore/connectors/endpoint.py

def set(self, key: EndpointKey, obj: BytesLike) -> None:
    """Set the object associated with a key.

    Note:
        The [`Connector`][proxystore.connectors.protocols.Connector]
        provides write-once, read-many semantics. Thus,
        [`set()`][proxystore.connectors.endpoint.EndpointConnector.set]
        should only be called once per key, otherwise unexpected behavior
        can occur.

    Args:
        key: Key that the object will be associated with.
        obj: Object to associate with the key.
    """
    self._request(
        'Set',
        lambda client: client.set(key.object_id, obj, key.endpoint_id),
    )

proxystore.connectors.endpoint¶

EndpointConnectorError ¶

EndpointKey ¶

EndpointConnector ¶

close ¶

config ¶

from_config classmethod ¶

evict ¶

exists ¶

get ¶

get_batch ¶

new_key ¶

put ¶

put_batch ¶

set ¶

from_config `classmethod` ¶