flwr 1.18.0__py3-none-any.whl → 1.20.0__py3-none-any.whl

flwr/cli/run/run.py

@@ -24,9 +24,8 @@ from typing import Annotated, Any, Optional
 import typer
 from rich.console import Console
 
-from flwr.cli.build import build
+from flwr.cli.build import build_fab, get_fab_filename
 from flwr.cli.config_utils import (
-    get_fab_metadata,
     load_and_validate,
     process_loaded_project_config,
     validate_federation_in_project_config,
@@ -34,6 +33,7 @@ from flwr.cli.config_utils import (
 from flwr.cli.constant import FEDERATION_CONFIG_HELP_MESSAGE
 from flwr.common.config import (
     flatten_dict,
+    get_metadata_from_config,
     parse_config_args,
     user_config_to_configrecord,
 )
@@ -45,11 +45,7 @@ from flwr.proto.exec_pb2 import StartRunRequest  # pylint: disable=E0611
 from flwr.proto.exec_pb2_grpc import ExecStub
 
 from ..log import start_stream
-from ..utils import (
-    init_channel,
-    try_obtain_cli_auth_plugin,
-    unauthenticated_exc_handler,
-)
+from ..utils import flwr_cli_grpc_exc_handler, init_channel, try_obtain_cli_auth_plugin
 
 CONN_REFRESH_PERIOD = 60  # Connection refresh period for log streaming (seconds)
 
@@ -154,54 +150,57 @@ def _run_with_exec_api(
     stream: bool,
     output_format: str,
 ) -> None:
-    auth_plugin = try_obtain_cli_auth_plugin(app, federation, federation_config)
-    channel = init_channel(app, federation_config, auth_plugin)
-    stub = ExecStub(channel)
-
-    fab_path, fab_hash = build(app)
-    content = Path(fab_path).read_bytes()
-    fab_id, fab_version = get_fab_metadata(Path(fab_path))
+    channel = None
+    try:
+        auth_plugin = try_obtain_cli_auth_plugin(app, federation, federation_config)
+        channel = init_channel(app, federation_config, auth_plugin)
+        stub = ExecStub(channel)
 
-    # Delete FAB file once the bytes is computed
-    Path(fab_path).unlink()
+        fab_bytes, fab_hash, config = build_fab(app)
+        fab_id, fab_version = get_metadata_from_config(config)
 
-    fab = Fab(fab_hash, content)
+        fab = Fab(fab_hash, fab_bytes)
 
-    # Construct a `ConfigRecord` out of a flattened `UserConfig`
-    fed_conf = flatten_dict(federation_config.get("options", {}))
-    c_record = user_config_to_configrecord(fed_conf)
+        # Construct a `ConfigRecord` out of a flattened `UserConfig`
+        fed_config = flatten_dict(federation_config.get("options", {}))
+        c_record = user_config_to_configrecord(fed_config)
 
-    req = StartRunRequest(
-        fab=fab_to_proto(fab),
-        override_config=user_config_to_proto(parse_config_args(config_overrides)),
-        federation_options=config_record_to_proto(c_record),
-    )
-    with unauthenticated_exc_handler():
-        res = stub.StartRun(req)
-
-    if res.HasField("run_id"):
-        typer.secho(f"🎊 Successfully started run {res.run_id}", fg=typer.colors.GREEN)
-    else:
-        typer.secho("❌ Failed to start run", fg=typer.colors.RED)
-        raise typer.Exit(code=1)
-
-    if output_format == CliOutputFormat.JSON:
-        run_output = json.dumps(
-            {
-                "success": res.HasField("run_id"),
-                "run-id": res.run_id if res.HasField("run_id") else None,
-                "fab-id": fab_id,
-                "fab-name": fab_id.rsplit("/", maxsplit=1)[-1],
-                "fab-version": fab_version,
-                "fab-hash": fab_hash[:8],
-                "fab-filename": fab_path,
-            }
+        req = StartRunRequest(
+            fab=fab_to_proto(fab),
+            override_config=user_config_to_proto(parse_config_args(config_overrides)),
+            federation_options=config_record_to_proto(c_record),
         )
-        restore_output()
-        Console().print_json(run_output)
+        with flwr_cli_grpc_exc_handler():
+            res = stub.StartRun(req)
 
-    if stream:
-        start_stream(res.run_id, channel, CONN_REFRESH_PERIOD)
+        if res.HasField("run_id"):
+            typer.secho(
+                f"🎊 Successfully started run {res.run_id}", fg=typer.colors.GREEN
+            )
+        else:
+            typer.secho("❌ Failed to start run", fg=typer.colors.RED)
+            raise typer.Exit(code=1)
+
+        if output_format == CliOutputFormat.JSON:
+            run_output = json.dumps(
+                {
+                    "success": res.HasField("run_id"),
+                    "run-id": res.run_id if res.HasField("run_id") else None,
+                    "fab-id": fab_id,
+                    "fab-name": fab_id.rsplit("/", maxsplit=1)[-1],
+                    "fab-version": fab_version,
+                    "fab-hash": fab_hash[:8],
+                    "fab-filename": get_fab_filename(config, fab_hash),
+                }
+            )
+            restore_output()
+            Console().print_json(run_output)
+
+        if stream:
+            start_stream(res.run_id, channel, CONN_REFRESH_PERIOD)
+    finally:
+        if channel:
+            channel.close()
 
 
 def _run_without_exec_api(

flwr/cli/stop.py

@@ -35,7 +35,7 @@ from flwr.common.logger import print_json_error, redirect_output, restore_output
 from flwr.proto.exec_pb2 import StopRunRequest, StopRunResponse  # pylint: disable=E0611
 from flwr.proto.exec_pb2_grpc import ExecStub
 
-from .utils import init_channel, try_obtain_cli_auth_plugin, unauthenticated_exc_handler
+from .utils import flwr_cli_grpc_exc_handler, init_channel, try_obtain_cli_auth_plugin
 
 
 def stop(  # pylint: disable=R0914
@@ -122,7 +122,7 @@ def stop(  # pylint: disable=R0914
 
 def _stop_run(stub: ExecStub, run_id: int, output_format: str) -> None:
     """Stop a run."""
-    with unauthenticated_exc_handler():
+    with flwr_cli_grpc_exc_handler():
         response: StopRunResponse = stub.StopRun(request=StopRunRequest(run_id=run_id))
     if response.success:
         typer.secho(f"✅ Run {run_id} successfully stopped.", fg=typer.colors.GREEN)

flwr/cli/utils.py

@@ -28,7 +28,12 @@ import typer
 
 from flwr.cli.cli_user_auth_interceptor import CliUserAuthInterceptor
 from flwr.common.auth_plugin import CliAuthPlugin
-from flwr.common.constant import AUTH_TYPE_JSON_KEY, CREDENTIALS_DIR, FLWR_DIR
+from flwr.common.constant import (
+    AUTH_TYPE_JSON_KEY,
+    CREDENTIALS_DIR,
+    FLWR_DIR,
+    RUN_ID_NOT_FOUND_MESSAGE,
+)
 from flwr.common.grpc import (
     GRPC_MAX_MESSAGE_LENGTH,
     create_channel,
@@ -288,11 +293,12 @@ def init_channel(
 
 
 @contextmanager
-def unauthenticated_exc_handler() -> Iterator[None]:
-    """Context manager to handle gRPC UNAUTHENTICATED errors.
+def flwr_cli_grpc_exc_handler() -> Iterator[None]:
+    """Context manager to handle specific gRPC errors.
 
-    It catches grpc.RpcError exceptions with UNAUTHENTICATED status, informs the user,
-    and exits the application. All other exceptions will be allowed to escape.
+    It catches grpc.RpcError exceptions with UNAUTHENTICATED, UNIMPLEMENTED,
+    UNAVAILABLE, and PERMISSION_DENIED statuses, informs the user, and exits the
+    application. All other exceptions will be allowed to escape.
     """
     try:
         yield
@@ -312,4 +318,31 @@ def unauthenticated_exc_handler() -> Iterator[None]:
                 bold=True,
             )
             raise typer.Exit(code=1) from None
+        if e.code() == grpc.StatusCode.PERMISSION_DENIED:
+            typer.secho(
+                "❌ Permission denied.",
+                fg=typer.colors.RED,
+                bold=True,
+            )
+            # pylint: disable=E1101
+            typer.secho(e.details(), fg=typer.colors.RED, bold=True)
+            raise typer.Exit(code=1) from None
+        if e.code() == grpc.StatusCode.UNAVAILABLE:
+            typer.secho(
+                "Connection to the SuperLink is unavailable. Please check your network "
+                "connection and 'address' in the federation configuration.",
+                fg=typer.colors.RED,
+                bold=True,
+            )
+            raise typer.Exit(code=1) from None
+        if (
+            e.code() == grpc.StatusCode.NOT_FOUND
+            and e.details() == RUN_ID_NOT_FOUND_MESSAGE
+        ):
+            typer.secho(
+                "❌ Run ID not found.",
+                fg=typer.colors.RED,
+                bold=True,
+            )
+            raise typer.Exit(code=1) from None
         raise

flwr/client/__init__.py

@@ -15,8 +15,8 @@
 """Flower client."""
 
 
-from .app import start_client as start_client
-from .app import start_numpy_client as start_numpy_client
+from ..compat.client.app import start_client as start_client  # Deprecated
+from ..compat.client.app import start_numpy_client as start_numpy_client  # Deprecated
 from .client import Client as Client
 from .client_app import ClientApp as ClientApp
 from .numpy_client import NumPyClient as NumPyClient

flwr/client/client_app.py

@@ -20,6 +20,7 @@ from collections.abc import Iterator
 from contextlib import contextmanager
 from typing import Callable, Optional
 
+from flwr.app.metadata import validate_message_type
 from flwr.client.client import Client
 from flwr.client.message_handler.message_handler import (
     handle_legacy_message_from_msgtype,
@@ -28,7 +29,6 @@ from flwr.client.mod.utils import make_ffn
 from flwr.client.typing import ClientFnExt, Mod
 from flwr.common import Context, Message, MessageType
 from flwr.common.logger import warn_deprecated_feature
-from flwr.common.message import validate_message_type
 
 from .typing import ClientAppCallable
 

flwr/client/clientapp/__init__.py

@@ -13,10 +13,3 @@
 # limitations under the License.
 # ==============================================================================
 """Flower AppIO service."""
-
-
-from .app import flwr_clientapp as flwr_clientapp
-
-__all__ = [
-    "flwr_clientapp",
-]

flwr/client/grpc_adapter_client/connection.py

@@ -29,6 +29,7 @@ from flwr.common.logger import log
 from flwr.common.message import Message
 from flwr.common.retry_invoker import RetryInvoker
 from flwr.common.typing import Fab, Run
+from flwr.proto.message_pb2 import ObjectTree  # pylint: disable=E0611
 
 
 @contextmanager
@@ -43,12 +44,15 @@ def grpc_adapter(  # pylint: disable=R0913,too-many-positional-arguments
     ] = None,
 ) -> Iterator[
     tuple[
-        Callable[[], Optional[Message]],
-        Callable[[Message], None],
-        Optional[Callable[[], Optional[int]]],
-        Optional[Callable[[], None]],
-        Optional[Callable[[int], Run]],
-        Optional[Callable[[str, int], Fab]],
+        Callable[[], Optional[tuple[Message, ObjectTree]]],
+        Callable[[Message, ObjectTree], set[str]],
+        Callable[[], Optional[int]],
+        Callable[[], None],
+        Callable[[int], Run],
+        Callable[[str, int], Fab],
+        Callable[[int, str], bytes],
+        Callable[[int, str, bytes], None],
+        Callable[[int, str], None],
     ]
 ]:
     """Primitives for request/response-based interaction with a server via GrpcAdapter.
@@ -77,12 +81,15 @@ def grpc_adapter(  # pylint: disable=R0913,too-many-positional-arguments
 
     Returns
     -------
-    receive : Callable
-    send : Callable
+    receive : Callable[[], Optional[tuple[Message, ObjectTree]]]
+    send : Callable[[Message, ObjectTree], set[str]]
     create_node : Optional[Callable]
     delete_node : Optional[Callable]
     get_run : Optional[Callable]
     get_fab : Optional[Callable]
+    pull_object : Callable[[str], bytes]
+    push_object : Callable[[str, bytes], None]
+    confirm_message_received : Callable[[str], None]
     """
     if authentication_keys is not None:
         log(ERROR, "Client authentication is not supported for this transport type.")

flwr/client/grpc_rere_client/connection.py

@@ -15,11 +15,8 @@
 """Contextmanager for a gRPC request-response channel to the Flower server."""
 
 
-import random
-import threading
 from collections.abc import Iterator, Sequence
 from contextlib import contextmanager
-from copy import copy
 from logging import ERROR
 from pathlib import Path
 from typing import Callable, Optional, Union, cast
@@ -27,19 +24,18 @@ from typing import Callable, Optional, Union, cast
 import grpc
 from cryptography.hazmat.primitives.asymmetric import ec
 
-from flwr.client.heartbeat import start_ping_loop
-from flwr.client.message_handler.message_handler import validate_out_message
 from flwr.common import GRPC_MAX_MESSAGE_LENGTH
-from flwr.common.constant import (
-    PING_BASE_MULTIPLIER,
-    PING_CALL_TIMEOUT,
-    PING_DEFAULT_INTERVAL,
-    PING_RANDOM_RANGE,
-)
+from flwr.common.constant import HEARTBEAT_CALL_TIMEOUT, HEARTBEAT_DEFAULT_INTERVAL
 from flwr.common.grpc import create_channel, on_channel_state_change
+from flwr.common.heartbeat import HeartbeatSender
+from flwr.common.inflatable_protobuf_utils import (
+    make_confirm_message_received_fn_protobuf,
+    make_pull_object_fn_protobuf,
+    make_push_object_fn_protobuf,
+)
 from flwr.common.logger import log
-from flwr.common.message import Message, Metadata
-from flwr.common.retry_invoker import RetryInvoker
+from flwr.common.message import Message, remove_content_from_message
+from flwr.common.retry_invoker import RetryInvoker, _wrap_stub
 from flwr.common.secure_aggregation.crypto.symmetric_encryption import (
     generate_key_pairs,
 )
@@ -49,13 +45,17 @@ from flwr.proto.fab_pb2 import GetFabRequest, GetFabResponse  # pylint: disable=
 from flwr.proto.fleet_pb2 import (  # pylint: disable=E0611
     CreateNodeRequest,
     DeleteNodeRequest,
-    PingRequest,
-    PingResponse,
     PullMessagesRequest,
     PullMessagesResponse,
     PushMessagesRequest,
+    PushMessagesResponse,
 )
 from flwr.proto.fleet_pb2_grpc import FleetStub  # pylint: disable=E0611
+from flwr.proto.heartbeat_pb2 import (  # pylint: disable=E0611
+    SendNodeHeartbeatRequest,
+    SendNodeHeartbeatResponse,
+)
+from flwr.proto.message_pb2 import ObjectTree  # pylint: disable=E0611
 from flwr.proto.node_pb2 import Node  # pylint: disable=E0611
 from flwr.proto.run_pb2 import GetRunRequest, GetRunResponse  # pylint: disable=E0611
 
@@ -76,12 +76,15 @@ def grpc_request_response(  # pylint: disable=R0913,R0914,R0915,R0917
     adapter_cls: Optional[Union[type[FleetStub], type[GrpcAdapter]]] = None,
 ) -> Iterator[
     tuple[
-        Callable[[], Optional[Message]],
-        Callable[[Message], None],
-        Optional[Callable[[], Optional[int]]],
-        Optional[Callable[[], None]],
-        Optional[Callable[[int], Run]],
-        Optional[Callable[[str, int], Fab]],
+        Callable[[], Optional[tuple[Message, ObjectTree]]],
+        Callable[[Message, ObjectTree], set[str]],
+        Callable[[], Optional[int]],
+        Callable[[], None],
+        Callable[[int], Run],
+        Callable[[str, int], Fab],
+        Callable[[int, str], bytes],
+        Callable[[int, str, bytes], None],
+        Callable[[int, str], None],
     ]
 ]:
     """Primitives for request/response-based interaction with a server.
@@ -124,6 +127,9 @@ def grpc_request_response(  # pylint: disable=R0913,R0914,R0915,R0917
     create_node : Optional[Callable]
     delete_node : Optional[Callable]
     get_run : Optional[Callable]
+    pull_object : Callable[[str], bytes]
+    push_object : Callable[[str, bytes], None]
+    confirm_message_received : Callable[[str], None]
     """
     if isinstance(root_certificates, str):
         root_certificates = Path(root_certificates).read_bytes()
@@ -149,10 +155,7 @@ def grpc_request_response(  # pylint: disable=R0913,R0914,R0915,R0917
     if adapter_cls is None:
         adapter_cls = FleetStub
     stub = adapter_cls(channel)
-    metadata: Optional[Metadata] = None
     node: Optional[Node] = None
-    ping_thread: Optional[threading.Thread] = None
-    ping_stop_event = threading.Event()
 
     def _should_giveup_fn(e: Exception) -> bool:
         if e.code() == grpc.StatusCode.PERMISSION_DENIED:  # type: ignore
@@ -165,46 +168,58 @@ def grpc_request_response(  # pylint: disable=R0913,R0914,R0915,R0917
     # If the status code is PERMISSION_DENIED, additionally raise RunNotRunningException
     retry_invoker.should_giveup = _should_giveup_fn
 
+    # Wrap stub
+    _wrap_stub(stub, retry_invoker)
     ###########################################################################
-    # ping/create_node/delete_node/receive/send/get_run functions
+    # send_node_heartbeat/create_node/delete_node/receive/send/get_run functions
     ###########################################################################
 
-    def ping() -> None:
+    def send_node_heartbeat() -> bool:
         # Get Node
         if node is None:
             log(ERROR, "Node instance missing")
-            return
+            return False
 
-        # Construct the ping request
-        req = PingRequest(node=node, ping_interval=PING_DEFAULT_INTERVAL)
+        # Construct the heartbeat request
+        req = SendNodeHeartbeatRequest(
+            node=node, heartbeat_interval=HEARTBEAT_DEFAULT_INTERVAL
+        )
 
         # Call FleetAPI
-        res: PingResponse = stub.Ping(req, timeout=PING_CALL_TIMEOUT)
+        try:
+            res: SendNodeHeartbeatResponse = stub.SendNodeHeartbeat(
+                req, timeout=HEARTBEAT_CALL_TIMEOUT
+            )
+        except grpc.RpcError as e:
+            status_code = e.code()
+            if status_code == grpc.StatusCode.UNAVAILABLE:
+                return False
+            if status_code == grpc.StatusCode.DEADLINE_EXCEEDED:
+                return False
+            raise
 
         # Check if success
         if not res.success:
-            raise RuntimeError("Ping failed unexpectedly.")
+            raise RuntimeError(
+                "Heartbeat failed unexpectedly. The SuperLink does not "
+                "recognize this SuperNode."
+            )
+        return True
 
-        # Wait
-        rd = random.uniform(*PING_RANDOM_RANGE)
-        next_interval: float = PING_DEFAULT_INTERVAL - PING_CALL_TIMEOUT
-        next_interval *= PING_BASE_MULTIPLIER + rd
-        if not ping_stop_event.is_set():
-            ping_stop_event.wait(next_interval)
+    heartbeat_sender = HeartbeatSender(send_node_heartbeat)
 
     def create_node() -> Optional[int]:
         """Set create_node."""
         # Call FleetAPI
-        create_node_request = CreateNodeRequest(ping_interval=PING_DEFAULT_INTERVAL)
-        create_node_response = retry_invoker.invoke(
-            stub.CreateNode,
-            request=create_node_request,
+        create_node_request = CreateNodeRequest(
+            heartbeat_interval=HEARTBEAT_DEFAULT_INTERVAL
         )
+        create_node_response = stub.CreateNode(request=create_node_request)
 
-        # Remember the node and the ping-loop thread
-        nonlocal node, ping_thread
+        # Remember the node and start the heartbeat sender
+        nonlocal node
         node = cast(Node, create_node_response.node)
-        ping_thread = start_ping_loop(ping, ping_stop_event)
+        heartbeat_sender.start()
         return node.node_id
 
     def delete_node() -> None:
@@ -215,83 +230,67 @@ def grpc_request_response(  # pylint: disable=R0913,R0914,R0915,R0917
             log(ERROR, "Node instance missing")
             return
 
-        # Stop the ping-loop thread
-        ping_stop_event.set()
+        # Stop the heartbeat sender
+        heartbeat_sender.stop()
 
         # Call FleetAPI
         delete_node_request = DeleteNodeRequest(node=node)
-        retry_invoker.invoke(stub.DeleteNode, request=delete_node_request)
+        stub.DeleteNode(request=delete_node_request)
 
         # Cleanup
         node = None
 
-    def receive() -> Optional[Message]:
-        """Receive next message from server."""
+    def receive() -> Optional[tuple[Message, ObjectTree]]:
+        """Pull a message with its ObjectTree from SuperLink."""
         # Get Node
         if node is None:
             log(ERROR, "Node instance missing")
             return None
 
-        # Request instructions (message) from server
+        # Try to pull a message with its object tree from SuperLink
         request = PullMessagesRequest(node=node)
-        response: PullMessagesResponse = retry_invoker.invoke(
-            stub.PullMessages, request=request
-        )
+        response: PullMessagesResponse = stub.PullMessages(request=request)
 
-        # Get the current Messages
-        message_proto = (
-            None if len(response.messages_list) == 0 else response.messages_list[0]
-        )
+        # If no messages are available, return None
+        if len(response.messages_list) == 0:
+            return None
 
-        # Discard the current message if not valid
-        if message_proto is not None and not (
-            message_proto.metadata.dst_node_id == node.node_id
-        ):
-            message_proto = None
+        # Get the current Message and its object tree
+        message_proto = response.messages_list[0]
+        object_tree = response.message_object_trees[0]
 
         # Construct the Message
-        in_message = message_from_proto(message_proto) if message_proto else None
-
-        # Remember `metadata` of the in message
-        nonlocal metadata
-        metadata = copy(in_message.metadata) if in_message else None
+        in_message = message_from_proto(message_proto)
 
-        # Return the message if available
-        return in_message
+        # Return the Message and its object tree
+        return in_message, object_tree
 
-    def send(message: Message) -> None:
-        """Send message reply to server."""
+    def send(message: Message, object_tree: ObjectTree) -> set[str]:
+        """Send the message with its ObjectTree to SuperLink."""
         # Get Node
         if node is None:
             log(ERROR, "Node instance missing")
-            return
-
-        # Get the metadata of the incoming message
-        nonlocal metadata
-        if metadata is None:
-            log(ERROR, "No current message")
-            return
+            return set()
 
-        # Validate out message
-        if not validate_out_message(message, metadata):
-            log(ERROR, "Invalid out message")
-            return
+        # Remove the content from the message if it has
+        if message.has_content():
+            message = remove_content_from_message(message)
 
-        # Serialize Message
-        message_proto = message_to_proto(message=message)
-        request = PushMessagesRequest(node=node, messages_list=[message_proto])
-        _ = retry_invoker.invoke(stub.PushMessages, request)
+        # Send the message with its ObjectTree to SuperLink
+        request = PushMessagesRequest(
+            node=node,
+            messages_list=[message_to_proto(message)],
+            message_object_trees=[object_tree],
+        )
+        response: PushMessagesResponse = stub.PushMessages(request=request)
 
-        # Cleanup
-        metadata = None
+        # Get and return the object IDs to push
+        return set(response.objects_to_push)
 
     def get_run(run_id: int) -> Run:
         # Call FleetAPI
         get_run_request = GetRunRequest(node=node, run_id=run_id)
-        get_run_response: GetRunResponse = retry_invoker.invoke(
-            stub.GetRun,
-            request=get_run_request,
-        )
+        get_run_response: GetRunResponse = stub.GetRun(request=get_run_request)
 
         # Return fab_id and fab_version
         return run_from_proto(get_run_response.run)
@@ -299,16 +298,62 @@ def grpc_request_response(  # pylint: disable=R0913,R0914,R0915,R0917
     def get_fab(fab_hash: str, run_id: int) -> Fab:
         # Call FleetAPI
         get_fab_request = GetFabRequest(node=node, hash_str=fab_hash, run_id=run_id)
-        get_fab_response: GetFabResponse = retry_invoker.invoke(
-            stub.GetFab,
-            request=get_fab_request,
-        )
+        get_fab_response: GetFabResponse = stub.GetFab(request=get_fab_request)
 
         return Fab(get_fab_response.fab.hash_str, get_fab_response.fab.content)
 
+    def pull_object(run_id: int, object_id: str) -> bytes:
+        """Pull the object from the SuperLink."""
+        # Check Node
+        if node is None:
+            raise RuntimeError("Node instance missing")
+
+        fn = make_pull_object_fn_protobuf(
+            pull_object_protobuf=stub.PullObject,
+            node=node,
+            run_id=run_id,
+        )
+        return fn(object_id)
+
+    def push_object(run_id: int, object_id: str, contents: bytes) -> None:
+        """Push the object to the SuperLink."""
+        # Check Node
+        if node is None:
+            raise RuntimeError("Node instance missing")
+
+        fn = make_push_object_fn_protobuf(
+            push_object_protobuf=stub.PushObject,
+            node=node,
+            run_id=run_id,
+        )
+        fn(object_id, contents)
+
+    def confirm_message_received(run_id: int, object_id: str) -> None:
+        """Confirm that the message has been received."""
+        # Check Node
+        if node is None:
+            raise RuntimeError("Node instance missing")
+
+        fn = make_confirm_message_received_fn_protobuf(
+            confirm_message_received_protobuf=stub.ConfirmMessageReceived,
+            node=node,
+            run_id=run_id,
+        )
+        fn(object_id)
+
     try:
         # Yield methods
-        yield (receive, send, create_node, delete_node, get_run, get_fab)
+        yield (
+            receive,
+            send,
+            create_node,
+            delete_node,
+            get_run,
+            get_fab,
+            pull_object,
+            push_object,
+            confirm_message_received,
+        )
     except Exception as exc:  # pylint: disable=broad-except
         log(ERROR, exc)
     # Cleanup