flwr 1.23.0__py3-none-any.whl → 1.25.0__py3-none-any.whl

flwr/server/superlink/linkstate/linkstate.py

@@ -17,20 +17,25 @@
 
 import abc
 from collections.abc import Sequence
-from typing import Optional
 
 from flwr.common import Context, Message
 from flwr.common.record import ConfigRecord
 from flwr.common.typing import Run, RunStatus, UserConfig
 from flwr.proto.node_pb2 import NodeInfo  # pylint: disable=E0611
 from flwr.supercore.corestate import CoreState
+from flwr.superlink.federation import FederationManager
 
 
 class LinkState(CoreState):  # pylint: disable=R0904
     """Abstract LinkState."""
 
+    @property
     @abc.abstractmethod
-    def store_message_ins(self, message: Message) -> Optional[str]:
+    def federation_manager(self) -> FederationManager:
+        """Return the FederationManager instance."""
+
+    @abc.abstractmethod
+    def store_message_ins(self, message: Message) -> str | None:
         """Store one Message.
 
         Usually, the ServerAppIo API calls this to schedule instructions.
@@ -48,7 +53,7 @@ class LinkState(CoreState):  # pylint: disable=R0904
         """
 
     @abc.abstractmethod
-    def get_message_ins(self, node_id: int, limit: Optional[int]) -> list[Message]:
+    def get_message_ins(self, node_id: int, limit: int | None) -> list[Message]:
         """Get zero or more `Message` objects for the provided `node_id`.
 
         Usually, the Fleet API calls this for Nodes planning to work on one or more
@@ -63,7 +68,7 @@ class LinkState(CoreState):  # pylint: disable=R0904
         """
 
     @abc.abstractmethod
-    def store_message_res(self, message: Message) -> Optional[str]:
+    def store_message_res(self, message: Message) -> str | None:
         """Store one Message.
 
         Usually, the Fleet API calls this for Nodes returning results.
@@ -131,7 +136,11 @@ class LinkState(CoreState):  # pylint: disable=R0904
 
     @abc.abstractmethod
     def create_node(
-        self, owner_aid: str, public_key: bytes, heartbeat_interval: float
+        self,
+        owner_aid: str,
+        owner_name: str,
+        public_key: bytes,
+        heartbeat_interval: float,
     ) -> int:
         """Create, store in the link state, and return `node_id`."""
 
@@ -189,7 +198,7 @@ class LinkState(CoreState):  # pylint: disable=R0904
         """
 
     @abc.abstractmethod
-    def get_node_id_by_public_key(self, public_key: bytes) -> Optional[int]:
+    def get_node_id_by_public_key(self, public_key: bytes) -> int | None:
         """Get `node_id` for the specified `public_key` if it exists and is not deleted.
 
         Parameters
@@ -208,9 +217,9 @@ class LinkState(CoreState):  # pylint: disable=R0904
     def get_node_info(
         self,
         *,
-        node_ids: Optional[Sequence[int]] = None,
-        owner_aids: Optional[Sequence[str]] = None,
-        statuses: Optional[Sequence[str]] = None,
+        node_ids: Sequence[int] | None = None,
+        owner_aids: Sequence[str] | None = None,
+        statuses: Sequence[str] | None = None,
     ) -> Sequence[NodeInfo]:
         """Retrieve information about nodes based on the specified filters.
 
@@ -259,24 +268,54 @@ class LinkState(CoreState):  # pylint: disable=R0904
     @abc.abstractmethod
     def create_run(  # pylint: disable=too-many-arguments,too-many-positional-arguments
         self,
-        fab_id: Optional[str],
-        fab_version: Optional[str],
-        fab_hash: Optional[str],
+        fab_id: str | None,
+        fab_version: str | None,
+        fab_hash: str | None,
         override_config: UserConfig,
+        federation: str,
         federation_options: ConfigRecord,
-        flwr_aid: Optional[str],
+        flwr_aid: str | None,
     ) -> int:
-        """Create a new run for the specified `fab_hash`."""
+        """Create a new run.
+
+        Parameters
+        ----------
+        fab_id : Optional[str]
+            The ID of the FAB, of format `<publisher>/<app-name>`.
+        fab_version : Optional[str]
+            The version of the FAB.
+        fab_hash : Optional[str]
+            The SHA256 hex hash of the FAB.
+        override_config : UserConfig
+            Configuration overrides for the run config.
+        federation : str
+            The federation this run belongs to.
+        federation_options : ConfigRecord
+            Federation configurations. For now, only `num-supernodes` for
+            the simulation runtime.
+        flwr_aid : Optional[str]
+            Flower Account ID of the creator.
+
+        Returns
+        -------
+        int
+            The run ID of the newly created run.
+
+        Notes
+        -----
+        This method will not verify if the account has permission to create
+        a run in the federation.
+        """
 
     @abc.abstractmethod
-    def get_run_ids(self, flwr_aid: Optional[str]) -> set[int]:
+    def get_run_ids(self, flwr_aid: str | None) -> set[int]:
         """Retrieve all run IDs if `flwr_aid` is not specified.
 
         Otherwise, retrieve all run IDs for the specified `flwr_aid`.
         """
 
     @abc.abstractmethod
-    def get_run(self, run_id: int) -> Optional[Run]:
+    def get_run(self, run_id: int) -> Run | None:
         """Retrieve information about the run with the specified `run_id`.
 
         Parameters
@@ -328,7 +367,7 @@ class LinkState(CoreState):  # pylint: disable=R0904
         """
 
     @abc.abstractmethod
