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

flwr/server/{driver/inmemory_driver.py → grid/inmemory_grid.py}

@@ -1,4 +1,4 @@
-# Copyright 2024 Flower Labs GmbH. All Rights Reserved.
+# Copyright 2025 Flower Labs GmbH. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -12,35 +12,37 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""Flower in-memory Driver."""
+"""Flower in-memory Grid."""
 
 
 import time
-import warnings
 from collections.abc import Iterable
 from typing import Optional, cast
 from uuid import UUID
 
-from flwr.common import DEFAULT_TTL, Message, Metadata, RecordSet
+from flwr.common import Message, RecordDict
 from flwr.common.constant import SUPERLINK_NODE_ID
+from flwr.common.logger import warn_deprecated_feature
 from flwr.common.typing import Run
 from flwr.proto.node_pb2 import Node  # pylint: disable=E0611
 from flwr.server.superlink.linkstate import LinkStateFactory
 
-from .driver import Driver
+from .grid import Grid
 
 
-class InMemoryDriver(Driver):
-    """`InMemoryDriver` class provides an interface to the ServerAppIo API.
+class InMemoryGrid(Grid):
+    """`InMemoryGrid` class provides an interface to the ServerAppIo API.
 
     Parameters
     ----------
     state_factory : StateFactory
-        A StateFactory embedding a state that this driver can interface with.
+        A StateFactory embedding a state that this grid can interface with.
     pull_interval : float (default=0.1)
         Sleep duration between calls to `pull_messages`.
     """
 
+    _deprecation_warning_logged = False
+
     def __init__(
         self,
         state_factory: LinkStateFactory,
@@ -54,10 +56,8 @@ class InMemoryDriver(Driver):
     def _check_message(self, message: Message) -> None:
         # Check if the message is valid
         if not (
-            message.metadata.run_id == cast(Run, self._run).run_id
-            and message.metadata.src_node_id == self.node.node_id
-            and message.metadata.message_id == ""
-            and message.metadata.reply_to_message == ""
+            message.metadata.message_id == ""
+            and message.metadata.reply_to_message_id == ""
             and message.metadata.ttl > 0
             and message.metadata.delivered_at == ""
         ):
@@ -77,7 +77,7 @@ class InMemoryDriver(Driver):
 
     def create_message(  # pylint: disable=too-many-arguments,R0917
         self,
-        content: RecordSet,
+        content: RecordDict,
         message_type: str,
         dst_node_id: int,
         group_id: str,
@@ -88,26 +88,13 @@ class InMemoryDriver(Driver):
         This method constructs a new `Message` with given content and metadata.
         The `run_id` and `src_node_id` will be set automatically.
         """
-        if ttl:
-            warnings.warn(
-                "A custom TTL was set, but note that the SuperLink does not enforce "
-                "the TTL yet. The SuperLink will start enforcing the TTL in a future "
-                "version of Flower.",
-                stacklevel=2,
+        if not InMemoryGrid._deprecation_warning_logged:
+            InMemoryGrid._deprecation_warning_logged = True
+            warn_deprecated_feature(
+                "`Driver.create_message` / `Grid.create_message` is deprecated."
+                "Use `Message` constructor instead."
             )
-        ttl_ = DEFAULT_TTL if ttl is None else ttl
-
-        metadata = Metadata(
-            run_id=cast(Run, self._run).run_id,
-            message_id="",  # Will be set by the server
-            src_node_id=self.node.node_id,
-            dst_node_id=dst_node_id,
-            reply_to_message="",
-            group_id=group_id,
-            ttl=ttl_,
-            message_type=message_type,
-        )
-        return Message(metadata=metadata, content=content)
+        return Message(content, dst_node_id, message_type, ttl=ttl, group_id=group_id)
 
     def get_node_ids(self) -> Iterable[int]:
         """Get node IDs."""
@@ -121,6 +108,9 @@ class InMemoryDriver(Driver):
         """
         msg_ids: list[str] = []
         for msg in messages:
+            # Populate metadata
+            msg.metadata.__dict__["_run_id"] = cast(Run, self._run).run_id
+            msg.metadata.__dict__["_src_node_id"] = self.node.node_id
             # Check message
             self._check_message(msg)
             # Store in state
@@ -141,7 +131,7 @@ class InMemoryDriver(Driver):
         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
+            UUID(msg_res.metadata.reply_to_message_id) for msg_res in message_res_list
         }
         # Delete
         self.state.delete_messages(message_ins_ids=message_ins_ids_to_delete)
@@ -170,7 +160,7 @@ class InMemoryDriver(Driver):
             res_msgs = self.pull_messages(msg_ids)
             ret.extend(res_msgs)
             msg_ids.difference_update(
-                {msg.metadata.reply_to_message for msg in res_msgs}
+                {msg.metadata.reply_to_message_id for msg in res_msgs}
             )
             if len(msg_ids) == 0:
                 break

flwr/server/run_serverapp.py

@@ -22,18 +22,18 @@ from flwr.common import Context
 from flwr.common.logger import log
 from flwr.common.object_ref import load_app
 
-from .driver import Driver
+from .grid import Grid
 from .server_app import LoadServerAppError, ServerApp
 
 
 def run(
-    driver: Driver,
+    grid: Grid,
     context: Context,
     server_app_dir: str,
     server_app_attr: Optional[str] = None,
     loaded_server_app: Optional[ServerApp] = None,
 ) -> Context:
-    """Run ServerApp with a given Driver."""
+    """Run ServerApp with a given Grid."""
     if not (server_app_attr is None) ^ (loaded_server_app is None):
         raise ValueError(
             "Either `server_app_attr` or `loaded_server_app` should be set "
@@ -59,7 +59,7 @@ def run(
     server_app = _load()
 
     # Call ServerApp
-    server_app(driver=driver, context=context)
+    server_app(grid=grid, context=context)
 
     log(DEBUG, "ServerApp finished running.")
     return context

flwr/server/server_app.py

@@ -15,20 +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
@@ -46,6 +44,21 @@ SERVER_FN_USAGE_EXAMPLE = """
         app = ServerApp(server_fn=server_fn)
 """
 
+GRID_USAGE_EXAMPLE = """
+                app = 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]:
@@ -57,7 +70,7 @@ class ServerApp:  # pylint: disable=too-many-instance-attributes
 
     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)
