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

flwr/{common → supercore}/heartbeat.py

@@ -16,24 +16,27 @@
 
 
 import random
+import signal
 import threading
-from typing import Callable, Union
+from collections.abc import Callable
 
 import grpc
 
+from flwr.common.constant import (
+    HEARTBEAT_BASE_MULTIPLIER,
+    HEARTBEAT_CALL_TIMEOUT,
+    HEARTBEAT_DEFAULT_INTERVAL,
+    HEARTBEAT_RANDOM_RANGE,
+)
+from flwr.common.retry_invoker import RetryInvoker, exponential
+from flwr.proto.clientappio_pb2_grpc import ClientAppIoStub
+
 # pylint: disable=E0611
 from flwr.proto.heartbeat_pb2 import SendAppHeartbeatRequest
 from flwr.proto.serverappio_pb2_grpc import ServerAppIoStub
 from flwr.proto.simulationio_pb2_grpc import SimulationIoStub
 
 # pylint: enable=E0611
-from .constant import (
-    HEARTBEAT_BASE_MULTIPLIER,
-    HEARTBEAT_CALL_TIMEOUT,
-    HEARTBEAT_DEFAULT_INTERVAL,
-    HEARTBEAT_RANDOM_RANGE,
-)
-from .retry_invoker import RetryInvoker, exponential
 
 
 class HeartbeatFailure(Exception):
@@ -116,24 +119,18 @@ class HeartbeatSender:
                 raise HeartbeatFailure
 
 
