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

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

@@ -12,39 +12,51 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""SQLite based implemenation of server state."""
+"""SQLite based implemenation of the link state."""
 
 # pylint: disable=too-many-lines
 
 import json
 import re
 import sqlite3
+import threading
 import time
 from collections.abc import Sequence
 from logging import DEBUG, ERROR, WARNING
 from typing import Any, Optional, Union, cast
 from uuid import UUID, uuid4
 
-from flwr.common import log, now
+from flwr.common import Context, log, now
 from flwr.common.constant import (
     MESSAGE_TTL_TOLERANCE,
     NODE_ID_NUM_BYTES,
     RUN_ID_NUM_BYTES,
+    Status,
 )
-from flwr.common.typing import Run, UserConfig
-from flwr.proto.node_pb2 import Node  # pylint: disable=E0611
-from flwr.proto.recordset_pb2 import RecordSet  # pylint: disable=E0611
-from flwr.proto.task_pb2 import Task, TaskIns, TaskRes  # pylint: disable=E0611
+from flwr.common.record import ConfigsRecord
+from flwr.common.typing import Run, RunStatus, UserConfig
+
+# pylint: disable=E0611
+from flwr.proto.node_pb2 import Node
+from flwr.proto.recordset_pb2 import RecordSet as ProtoRecordSet
+from flwr.proto.task_pb2 import Task, TaskIns, TaskRes
+
+# pylint: enable=E0611
 from flwr.server.utils.validator import validate_task_ins_or_res
 
-from .state import State
+from .linkstate import LinkState
 from .utils import (
+    configsrecord_from_bytes,
+    configsrecord_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,
-    make_node_unavailable_taskres,
+    has_valid_sub_status,
+    is_valid_transition,
 )
 
 SQL_CREATE_TABLE_NODE = """
@@ -79,7 +91,33 @@ CREATE TABLE IF NOT EXISTS run(
     fab_id                TEXT,
     fab_version           TEXT,
     fab_hash              TEXT,
-    override_config       TEXT
+    override_config       TEXT,
+    pending_at            TEXT,
+    starting_at           TEXT,
+    running_at            TEXT,
+    finished_at           TEXT,
+    sub_status            TEXT,
+    details               TEXT,
+    federation_options    BLOB
+);
+"""
+
+SQL_CREATE_TABLE_LOGS = """
+CREATE TABLE IF NOT EXISTS logs (
+    timestamp             REAL,
+    run_id                INTEGER,
+    node_id               INTEGER,
+    log                   TEXT,
+    PRIMARY KEY (timestamp, run_id, node_id),
+    FOREIGN KEY (run_id) REFERENCES run(run_id)
+);
+"""
+
+SQL_CREATE_TABLE_CONTEXT = """
+CREATE TABLE IF NOT EXISTS context(
+    run_id                INTEGER UNIQUE,
+    context               BLOB,
+    FOREIGN KEY(run_id) REFERENCES run(run_id)
 );
 """
 
@@ -126,14 +164,14 @@ CREATE TABLE IF NOT EXISTS task_res(
 DictOrTuple = Union[tuple[Any, ...], dict[str, Any]]
 
 
-class SqliteState(State):  # pylint: disable=R0904
-    """SQLite-based state implementation."""
+class SqliteLinkState(LinkState):  # pylint: disable=R0904
+    """SQLite-based LinkState implementation."""
 
     def __init__(
         self,
         database_path: str,
     ) -> None:
-        """Initialize an SqliteState.
+        """Initialize an SqliteLinkState.
 
         Parameters
         ----------
@@ -143,6 +181,7 @@ class SqliteState(State):  # pylint: disable=R0904
         """
         self.database_path = database_path
         self.conn: Optional[sqlite3.Connection] = None
