wandb 0.19.6rc4__py3-none-win_amd64.whl → 0.19.8__py3-none-win_amd64.whl

wandb/sdk/interface/interface_shared.py

@@ -5,64 +5,23 @@ See interface.py for how interface classes relate to each other.
 """
 
 import logging
-import time
 from abc import abstractmethod
-from multiprocessing.process import BaseProcess
 from typing import Any, Optional, cast
 
-import wandb
 from wandb.proto import wandb_internal_pb2 as pb
 from wandb.proto import wandb_telemetry_pb2 as tpb
+from wandb.sdk.mailbox import Mailbox, MailboxHandle
 from wandb.util import json_dumps_safer, json_friendly
 
-from ..lib.mailbox import Mailbox, MailboxHandle
 from .interface import InterfaceBase
-from .message_future import MessageFuture
-from .router import MessageRouter
 
 logger = logging.getLogger("wandb")
 
 
 class InterfaceShared(InterfaceBase):
-    process: Optional[BaseProcess]
-    _process_check: bool
-    _router: Optional[MessageRouter]
-    _mailbox: Optional[Mailbox]
-    _transport_success_timestamp: float
-    _transport_failed: bool
-
-    def __init__(
-        self,
-        process: Optional[BaseProcess] = None,
-        process_check: bool = True,
-        mailbox: Optional[Any] = None,
-    ) -> None:
+    def __init__(self, mailbox: Optional[Mailbox] = None) -> None:
         super().__init__()
-        self._transport_success_timestamp = time.monotonic()
-        self._transport_failed = False
-        self._process = process
-        self._router = None
-        self._process_check = process_check
         self._mailbox = mailbox
-        self._init_router()
-
-    @abstractmethod
-    def _init_router(self) -> None:
-        raise NotImplementedError
-
-    @property
-    def transport_failed(self) -> bool:
-        return self._transport_failed
-
-    @property
-    def transport_success_timestamp(self) -> float:
-        return self._transport_success_timestamp
-
-    def _transport_mark_failed(self) -> None:
-        self._transport_failed = True
-
-    def _transport_mark_success(self) -> None:
-        self._transport_success_timestamp = time.monotonic()
 
     def _publish_output(self, outdata: pb.OutputRecord) -> None:
         rec = pb.Record()
@@ -100,7 +59,9 @@ class InterfaceShared(InterfaceBase):
         rec = self._make_record(telemetry=telem)
         self._publish(rec)
 
-    def _publish_job_input(self, job_input: pb.JobInputRequest) -> MailboxHandle:
+    def _publish_job_input(
+        self, job_input: pb.JobInputRequest
+    ) -> MailboxHandle[pb.Result]:
         record = self._make_request(job_input=job_input)
         return self._deliver_record(record)
 
@@ -114,15 +75,8 @@ class InterfaceShared(InterfaceBase):
             item.value_json = json_dumps_safer(json_friendly(v)[0])
         return stats
 
-    def _make_login(self, api_key: Optional[str] = None) -> pb.LoginRequest:
-        login = pb.LoginRequest()
-        if api_key:
-            login.api_key = api_key
-        return login
-
     def _make_request(  # noqa: C901
         self,
-        login: Optional[pb.LoginRequest] = None,
         get_summary: Optional[pb.GetSummaryRequest] = None,
         pause: Optional[pb.PauseRequest] = None,
         resume: Optional[pb.ResumeRequest] = None,
@@ -130,6 +84,7 @@ class InterfaceShared(InterfaceBase):
         stop_status: Optional[pb.StopStatusRequest] = None,
         internal_messages: Optional[pb.InternalMessagesRequest] = None,
         network_status: Optional[pb.NetworkStatusRequest] = None,
+        operation_stats: Optional[pb.OperationStatsRequest] = None,
         poll_exit: Optional[pb.PollExitRequest] = None,
         partial_history: Optional[pb.PartialHistoryRequest] = None,
         sampled_history: Optional[pb.SampledHistoryRequest] = None,
@@ -158,9 +113,7 @@ class InterfaceShared(InterfaceBase):
         metadata: Optional[pb.MetadataRequest] = None,
     ) -> pb.Record:
         request = pb.Request()
-        if login:
-            request.login.CopyFrom(login)
-        elif get_summary:
+        if get_summary:
             request.get_summary.CopyFrom(get_summary)
         elif pause:
             request.pause.CopyFrom(pause)
@@ -174,6 +127,8 @@ class InterfaceShared(InterfaceBase):
             request.internal_messages.CopyFrom(internal_messages)
         elif network_status:
             request.network_status.CopyFrom(network_status)
+        elif operation_stats:
+            request.operations.CopyFrom(operation_stats)
         elif poll_exit:
             request.poll_exit.CopyFrom(poll_exit)
         elif partial_history:
@@ -307,35 +262,6 @@ class InterfaceShared(InterfaceBase):
     def _publish(self, record: pb.Record, local: Optional[bool] = None) -> None:
         raise NotImplementedError
 
-    def _communicate(
-        self, rec: pb.Record, timeout: Optional[int] = 30, local: Optional[bool] = None
-    ) -> Optional[pb.Result]:
-        return self._communicate_async(rec, local=local).get(timeout=timeout)
-
-    def _communicate_async(
-        self, rec: pb.Record, local: Optional[bool] = None
-    ) -> MessageFuture:
-        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
-
-    def communicate_login(
-        self, api_key: Optional[str] = None, timeout: Optional[int] = 15
-    ) -> pb.LoginResponse:
-        login = self._make_login(api_key)
-        rec = self._make_request(login=login)
-        result = self._communicate(rec, timeout=timeout)
-        if result is None:
-            # TODO: friendlier error message here
-            raise wandb.Error(
-                "Couldn't communicate with backend after {} seconds".format(timeout)
-            )
-        login_response = result.response.login_response
-        assert login_response
-        return login_response
-
     def _publish_defer(self, state: "pb.DeferRequest.DeferState.V") -> None:
         defer = pb.DeferRequest(state=state)
         rec = self._make_request(defer=defer)
@@ -358,11 +284,6 @@ class InterfaceShared(InterfaceBase):
         rec = self._make_record(final=final)
         self._publish(rec)
 
-    def publish_login(self, api_key: Optional[str] = None) -> None:
-        login = self._make_login(api_key)
-        rec = self._make_request(login=login)
-        self._publish(rec)
-
     def _publish_pause(self, pause: pb.PauseRequest) -> None:
         rec = self._make_request(pause=pause)
         self._publish(rec)
@@ -410,19 +331,22 @@ class InterfaceShared(InterfaceBase):
         rec = self._make_record(use_artifact=use_artifact)
         self._publish(rec)
 
-    def _communicate_artifact(self, log_artifact: pb.LogArtifactRequest) -> Any:
+    def _deliver_artifact(
+        self,
+        log_artifact: pb.LogArtifactRequest,
+    ) -> MailboxHandle[pb.Result]:
         rec = self._make_request(log_artifact=log_artifact)
-        return self._communicate_async(rec)
+        return self._deliver_record(rec)
 
     def _deliver_download_artifact(
         self, download_artifact: pb.DownloadArtifactRequest
-    ) -> MailboxHandle:
+    ) -> MailboxHandle[pb.Result]:
         rec = self._make_request(download_artifact=download_artifact)
         return self._deliver_record(rec)
 
     def _deliver_link_artifact(
         self, link_artifact: pb.LinkArtifactRequest
-    ) -> MailboxHandle:
+    ) -> MailboxHandle[pb.Result]:
         rec = self._make_request(link_artifact=link_artifact)
         return self._deliver_record(rec)
 
@@ -437,7 +361,7 @@ class InterfaceShared(InterfaceBase):
     def _deliver_status(
         self,
         status: pb.StatusRequest,
-    ) -> MailboxHandle:
+    ) -> MailboxHandle[pb.Result]:
         req = self._make_request(status=status)
         return self._deliver_record(req)
 
@@ -449,114 +373,119 @@ class InterfaceShared(InterfaceBase):
         record = self._make_request(keepalive=keepalive)
         self._publish(record)
 
-    def _communicate_shutdown(self) -> None:
-        # shutdown
+    def _deliver_shutdown(self) -> MailboxHandle[pb.Result]:
         request = pb.Request(shutdown=pb.ShutdownRequest())
         record = self._make_record(request=request)
-        _ = self._communicate(record)
+        return self._deliver_record(record)
 
     def _get_mailbox(self) -> Mailbox:
         mailbox = self._mailbox
         assert mailbox
         return mailbox
 
-    def _deliver_record(self, record: pb.Record) -> MailboxHandle:
+    def _deliver_record(self, record: pb.Record) -> MailboxHandle[pb.Result]:
         mailbox = self._get_mailbox()
-        handle = mailbox._deliver_record(record, interface=self)
-        return handle
 
-    def _deliver_run(self, run: pb.RunRecord) -> MailboxHandle:
+        handle = mailbox.require_response(record)
+        self._publish(record)
+
+        return handle.map(lambda resp: resp.result_communicate)
+
+    def _deliver_run(self, run: pb.RunRecord) -> MailboxHandle[pb.Result]:
         record = self._make_record(run=run)
         return self._deliver_record(record)
 
-    def _deliver_finish_sync(self, sync_finish: pb.SyncFinishRequest) -> MailboxHandle:
+    def _deliver_finish_sync(
+        self,
+        sync_finish: pb.SyncFinishRequest,
+    ) -> MailboxHandle[pb.Result]:
         record = self._make_request(sync_finish=sync_finish)
         return self._deliver_record(record)
 
-    def _deliver_run_start(self, run_start: pb.RunStartRequest) -> MailboxHandle:
+    def _deliver_run_start(
+        self,
+        run_start: pb.RunStartRequest,
+    ) -> MailboxHandle[pb.Result]:
         record = self._make_request(run_start=run_start)
         return self._deliver_record(record)
 
-    def _deliver_get_summary(self, get_summary: pb.GetSummaryRequest) -> MailboxHandle:
+    def _deliver_get_summary(
+        self,
+        get_summary: pb.GetSummaryRequest,
+    ) -> MailboxHandle[pb.Result]:
         record = self._make_request(get_summary=get_summary)
         return self._deliver_record(record)
 
     def _deliver_get_system_metrics(
         self, get_system_metrics: pb.GetSystemMetricsRequest
-    ) -> MailboxHandle:
+    ) -> MailboxHandle[pb.Result]:
         record = self._make_request(get_system_metrics=get_system_metrics)
         return self._deliver_record(record)
 
     def _deliver_get_system_metadata(
         self, get_system_metadata: pb.GetSystemMetadataRequest
-    ) -> MailboxHandle:
+    ) -> MailboxHandle[pb.Result]:
         record = self._make_request(get_system_metadata=get_system_metadata)
         return self._deliver_record(record)
 
-    def _deliver_exit(self, exit_data: pb.RunExitRecord) -> MailboxHandle:
+    def _deliver_exit(
+        self,
+        exit_data: pb.RunExitRecord,
+    ) -> MailboxHandle[pb.Result]:
         record = self._make_record(exit=exit_data)
         return self._deliver_record(record)
 
-    def _deliver_poll_exit(self, poll_exit: pb.PollExitRequest) -> MailboxHandle:
+    def deliver_operation_stats(self):
+        record = self._make_request(operation_stats=pb.OperationStatsRequest())
+        return self._deliver_record(record)
+
+    def _deliver_poll_exit(
+        self,
+        poll_exit: pb.PollExitRequest,
+    ) -> MailboxHandle[pb.Result]:
         record = self._make_request(poll_exit=poll_exit)
         return self._deliver_record(record)
 
     def _deliver_finish_without_exit(
         self, run_finish_without_exit: pb.RunFinishWithoutExitRequest
-    ) -> MailboxHandle:
+    ) -> MailboxHandle[pb.Result]:
         record = self._make_request(run_finish_without_exit=run_finish_without_exit)
         return self._deliver_record(record)
 
-    def _deliver_stop_status(self, stop_status: pb.StopStatusRequest) -> MailboxHandle:
+    def _deliver_stop_status(
+        self,
+        stop_status: pb.StopStatusRequest,
+    ) -> MailboxHandle[pb.Result]:
         record = self._make_request(stop_status=stop_status)
         return self._deliver_record(record)
 
-    def _deliver_attach(self, attach: pb.AttachRequest) -> MailboxHandle:
+    def _deliver_attach(
+        self,
+        attach: pb.AttachRequest,
+    ) -> MailboxHandle[pb.Result]:
         record = self._make_request(attach=attach)
         return self._deliver_record(record)
 
     def _deliver_network_status(
         self, network_status: pb.NetworkStatusRequest
-    ) -> MailboxHandle:
+    ) -> MailboxHandle[pb.Result]:
         record = self._make_request(network_status=network_status)
         return self._deliver_record(record)
 
     def _deliver_internal_messages(
         self, internal_message: pb.InternalMessagesRequest
-    ) -> MailboxHandle:
+    ) -> MailboxHandle[pb.Result]:
         record = self._make_request(internal_messages=internal_message)
         return self._deliver_record(record)
 
     def _deliver_request_sampled_history(
         self, sampled_history: pb.SampledHistoryRequest
-    ) -> MailboxHandle:
+    ) -> MailboxHandle[pb.Result]:
         record = self._make_request(sampled_history=sampled_history)
         return self._deliver_record(record)
 
     def _deliver_request_run_status(
         self, run_status: pb.RunStatusRequest
-    ) -> MailboxHandle:
+    ) -> MailboxHandle[pb.Result]:
         record = self._make_request(run_status=run_status)
         return self._deliver_record(record)
-
-    def _transport_keepalive_failed(self, keepalive_interval: int = 5) -> bool:
-        if self._transport_failed:
-            return True
-
-        now = time.monotonic()
-        if now < self._transport_success_timestamp + keepalive_interval:
-            return False
-
-        try:
-            self.publish_keepalive()
-        except Exception:
-            self._transport_mark_failed()
-        else:
-            self._transport_mark_success()
-        return self._transport_failed
-
-    def join(self) -> None:
-        super().join()
-
-        if self._router:
-            self._router.join()

wandb/sdk/interface/interface_sock.py

@@ -7,11 +7,10 @@ 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 wandb.sdk.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
@@ -21,23 +20,16 @@ logger = logging.getLogger("wandb")
 
 
 class InterfaceSock(InterfaceShared):
-    _mailbox: Mailbox
-
     def __init__(
         self,
         sock_client: SockClient,
         mailbox: Mailbox,
         stream_id: str,
     ) -> 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._sock_client = sock_client
         self._stream_id = stream_id
 
-    def _init_router(self) -> None:
-        self._router = MessageSockRouter(self._sock_client, mailbox=self._mailbox)
-
     def _assign(self, record: Any) -> None:
         assert self._stream_id
         record._info.stream_id = self._stream_id
@@ -45,13 +37,3 @@ class InterfaceSock(InterfaceShared):
     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/router.py

@@ -4,20 +4,20 @@ Router to manage responses.
 
 """
 
