wandb 0.15.9__py3-none-any.whl → 0.15.11__py3-none-any.whl

wandb/sdk/artifacts/storage_policies/__init__.py

@@ -0,0 +1,4 @@
+from wandb.sdk.artifacts.storage_policies.register import WANDB_STORAGE_POLICY
+from wandb.sdk.artifacts.storage_policies.wandb_storage_policy import WandbStoragePolicy
+
+__all__ = ["WANDB_STORAGE_POLICY", "WandbStoragePolicy"]

wandb/sdk/artifacts/storage_policies/register.py

@@ -0,0 +1 @@
+WANDB_STORAGE_POLICY = "wandb-storage-policy-v1"

wandb/sdk/artifacts/storage_policies/wandb_storage_policy.py

@@ -23,6 +23,7 @@ from wandb.sdk.artifacts.storage_handlers.wb_local_artifact_handler import (
     WBLocalArtifactHandler,
 )
 from wandb.sdk.artifacts.storage_layout import StorageLayout
+from wandb.sdk.artifacts.storage_policies.register import WANDB_STORAGE_POLICY
 from wandb.sdk.artifacts.storage_policy import StoragePolicy
 from wandb.sdk.internal.thread_local_settings import _thread_local_api_settings
 from wandb.sdk.lib.hashutil import B64MD5, b64_to_hex_id, hex_to_b64_id
@@ -53,7 +54,7 @@ S3_MAX_MULTI_UPLOAD_SIZE = 5 * 1024**4
 class WandbStoragePolicy(StoragePolicy):
     @classmethod
     def name(cls) -> str:
-        return "wandb-storage-policy-v1"
+        return WANDB_STORAGE_POLICY
 
     @classmethod
     def from_config(cls, config: Dict) -> "WandbStoragePolicy":
@@ -380,7 +381,7 @@ class WandbStoragePolicy(StoragePolicy):
         )
         if not hit:
             try:
-                with cache_open() as f:
-                    shutil.copyfile(entry.local_path, f.name)
+                with cache_open("wb") as f, open(entry.local_path, "rb") as src:
+                    shutil.copyfileobj(src, f)
             except OSError as e:
                 termwarn(f"Failed to cache {entry.local_path}, ignoring {e}")

wandb/sdk/artifacts/storage_policy.py

@@ -12,11 +12,13 @@ if TYPE_CHECKING:
 
 class StoragePolicy:
     @classmethod
-    def lookup_by_name(cls, name: str) -> Optional[Type["StoragePolicy"]]:
+    def lookup_by_name(cls, name: str) -> Type["StoragePolicy"]:
+        import wandb.sdk.artifacts.storage_policies  # noqa: F401
+
         for sub in cls.__subclasses__():
             if sub.name() == name:
                 return sub
-        return None
+        raise NotImplementedError(f"Failed to find storage policy '{name}'")
 
     @classmethod
     def name(cls) -> str:

wandb/sdk/backend/backend.py

@@ -25,7 +25,6 @@ from wandb.sdk.wandb_settings import Settings
 if TYPE_CHECKING:
     from wandb.proto.wandb_internal_pb2 import Record, Result
 
-    from ..service.service_grpc import ServiceGrpcInterface
     from ..service.service_sock import ServiceSockInterface
     from ..wandb_run import Run
 
@@ -141,13 +140,6 @@ class Backend:
             main_module.__file__ = self._save_mod_path
 
     def _ensure_launched_manager(self) -> None:
-        # grpc_port: Optional[int] = None
-        # attach_id = self._settings._attach_id if self._settings else None
-        # if attach_id:
-        #     # TODO(attach): implement
-        #     # already have a server, assume it is already up
-        #     grpc_port = int(attach_id)
-
         assert self._manager
         svc = self._manager._get_service()
         assert svc
@@ -161,14 +153,6 @@ class Backend:
             sock_client = svc_iface_sock._get_sock_client()
             sock_interface = InterfaceSock(sock_client, mailbox=self._mailbox)
             self.interface = sock_interface
