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

flwr/server/superlink/linkstate/sqlite_linkstate.py

@@ -18,11 +18,10 @@
 # pylint: disable=too-many-lines
 
 import json
-import re
 import secrets
 import sqlite3
 from collections.abc import Sequence
-from logging import DEBUG, ERROR, WARNING
+from logging import ERROR, WARNING
 from typing import Any, Optional, Union, cast
 
 from flwr.common import Context, Message, Metadata, log, now
@@ -52,6 +51,8 @@ from flwr.proto.recorddict_pb2 import RecordDict as ProtoRecordDict
 # pylint: enable=E0611
 from flwr.server.utils.validator import validate_message
 from flwr.supercore.constant import NodeStatus
+from flwr.supercore.sqlite_mixin import SqliteMixin
+from flwr.supercore.utils import int64_to_uint64, uint64_to_int64
 
 from .linkstate import LinkState
 from .utils import (
@@ -60,9 +61,7 @@ from .utils import (
     configrecord_to_bytes,
     context_from_bytes,
     context_to_bytes,
-    convert_sint64_to_uint64,
     convert_sint64_values_in_dict_to_uint64,
-    convert_uint64_to_sint64,
     convert_uint64_values_in_dict_to_sint64,
     generate_rand_int_from_bytes,
     has_valid_sub_status,
@@ -183,95 +182,25 @@ CREATE TABLE IF NOT EXISTS token_store (
 );
 """
 
-DictOrTuple = Union[tuple[Any, ...], dict[str, Any]]
 
-
-class SqliteLinkState(LinkState):  # pylint: disable=R0904
+class SqliteLinkState(LinkState, SqliteMixin):  # pylint: disable=R0904
     """SQLite-based LinkState implementation."""
 
-    def __init__(
-        self,
-        database_path: str,
-    ) -> None:
-        """Initialize an SqliteLinkState.
-
-        Parameters
-        ----------
-        database : (path-like object)
-            The path to the database file to be opened. Pass ":memory:" to open
-            a connection to a database that is in RAM, instead of on disk.
-        """
-        self.database_path = database_path
-        self.conn: Optional[sqlite3.Connection] = None
-
     def initialize(self, log_queries: bool = False) -> list[tuple[str]]:
-        """Create tables if they don't exist yet.
-
-        Parameters
-        ----------
-        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)
-        self.conn.execute("PRAGMA foreign_keys = ON;")
-        self.conn.row_factory = dict_factory
-        if log_queries:
-            self.conn.set_trace_callback(lambda query: log(DEBUG, query))
-        cur = self.conn.cursor()
-
-        # Create each table if not exists queries
-        cur.execute(SQL_CREATE_TABLE_RUN)
-        cur.execute(SQL_CREATE_TABLE_LOGS)
-        cur.execute(SQL_CREATE_TABLE_CONTEXT)
-        cur.execute(SQL_CREATE_TABLE_MESSAGE_INS)
-        cur.execute(SQL_CREATE_TABLE_MESSAGE_RES)
-        cur.execute(SQL_CREATE_TABLE_NODE)
-        cur.execute(SQL_CREATE_TABLE_PUBLIC_KEY)
-        cur.execute(SQL_CREATE_TABLE_TOKEN_STORE)
-        cur.execute(SQL_CREATE_INDEX_ONLINE_UNTIL)
-        cur.execute(SQL_CREATE_INDEX_OWNER_AID)
-        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."""
-        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
+        """Connect to the DB, enable FK support, and create tables if needed."""
+        return self._ensure_initialized(
+            SQL_CREATE_TABLE_RUN,
+            SQL_CREATE_TABLE_LOGS,
+            SQL_CREATE_TABLE_CONTEXT,
+            SQL_CREATE_TABLE_MESSAGE_INS,
+            SQL_CREATE_TABLE_MESSAGE_RES,
+            SQL_CREATE_TABLE_NODE,
+            SQL_CREATE_TABLE_PUBLIC_KEY,
+            SQL_CREATE_TABLE_TOKEN_STORE,
+            SQL_CREATE_INDEX_ONLINE_UNTIL,
+            SQL_CREATE_INDEX_OWNER_AID,
+            log_queries=log_queries,
+        )
 
     def store_message_ins(self, message: Message) -> Optional[str]:
         """Store one Message."""
@@ -335,7 +264,7 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
         data: dict[str, Union[str, int]] = {}
 
         # Convert the uint64 value to sint64 for SQLite
-        data["node_id"] = convert_uint64_to_sint64(node_id)
+        data["node_id"] = uint64_to_int64(node_id)
 
         # Retrieve all Messages for node_id
         query = """
@@ -410,8 +339,7 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
         if (
             msg_ins
             and message
-            and convert_sint64_to_uint64(msg_ins["dst_node_id"])
-            != res_metadata.src_node_id
+            and int64_to_uint64(msg_ins["dst_node_id"]) != res_metadata.src_node_id
         ):
             return None
 
@@ -487,20 +415,20 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
         dst_node_ids: set[int] = set()
         for message_id in message_ids:
             in_message = found_message_ins_dict[message_id]
-            sint_node_id = convert_uint64_to_sint64(in_message.metadata.dst_node_id)
+            sint_node_id = uint64_to_int64(in_message.metadata.dst_node_id)
             dst_node_ids.add(sint_node_id)
         query = f"""
-                    SELECT node_id, online_until
-                    FROM node
-                    WHERE node_id IN ({",".join(["?"] * len(dst_node_ids))});
-                """
-        rows = self.query(query, tuple(dst_node_ids))
+            SELECT node_id, online_until
+            FROM node
+            WHERE node_id IN ({",".join(["?"] * len(dst_node_ids))})
+            AND status != ?
+        """
+        rows = self.query(query, tuple(dst_node_ids) + (NodeStatus.UNREGISTERED,))
         tmp_ret_dict = check_node_availability_for_in_message(
             inquired_in_message_ids=message_ids,
             found_in_message_dict=found_message_ins_dict,
             node_id_to_online_until={
-                convert_sint64_to_uint64(row["node_id"]): row["online_until"]
-                for row in rows
+                int64_to_uint64(row["node_id"]): row["online_until"] for row in rows
             },
             current_time=current,
         )
@@ -601,7 +529,7 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
             WHERE run_id = :run_id;
         """
 
-        sint64_run_id = convert_uint64_to_sint64(run_id)
+        sint64_run_id = uint64_to_int64(run_id)
         data = {"run_id": sint64_run_id}
 
         with self.conn:
@@ -619,7 +547,7 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
         )
 
         # Convert the uint64 value to sint64 for SQLite
