proxystore.endpoint.server

proxystore/endpoint/server.py

Server that handles client connections to an endpoint.

Clients connect to their local endpoint over TCP using the protocol defined in proxystore.endpoint.protocol. The ClientHandler authenticates clients and forwards their requests to an Endpoint.

HANDSHAKE_TIMEOUT module-attribute

HANDSHAKE_TIMEOUT = 10

Seconds a client has to complete the handshake after connecting.

ClientHandler

ClientHandler(
    endpoint: Endpoint,
    token: bytes,
    *,
    max_object_size: int | None = None,
    handshake_timeout: float = HANDSHAKE_TIMEOUT
)

Handles client connections to an endpoint.

The handler authenticates each client with the handshake defined in proxystore.endpoint.protocol then forwards the client's requests to the endpoint.

Example

handler = ClientHandler(endpoint, token)
server = await handler.start_server('localhost', 8765)
...
server.close()
await handler.close_connections()
await server.wait_closed()

Parameters:

• endpoint (Endpoint) –

  Endpoint to forward client requests to.

• token (bytes) –

  Token that clients must prove they know.

• max_object_size (int | None, default: None ) –

  Optional maximum size in bytes of objects that clients can set. Requests exceeding this size are rejected before the data is read. This should match the maximum object size of the endpoint's storage, which rejects objects only after the data is read.

• handshake_timeout (float, default: HANDSHAKE_TIMEOUT ) –

  Seconds a client has to complete the handshake.

Source code in proxystore/endpoint/server.py

def __init__(
    self,
    endpoint: Endpoint,
    token: bytes,
    *,
    max_object_size: int | None = None,
    handshake_timeout: float = HANDSHAKE_TIMEOUT,
) -> None:
    self.endpoint = endpoint
    self.token = token
    self.max_object_size = max_object_size
    self.handshake_timeout = handshake_timeout
    self._connections: set[_ClientConnection] = set()
    self._tasks: set[asyncio.Task[None]] = set()
    self._warned_versions: set[Versions] = set()

start_server async

start_server(
    host: str,
    port: int,
    *,
    ssl_context: SSLContext | None = None
) -> Server

Start a server that handles connections on the host and port.

Parameters:

• host (str) –

  Address to listen on.

• port (int) –

  Port to listen on.

• ssl_context (SSLContext | None, default: None ) –

  Optional SSL context to encrypt connections with TLS.

Source code in proxystore/endpoint/server.py

async def start_server(
    self,
    host: str,
    port: int,
    *,
    ssl_context: ssl.SSLContext | None = None,
) -> asyncio.Server:
    """Start a server that handles connections on the host and port.

    Args:
        host: Address to listen on.
        port: Port to listen on.
        ssl_context: Optional SSL context to encrypt connections with TLS.
    """
    loop = asyncio.get_running_loop()
    return await loop.create_server(
        lambda: _ClientConnection(self._handle_connection, self._tasks),
        host=host,
        port=port,
        ssl=ssl_context,
    )

close_connections async

close_connections(timeout: float = 1) -> None

Close all open client connections.

Connection handlers waiting on the client finish once their connection is closed. Handlers that do not finish within timeout seconds (e.g., because a request is waiting on a peer endpoint) are cancelled.

Parameters:

• timeout (float, default: 1 ) –

  Seconds to wait for connection handlers to finish before cancelling them.

Source code in proxystore/endpoint/server.py

async def close_connections(self, timeout: float = 1) -> None:
    """Close all open client connections.

    Connection handlers waiting on the client finish once their
    connection is closed. Handlers that do not finish within `timeout`
    seconds (e.g., because a request is waiting on a peer endpoint) are
    cancelled.

    Args:
        timeout: Seconds to wait for connection handlers to finish
            before cancelling them.
    """
    for conn in list(self._connections):
        conn.close()
    tasks = list(self._tasks)
    if len(tasks) == 0:
        return
    _, pending = await asyncio.wait(tasks, timeout=timeout)
    for task in pending:
        task.cancel()
    await asyncio.gather(*pending, return_exceptions=True)