-        elif svc_transport == "grpc":
-            from ..interface.interface_grpc import InterfaceGrpc
-
-            svc_iface_grpc = cast("ServiceGrpcInterface", svc_iface)
-            stub = svc_iface_grpc._get_stub()
-            grpc_interface = InterfaceGrpc(mailbox=self._mailbox)
-            grpc_interface._connect(stub=stub)
-            self.interface = grpc_interface
         else:
             raise AssertionError(f"Unsupported service transport: {svc_transport}")
 

wandb/sdk/data_types/image.py

@@ -69,7 +69,7 @@ class Image(BatchableMedia):
             image data, or a PIL image. The class attempts to infer
             the data format and converts it.
         mode: (string) The PIL mode for an image. Most common are "L", "RGB",
-            "RGBA". Full explanation at https://pillow.readthedocs.io/en/4.2.x/handbook/concepts.html#concept-modes.
+            "RGBA". Full explanation at https://pillow.readthedocs.io/en/stable/handbook/concepts.html#modes.
         caption: (string) Label for display of image.
 
     Note : When logging a `torch.Tensor` as a `wandb.Image`, images are normalized. If you do not want to normalize your images, please convert your tensors to a PIL Image.
@@ -287,6 +287,8 @@ class Image(BatchableMedia):
             )
             if hasattr(data, "requires_grad") and data.requires_grad:
                 data = data.detach()  # type: ignore
