flwr-nightly 1.13.0.dev20241022__py3-none-any.whl → 1.13.0.dev20241024__py3-none-any.whl

flwr/server/superlink/linkstate/in_memory_linkstate.py

@@ -17,22 +17,41 @@
 
 import threading
 import time
+from dataclasses import dataclass
 from logging import ERROR, WARNING
 from typing import Optional
 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.common.typing import Run, RunStatus, UserConfig
 from flwr.proto.task_pb2 import TaskIns, TaskRes  # pylint: disable=E0611
 from flwr.server.superlink.linkstate.linkstate import LinkState
 from flwr.server.utils import validate_task_ins_or_res
 
-from .utils import generate_rand_int_from_bytes, make_node_unavailable_taskres
+from .utils import (
+    generate_rand_int_from_bytes,
+    has_valid_sub_status,
+    is_valid_transition,
+    make_node_unavailable_taskres,
+)
+
+
+@dataclass
+class RunRecord:
+    """The record of a specific run, including its status and timestamps."""
+
+    run: Run
+    status: RunStatus
+    pending_at: str = ""
+    starting_at: str = ""
+    running_at: str = ""
+    finished_at: str = ""
 
 
 class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
@@ -44,8 +63,9 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
         self.node_ids: dict[int, tuple[float, float]] = {}
         self.public_key_to_node_id: dict[bytes, int] = {}
 
-        # Map run_id to (fab_id, fab_version)
-        self.run_ids: dict[int, Run] = {}
+        # Map run_id to RunRecord
+        self.run_ids: dict[int, RunRecord] = {}
+        self.contexts: dict[int, Context] = {}
         self.task_ins_store: dict[UUID, TaskIns] = {}
         self.task_res_store: dict[UUID, TaskRes] = {}
 
@@ -351,13 +371,22 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
             run_id = generate_rand_int_from_bytes(RUN_ID_NUM_BYTES)
 
             if run_id not in self.run_ids:
-                self.run_ids[run_id] = Run(
-                    run_id=run_id,
-                    fab_id=fab_id if fab_id else "",
-                    fab_version=fab_version if fab_version else "",
-                    fab_hash=fab_hash if fab_hash else "",
-                    override_config=override_config,
+                run_record = RunRecord(
+                    run=Run(
+                        run_id=run_id,
+                        fab_id=fab_id if fab_id else "",
+                        fab_version=fab_version if fab_version else "",
+                        fab_hash=fab_hash if fab_hash else "",
+                        override_config=override_config,
+                    ),
+                    status=RunStatus(
+                        status=Status.PENDING,
+                        sub_status="",
+                        details="",
+                    ),
+                    pending_at=now().isoformat(),
                 )
+                self.run_ids[run_id] = run_record
                 return run_id
         log(ERROR, "Unexpected run creation failure.")
         return 0
@@ -401,7 +430,69 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
             if run_id not in self.run_ids:
                 log(ERROR, "`run_id` is invalid")
                 return None
-            return self.run_ids[run_id]
+            return self.run_ids[run_id].run
+
+    def get_run_status(self, run_ids: set[int]) -> dict[int, RunStatus]:
+        """Retrieve the statuses for the specified runs."""
+        with self.lock:
+            return {
+                run_id: self.run_ids[run_id].status
+                for run_id in set(run_ids)
+                if run_id in self.run_ids
+            }
+
+    def update_run_status(self, run_id: int, new_status: RunStatus) -> bool:
+        """Update the status of the run with the specified `run_id`."""
+        with self.lock:
+            # Check if the run_id exists
+            if run_id not in self.run_ids:
+                log(ERROR, "`run_id` is invalid")
+                return False
+
+            # Check if the status transition is valid
+            current_status = self.run_ids[run_id].status
+            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
+            run_record = self.run_ids[run_id]
+            if new_status.status == Status.STARTING:
+                run_record.starting_at = now().isoformat()
+            elif new_status.status == Status.RUNNING:
+                run_record.running_at = now().isoformat()
+            elif new_status.status == Status.FINISHED:
+                run_record.finished_at = now().isoformat()
+            run_record.status = new_status
+            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
+
+        # Loop through all registered runs
+        for run_id, run_rec in self.run_ids.items():
+            # Break once a pending run is found
+            if run_rec.status.status == Status.PENDING:
+                pending_run_id = run_id
+                break
+
+        return pending_run_id
 
     def acknowledge_ping(self, node_id: int, ping_interval: float) -> bool:
         """Acknowledge a ping received from a node, serving as a heartbeat."""
@@ -410,3 +501,13 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
                 self.node_ids[node_id] = (time.time() + ping_interval, ping_interval)
                 return True
         return False
+
+    def get_serverapp_context(self, run_id: int) -> Optional[Context]:
+        """Get the context for the specified `run_id`."""
+        return self.contexts.get(run_id)
+
+    def set_serverapp_context(self, run_id: int, context: Context) -> None:
+        """Set the context for the specified `run_id`."""
+        if run_id not in self.run_ids:
+            raise ValueError(f"Run {run_id} not found")
+        self.contexts[run_id] = context

flwr/server/superlink/linkstate/linkstate.py

@@ -19,7 +19,8 @@ import abc
 from typing import Optional
 from uuid import UUID
 
-from flwr.common.typing import Run, UserConfig
+from flwr.common import Context
+from flwr.common.typing import Run, RunStatus, UserConfig
 from flwr.proto.task_pb2 import TaskIns, TaskRes  # pylint: disable=E0611
 
 
@@ -178,6 +179,54 @@ class LinkState(abc.ABC):  # pylint: disable=R0904
             - `fab_version`: The version of the FAB used in the specified run.
         """
 
+    @abc.abstractmethod
+    def get_run_status(self, run_ids: set[int]) -> dict[int, RunStatus]:
+        """Retrieve the statuses for the specified runs.
+
+        Parameters
+        ----------
+        run_ids : set[int]
+            A set of run identifiers for which to retrieve statuses.
+
+        Returns
+        -------
+        dict[int, RunStatus]
+            A dictionary mapping each valid run ID to its corresponding status.
+
+        Notes
+        -----
+        Only valid run IDs that exist in the State will be included in the returned
+        dictionary. If a run ID is not found, it will be omitted from the result.
+        """
+
+    @abc.abstractmethod
+    def update_run_status(self, run_id: int, new_status: RunStatus) -> bool:
+        """Update the status of the run with the specified `run_id`.
+
+        Parameters
+        ----------
+        run_id : int
+            The identifier of the run.
+        new_status : RunStatus
+            The new status to be assigned to the run.
+
+        Returns
+        -------
+        bool
+            True if the status update is successful; False otherwise.
+        """
+
+    @abc.abstractmethod
+    def get_pending_run_id(self) -> Optional[int]:
+        """Get the `run_id` of a run with `Status.PENDING` status.
+
+        Returns
+        -------
+        Optional[int]
+            The `run_id` of a `Run` that is pending to be started; None if
+            there is no Run pending.
+        """
+
     @abc.abstractmethod
     def store_server_private_public_key(
         self, private_key: bytes, public_key: bytes
@@ -222,3 +271,31 @@ class LinkState(abc.ABC):  # pylint: disable=R0904
         is_acknowledged : bool
             True if the ping is successfully acknowledged; otherwise, False.
         """
+
+    @abc.abstractmethod
+    def get_serverapp_context(self, run_id: int) -> Optional[Context]:
+        """Get the context for the specified `run_id`.
+
+        Parameters
+        ----------
+        run_id : int
+            The identifier of the run for which to retrieve the context.
+
+        Returns
+        -------
+        Optional[Context]
+            The context associated with the specified `run_id`, or `None` if no context
+            exists for the given `run_id`.
+        """
+
+    @abc.abstractmethod
+    def set_serverapp_context(self, run_id: int, context: Context) -> None:
+        """Set the context for the specified `run_id`.
+
+        Parameters
+        ----------
+        run_id : int
+            The identifier of the run for which to set the context.
+        context : Context
+            The context to be associated with the specified `run_id`.
+        """

flwr/server/superlink/linkstate/sqlite_linkstate.py

@@ -19,31 +19,41 @@
 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.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 .linkstate import LinkState
 from .utils import (
+    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,
+    is_valid_transition,
     make_node_unavailable_taskres,
 )
 
@@ -79,7 +89,21 @@ 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
+);
+"""
+
+SQL_CREATE_TABLE_CONTEXT = """
+CREATE TABLE IF NOT EXISTS context(
+    run_id                INTEGER UNIQUE,
+    context               BLOB,
+    FOREIGN KEY(run_id) REFERENCES run(run_id)
 );
 """
 
@@ -133,7 +157,7 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
         self,
         database_path: str,
     ) -> None:
-        """Initialize an SqliteState.
+        """Initialize an SqliteLinkState.
 
         Parameters
         ----------
@@ -143,6 +167,7 @@ class SqliteLinkState(LinkState):  # 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 +191,7 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
 
         # Create each table if not exists queries
         cur.execute(SQL_CREATE_TABLE_RUN)
+        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)
@@ -773,26 +799,16 @@ class SqliteLinkState(LinkState):  # 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, 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]
+            data += [now().isoformat(), "", "", "", "", ""]
+            self.query(query, tuple(data))
             return uint64_run_id
         log(ERROR, "Unexpected run creation failure.")
         return 0
@@ -868,6 +884,94 @@ class SqliteLinkState(LinkState):  # 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 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 +987,34 @@ class SqliteLinkState(LinkState):  # 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 get_valid_task_ins(self, task_id: str) -> Optional[dict[str, Any]]:
         """Check if the TaskIns exists and is valid (not expired).
 
@@ -967,7 +1099,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 +1129,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 +1155,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.")

flwr/server/superlink/linkstate/utils.py

@@ -20,9 +20,11 @@ from logging import ERROR
 from os import urandom
 from uuid import uuid4
 
-from flwr.common import log
-from flwr.common.constant import ErrorCode
+from flwr.common import Context, log, serde
+from flwr.common.constant import ErrorCode, Status, SubStatus
+from flwr.common.typing import RunStatus
 from flwr.proto.error_pb2 import Error  # pylint: disable=E0611
+from flwr.proto.message_pb2 import Context as ProtoContext  # pylint: disable=E0611
 from flwr.proto.node_pb2 import Node  # pylint: disable=E0611
 from flwr.proto.task_pb2 import Task, TaskIns, TaskRes  # pylint: disable=E0611
 
@@ -31,6 +33,17 @@ NODE_UNAVAILABLE_ERROR_REASON = (
     "It exceeds the time limit specified in its last ping."
 )
 
+VALID_RUN_STATUS_TRANSITIONS = {
+    (Status.PENDING, Status.STARTING),
+    (Status.STARTING, Status.RUNNING),
+    (Status.RUNNING, Status.FINISHED),
+}
+VALID_RUN_SUB_STATUSES = {
+    SubStatus.COMPLETED,
+    SubStatus.FAILED,
+    SubStatus.STOPPED,
+}
+
 
 def generate_rand_int_from_bytes(num_bytes: int) -> int:
     """Generate a random unsigned integer from `num_bytes` bytes."""
@@ -123,6 +136,16 @@ def convert_sint64_values_in_dict_to_uint64(
             data_dict[key] = convert_sint64_to_uint64(data_dict[key])
 
 
+def context_to_bytes(context: Context) -> bytes:
+    """Serialize `Context` to bytes."""
+    return serde.context_to_proto(context).SerializeToString()
+
+
+def context_from_bytes(context_bytes: bytes) -> Context:
+    """Deserialize `Context` from bytes."""
+    return serde.context_from_proto(ProtoContext.FromString(context_bytes))
+
+
 def make_node_unavailable_taskres(ref_taskins: TaskIns) -> TaskRes:
     """Generate a TaskRes with a node unavailable error from a TaskIns."""
     current_time = time.time()
@@ -146,3 +169,47 @@ def make_node_unavailable_taskres(ref_taskins: TaskIns) -> TaskRes:
             ),
         ),
     )
+
+
+def is_valid_transition(current_status: RunStatus, new_status: RunStatus) -> bool:
+    """Check if a transition between two run statuses is valid.
+
+    Parameters
+    ----------
+    current_status : RunStatus
+        The current status of the run.
+    new_status : RunStatus
+        The new status to transition to.
+
+    Returns
+    -------
+    bool
+        True if the transition is valid, False otherwise.
+    """
+    return (
+        current_status.status,
+        new_status.status,
+    ) in VALID_RUN_STATUS_TRANSITIONS
+
+
+def has_valid_sub_status(status: RunStatus) -> bool:
+    """Check if the 'sub_status' field of the given status is valid.
+
+    Parameters
+    ----------
+    status : RunStatus
+        The status object to be checked.
+
+    Returns
+    -------
+    bool
+        True if the status object has a valid sub-status, False otherwise.
+
+    Notes
+    -----
+    Only an empty string (i.e., "") is considered a valid sub-status for
+    non-finished statuses. The sub-status of a finished status cannot be empty.
+    """
+    if status.status == Status.FINISHED:
+        return status.sub_status in VALID_RUN_SUB_STATUSES
+    return status.sub_status == ""