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

flwr/supercore/corestate/sqlite_corestate.py

@@ -0,0 +1,157 @@
+# 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.
+# ==============================================================================
+"""SQLite-based CoreState implementation."""
+
+
+import secrets
+import sqlite3
+from typing import cast
+
+from flwr.common import now
+from flwr.common.constant import (
+    FLWR_APP_TOKEN_LENGTH,
+    HEARTBEAT_DEFAULT_INTERVAL,
+    HEARTBEAT_PATIENCE,
+)
+from flwr.supercore.sqlite_mixin import SqliteMixin
+from flwr.supercore.utils import int64_to_uint64, uint64_to_int64
+
+from ..object_store import ObjectStore
+from .corestate import CoreState
+
+SQL_CREATE_TABLE_TOKEN_STORE = """
+CREATE TABLE IF NOT EXISTS token_store (
+    run_id                  INTEGER PRIMARY KEY,
+    token                   TEXT UNIQUE NOT NULL,
+    active_until            REAL
+);
+"""
+
+
+class SqliteCoreState(CoreState, SqliteMixin):
+    """SQLite-based CoreState implementation."""
+
+    def __init__(self, database_path: str, object_store: ObjectStore) -> None:
+        super().__init__(database_path)
+        self._object_store = object_store
+
+    @property
+    def object_store(self) -> ObjectStore:
+        """Return the ObjectStore instance used by this CoreState."""
+        return self._object_store
+
+    def get_sql_statements(self) -> tuple[str, ...]:
+        """Return SQL statements needed for CoreState tables."""
+        return (SQL_CREATE_TABLE_TOKEN_STORE,)
+
+    def create_token(self, run_id: int) -> str | None:
+        """Create a token for the given run ID."""
+        token = secrets.token_hex(FLWR_APP_TOKEN_LENGTH)  # Generate a random token
+        current = now().timestamp()
+        active_until = current + HEARTBEAT_DEFAULT_INTERVAL
+        query = """
+            INSERT INTO token_store (run_id, token, active_until)
+            VALUES (:run_id, :token, :active_until);
+        """
+        data = {
+            "run_id": uint64_to_int64(run_id),
+            "token": token,
+            "active_until": active_until,
+        }
+        try:
+            self.query(query, data)
+        except sqlite3.IntegrityError:
+            return None  # Token already created for this run ID
+        return token
+
+    def verify_token(self, run_id: int, token: str) -> bool:
+        """Verify a token for the given run ID."""
+        self._cleanup_expired_tokens()
+        query = "SELECT token FROM token_store WHERE run_id = :run_id;"
+        data = {"run_id": uint64_to_int64(run_id)}
+        rows = self.query(query, data)
+        if not rows:
+            return False
+        return cast(str, rows[0]["token"]) == token
+
+    def delete_token(self, run_id: int) -> None:
+        """Delete the token for the given run ID."""
+        query = "DELETE FROM token_store WHERE run_id = :run_id;"
+        data = {"run_id": uint64_to_int64(run_id)}
+        self.query(query, data)
+
+    def get_run_id_by_token(self, token: str) -> int | None:
+        """Get the run ID associated with a given token."""
+        self._cleanup_expired_tokens()
+        query = "SELECT run_id FROM token_store WHERE token = :token;"
+        data = {"token": token}
+        rows = self.query(query, data)
+        if not rows:
+            return None
+        return int64_to_uint64(rows[0]["run_id"])
+
+    def acknowledge_app_heartbeat(self, token: str) -> bool:
+        """Acknowledge an app heartbeat with the provided token."""
+        # Clean up expired tokens
+        self._cleanup_expired_tokens()
+
+        # Update the active_until field
+        current = now().timestamp()
+        active_until = current + HEARTBEAT_PATIENCE * HEARTBEAT_DEFAULT_INTERVAL
+        query = """
+            UPDATE token_store
+            SET active_until = :active_until
+            WHERE token = :token
+            RETURNING run_id;
+        """
+        data = {"active_until": active_until, "token": token}
+        rows = self.query(query, data)
+        return len(rows) > 0
+
+    def _cleanup_expired_tokens(self) -> None:
+        """Remove expired tokens and perform additional cleanup.
+
+        This method is called before token operations to ensure integrity.
+        Subclasses can override `_on_tokens_expired` to add custom cleanup logic.
+        """
+        current = now().timestamp()
+
+        with self.conn:
+            # Delete expired tokens and get their run_ids and active_until timestamps
+            query = """
+                DELETE FROM token_store
+                WHERE active_until < :current
+                RETURNING run_id, active_until;
+            """
+            rows = self.conn.execute(query, {"current": current}).fetchall()
+            expired_records = [
+                (int64_to_uint64(row["run_id"]), row["active_until"]) for row in rows
+            ]
+
+            # Hook for subclasses
+            if expired_records:
+                self._on_tokens_expired(expired_records)
+
+    def _on_tokens_expired(self, expired_records: list[tuple[int, float]]) -> None:
+        """Handle cleanup of expired tokens.
+
+        Override in subclasses to add custom cleanup logic.
+
+        Parameters
+        ----------
+        expired_records : list[tuple[int, float]]
+            List of tuples containing (run_id, active_until timestamp)
+            for expired tokens.
+        """

