wandb 0.19.6rc4__py3-none-any.whl → 0.19.7__py3-none-any.whl

wandb/sdk/data_types/saved_model.py

@@ -1,18 +1,9 @@
+from __future__ import annotations
+
 import os
 import shutil
 import sys
-from typing import (
-    TYPE_CHECKING,
-    Any,
-    ClassVar,
-    Generic,
-    List,
-    Optional,
-    Type,
-    TypeVar,
-    Union,
-    cast,
-)
+from typing import TYPE_CHECKING, Any, ClassVar, Generic, TypeVar, cast
 
 import wandb
 from wandb import util
@@ -23,22 +14,24 @@ from wandb.sdk.lib.paths import LogicalPath
 from ._private import MEDIA_TMP
 from .base_types.wb_value import WBValue
 
-if TYPE_CHECKING:  # pragma: no cover
+if TYPE_CHECKING:
+    from types import ModuleType
+
     import cloudpickle  # type: ignore
     import sklearn  # type: ignore
     import tensorflow  # type: ignore
     import torch  # type: ignore
+    from typing_extensions import Self
 
     from wandb.sdk.artifacts.artifact import Artifact
-
-    from ..wandb_run import Run as LocalRun
+    from wandb.sdk.wandb_run import Run as LocalRun
 
 
 DEBUG_MODE = False
 
 
 def _add_deterministic_dir_to_artifact(
-    artifact: "Artifact", dir_name: str, target_dir_root: str
+    artifact: Artifact, dir_name: str, target_dir_root: str
 ) -> str:
     file_paths = []
     for dirpath, _, filenames in os.walk(dir_name, topdown=True):