+from __future__ import annotations
+
 import logging
 import threading
-import uuid
 from abc import abstractmethod
-from typing import TYPE_CHECKING, Dict, Optional
+from typing import TYPE_CHECKING
 
-from ..lib import mailbox
-from .message_future import MessageFuture
+from wandb.proto import wandb_internal_pb2 as pb
+from wandb.proto import wandb_server_pb2 as spb
+from wandb.sdk import mailbox
 
 if TYPE_CHECKING:
     from queue import Queue
 
-    from wandb.proto import wandb_internal_pb2 as pb
-
 
 logger = logging.getLogger("wandb")
 
@@ -26,26 +26,13 @@ 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]
+    _request_queue: Queue[pb.Record]
+    _response_queue: Queue[pb.Result]
+    _mailbox: mailbox.Mailbox | None
 
-    def __init__(self, mailbox: Optional[mailbox.Mailbox] = None) -> None:
+    def __init__(self, mailbox: mailbox.Mailbox | None = None) -> None:
         self._mailbox = mailbox
-        self._pending_reqs = {}
         self._lock = threading.Lock()
 
         self._join_event = threading.Event()
@@ -55,61 +42,48 @@ class MessageRouter:
         self._thread.start()
 
     @abstractmethod
-    def _read_message(self) -> Optional["pb.Result"]:
+    def _read_message(self) -> pb.Result | spb.ServerResponse | None:
         raise NotImplementedError
 
     @abstractmethod