flwr/supercore/ffs/disk_ffs.py

@@ -18,7 +18,6 @@
 import hashlib
 import json
 from pathlib import Path
-from typing import Optional
 
 from .ffs import Ffs
 
@@ -59,7 +58,7 @@ class DiskFfs(Ffs):  # pylint: disable=R0904
 
         return content_hash
 
-    def get(self, key: str) -> Optional[tuple[bytes, dict[str, str]]]:
+    def get(self, key: str) -> tuple[bytes, dict[str, str]] | None:
         """Return tuple containing the object content and metadata.
 
         Parameters

flwr/supercore/ffs/ffs.py

@@ -16,7 +16,6 @@
 
 
 import abc
-from typing import Optional
 
 
 class Ffs(abc.ABC):  # pylint: disable=R0904
@@ -40,7 +39,7 @@ class Ffs(abc.ABC):  # pylint: disable=R0904
         """
 
     @abc.abstractmethod
-    def get(self, key: str) -> Optional[tuple[bytes, dict[str, str]]]:
+    def get(self, key: str) -> tuple[bytes, dict[str, str]] | None:
         """Return tuple containing the object content and metadata.
 
         Parameters

flwr/supercore/ffs/ffs_factory.py

@@ -16,7 +16,6 @@
 
 
 from logging import DEBUG
-from typing import Optional
 
 from flwr.common.logger import log
 
@@ -35,7 +34,7 @@ class FfsFactory:
 
     def __init__(self, base_dir: str) -> None:
         self.base_dir = base_dir
-        self.ffs_instance: Optional[Ffs] = None
+        self.ffs_instance: Ffs | None = None
 
     def ffs(self) -> Ffs:
         """Return a Ffs instance and create it, if necessary."""

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,
@@ -48,9 +47,6 @@ class InMemoryObjectStore(ObjectStore):
         self.verify = verify
         self.store: dict[str, ObjectEntry] = {}
         self.lock_store = threading.RLock()
-        # Mapping the Object ID of a message to the list of descendant object IDs
-        self.msg_descendant_objects_mapping: dict[str, list[str]] = {}
-        self.lock_msg_mapping = threading.RLock()
         # Mapping each run ID to a set of object IDs that are used in that run
         self.run_objects_mapping: dict[int, set[str]] = {}
 
@@ -157,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
@@ -215,7 +211,6 @@ class InMemoryObjectStore(ObjectStore):
         """Clear the store."""
         with self.lock_store:
             self.store.clear()
-            self.msg_descendant_objects_mapping.clear()
             self.run_objects_mapping.clear()
 
     def __contains__(self, object_id: str) -> bool:

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,19 +16,30 @@
 
 
 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 .in_memory_object_store import InMemoryObjectStore
 from .object_store import ObjectStore
+from .sqlite_object_store import SqliteObjectStore
 
 
 class ObjectStoreFactory:
-    """Factory class that creates ObjectStore instances."""
+    """Factory class that creates ObjectStore instances.
 
-    def __init__(self) -> None:
-        self.store_instance: Optional[ObjectStore] = None
+    Parameters
+    ----------
+    database : str (default: FLWR_IN_MEMORY_DB_NAME)
+        A string representing the path to the database file that will be opened.
+        Note that passing ":memory:" will open a connection to a database that is
+        in RAM, instead of on disk. And FLWR_IN_MEMORY_DB_NAME will create an
+        Python-based in-memory ObjectStore.
+    """
+
+    def __init__(self, database: str = FLWR_IN_MEMORY_DB_NAME) -> None:
+        self.database = database
+        self.store_instance: ObjectStore | None = None
 
     def store(self) -> ObjectStore:
         """Return an ObjectStore instance and create it, if necessary.
