proxystore.p2p.relay.messages

proxystore/p2p/relay/messages.py

Message types for relay client and relay server communication.

RelayMessageType

Bases: Enum

Types of messages supported.

relay_response class-attribute instance-attribute

relay_response = 'RelayResponse'

Relay response message.

relay_registration class-attribute instance-attribute

relay_registration = 'RelayRegistrationRequest'

Relay registration request message.

peer_connection class-attribute instance-attribute

peer_connection = 'PeerConnectionRequest'

Peer connection request message.

RelayMessage dataclass

RelayMessage()

Base message.

RelayRegistrationRequest dataclass

RelayRegistrationRequest(
    name: str,
    uuid: UUID,
    message_type: str = RelayMessageType.relay_registration.name,
)

Bases: RelayMessage

Register with relay server as peer.

Attributes:

• name (str) –

  Name of peer requesting to register.

• uuid (UUID) –

  UUID of peer requesting to register.

RelayResponse dataclass

RelayResponse(
    success: bool = True,
    message: str | None = None,
    error: bool = False,
    message_type: str = RelayMessageType.relay_response.name,
)

Bases: RelayMessage

Message returned by relay server on success or error.

Attributes:

• success (bool) –

  If the registration was successful.

• message (str | None) –

  Message from server.

• error (bool) –

  If message is an error message.

PeerConnectionRequest dataclass

PeerConnectionRequest(
    source_uuid: UUID,
    source_name: str,
    peer_uuid: UUID,
    description_type: Literal["answer", "offer"],
    description: str,
    error: str | None = None,
    message_type: str = RelayMessageType.peer_connection.name,
)

Bases: RelayMessage

Message used to request a peer-to-peer connection from a relay.

Attributes:

• source_uuid (UUID) –

  UUID of sending peer.

• source_name (str) –

  Name of sending peer.

• peer_uuid (UUID) –

  UUID of destination peer.

• description_type (Literal['answer', 'offer']) –

  One of 'answer' or 'offer' indicating the type of message being sent.

• description (str) –

  Session description protocol message.

• error (str | None) –

  Error string if a problem occurs.

RelayMessageError

Bases: Exception

Base exception type for relay messages.

RelayMessageDecodeError

Bases: RelayMessageError

Exception raised when a message cannot be decoded.

RelayMessageEncodeError

Bases: RelayMessageError

Exception raised when an message cannot be encoded.

uuid_to_str

uuid_to_str(data: dict[str, Any]) -> dict[str, Any]

Cast any UUIDs to strings.

Scans the input dictionary for any values where the associated key contains 'uuid' and value is a UUID instance and converts it to a string for jsonification.

Returns:

• dict[str, Any] –

  Shallow copy of the input dictionary with values cast from UUID to str if their key also contains UUID.

Source code in proxystore/p2p/relay/messages.py

def uuid_to_str(data: dict[str, Any]) -> dict[str, Any]:
    """Cast any UUIDs to strings.

    Scans the input dictionary for any values where the associated key
    contains 'uuid' and value is a UUID instance and converts it to a
    string for jsonification.

    Returns:
        Shallow copy of the input dictionary with values cast from UUID \
        to str if their key also contains UUID.
    """
    data = data.copy()
    for key in data:
        if 'uuid' in key.lower() and isinstance(data[key], uuid.UUID):
            data[key] = str(data[key])
    return data

str_to_uuid

str_to_uuid(data: dict[str, Any]) -> dict[str, Any]

Cast any possible UUID strings to UUID objects.

The inverse operation of uuid_to_str().

Returns:

• dict[str, Any] –

  Shallow copy of the input dictionary with values cast from str to UUID if the key also contains UUID.

Raises:

• RelayMessageDecodeError –

  If a key contains 'uuid' but the value cannot be cast to a UUID.

Source code in proxystore/p2p/relay/messages.py

def str_to_uuid(data: dict[str, Any]) -> dict[str, Any]:
    """Cast any possible UUID strings to UUID objects.

    The inverse operation of
    [uuid_to_str()][proxystore.p2p.relay.messages.uuid_to_str].

    Returns:
        Shallow copy of the input dictionary with values cast from \
        str to UUID if the key also contains UUID.

    Raises:
        RelayMessageDecodeError: If a key contains 'uuid' but the value cannot
            be cast to a UUID.
    """
    data = data.copy()
    for key in data:
        if 'uuid' in key.lower():
            try:
                data[key] = uuid.UUID(data[key])
            except (AttributeError, TypeError, ValueError) as e:
                raise RelayMessageDecodeError(
                    f'Failed to convert key {key} to UUID.',
                ) from e
    return data

decode_relay_message

decode_relay_message(message: str) -> RelayMessage

Decode JSON string into correct relay message type.

Parameters:

• message (str) –

  JSON string to decode.

Returns:

• RelayMessage –

  Parsed message.

Raises:

• RelayMessageDecodeError –

  If the message cannot be decoded.

Source code in proxystore/p2p/relay/messages.py

def decode_relay_message(message: str) -> RelayMessage:
    """Decode JSON string into correct relay message type.

    Args:
        message: JSON string to decode.

    Returns:
        Parsed message.

    Raises:
        RelayMessageDecodeError: If the message cannot be decoded.
    """
    try:
        data = json.loads(message)
    except json.JSONDecodeError as e:
        raise RelayMessageDecodeError('Failed to load string as JSON.') from e

    try:
        message_type_name = data.pop('message_type')
    except KeyError as e:
        raise RelayMessageDecodeError(
            'Message does not contain a message_type key.',
        ) from e

    try:
        message_type = getattr(
            sys.modules[__name__],
            RelayMessageType[message_type_name].value,
        )
    except (AttributeError, KeyError) as e:
        raise RelayMessageDecodeError(
            'The message is of an unknown message type: '
            f'{message_type_name}.',
        ) from e

    data = str_to_uuid(data)

    try:
        return message_type(**data)
    except TypeError as e:
        raise RelayMessageDecodeError(
            f'Failed to convert message to {message_type.__name__}: {e}',
        ) from e

encode_relay_message

encode_relay_message(message: RelayMessage) -> str

Encode message as JSON string.

Parameters:

• message (RelayMessage) –

  Message to JSON encode.

Raises:

• RelayMessageEncodeError –

  If the message cannot be JSON encoded.

Source code in proxystore/p2p/relay/messages.py

def encode_relay_message(message: RelayMessage) -> str:
    """Encode message as JSON string.

    Args:
        message: Message to JSON encode.

    Raises:
        RelayMessageEncodeError: If the message cannot be JSON encoded.
    """
    if not isinstance(message, RelayMessage):
        raise RelayMessageEncodeError(
            f'Message is not an instance of {RelayMessage.__name__}. '
            f'Got {type(message).__name__}.',
        )

    data = dataclasses.asdict(message)
    data = uuid_to_str(data)

    try:
        return json.dumps(data)
    except TypeError as e:
        raise RelayMessageEncodeError('Error encoding message.') from e