-    def get_pending_run_id(self) -> Optional[int]:
+    def get_pending_run_id(self) -> int | None:
         """Get the `run_id` of a run with `Status.PENDING` status.
 
         Returns
@@ -339,7 +378,7 @@ class LinkState(CoreState):  # pylint: disable=R0904
         """
 
     @abc.abstractmethod
-    def get_federation_options(self, run_id: int) -> Optional[ConfigRecord]:
+    def get_federation_options(self, run_id: int) -> ConfigRecord | None:
         """Retrieve the federation options for the specified `run_id`.
 
         Parameters
@@ -380,30 +419,7 @@ class LinkState(CoreState):  # pylint: disable=R0904
         """
 
     @abc.abstractmethod
-    def acknowledge_app_heartbeat(self, run_id: int, heartbeat_interval: float) -> bool:
-        """Acknowledge a heartbeat received from a ServerApp for a given run.
-
-        A run with status `"running"` is considered alive as long as it sends heartbeats
-        within the tolerated interval: HEARTBEAT_PATIENCE × heartbeat_interval.
-        HEARTBEAT_PATIENCE = N allows for N-1 missed heartbeat before the run is
-        marked as `"completed:failed"`.
-
-        Parameters
-        ----------
-        run_id : int
-            The `run_id` from which the heartbeat was received.
-        heartbeat_interval : float
-            The interval (in seconds) from the current timestamp within which the next
-            heartbeat from the ServerApp for this run must be received.
-
-        Returns
-        -------
-        is_acknowledged : bool
-            True if the heartbeat is successfully acknowledged; otherwise, False.
-        """
-
-    @abc.abstractmethod
-    def get_serverapp_context(self, run_id: int) -> Optional[Context]:
+    def get_serverapp_context(self, run_id: int) -> Context | None:
         """Get the context for the specified `run_id`.
 
         Parameters
@@ -444,7 +460,7 @@ class LinkState(CoreState):  # pylint: disable=R0904
 
     @abc.abstractmethod
     def get_serverapp_log(
-        self, run_id: int, after_timestamp: Optional[float]
+        self, run_id: int, after_timestamp: float | None
     ) -> tuple[str, float]:
         """Get the ServerApp logs for the specified `run_id`.
 
@@ -464,3 +480,35 @@ class LinkState(CoreState):  # pylint: disable=R0904
             - The timestamp of the latest log entry in the returned logs.
               Returns `0` if no logs are returned.
         """
+
+    @abc.abstractmethod
+    def store_traffic(self, run_id: int, *, bytes_sent: int, bytes_recv: int) -> None:
+        """Store traffic data for the specified `run_id`.
+
+        Parameters
+        ----------
+        run_id : int
+            The identifier of the run for which to store traffic data.
+        bytes_sent : int
+            The number of bytes pulled by SuperNodes from the SuperLink to add to the
+            run's total.
+        bytes_recv : int
+            The number of bytes received by SuperLink from SuperNodes to add to the
+            run's total.
+        """
+
+    @abc.abstractmethod
+    def add_clientapp_runtime(self, run_id: int, runtime: float) -> None:
+        """Add ClientApp runtime to the cumulative total for the specified `run_id`.
+
+        This method accumulates the runtime by adding the provided value to the
+        existing total runtime for the run. Multiple ClientApps can contribute
+        to the same run's total runtime.
+
+        Parameters
+        ----------
+        run_id : int
+            The identifier of the run for which to store each ClientApp's runtime.
+        runtime : float
+            The runtime in seconds to add to the `run_id`'s cumulative total.
+        """

flwr/server/superlink/linkstate/linkstate_factory.py

@@ -16,10 +16,11 @@
 
 
 from logging import DEBUG
-from typing import Optional
 
 from flwr.common.logger import log
 from flwr.supercore.constant import FLWR_IN_MEMORY_DB_NAME
+from flwr.supercore.object_store import ObjectStoreFactory
+from flwr.superlink.federation import FederationManager
 
 from .in_memory_linkstate import InMemoryLinkState
 from .linkstate import LinkState
@@ -36,23 +37,39 @@ class LinkStateFactory:
         Note that passing ':memory:' will open a connection to a database that is
         in RAM, instead of on disk. For more information on special in-memory
         databases, please refer to https://sqlite.org/inmemorydb.html.
+    federation_manager : FederationManager
+        An instance of FederationManager to manage federations.
+    objectstore_factory : ObjectStoreFactory
+        An instance of ObjectStoreFactory to create object stores.
     """
 
-    def __init__(self, database: str) -> None:
+    def __init__(
+        self,
+        database: str,
+        federation_manager: FederationManager,
+        objectstore_factory: ObjectStoreFactory,
+    ) -> None:
         self.database = database
-        self.state_instance: Optional[LinkState] = None
+        self.state_instance: LinkState | None = None
+        self.federation_manager = federation_manager
+        self.objectstore_factory = objectstore_factory
 
     def state(self) -> LinkState:
         """Return a State instance and create it, if necessary."""
+        # Get the ObjectStore instance
+        object_store = self.objectstore_factory.store()
+
         # InMemoryState
         if self.database == FLWR_IN_MEMORY_DB_NAME:
             if self.state_instance is None:
-                self.state_instance = InMemoryLinkState()
+                self.state_instance = InMemoryLinkState(
+                    self.federation_manager, object_store
+                )
             log(DEBUG, "Using InMemoryState")
             return self.state_instance
 
         # SqliteState
-        state = SqliteLinkState(self.database)
+        state = SqliteLinkState(self.database, self.federation_manager, object_store)
         state.initialize()
         log(DEBUG, "Using SqliteState")
         return state