feat(framework) Add PushLogs RPC to Driver service (#4390)

Co-authored-by: Chong Shen Ng <[email protected]>

src/proto/flwr/proto/driver.proto

@@ -17,6 +17,7 @@ syntax = "proto3";
 
 package flwr.proto;
 
+import "flwr/proto/log.proto";
 import "flwr/proto/node.proto";
 import "flwr/proto/message.proto";
 import "flwr/proto/task.proto";
@@ -53,6 +54,9 @@ service Driver {
   // Update the status of a given run
   rpc UpdateRunStatus(UpdateRunStatusRequest)
       returns (UpdateRunStatusResponse) {}
+
+  // Push ServerApp logs
+  rpc PushLogs(PushLogsRequest) returns (PushLogsResponse) {}
 }
 
 // GetNodes messages

src/proto/flwr/proto/log.proto

@@ -0,0 +1,27 @@
+// Copyright 2024 Flower Labs GmbH. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// ==============================================================================
+
+syntax = "proto3";
+
+package flwr.proto;
+
+import "flwr/proto/node.proto";
+
+message PushLogsRequest {
+  Node node = 1;
+  uint64 run_id = 2;
+  repeated string logs = 3;
+}
+message PushLogsResponse {}

src/py/flwr/proto/driver_pb2_grpc.py

@@ -4,6 +4,7 @@
 
 from flwr.proto import driver_pb2 as flwr_dot_proto_dot_driver__pb2
 from flwr.proto import fab_pb2 as flwr_dot_proto_dot_fab__pb2
+from flwr.proto import log_pb2 as flwr_dot_proto_dot_log__pb2
 from flwr.proto import run_pb2 as flwr_dot_proto_dot_run__pb2
 
 
@@ -61,6 +62,11 @@ def __init__(self, channel):
                 request_serializer=flwr_dot_proto_dot_run__pb2.UpdateRunStatusRequest.SerializeToString,
                 response_deserializer=flwr_dot_proto_dot_run__pb2.UpdateRunStatusResponse.FromString,
                 )
+        self.PushLogs = channel.unary_unary(
+                '/flwr.proto.Driver/PushLogs',
+                request_serializer=flwr_dot_proto_dot_log__pb2.PushLogsRequest.SerializeToString,
+                response_deserializer=flwr_dot_proto_dot_log__pb2.PushLogsResponse.FromString,
+                )
 
 
 class DriverServicer(object):
@@ -129,6 +135,13 @@ def UpdateRunStatus(self, request, context):
         context.set_details('Method not implemented!')
         raise NotImplementedError('Method not implemented!')
 
+    def PushLogs(self, request, context):
+        """Push ServerApp logs
+        """
+        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+        context.set_details('Method not implemented!')
+        raise NotImplementedError('Method not implemented!')
+
 
 def add_DriverServicer_to_server(servicer, server):
     rpc_method_handlers = {
@@ -177,6 +190,11 @@ def add_DriverServicer_to_server(servicer, server):
                     request_deserializer=flwr_dot_proto_dot_run__pb2.UpdateRunStatusRequest.FromString,
                     response_serializer=flwr_dot_proto_dot_run__pb2.UpdateRunStatusResponse.SerializeToString,
             ),
+            'PushLogs': grpc.unary_unary_rpc_method_handler(
+                    servicer.PushLogs,
+                    request_deserializer=flwr_dot_proto_dot_log__pb2.PushLogsRequest.FromString,
+                    response_serializer=flwr_dot_proto_dot_log__pb2.PushLogsResponse.SerializeToString,
+            ),
     }
     generic_handler = grpc.method_handlers_generic_handler(
             'flwr.proto.Driver', rpc_method_handlers)
@@ -339,3 +357,20 @@ def UpdateRunStatus(request,
             flwr_dot_proto_dot_run__pb2.UpdateRunStatusResponse.FromString,
             options, channel_credentials,
             insecure, call_credentials, compression, wait_for_ready, timeout, metadata)
+
+    @staticmethod
+    def PushLogs(request,
+            target,
+            options=(),
+            channel_credentials=None,
+            call_credentials=None,
+            insecure=False,
+            compression=None,
+            wait_for_ready=None,
+            timeout=None,
+            metadata=None):
+        return grpc.experimental.unary_unary(request, target, '/flwr.proto.Driver/PushLogs',
+            flwr_dot_proto_dot_log__pb2.PushLogsRequest.SerializeToString,
+            flwr_dot_proto_dot_log__pb2.PushLogsResponse.FromString,
+            options, channel_credentials,
+            insecure, call_credentials, compression, wait_for_ready, timeout, metadata)

src/py/flwr/proto/driver_pb2_grpc.pyi

