flwr-nightly 1.13.0.dev20241021__py3-none-any.whl → 1.13.0.dev20241111__py3-none-any.whl

flwr/server/superlink/{state/state.py → linkstate/linkstate.py}

@@ -12,28 +12,30 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""Abstract base class State."""
+"""Abstract base class LinkState."""
 
 
 import abc
 from typing import Optional
 from uuid import UUID
 
-from flwr.common.typing import Run, UserConfig
+from flwr.common import Context
+from flwr.common.record import ConfigsRecord
+from flwr.common.typing import Run, RunStatus, UserConfig
 from flwr.proto.task_pb2 import TaskIns, TaskRes  # pylint: disable=E0611
 
 
-class State(abc.ABC):  # pylint: disable=R0904
-    """Abstract State."""
+class LinkState(abc.ABC):  # pylint: disable=R0904
+    """Abstract LinkState."""
 
     @abc.abstractmethod
     def store_task_ins(self, task_ins: TaskIns) -> Optional[UUID]:
         """Store one TaskIns.
 
-        Usually, the Driver API calls this to schedule instructions.
+        Usually, the ServerAppIo API calls this to schedule instructions.
 
-        Stores the value of the `task_ins` in the state and, if successful, returns the
-        `task_id` (UUID) of the `task_ins`. If, for any reason,
+        Stores the value of the `task_ins` in the link state and, if successful,
+        returns the `task_id` (UUID) of the `task_ins`. If, for any reason,
         storing the `task_ins` fails, `None` is returned.
 
         Constraints
@@ -101,8 +103,8 @@ class State(abc.ABC):  # pylint: disable=R0904
     def get_task_res(self, task_ids: set[UUID]) -> list[TaskRes]:
         """Get TaskRes for task_ids.
 
-        Usually, the Driver API calls this method to get results for instructions it has
-        previously scheduled.
+        Usually, the ServerAppIo API calls this method to get results for instructions
+        it has previously scheduled.
 
         Retrieves all TaskRes for the given `task_ids` and returns and empty list of
         none could be found.
@@ -130,11 +132,11 @@ class State(abc.ABC):  # pylint: disable=R0904
     def create_node(
         self, ping_interval: float, public_key: Optional[bytes] = None
     ) -> int:
-        """Create, store in state, and return `node_id`."""
+        """Create, store in the link state, and return `node_id`."""
 
     @abc.abstractmethod
     def delete_node(self, node_id: int, public_key: Optional[bytes] = None) -> None:
-        """Remove `node_id` from state."""
+        """Remove `node_id` from the link state."""
 
     @abc.abstractmethod
     def get_nodes(self, run_id: int) -> set[int]:
@@ -151,15 +153,20 @@ class State(abc.ABC):  # pylint: disable=R0904
         """Retrieve stored `node_id` filtered by `node_public_keys`."""
 
     @abc.abstractmethod
-    def create_run(
+    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],
         override_config: UserConfig,
+        federation_options: ConfigsRecord,
     ) -> int:
         """Create a new run for the specified `fab_hash`."""
 
+    @abc.abstractmethod
+    def get_run_ids(self) -> set[int]:
+        """Retrieve all run IDs."""
+
     @abc.abstractmethod
     def get_run(self, run_id: int) -> Optional[Run]:
         """Retrieve information about the run with the specified `run_id`.
@@ -172,17 +179,77 @@ class State(abc.ABC):  # pylint: disable=R0904
         Returns
         -------
         Optional[Run]