@@ -69,12 +82,12 @@ class ServerApp:  # pylint: disable=too-many-instance-attributes
     >>>
     >>> 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")
     """
 
@@ -114,7 +127,7 @@ class ServerApp:  # pylint: disable=too-many-instance-attributes
         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`."""
         with self._lifespan(context):
             # Compatibility mode
@@ -126,17 +139,17 @@ class ServerApp:  # pylint: disable=too-many-instance-attributes
                     self._config = components.config
                     self._strategy = components.strategy
                     self._client_manager = components.client_manager
-                start_driver(
+                start_grid(
                     server=self._server,
                     config=self._config,
                     strategy=self._strategy,
                     client_manager=self._client_manager,
-                    driver=driver,
+                    grid=grid,
                 )
                 return
 
             # New execution mode
-            self._main(driver, context)
+            self._main(grid, context)
 
     def main(self) -> Callable[[ServerAppCallable], ServerAppCallable]:
         """Return a decorator that registers the main fn with the server app.
@@ -146,7 +159,7 @@ class ServerApp:  # pylint: disable=too-many-instance-attributes
         >>> app = ServerApp()
         >>>
         >>> @app.main()
-        >>> def main(driver: Driver, context: Context) -> None:
+        >>> def main(grid: Grid, context: Context) -> None:
         >>>    print("ServerApp running")
         """
 
@@ -171,12 +184,20 @@ class ServerApp:  # pylint: disable=too-many-instance-attributes
                     >>> 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
@@ -212,10 +233,9 @@ class ServerApp:  # pylint: disable=too-many-instance-attributes
         """
 
         def lifespan_decorator(
-            lifespan_fn: Callable[[Context], Iterator[None]]
+            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]:

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/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):

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

@@ -130,9 +130,7 @@ 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:

flwr/server/superlink/linkstate/in_memory_linkstate.py

@@ -27,16 +27,18 @@ from flwr.common import Context, Message, log, now
 from flwr.common.constant import (
     MESSAGE_TTL_TOLERANCE,
     NODE_ID_NUM_BYTES,
+    PING_PATIENCE,
     RUN_ID_NUM_BYTES,
     SUPERLINK_NODE_ID,
     Status,
 )
-from flwr.common.record import ConfigsRecord
+from flwr.common.record import ConfigRecord
 from flwr.common.typing import Run, RunStatus, UserConfig
 from flwr.server.superlink.linkstate.linkstate import LinkState
 from flwr.server.utils import validate_message
 
 from .utils import (
+    check_node_availability_for_in_message,
     generate_rand_int_from_bytes,
     has_valid_sub_status,
     is_valid_transition,
@@ -67,7 +69,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.federation_options: dict[int, ConfigsRecord] = {}
+        self.federation_options: dict[int, ConfigRecord] = {}
         self.message_ins_store: dict[UUID, Message] = {}
         self.message_res_store: dict[UUID, Message] = {}
         self.message_ins_id_to_message_res_id: dict[UUID, UUID] = {}
@@ -156,7 +158,7 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
         res_metadata = message.metadata
         with self.lock:
             # Check if the Message it is replying to exists and is valid
-            msg_ins_id = res_metadata.reply_to_message
+            msg_ins_id = res_metadata.reply_to_message_id
             msg_ins = self.message_ins_store.get(UUID(msg_ins_id))
 
             # Ensure that dst_node_id of original Message matches the src_node_id of
@@ -232,13 +234,28 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
         with self.lock:
             current = time.time()
 
-            # Verify Messge IDs
+            # Verify Message IDs
             ret = verify_message_ids(
                 inquired_message_ids=message_ids,
                 found_message_ins_dict=self.message_ins_store,
                 current_time=current,
             )
 
+            # Check node availability
+            dst_node_ids = {
+                self.message_ins_store[message_id].metadata.dst_node_id
+                for message_id in message_ids
+            }
+            tmp_ret_dict = check_node_availability_for_in_message(
+                inquired_in_message_ids=message_ids,
+                found_in_message_dict=self.message_ins_store,
+                node_id_to_online_until={
+                    node_id: self.node_ids[node_id][0] for node_id in dst_node_ids
+                },
+                current_time=current,
+            )
+            ret.update(tmp_ret_dict)
+
             # Find all reply Messages
             message_res_found: list[Message] = []
             for message_id in message_ids:
@@ -317,6 +334,7 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
                 log(ERROR, "Unexpected node registration failure.")
                 return 0
 
+            # Mark the node online util time.time() + ping_interval
             self.node_ids[node_id] = (time.time() + ping_interval, ping_interval)
             return node_id
 
@@ -381,7 +399,7 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
         fab_version: Optional[str],
         fab_hash: Optional[str],
         override_config: UserConfig,
-        federation_options: ConfigsRecord,
+        federation_options: ConfigRecord,
     ) -> int:
         """Create a new run for the specified `fab_hash`."""
         # Sample a random int64 as run_id
@@ -510,7 +528,7 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
 
         return pending_run_id
 
-    def get_federation_options(self, run_id: int) -> Optional[ConfigsRecord]:
+    def get_federation_options(self, run_id: int) -> Optional[ConfigRecord]:
         """Retrieve the federation options for the specified `run_id`."""
         with self.lock:
             if run_id not in self.run_ids:
@@ -519,10 +537,17 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
             return self.federation_options[run_id]
 
     def acknowledge_ping(self, node_id: int, ping_interval: float) -> bool:
-        """Acknowledge a ping received from a node, serving as a heartbeat."""
+        """Acknowledge a ping received from a node, serving as a heartbeat.
+
+        It allows for one missed ping (in a PING_PATIENCE * ping_interval) before
+        marking the node as offline, where PING_PATIENCE = 2 in default.
+        """
         with self.lock:
             if node_id in self.node_ids:
-                self.node_ids[node_id] = (time.time() + ping_interval, ping_interval)
+                self.node_ids[node_id] = (
+                    time.time() + PING_PATIENCE * ping_interval,
+                    ping_interval,
+                )
                 return True
         return False
 

flwr/server/superlink/linkstate/linkstate.py

@@ -20,7 +20,7 @@ from typing import Optional
 from uuid import UUID
 
 from flwr.common import Context, Message
-from flwr.common.record import ConfigsRecord
+from flwr.common.record import ConfigRecord
 from flwr.common.typing import Run, RunStatus, UserConfig
 
 
@@ -164,7 +164,7 @@ class LinkState(abc.ABC):  # pylint: disable=R0904
         fab_version: Optional[str],
         fab_hash: Optional[str],
         override_config: UserConfig,
-        federation_options: ConfigsRecord,
+        federation_options: ConfigRecord,
     ) -> int:
         """Create a new run for the specified `fab_hash`."""
 
@@ -236,7 +236,7 @@ class LinkState(abc.ABC):  # pylint: disable=R0904
         """
 
     @abc.abstractmethod
-    def get_federation_options(self, run_id: int) -> Optional[ConfigsRecord]:
+    def get_federation_options(self, run_id: int) -> Optional[ConfigRecord]:
         """Retrieve the federation options for the specified `run_id`.
 
         Parameters
@@ -246,7 +246,7 @@ class LinkState(abc.ABC):  # pylint: disable=R0904
 
         Returns
         -------
-        Optional[ConfigsRecord]
+        Optional[ConfigRecord]
             The federation options for the run if it exists; None otherwise.
         """