-def get_grpc_app_heartbeat_fn(
-    stub: Union[ServerAppIoStub, SimulationIoStub],
-    run_id: int,
-    *,
-    failure_message: str,
+def make_app_heartbeat_fn_grpc(
+    stub: ServerAppIoStub | SimulationIoStub | ClientAppIoStub,
+    token: str,
 ) -> Callable[[], bool]:
-    """Get the function to send a heartbeat to gRPC endpoint.
-
-    This function is for app heartbeats only. It is not used for node heartbeats.
+    """Get the function to send a heartbeat to gRPC endpoint from an app process.
 
     Parameters
     ----------
     stub : Union[ServerAppIoStub, SimulationIoStub]
         gRPC stub to send the heartbeat.
-    run_id : int
-        The run ID to use in the heartbeat request.
-    failure_message : str
-        Error message to raise if the heartbeat fails.
+    token : str
+        The token to use in the heartbeat request.
 
     Returns
     -------
@@ -141,9 +138,7 @@ def get_grpc_app_heartbeat_fn(
         Function that sends a heartbeat to the gRPC endpoint.
     """
     # Construct the heartbeat request
-    req = SendAppHeartbeatRequest(
-        run_id=run_id, heartbeat_interval=HEARTBEAT_DEFAULT_INTERVAL
-    )
+    req = SendAppHeartbeatRequest(token=token)
 
     def fn() -> bool:
         # Call ServerAppIo API
@@ -157,9 +152,9 @@ def get_grpc_app_heartbeat_fn(
                 return False
             raise
 
-        # Check if not successful
+        # Raise SIGINT to trigger graceful shutdown if heartbeat failed
         if not res.success:
-            raise RuntimeError(failure_message)
+            signal.raise_signal(signal.SIGINT)
         return True
 
     return fn

flwr/supercore/object_store/in_memory_object_store.py

@@ -17,7 +17,6 @@
 
 import threading
 from dataclasses import dataclass
-from typing import Optional
 
 from flwr.common.inflatable import (
     get_object_id,
@@ -154,7 +153,7 @@ class InMemoryObjectStore(ObjectStore):
             self.store[object_id].content = object_content
             self.store[object_id].is_available = True
 
-    def get(self, object_id: str) -> Optional[bytes]:
+    def get(self, object_id: str) -> bytes | None:
         """Get an object from the store."""
         with self.lock_store:
             # Check if the object ID is pre-registered

flwr/supercore/object_store/object_store.py

@@ -16,7 +16,6 @@
 
 
 import abc
-from typing import Optional
 
 from flwr.proto.message_pb2 import ObjectTree  # pylint: disable=E0611
 
@@ -89,7 +88,7 @@ class ObjectStore(abc.ABC):
         """
 
     @abc.abstractmethod
-    def get(self, object_id: str) -> Optional[bytes]:
+    def get(self, object_id: str) -> bytes | None:
         """Get an object from the store.
 
         Parameters

flwr/supercore/object_store/object_store_factory.py

@@ -16,7 +16,6 @@
 
 
 from logging import DEBUG
-from typing import Optional
 
 from flwr.common.logger import log
 from flwr.supercore.constant import FLWR_IN_MEMORY_DB_NAME
@@ -40,7 +39,7 @@ class ObjectStoreFactory:
 
     def __init__(self, database: str = FLWR_IN_MEMORY_DB_NAME) -> None:
         self.database = database
-        self.store_instance: Optional[ObjectStore] = None
+        self.store_instance: ObjectStore | None = None
 
     def store(self) -> ObjectStore:
         """Return an ObjectStore instance and create it, if necessary.

flwr/supercore/object_store/sqlite_object_store.py

@@ -15,7 +15,7 @@
 """Flower SQLite ObjectStore implementation."""
 
 
-from typing import Optional, cast
+from typing import cast
 
 from flwr.common.inflatable import (
     get_object_id,
@@ -63,13 +63,12 @@ class SqliteObjectStore(ObjectStore, SqliteMixin):
         super().__init__(database_path)
         self.verify = verify
 
-    def initialize(self, log_queries: bool = False) -> list[tuple[str]]:
-        """Connect to the DB, enable FK support, and create tables if needed."""
-        return self._ensure_initialized(
+    def get_sql_statements(self) -> tuple[str, ...]:
+        """Return SQL statements for ObjectStore tables."""
+        return (
             SQL_CREATE_OBJECTS,
             SQL_CREATE_OBJECT_CHILDREN,
             SQL_CREATE_RUN_OBJECTS,
-            log_queries=log_queries,
         )
 
     def preregister(self, run_id: int, object_tree: ObjectTree) -> list[str]:
@@ -126,7 +125,9 @@ class SqliteObjectStore(ObjectStore, SqliteMixin):
                 "SELECT object_id FROM objects WHERE object_id=?", (object_id,)
             ).fetchone()
             if not row:
-                raise NoObjectInStoreError(f"Object {object_id} not found.")
+                raise NoObjectInStoreError(
+                    f"Object {object_id} was not pre-registered."
+                )
             children = self.query(
                 "SELECT child_id FROM object_children WHERE parent_id=?", (object_id,)
             )
@@ -176,7 +177,7 @@ class SqliteObjectStore(ObjectStore, SqliteMixin):
                 (object_content, object_id),
             )
 
-    def get(self, object_id: str) -> Optional[bytes]:
+    def get(self, object_id: str) -> bytes | None:
         """Get an object from the store."""
         rows = self.query("SELECT content FROM objects WHERE object_id=?", (object_id,))
         return rows[0]["content"] if rows else None

flwr/supercore/primitives/asymmetric.py

@@ -113,5 +113,5 @@ def uses_nist_ec_curve(public_key: ec.EllipticCurvePublicKey) -> bool:
     """Return True if the provided key uses a NIST EC curve."""
     return isinstance(
         public_key.curve,
-        (ec.SECP192R1, ec.SECP224R1, ec.SECP256R1, ec.SECP384R1, ec.SECP521R1),
+        (ec.SECP192R1 | ec.SECP224R1 | ec.SECP256R1 | ec.SECP384R1 | ec.SECP521R1),
     )

flwr/supercore/primitives/asymmetric_ed25519.py

@@ -15,6 +15,7 @@
 """Ed25519-only asymmetric cryptography utilities."""
 
 import base64
+from pathlib import Path
 
 from cryptography.exceptions import InvalidSignature
 from cryptography.hazmat.primitives import serialization
@@ -150,7 +151,7 @@ def verify_signature(
         return False
 
 
-def create_signed_message(fab_digest: bytes, timestamp: int) -> bytes:
+def create_message_to_sign(fab_digest: bytes, timestamp: int) -> bytes:
     """Create a canonical message:
     timestamp (8 bytes big-endian) + fab_digest.
     """
@@ -163,3 +164,12 @@ def decode_base64url(sig: str) -> bytes:
     # add missing padding (=) to a multiple of 4
     pad = (-len(sig)) % 4
     return base64.urlsafe_b64decode(sig + ("=" * pad))
+
+
+def load_private_key(path: Path) -> ed25519.Ed25519PrivateKey:
+    """Load an SSH-format private key (Ed25519) using cryptography."""
+    key_bytes = path.read_bytes()
+    private_key = serialization.load_ssh_private_key(key_bytes, password=None)
+    if not isinstance(private_key, ed25519.Ed25519PrivateKey):
+        raise ValueError("Private key is not Ed25519")
+    return private_key

flwr/supercore/sqlite_mixin.py

@@ -17,14 +17,14 @@
 
 import re
 import sqlite3
-from abc import ABC, abstractmethod
+from abc import ABC
 from collections.abc import Sequence
 from logging import DEBUG, ERROR
-from typing import Any, Optional, Union
+from typing import Any
 
 from flwr.common.logger import log
 
-DictOrTuple = Union[tuple[Any, ...], dict[str, Any]]
+DictOrTuple = tuple[Any, ...] | dict[str, Any]
 
 
 class SqliteMixin(ABC):
@@ -32,7 +32,7 @@ class SqliteMixin(ABC):
 
     def __init__(self, database_path: str) -> None:
         self.database_path = database_path
-        self._conn: Optional[sqlite3.Connection] = None
+        self._conn: sqlite3.Connection | None = None
 
     @property
     def conn(self) -> sqlite3.Connection:
@@ -41,10 +41,24 @@ class SqliteMixin(ABC):
             raise AttributeError("Database not initialized. Call initialize() first.")
         return self._conn
 
-    @abstractmethod
+    def get_sql_statements(self) -> tuple[str, ...]:
+        """Return SQL statements for this class.
+
+        Subclasses can override this to provide their SQL CREATE statements.
+        The base implementation returns an empty tuple.
+
+        Returns
+        -------
+        tuple[str, ...]
+            SQL CREATE TABLE/INDEX statements for this class.
+        """
+        return ()
+
     def initialize(self, log_queries: bool = False) -> list[tuple[str]]:
         """Connect to the DB, enable FK support, and create tables if needed.
 
+        This method executes SQL statements returned by `get_sql_statements()`.
+
         Parameters
         ----------
         log_queries : bool
@@ -57,45 +71,34 @@ class SqliteMixin(ABC):
 
         Examples
         --------
-        Implement in subclass:
+        Override `get_sql_statements()` in your subclass:
 
         .. code:: python
 
-            def initialize(self, log_queries: bool = False) -> list[tuple[str]]:
-                return self._ensure_initialized(
+            def get_sql_statements(self) -> tuple[str, ...]:
+                return (
                     SQL_CREATE_TABLE_FOO,
                     SQL_CREATE_TABLE_BAR,
-                    log_queries=log_queries
                 )
-        """
-
-    def _ensure_initialized(
-        self,
-        *create_statements: str,
-        log_queries: bool = False,
-    ) -> list[tuple[str]]:
-        """Connect to the DB, enable FK support, and create tables if needed.
 
-        Subclasses should call this with their own CREATE TABLE/INDEX statements in
-        their `.initialize()` methods.
+        To include parent SQL statements, call super():
 
-        Parameters
-        ----------
-        create_statements : str
-            SQL statements to create tables and indexes.
-        log_queries : bool
-            Log each query which is executed.
+        .. code:: python
 
-        Returns
-        -------
-        list[tuple[str]]
-            The list of all tables in the DB.
+            def get_sql_statements(self) -> tuple[str, ...]:
+                return super().get_sql_statements() + (
+                    SQL_CREATE_TABLE_FOO,
+                    SQL_CREATE_TABLE_BAR,
+                )
         """
         self._conn = sqlite3.connect(self.database_path)
         # Enable Write-Ahead Logging (WAL) for better concurrency
         self._conn.execute("PRAGMA journal_mode = WAL;")
         self._conn.execute("PRAGMA synchronous = NORMAL;")
         self._conn.execute("PRAGMA foreign_keys = ON;")
+        self._conn.execute("PRAGMA cache_size = -64000;")  # 64MB cache
+        self._conn.execute("PRAGMA temp_store = MEMORY;")  # In-memory temp tables
+        self._conn.execute("PRAGMA mmap_size = 268435456;")  # 256MB memory-mapped I/O
         self._conn.row_factory = dict_factory
 
         if log_queries:
@@ -103,7 +106,7 @@ class SqliteMixin(ABC):
 
         # Create tables and indexes
         cur = self._conn.cursor()
-        for sql in create_statements:
+        for sql in self.get_sql_statements():
             cur.execute(sql)
         res = cur.execute("SELECT name FROM sqlite_schema;")
         return res.fetchall()
@@ -111,7 +114,7 @@ class SqliteMixin(ABC):
     def query(
         self,
         query: str,
-        data: Optional[Union[Sequence[DictOrTuple], DictOrTuple]] = None,
+        data: Sequence[DictOrTuple] | DictOrTuple | None = None,
     ) -> list[dict[str, Any]]:
         """Execute a SQL query and return the results as list of dicts."""
         if self._conn is None:
@@ -127,8 +130,8 @@ class SqliteMixin(ABC):
             with self._conn:
                 if (
                     len(data) > 0
-                    and isinstance(data, (tuple, list))
-                    and isinstance(data[0], (tuple, dict))
+                    and isinstance(data, (tuple | list))
+                    and isinstance(data[0], (tuple | dict))
                 ):
                     rows = self._conn.executemany(query, data)
                 else:
@@ -153,4 +156,4 @@ def dict_factory(
     Less efficent for retrival of large amounts of data but easier to use.
     """
     fields = [column[0] for column in cursor.description]
-    return dict(zip(fields, row))
+    return dict(zip(fields, row, strict=True))

flwr/supercore/superexec/plugin/base_exec_plugin.py

@@ -18,7 +18,6 @@
 import os
 import subprocess
 from collections.abc import Sequence
-from typing import Optional
 
 from .exec_plugin import ExecPlugin
 
@@ -33,7 +32,7 @@ class BaseExecPlugin(ExecPlugin):
     command = ""
     appio_api_address_arg = ""
 
-    def select_run_id(self, candidate_run_ids: Sequence[int]) -> Optional[int]:
+    def select_run_id(self, candidate_run_ids: Sequence[int]) -> int | None:
         """Select a run ID to execute from a sequence of candidates."""
         if not candidate_run_ids:
             return None

flwr/supercore/superexec/plugin/exec_plugin.py

@@ -16,8 +16,8 @@
 
 
 from abc import ABC, abstractmethod
-from collections.abc import Sequence
-from typing import Any, Callable, Optional
+from collections.abc import Callable, Sequence
+from typing import Any
 
 from flwr.common.typing import Run
 
@@ -36,7 +36,7 @@ class ExecPlugin(ABC):
         self.get_run = get_run
 
     @abstractmethod
-    def select_run_id(self, candidate_run_ids: Sequence[int]) -> Optional[int]:
+    def select_run_id(self, candidate_run_ids: Sequence[int]) -> int | None:
         """Select a run ID to execute from a sequence of candidates.
 
         A candidate run ID is one that has at least one pending message and is

flwr/supercore/superexec/run_superexec.py

@@ -17,7 +17,7 @@
 
 import time
 from logging import WARN
-from typing import Any, Optional, Union
+from typing import Any
 
 from flwr.common.config import get_flwr_dir
 from flwr.common.exit import ExitCode, flwr_exit, register_signal_handlers
@@ -43,14 +43,12 @@ from .plugin import ExecPlugin
 
 def run_superexec(  # pylint: disable=R0913,R0914,R0917
     plugin_class: type[ExecPlugin],
-    stub_class: Union[
-        type[ClientAppIoStub], type[ServerAppIoStub], type[SimulationIoStub]
-    ],
+    stub_class: type[ClientAppIoStub] | type[ServerAppIoStub] | type[SimulationIoStub],
     appio_api_address: str,
-    plugin_config: Optional[dict[str, Any]] = None,
-    flwr_dir: Optional[str] = None,
-    parent_pid: Optional[int] = None,
-    health_server_address: Optional[str] = None,
+    plugin_config: dict[str, Any] | None = None,
+    flwr_dir: str | None = None,
+    parent_pid: int | None = None,
+    health_server_address: str | None = None,
 ) -> None:
     """Run Flower SuperExec.
 
@@ -158,12 +156,10 @@ def run_with_deprecation_warning(  # pylint: disable=R0913, R0917
     cmd: str,
     plugin_type: str,
     plugin_class: type[ExecPlugin],
-    stub_class: Union[
-        type[ClientAppIoStub], type[ServerAppIoStub], type[SimulationIoStub]
-    ],
+    stub_class: type[ClientAppIoStub] | type[ServerAppIoStub] | type[SimulationIoStub],
     appio_api_address: str,
-    flwr_dir: Optional[str],
-    parent_pid: Optional[int],
+    flwr_dir: str | None,
+    parent_pid: int | None,
     warn_run_once: bool,
 ) -> None:
     """Log a deprecation warning and run the equivalent `flower-superexec` command.

flwr/superlink/artifact_provider/artifact_provider.py

@@ -16,14 +16,13 @@
 
 
 from abc import ABC, abstractmethod
-from typing import Optional
 
 
 class ArtifactProvider(ABC):
     """ArtifactProvider interface for providing artifact download links."""
 
     @abstractmethod
-    def get_url(self, run_id: int) -> Optional[str]:
+    def get_url(self, run_id: int) -> str | None:
         """Return the artifact download link for the given run ID."""
 
     @property

flwr/superlink/auth_plugin/auth_plugin.py

@@ -18,7 +18,6 @@
 from abc import ABC, abstractmethod
 from collections.abc import Sequence
 from pathlib import Path
-from typing import Optional, Union
 
 from flwr.common.typing import (
     AccountAuthCredentials,
@@ -48,25 +47,23 @@ class ControlAuthnPlugin(ABC):
         """Abstract constructor."""
 
     @abstractmethod
-    def get_login_details(self) -> Optional[AccountAuthLoginDetails]:
+    def get_login_details(self) -> AccountAuthLoginDetails | None:
         """Get the login details."""
 
     @abstractmethod
     def validate_tokens_in_metadata(
-        self, metadata: Sequence[tuple[str, Union[str, bytes]]]
-    ) -> tuple[bool, Optional[AccountInfo]]:
+        self, metadata: Sequence[tuple[str, str | bytes]]
+    ) -> tuple[bool, AccountInfo | None]:
         """Validate authentication tokens in the provided metadata."""
 
     @abstractmethod
-    def get_auth_tokens(self, device_code: str) -> Optional[AccountAuthCredentials]:
+    def get_auth_tokens(self, device_code: str) -> AccountAuthCredentials | None:
         """Get authentication tokens."""
 
     @abstractmethod
     def refresh_tokens(
-        self, metadata: Sequence[tuple[str, Union[str, bytes]]]
-    ) -> tuple[
-        Optional[Sequence[tuple[str, Union[str, bytes]]]], Optional[AccountInfo]
-    ]:
+        self, metadata: Sequence[tuple[str, str | bytes]]
+    ) -> tuple[Sequence[tuple[str, str | bytes]] | None, AccountInfo | None]:
         """Refresh authentication tokens in the provided metadata."""
 
 

flwr/superlink/auth_plugin/noop_auth_plugin.py

@@ -18,7 +18,6 @@ authorization plugins."""
 
 from collections.abc import Sequence
 from pathlib import Path
-from typing import Optional, Union
 
 from flwr.common.constant import NOOP_ACCOUNT_NAME, NOOP_FLWR_AID, AuthnType
 from flwr.common.typing import (
@@ -45,7 +44,7 @@ class NoOpControlAuthnPlugin(ControlAuthnPlugin):
     ):
         pass
 
-    def get_login_details(self) -> Optional[AccountAuthLoginDetails]:
+    def get_login_details(self) -> AccountAuthLoginDetails | None:
         """Get the login details."""
         # This allows the `flwr login` command to load the NoOp plugin accordingly,
         # which then raises a LoginError when attempting to login.
@@ -58,20 +57,18 @@ class NoOpControlAuthnPlugin(ControlAuthnPlugin):
         )
 
     def validate_tokens_in_metadata(
-        self, metadata: Sequence[tuple[str, Union[str, bytes]]]
-    ) -> tuple[bool, Optional[AccountInfo]]:
+        self, metadata: Sequence[tuple[str, str | bytes]]
+    ) -> tuple[bool, AccountInfo | None]:
         """Return valid for no-op plugin."""
         return True, NOOP_ACCOUNT_INFO
 
-    def get_auth_tokens(self, device_code: str) -> Optional[AccountAuthCredentials]:
+    def get_auth_tokens(self, device_code: str) -> AccountAuthCredentials | None:
         """Get authentication tokens."""
         raise RuntimeError("NoOp plugin does not support getting auth tokens.")
 
     def refresh_tokens(
-        self, metadata: Sequence[tuple[str, Union[str, bytes]]]
-    ) -> tuple[
-        Optional[Sequence[tuple[str, Union[str, bytes]]]], Optional[AccountInfo]
-    ]:
+        self, metadata: Sequence[tuple[str, str | bytes]]
+    ) -> tuple[Sequence[tuple[str, str | bytes]] | None, AccountInfo | None]:
         """Refresh authentication tokens in the provided metadata."""
         return metadata, NOOP_ACCOUNT_INFO
 

flwr/superlink/federation/__init__.py

@@ -0,0 +1,24 @@
+# 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.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Federation Managers."""
+
+
+from .federation_manager import FederationManager
+from .noop_federation_manager import NoOpFederationManager
+
+__all__ = [
+    "FederationManager",
+    "NoOpFederationManager",
+]

flwr/superlink/federation/federation_manager.py

@@ -0,0 +1,64 @@
+# 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.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Abstract base class FederationManager."""
+
+
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING
+
+from flwr.common.typing import Federation
+
+if TYPE_CHECKING:
+    from flwr.server.superlink.linkstate.linkstate import LinkState
+
+
+class FederationManager(ABC):
+    """Abstract base class for FederationManager."""
+
+    @property
+    def linkstate(self) -> "LinkState":
+        """Return the LinkState instance."""
+        if not (ret := getattr(self, "_linkstate", None)):
+            raise RuntimeError("linkstate not set. Assign to linkstate property first.")
+        return ret  # type: ignore
+
+    @linkstate.setter
+    def linkstate(self, linkstate: "LinkState") -> None:
+        """Set the LinkState instance."""
+        self._linkstate = linkstate
+
+    @abstractmethod
+    def exists(self, federation: str) -> bool:
+        """Check if a federation exists."""
+
+    @abstractmethod
+    def has_member(self, flwr_aid: str, federation: str) -> bool:
+        """Check if the given account is a member of the federation."""
+
+    @abstractmethod
+    def filter_nodes(self, node_ids: set[int], federation: str) -> set[int]:
+        """Given a list of node IDs, return sublist with nodes in federation."""
+
+    @abstractmethod
+    def has_node(self, node_id: int, federation: str) -> bool:
+        """Given a node ID, check if it is in the federation."""
+
+    @abstractmethod
+    def get_federations(self, flwr_aid: str) -> list[str]:
+        """Get federations of which the account is a member."""
+
+    @abstractmethod
+    def get_details(self, federation: str) -> Federation:
+        """Get details of the federation."""