flwr 1.15.1__py3-none-any.whl → 1.16.0__py3-none-any.whl

flwr/server/driver/inmemory_driver.py

@@ -23,7 +23,6 @@ from uuid import UUID
 
 from flwr.common import DEFAULT_TTL, Message, Metadata, RecordSet
 from flwr.common.constant import SUPERLINK_NODE_ID
-from flwr.common.serde import message_from_taskres, message_to_taskins
 from flwr.common.typing import Run
 from flwr.proto.node_pb2 import Node  # pylint: disable=E0611
 from flwr.server.superlink.linkstate import LinkStateFactory
@@ -60,6 +59,7 @@ class InMemoryDriver(Driver):
             and message.metadata.message_id == ""
             and message.metadata.reply_to_message == ""
             and message.metadata.ttl > 0
+            and message.metadata.delivered_at == ""
         ):
             raise ValueError(f"Invalid message: {message}")
 
@@ -109,9 +109,9 @@ class InMemoryDriver(Driver):
         )
         return Message(metadata=metadata, content=content)
 
-    def get_node_ids(self) -> list[int]:
+    def get_node_ids(self) -> Iterable[int]:
         """Get node IDs."""
-        return list(self.state.get_nodes(cast(Run, self._run).run_id))
+        return self.state.get_nodes(cast(Run, self._run).run_id)
 
     def push_messages(self, messages: Iterable[Message]) -> Iterable[str]:
         """Push messages to specified node IDs.
@@ -119,19 +119,16 @@ class InMemoryDriver(Driver):
         This method takes an iterable of messages and sends each message
         to the node specified in `dst_node_id`.
         """
-        task_ids: list[str] = []
+        msg_ids: list[str] = []
         for msg in messages:
             # Check message
             self._check_message(msg)
-            # Convert Message to TaskIns
-            taskins = message_to_taskins(msg)
             # Store in state
-            taskins.task.pushed_at = time.time()
-            task_id = self.state.store_task_ins(taskins)
-            if task_id:
-                task_ids.append(str(task_id))
+            msg_id = self.state.store_message_ins(msg)
+            if msg_id:
+                msg_ids.append(str(msg_id))
 
-        return task_ids
+        return msg_ids
 
     def pull_messages(self, message_ids: Iterable[str]) -> Iterable[Message]:
         """Pull messages based on message IDs.
@@ -140,17 +137,16 @@ class InMemoryDriver(Driver):
         set of given message IDs.
         """
         msg_ids = {UUID(msg_id) for msg_id in message_ids}
