flwr-nightly 1.13.0.dev20241023__py3-none-any.whl → 1.13.0.dev20241025__py3-none-any.whl

flwr/cli/new/templates/app/pyproject.tensorflow.toml.tpl

@@ -10,7 +10,7 @@ license = "Apache-2.0"
 dependencies = [
     "flwr[simulation]>=1.12.0",
     "flwr-datasets[vision]>=0.3.0",
-    "tensorflow>=2.11.1",
+    "tensorflow>=2.11.1,<2.18.0",
 ]
 
 [tool.hatch.build.targets.wheel]

flwr/proto/driver_pb2.py

@@ -13,30 +13,39 @@ _sym_db = _symbol_database.Default()
 
 
 from flwr.proto import node_pb2 as flwr_dot_proto_dot_node__pb2
+from flwr.proto import message_pb2 as flwr_dot_proto_dot_message__pb2
 from flwr.proto import task_pb2 as flwr_dot_proto_dot_task__pb2
 from flwr.proto import run_pb2 as flwr_dot_proto_dot_run__pb2
 from flwr.proto import fab_pb2 as flwr_dot_proto_dot_fab__pb2
 
 
-DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x17\x66lwr/proto/driver.proto\x12\nflwr.proto\x1a\x15\x66lwr/proto/node.proto\x1a\x15\x66lwr/proto/task.proto\x1a\x14\x66lwr/proto/run.proto\x1a\x14\x66lwr/proto/fab.proto\"!\n\x0fGetNodesRequest\x12\x0e\n\x06run_id\x18\x01 \x01(\x04\"3\n\x10GetNodesResponse\x12\x1f\n\x05nodes\x18\x01 \x03(\x0b\x32\x10.flwr.proto.Node\"@\n\x12PushTaskInsRequest\x12*\n\rtask_ins_list\x18\x01 \x03(\x0b\x32\x13.flwr.proto.TaskIns\"\'\n\x13PushTaskInsResponse\x12\x10\n\x08task_ids\x18\x02 \x03(\t\"F\n\x12PullTaskResRequest\x12\x1e\n\x04node\x18\x01 \x01(\x0b\x32\x10.flwr.proto.Node\x12\x10\n\x08task_ids\x18\x02 \x03(\t\"A\n\x13PullTaskResResponse\x12*\n\rtask_res_list\x18\x01 \x03(\x0b\x32\x13.flwr.proto.TaskRes2\xc7\x03\n\x06\x44river\x12J\n\tCreateRun\x12\x1c.flwr.proto.CreateRunRequest\x1a\x1d.flwr.proto.CreateRunResponse\"\x00\x12G\n\x08GetNodes\x12\x1b.flwr.proto.GetNodesRequest\x1a\x1c.flwr.proto.GetNodesResponse\"\x00\x12P\n\x0bPushTaskIns\x12\x1e.flwr.proto.PushTaskInsRequest\x1a\x1f.flwr.proto.PushTaskInsResponse\"\x00\x12P\n\x0bPullTaskRes\x12\x1e.flwr.proto.PullTaskResRequest\x1a\x1f.flwr.proto.PullTaskResResponse\"\x00\x12\x41\n\x06GetRun\x12\x19.flwr.proto.GetRunRequest\x1a\x1a.flwr.proto.GetRunResponse\"\x00\x12\x41\n\x06GetFab\x12\x19.flwr.proto.GetFabRequest\x1a\x1a.flwr.proto.GetFabResponse\"\x00\x62\x06proto3')
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x17\x66lwr/proto/driver.proto\x12\nflwr.proto\x1a\x15\x66lwr/proto/node.proto\x1a\x18\x66lwr/proto/message.proto\x1a\x15\x66lwr/proto/task.proto\x1a\x14\x66lwr/proto/run.proto\x1a\x14\x66lwr/proto/fab.proto\"!\n\x0fGetNodesRequest\x12\x0e\n\x06run_id\x18\x01 \x01(\x04\"3\n\x10GetNodesResponse\x12\x1f\n\x05nodes\x18\x01 \x03(\x0b\x32\x10.flwr.proto.Node\"@\n\x12PushTaskInsRequest\x12*\n\rtask_ins_list\x18\x01 \x03(\x0b\x32\x13.flwr.proto.TaskIns\"\'\n\x13PushTaskInsResponse\x12\x10\n\x08task_ids\x18\x02 \x03(\t\"F\n\x12PullTaskResRequest\x12\x1e\n\x04node\x18\x01 \x01(\x0b\x32\x10.flwr.proto.Node\x12\x10\n\x08task_ids\x18\x02 \x03(\t\"A\n\x13PullTaskResResponse\x12*\n\rtask_res_list\x18\x01 \x03(\x0b\x32\x13.flwr.proto.TaskRes\",\n\x1aPullServerAppInputsRequest\x12\x0e\n\x06run_id\x18\x01 \x01(\x04\"\x7f\n\x1bPullServerAppInputsResponse\x12$\n\x07\x63ontext\x18\x01 \x01(\x0b\x32\x13.flwr.proto.Context\x12\x1c\n\x03run\x18\x02 \x01(\x0b\x32\x0f.flwr.proto.Run\x12\x1c\n\x03\x66\x61\x62\x18\x03 \x01(\x0b\x32\x0f.flwr.proto.Fab\"S\n\x1bPushServerAppOutputsRequest\x12\x0e\n\x06run_id\x18\x01 \x01(\x04\x12$\n\x07\x63ontext\x18\x02 \x01(\x0b\x32\x13.flwr.proto.Context\"\x1e\n\x1cPushServerAppOutputsResponse2\x9e\x05\n\x06\x44river\x12J\n\tCreateRun\x12\x1c.flwr.proto.CreateRunRequest\x1a\x1d.flwr.proto.CreateRunResponse\"\x00\x12G\n\x08GetNodes\x12\x1b.flwr.proto.GetNodesRequest\x1a\x1c.flwr.proto.GetNodesResponse\"\x00\x12P\n\x0bPushTaskIns\x12\x1e.flwr.proto.PushTaskInsRequest\x1a\x1f.flwr.proto.PushTaskInsResponse\"\x00\x12P\n\x0bPullTaskRes\x12\x1e.flwr.proto.PullTaskResRequest\x1a\x1f.flwr.proto.PullTaskResResponse\"\x00\x12\x41\n\x06GetRun\x12\x19.flwr.proto.GetRunRequest\x1a\x1a.flwr.proto.GetRunResponse\"\x00\x12\x41\n\x06GetFab\x12\x19.flwr.proto.GetFabRequest\x1a\x1a.flwr.proto.GetFabResponse\"\x00\x12h\n\x13PullServerAppInputs\x12&.flwr.proto.PullServerAppInputsRequest\x1a\'.flwr.proto.PullServerAppInputsResponse\"\x00\x12k\n\x14PushServerAppOutputs\x12\'.flwr.proto.PushServerAppOutputsRequest\x1a(.flwr.proto.PushServerAppOutputsResponse\"\x00\x62\x06proto3')
 
 _globals = globals()
 _builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
 _builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'flwr.proto.driver_pb2', _globals)
 if _descriptor._USE_C_DESCRIPTORS == False:
   DESCRIPTOR._options = None
