flwr 1.15.2__py3-none-any.whl → 1.17.0__py3-none-any.whl

flwr/server/server_app.py

@@ -15,18 +15,18 @@
 """Flower ServerApp."""
 
 
+import inspect
+from collections.abc import Iterator
+from contextlib import contextmanager
 from typing import Callable, Optional
 
 from flwr.common import Context
-from flwr.common.logger import (
-    warn_deprecated_feature_with_example,
-    warn_preview_feature,
-)
+from flwr.common.logger import warn_deprecated_feature_with_example
 from flwr.server.strategy import Strategy
 
 from .client_manager import ClientManager
-from .compat import start_driver
-from .driver import Driver
+from .compat import start_grid
+from .grid import Driver, Grid
 from .server import Server
 from .server_config import ServerConfig
 from .typing import ServerAppCallable, ServerFn
@@ -44,13 +44,33 @@ SERVER_FN_USAGE_EXAMPLE = """
         app = ServerApp(server_fn=server_fn)
 """
 
+GRID_USAGE_EXAMPLE = """
+                app = ServerApp()
 
-class ServerApp:
+                @app.main()
+                def main(grid: Grid, context: Context) -> None:
+                    # Your existing ServerApp code ...
+"""
+
+DRIVER_DEPRECATION_MSG = """
+            The `Driver` class is deprecated, it will be removed in a future release.
+"""
+DRIVER_EXAMPLE_MSG = """
+            Instead, use `Grid` in the signature of your `ServerApp`. For example:
+"""
+
+
+@contextmanager
+def _empty_lifespan(_: Context) -> Iterator[None]:
+    yield
+
+
+class ServerApp:  # pylint: disable=too-many-instance-attributes
     """Flower ServerApp.
 
     Examples
     --------
-    Use the `ServerApp` with an existing `Strategy`:
+    Use the ``ServerApp`` with an existing ``Strategy``:
 
     >>> def server_fn(context: Context):
     >>>     server_config = ServerConfig(num_rounds=3)
@@ -62,12 +82,12 @@ class ServerApp:
     >>>
     >>> app = ServerApp(server_fn=server_fn)
 
-    Use the `ServerApp` with a custom main function:
+    Use the ``ServerApp`` with a custom main function:
 
     >>> app = ServerApp()
     >>>
     >>> @app.main()
-    >>> def main(driver: Driver, context: Context) -> None:
+    >>> def main(grid: Grid, context: Context) -> None:
     >>>    print("ServerApp running")
     """
 
@@ -105,29 +125,31 @@ class ServerApp:
         self._client_manager = client_manager
         self._server_fn = server_fn
         self._main: Optional[ServerAppCallable] = None
+        self._lifespan = _empty_lifespan
 
-    def __call__(self, driver: Driver, context: Context) -> None:
+    def __call__(self, grid: Grid, context: Context) -> None:
         """Execute `ServerApp`."""
-        # Compatibility mode
-        if not self._main:
-            if self._server_fn:
-                # Execute server_fn()
-                components = self._server_fn(context)
-                self._server = components.server
-                self._config = components.config
-                self._strategy = components.strategy
-                self._client_manager = components.client_manager
-            start_driver(
-                server=self._server,
-                config=self._config,
-                strategy=self._strategy,
-                client_manager=self._client_manager,
-                driver=driver,
-            )
-            return
+        with self._lifespan(context):
+            # Compatibility mode
+            if not self._main:
+                if self._server_fn:
+                    # Execute server_fn()
+                    components = self._server_fn(context)
+                    self._server = components.server
+                    self._config = components.config
+                    self._strategy = components.strategy
+                    self._client_manager = components.client_manager
+                start_grid(
+                    server=self._server,
+                    config=self._config,
+                    strategy=self._strategy,
+                    client_manager=self._client_manager,
+                    grid=grid,
+                )
+                return
 
-        # New execution mode
-        self._main(driver, context)
+            # New execution mode
+            self._main(grid, context)
 
     def main(self) -> Callable[[ServerAppCallable], ServerAppCallable]:
         """Return a decorator that registers the main fn with the server app.
@@ -137,7 +159,7 @@ class ServerApp:
         >>> app = ServerApp()
         >>>
         >>> @app.main()
-        >>> def main(driver: Driver, context: Context) -> None:
+        >>> def main(grid: Grid, context: Context) -> None:
         >>>    print("ServerApp running")
         """
 