-        sint64_node_id = convert_uint64_to_sint64(uint64_node_id)
+        sint64_node_id = uint64_to_int64(uint64_node_id)
 
         query = """
             INSERT INTO node
@@ -658,17 +586,21 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
 
     def delete_node(self, owner_aid: str, node_id: int) -> None:
         """Delete a node."""
-        sint64_node_id = convert_uint64_to_sint64(node_id)
+        sint64_node_id = uint64_to_int64(node_id)
 
         query = """
             UPDATE node
-            SET status = ?, unregistered_at = ?
+            SET status = ?, unregistered_at = ?,
+            online_until = IIF(online_until > ?, ?, online_until)
             WHERE node_id = ? AND status != ? AND owner_aid = ?
             RETURNING node_id
         """
+        current = now()
         params = (
             NodeStatus.UNREGISTERED,
-            now().isoformat(),
+            current.isoformat(),
+            current.timestamp(),
+            current.timestamp(),
             sint64_node_id,
             NodeStatus.UNREGISTERED,
             owner_aid,
@@ -693,7 +625,7 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
             raise AttributeError("LinkState not initialized")
 
         # Convert the uint64 value to sint64 for SQLite
-        sint64_run_id = convert_uint64_to_sint64(run_id)
+        sint64_run_id = uint64_to_int64(run_id)
 
         # Validate run ID
         query = "SELECT COUNT(*) FROM run WHERE run_id = ?"
@@ -738,9 +670,7 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
             conditions = []
             params = []
             if node_ids is not None:
-                sint64_node_ids = [
-                    convert_uint64_to_sint64(node_id) for node_id in node_ids
-                ]
+                sint64_node_ids = [uint64_to_int64(node_id) for node_id in node_ids]
                 placeholders = ",".join(["?"] * len(sint64_node_ids))
                 conditions.append(f"node_id IN ({placeholders})")
                 params.extend(sint64_node_ids)
@@ -763,7 +693,7 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
             result: list[NodeInfo] = []
             for row in rows:
                 # Convert sint64 node_id to uint64
-                row["node_id"] = convert_sint64_to_uint64(row["node_id"])
+                row["node_id"] = int64_to_uint64(row["node_id"])
                 result.append(NodeInfo(**row))
 
             return result
@@ -771,7 +701,7 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
     def get_node_public_key(self, node_id: int) -> bytes:
         """Get `public_key` for the specified `node_id`."""
         # Convert the uint64 value to sint64 for SQLite
-        sint64_node_id = convert_uint64_to_sint64(node_id)
+        sint64_node_id = uint64_to_int64(node_id)
 
         # Query the public key for the given node_id
         query = "SELECT public_key FROM node WHERE node_id = ? AND status != ?;"
@@ -795,7 +725,7 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
             return None
 
         # Convert sint64 node_id to uint64
-        node_id = convert_sint64_to_uint64(rows[0]["node_id"])
+        node_id = int64_to_uint64(rows[0]["node_id"])
         return node_id
 
     # pylint: disable=too-many-arguments,too-many-positional-arguments
@@ -813,7 +743,7 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
         uint64_run_id = generate_rand_int_from_bytes(RUN_ID_NUM_BYTES)
 
         # Convert the uint64 value to sint64 for SQLite
-        sint64_run_id = convert_uint64_to_sint64(uint64_run_id)
+        sint64_run_id = uint64_to_int64(uint64_run_id)
 
         # Check conflicts
         query = "SELECT COUNT(*) FROM run WHERE run_id = ?;"
@@ -861,7 +791,7 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
             )
         else:
             rows = self.query("SELECT run_id FROM run;", ())
-        return {convert_sint64_to_uint64(row["run_id"]) for row in rows}
+        return {int64_to_uint64(row["run_id"]) for row in rows}
 
     def _check_and_tag_inactive_run(self, run_ids: set[int]) -> None:
         """Check if any runs are no longer active.