-  _globals['_GETNODESREQUEST']._serialized_start=129
-  _globals['_GETNODESREQUEST']._serialized_end=162
-  _globals['_GETNODESRESPONSE']._serialized_start=164
-  _globals['_GETNODESRESPONSE']._serialized_end=215
-  _globals['_PUSHTASKINSREQUEST']._serialized_start=217
-  _globals['_PUSHTASKINSREQUEST']._serialized_end=281
-  _globals['_PUSHTASKINSRESPONSE']._serialized_start=283
-  _globals['_PUSHTASKINSRESPONSE']._serialized_end=322
-  _globals['_PULLTASKRESREQUEST']._serialized_start=324
-  _globals['_PULLTASKRESREQUEST']._serialized_end=394
-  _globals['_PULLTASKRESRESPONSE']._serialized_start=396
-  _globals['_PULLTASKRESRESPONSE']._serialized_end=461
-  _globals['_DRIVER']._serialized_start=464
-  _globals['_DRIVER']._serialized_end=919
+  _globals['_GETNODESREQUEST']._serialized_start=155
+  _globals['_GETNODESREQUEST']._serialized_end=188
+  _globals['_GETNODESRESPONSE']._serialized_start=190
+  _globals['_GETNODESRESPONSE']._serialized_end=241
+  _globals['_PUSHTASKINSREQUEST']._serialized_start=243
+  _globals['_PUSHTASKINSREQUEST']._serialized_end=307
+  _globals['_PUSHTASKINSRESPONSE']._serialized_start=309
+  _globals['_PUSHTASKINSRESPONSE']._serialized_end=348
+  _globals['_PULLTASKRESREQUEST']._serialized_start=350
+  _globals['_PULLTASKRESREQUEST']._serialized_end=420
+  _globals['_PULLTASKRESRESPONSE']._serialized_start=422
+  _globals['_PULLTASKRESRESPONSE']._serialized_end=487
+  _globals['_PULLSERVERAPPINPUTSREQUEST']._serialized_start=489
+  _globals['_PULLSERVERAPPINPUTSREQUEST']._serialized_end=533
+  _globals['_PULLSERVERAPPINPUTSRESPONSE']._serialized_start=535
+  _globals['_PULLSERVERAPPINPUTSRESPONSE']._serialized_end=662
+  _globals['_PUSHSERVERAPPOUTPUTSREQUEST']._serialized_start=664
+  _globals['_PUSHSERVERAPPOUTPUTSREQUEST']._serialized_end=747
+  _globals['_PUSHSERVERAPPOUTPUTSRESPONSE']._serialized_start=749
+  _globals['_PUSHSERVERAPPOUTPUTSRESPONSE']._serialized_end=779
+  _globals['_DRIVER']._serialized_start=782
+  _globals['_DRIVER']._serialized_end=1452
 # @@protoc_insertion_point(module_scope)