+        self.lock = threading.RLock()
 
     def initialize(self, log_queries: bool = False) -> list[tuple[str]]:
         """Create tables if they don't exist yet.
@@ -166,6 +205,8 @@ class SqliteState(State):  # pylint: disable=R0904
 
         # 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_TASK_INS)
         cur.execute(SQL_CREATE_TABLE_TASK_RES)
         cur.execute(SQL_CREATE_TABLE_NODE)
@@ -183,7 +224,7 @@ class SqliteState(State):  # pylint: disable=R0904
     ) -> list[dict[str, Any]]:
         """Execute a SQL query."""
         if self.conn is None:
-            raise AttributeError("State is not initialized.")
+            raise AttributeError("LinkState is not initialized.")
 
         if data is None:
             data = []
@@ -214,11 +255,11 @@ class SqliteState(State):  # pylint: disable=R0904
     def store_task_ins(self, task_ins: TaskIns) -> Optional[UUID]:
         """Store one TaskIns.
 
-        Usually, the Driver API calls this to schedule instructions.
+        Usually, the ServerAppIo API calls this to schedule instructions.
 
-        Stores the value of the task_ins in the state and, if successful, returns the
-        task_id (UUID) of the task_ins. If, for any reason, storing the task_ins fails,
-        `None` is returned.
+        Stores the value of the task_ins in the link state and, if successful,
+        returns the task_id (UUID) of the task_ins. If, for any reason, storing
+        the task_ins fails, `None` is returned.
 
         Constraints
         -----------
@@ -233,7 +274,6 @@ class SqliteState(State):  # pylint: disable=R0904
         if any(errors):
             log(ERROR, errors)
             return None
-
         # Create task_id
         task_id = uuid4()
 
@@ -246,16 +286,36 @@ class SqliteState(State):  # pylint: disable=R0904
             data[0], ["run_id", "producer_node_id", "consumer_node_id"]
         )
 
+        # Validate run_id
+        query = "SELECT run_id FROM run WHERE run_id = ?;"
+        if not self.query(query, (data[0]["run_id"],)):
+            log(ERROR, "Invalid run ID for TaskIns: %s", task_ins.run_id)
+            return None
+        # Validate source node ID
+        if task_ins.task.producer.node_id != 0:
+            log(
+                ERROR,
+                "Invalid source node ID for TaskIns: %s",
+                task_ins.task.producer.node_id,
+            )
+            return None
+        # Validate destination node ID
+        query = "SELECT node_id FROM node WHERE node_id = ?;"
+        if not task_ins.task.consumer.anonymous:
+            if not self.query(query, (data[0]["consumer_node_id"],)):
+                log(
+                    ERROR,
+                    "Invalid destination node ID for TaskIns: %s",
+                    task_ins.task.consumer.node_id,
+                )
+                return None
+
         columns = ", ".join([f":{key}" for key in data[0]])
         query = f"INSERT INTO task_ins VALUES({columns});"
 
         # Only invalid run_id can trigger IntegrityError.
         # This may need to be changed in the future version with more integrity checks.
-        try:
-            self.query(query, data)
-        except sqlite3.IntegrityError:
-            log(ERROR, "`run` is invalid")
-            return None
+        self.query(query, data)
 
         return task_id
 
@@ -452,8 +512,8 @@ class SqliteState(State):  # pylint: disable=R0904
     def get_task_res(self, task_ids: set[UUID]) -> list[TaskRes]:
         """Get TaskRes for task_ids.
 
-        Usually, the Driver API calls this method to get results for instructions it has
-        previously scheduled.
+        Usually, the ServerAppIo API calls this method to get results for instructions
+        it has previously scheduled.
 
         Retrieves all TaskRes for the given `task_ids` and returns and empty list if
         none could be found.
@@ -579,20 +639,6 @@ class SqliteState(State):  # pylint: disable=R0904
         data = {f"id_{i}": str(node_id) for i, node_id in enumerate(offline_node_ids)}
         task_ins_rows = self.query(query, data)
 
-        # Make TaskRes containing node unavailabe error
-        for row in task_ins_rows:
-            for row in rows:
-                # Convert values from sint64 to uint64
-                convert_sint64_values_in_dict_to_uint64(
-                    row, ["run_id", "producer_node_id", "consumer_node_id"]
-                )
-
-            task_ins = dict_to_task_ins(row)
-            err_taskres = make_node_unavailable_taskres(
-                ref_taskins=task_ins,
-            )
-            result.append(err_taskres)
-
         return result
 
     def num_task_ins(self) -> int:
@@ -645,7 +691,7 @@ class SqliteState(State):  # pylint: disable=R0904
         """
 
         if self.conn is None:
-            raise AttributeError("State not intitialized")
+            raise AttributeError("LinkState not intitialized")
 
         with self.conn:
             self.conn.execute(query_1, data)
@@ -656,7 +702,7 @@ class SqliteState(State):  # pylint: disable=R0904
     def create_node(
         self, ping_interval: float, public_key: Optional[bytes] = None
     ) -> int:
-        """Create, store in state, and return `node_id`."""
+        """Create, store in the link state, and return `node_id`."""
         # Sample a random uint64 as node_id
         uint64_node_id = generate_rand_int_from_bytes(NODE_ID_NUM_BYTES)
 
@@ -706,7 +752,7 @@ class SqliteState(State):  # pylint: disable=R0904
             params += (public_key,)  # type: ignore
 
         if self.conn is None:
-            raise AttributeError("State is not initialized.")
+            raise AttributeError("LinkState is not initialized.")
 
         try:
             with self.conn:
@@ -753,12 +799,14 @@ class SqliteState(State):  # pylint: disable=R0904
             return uint64_node_id
         return None
 
+    # pylint: disable=too-many-arguments,too-many-positional-arguments
     def create_run(
         self,
         fab_id: Optional[str],
         fab_version: Optional[str],
         fab_hash: Optional[str],
         override_config: UserConfig,
+        federation_options: ConfigsRecord,
     ) -> int:
         """Create a new run for the specified `fab_id` and `fab_version`."""
         # Sample a random int64 as run_id
@@ -773,26 +821,30 @@ class SqliteState(State):  # pylint: disable=R0904
         if self.query(query, (sint64_run_id,))[0]["COUNT(*)"] == 0:
             query = (
                 "INSERT INTO run "
-                "(run_id, fab_id, fab_version, fab_hash, override_config)"
-                "VALUES (?, ?, ?, ?, ?);"
+                "(run_id, fab_id, fab_version, fab_hash, override_config, "
+                "federation_options, pending_at, starting_at, running_at, finished_at, "
+                "sub_status, details) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?);"
             )
             if fab_hash:
-                self.query(
-                    query,
-                    (sint64_run_id, "", "", fab_hash, json.dumps(override_config)),
-                )
-            else:
-                self.query(
-                    query,
-                    (
-                        sint64_run_id,
-                        fab_id,
-                        fab_version,
-                        "",
-                        json.dumps(override_config),
-                    ),
-                )
-            # Note: we need to return the uint64 value of the run_id
+                fab_id, fab_version = "", ""
+            override_config_json = json.dumps(override_config)
+            data = [
+                sint64_run_id,
+                fab_id,
+                fab_version,
+                fab_hash,
+                override_config_json,
+                configsrecord_to_bytes(federation_options),
+            ]
+            data += [
+                now().isoformat(),
+                "",
+                "",
+                "",
+                "",
+                "",
+            ]
+            self.query(query, tuple(data))
             return uint64_run_id
         log(ERROR, "Unexpected run creation failure.")
         return 0
@@ -800,7 +852,7 @@ class SqliteState(State):  # pylint: disable=R0904
     def store_server_private_public_key(
         self, private_key: bytes, public_key: bytes
     ) -> None:
-        """Store `server_private_key` and `server_public_key` in state."""
+        """Store `server_private_key` and `server_public_key` in the link state."""
         query = "SELECT COUNT(*) FROM credential"
         count = self.query(query)[0]["COUNT(*)"]
         if count < 1:
@@ -833,13 +885,13 @@ class SqliteState(State):  # pylint: disable=R0904
         return public_key
 
     def store_node_public_keys(self, public_keys: set[bytes]) -> None:
-        """Store a set of `node_public_keys` in state."""
+        """Store a set of `node_public_keys` in the link state."""
         query = "INSERT INTO public_key (public_key) VALUES (?)"
         data = [(key,) for key in public_keys]
         self.query(query, data)
 
     def store_node_public_key(self, public_key: bytes) -> None:
-        """Store a `node_public_key` in state."""
+        """Store a `node_public_key` in the link state."""
         query = "INSERT INTO public_key (public_key) VALUES (:public_key)"
         self.query(query, {"public_key": public_key})
 
@@ -850,6 +902,12 @@ class SqliteState(State):  # pylint: disable=R0904
         result: set[bytes] = {row["public_key"] for row in rows}
         return result
 
+    def get_run_ids(self) -> set[int]:
+        """Retrieve all run IDs."""
+        query = "SELECT run_id FROM run;"
+        rows = self.query(query)
+        return {convert_sint64_to_uint64(row["run_id"]) for row in rows}
+
     def get_run(self, run_id: int) -> Optional[Run]:
         """Retrieve information about the run with the specified `run_id`."""
         # Convert the uint64 value to sint64 for SQLite
@@ -868,6 +926,109 @@ class SqliteState(State):  # pylint: disable=R0904
         log(ERROR, "`run_id` does not exist.")
         return None
 
+    def get_run_status(self, run_ids: set[int]) -> dict[int, RunStatus]:
+        """Retrieve the statuses for the specified runs."""
+        # Convert the uint64 value to sint64 for SQLite
+        sint64_run_ids = (convert_uint64_to_sint64(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(
+                status=determine_run_status(row),
+                sub_status=row["sub_status"],
+                details=row["details"],
+            )
+            for row in rows
+        }
+
+    def update_run_status(self, run_id: int, new_status: RunStatus) -> bool:
+        """Update the status of the run with the specified `run_id`."""
+        # Convert the uint64 value to sint64 for SQLite
+        sint64_run_id = convert_uint64_to_sint64(run_id)
+        query = "SELECT * FROM run WHERE run_id = ?;"
+        rows = self.query(query, (sint64_run_id,))
+
+        # Check if the run_id exists
+        if not rows:
+            log(ERROR, "`run_id` is invalid")
+            return False
+
+        # Check if the status transition is valid
+        row = rows[0]
+        current_status = RunStatus(
+            status=determine_run_status(row),
+            sub_status=row["sub_status"],
+            details=row["details"],
+        )
+        if not is_valid_transition(current_status, new_status):
+            log(
+                ERROR,
+                'Invalid status transition: from "%s" to "%s"',
+                current_status.status,
+                new_status.status,
+            )
+            return False
+
+        # Check if the sub-status is valid
+        if not has_valid_sub_status(current_status):
+            log(
+                ERROR,
+                'Invalid sub-status "%s" for status "%s"',
+                current_status.sub_status,
+                current_status.status,
+            )
+            return False
+
+        # Update the status
+        query = "UPDATE run SET %s= ?, sub_status = ?, details = ? "
+        query += "WHERE run_id = ?;"
+
+        timestamp_fld = ""
+        if new_status.status == Status.STARTING:
+            timestamp_fld = "starting_at"
+        elif new_status.status == Status.RUNNING:
+            timestamp_fld = "running_at"
+        elif new_status.status == Status.FINISHED:
+            timestamp_fld = "finished_at"
+
+        data = (
+            now().isoformat(),
+            new_status.sub_status,
+            new_status.details,
+            sint64_run_id,
+        )
+        self.query(query % timestamp_fld, data)
+        return True
+
+    def get_pending_run_id(self) -> Optional[int]:
+        """Get the `run_id` of a run with `Status.PENDING` status, if any."""
+        pending_run_id = None
+
+        # Fetch all runs with unset `starting_at` (i.e. they are in PENDING status)
+        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"])
+
+        return pending_run_id
+
+    def get_federation_options(self, run_id: int) -> Optional[ConfigsRecord]:
+        """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)
+        query = "SELECT federation_options FROM run WHERE run_id = ?;"
+        rows = self.query(query, (sint64_run_id,))
+
+        # Check if the run_id exists
+        if not rows:
+            log(ERROR, "`run_id` is invalid")
+            return None
+
+        row = rows[0]
+        return configsrecord_from_bytes(row["federation_options"])
+
     def acknowledge_ping(self, node_id: int, ping_interval: float) -> bool:
         """Acknowledge a ping received from a node, serving as a heartbeat."""
         sint64_node_id = convert_uint64_to_sint64(node_id)
@@ -883,6 +1044,72 @@ class SqliteState(State):  # pylint: disable=R0904
             log(ERROR, "`node_id` does not exist.")
             return False
 
+    def get_serverapp_context(self, run_id: int) -> Optional[Context]:
+        """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),))
+        context = context_from_bytes(rows[0]["context"]) if rows else None
+        return context
+
+    def set_serverapp_context(self, run_id: int, context: Context) -> None:
+        """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)
+
+        # Check if any existing Context assigned to the run_id
+        query = "SELECT COUNT(*) FROM context WHERE run_id = ?;"
+        if self.query(query, (sint_run_id,))[0]["COUNT(*)"] > 0:
+            # Update context
+            query = "UPDATE context SET context = ? WHERE run_id = ?;"
+            self.query(query, (context_bytes, sint_run_id))
+        else:
+            try:
+                # Store context
+                query = "INSERT INTO context (run_id, context) VALUES (?, ?);"
+                self.query(query, (sint_run_id, context_bytes))
+            except sqlite3.IntegrityError:
+                raise ValueError(f"Run {run_id} not found") from None
+
+    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)
+
+        # Store log
+        try:
+            query = """
+                INSERT INTO logs (timestamp, run_id, node_id, log) VALUES (?, ?, ?, ?);
+            """
+            self.query(query, (now().timestamp(), sint64_run_id, 0, log_message))
+        except sqlite3.IntegrityError:
+            raise ValueError(f"Run {run_id} not found") from None
+
+    def get_serverapp_log(
+        self, run_id: int, after_timestamp: Optional[float]
+    ) -> 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)
+
+        # Check if the run_id exists
+        query = "SELECT run_id FROM run WHERE run_id = ?;"
+        if not self.query(query, (sint64_run_id,)):
+            raise ValueError(f"Run {run_id} not found")
+
+        # Retrieve logs
+        if after_timestamp is None:
+            after_timestamp = 0.0
+        query = """
+            SELECT log, timestamp FROM logs
+            WHERE run_id = ? AND node_id = ? AND timestamp > ?;
+        """
+        rows = self.query(query, (sint64_run_id, 0, after_timestamp))
+        rows.sort(key=lambda x: x["timestamp"])
+        latest_timestamp = rows[-1]["timestamp"] if rows else 0.0
+        return "".join(row["log"] for row in rows), latest_timestamp
+
     def get_valid_task_ins(self, task_id: str) -> Optional[dict[str, Any]]:
         """Check if the TaskIns exists and is valid (not expired).
 
