"""InterfaceSock - Derived from InterfaceShared using a socket to send to internal thread. See interface.py for how interface classes relate to each other. """ import logging from typing import TYPE_CHECKING, Any, Optional from wandb.sdk.mailbox import Mailbox from ..lib.sock_client import SockClient from .interface_shared import InterfaceShared if TYPE_CHECKING: from wandb.proto import wandb_internal_pb2 as pb logger = logging.getLogger("wandb") class InterfaceSock(InterfaceShared): def __init__( self, sock_client: SockClient, mailbox: Mailbox, stream_id: str, ) -> None: super().__init__(mailbox=mailbox) self._sock_client = sock_client self._stream_id = stream_id def _assign(self, record: Any) -> None: assert self._stream_id record._info.stream_id = self._stream_id def _publish(self, record: "pb.Record", local: Optional[bool] = None) -> None: self._assign(record) self._sock_client.send_record_publish(record)