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 passes their requests to a Dispatcher.

HANDSHAKE_TIMEOUT module-attribute

HANDSHAKE_TIMEOUT = 10

Seconds a client has to complete the handshake after connecting.

ClientHandler

ClientHandler(
    dispatcher: Dispatcher,
    token: EndpointToken,
    *,
    name: str,
    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 passes the client's requests to the dispatcher of the endpoint.

Example

handler = ClientHandler(dispatcher, token, name='my-endpoint')
await handler.start('localhost', 8765)
...
await handler.close()

Parameters:

• dispatcher (Dispatcher) –

  Dispatcher which handles client requests.

• token (EndpointToken) –

  Token that clients must prove they know.

• name (str) –

  Name of the endpoint sent to clients.

• 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.

• handshake_timeout (float, default: HANDSHAKE_TIMEOUT ) –

  Seconds a client has to complete the handshake.

Source code in proxystore/endpoint/server.py

def __init__(
    self,
    dispatcher: Dispatcher,
    token: EndpointToken,
    *,
    name: str,
    max_object_size: int | None = None,
    handshake_timeout: float = HANDSHAKE_TIMEOUT,
) -> None:
    self.dispatcher = dispatcher
    self.token = token
    self.name = name
    self.max_object_size = max_object_size
    self.handshake_timeout = handshake_timeout
    self._server: asyncio.Server | None = None
    self._connections: set[_ClientConnection] = set()
    self._tasks: set[asyncio.Task[None]] = set()
    self._warned_versions: set[Versions] = set()

start async

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

Start accepting client connections on the host and port.

Parameters:

• host (str) –

  Address to listen on.

• port (int) –

  Port to listen on or 0 for a random port.

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

  Optional SSL context to encrypt connections with TLS.

Returns:

• int –

  The port the server is listening on.

Raises:

• RuntimeError –

  If the handler has already been started.

• OSError –

  If the server cannot listen on the host and port.

Source code in proxystore/endpoint/server.py

async def start(
    self,
    host: str,
    port: int,
    *,
    ssl_context: ssl.SSLContext | None = None,
) -> int:
    """Start accepting client connections on the host and port.

    Args:
        host: Address to listen on.
        port: Port to listen on or `0` for a random port.
        ssl_context: Optional SSL context to encrypt connections with TLS.

    Returns:
        The port the server is listening on.

    Raises:
        RuntimeError: If the handler has already been started.
        OSError: If the server cannot listen on the host and port.
    """
    if self._server is not None:
        raise RuntimeError('The client handler has already been started.')
    loop = asyncio.get_running_loop()
    self._server = await loop.create_server(
        lambda: _ClientConnection(self._handle_connection, self._tasks),
        host=host,
        port=port,
        ssl=ssl_context,
    )
    return self._server.sockets[0].getsockname()[1]

close async

close(timeout: float = 1) -> None

Stop accepting connections and close all 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, and connections still open are aborted (e.g., TLS connections to idle clients). This is idempotent so it is safe to call multiple times.

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(self, timeout: float = 1) -> None:
    """Stop accepting connections and close all 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, and connections still open are aborted (e.g., TLS
    connections to idle clients). This is idempotent so it is safe to
    call multiple times.

    Args:
        timeout: Seconds to wait for connection handlers to finish
            before cancelling them.
    """
    if self._server is None:
        return
    server, self._server = self._server, None
    server.close()
    await self._close_connections(timeout)
    await server.wait_closed()