@@ -38,7 +49,15 @@ class ObjectStoreFactory:
         ObjectStore
             An ObjectStore instance for storing objects by object_id.
         """
-        if self.store_instance is None:
-            self.store_instance = InMemoryObjectStore()
-        log(DEBUG, "Using InMemoryObjectStore")
-        return self.store_instance
+        # InMemoryObjectStore
+        if self.database == FLWR_IN_MEMORY_DB_NAME:
+            if self.store_instance is None:
+                self.store_instance = InMemoryObjectStore()
+            log(DEBUG, "Using InMemoryObjectStore")
+            return self.store_instance
+
+        # SqliteObjectStore
+        store = SqliteObjectStore(self.database)
+        store.initialize()
+        log(DEBUG, "Using SqliteObjectStore")
+        return store

flwr/supercore/object_store/sqlite_object_store.py

@@ -0,0 +1,253 @@
+# 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.
+# ==============================================================================
+"""Flower SQLite ObjectStore implementation."""
+
+
+from typing import cast
+
+from flwr.common.inflatable import (
+    get_object_id,
+    is_valid_sha256_hash,
+    iterate_object_tree,
+)
+from flwr.common.inflatable_utils import validate_object_content
+from flwr.proto.message_pb2 import ObjectTree  # pylint: disable=E0611
+from flwr.supercore.sqlite_mixin import SqliteMixin
+from flwr.supercore.utils import uint64_to_int64
+
+from .object_store import NoObjectInStoreError, ObjectStore
+
+SQL_CREATE_OBJECTS = """
+CREATE TABLE IF NOT EXISTS objects (
+    object_id               TEXT PRIMARY KEY,
+    content                 BLOB,
+    is_available            INTEGER NOT NULL CHECK (is_available IN (0,1)),
+    ref_count               INTEGER NOT NULL
+);
+"""
+SQL_CREATE_OBJECT_CHILDREN = """
+CREATE TABLE IF NOT EXISTS object_children (
+    parent_id               TEXT NOT NULL,
+    child_id                TEXT NOT NULL,
+    FOREIGN KEY (parent_id) REFERENCES objects(object_id) ON DELETE CASCADE,
+    FOREIGN KEY (child_id)  REFERENCES objects(object_id) ON DELETE CASCADE,
+    PRIMARY KEY (parent_id, child_id)
+);
+"""
+SQL_CREATE_RUN_OBJECTS = """
+CREATE TABLE IF NOT EXISTS run_objects (
+    run_id                  INTEGER NOT NULL,
+    object_id               TEXT NOT NULL,
+    FOREIGN KEY (object_id) REFERENCES objects(object_id) ON DELETE CASCADE,
+    PRIMARY KEY (run_id, object_id)
+);
+"""
+
+
+class SqliteObjectStore(ObjectStore, SqliteMixin):
+    """SQLite-based implementation of the ObjectStore interface."""
+
+    def __init__(self, database_path: str, verify: bool = True) -> None:
+        super().__init__(database_path)
+        self.verify = verify
+
+    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,
+        )
+
+    def preregister(self, run_id: int, object_tree: ObjectTree) -> list[str]:
+        """Identify and preregister missing objects in the `ObjectStore`."""
+        new_objects = []
+        for tree_node in iterate_object_tree(object_tree):
+            obj_id = tree_node.object_id
+            if not is_valid_sha256_hash(obj_id):
+                raise ValueError(f"Invalid object ID format: {obj_id}")
+
+            child_ids = [child.object_id for child in tree_node.children]
+            with self.conn:
+                row = self.conn.execute(
+                    "SELECT object_id, is_available FROM objects WHERE object_id=?",
+                    (obj_id,),
+                ).fetchone()
+                if row is None:
+                    # Insert new object
+                    self.conn.execute(
+                        "INSERT INTO objects"
+                        "(object_id, content, is_available, ref_count) "
+                        "VALUES (?, ?, ?, ?)",
+                        (obj_id, b"", 0, 0),
+                    )
+                    for cid in child_ids:
+                        self.conn.execute(
+                            "INSERT INTO object_children(parent_id, child_id) "
+                            "VALUES (?, ?)",
+                            (obj_id, cid),
+                        )
+                        self.conn.execute(
+                            "UPDATE objects SET ref_count = ref_count + 1 "
+                            "WHERE object_id = ?",
+                            (cid,),
+                        )
+                    new_objects.append(obj_id)
+                else:
+                    # Add to the list of new objects if not available
+                    if not row["is_available"]:
+                        new_objects.append(obj_id)
+
+                # Ensure run mapping
+                self.conn.execute(
+                    "INSERT OR IGNORE INTO run_objects(run_id, object_id) "
+                    "VALUES (?, ?)",
+                    (uint64_to_int64(run_id), obj_id),
+                )
+        return new_objects
+
+    def get_object_tree(self, object_id: str) -> ObjectTree:
+        """Get the object tree for a given object ID."""
+        with self.conn:
+            row = self.conn.execute(
+                "SELECT object_id FROM objects WHERE object_id=?", (object_id,)
+            ).fetchone()
+            if not row:
+                raise NoObjectInStoreError(
+                    f"Object {object_id} was not pre-registered."
+                )
+            children = self.query(
+                "SELECT child_id FROM object_children WHERE parent_id=?", (object_id,)
+            )
+
+            # Build the object trees of all children
+            try:
+                child_trees = [self.get_object_tree(ch["child_id"]) for ch in children]
+            except NoObjectInStoreError as e:
+                # Raise an error if any child object is missing
+                # This indicates an integrity issue
+                raise NoObjectInStoreError(
+                    f"Object tree for object ID '{object_id}' contains missing "
+                    "children. This may indicate a corrupted object store."
+                ) from e
+
+            # Create and return the ObjectTree for the current object
+            return ObjectTree(object_id=object_id, children=child_trees)
+
+    def put(self, object_id: str, object_content: bytes) -> None:
+        """Put an object into the store."""
+        if self.verify:
+            # Verify object_id and object_content match
+            object_id_from_content = get_object_id(object_content)
+            if object_id != object_id_from_content:
+                raise ValueError(f"Object ID {object_id} does not match content hash")
+
+            # Validate object content
+            validate_object_content(content=object_content)
+
+        with self.conn:
+            # Only allow adding the object if it has been preregistered
+            row = self.conn.execute(
+                "SELECT is_available FROM objects WHERE object_id=?", (object_id,)
+            ).fetchone()
+            if row is None:
+                raise NoObjectInStoreError(
+                    f"Object with ID '{object_id}' was not pre-registered."
+                )
+
+            # Return if object is already present in the store
+            if row["is_available"]:
+                return
+
+            # Update the object entry in the store
+            self.conn.execute(
+                "UPDATE objects SET content=?, is_available=1 WHERE object_id=?",
+                (object_content, object_id),
+            )
+
+    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
+
+    def delete(self, object_id: str) -> None:
+        """Delete an object and its unreferenced descendants from the store."""
+        with self.conn:
+            row = self.conn.execute(
+                "SELECT ref_count FROM objects WHERE object_id=?", (object_id,)
+            ).fetchone()
+
+            # If the object is not in the store, nothing to delete
+            if row is None:
+                return
+
+            # Skip deletion if there are still references
+            if row["ref_count"] > 0:
+                return
+
+            # Deleting will cascade via FK, but we need to decrement children first
+            children = self.conn.execute(
+                "SELECT child_id FROM object_children WHERE parent_id=?", (object_id,)
+            ).fetchall()
+            child_ids = [child["child_id"] for child in children]
+
+            if child_ids:
+                placeholders = ", ".join("?" for _ in child_ids)
+                query = f"""
+                    UPDATE objects SET ref_count = ref_count - 1
+                    WHERE object_id IN ({placeholders})
+                """
+                self.conn.execute(query, child_ids)
+
+            self.conn.execute("DELETE FROM objects WHERE object_id=?", (object_id,))
+
+            # Recursively clean children
+            for child_id in child_ids:
+                self.delete(child_id)
+
+    def delete_objects_in_run(self, run_id: int) -> None:
+        """Delete all objects that were registered in a specific run."""
+        run_id_sint = uint64_to_int64(run_id)
+        with self.conn:
+            objs = self.conn.execute(
+                "SELECT object_id FROM run_objects WHERE run_id=?", (run_id_sint,)
+            ).fetchall()
+            for obj in objs:
+                object_id = obj["object_id"]
+                row = self.conn.execute(
+                    "SELECT ref_count FROM objects WHERE object_id=?", (object_id,)
+                ).fetchone()
+                if row and row["ref_count"] == 0:
+                    self.delete(object_id)
+            self.conn.execute("DELETE FROM run_objects WHERE run_id=?", (run_id_sint,))
+
+    def clear(self) -> None:
+        """Clear the store."""
+        with self.conn:
+            self.conn.execute("DELETE FROM object_children;")
+            self.conn.execute("DELETE FROM run_objects;")
+            self.conn.execute("DELETE FROM objects;")
+
+    def __contains__(self, object_id: str) -> bool:
+        """Check if an object_id is in the store."""
+        row = self.conn.execute(
+            "SELECT 1 FROM objects WHERE object_id=?", (object_id,)
+        ).fetchone()
+        return row is not None
+
+    def __len__(self) -> int:
+        """Return the number of objects in the store."""
+        row = self.conn.execute("SELECT COUNT(*) AS cnt FROM objects;").fetchone()
+        return cast(int, row["cnt"])

flwr/{client/clientapp → supercore/primitives}/__init__.py

@@ -12,4 +12,4 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""Flower AppIO service."""
+"""Cryptographic primitives for the Flower infrastructure."""