@@ -162,12 +184,20 @@ class ServerApp:
                     >>> app = ServerApp()
                     >>>
                     >>> @app.main()
-                    >>> def main(driver: Driver, context: Context) -> None:
+                    >>> def main(grid: Grid, context: Context) -> None:
                     >>>    print("ServerApp running")
                     """,
                 )
 
-            warn_preview_feature("ServerApp-register-main-function")
+            sig = inspect.signature(main_fn)
+            param = list(sig.parameters.values())[0]
+            # Check if parameter name or the annotation should be updated
+            if param.name == "driver" or param.annotation is Driver:
+                warn_deprecated_feature_with_example(
+                    deprecation_message=DRIVER_DEPRECATION_MSG,
+                    example_message=DRIVER_EXAMPLE_MSG,
+                    code_example=GRID_USAGE_EXAMPLE,
+                )
 
             # Register provided function with the ServerApp object
             self._main = main_fn
@@ -177,6 +207,69 @@ class ServerApp:
 
         return main_decorator
 
+    def lifespan(
+        self,
+    ) -> Callable[
+        [Callable[[Context], Iterator[None]]], Callable[[Context], Iterator[None]]
+    ]:
+        """Return a decorator that registers the lifespan fn with the server app.
+
+        The decorated function should accept a `Context` object and use `yield`
+        to define enter and exit behavior.
+
+        Examples
+        --------
+        >>> app = ServerApp()
+        >>>
+        >>> @app.lifespan()
+        >>> def lifespan(context: Context) -> None:
+        >>>     # Perform initialization tasks before the app starts
+        >>>     print("Initializing ServerApp")
+        >>>
+        >>>     yield  # ServerApp is running
+        >>>
+        >>>     # Perform cleanup tasks after the app stops
+        >>>     print("Cleaning up ServerApp")
+        """
+
+        def lifespan_decorator(
+            lifespan_fn: Callable[[Context], Iterator[None]],
+        ) -> Callable[[Context], Iterator[None]]:
+            """Register the lifespan fn with the ServerApp object."""
+
+            @contextmanager
+            def decorated_lifespan(context: Context) -> Iterator[None]:
+                # Execute the code before `yield` in lifespan_fn
+                try:
+                    if not isinstance(it := lifespan_fn(context), Iterator):
+                        raise StopIteration
+                    next(it)
+                except StopIteration:
+                    raise RuntimeError(
+                        "lifespan function should yield at least once."
+                    ) from None
+
+                try:
+                    # Enter the context
+                    yield
+                finally:
+                    try:
+                        # Execute the code after `yield` in lifespan_fn
+                        next(it)
+                    except StopIteration:
+                        pass
+                    else:
+                        raise RuntimeError("lifespan function should only yield once.")
+
+            # Register provided function with the ServerApp object
+            # Ignore mypy error because of different argument names (`_` vs `context`)
+            self._lifespan = decorated_lifespan  # type: ignore
+
+            # Return provided function unmodified
+            return lifespan_fn
+
+        return lifespan_decorator
+
 
 class LoadServerAppError(Exception):
     """Error when trying to load `ServerApp`."""

flwr/server/serverapp/app.py

@@ -60,7 +60,7 @@ from flwr.proto.serverappio_pb2 import (  # pylint: disable=E0611
     PullServerAppInputsResponse,
     PushServerAppOutputsRequest,
 )
-from flwr.server.driver.grpc_driver import GrpcDriver
+from flwr.server.grid.grpc_grid import GrpcGrid
 from flwr.server.run_serverapp import run as run_
 
 
@@ -106,7 +106,7 @@ def run_serverapp(  # pylint: disable=R0914, disable=W0212, disable=R0915
     certificates: Optional[bytes] = None,
 ) -> None:
     """Run Flower ServerApp process."""
-    driver = GrpcDriver(
+    grid = GrpcGrid(
         serverappio_service_address=serverappio_api_address,
         root_certificates=certificates,
     )
@@ -123,7 +123,7 @@ def run_serverapp(  # pylint: disable=R0914, disable=W0212, disable=R0915
             # Pull ServerAppInputs from LinkState
             req = PullServerAppInputsRequest()
             log(DEBUG, "[flwr-serverapp] Pull ServerAppInputs")
-            res: PullServerAppInputsResponse = driver._stub.PullServerAppInputs(req)
+            res: PullServerAppInputsResponse = grid._stub.PullServerAppInputs(req)
             if not res.HasField("run"):
                 sleep(3)
                 run_status = None
@@ -135,14 +135,14 @@ def run_serverapp(  # pylint: disable=R0914, disable=W0212, disable=R0915
 
             hash_run_id = get_sha256_hash(run.run_id)
 
-            driver.set_run(run.run_id)
+            grid.set_run(run.run_id)
 
             # Start log uploader for this run
             log_uploader = start_log_uploader(
                 log_queue=log_queue,
                 node_id=0,
                 run_id=run.run_id,
-                stub=driver._stub,
+                stub=grid._stub,
             )
 
             log(DEBUG, "[flwr-serverapp] Start FAB installation.")
@@ -173,7 +173,7 @@ def run_serverapp(  # pylint: disable=R0914, disable=W0212, disable=R0915
 
             # Change status to Running
             run_status_proto = run_status_to_proto(RunStatus(Status.RUNNING, "", ""))
-            driver._stub.UpdateRunStatus(
+            grid._stub.UpdateRunStatus(
                 UpdateRunStatusRequest(run_id=run.run_id, run_status=run_status_proto)
             )
 
@@ -182,9 +182,9 @@ def run_serverapp(  # pylint: disable=R0914, disable=W0212, disable=R0915
                 event_details={"run-id-hash": hash_run_id},
             )
 
-            # Load and run the ServerApp with the Driver
+            # Load and run the ServerApp with the Grid
             updated_context = run_(
-                driver=driver,
+                grid=grid,
                 server_app_dir=app_path,
                 server_app_attr=server_app_attr,
                 context=context,
@@ -196,7 +196,7 @@ def run_serverapp(  # pylint: disable=R0914, disable=W0212, disable=R0915
             out_req = PushServerAppOutputsRequest(
                 run_id=run.run_id, context=context_proto
             )
-            _ = driver._stub.PushServerAppOutputs(out_req)
+            _ = grid._stub.PushServerAppOutputs(out_req)
 
             run_status = RunStatus(Status.FINISHED, SubStatus.COMPLETED, "")
         except RunNotRunningException:
@@ -221,7 +221,7 @@ def run_serverapp(  # pylint: disable=R0914, disable=W0212, disable=R0915
             # Update run status
             if run_status:
                 run_status_proto = run_status_to_proto(run_status)
-                driver._stub.UpdateRunStatus(
+                grid._stub.UpdateRunStatus(
                     UpdateRunStatusRequest(
                         run_id=run.run_id, run_status=run_status_proto
                     )

flwr/server/superlink/fleet/grpc_rere/fleet_servicer.py

@@ -103,11 +103,11 @@ class FleetServicer(fleet_pb2_grpc.FleetServicer):
         if request.messages_list:
             log(
                 INFO,
-                "[Fleet.PushMessages] Push results from node_id=%s",
+                "[Fleet.PushMessages] Push replies from node_id=%s",
                 request.messages_list[0].metadata.src_node_id,
             )
         else:
-            log(INFO, "[Fleet.PushMessages] No task results to push")
+            log(INFO, "[Fleet.PushMessages] No replies to push")
 
         try:
             res = message_handler.push_messages(

flwr/server/superlink/fleet/message_handler/message_handler.py

@@ -18,13 +18,12 @@
 from typing import Optional
 from uuid import UUID
 
+from flwr.common import Message
 from flwr.common.constant import Status
 from flwr.common.serde import (
     fab_to_proto,
     message_from_proto,
-    message_from_taskins,
     message_to_proto,
-    message_to_taskres,
     user_config_to_proto,
 )
 from flwr.common.typing import Fab, InvalidRunStatusException
@@ -48,7 +47,6 @@ from flwr.proto.run_pb2 import (  # pylint: disable=E0611
     GetRunResponse,
     Run,
 )
-from flwr.proto.task_pb2 import TaskIns  # pylint: disable=E0611
 from flwr.server.superlink.ffs.ffs import Ffs
 from flwr.server.superlink.linkstate import LinkState
 from flwr.server.superlink.utils import check_abort
@@ -92,13 +90,12 @@ def pull_messages(
     node = request.node  # pylint: disable=no-member
     node_id: int = node.node_id
 
-    # Retrieve TaskIns from State
-    task_ins_list: list[TaskIns] = state.get_task_ins(node_id=node_id, limit=1)
+    # Retrieve Message from State
+    message_list: list[Message] = state.get_message_ins(node_id=node_id, limit=1)
 
     # Convert to Messages
     msg_proto = []
-    for task_ins in task_ins_list:
-        msg = message_from_taskins(task_ins)
+    for msg in message_list:
         msg_proto.append(message_to_proto(msg))
 
     return PullMessagesResponse(messages_list=msg_proto)
@@ -108,21 +105,20 @@ def push_messages(
     request: PushMessagesRequest, state: LinkState
 ) -> PushMessagesResponse:
     """Push Messages handler."""
-    # Convert Message to TaskRes
+    # Convert Message from proto
     msg = message_from_proto(message_proto=request.messages_list[0])
-    task_res = message_to_taskres(msg)
 
     # Abort if the run is not running
     abort_msg = check_abort(
-        task_res.run_id,
+        msg.metadata.run_id,
         [Status.PENDING, Status.STARTING, Status.FINISHED],
         state,
     )
     if abort_msg:
         raise InvalidRunStatusException(abort_msg)
 
-    # Store TaskRes in State
-    message_id: Optional[UUID] = state.store_task_res(task_res=task_res)
+    # Store Message in State
+    message_id: Optional[UUID] = state.store_message_res(message=msg)
 
     # Build response
     response = PushMessagesResponse(

flwr/server/superlink/fleet/vce/backend/backend.py

@@ -21,9 +21,9 @@ from typing import Callable
 from flwr.client.client_app import ClientApp
 from flwr.common.context import Context
 from flwr.common.message import Message
-from flwr.common.typing import ConfigsRecordValues
+from flwr.common.typing import ConfigRecordValues
 
-BackendConfig = dict[str, dict[str, ConfigsRecordValues]]
+BackendConfig = dict[str, dict[str, ConfigRecordValues]]
 
 
 class Backend(ABC):
@@ -45,7 +45,7 @@ class Backend(ABC):
     def num_workers(self) -> int:
         """Return number of workers in the backend.
 
