flwr-nightly 1.13.0.dev20241023__tar.gz → 1.13.0.dev20241024__tar.gz

{flwr_nightly-1.13.0.dev20241023 → flwr_nightly-1.13.0.dev20241024}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: flwr-nightly
-Version: 1.13.0.dev20241023
+Version: 1.13.0.dev20241024
 Summary: Flower: A Friendly Federated Learning Framework
 Home-page: https://flower.ai
 License: Apache-2.0

{flwr_nightly-1.13.0.dev20241023 → flwr_nightly-1.13.0.dev20241024}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
 name = "flwr-nightly"
-version = "1.13.0.dev20241023"
+version = "1.13.0.dev20241024"
 description = "Flower: A Friendly Federated Learning Framework"
 license = "Apache-2.0"
 authors = ["The Flower Authors <hello@flower.ai>"]

flwr_nightly-1.13.0.dev20241024/src/py/flwr/proto/driver_pb2.py

@@ -0,0 +1,51 @@
+# -*- coding: utf-8 -*-
+# Generated by the protocol buffer compiler.  DO NOT EDIT!
+# source: flwr/proto/driver.proto
+# Protobuf Python Version: 4.25.0
+"""Generated protocol buffer code."""
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import descriptor_pool as _descriptor_pool
+from google.protobuf import symbol_database as _symbol_database
+from google.protobuf.internal import builder as _builder
+# @@protoc_insertion_point(imports)
+
+_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\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=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_nightly-1.13.0.dev20241023 → flwr_nightly-1.13.0.dev20241024}/src/py/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_nightly-1.13.0.dev20241023 → flwr_nightly-1.13.0.dev20241024}/src/py/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_nightly-1.13.0.dev20241023 → flwr_nightly-1.13.0.dev20241024}/src/py/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_nightly-1.13.0.dev20241023 → flwr_nightly-1.13.0.dev20241024}/src/py/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_nightly-1.13.0.dev20241023 → flwr_nightly-1.13.0.dev20241024}/src/py/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_nightly-1.13.0.dev20241023 → flwr_nightly-1.13.0.dev20241024}/src/py/flwr/server/superlink/driver/driver_servicer.py

@@ -34,8 +34,12 @@ 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,
 )
@@ -200,6 +204,18 @@ 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."""
+        raise NotImplementedError()
+
+    def PushServerAppOutputs(
+        self, request: PushServerAppOutputsRequest, context: grpc.ServicerContext
+    ) -> PushServerAppOutputsResponse:
+        """Push ServerApp process outputs."""
+        raise NotImplementedError()
+
 
 def _raise_if(validation_error: bool, detail: str) -> None:
     if validation_error:

{flwr_nightly-1.13.0.dev20241023 → flwr_nightly-1.13.0.dev20241024}/src/py/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_nightly-1.13.0.dev20241023 → flwr_nightly-1.13.0.dev20241024}/src/py/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`.
+        """

{flwr_nightly-1.13.0.dev20241023 → flwr_nightly-1.13.0.dev20241024}/src/py/flwr/server/superlink/linkstate/sqlite_linkstate.py

@@ -19,13 +19,14 @@
 import json
 import re
 import sqlite3
+import threading
 import time
 from collections.abc import Sequence
 from logging import DEBUG, ERROR, WARNING
 from typing import Any, Optional, Union, cast
 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,
@@ -33,13 +34,19 @@ from flwr.common.constant import (
     Status,
 )
 from flwr.common.typing import Run, RunStatus, UserConfig