-    def _send_message(self, record: "pb.Record") -> None:
+    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
+        try:
+            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 as e:
+                    logger.warning("message_loop has been closed", exc_info=e)
+                    break
+                if not msg:
+                    continue
+                self._handle_msg_rcv(msg)
+
+        finally:
+            if self._mailbox:
+                self._mailbox.close()
 
     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)
+    def _handle_msg_rcv(self, msg: pb.Result | spb.ServerResponse) -> None:
+        if not self._mailbox:
             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
+
+        if isinstance(msg, pb.Result) and msg.control.mailbox_slot:
+            self._mailbox.deliver(
+                spb.ServerResponse(
+                    request_id=msg.control.mailbox_slot,
+                    result_communicate=msg,
                 )
-            return
-        future._set_object(msg)
+            )
+        elif isinstance(msg, spb.ServerResponse) and msg.request_id:
+            self._mailbox.deliver(msg)

wandb/sdk/interface/router_queue.py

@@ -4,38 +4,40 @@ Router to manage responses from a queue.
 
 """
 
+from __future__ import annotations
+
 import queue
-from typing import TYPE_CHECKING, Optional
+from typing import TYPE_CHECKING
+
+from wandb.proto import wandb_internal_pb2 as pb
+from wandb.sdk.mailbox import Mailbox
 
-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]"
+    _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,
+        request_queue: Queue[pb.Record],
+        response_queue: Queue[pb.Result],
+        mailbox: Mailbox | None = None,
     ) -> None:
         self._request_queue = request_queue
         self._response_queue = response_queue
         super().__init__(mailbox=mailbox)
 
-    def _read_message(self) -> Optional["pb.Result"]:
+    def _read_message(self) -> pb.Result | None:
         try:
             msg = self._response_queue.get(timeout=1)
         except queue.Empty:
             return None
         return msg
 
-    def _send_message(self, record: "pb.Record") -> None:
+    def _send_message(self, record: pb.Record) -> None:
         self._request_queue.put(record)

wandb/sdk/interface/router_relay.py

@@ -4,34 +4,47 @@ Router to manage responses from a queue with relay.
 
 """
 
+from __future__ import annotations
+
 from typing import TYPE_CHECKING
 
-from ..lib.mailbox import Mailbox
+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
 
-    from wandb.proto import wandb_internal_pb2 as pb
-
 
 class MessageRelayRouter(MessageQueueRouter):
-    _relay_queue: "Queue[pb.Result]"
+    _relay_queue: Queue[pb.Result]
 
     def __init__(
         self,
-        request_queue: "Queue[pb.Record]",
-        response_queue: "Queue[pb.Result]",
-        relay_queue: "Queue[pb.Result]",
+        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
+            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)
+    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)