@@ -5,6 +5,7 @@ isort:skip_file
 import abc
 import flwr.proto.driver_pb2
 import flwr.proto.fab_pb2
+import flwr.proto.log_pb2
 import flwr.proto.run_pb2
 import grpc
 
@@ -55,6 +56,11 @@ class DriverStub:
         flwr.proto.run_pb2.UpdateRunStatusResponse]
     """Update the status of a given run"""
 
+    PushLogs: grpc.UnaryUnaryMultiCallable[
+        flwr.proto.log_pb2.PushLogsRequest,
+        flwr.proto.log_pb2.PushLogsResponse]
+    """Push ServerApp logs"""
+
 
 class DriverServicer(metaclass=abc.ABCMeta):
     @abc.abstractmethod
@@ -129,5 +135,13 @@ class DriverServicer(metaclass=abc.ABCMeta):
         """Update the status of a given run"""
         pass
 
+    @abc.abstractmethod
+    def PushLogs(self,
+        request: flwr.proto.log_pb2.PushLogsRequest,
+        context: grpc.ServicerContext,
+    ) -> flwr.proto.log_pb2.PushLogsResponse:
+        """Push ServerApp logs"""
+        pass
+
 
 def add_DriverServicer_to_server(servicer: DriverServicer, server: grpc.Server) -> None: ...

src/py/flwr/proto/log_pb2.pyi

@@ -0,0 +1,39 @@
+"""
+@generated by mypy-protobuf.  Do not edit manually!
+isort:skip_file
+"""
+import builtins
+import flwr.proto.node_pb2
+import google.protobuf.descriptor
+import google.protobuf.internal.containers
+import google.protobuf.message
+import typing
+import typing_extensions
+
+DESCRIPTOR: google.protobuf.descriptor.FileDescriptor
+
+class PushLogsRequest(google.protobuf.message.Message):
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+    NODE_FIELD_NUMBER: builtins.int
+    RUN_ID_FIELD_NUMBER: builtins.int
+    LOGS_FIELD_NUMBER: builtins.int
+    @property
+    def node(self) -> flwr.proto.node_pb2.Node: ...
+    run_id: builtins.int
+    @property
+    def logs(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[typing.Text]: ...
+    def __init__(self,
+        *,
+        node: typing.Optional[flwr.proto.node_pb2.Node] = ...,
+        run_id: builtins.int = ...,
+        logs: typing.Optional[typing.Iterable[typing.Text]] = ...,
+        ) -> None: ...
+    def HasField(self, field_name: typing_extensions.Literal["node",b"node"]) -> builtins.bool: ...
+    def ClearField(self, field_name: typing_extensions.Literal["logs",b"logs","node",b"node","run_id",b"run_id"]) -> None: ...
+global___PushLogsRequest = PushLogsRequest
+
+class PushLogsResponse(google.protobuf.message.Message):
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+    def __init__(self,
+        ) -> None: ...
+global___PushLogsResponse = PushLogsResponse

src/py/flwr/proto/log_pb2_grpc.py

@@ -0,0 +1,4 @@
+# Generated by the gRPC Python protocol compiler plugin. DO NOT EDIT!
+"""Client and server classes corresponding to protobuf-defined services."""
+import grpc
+

src/py/flwr/proto/log_pb2_grpc.pyi

@@ -0,0 +1,4 @@
+"""
+@generated by mypy-protobuf.  Do not edit manually!
+isort:skip_file
+"""

src/py/flwr/server/superlink/driver/driver_servicer.py

@@ -50,6 +50,10 @@
     PushTaskInsResponse,
 )
 from flwr.proto.fab_pb2 import GetFabRequest, GetFabResponse  # pylint: disable=E0611
+from flwr.proto.log_pb2 import (  # pylint: disable=E0611
+    PushLogsRequest,
+    PushLogsResponse,
+)
 from flwr.proto.node_pb2 import Node  # pylint: disable=E0611
 from flwr.proto.run_pb2 import (  # pylint: disable=E0611
     CreateRunRequest,
@@ -269,6 +273,13 @@ def UpdateRunStatus(
         )
         return UpdateRunStatusResponse()
 
+    def PushLogs(
+        self, request: PushLogsRequest, context: grpc.ServicerContext
+    ) -> PushLogsResponse:
+        """Push logs."""
+        log(DEBUG, "DriverServicer.PushLogs")
+        raise NotImplementedError()
+
 
 def _raise_if(validation_error: bool, detail: str) -> None:
     if validation_error:

src/py/flwr_tool/protoc_test.py

@@ -28,4 +28,4 @@ def test_directories() -> None:
 
 def test_proto_file_count() -> None:
     """Test if the correct number of proto files were captured by the glob."""
-    assert len(PROTO_FILES) == 14
+    assert len(PROTO_FILES) == 15