flwr-nightly 1.23.0.dev20251020__py3-none-any.whl → 1.23.0.dev20251022__py3-none-any.whl

flwr/supercore/object_store/sqlite_object_store.py

@@ -0,0 +1,252 @@
+# 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 Optional, 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 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(
+            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]:
+        """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} not found.")
+            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) -> Optional[bytes]:
+        """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/supercore/sqlite_mixin.py

@@ -0,0 +1,156 @@
+# 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.
+# ==============================================================================
+"""Mixin providing common SQLite connection and initialization logic."""
+
+
+import re
+import sqlite3
+from abc import ABC, abstractmethod
+from collections.abc import Sequence
+from logging import DEBUG, ERROR
+from typing import Any, Optional, Union
+
+from flwr.common.logger import log
+
+DictOrTuple = Union[tuple[Any, ...], dict[str, Any]]
+
+
+class SqliteMixin(ABC):
+    """Mixin providing common SQLite connection and initialization logic."""
+
+    def __init__(self, database_path: str) -> None:
+        self.database_path = database_path
+        self._conn: Optional[sqlite3.Connection] = None
+
+    @property
+    def conn(self) -> sqlite3.Connection:
+        """Get the SQLite connection."""
+        if self._conn is None:
+            raise AttributeError("Database not initialized. Call initialize() first.")
+        return self._conn
+
+    @abstractmethod
+    def initialize(self, log_queries: bool = False) -> list[tuple[str]]:
+        """Connect to the DB, enable FK support, and create tables if needed.
+
+        Parameters
+        ----------
+        log_queries : bool
+            Log each query which is executed.
+
+        Returns
+        -------
+        list[tuple[str]]
+            The list of all tables in the DB.
+
+        Examples
+        --------
+        Implement in subclass:
+
+        .. code:: python
+
+            def initialize(self, log_queries: bool = False) -> list[tuple[str]]:
+                return self._ensure_initialized(
+                    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.
+
+        Parameters
+        ----------
+        create_statements : str
+            SQL statements to create tables and indexes.
+        log_queries : bool
+            Log each query which is executed.
+
+        Returns
+        -------
+        list[tuple[str]]
+            The list of all tables in the DB.
+        """
+        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.row_factory = dict_factory
+
+        if log_queries:
+            self._conn.set_trace_callback(lambda q: log(DEBUG, q))
+
+        # Create tables and indexes
+        cur = self._conn.cursor()
+        for sql in create_statements:
+            cur.execute(sql)
+        res = cur.execute("SELECT name FROM sqlite_schema;")
+        return res.fetchall()
+
+    def query(
+        self,
+        query: str,
+        data: Optional[Union[Sequence[DictOrTuple], DictOrTuple]] = None,
+    ) -> list[dict[str, Any]]:
+        """Execute a SQL query and return the results as list of dicts."""
+        if self._conn is None:
+            raise AttributeError("LinkState is not initialized.")
+
+        if data is None:
+            data = []
+
+        # Clean up whitespace to make the logs nicer
+        query = re.sub(r"\s+", " ", query)
+
+        try:
+            with self._conn:
+                if (
+                    len(data) > 0
+                    and isinstance(data, (tuple, list))
+                    and isinstance(data[0], (tuple, dict))
+                ):
+                    rows = self._conn.executemany(query, data)
+                else:
+                    rows = self._conn.execute(query, data)
+
+                # Extract results before committing to support
+                #   INSERT/UPDATE ... RETURNING
+                # style queries
+                result = rows.fetchall()
+        except KeyError as exc:
+            log(ERROR, {"query": query, "data": data, "exception": exc})
+
+        return result
+
+
+def dict_factory(
+    cursor: sqlite3.Cursor,
+    row: sqlite3.Row,
+) -> dict[str, Any]:
+    """Turn SQLite results into dicts.
+
+    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))

flwr/supercore/utils.py

@@ -30,3 +30,23 @@ def mask_string(value: str, head: int = 4, tail: int = 4) -> str:
     if len(value) <= head + tail:
         return value
     return f"{value[:head]}...{value[-tail:]}"
+
+
+def uint64_to_int64(unsigned: int) -> int:
+    """Convert a uint64 integer to a sint64 with the same bit pattern.
+
+    For values >= 2^63, wraps around by subtracting 2^64.
+    """
+    if unsigned >= (1 << 63):
+        return unsigned - (1 << 64)
+    return unsigned
+
+
+def int64_to_uint64(signed: int) -> int:
+    """Convert a sint64 integer to a uint64 with the same bit pattern.
+
+    For negative values, wraps around by adding 2^64.
+    """
+    if signed < 0:
+        return signed + (1 << 64)
+    return signed

flwr/superlink/servicer/control/control_servicer.py

@@ -18,7 +18,6 @@
 import hashlib
 import time
 from collections.abc import Generator, Sequence
-from datetime import timedelta
 from logging import ERROR, INFO
 from typing import Any, Optional, cast
 
@@ -72,7 +71,6 @@ from flwr.proto.control_pb2 import (  # pylint: disable=E0611
 )
 from flwr.proto.node_pb2 import NodeInfo  # pylint: disable=E0611
 from flwr.server.superlink.linkstate import LinkState, LinkStateFactory
-from flwr.supercore.constant import NodeStatus
 from flwr.supercore.ffs import FfsFactory
 from flwr.supercore.object_store import ObjectStore, ObjectStoreFactory
 from flwr.supercore.primitives.asymmetric import bytes_to_public_key, uses_nist_ec_curve
@@ -466,79 +464,17 @@ class ControlServicer(control_pb2_grpc.ControlServicer):
             raise grpc.RpcError()  # This line is unreachable
 
         nodes_info: Sequence[NodeInfo] = []
-        # If dry run is enabled, create dummy NodeInfo data
-        if request.dry_run:
-            nodes_info = _create_list_nodeif_for_dry_run()
-
-        else:
-            # Init link state
-            state = self.linkstate_factory.state()
+        # Init link state
+        state = self.linkstate_factory.state()
 
-            flwr_aid = shared_account_info.get().flwr_aid
-            flwr_aid = _check_flwr_aid_exists(flwr_aid, context)
-            # Retrieve all nodes for the account
-            nodes_info = state.get_node_info(owner_aids=[flwr_aid])
+        flwr_aid = shared_account_info.get().flwr_aid
+        flwr_aid = _check_flwr_aid_exists(flwr_aid, context)
+        # Retrieve all nodes for the account
+        nodes_info = state.get_node_info(owner_aids=[flwr_aid])
 
         return ListNodesResponse(nodes_info=nodes_info, now=now().isoformat())
 
 
-def _create_list_nodeif_for_dry_run() -> Sequence[NodeInfo]:
-    """Create a list of NodeInfo for dry run testing."""
-    nodes_info: list[NodeInfo] = []
-    # A node registered (but not connected)
-    nodes_info.append(
-        NodeInfo(
-            node_id=15390646978706312628,
-            owner_aid="owner_aid_1",
-            status=NodeStatus.REGISTERED,
-            registered_at=(now()).isoformat(),
-            last_activated_at="",
-            last_deactivated_at="",
-            unregistered_at="",
-        )
-    )
-
-    # A node registered and connected
-    nodes_info.append(
-        NodeInfo(
-            node_id=2941141058168602545,
-            owner_aid="owner_aid_2",
-            status=NodeStatus.ONLINE,
-            registered_at=(now()).isoformat(),
-            last_activated_at=(now() + timedelta(hours=0.5)).isoformat(),
-            last_deactivated_at="",
-            unregistered_at="",
-        )
-    )
-
-    # A node registered and unregistered (never connected)
-    nodes_info.append(
-        NodeInfo(
-            node_id=906971720890549292,
-            owner_aid="owner_aid_3",
-            status=NodeStatus.UNREGISTERED,
-            registered_at=(now()).isoformat(),
-            last_activated_at="",
-            last_deactivated_at="",
-            unregistered_at=(now() + timedelta(hours=1)).isoformat(),
-        )
-    )
-
-    # A node registered, deactivate and then unregistered
-    nodes_info.append(
-        NodeInfo(
-            node_id=1781174086018058152,
-            owner_aid="owner_aid_4",
-            status=NodeStatus.OFFLINE,
-            registered_at=(now()).isoformat(),
-            last_activated_at=(now() + timedelta(hours=0.5)).isoformat(),
-            last_deactivated_at=(now() + timedelta(hours=1)).isoformat(),
-            unregistered_at=(now() + timedelta(hours=1.5)).isoformat(),
-        )
-    )
-    return nodes_info
-
-
 def _create_list_runs_response(
     run_ids: set[int], state: LinkState, store: ObjectStore
 ) -> ListRunsResponse:

{flwr_nightly-1.23.0.dev20251020.dist-info → flwr_nightly-1.23.0.dev20251022.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: flwr-nightly
-Version: 1.23.0.dev20251020
+Version: 1.23.0.dev20251022
 Summary: Flower: A Friendly Federated AI Framework
 License: Apache-2.0
 Keywords: Artificial Intelligence,Federated AI,Federated Analytics,Federated Evaluation,Federated Learning,Flower,Machine Learning