-        This is the number of TaskIns that can be processed concurrently.
+        This is the number of Messages that can be processed concurrently.
         """
         return 0
 

flwr/server/superlink/fleet/vce/backend/raybackend.py

@@ -26,7 +26,7 @@ from flwr.common.constant import PARTITION_ID_KEY
 from flwr.common.context import Context
 from flwr.common.logger import log
 from flwr.common.message import Message
-from flwr.common.typing import ConfigsRecordValues
+from flwr.common.typing import ConfigRecordValues
 from flwr.simulation.ray_transport.ray_actor import BasicActorPool, ClientAppActor
 from flwr.simulation.ray_transport.utils import enable_tf_gpu_growth
 
@@ -104,7 +104,7 @@ class RayBackend(Backend):
         if not ray.is_initialized():
             ray_init_args: dict[
                 str,
-                ConfigsRecordValues,
+                ConfigRecordValues,
             ] = {}
 
             if backend_config.get(self.init_args_key):

flwr/server/superlink/fleet/vce/vce_api.py

@@ -29,6 +29,7 @@ from typing import Callable, Optional
 from flwr.client.client_app import ClientApp, ClientAppException, LoadClientAppError
 from flwr.client.clientapp.utils import get_load_client_app_fn
 from flwr.client.run_info_store import DeprecatedRunInfoStore
+from flwr.common import Message
 from flwr.common.constant import (
     NUM_PARTITIONS_KEY,
     PARTITION_ID_KEY,
@@ -37,9 +38,7 @@ from flwr.common.constant import (
 )
 from flwr.common.logger import log
 from flwr.common.message import Error
-from flwr.common.serde import message_from_taskins, message_to_taskres
 from flwr.common.typing import Run
-from flwr.proto.task_pb2 import TaskIns, TaskRes  # pylint: disable=E0611
 from flwr.server.superlink.linkstate import LinkState, LinkStateFactory
 
 from .backend import Backend, error_messages_backends, supported_backends
@@ -87,33 +86,33 @@ def _register_node_info_stores(
 
 # pylint: disable=too-many-arguments,too-many-locals
 def worker(
-    taskins_queue: "Queue[TaskIns]",
-    taskres_queue: "Queue[TaskRes]",
+    messageins_queue: Queue[Message],
+    messageres_queue: Queue[Message],
     node_info_store: dict[int, DeprecatedRunInfoStore],
     backend: Backend,
     f_stop: threading.Event,
 ) -> None:
-    """Get TaskIns from queue and pass it to an actor in the pool to execute it."""
+    """Process messages from the queue, execute them, update context, and enqueue
+    replies."""
     while not f_stop.is_set():
         out_mssg = None
         try:
             # Fetch from queue with timeout. We use a timeout so
             # the stopping event can be evaluated even when the queue is empty.
-            task_ins: TaskIns = taskins_queue.get(timeout=1.0)
-            node_id = task_ins.task.consumer.node_id
+            message: Message = messageins_queue.get(timeout=1.0)
+            node_id = message.metadata.dst_node_id
 
             # Retrieve context
-            context = node_info_store[node_id].retrieve_context(run_id=task_ins.run_id)
-
-            # Convert TaskIns to Message
-            message = message_from_taskins(task_ins)
+            context = node_info_store[node_id].retrieve_context(
+                run_id=message.metadata.run_id
+            )
 
             # Let backend process message
             out_mssg, updated_context = backend.process_message(message, context)
 
             # Update Context
             node_info_store[node_id].update_context(
-                task_ins.run_id, context=updated_context
+                message.metadata.run_id, context=updated_context
             )
         except Empty:
             # An exception raised if queue.get times out
@@ -131,41 +130,37 @@ def worker(
                 e_code = ErrorCode.UNKNOWN
 
             reason = str(type(ex)) + ":<'" + str(ex) + "'>"
-            out_mssg = message.create_error_reply(
-                error=Error(code=e_code, reason=reason)
-            )
+            out_mssg = Message(Error(code=e_code, reason=reason), reply_to=message)
 
         finally:
             if out_mssg:
-                # Convert to TaskRes
-                task_res = message_to_taskres(out_mssg)
-                # Store TaskRes in state
-                taskres_queue.put(task_res)
+                # Store reply Messages in state
+                messageres_queue.put(out_mssg)
 
 
-def add_taskins_to_queue(
+def add_messages_to_queue(
     state: LinkState,
-    queue: "Queue[TaskIns]",
+    queue: Queue[Message],
     nodes_mapping: NodeToPartitionMapping,
     f_stop: threading.Event,
 ) -> None:
-    """Put TaskIns in a queue from State."""
+    """Put Messages in the queue from the LinkState."""
     while not f_stop.is_set():
         for node_id in nodes_mapping.keys():
-            task_ins_list = state.get_task_ins(node_id=node_id, limit=1)
-            for task_ins in task_ins_list:
-                queue.put(task_ins)
+            message_ins_list = state.get_message_ins(node_id=node_id, limit=1)
+            for msg in message_ins_list:
+                queue.put(msg)
         sleep(0.1)
 
 
-def put_taskres_into_state(
-    state: LinkState, queue: "Queue[TaskRes]", f_stop: threading.Event
+def put_message_into_state(
+    state: LinkState, queue: Queue[Message], f_stop: threading.Event
 ) -> None:
-    """Put TaskRes into State from a queue."""
+    """Store reply Messages into the LinkState from the queue."""
     while not f_stop.is_set():
         try:
-            taskres = queue.get(timeout=1.0)
-            state.store_task_res(taskres)
+            message_reply = queue.get(timeout=1.0)
+            state.store_message_res(message_reply)
         except Empty:
             # queue is empty when timeout was triggered
             pass
@@ -181,8 +176,8 @@ def run_api(
     f_stop: threading.Event,
 ) -> None:
     """Run the VCE."""
-    taskins_queue: Queue[TaskIns] = Queue()
-    taskres_queue: Queue[TaskRes] = Queue()
+    messageins_queue: Queue[Message] = Queue()
+    messageres_queue: Queue[Message] = Queue()
 
     try:
 
@@ -196,10 +191,10 @@ def run_api(
         state = state_factory.state()
 
         extractor_th = threading.Thread(
-            target=add_taskins_to_queue,
+            target=add_messages_to_queue,
             args=(
                 state,
-                taskins_queue,
+                messageins_queue,
                 nodes_mapping,
                 f_stop,
             ),
@@ -207,10 +202,10 @@ def run_api(
         extractor_th.start()
 
         injector_th = threading.Thread(
-            target=put_taskres_into_state,
+            target=put_message_into_state,
             args=(
                 state,
-                taskres_queue,
+                messageres_queue,
                 f_stop,
             ),
         )
@@ -220,8 +215,8 @@ def run_api(
             _ = [
                 executor.submit(
                     worker,
-                    taskins_queue,
-                    taskres_queue,
+                    messageins_queue,
+                    messageres_queue,
                     node_info_stores,
                     backend,
                     f_stop,