-from flwr.proto.node_pb2 import Node  # pylint: disable=E0611
-from flwr.proto.recordset_pb2 import RecordSet  # pylint: disable=E0611
-from flwr.proto.task_pb2 import Task, TaskIns, TaskRes  # pylint: disable=E0611
+
+# pylint: disable=E0611
+from flwr.proto.node_pb2 import Node
+from flwr.proto.recordset_pb2 import RecordSet as ProtoRecordSet
+from flwr.proto.task_pb2 import Task, TaskIns, TaskRes
+
+# pylint: enable=E0611
 from flwr.server.utils.validator import validate_task_ins_or_res
 
 from .linkstate import LinkState
 from .utils import (
+    context_from_bytes,
+    context_to_bytes,
     convert_sint64_to_uint64,
     convert_sint64_values_in_dict_to_uint64,
     convert_uint64_to_sint64,
@@ -92,6 +99,14 @@ CREATE TABLE IF NOT EXISTS run(
 );
 """
 
+SQL_CREATE_TABLE_CONTEXT = """
+CREATE TABLE IF NOT EXISTS context(
+    run_id                INTEGER UNIQUE,
+    context               BLOB,
+    FOREIGN KEY(run_id) REFERENCES run(run_id)
+);
+"""
+
 SQL_CREATE_TABLE_TASK_INS = """
 CREATE TABLE IF NOT EXISTS task_ins(
     task_id                 TEXT UNIQUE,
@@ -152,6 +167,7 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
         """
         self.database_path = database_path
         self.conn: Optional[sqlite3.Connection] = None
+        self.lock = threading.RLock()
 
     def initialize(self, log_queries: bool = False) -> list[tuple[str]]:
         """Create tables if they don't exist yet.
@@ -175,6 +191,7 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
 
         # Create each table if not exists queries
         cur.execute(SQL_CREATE_TABLE_RUN)
+        cur.execute(SQL_CREATE_TABLE_CONTEXT)
         cur.execute(SQL_CREATE_TABLE_TASK_INS)
         cur.execute(SQL_CREATE_TABLE_TASK_RES)
         cur.execute(SQL_CREATE_TABLE_NODE)
@@ -970,6 +987,34 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
             log(ERROR, "`node_id` does not exist.")
             return False
 
+    def get_serverapp_context(self, run_id: int) -> Optional[Context]:
+        """Get the context for the specified `run_id`."""
+        # Retrieve context if any
+        query = "SELECT context FROM context WHERE run_id = ?;"
+        rows = self.query(query, (convert_uint64_to_sint64(run_id),))
+        context = context_from_bytes(rows[0]["context"]) if rows else None
+        return context
+
+    def set_serverapp_context(self, run_id: int, context: Context) -> None:
+        """Set the context for the specified `run_id`."""
+        # Convert context to bytes
+        context_bytes = context_to_bytes(context)
+        sint_run_id = convert_uint64_to_sint64(run_id)
+
+        # Check if any existing Context assigned to the run_id
+        query = "SELECT COUNT(*) FROM context WHERE run_id = ?;"
+        if self.query(query, (sint_run_id,))[0]["COUNT(*)"] > 0:
+            # Update context
+            query = "UPDATE context SET context = ? WHERE run_id = ?;"
+            self.query(query, (context_bytes, sint_run_id))
+        else:
+            try:
+                # Store context
+                query = "INSERT INTO context (run_id, context) VALUES (?, ?);"
+                self.query(query, (sint_run_id, context_bytes))
+            except sqlite3.IntegrityError:
+                raise ValueError(f"Run {run_id} not found") from None
+
     def get_valid_task_ins(self, task_id: str) -> Optional[dict[str, Any]]:
         """Check if the TaskIns exists and is valid (not expired).
 
@@ -1054,7 +1099,7 @@ def task_res_to_dict(task_msg: TaskRes) -> dict[str, Any]:
 
 def dict_to_task_ins(task_dict: dict[str, Any]) -> TaskIns:
     """Turn task_dict into protobuf message."""
-    recordset = RecordSet()
+    recordset = ProtoRecordSet()
     recordset.ParseFromString(task_dict["recordset"])
 
     result = TaskIns(
@@ -1084,7 +1129,7 @@ def dict_to_task_ins(task_dict: dict[str, Any]) -> TaskIns:
 
 def dict_to_task_res(task_dict: dict[str, Any]) -> TaskRes:
     """Turn task_dict into protobuf message."""
-    recordset = RecordSet()
+    recordset = ProtoRecordSet()
     recordset.ParseFromString(task_dict["recordset"])
 
     result = TaskRes(

{flwr_nightly-1.13.0.dev20241023 → flwr_nightly-1.13.0.dev20241024}/src/py/flwr/server/superlink/linkstate/utils.py

@@ -20,10 +20,11 @@ from logging import ERROR
 from os import urandom
 from uuid import uuid4
 
-from flwr.common import log
+from flwr.common import Context, log, serde
 from flwr.common.constant import ErrorCode, Status, SubStatus
 from flwr.common.typing import RunStatus
 from flwr.proto.error_pb2 import Error  # pylint: disable=E0611
+from flwr.proto.message_pb2 import Context as ProtoContext  # pylint: disable=E0611
 from flwr.proto.node_pb2 import Node  # pylint: disable=E0611
 from flwr.proto.task_pb2 import Task, TaskIns, TaskRes  # pylint: disable=E0611
 
@@ -135,6 +136,16 @@ def convert_sint64_values_in_dict_to_uint64(
             data_dict[key] = convert_sint64_to_uint64(data_dict[key])
 
 
+def context_to_bytes(context: Context) -> bytes:
+    """Serialize `Context` to bytes."""
+    return serde.context_to_proto(context).SerializeToString()
+
+
+def context_from_bytes(context_bytes: bytes) -> Context:
+    """Deserialize `Context` from bytes."""
+    return serde.context_from_proto(ProtoContext.FromString(context_bytes))
+
+
 def make_node_unavailable_taskres(ref_taskins: TaskIns) -> TaskRes:
     """Generate a TaskRes with a node unavailable error from a TaskIns."""
     current_time = time.time()

{flwr_nightly-1.13.0.dev20241023 → flwr_nightly-1.13.0.dev20241024}/src/py/flwr/simulation/run_simulation.py

@@ -29,7 +29,7 @@ from typing import Any, Optional
 
 from flwr.cli.config_utils import load_and_validate
 from flwr.client import ClientApp
-from flwr.common import EventType, event, log, now
+from flwr.common import Context, EventType, RecordSet, event, log, now
 from flwr.common.config import get_fused_config_from_dir, parse_config_args
 from flwr.common.constant import RUN_ID_NUM_BYTES, Status
 from flwr.common.logger import (
@@ -40,7 +40,7 @@ from flwr.common.logger import (
 )
 from flwr.common.typing import Run, RunStatus, UserConfig
 from flwr.server.driver import Driver, InMemoryDriver
-from flwr.server.run_serverapp import run as run_server_app
+from flwr.server.run_serverapp import run as _run
 from flwr.server.server_app import ServerApp
 from flwr.server.superlink.fleet import vce
 from flwr.server.superlink.fleet.vce.backend.backend import BackendConfig
@@ -333,11 +333,19 @@ def run_serverapp_th(
                 log(INFO, "Enabling GPU growth for Tensorflow on the server thread.")
                 enable_gpu_growth()
 
+            # Initialize Context
+            context = Context(
+                node_id=0,
+                node_config={},
+                state=RecordSet(),
+                run_config=_server_app_run_config,
+            )
+
             # Run ServerApp
-            run_server_app(
+            _run(
                 driver=_driver,
+                context=context,
                 server_app_dir=_server_app_dir,
-                server_app_run_config=_server_app_run_config,
                 server_app_attr=_server_app_attr,
                 loaded_server_app=_server_app,
             )