@@ -869,7 +799,7 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
         Marks runs with status 'starting' or 'running' as failed
         if they have not sent a heartbeat before `active_until`.
         """
-        sint_run_ids = [convert_uint64_to_sint64(run_id) for run_id in run_ids]
+        sint_run_ids = [uint64_to_int64(run_id) for run_id in run_ids]
         query = "UPDATE run SET finished_at = ?, sub_status = ?, details = ? "
         query += "WHERE starting_at != '' AND finished_at = '' AND active_until < ?"
         query += f" AND run_id IN ({','.join(['?'] * len(run_ids))});"
@@ -891,13 +821,13 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
         self._check_and_tag_inactive_run(run_ids={run_id})
 
         # Convert the uint64 value to sint64 for SQLite
-        sint64_run_id = convert_uint64_to_sint64(run_id)
+        sint64_run_id = uint64_to_int64(run_id)
         query = "SELECT * FROM run WHERE run_id = ?;"
         rows = self.query(query, (sint64_run_id,))
         if rows:
             row = rows[0]
             return Run(
-                run_id=convert_sint64_to_uint64(row["run_id"]),
+                run_id=int64_to_uint64(row["run_id"]),
                 fab_id=row["fab_id"],
                 fab_version=row["fab_version"],
                 fab_hash=row["fab_hash"],
@@ -922,13 +852,13 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
         self._check_and_tag_inactive_run(run_ids=run_ids)
 
         # Convert the uint64 value to sint64 for SQLite
-        sint64_run_ids = (convert_uint64_to_sint64(run_id) for run_id in set(run_ids))
+        sint64_run_ids = (uint64_to_int64(run_id) for run_id in set(run_ids))
         query = f"SELECT * FROM run WHERE run_id IN ({','.join(['?'] * len(run_ids))});"
         rows = self.query(query, tuple(sint64_run_ids))
 
         return {
             # Restore uint64 run IDs
-            convert_sint64_to_uint64(row["run_id"]): RunStatus(
+            int64_to_uint64(row["run_id"]): RunStatus(
                 status=determine_run_status(row),
                 sub_status=row["sub_status"],
                 details=row["details"],
@@ -942,7 +872,7 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
         self._check_and_tag_inactive_run(run_ids={run_id})
 
         # Convert the uint64 value to sint64 for SQLite
-        sint64_run_id = convert_uint64_to_sint64(run_id)
+        sint64_run_id = uint64_to_int64(run_id)
         query = "SELECT * FROM run WHERE run_id = ?;"
         rows = self.query(query, (sint64_run_id,))
 
@@ -1008,7 +938,7 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
             new_status.details,
             active_until,
             heartbeat_interval,
-            convert_uint64_to_sint64(run_id),
+            uint64_to_int64(run_id),
         )
         self.query(query % timestamp_fld, data)
         return True
@@ -1021,14 +951,14 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
         query = "SELECT * FROM run WHERE starting_at = '' LIMIT 1;"
         rows = self.query(query)
         if rows:
-            pending_run_id = convert_sint64_to_uint64(rows[0]["run_id"])
+            pending_run_id = int64_to_uint64(rows[0]["run_id"])
 
         return pending_run_id
 
     def get_federation_options(self, run_id: int) -> Optional[ConfigRecord]:
         """Retrieve the federation options for the specified `run_id`."""
         # Convert the uint64 value to sint64 for SQLite
-        sint64_run_id = convert_uint64_to_sint64(run_id)
+        sint64_run_id = uint64_to_int64(run_id)
         query = "SELECT federation_options FROM run WHERE run_id = ?;"
         rows = self.query(query, (sint64_run_id,))
 
@@ -1053,7 +983,7 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
         if self.conn is None:
             raise AttributeError("LinkState not initialized")
 
-        sint64_node_id = convert_uint64_to_sint64(node_id)
+        sint64_node_id = uint64_to_int64(node_id)
 
         with self.conn:
             # Check if node exists and not deleted
@@ -1095,7 +1025,7 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
         self._check_and_tag_inactive_run(run_ids={run_id})
 
         # Search for the run
-        sint_run_id = convert_uint64_to_sint64(run_id)
+        sint_run_id = uint64_to_int64(run_id)
         query = "SELECT * FROM run WHERE run_id = ?;"
         rows = self.query(query, (sint_run_id,))
 
@@ -1125,7 +1055,7 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
         """Get the context for the specified `run_id`."""
         # Retrieve context if any
         query = "SELECT context FROM context WHERE run_id = ?;"
-        rows = self.query(query, (convert_uint64_to_sint64(run_id),))
+        rows = self.query(query, (uint64_to_int64(run_id),))
         context = context_from_bytes(rows[0]["context"]) if rows else None
         return context
 
@@ -1133,7 +1063,7 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
         """Set the context for the specified `run_id`."""
         # Convert context to bytes
         context_bytes = context_to_bytes(context)
-        sint_run_id = convert_uint64_to_sint64(run_id)
+        sint_run_id = uint64_to_int64(run_id)
 
         # Check if any existing Context assigned to the run_id
         query = "SELECT COUNT(*) FROM context WHERE run_id = ?;"
@@ -1152,7 +1082,7 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
     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`."""
         # Convert the uint64 value to sint64 for SQLite