-        # Pull TaskRes
-        task_res_list = self.state.get_task_res(task_ids=msg_ids)
-        # Delete tasks in state
-        # Delete the TaskIns/TaskRes pairs if TaskRes is found
-        task_ins_ids_to_delete = {
-            UUID(task_res.task.ancestry[0]) for task_res in task_res_list
+        # Pull Messages
+        message_res_list = self.state.get_message_res(message_ids=msg_ids)
+        # Get IDs of Messages these replies are for
+        message_ins_ids_to_delete = {
+            UUID(msg_res.metadata.reply_to_message) for msg_res in message_res_list
         }
-        self.state.delete_tasks(task_ins_ids=task_ins_ids_to_delete)
-        # Convert TaskRes to Message
-        msgs = [message_from_taskres(taskres) for taskres in task_res_list]
-        return msgs
+        # Delete
+        self.state.delete_messages(message_ins_ids=message_ins_ids_to_delete)
+
+        return message_res_list
 
     def send_and_receive(
         self,

flwr/server/run_serverapp.py

@@ -15,11 +15,10 @@
 """Run ServerApp."""
 
 
-from logging import DEBUG, ERROR
+from logging import DEBUG
 from typing import Optional
 
-from flwr.common import Context, EventType, event
-from flwr.common.exit_handlers import register_exit_handlers
+from flwr.common import Context
 from flwr.common.logger import log
 from flwr.common.object_ref import load_app
 
@@ -64,13 +63,3 @@ def run(
 
     log(DEBUG, "ServerApp finished running.")
     return context
-
-
-def run_server_app() -> None:
-    """Run Flower server app."""
-    event(EventType.RUN_SERVER_APP_ENTER)
-    log(
-        ERROR,
-        "The command `flower-server-app` has been replaced by `flwr run`.",
-    )
-    register_exit_handlers(event_type=EventType.RUN_SERVER_APP_LEAVE)

flwr/server/server_app.py

@@ -15,6 +15,8 @@
 """Flower ServerApp."""
 
 
+from collections.abc import Iterator
+from contextlib import contextmanager
 from typing import Callable, Optional
 
 from flwr.common import Context
@@ -45,7 +47,12 @@ SERVER_FN_USAGE_EXAMPLE = """
 """
 
 
-class ServerApp:
+@contextmanager
+def _empty_lifespan(_: Context) -> Iterator[None]:
+    yield
+
+
+class ServerApp:  # pylint: disable=too-many-instance-attributes
     """Flower ServerApp.
 
     Examples
@@ -105,29 +112,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:
         """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_driver(
+                    server=self._server,
+                    config=self._config,
+                    strategy=self._strategy,
+                    client_manager=self._client_manager,
+                    driver=driver,
+                )
+                return
 
-        # New execution mode
-        self._main(driver, context)
+            # New execution mode
+            self._main(driver, context)
 
     def main(self) -> Callable[[ServerAppCallable], ServerAppCallable]:
         """Return a decorator that registers the main fn with the server app.
@@ -177,6 +186,70 @@ 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."""
+            warn_preview_feature("ServerApp-register-lifespan-function")
+
+            @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/superlink/driver/serverappio_servicer.py

@@ -22,8 +22,8 @@ from uuid import UUID
 
 import grpc
 
-from flwr.common import ConfigsRecord, now
-from flwr.common.constant import Status
+from flwr.common import ConfigsRecord, Message
+from flwr.common.constant import SUPERLINK_NODE_ID, Status
 from flwr.common.logger import log
 from flwr.common.serde import (
     context_from_proto,
@@ -31,9 +31,7 @@ from flwr.common.serde import (
     fab_from_proto,
     fab_to_proto,
     message_from_proto,
-    message_from_taskres,
     message_to_proto,
-    message_to_taskins,
     run_status_from_proto,
     run_status_to_proto,
     run_to_proto,
@@ -69,12 +67,11 @@ from flwr.proto.serverappio_pb2 import (  # pylint: disable=E0611
     PushServerAppOutputsRequest,
     PushServerAppOutputsResponse,
 )
-from flwr.proto.task_pb2 import TaskRes  # pylint: disable=E0611
 from flwr.server.superlink.ffs.ffs import Ffs
 from flwr.server.superlink.ffs.ffs_factory import FfsFactory
 from flwr.server.superlink.linkstate import LinkState, LinkStateFactory
 from flwr.server.superlink.utils import abort_if
-from flwr.server.utils.validator import validate_task_ins_or_res
+from flwr.server.utils.validator import validate_message
 
 
 class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
@@ -151,9 +148,6 @@ class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
             context,
         )
 
-        # Set pushed_at (timestamp in seconds)
-        pushed_at = now().timestamp()
-
         # Validate request and insert in State
         _raise_if(
             validation_error=len(request.messages_list) == 0,
@@ -164,21 +158,19 @@ class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
         while request.messages_list:
             message_proto = request.messages_list.pop(0)
             message = message_from_proto(message_proto=message_proto)
-            task_ins = message_to_taskins(message=message)
-            task_ins.task.pushed_at = pushed_at
-            validation_errors = validate_task_ins_or_res(task_ins)
+            validation_errors = validate_message(message, is_reply_message=False)
             _raise_if(
                 validation_error=bool(validation_errors),
                 request_name="PushMessages",
                 detail=", ".join(validation_errors),
             )
             _raise_if(
-                validation_error=request.run_id != task_ins.run_id,
+                validation_error=request.run_id != message.metadata.run_id,
                 request_name="PushMessages",
-                detail="`task_ins` has mismatched `run_id`",
+                detail="`Message.metadata` has mismatched `run_id`",
             )
             # Store
-            message_id: Optional[UUID] = state.store_task_ins(task_ins=task_ins)
+            message_id: Optional[UUID] = state.store_message_ins(message=message)
             message_ids.append(message_id)
 
         return PushInsMessagesResponse(
@@ -204,32 +196,34 @@ class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
             context,
         )
 
-        # Convert each task_id str to UUID
+        # Convert each message_id str to UUID
         message_ids: set[UUID] = {
             UUID(message_id) for message_id in request.message_ids
         }
 
         # Read from state
-        task_res_list: list[TaskRes] = state.get_task_res(task_ids=message_ids)
+        messages_res: list[Message] = state.get_message_res(message_ids=message_ids)
 
-        # Convert to Messages
-        messages_list = []
-        while task_res_list:
-            task_res = task_res_list.pop(0)
-            _raise_if(
-                validation_error=request.run_id != task_res.run_id,
-                request_name="PullMessages",
-                detail="`task_res` has mismatched `run_id`",
-            )
-            message = message_from_taskres(taskres=task_res)
-            messages_list.append(message_to_proto(message))
-
-        # Delete the TaskIns/TaskRes pairs if TaskRes is found
-        task_ins_ids_to_delete = {
-            UUID(task_res.task.ancestry[0]) for task_res in task_res_list
+        # Delete the instruction Messages and their replies if found
+        message_ins_ids_to_delete = {
+            UUID(msg_res.metadata.reply_to_message) for msg_res in messages_res
         }
 
-        state.delete_tasks(task_ins_ids=task_ins_ids_to_delete)
+        state.delete_messages(message_ins_ids=message_ins_ids_to_delete)
+
+        # Convert Messages to proto
+        messages_list = []
+        while messages_res:
+            msg = messages_res.pop(0)
+
+            # Skip `run_id` check for SuperLink generated replies
+            if msg.metadata.src_node_id != SUPERLINK_NODE_ID:
+                _raise_if(
+                    validation_error=request.run_id != msg.metadata.run_id,
+                    request_name="PullMessages",
+                    detail="`message.metadata` has mismatched `run_id`",
+                )
+            messages_list.append(message_to_proto(msg))
 
         return PullResMessagesResponse(messages_list=messages_list)
 

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

@@ -15,17 +15,15 @@
 """Fleet API message handlers."""
 
 
-import time
 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
@@ -49,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
@@ -93,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)
@@ -109,24 +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)
 
-    # Set pushed_at (timestamp in seconds)
-    task_res.task.pushed_at = time.time()
-
-    # 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

@@ -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/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
@@ -137,36 +136,33 @@ def worker(
 
         finally:
             if out_mssg:
-                # Convert to TaskRes
-                task_res = message_to_taskres(out_mssg)
-                # Store TaskRes in state
-                task_res.task.pushed_at = time.time()
-                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
@@ -182,8 +178,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:
 
@@ -197,10 +193,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,
             ),
@@ -208,10 +204,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,
             ),
         )
@@ -221,8 +217,8 @@ def run_api(
             _ = [
                 executor.submit(
                     worker,
-                    taskins_queue,
-                    taskres_queue,
+                    messageins_queue,
+                    messageres_queue,
                     node_info_stores,
                     backend,
                     f_stop,