@@ -967,7 +1194,7 @@ def task_res_to_dict(task_msg: TaskRes) -> dict[str, Any]:
 
 def dict_to_task_ins(task_dict: dict[str, Any]) -> TaskIns:
     """Turn task_dict into protobuf message."""
-    recordset = RecordSet()
+    recordset = ProtoRecordSet()
     recordset.ParseFromString(task_dict["recordset"])
 
     result = TaskIns(
@@ -997,7 +1224,7 @@ def dict_to_task_ins(task_dict: dict[str, Any]) -> TaskIns:
 
 def dict_to_task_res(task_dict: dict[str, Any]) -> TaskRes:
     """Turn task_dict into protobuf message."""
-    recordset = RecordSet()
+    recordset = ProtoRecordSet()
     recordset.ParseFromString(task_dict["recordset"])
 
     result = TaskRes(
@@ -1023,3 +1250,17 @@ def dict_to_task_res(task_dict: dict[str, Any]) -> TaskRes:
         ),
     )
     return result
+
+
+def determine_run_status(row: dict[str, Any]) -> str:
+    """Determine the status of the run based on timestamp fields."""
+    if row["pending_at"]:
+        if row["starting_at"]:
+            if row["running_at"]:
+                if row["finished_at"]:
+                    return Status.FINISHED
+                return Status.RUNNING
+            return Status.STARTING
+        return Status.PENDING
+    run_id = convert_sint64_to_uint64(row["run_id"])
+    raise sqlite3.IntegrityError(f"The run {run_id} does not have a valid status.")