venv/lib/python3.11/site-packages/wandb/sdk/interface/router_relay.py

"""Router - handle message router (relay).

Router to manage responses from a queue with relay.

"""

from __future__ import annotations

from typing import TYPE_CHECKING

from wandb.proto import wandb_internal_pb2 as pb
from wandb.proto import wandb_server_pb2 as spb
from wandb.sdk.mailbox import Mailbox

from .router_queue import MessageQueueRouter

if TYPE_CHECKING:
    from queue import Queue


class MessageRelayRouter(MessageQueueRouter):
    _relay_queue: Queue[pb.Result]

    def __init__(
        self,
        request_queue: Queue[pb.Record],
        response_queue: Queue[pb.Result],
        relay_queue: Queue[pb.Result],
        mailbox: Mailbox,
    ) -> None:
        self._relay_queue = relay_queue
        super().__init__(
            request_queue=request_queue,
            response_queue=response_queue,
            mailbox=mailbox,
        )

    def _handle_msg_rcv(self, msg: pb.Result | spb.ServerResponse) -> None:
        if isinstance(msg, pb.Result):
            relay_msg = msg
        else:
            relay_msg = msg.result_communicate

        # This is legacy-service logic for returning responses to the client.
        # A different thread reads the "relay queue" and writes responses on
        # the socket.
        if relay_msg.control.relay_id:
            self._relay_queue.put(relay_msg)
        else:
            super()._handle_msg_rcv(msg)