flwr/proto/driver_pb2.pyi

@@ -3,7 +3,10 @@
 isort:skip_file
 """
 import builtins
+import flwr.proto.fab_pb2
+import flwr.proto.message_pb2
 import flwr.proto.node_pb2
+import flwr.proto.run_pb2
 import flwr.proto.task_pb2
 import google.protobuf.descriptor
 import google.protobuf.internal.containers
@@ -91,3 +94,59 @@ class PullTaskResResponse(google.protobuf.message.Message):
         ) -> None: ...
     def ClearField(self, field_name: typing_extensions.Literal["task_res_list",b"task_res_list"]) -> None: ...
 global___PullTaskResResponse = PullTaskResResponse
+
+class PullServerAppInputsRequest(google.protobuf.message.Message):
+    """PullServerAppInputs messages"""
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+    RUN_ID_FIELD_NUMBER: builtins.int
+    run_id: builtins.int
+    def __init__(self,
+        *,
+        run_id: builtins.int = ...,
+        ) -> None: ...
+    def ClearField(self, field_name: typing_extensions.Literal["run_id",b"run_id"]) -> None: ...
+global___PullServerAppInputsRequest = PullServerAppInputsRequest
+
+class PullServerAppInputsResponse(google.protobuf.message.Message):
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+    CONTEXT_FIELD_NUMBER: builtins.int
+    RUN_FIELD_NUMBER: builtins.int
+    FAB_FIELD_NUMBER: builtins.int
+    @property
+    def context(self) -> flwr.proto.message_pb2.Context: ...
+    @property
+    def run(self) -> flwr.proto.run_pb2.Run: ...
+    @property
+    def fab(self) -> flwr.proto.fab_pb2.Fab: ...
+    def __init__(self,
+        *,
+        context: typing.Optional[flwr.proto.message_pb2.Context] = ...,
+        run: typing.Optional[flwr.proto.run_pb2.Run] = ...,
+        fab: typing.Optional[flwr.proto.fab_pb2.Fab] = ...,
+        ) -> None: ...
+    def HasField(self, field_name: typing_extensions.Literal["context",b"context","fab",b"fab","run",b"run"]) -> builtins.bool: ...
+    def ClearField(self, field_name: typing_extensions.Literal["context",b"context","fab",b"fab","run",b"run"]) -> None: ...
+global___PullServerAppInputsResponse = PullServerAppInputsResponse
+
+class PushServerAppOutputsRequest(google.protobuf.message.Message):
+    """PushServerAppOutputs messages"""
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+    RUN_ID_FIELD_NUMBER: builtins.int
+    CONTEXT_FIELD_NUMBER: builtins.int
+    run_id: builtins.int
+    @property
+    def context(self) -> flwr.proto.message_pb2.Context: ...
+    def __init__(self,
+        *,
+        run_id: builtins.int = ...,
+        context: typing.Optional[flwr.proto.message_pb2.Context] = ...,
+        ) -> None: ...
+    def HasField(self, field_name: typing_extensions.Literal["context",b"context"]) -> builtins.bool: ...
+    def ClearField(self, field_name: typing_extensions.Literal["context",b"context","run_id",b"run_id"]) -> None: ...
+global___PushServerAppOutputsRequest = PushServerAppOutputsRequest
+
+class PushServerAppOutputsResponse(google.protobuf.message.Message):
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+    def __init__(self,
+        ) -> None: ...
+global___PushServerAppOutputsResponse = PushServerAppOutputsResponse

flwr/proto/driver_pb2_grpc.py

@@ -46,6 +46,16 @@ class DriverStub(object):
                 request_serializer=flwr_dot_proto_dot_fab__pb2.GetFabRequest.SerializeToString,
                 response_deserializer=flwr_dot_proto_dot_fab__pb2.GetFabResponse.FromString,
                 )
+        self.PullServerAppInputs = channel.unary_unary(
+                '/flwr.proto.Driver/PullServerAppInputs',
+                request_serializer=flwr_dot_proto_dot_driver__pb2.PullServerAppInputsRequest.SerializeToString,
+                response_deserializer=flwr_dot_proto_dot_driver__pb2.PullServerAppInputsResponse.FromString,
+                )
+        self.PushServerAppOutputs = channel.unary_unary(
+                '/flwr.proto.Driver/PushServerAppOutputs',
+                request_serializer=flwr_dot_proto_dot_driver__pb2.PushServerAppOutputsRequest.SerializeToString,
+                response_deserializer=flwr_dot_proto_dot_driver__pb2.PushServerAppOutputsResponse.FromString,
+                )
 
 
 class DriverServicer(object):
@@ -93,6 +103,20 @@ class DriverServicer(object):
         context.set_details('Method not implemented!')
         raise NotImplementedError('Method not implemented!')
 
+    def PullServerAppInputs(self, request, context):
+        """Pull ServerApp inputs
+        """
+        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+        context.set_details('Method not implemented!')
+        raise NotImplementedError('Method not implemented!')
+
+    def PushServerAppOutputs(self, request, context):
+        """Push ServerApp outputs
+        """
+        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 = {
@@ -126,6 +150,16 @@ def add_DriverServicer_to_server(servicer, server):
                     request_deserializer=flwr_dot_proto_dot_fab__pb2.GetFabRequest.FromString,
                     response_serializer=flwr_dot_proto_dot_fab__pb2.GetFabResponse.SerializeToString,
             ),
+            'PullServerAppInputs': grpc.unary_unary_rpc_method_handler(
+                    servicer.PullServerAppInputs,
+                    request_deserializer=flwr_dot_proto_dot_driver__pb2.PullServerAppInputsRequest.FromString,
+                    response_serializer=flwr_dot_proto_dot_driver__pb2.PullServerAppInputsResponse.SerializeToString,
+            ),
+            'PushServerAppOutputs': grpc.unary_unary_rpc_method_handler(
+                    servicer.PushServerAppOutputs,
+                    request_deserializer=flwr_dot_proto_dot_driver__pb2.PushServerAppOutputsRequest.FromString,
+                    response_serializer=flwr_dot_proto_dot_driver__pb2.PushServerAppOutputsResponse.SerializeToString,
+            ),
     }
     generic_handler = grpc.method_handlers_generic_handler(
             'flwr.proto.Driver', rpc_method_handlers)
@@ -237,3 +271,37 @@ class Driver(object):
             flwr_dot_proto_dot_fab__pb2.GetFabResponse.FromString,
             options, channel_credentials,
             insecure, call_credentials, compression, wait_for_ready, timeout, metadata)
+
+    @staticmethod
+    def PullServerAppInputs(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/PullServerAppInputs',
+            flwr_dot_proto_dot_driver__pb2.PullServerAppInputsRequest.SerializeToString,
+            flwr_dot_proto_dot_driver__pb2.PullServerAppInputsResponse.FromString,
+            options, channel_credentials,
+            insecure, call_credentials, compression, wait_for_ready, timeout, metadata)
+
+    @staticmethod
+    def PushServerAppOutputs(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/PushServerAppOutputs',
+            flwr_dot_proto_dot_driver__pb2.PushServerAppOutputsRequest.SerializeToString,
+            flwr_dot_proto_dot_driver__pb2.PushServerAppOutputsResponse.FromString,
+            options, channel_credentials,
+            insecure, call_credentials, compression, wait_for_ready, timeout, metadata)

flwr/proto/driver_pb2_grpc.pyi

@@ -40,6 +40,16 @@ class DriverStub:
         flwr.proto.fab_pb2.GetFabResponse]
     """Get FAB"""
 
+    PullServerAppInputs: grpc.UnaryUnaryMultiCallable[
+        flwr.proto.driver_pb2.PullServerAppInputsRequest,
+        flwr.proto.driver_pb2.PullServerAppInputsResponse]
+    """Pull ServerApp inputs"""
+
+    PushServerAppOutputs: grpc.UnaryUnaryMultiCallable[
+        flwr.proto.driver_pb2.PushServerAppOutputsRequest,
+        flwr.proto.driver_pb2.PushServerAppOutputsResponse]
+    """Push ServerApp outputs"""
+
 
 class DriverServicer(metaclass=abc.ABCMeta):
     @abc.abstractmethod
@@ -90,5 +100,21 @@ class DriverServicer(metaclass=abc.ABCMeta):
         """Get FAB"""
         pass
 
+    @abc.abstractmethod
+    def PullServerAppInputs(self,
+        request: flwr.proto.driver_pb2.PullServerAppInputsRequest,
+        context: grpc.ServicerContext,
+    ) -> flwr.proto.driver_pb2.PullServerAppInputsResponse:
+        """Pull ServerApp inputs"""
+        pass
+
+    @abc.abstractmethod
+    def PushServerAppOutputs(self,
+        request: flwr.proto.driver_pb2.PushServerAppOutputsRequest,
+        context: grpc.ServicerContext,
+    ) -> flwr.proto.driver_pb2.PushServerAppOutputsResponse:
+        """Push ServerApp outputs"""
+        pass
+
 
 def add_DriverServicer_to_server(servicer: DriverServicer, server: grpc.Server) -> None: ...

flwr/server/app.py

@@ -64,7 +64,7 @@ from flwr.proto.fleet_pb2_grpc import (  # pylint: disable=E0611
 )
 from flwr.proto.grpcadapter_pb2_grpc import add_GrpcAdapterServicer_to_server
 from flwr.superexec.app import load_executor
-from flwr.superexec.exec_grpc import run_superexec_api_grpc
+from flwr.superexec.exec_grpc import run_exec_api_grpc
 
 from .client_manager import ClientManager
 from .history import History
@@ -329,8 +329,10 @@ def run_superlink() -> None:
         raise ValueError(f"Unknown fleet_api_type: {args.fleet_api_type}")
 
     # Start Exec API
-    exec_server: grpc.Server = run_superexec_api_grpc(
+    exec_server: grpc.Server = run_exec_api_grpc(
         address=exec_address,
+        state_factory=state_factory,
+        ffs_factory=ffs_factory,
         executor=load_executor(args),
         certificates=certificates,
         config=parse_config_args(

flwr/server/run_serverapp.py

@@ -34,7 +34,6 @@ from flwr.common.config import (
 from flwr.common.constant import DRIVER_API_DEFAULT_ADDRESS
 from flwr.common.logger import log, update_console_handler, warn_deprecated_feature
 from flwr.common.object_ref import load_app
-from flwr.common.typing import UserConfig
 from flwr.proto.fab_pb2 import GetFabRequest, GetFabResponse  # pylint: disable=E0611
 from flwr.proto.run_pb2 import (  # pylint: disable=E0611
     CreateRunRequest,
@@ -46,13 +45,14 @@ from .driver.grpc_driver import GrpcDriver
 from .server_app import LoadServerAppError, ServerApp
 
 
+# pylint: disable-next=too-many-arguments,too-many-positional-arguments
 def run(
     driver: Driver,
+    context: Context,
     server_app_dir: str,
-    server_app_run_config: UserConfig,
     server_app_attr: Optional[str] = None,
     loaded_server_app: Optional[ServerApp] = None,
-) -> None:
+) -> Context:
     """Run ServerApp with a given Driver."""
     if not (server_app_attr is None) ^ (loaded_server_app is None):
         raise ValueError(
@@ -78,15 +78,11 @@ def run(
 
     server_app = _load()
 
-    # Initialize Context
-    context = Context(
-        node_id=0, node_config={}, state=RecordSet(), run_config=server_app_run_config
-    )
-
     # Call ServerApp
     server_app(driver=driver, context=context)
 
     log(DEBUG, "ServerApp finished running.")
+    return context
 
 
 # pylint: disable-next=too-many-branches,too-many-statements,too-many-locals
@@ -225,11 +221,19 @@ def run_server_app() -> None:
         root_certificates,
     )
 
+    # Initialize Context
+    context = Context(
+        node_id=0,
+        node_config={},
+        state=RecordSet(),
+        run_config=server_app_run_config,
+    )
+
     # Run the ServerApp with the Driver
     run(
         driver=driver,
+        context=context,
         server_app_dir=app_path,
-        server_app_run_config=server_app_run_config,
         server_app_attr=server_app_attr,
     )
 

flwr/server/superlink/driver/driver_servicer.py

@@ -15,27 +15,35 @@
 """Driver API servicer."""
 
 
+import threading
 import time
-from logging import DEBUG
+from logging import DEBUG, INFO
 from typing import Optional
 from uuid import UUID
 
 import grpc
 
+from flwr.common.constant import Status
 from flwr.common.logger import log
 from flwr.common.serde import (
+    context_from_proto,
+    context_to_proto,
     fab_from_proto,
     fab_to_proto,
+    run_to_proto,
     user_config_from_proto,
-    user_config_to_proto,
 )
-from flwr.common.typing import Fab
+from flwr.common.typing import Fab, RunStatus
 from flwr.proto import driver_pb2_grpc  # pylint: disable=E0611
 from flwr.proto.driver_pb2 import (  # pylint: disable=E0611
     GetNodesRequest,
     GetNodesResponse,
+    PullServerAppInputsRequest,
+    PullServerAppInputsResponse,
     PullTaskResRequest,
     PullTaskResResponse,
+    PushServerAppOutputsRequest,
+    PushServerAppOutputsResponse,
     PushTaskInsRequest,
     PushTaskInsResponse,
 )
@@ -46,7 +54,6 @@ from flwr.proto.run_pb2 import (  # pylint: disable=E0611
     CreateRunResponse,
     GetRunRequest,
     GetRunResponse,
-    Run,
 )
 from flwr.proto.task_pb2 import TaskRes  # pylint: disable=E0611
 from flwr.server.superlink.ffs.ffs import Ffs
@@ -63,6 +70,7 @@ class DriverServicer(driver_pb2_grpc.DriverServicer):
     ) -> None:
         self.state_factory = state_factory
         self.ffs_factory = ffs_factory
+        self.lock = threading.RLock()
 
     def GetNodes(
         self, request: GetNodesRequest, context: grpc.ServicerContext
@@ -177,15 +185,7 @@ class DriverServicer(driver_pb2_grpc.DriverServicer):
         if run is None:
             return GetRunResponse()
 
-        return GetRunResponse(
-            run=Run(
-                run_id=run.run_id,
-                fab_id=run.fab_id,
-                fab_version=run.fab_version,
-                override_config=user_config_to_proto(run.override_config),
-                fab_hash=run.fab_hash,
-            )
-        )
+        return GetRunResponse(run=run_to_proto(run))
 
     def GetFab(
         self, request: GetFabRequest, context: grpc.ServicerContext
@@ -200,6 +200,58 @@ class DriverServicer(driver_pb2_grpc.DriverServicer):
 
         raise ValueError(f"Found no FAB with hash: {request.hash_str}")
 
+    def PullServerAppInputs(
+        self, request: PullServerAppInputsRequest, context: grpc.ServicerContext
+    ) -> PullServerAppInputsResponse:
+        """Pull ServerApp process inputs."""
+        log(DEBUG, "DriverServicer.PullServerAppInputs")
+        # Init access to LinkState and Ffs
+        state = self.state_factory.state()
+        ffs = self.ffs_factory.ffs()
+
+        # Lock access to LinkState, preventing obtaining the same pending run_id
+        with self.lock:
+            # If run_id is provided, use it, otherwise use the pending run_id
+            if request.HasField("run_id"):
+                run_id: Optional[int] = request.run_id
+            else:
+                run_id = state.get_pending_run_id()
+            # If there's no pending run, return an empty response
+            if run_id is None:
+                return PullServerAppInputsResponse()
+
+            # Retrieve Context, Run and Fab for the run_id
+            serverapp_ctxt = state.get_serverapp_context(run_id)
+            run = state.get_run(run_id)
+            fab = None
+            if run and run.fab_hash:
+                if result := ffs.get(run.fab_hash):
+                    fab = Fab(run.fab_hash, result[0])
+            if run and fab:
+                # Update run status to STARTING
+                if state.update_run_status(run_id, RunStatus(Status.STARTING, "", "")):
+                    log(INFO, "Starting run %d", run_id)
+                    return PullServerAppInputsResponse(
+                        context=(
+                            context_to_proto(serverapp_ctxt) if serverapp_ctxt else None
+                        ),
+                        run=run_to_proto(run),
+                        fab=fab_to_proto(fab),
+                    )
+
+        # Raise an exception if the Run or Fab is not found,
+        # or if the status cannot be updated to STARTING
+        raise RuntimeError(f"Failed to start run {run_id}")
+
+    def PushServerAppOutputs(
+        self, request: PushServerAppOutputsRequest, context: grpc.ServicerContext
+    ) -> PushServerAppOutputsResponse:
+        """Push ServerApp process outputs."""
+        log(DEBUG, "DriverServicer.PushServerAppOutputs")
+        state = self.state_factory.state()
+        state.set_serverapp_context(request.run_id, context_from_proto(request.context))
+        return PushServerAppOutputsResponse()
+
 
 def _raise_if(validation_error: bool, detail: str) -> None:
     if validation_error:

flwr/server/superlink/linkstate/in_memory_linkstate.py

@@ -22,7 +22,7 @@ from logging import ERROR, WARNING
 from typing import Optional
 from uuid import UUID, uuid4
 
-from flwr.common import log, now
+from flwr.common import Context, log, now
 from flwr.common.constant import (
     MESSAGE_TTL_TOLERANCE,
     NODE_ID_NUM_BYTES,
@@ -65,6 +65,7 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
 
         # Map run_id to RunRecord
         self.run_ids: dict[int, RunRecord] = {}
+        self.contexts: dict[int, Context] = {}
         self.task_ins_store: dict[UUID, TaskIns] = {}
         self.task_res_store: dict[UUID, TaskRes] = {}
 
@@ -500,3 +501,13 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
                 self.node_ids[node_id] = (time.time() + ping_interval, ping_interval)
                 return True
         return False
+
+    def get_serverapp_context(self, run_id: int) -> Optional[Context]:
+        """Get the context for the specified `run_id`."""
+        return self.contexts.get(run_id)
+
+    def set_serverapp_context(self, run_id: int, context: Context) -> None:
+        """Set the context for the specified `run_id`."""
+        if run_id not in self.run_ids:
+            raise ValueError(f"Run {run_id} not found")
+        self.contexts[run_id] = context

flwr/server/superlink/linkstate/linkstate.py

@@ -19,6 +19,7 @@ import abc
 from typing import Optional
 from uuid import UUID
 
+from flwr.common import Context
 from flwr.common.typing import Run, RunStatus, UserConfig
 from flwr.proto.task_pb2 import TaskIns, TaskRes  # pylint: disable=E0611
 
@@ -270,3 +271,31 @@ class LinkState(abc.ABC):  # pylint: disable=R0904
         is_acknowledged : bool
             True if the ping is successfully acknowledged; otherwise, False.
         """
+
+    @abc.abstractmethod
+    def get_serverapp_context(self, run_id: int) -> Optional[Context]:
+        """Get the context for the specified `run_id`.
+
+        Parameters
+        ----------
+        run_id : int
+            The identifier of the run for which to retrieve the context.
+
+        Returns
+        -------
+        Optional[Context]
+            The context associated with the specified `run_id`, or `None` if no context
+            exists for the given `run_id`.
+        """
+
+    @abc.abstractmethod
+    def set_serverapp_context(self, run_id: int, context: Context) -> None:
+        """Set the context for the specified `run_id`.
+
+        Parameters
+        ----------
+        run_id : int
+            The identifier of the run for which to set the context.
+        context : Context
+            The context to be associated with the specified `run_id`.
+        """