-        sint64_run_id = convert_uint64_to_sint64(run_id)
+        sint64_run_id = uint64_to_int64(run_id)
 
         # Store log
         try:
@@ -1168,7 +1098,7 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
     ) -> tuple[str, float]:
         """Get the ServerApp logs for the specified `run_id`."""
         # Convert the uint64 value to sint64 for SQLite
-        sint64_run_id = convert_uint64_to_sint64(run_id)
+        sint64_run_id = uint64_to_int64(run_id)
 
         # Check if the run_id exists
         query = "SELECT run_id FROM run WHERE run_id = ?;"
@@ -1218,7 +1148,7 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
         """Create a token for the given run ID."""
         token = secrets.token_hex(FLWR_APP_TOKEN_LENGTH)  # Generate a random token
         query = "INSERT INTO token_store (run_id, token) VALUES (:run_id, :token);"
-        data = {"run_id": convert_uint64_to_sint64(run_id), "token": token}
+        data = {"run_id": uint64_to_int64(run_id), "token": token}
         try:
             self.query(query, data)
         except sqlite3.IntegrityError:
@@ -1228,7 +1158,7 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
     def verify_token(self, run_id: int, token: str) -> bool:
         """Verify a token for the given run ID."""
         query = "SELECT token FROM token_store WHERE run_id = :run_id;"
-        data = {"run_id": convert_uint64_to_sint64(run_id)}
+        data = {"run_id": uint64_to_int64(run_id)}
         rows = self.query(query, data)
         if not rows:
             return False
@@ -1237,7 +1167,7 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
     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": convert_uint64_to_sint64(run_id)}
+        data = {"run_id": uint64_to_int64(run_id)}
         self.query(query, data)
 
     def get_run_id_by_token(self, token: str) -> Optional[int]:
@@ -1247,19 +1177,7 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
         rows = self.query(query, data)
         if not rows:
             return None
-        return convert_sint64_to_uint64(rows[0]["run_id"])
-
-
-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))
+        return int64_to_uint64(rows[0]["run_id"])
 
 
 def message_to_dict(message: Message) -> dict[str, Any]:
@@ -1314,5 +1232,5 @@ def determine_run_status(row: dict[str, Any]) -> str:
                 return Status.RUNNING
             return Status.STARTING
         return Status.PENDING
-    run_id = convert_sint64_to_uint64(row["run_id"])
+    run_id = int64_to_uint64(row["run_id"])
     raise sqlite3.IntegrityError(f"The run {run_id} does not have a valid status.")

