wandb 0.19.7__py3-none-macosx_11_0_arm64.whl → 0.19.8__py3-none-macosx_11_0_arm64.whl

wandb/sdk/interface/interface_shared.py

@@ -14,24 +14,14 @@ from wandb.sdk.mailbox import Mailbox, MailboxHandle
 from wandb.util import json_dumps_safer, json_friendly
 
 from .interface import InterfaceBase
-from .router import MessageRouter
 
 logger = logging.getLogger("wandb")
 
 
 class InterfaceShared(InterfaceBase):
-    _router: Optional[MessageRouter]
-    _mailbox: Optional[Mailbox]
-
-    def __init__(self, mailbox: Optional[Any] = None) -> None:
+    def __init__(self, mailbox: Optional[Mailbox] = None) -> None:
         super().__init__()
-        self._router = None
         self._mailbox = mailbox
-        self._init_router()
-
-    @abstractmethod
-    def _init_router(self) -> None:
-        raise NotImplementedError
 
     def _publish_output(self, outdata: pb.OutputRecord) -> None:
         rec = pb.Record()
@@ -69,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)
 
@@ -342,19 +334,19 @@ class InterfaceShared(InterfaceBase):
     def _deliver_artifact(
         self,
         log_artifact: pb.LogArtifactRequest,
-    ) -> MailboxHandle:
+    ) -> MailboxHandle[pb.Result]:
         rec = self._make_request(log_artifact=log_artifact)
         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)
 
@@ -369,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)
 
@@ -381,7 +373,7 @@ class InterfaceShared(InterfaceBase):
         record = self._make_request(keepalive=keepalive)
         self._publish(record)
 
-    def _deliver_shutdown(self) -> MailboxHandle:
+    def _deliver_shutdown(self) -> MailboxHandle[pb.Result]:
         request = pb.Request(shutdown=pb.ShutdownRequest())
         record = self._make_record(request=request)
         return self._deliver_record(record)
@@ -391,43 +383,55 @@ class InterfaceShared(InterfaceBase):
         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.require_response(record)
         self._publish(record)
 
-        return handle
+        return handle.map(lambda resp: resp.result_communicate)
 
-    def _deliver_run(self, run: pb.RunRecord) -> MailboxHandle:
+    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)
 
@@ -435,50 +439,53 @@ class InterfaceShared(InterfaceBase):
         record = self._make_request(operation_stats=pb.OperationStatsRequest())
         return self._deliver_record(record)
 
-    def _deliver_poll_exit(self, poll_exit: pb.PollExitRequest) -> MailboxHandle:
+    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 join(self) -> None:
-        super().join()
-
-        if self._router:
-            self._router.join()

wandb/sdk/interface/interface_sock.py

@@ -11,7 +11,6 @@ from wandb.sdk.mailbox import Mailbox
 
 from ..lib.sock_client import SockClient
 from .interface_shared import InterfaceShared
-from .router_sock import MessageSockRouter
 
 if TYPE_CHECKING:
     from wandb.proto import wandb_internal_pb2 as pb
@@ -21,22 +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._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

wandb/sdk/interface/router.py

@@ -4,21 +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 wandb.proto import wandb_internal_pb2 as pb
+from wandb.proto import wandb_server_pb2 as spb
 from wandb.sdk 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")
 
@@ -27,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()
@@ -56,11 +42,11 @@ 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:
@@ -84,38 +70,20 @@ class MessageRouter:
             if self._mailbox:
                 self._mailbox.close()
 
-    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)
+    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,9 +4,12 @@ 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 .router import MessageRouter
@@ -14,29 +17,27 @@ 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,8 +4,12 @@ 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
@@ -13,26 +17,34 @@ 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)

wandb/sdk/interface/router_sock.py

@@ -4,16 +4,15 @@ Router to manage responses from a socket client.
 
 """
 
-from typing import TYPE_CHECKING, Optional
+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
 
-if TYPE_CHECKING:
-    from wandb.proto import wandb_internal_pb2 as pb
-
 
 class MessageSockRouter(MessageRouter):
     _sock_client: SockClient
@@ -23,15 +22,11 @@ class MessageSockRouter(MessageRouter):
         self._sock_client = sock_client
         super().__init__(mailbox=mailbox)
 
-    def _read_message(self) -> Optional["pb.Result"]:
+    def _read_message(self) -> spb.ServerResponse | None:
         try:
-            resp = self._sock_client.read_server_response(timeout=1)
+            return self._sock_client.read_server_response(timeout=1)
         except SockClientClosedError as e:
             raise MessageRouterClosedError from e
-        if not resp:
-            return None
-        msg = resp.result_communicate
-        return msg
 
-    def _send_message(self, record: "pb.Record") -> None:
+    def _send_message(self, record: pb.Record) -> None:
         self._sock_client.send_record_communicate(record)

wandb/sdk/internal/sender.py

@@ -1,6 +1,7 @@
 """sender."""
 
 import contextlib
+import glob
 import gzip
 import json
 import logging
@@ -1408,7 +1409,8 @@ class SendManager:
         for k in files.files:
             # TODO(jhr): fix paths with directories
             self._save_file(
-                interface.GlobStr(k.path), interface.file_enum_to_policy(k.policy)
+                interface.GlobStr(glob.escape(k.path)),
+                interface.file_enum_to_policy(k.policy),
             )
 
     def send_header(self, record: "Record") -> None: