wandb 0.19.1__py3-none-musllinux_1_2_aarch64.whl

wandb/sdk/interface/interface_sock.py

@@ -0,0 +1,61 @@
+"""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 ..lib.mailbox import Mailbox
+from ..lib.sock_client import SockClient
+from .interface_shared import InterfaceShared
+from .message_future import MessageFuture
+from .router_sock import MessageSockRouter
+
+if TYPE_CHECKING:
+    from wandb.proto import wandb_internal_pb2 as pb
+
+    from ..wandb_run import Run
+
+
+logger = logging.getLogger("wandb")
+
+
+class InterfaceSock(InterfaceShared):
+    _stream_id: Optional[str]
+    _sock_client: SockClient
+    _mailbox: Mailbox
+
+    def __init__(self, sock_client: SockClient, mailbox: Mailbox) -> None:
+        # _sock_client is used when abstract method _init_router() is called by constructor
+        self._sock_client = sock_client
+        super().__init__(mailbox=mailbox)
+        self._process_check = False
+        self._stream_id = None
+
+    def _init_router(self) -> None:
+        self._router = MessageSockRouter(self._sock_client, mailbox=self._mailbox)
+
+    def _hack_set_run(self, run: "Run") -> None:
+        super()._hack_set_run(run)
+        assert run._settings.run_id
+        self._stream_id = run._settings.run_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)
+
+    def _communicate_async(
+        self, rec: "pb.Record", local: Optional[bool] = None
+    ) -> MessageFuture:
+        self._assign(rec)
+        assert self._router
+        if self._process_check and self._process and not self._process.is_alive():
+            raise Exception("The wandb backend process has shutdown")
+        future = self._router.send_and_receive(rec, local=local)
+        return future

wandb/sdk/interface/message_future.py

@@ -0,0 +1,27 @@
+"""MessageFuture - represents a message result of an asynchronous operation.
+
+Base class MessageFuture for MessageFutureObject and MessageFuturePoll
+
+"""
+
+import threading
+from abc import abstractmethod
+from typing import Optional
+
+from wandb.proto import wandb_internal_pb2 as pb
+
+
+class MessageFuture:
+    _object: Optional[pb.Result]
+
+    def __init__(self) -> None:
+        self._object = None
+        self._object_ready = threading.Event()
+
+    def _set_object(self, obj: pb.Result) -> None:
+        self._object = obj
+        self._object_ready.set()
+
+    @abstractmethod
+    def get(self, timeout: Optional[int] = None) -> Optional[pb.Result]:
+        raise NotImplementedError

wandb/sdk/interface/message_future_poll.py

@@ -0,0 +1,50 @@
+"""MessageFuturePoll - Derived from MessageFuture but implementing polling loop.
+
+MessageFuture represents a message result of an asynchronous operation.
+
+MessageFuturePoll implements a polling loop to periodically query for a
+completed async operation.
+
+"""
+
+import time
+from typing import Any, Optional
+
+from wandb.proto import wandb_internal_pb2 as pb
+
+from .message_future import MessageFuture
+
+
+class MessageFuturePoll(MessageFuture):
+    _fn: Any
+    _xid: str
+
+    def __init__(self, fn: Any, xid: str) -> None:
+        super().__init__()
+        self._fn = fn
+        self._xid = xid
+
+    def get(self, timeout: Optional[int] = None) -> Optional[pb.Result]:
+        self._poll(timeout=timeout)
+        if self._object_ready.is_set():
+            return self._object
+        return None
+
+    def _poll(self, timeout: Optional[int] = None) -> None:
+        if self._object_ready.is_set():
+            return
+        done = False
+        start_time = time.time()
+        sleep_time = 0.5
+        while not done:
+            result = self._fn(xid=self._xid)
+            if result:
+                self._set_object(result)
+                done = True
+                continue
+            now_time = time.time()
+            if timeout and start_time - now_time > timeout:
+                done = True
+                continue
+            time.sleep(sleep_time)
+            sleep_time = min(sleep_time * 2, 5)

wandb/sdk/interface/router.py

@@ -0,0 +1,115 @@
+"""Router - handle message router (base class).
+
+Router to manage responses.
+
+"""
+
+import logging
+import threading
+import uuid
+from abc import abstractmethod
+from typing import TYPE_CHECKING, Dict, Optional
+
+from ..lib import mailbox
+from .message_future import MessageFuture
+
+if TYPE_CHECKING:
+    from queue import Queue
+
+    from wandb.proto import wandb_internal_pb2 as pb
+
+
+logger = logging.getLogger("wandb")
+
+
+class MessageRouterClosedError(Exception):
+    """Router has been closed."""
+
+
+class MessageFutureObject(MessageFuture):
+    def __init__(self) -> None:
+        super().__init__()
+
+    def get(self, timeout: Optional[int] = None) -> Optional["pb.Result"]:
+        is_set = self._object_ready.wait(timeout)
+        if is_set and self._object:
+            return self._object
+        return None
+
+
+class MessageRouter:
+    _pending_reqs: Dict[str, MessageFutureObject]
+    _request_queue: "Queue[pb.Record]"
+    _response_queue: "Queue[pb.Result]"
+    _mailbox: Optional[mailbox.Mailbox]
+
+    def __init__(self, mailbox: Optional[mailbox.Mailbox] = None) -> None:
+        self._mailbox = mailbox
+        self._pending_reqs = {}
+        self._lock = threading.Lock()
+
+        self._join_event = threading.Event()
+        self._thread = threading.Thread(target=self.message_loop)
+        self._thread.name = "MsgRouterThr"
+        self._thread.daemon = True
+        self._thread.start()
+
+    @abstractmethod
+    def _read_message(self) -> Optional["pb.Result"]:
+        raise NotImplementedError
+
+    @abstractmethod
+    def _send_message(self, record: "pb.Record") -> None:
+        raise NotImplementedError
+
+    def message_loop(self) -> None:
+        while not self._join_event.is_set():
+            try:
+                msg = self._read_message()
+            except EOFError:
+                # On abnormal shutdown the queue will be destroyed underneath
+                # resulting in EOFError.  message_loop needs to exit..
+                logger.warning("EOFError seen in message_loop")
+                break
+            except MessageRouterClosedError:
+                logger.warning("message_loop has been closed")
+                break
+            if not msg:
+                continue
+            self._handle_msg_rcv(msg)
+
+    def send_and_receive(
+        self, rec: "pb.Record", local: Optional[bool] = None
+    ) -> MessageFuture:
+        rec.control.req_resp = True
+        if local:
+            rec.control.local = local
+        rec.uuid = uuid.uuid4().hex
+        future = MessageFutureObject()
+        with self._lock:
+            self._pending_reqs[rec.uuid] = future
+
+        self._send_message(rec)
+
+        return future
+
+    def join(self) -> None:
+        self._join_event.set()
+        self._thread.join()
+
+    def _handle_msg_rcv(self, msg: "pb.Result") -> None:
+        # deliver mailbox addressed messages to mailbox
+        if self._mailbox and msg.control.mailbox_slot:
+            self._mailbox.deliver(msg)
+            return
+        with self._lock:
+            future = self._pending_reqs.pop(msg.uuid, None)
+        if future is None:
+            # TODO (cvp): saw this in tests, seemed benign enough to ignore, but
+            # could point to other issues.
+            if msg.uuid != "":
+                logger.warning(
+                    "No listener found for msg with uuid %s (%s)", msg.uuid, msg
+                )
+            return
+        future._set_object(msg)

wandb/sdk/interface/router_queue.py

@@ -0,0 +1,41 @@
+"""Router - handle message router (queue).
+
+Router to manage responses from a queue.
+
+"""
+
+import queue
+from typing import TYPE_CHECKING, Optional
+
+from ..lib.mailbox import Mailbox
+from .router import MessageRouter
+
+if TYPE_CHECKING:
+    from queue import Queue
+
+    from wandb.proto import wandb_internal_pb2 as pb
+
+
+class MessageQueueRouter(MessageRouter):
+    _request_queue: "Queue[pb.Record]"
+    _response_queue: "Queue[pb.Result]"
+
+    def __init__(
+        self,
+        request_queue: "Queue[pb.Record]",
+        response_queue: "Queue[pb.Result]",
+        mailbox: Optional[Mailbox] = None,
+    ) -> None:
+        self._request_queue = request_queue
+        self._response_queue = response_queue
+        super().__init__(mailbox=mailbox)
+
+    def _read_message(self) -> Optional["pb.Result"]:
+        try:
+            msg = self._response_queue.get(timeout=1)
+        except queue.Empty:
+            return None
+        return msg
+
+    def _send_message(self, record: "pb.Record") -> None:
+        self._request_queue.put(record)

wandb/sdk/interface/router_relay.py

@@ -0,0 +1,37 @@
+"""Router - handle message router (relay).
+
+Router to manage responses from a queue with relay.
+
+"""
+
+from typing import TYPE_CHECKING
+
+from ..lib.mailbox import Mailbox
+from .router_queue import MessageQueueRouter
+
+if TYPE_CHECKING:
+    from queue import Queue
+
+    from wandb.proto import wandb_internal_pb2 as pb
+
+
+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") -> None:
+        if msg.control.relay_id:
+            self._relay_queue.put(msg)
+            return
+        super()._handle_msg_rcv(msg)

wandb/sdk/interface/router_sock.py

@@ -0,0 +1,36 @@
+"""Router - handle message router (sock).
+
+Router to manage responses from a socket client.
+
+"""
+
+from typing import TYPE_CHECKING, Optional
+
+from ..lib.mailbox import Mailbox
+from ..lib.sock_client import SockClient, SockClientClosedError
+from .router import MessageRouter, MessageRouterClosedError
+
+if TYPE_CHECKING:
+    from wandb.proto import wandb_internal_pb2 as pb
+
+
+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) -> Optional["pb.Result"]:
+        try:
+            resp = self._sock_client.read_server_response(timeout=1)
+        except SockClientClosedError:
+            raise MessageRouterClosedError
+        if not resp:
+            return None
+        msg = resp.result_communicate
+        return msg
+
+    def _send_message(self, record: "pb.Record") -> None:
+        self._sock_client.send_record_communicate(record)

wandb/sdk/interface/summary_record.py

@@ -0,0 +1,67 @@
+"""Summary Record.
+
+This module implements a summary record as an intermediate format before being converted
+to a protocol buffer.
+"""
+
+import typing as t
+
+
+class SummaryRecord:
+    """Encodes a diff -- analogous to the SummaryRecord protobuf message."""
+
+    update: t.List["SummaryItem"]
+    remove: t.List["SummaryItem"]
+
+    def __init__(self):
+        self.update = []
+        self.remove = []
+
+    def __str__(self):
+        s = "SummaryRecord:\n  Update:\n    "
+        s += "\n    ".join([str(item) for item in self.update])
+        s += "\n  Remove:\n    "
+        s += "\n    ".join([str(item) for item in self.remove])
+        s += "\n"
+        return s
+
+    __repr__ = __str__
+
+    def _add_next_parent(self, parent_key):
+        with_next_parent = SummaryRecord()
+        with_next_parent.update = [
+            item._add_next_parent(parent_key) for item in self.update
+        ]
+        with_next_parent.remove = [
+            item._add_next_parent(parent_key) for item in self.remove
+        ]
+
+        return with_next_parent
+
+
+class SummaryItem:
+    """Analogous to the SummaryItem protobuf message."""
+
+    key: t.Tuple[str]
+    value: t.Any
+
+    def __init__(self):
+        self.key = tuple()
+        self.value = None
+
+    def __str__(self):
+        return "SummaryItem: key: " + str(self.key) + " value: " + str(self.value)
+
+    __repr__ = __str__
+
+    def _add_next_parent(self, parent_key):
+        with_next_parent = SummaryItem()
+
+        key = self.key
+        if not isinstance(key, tuple):
+            key = (key,)
+
+        with_next_parent.key = (parent_key,) + self.key
+        with_next_parent.value = self.value
+
+        return with_next_parent

wandb/sdk/internal/context.py

@@ -0,0 +1,89 @@
+"""Context Keeper."""
+
+import logging
+import threading
+from typing import Dict, Optional
+
+from wandb.proto.wandb_internal_pb2 import Record, Result
+
+logger = logging.getLogger(__name__)
+
+
+class Context:
+    _cancel_event: threading.Event
+    # TODO(debug_context) add debug setting to enable this
+    # _debug_record: Optional[Record]
+
+    def __init__(self) -> None:
+        self._cancel_event = threading.Event()
+        # TODO(debug_context) see above
+        # self._debug_record = None
+
+    def cancel(self) -> None:
+        self._cancel_event.set()
+
+    @property
+    def cancel_event(self) -> threading.Event:
+        return self._cancel_event
+
+
+def context_id_from_record(record: Record) -> str:
+    context_id = record.control.mailbox_slot
+    return context_id
+
+
+def context_id_from_result(result: Result) -> str:
+    context_id = result.control.mailbox_slot
+    return context_id
+
+
+class ContextKeeper:
+    _active_items: Dict[str, Context]
+
+    def __init__(self) -> None:
+        self._active_items = {}
+
+    def add_from_record(self, record: Record) -> Optional[Context]:
+        context_id = context_id_from_record(record)
+        if not context_id:
+            return None
+        context_obj = self.add(context_id)
+
+        # TODO(debug_context) see above
+        # context_obj._debug_record = record
+
+        return context_obj
+
+    def add(self, context_id: str) -> Context:
+        assert context_id
+        context_obj = Context()
+        self._active_items[context_id] = context_obj
+        return context_obj
+
+    def get(self, context_id: str) -> Optional[Context]:
+        item = self._active_items.get(context_id)
+        return item
+
+    def release(self, context_id: str) -> None:
+        if not context_id:
+            return
+        _ = self._active_items.pop(context_id, None)
+
+    def cancel(self, context_id: str) -> bool:
+        item = self.get(context_id)
+        if item:
+            item.cancel()
+            return True
+        return False
+
+    # TODO(debug_context) see above
+    # def _debug_print_orphans(self, print_to_stdout: bool) -> None:
+    #     for context_id, context in self._active_items.items():
+    #         record = context._debug_record
+    #         record_type = record.WhichOneof("record_type") if record else "unknown"
+    #         message = (
+    #             f"Context: {context_id} {context.cancel_event.is_set()} {record_type}"
+    #         )
+    #         logger.warning(message)
+    #         if print_to_stdout:
+    #             print(message)