@@ -50,7 +43,7 @@ def _add_deterministic_dir_to_artifact(
     return target_path
 
 
-def _load_dir_from_artifact(source_artifact: "Artifact", path: str) -> str:
+def _load_dir_from_artifact(source_artifact: Artifact, path: str) -> str:
     dl_path = None
 
     # Look through the entire manifest to find all of the files in the directory.
@@ -79,14 +72,12 @@ class _SavedModel(WBValue, Generic[SavedModelObjType]):
     _log_type: ClassVar[str]
     _path_extension: ClassVar[str]
 
-    _model_obj: Optional["SavedModelObjType"]
-    _path: Optional[str]
-    _input_obj_or_path: Union[SavedModelObjType, str]
+    _model_obj: SavedModelObjType | None
+    _path: str | None
+    _input_obj_or_path: SavedModelObjType | str
 
     # Public Methods
-    def __init__(
-        self, obj_or_path: Union[SavedModelObjType, str], **kwargs: Any
-    ) -> None:
+    def __init__(self, obj_or_path: SavedModelObjType | str, **kwargs: Any) -> None:
         super().__init__()
         if self.__class__ == _SavedModel:
             raise TypeError(
@@ -115,7 +106,7 @@ class _SavedModel(WBValue, Generic[SavedModelObjType]):
             self._unset_obj()
 
     @staticmethod
-    def init(obj_or_path: Any, **kwargs: Any) -> "_SavedModel":
+    def init(obj_or_path: Any, **kwargs: Any) -> _SavedModel:
         maybe_instance = _SavedModel._maybe_init(obj_or_path, **kwargs)
         if maybe_instance is None:
             raise ValueError(
@@ -125,8 +116,8 @@ class _SavedModel(WBValue, Generic[SavedModelObjType]):
 
     @classmethod
     def from_json(
-        cls: Type["_SavedModel"], json_obj: dict, source_artifact: "Artifact"
-    ) -> "_SavedModel":
+        cls: type[_SavedModel], json_obj: dict, source_artifact: Artifact
+    ) -> _SavedModel:
         path = json_obj["path"]
 
         # First, if the entry is a file, the download it.
@@ -143,7 +134,7 @@ class _SavedModel(WBValue, Generic[SavedModelObjType]):
         # and specified adapter.
         return cls(dl_path)
 
-    def to_json(self, run_or_artifact: Union["LocalRun", "Artifact"]) -> dict:
+    def to_json(self, run_or_artifact: LocalRun | Artifact) -> dict:
         # Unlike other data types, we do not allow adding to a Run directly. There is a
         # bit of tech debt in the other data types which requires the input to `to_json`
         # to accept a Run or Artifact. However, Run additions should be deprecated in the future.
@@ -218,8 +209,8 @@ class _SavedModel(WBValue, Generic[SavedModelObjType]):
     # Private Class Methods
     @classmethod
     def _maybe_init(
-        cls: Type["_SavedModel"], obj_or_path: Any, **kwargs: Any
-    ) -> Optional["_SavedModel"]:
+        cls: type[_SavedModel], obj_or_path: Any, **kwargs: Any
+    ) -> _SavedModel | None:
         # _maybe_init is an exception-safe method that will return an instance of this class
         # (or any subclass of this class - recursively) OR None if no subclass constructor is found.
         # We first try the current class, then recursively call this method on children classes. This pattern
@@ -241,7 +232,7 @@ class _SavedModel(WBValue, Generic[SavedModelObjType]):
         return None
 
     @classmethod
-    def _tmp_path(cls: Type["_SavedModel"]) -> str:
+    def _tmp_path(cls: type[_SavedModel]) -> str:
         # Generates a tmp path under our MEDIA_TMP directory which confirms to the file
         # or folder preferences of the class.
         assert isinstance(cls._path_extension, str), "_path_extension must be a string"
@@ -286,13 +277,13 @@ PicklingSavedModelObjType = TypeVar("PicklingSavedModelObjType")
 
 
 class _PicklingSavedModel(_SavedModel[SavedModelObjType]):
-    _dep_py_files: Optional[List[str]] = None
-    _dep_py_files_path: Optional[str] = None
+    _dep_py_files: list[str] | None = None
+    _dep_py_files_path: str | None = None
 
     def __init__(
         self,
-        obj_or_path: Union[SavedModelObjType, str],
-        dep_py_files: Optional[List[str]] = None,
+        obj_or_path: SavedModelObjType | str,
+        dep_py_files: list[str] | None = None,
     ):
         super().__init__(obj_or_path)
         if self.__class__ == _PicklingSavedModel:
@@ -319,9 +310,7 @@ class _PicklingSavedModel(_SavedModel[SavedModelObjType]):
                     raise ValueError(f"Invalid dependency file: {extra_file}")
 
     @classmethod
-    def from_json(
-        cls: Type["_SavedModel"], json_obj: dict, source_artifact: "Artifact"
-    ) -> "_PicklingSavedModel":
+    def from_json(cls, json_obj: dict, source_artifact: Artifact) -> Self:
         backup_path = [p for p in sys.path]
         if (
             "dep_py_files_path" in json_obj
@@ -337,7 +326,7 @@ class _PicklingSavedModel(_SavedModel[SavedModelObjType]):
 
         return inst  # type: ignore
 
-    def to_json(self, run_or_artifact: Union["LocalRun", "Artifact"]) -> dict:
+    def to_json(self, run_or_artifact: LocalRun | Artifact) -> dict:
         json_obj = super().to_json(run_or_artifact)
         assert isinstance(run_or_artifact, wandb.Artifact)
         if self._dep_py_files_path is not None:
@@ -361,7 +350,7 @@ class _PytorchSavedModel(_PicklingSavedModel["torch.nn.Module"]):
     _path_extension = "pt"
 
     @staticmethod
-    def _deserialize(dir_or_file_path: str) -> "torch.nn.Module":
+    def _deserialize(dir_or_file_path: str) -> torch.nn.Module:
         return _get_torch().load(dir_or_file_path, weights_only=False)
 
     @staticmethod
@@ -369,7 +358,7 @@ class _PytorchSavedModel(_PicklingSavedModel["torch.nn.Module"]):
         return isinstance(obj, _get_torch().nn.Module)
 
     @staticmethod
-    def _serialize(model_obj: "torch.nn.Module", dir_or_file_path: str) -> None:
+    def _serialize(model_obj: torch.nn.Module, dir_or_file_path: str) -> None:
         _get_torch().save(
             model_obj,
             dir_or_file_path,
@@ -391,7 +380,7 @@ class _SklearnSavedModel(_PicklingSavedModel["sklearn.base.BaseEstimator"]):
     @staticmethod
     def _deserialize(
         dir_or_file_path: str,
-    ) -> "sklearn.base.BaseEstimator":
+    ) -> sklearn.base.BaseEstimator:
         with open(dir_or_file_path, "rb") as file:
             model = _get_cloudpickle().load(file)
         return model
@@ -410,16 +399,16 @@ class _SklearnSavedModel(_PicklingSavedModel["sklearn.base.BaseEstimator"]):
 
     @staticmethod
     def _serialize(
-        model_obj: "sklearn.base.BaseEstimator", dir_or_file_path: str
+        model_obj: sklearn.base.BaseEstimator, dir_or_file_path: str
     ) -> None:
         dynamic_cloudpickle = _get_cloudpickle()
         with open(dir_or_file_path, "wb") as file:
             dynamic_cloudpickle.dump(model_obj, file)
 
 
-def _get_tf_keras() -> "tensorflow.keras":
+def _get_tf_keras() -> ModuleType:
     return cast(
-        "tensorflow",
+        ModuleType,
         util.get_module("tensorflow", "ModelAdapter requires `tensorflow`"),
     ).keras
 
@@ -431,7 +420,7 @@ class _TensorflowKerasSavedModel(_SavedModel["tensorflow.keras.Model"]):
     @staticmethod
     def _deserialize(
         dir_or_file_path: str,
-    ) -> "tensorflow.keras.Model":
+    ) -> tensorflow.keras.Model:
         return _get_tf_keras().models.load_model(dir_or_file_path)
 
     @staticmethod
@@ -439,7 +428,7 @@ class _TensorflowKerasSavedModel(_SavedModel["tensorflow.keras.Model"]):
         return isinstance(obj, _get_tf_keras().models.Model)
 
     @staticmethod
-    def _serialize(model_obj: "tensorflow.keras.Model", dir_or_file_path: str) -> None:
+    def _serialize(model_obj: tensorflow.keras.Model, dir_or_file_path: str) -> None:
         _get_tf_keras().models.save_model(
             model_obj, dir_or_file_path, include_optimizer=True
         )

wandb/sdk/data_types/video.py

@@ -71,10 +71,10 @@ class Video(BatchableMedia):
         import numpy as np
         import wandb
 
-        wandb.init()
+        run = wandb.init()
         # axes are (time, channel, height, width)
         frames = np.random.randint(low=0, high=256, size=(10, 3, 100, 100), dtype=np.uint8)
-        wandb.log({"video": wandb.Video(frames, fps=4)})
+        run.log({"video": wandb.Video(frames, fps=4)})
         ```
     """
 
@@ -138,20 +138,11 @@ class Video(BatchableMedia):
             self.encode(fps=fps)
 
     def encode(self, fps: int = 4) -> None:
-        # Try to import ImageSequenceClip from the appropriate MoviePy module
-        mpy = None
-        try:
-            # Attempt to load moviepy.editor for MoviePy < 2.0
-            mpy = util.get_module(
-                "moviepy.editor",
-                required='wandb.Video requires moviepy when passing raw data. Install with "pip install wandb[media]"',
-            )
-        except wandb.Error:
-            # Fallback to moviepy for MoviePy >= 2.0
-            mpy = util.get_module(
-                "moviepy",
-                required='wandb.Video requires moviepy when passing raw data. Install with "pip install wandb[media]"',
-            )
+        # import ImageSequenceClip from the appropriate MoviePy module
+        mpy = util.get_module(
+            "moviepy.video.io.ImageSequenceClip",
+            required='wandb.Video requires moviepy when passing raw data. Install with "pip install wandb[media]"',
+        )
 
         tensor = self._prepare_video(self.data)
         _, self._height, self._width, self._channels = tensor.shape  # type: ignore

wandb/sdk/interface/interface.py

@@ -35,6 +35,7 @@ from wandb.sdk.artifacts.artifact import Artifact
 from wandb.sdk.artifacts.artifact_manifest import ArtifactManifest
 from wandb.sdk.artifacts.staging import get_staging_dir
 from wandb.sdk.lib import json_util as json
+from wandb.sdk.mailbox import HandleAbandonedError, MailboxHandle
 from wandb.util import (
     WandBJSONEncoderOld,
     get_h5_typename,
@@ -46,9 +47,7 @@ from wandb.util import (
 )
 
 from ..data_types.utils import history_dict_to_json, val_to_json
-from ..lib.mailbox import MailboxHandle
 from . import summary_record as sr
-from .message_future import MessageFuture
 
 MANIFEST_FILE_SIZE_THRESHOLD = 100_000
 
@@ -561,7 +560,7 @@ class InterfaceBase:
     def _publish_use_artifact(self, proto_artifact: pb.UseArtifactRecord) -> None:
         raise NotImplementedError
 
-    def communicate_artifact(
+    def deliver_artifact(
         self,
         run: "Run",
         artifact: "Artifact",
@@ -571,7 +570,7 @@ class InterfaceBase:
         is_user_created: bool = False,
         use_after_commit: bool = False,
         finalize: bool = True,
-    ) -> MessageFuture:
+    ) -> MailboxHandle:
         proto_run = self._make_run(run)
         proto_artifact = self._make_artifact(artifact)
         proto_artifact.run_id = proto_run.run_id
@@ -589,13 +588,14 @@ class InterfaceBase:
         if history_step is not None:
             log_artifact.history_step = history_step
         log_artifact.staging_dir = get_staging_dir()
-        resp = self._communicate_artifact(log_artifact)
+        resp = self._deliver_artifact(log_artifact)
         return resp
 
     @abstractmethod
-    def _communicate_artifact(
-        self, log_artifact: pb.LogArtifactRequest
-    ) -> MessageFuture:
+    def _deliver_artifact(
+        self,
+        log_artifact: pb.LogArtifactRequest,
+    ) -> MailboxHandle:
         raise NotImplementedError
 
     def deliver_download_artifact(
@@ -877,10 +877,22 @@ class InterfaceBase:
         # Drop indicates that the internal process has already been shutdown
         if self._drop:
             return
-        _ = self._communicate_shutdown()
+
+        handle = self._deliver_shutdown()
+
+        try:
+            handle.wait_or(timeout=30)
+        except TimeoutError:
+            # This can happen if the server fails to respond due to a bug
+            # or due to being very busy.
+            logger.warning("timed out communicating shutdown")
+        except HandleAbandonedError:
+            # This can happen if the connection to the server is closed
+            # before a response is read.
+            logger.warning("handle abandoned while communicating shutdown")
 
     @abstractmethod
-    def _communicate_shutdown(self) -> None:
+    def _deliver_shutdown(self) -> MailboxHandle:
         raise NotImplementedError
 
     def deliver_run(self, run: "Run") -> MailboxHandle:
@@ -979,6 +991,10 @@ class InterfaceBase:
     def _deliver_exit(self, exit_data: pb.RunExitRecord) -> MailboxHandle:
         raise NotImplementedError
 
+    @abstractmethod
+    def deliver_operation_stats(self) -> MailboxHandle:
+        raise NotImplementedError
+
     def deliver_poll_exit(self) -> MailboxHandle:
         poll_exit = pb.PollExitRequest()
         return self._deliver_poll_exit(poll_exit)

wandb/sdk/interface/interface_queue.py

@@ -8,7 +8,8 @@ import logging
 from multiprocessing.process import BaseProcess
 from typing import TYPE_CHECKING, Optional
 
-from ..lib.mailbox import Mailbox
+from wandb.sdk.mailbox import Mailbox
+
 from .interface_shared import InterfaceShared
 from .router_queue import MessageQueueRouter
 
@@ -22,21 +23,17 @@ logger = logging.getLogger("wandb")
 
 
 class InterfaceQueue(InterfaceShared):
-    record_q: Optional["Queue[pb.Record]"]
-    result_q: Optional["Queue[pb.Result]"]
-    _mailbox: Optional[Mailbox]
-
     def __init__(
         self,
         record_q: Optional["Queue[pb.Record]"] = None,
         result_q: Optional["Queue[pb.Result]"] = None,
         process: Optional[BaseProcess] = None,
-        process_check: bool = True,
         mailbox: Optional[Mailbox] = None,
     ) -> None:
         self.record_q = record_q
         self.result_q = result_q
-        super().__init__(process=process, process_check=process_check, mailbox=mailbox)
+        self._process = process
+        super().__init__(mailbox=mailbox)
 
     def _init_router(self) -> None:
         if self.record_q and self.result_q:
@@ -45,7 +42,7 @@ class InterfaceQueue(InterfaceShared):
             )
 
     def _publish(self, record: "pb.Record", local: Optional[bool] = None) -> None:
-        if self._process_check and self._process and not self._process.is_alive():
+        if self._process and not self._process.is_alive():
             raise Exception("The wandb backend process has shutdown")
         if local:
             record.control.local = local

wandb/sdk/interface/interface_relay.py

@@ -5,12 +5,11 @@ See interface.py for how interface classes relate to each other.
 """
 
 import logging
-from multiprocessing.process import BaseProcess
 from typing import TYPE_CHECKING, Optional
 
 from wandb.proto import wandb_internal_pb2 as pb
+from wandb.sdk.mailbox import Mailbox
 
-from ..lib.mailbox import Mailbox
 from .interface_queue import InterfaceQueue
 from .router_relay import MessageRelayRouter
 
@@ -31,15 +30,11 @@ class InterfaceRelay(InterfaceQueue):
         record_q: Optional["Queue[pb.Record]"] = None,
         result_q: Optional["Queue[pb.Result]"] = None,
         relay_q: Optional["Queue[pb.Result]"] = None,
-        process: Optional[BaseProcess] = None,
-        process_check: bool = True,
     ) -> None:
         self.relay_q = relay_q
         super().__init__(
             record_q=record_q,
             result_q=result_q,
-            process=process,
-            process_check=process_check,
             mailbox=mailbox,
         )
 

wandb/sdk/interface/interface_shared.py

@@ -5,44 +5,27 @@ 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[Any] = 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()
 
@@ -50,20 +33,6 @@ class InterfaceShared(InterfaceBase):
     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()
         rec.output.CopyFrom(outdata)
@@ -114,15 +83,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 +92,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 +121,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 +135,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 +270,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 +292,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,9 +339,12 @@ 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:
         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
@@ -449,11 +381,10 @@ class InterfaceShared(InterfaceBase):
         record = self._make_request(keepalive=keepalive)
         self._publish(record)
 
-    def _communicate_shutdown(self) -> None:
-        # shutdown
+    def _deliver_shutdown(self) -> MailboxHandle:
         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
@@ -462,7 +393,10 @@ class InterfaceShared(InterfaceBase):
 
     def _deliver_record(self, record: pb.Record) -> MailboxHandle:
         mailbox = self._get_mailbox()
-        handle = mailbox._deliver_record(record, interface=self)
+
+        handle = mailbox.require_response(record)
+        self._publish(record)
+
         return handle
 
     def _deliver_run(self, run: pb.RunRecord) -> MailboxHandle:
@@ -497,6 +431,10 @@ class InterfaceShared(InterfaceBase):
         record = self._make_record(exit=exit_data)
         return self._deliver_record(record)
 
+    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:
         record = self._make_request(poll_exit=poll_exit)
         return self._deliver_record(record)
@@ -539,22 +477,6 @@ class InterfaceShared(InterfaceBase):
         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()