+            if hasattr(data, "dtype") and str(data.dtype) == "torch.uint8":
+                data = data.to(float)
             data = vis_util.make_grid(data, normalize=True)
             self._image = pil_image.fromarray(
                 data.mul(255).clamp(0, 255).byte().permute(1, 2, 0).cpu().numpy()

wandb/sdk/integration_utils/auto_logging.py

@@ -1,4 +1,6 @@
+import asyncio
 import functools
+import inspect
 import logging
 import sys
 from typing import Any, Dict, Optional, Sequence, TypeVar
@@ -86,21 +88,44 @@ class PatchAPI:
             original = functools.reduce(getattr, symbol_parts, self.set_api)
 
             def method_factory(original_method: Any):
-                @functools.wraps(original_method)
-                def method(*args, **kwargs):
+                async def async_method(*args, **kwargs):
+                    future = asyncio.Future()
+
+                    async def callback(coro):
+                        try:
+                            result = await coro
+                            loggable_dict = self.resolver(
+                                args, kwargs, result, timer.start_time, timer.elapsed
+                            )
+                            if loggable_dict is not None:
+                                run.log(loggable_dict)
+                            future.set_result(result)
+                        except Exception as e:
+                            logger.warning(e)
+
+                    with Timer() as timer:
+                        coro = original_method(*args, **kwargs)
+                        asyncio.ensure_future(callback(coro))
+
+                    return await future
+
+                def sync_method(*args, **kwargs):
                     with Timer() as timer:
                         result = original_method(*args, **kwargs)
-                    try:
-                        loggable_dict = self.resolver(
-                            args, kwargs, result, timer.start_time, timer.elapsed
-                        )
-                        if loggable_dict is not None:
-                            run.log(loggable_dict)
-                    except Exception as e:
-                        logger.warning(e)
-                    return result
-
-                return method
+                        try:
+                            loggable_dict = self.resolver(
+                                args, kwargs, result, timer.start_time, timer.elapsed
+                            )
+                            if loggable_dict is not None:
+                                run.log(loggable_dict)
+                        except Exception as e:
+                            logger.warning(e)
+                        return result
+
+                if inspect.iscoroutinefunction(original_method):
+                    return functools.wraps(original_method)(async_method)
+                else:
+                    return functools.wraps(original_method)(sync_method)
 
             # save original method
             self.original_methods[symbol] = original

wandb/sdk/interface/interface.py

@@ -1,7 +1,6 @@
 """Interface base class - Used to send messages to the internal process.
 
 InterfaceBase: The abstract class
-InterfaceGrpc: Use gRPC to send and receive messages
 InterfaceShared: Common routines for socket and queue based implementations
 InterfaceQueue: Use multiprocessing queues to send and receive messages
 InterfaceSock: Use socket to send and receive messages
@@ -16,10 +15,11 @@ import time
 from abc import abstractmethod
 from typing import TYPE_CHECKING, Any, Dict, Iterable, NewType, Optional, Tuple, Union
 
-import wandb.sdk.lib.json_util as json
 from wandb.proto import wandb_internal_pb2 as pb
 from wandb.proto import wandb_telemetry_pb2 as tpb
+from wandb.sdk.artifacts.artifact import Artifact
 from wandb.sdk.artifacts.artifact_manifest import ArtifactManifest
+from wandb.sdk.lib import json_util as json
 from wandb.util import (
     WandBJSONEncoderOld,
     get_h5_typename,
@@ -38,8 +38,6 @@ from .message_future import MessageFuture
 GlobStr = NewType("GlobStr", str)
 
 if TYPE_CHECKING:
-    from wandb.sdk.artifacts.artifact import Artifact
-
     from ..wandb_run import Run
 
     if sys.version_info >= (3, 8):
@@ -97,21 +95,6 @@ class InterfaceBase:
     def _publish_header(self, header: pb.HeaderRecord) -> None:
         raise NotImplementedError
 
-    def communicate_check_version(
-        self, current_version: Optional[str] = None
-    ) -> Optional[pb.CheckVersionResponse]:
-        check_version = pb.CheckVersionRequest()
-        if current_version:
-            check_version.current_version = current_version
-        ret = self._communicate_check_version(check_version)
-        return ret
-
-    @abstractmethod
-    def _communicate_check_version(
-        self, current_version: pb.CheckVersionRequest
-    ) -> Optional[pb.CheckVersionResponse]:
-        raise NotImplementedError
-
     def communicate_status(self) -> Optional[pb.StatusResponse]:
         status = pb.StatusRequest()
         resp = self._communicate_status(status)
@@ -123,28 +106,6 @@ class InterfaceBase:
     ) -> Optional[pb.StatusResponse]:
         raise NotImplementedError
 
-    def communicate_stop_status(self) -> Optional[pb.StopStatusResponse]:
-        status = pb.StopStatusRequest()
-        resp = self._communicate_stop_status(status)
-        return resp
-
-    @abstractmethod
-    def _communicate_stop_status(
-        self, status: pb.StopStatusRequest
-    ) -> Optional[pb.StopStatusResponse]:
-        raise NotImplementedError
-
-    def communicate_network_status(self) -> Optional[pb.NetworkStatusResponse]:
-        status = pb.NetworkStatusRequest()
-        resp = self._communicate_network_status(status)
-        return resp
-
-    @abstractmethod
-    def _communicate_network_status(
-        self, status: pb.NetworkStatusRequest
-    ) -> Optional[pb.NetworkStatusResponse]:
-        raise NotImplementedError
-
     def _make_config(
         self,
         data: Optional[dict] = None,
@@ -214,41 +175,6 @@ class InterfaceBase:
     def _publish_metric(self, metric: pb.MetricRecord) -> None:
         raise NotImplementedError
 
-    def communicate_attach(self, attach_id: str) -> Optional[pb.AttachResponse]:
-        attach = pb.AttachRequest(attach_id=attach_id)
-        resp = self._communicate_attach(attach)
-        return resp
-
-    @abstractmethod
-    def _communicate_attach(
-        self, attach: pb.AttachRequest
-    ) -> Optional[pb.AttachResponse]:
-        raise NotImplementedError
-
-    def communicate_run(
-        self, run_obj: "Run", timeout: Optional[int] = None
-    ) -> Optional[pb.RunUpdateResult]:
-        run = self._make_run(run_obj)
-        return self._communicate_run(run, timeout=timeout)
-
-    @abstractmethod
-    def _communicate_run(
-        self, run: pb.RunRecord, timeout: Optional[int] = None
-    ) -> Optional[pb.RunUpdateResult]:
-        raise NotImplementedError
-
-    def communicate_run_start(self, run_pb: pb.RunRecord) -> bool:
-        run_start = pb.RunStartRequest()
-        run_start.run.CopyFrom(run_pb)
-        result = self._communicate_run_start(run_start)
-        return result is not None
-
-    @abstractmethod
-    def _communicate_run_start(
-        self, run_start: pb.RunStartRequest
-    ) -> Optional[pb.RunStartResponse]:
-        raise NotImplementedError
-
     def _make_summary_from_dict(self, summary_dict: dict) -> pb.SummaryRecord:
         summary = pb.SummaryRecord()
         for k, v in summary_dict.items():
@@ -336,27 +262,6 @@ class InterfaceBase:
     def _publish_summary(self, summary: pb.SummaryRecord) -> None:
         raise NotImplementedError
 
-    def communicate_get_summary(self) -> Optional[pb.GetSummaryResponse]:
-        get_summary = pb.GetSummaryRequest()
-        return self._communicate_get_summary(get_summary)
-
-    @abstractmethod
-    def _communicate_get_summary(
-        self, get_summary: pb.GetSummaryRequest
-    ) -> Optional[pb.GetSummaryResponse]:
-        raise NotImplementedError
-
-    def communicate_sampled_history(self) -> Optional[pb.SampledHistoryResponse]:
-        sampled_history = pb.SampledHistoryRequest()
-        resp = self._communicate_sampled_history(sampled_history)
-        return resp
-
-    @abstractmethod
-    def _communicate_sampled_history(
-        self, sampled_history: pb.SampledHistoryRequest
-    ) -> Optional[pb.SampledHistoryResponse]:
-        raise NotImplementedError
-
     def _make_files(self, files_dict: "FilesDict") -> pb.FilesRecord:
         files = pb.FilesRecord()
         for path, policy in files_dict["files"]:
@@ -388,6 +293,10 @@ class InterfaceBase:
             proto_artifact.metadata = json.dumps(json_friendly_val(artifact.metadata))
         if artifact._base_id:
             proto_artifact.base_id = artifact._base_id
+
+        ttl_duration_input = artifact._ttl_duration_seconds_to_gql()
+        if ttl_duration_input:
+            proto_artifact.ttl_duration_seconds = ttl_duration_input
         proto_artifact.incremental_beta1 = artifact.incremental
         self._make_artifact_manifest(artifact.manifest, obj=proto_artifact.manifest)
         return proto_artifact
@@ -565,22 +474,6 @@ class InterfaceBase:
     ) -> MessageFuture:
         raise NotImplementedError
 
-    @abstractmethod
-    def _communicate_artifact_send(
-        self, artifact_send: pb.ArtifactSendRequest
-    ) -> Optional[pb.ArtifactSendResponse]:
-        raise NotImplementedError
-
-    @abstractmethod
-    def _communicate_artifact_poll(
-        self, art_poll: pb.ArtifactPollRequest
-    ) -> Optional[pb.ArtifactPollResponse]:
-        raise NotImplementedError
-
-    @abstractmethod
-    def _publish_artifact_done(self, artifact_done: pb.ArtifactDoneRequest) -> None:
-        raise NotImplementedError
-
     def publish_artifact(
         self,
         run: "Run",
@@ -772,17 +665,6 @@ class InterfaceBase:
     def _publish_exit(self, exit_data: pb.RunExitRecord) -> None:
         raise NotImplementedError
 
-    def communicate_poll_exit(self) -> Optional[pb.PollExitResponse]:
-        poll_exit = pb.PollExitRequest()
-        resp = self._communicate_poll_exit(poll_exit)
-        return resp
-
-    @abstractmethod
-    def _communicate_poll_exit(
-        self, poll_exit: pb.PollExitRequest
-    ) -> Optional[pb.PollExitResponse]:
-        raise NotImplementedError
-
     def publish_keepalive(self) -> None:
         keepalive = pb.KeepaliveRequest()
         self._publish_keepalive(keepalive)
@@ -791,17 +673,6 @@ class InterfaceBase:
     def _publish_keepalive(self, keepalive: pb.KeepaliveRequest) -> None:
         raise NotImplementedError
 
-    def communicate_server_info(self) -> Optional[pb.ServerInfoResponse]:
-        server_info = pb.ServerInfoRequest()
-        resp = self._communicate_server_info(server_info)
-        return resp
-
-    @abstractmethod
-    def _communicate_server_info(
-        self, server_info: pb.ServerInfoRequest
-    ) -> Optional[pb.ServerInfoResponse]:
-        raise NotImplementedError
-
     def join(self) -> None:
         # Drop indicates that the internal process has already been shutdown
         if self._drop:
@@ -867,6 +738,16 @@ class InterfaceBase:
     def _deliver_network_status(self, status: pb.NetworkStatusRequest) -> MailboxHandle:
         raise NotImplementedError
 
+    def deliver_internal_messages(self) -> MailboxHandle:
+        internal_message = pb.InternalMessagesRequest()
+        return self._deliver_internal_messages(internal_message)
+
+    @abstractmethod
+    def _deliver_internal_messages(
+        self, internal_message: pb.InternalMessagesRequest
+    ) -> MailboxHandle:
+        raise NotImplementedError
+
     def deliver_get_summary(self) -> MailboxHandle:
         get_summary = pb.GetSummaryRequest()
         return self._deliver_get_summary(get_summary)

wandb/sdk/interface/interface_shared.py

@@ -124,6 +124,7 @@ class InterfaceShared(InterfaceBase):
         resume: Optional[pb.ResumeRequest] = None,
         status: Optional[pb.StatusRequest] = None,
         stop_status: Optional[pb.StopStatusRequest] = None,
+        internal_messages: Optional[pb.InternalMessagesRequest] = None,
         network_status: Optional[pb.NetworkStatusRequest] = None,
         poll_exit: Optional[pb.PollExitRequest] = None,
         partial_history: Optional[pb.PartialHistoryRequest] = None,
@@ -133,9 +134,6 @@ class InterfaceShared(InterfaceBase):
         log_artifact: Optional[pb.LogArtifactRequest] = None,
         defer: Optional[pb.DeferRequest] = None,
         attach: Optional[pb.AttachRequest] = None,
-        artifact_send: Optional[pb.ArtifactSendRequest] = None,
-        artifact_poll: Optional[pb.ArtifactPollRequest] = None,
-        artifact_done: Optional[pb.ArtifactDoneRequest] = None,
         server_info: Optional[pb.ServerInfoRequest] = None,
         keepalive: Optional[pb.KeepaliveRequest] = None,
         run_status: Optional[pb.RunStatusRequest] = None,
@@ -160,6 +158,8 @@ class InterfaceShared(InterfaceBase):
             request.status.CopyFrom(status)
         elif stop_status:
             request.stop_status.CopyFrom(stop_status)
+        elif internal_messages:
+            request.internal_messages.CopyFrom(internal_messages)
         elif network_status:
             request.network_status.CopyFrom(network_status)
         elif poll_exit:
@@ -178,12 +178,6 @@ class InterfaceShared(InterfaceBase):
             request.defer.CopyFrom(defer)
         elif attach:
             request.attach.CopyFrom(attach)
-        elif artifact_send:
-            request.artifact_send.CopyFrom(artifact_send)
-        elif artifact_poll:
-            request.artifact_poll.CopyFrom(artifact_poll)
-        elif artifact_done:
-            request.artifact_done.CopyFrom(artifact_done)
         elif server_info:
             request.server_info.CopyFrom(server_info)
         elif keepalive:
@@ -368,36 +362,6 @@ class InterfaceShared(InterfaceBase):
         rec = self._make_record(metric=metric)
         self._publish(rec)
 
-    def _communicate_attach(
-        self, attach: pb.AttachRequest
-    ) -> Optional[pb.AttachResponse]:
-        req = self._make_request(attach=attach)
-        resp = self._communicate(req)
-        if resp is None:
-            return None
-        return resp.response.attach_response
-
-    def _communicate_run(
-        self, run: pb.RunRecord, timeout: Optional[int] = None
-    ) -> Optional[pb.RunUpdateResult]:
-        """Send synchronous run object waiting for a response.
-
-        Arguments:
-            run: RunRecord object
-            timeout: number of seconds to wait
-
-        Returns:
-            RunRecord object
-        """
-        req = self._make_record(run=run)
-        resp = self._communicate(req, timeout=timeout)
-        if resp is None:
-            logger.info("couldn't get run from backend")
-            # Note: timeouts handled by callers: wandb_init.py
-            return None
-        assert resp.HasField("run_result")
-        return resp.run_result
-
     def publish_stats(self, stats_dict: dict) -> None:
         stats = self._make_stats(stats_dict)
         rec = self._make_record(stats=stats)
@@ -419,30 +383,6 @@ class InterfaceShared(InterfaceBase):
         rec = self._make_request(log_artifact=log_artifact)
         return self._communicate_async(rec)
 
-    def _communicate_artifact_send(
-        self, artifact_send: pb.ArtifactSendRequest
-    ) -> Optional[pb.ArtifactSendResponse]:
-        rec = self._make_request(artifact_send=artifact_send)
-        result = self._communicate(rec)
-        if result is None:
-            return None
-        artifact_send_resp = result.response.artifact_send_response
-        return artifact_send_resp
-
-    def _communicate_artifact_poll(
-        self, artifact_poll: pb.ArtifactPollRequest
-    ) -> Optional[pb.ArtifactPollResponse]:
-        rec = self._make_request(artifact_poll=artifact_poll)
-        result = self._communicate(rec)
-        if result is None:
-            return None
-        artifact_poll_resp = result.response.artifact_poll_response
-        return artifact_poll_resp
-
-    def _publish_artifact_done(self, artifact_done: pb.ArtifactDoneRequest) -> None:
-        rec = self._make_request(artifact_done=artifact_done)
-        self._publish(rec)
-
     def _publish_artifact(self, proto_artifact: pb.ArtifactRecord) -> None:
         rec = self._make_record(artifact=proto_artifact)
         self._publish(rec)
@@ -461,98 +401,14 @@ class InterfaceShared(InterfaceBase):
         assert resp.response.status_response
         return resp.response.status_response
 
-    def _communicate_stop_status(
-        self, status: pb.StopStatusRequest
-    ) -> Optional[pb.StopStatusResponse]:
-        req = self._make_request(stop_status=status)
-        resp = self._communicate(req, local=True)
-        if resp is None:
-            return None
-        assert resp.response.stop_status_response
-        return resp.response.stop_status_response
-
-    def _communicate_network_status(
-        self, status: pb.NetworkStatusRequest
-    ) -> Optional[pb.NetworkStatusResponse]:
-        req = self._make_request(network_status=status)
-        resp = self._communicate(req, local=True)
-        if resp is None:
-            return None
-        assert resp.response.network_status_response
-        return resp.response.network_status_response
-
     def _publish_exit(self, exit_data: pb.RunExitRecord) -> None:
         rec = self._make_record(exit=exit_data)
         self._publish(rec)
 
-    def _communicate_poll_exit(
-        self, poll_exit: pb.PollExitRequest
-    ) -> Optional[pb.PollExitResponse]:
-        rec = self._make_request(poll_exit=poll_exit)
-        result = self._communicate(rec)
-        if result is None:
-            return None
-        poll_exit_response = result.response.poll_exit_response
-        assert poll_exit_response
-        return poll_exit_response
-
     def _publish_keepalive(self, keepalive: pb.KeepaliveRequest) -> None:
         record = self._make_request(keepalive=keepalive)
         self._publish(record)
 
-    def _communicate_server_info(
-        self, server_info: pb.ServerInfoRequest
-    ) -> Optional[pb.ServerInfoResponse]:
-        rec = self._make_request(server_info=server_info)
-        result = self._communicate(rec)
-        if result is None:
-            return None
-        server_info_response = result.response.server_info_response
-        assert server_info_response
-        return server_info_response
-
-    def _communicate_check_version(
-        self, check_version: pb.CheckVersionRequest
-    ) -> Optional[pb.CheckVersionResponse]:
-        rec = self._make_request(check_version=check_version)
-        result = self._communicate(rec)
-        if result is None:
-            # Note: timeouts handled by callers: wandb_init.py
-            return None
-        return result.response.check_version_response
-
-    def _communicate_run_start(
-        self, run_start: pb.RunStartRequest
-    ) -> Optional[pb.RunStartResponse]:
-        rec = self._make_request(run_start=run_start)
-        result = self._communicate(rec)
-        if result is None:
-            return None
-        run_start_response = result.response.run_start_response
-        return run_start_response
-
-    def _communicate_get_summary(
-        self, get_summary: pb.GetSummaryRequest
-    ) -> Optional[pb.GetSummaryResponse]:
-        record = self._make_request(get_summary=get_summary)
-        result = self._communicate(record, timeout=10)
-        if result is None:
-            return None
-        get_summary_response = result.response.get_summary_response
-        assert get_summary_response
-        return get_summary_response
-
-    def _communicate_sampled_history(
-        self, sampled_history: pb.SampledHistoryRequest
-    ) -> Optional[pb.SampledHistoryResponse]:
-        record = self._make_request(sampled_history=sampled_history)
-        result = self._communicate(record)
-        if result is None:
-            return None
-        sampled_history_response = result.response.sampled_history_response
-        assert sampled_history_response
-        return sampled_history_response
-
     def _communicate_shutdown(self) -> None:
         # shutdown
         request = pb.Request(shutdown=pb.ShutdownRequest())
@@ -609,6 +465,12 @@ class InterfaceShared(InterfaceBase):
         record = self._make_request(network_status=network_status)
         return self._deliver_record(record)
 
+    def _deliver_internal_messages(
+        self, internal_message: pb.InternalMessagesRequest
+    ) -> MailboxHandle:
+        record = self._make_request(internal_messages=internal_message)
+        return self._deliver_record(record)
+
     def _deliver_request_server_info(
         self, server_info: pb.ServerInfoRequest
     ) -> MailboxHandle:

wandb/sdk/interface/interface_sock.py

@@ -59,29 +59,3 @@ class InterfaceSock(InterfaceShared):
             raise Exception("The wandb backend process has shutdown")
         future = self._router.send_and_receive(rec, local=local)
         return future
-
-    def _communicate_stop_status(
-        self, status: "pb.StopStatusRequest"
-    ) -> Optional["pb.StopStatusResponse"]:
-        # Message stop_status is called from a daemon thread started by wandb_run
-        # The underlying socket might go away while the thread is still running.
-        # Handle this like a timed-out message as the daemon thread will eventually
-        # be killed.
-        try:
-            data = super()._communicate_stop_status(status)
-        except BrokenPipeError:
-            data = None
-        return data
-
-    def _communicate_network_status(
-        self, status: "pb.NetworkStatusRequest"
-    ) -> Optional["pb.NetworkStatusResponse"]:
-        # Message network_status is called from a daemon thread started by wandb_run
-        # The underlying socket might go away while the thread is still running.
-        # Handle this like a timed-out message as the daemon thread will eventually
-        # be killed.
-        try:
-            data = super()._communicate_network_status(status)
-        except BrokenPipeError:
-            data = None
-        return data