-            A dataclass instance containing three elements if `run_id` is valid:
-            - `run_id`: The identifier of the run, same as the specified `run_id`.
-            - `fab_id`: The identifier of the FAB used in the specified run.
-            - `fab_version`: The version of the FAB used in the specified run.
+            The `Run` instance if found; otherwise, `None`.
+        """
+
+    @abc.abstractmethod
+    def get_run_status(self, run_ids: set[int]) -> dict[int, RunStatus]:
+        """Retrieve the statuses for the specified runs.
+
+        Parameters
+        ----------
+        run_ids : set[int]
+            A set of run identifiers for which to retrieve statuses.
+
+        Returns
+        -------
+        dict[int, RunStatus]
+            A dictionary mapping each valid run ID to its corresponding status.
+
+        Notes
+        -----
+        Only valid run IDs that exist in the State will be included in the returned
+        dictionary. If a run ID is not found, it will be omitted from the result.
+        """
+
+    @abc.abstractmethod
+    def update_run_status(self, run_id: int, new_status: RunStatus) -> bool:
+        """Update the status of the run with the specified `run_id`.
+
+        Parameters
+        ----------
+        run_id : int
+            The identifier of the run.
+        new_status : RunStatus
+            The new status to be assigned to the run.
+
+        Returns
+        -------
+        bool
+            True if the status update is successful; False otherwise.
+        """
+
+    @abc.abstractmethod
+    def get_pending_run_id(self) -> Optional[int]:
+        """Get the `run_id` of a run with `Status.PENDING` status.
+
+        Returns
+        -------
+        Optional[int]
+            The `run_id` of a `Run` that is pending to be started; None if
+            there is no Run pending.
+        """
+
+    @abc.abstractmethod
+    def get_federation_options(self, run_id: int) -> Optional[ConfigsRecord]:
+        """Retrieve the federation options for the specified `run_id`.
+
+        Parameters
+        ----------
+        run_id : int
+            The identifier of the run.
+
+        Returns
+        -------
+        Optional[ConfigsRecord]
+            The federation options for the run if it exists; None otherwise.
         """
 
     @abc.abstractmethod
     def store_server_private_public_key(
         self, private_key: bytes, public_key: bytes
     ) -> None:
-        """Store `server_private_key` and `server_public_key` in state."""
+        """Store `server_private_key` and `server_public_key` in the link state."""
 
     @abc.abstractmethod
     def get_server_private_key(self) -> Optional[bytes]:
@@ -194,11 +261,11 @@ class State(abc.ABC):  # pylint: disable=R0904
 
     @abc.abstractmethod
     def store_node_public_keys(self, public_keys: set[bytes]) -> None:
-        """Store a set of `node_public_keys` in state."""
+        """Store a set of `node_public_keys` in the link state."""
 
     @abc.abstractmethod
     def store_node_public_key(self, public_key: bytes) -> None:
-        """Store a `node_public_key` in state."""
+        """Store a `node_public_key` in the link state."""
 
     @abc.abstractmethod
     def get_node_public_keys(self) -> set[bytes]:
@@ -222,3 +289,66 @@ class State(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`.
+        """
+
+    @abc.abstractmethod
+    def add_serverapp_log(self, run_id: int, log_message: str) -> None:
+        """Add a log entry to the ServerApp logs for the specified `run_id`.
+
+        Parameters
+        ----------
+        run_id : int
+            The identifier of the run for which to add a log entry.
+        log_message : str
+            The log entry to be added to the ServerApp logs.
+        """
+
+    @abc.abstractmethod
+    def get_serverapp_log(
+        self, run_id: int, after_timestamp: Optional[float]
+    ) -> tuple[str, float]:
+        """Get the ServerApp logs for the specified `run_id`.
+
+        Parameters
+        ----------
+        run_id : int
+            The identifier of the run for which to retrieve the ServerApp logs.
+
+        after_timestamp : Optional[float]
+            Retrieve logs after this timestamp. If set to `None`, retrieve all logs.
+
+        Returns
+        -------
+        tuple[str, float]
+            A tuple containing:
+            - The ServerApp logs associated with the specified `run_id`.
+            - The timestamp of the latest log entry in the returned logs.
+              Returns `0` if no logs are returned.
+        """

flwr/server/superlink/{state/state_factory.py → linkstate/linkstate_factory.py}

@@ -20,13 +20,13 @@ from typing import Optional
 
 from flwr.common.logger import log
 
-from .in_memory_state import InMemoryState
-from .sqlite_state import SqliteState
-from .state import State
+from .in_memory_linkstate import InMemoryLinkState
+from .linkstate import LinkState
+from .sqlite_linkstate import SqliteLinkState
 
 
-class StateFactory:
-    """Factory class that creates State instances.
+class LinkStateFactory:
+    """Factory class that creates LinkState instances.
 
     Parameters
     ----------
@@ -39,19 +39,19 @@ class StateFactory:
 
     def __init__(self, database: str) -> None:
         self.database = database
-        self.state_instance: Optional[State] = None
+        self.state_instance: Optional[LinkState] = None
 
-    def state(self) -> State:
+    def state(self) -> LinkState:
         """Return a State instance and create it, if necessary."""
         # InMemoryState
         if self.database == ":flwr-in-memory-state:":
             if self.state_instance is None:
-                self.state_instance = InMemoryState()
+                self.state_instance = InMemoryLinkState()
             log(DEBUG, "Using InMemoryState")
             return self.state_instance
 
         # SqliteState
-        state = SqliteState(self.database)
+        state = SqliteLinkState(self.database)
         state.initialize()
         log(DEBUG, "Using SqliteState")
         return state