flwr/server/superlink/linkstate/utils.py

@@ -33,6 +33,7 @@ from flwr.common.typing import RunStatus
 # pylint: disable=E0611
 from flwr.proto.message_pb2 import Context as ProtoContext
 from flwr.proto.recorddict_pb2 import ConfigRecord as ProtoConfigRecord
+from flwr.supercore.utils import int64_to_uint64, uint64_to_int64
 
 # pylint: enable=E0611
 VALID_RUN_STATUS_TRANSITIONS = {
@@ -76,58 +77,6 @@ def generate_rand_int_from_bytes(
     return num
 
 
-def convert_uint64_to_sint64(u: int) -> int:
-    """Convert a uint64 value to a sint64 value with the same bit sequence.
-
-    Parameters
-    ----------
-    u : int
-        The unsigned 64-bit integer to convert.
-
-    Returns
-    -------
-    int
-        The signed 64-bit integer equivalent.
-
-        The signed 64-bit integer will have the same bit pattern as the
-        unsigned 64-bit integer but may have a different decimal value.
-
-        For numbers within the range [0, `sint64` max value], the decimal
-        value remains the same. However, for numbers greater than the `sint64`
-        max value, the decimal value will differ due to the wraparound caused
-        by the sign bit.
-    """
-    if u >= (1 << 63):
-        return u - (1 << 64)
-    return u
-
-
-def convert_sint64_to_uint64(s: int) -> int:
-    """Convert a sint64 value to a uint64 value with the same bit sequence.
-
-    Parameters
-    ----------
-    s : int
-        The signed 64-bit integer to convert.
-
-    Returns
-    -------
-    int
-        The unsigned 64-bit integer equivalent.
-
-        The unsigned 64-bit integer will have the same bit pattern as the
-        signed 64-bit integer but may have a different decimal value.
-
-        For negative `sint64` values, the conversion adds 2^64 to the
-        signed value to obtain the equivalent `uint64` value. For non-negative
-        `sint64` values, the decimal value remains unchanged in the `uint64`
-        representation.
-    """
-    if s < 0:
-        return s + (1 << 64)
-    return s
-
-
 def convert_uint64_values_in_dict_to_sint64(
     data_dict: dict[str, int], keys: list[str]
 ) -> None:
@@ -142,7 +91,7 @@ def convert_uint64_values_in_dict_to_sint64(
     """
     for key in keys:
         if key in data_dict:
-            data_dict[key] = convert_uint64_to_sint64(data_dict[key])
+            data_dict[key] = uint64_to_int64(data_dict[key])
 
 
 def convert_sint64_values_in_dict_to_uint64(
@@ -159,7 +108,7 @@ def convert_sint64_values_in_dict_to_uint64(
     """
     for key in keys:
         if key in data_dict:
-            data_dict[key] = convert_sint64_to_uint64(data_dict[key])
+            data_dict[key] = int64_to_uint64(data_dict[key])
 
 
 def context_to_bytes(context: Context) -> bytes:

flwr/simulation/run_simulation.py

@@ -51,6 +51,7 @@ from flwr.server.superlink.linkstate.utils import generate_rand_int_from_bytes
 from flwr.simulation.ray_transport.utils import (
     enable_tf_gpu_growth as enable_gpu_growth,
 )
+from flwr.supercore.constant import FLWR_IN_MEMORY_DB_NAME
 
 
 def _replace_keys(d: Any, match: str, target: str) -> Any:
@@ -336,7 +337,7 @@ def _main_loop(
 ) -> Context:
     """Start ServerApp on a separate thread, then launch Simulation Engine."""
     # Initialize StateFactory
-    state_factory = LinkStateFactory(":flwr-in-memory-state:")
+    state_factory = LinkStateFactory(FLWR_IN_MEMORY_DB_NAME)
 
     f_stop = threading.Event()
     # A Threading event to indicate if an exception was raised in the ServerApp thread

flwr/supercore/constant.py

@@ -20,6 +20,9 @@ from __future__ import annotations
 # Top-level key in YAML config for exec plugin settings
 EXEC_PLUGIN_SECTION = "exec_plugin"
 
+# Flower in-memory Python-based database name
+FLWR_IN_MEMORY_DB_NAME = ":flwr-in-memory:"
+
 
 class NodeStatus:
     """Event log writer types."""

flwr/supercore/object_store/object_store_factory.py

@@ -19,15 +19,27 @@ 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:
+    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: Optional[ObjectStore] = None
 
     def store(self) -> ObjectStore:
@@ -38,7 +50,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