File size: 1,002 Bytes
9c6594c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 |
"""Router - handle message router (sock).
Router to manage responses from a socket client.
"""
from __future__ import annotations
from wandb.proto import wandb_internal_pb2 as pb
from wandb.proto import wandb_server_pb2 as spb
from wandb.sdk.lib.sock_client import SockClient, SockClientClosedError
from wandb.sdk.mailbox import Mailbox
from .router import MessageRouter, MessageRouterClosedError
class MessageSockRouter(MessageRouter):
_sock_client: SockClient
_mailbox: Mailbox
def __init__(self, sock_client: SockClient, mailbox: Mailbox) -> None:
self._sock_client = sock_client
super().__init__(mailbox=mailbox)
def _read_message(self) -> spb.ServerResponse | None:
try:
return self._sock_client.read_server_response(timeout=1)
except SockClientClosedError as e:
raise MessageRouterClosedError from e
def _send_message(self, record: pb.Record) -> None:
self._sock_client.send_record_communicate(record)
|