flwr 1.13.1__py3-none-any.whl → 1.15.0__py3-none-any.whl

flwr/server/superlink/fleet/rest_rere/rest_api.py

@@ -17,13 +17,12 @@
 
 from __future__ import annotations
 
-import sys
 from collections.abc import Awaitable
 from typing import Callable, TypeVar, cast
 
 from google.protobuf.message import Message as GrpcMessage
 
-from flwr.common.constant import MISSING_EXTRA_REST
+from flwr.common.exit import ExitCode, flwr_exit
 from flwr.proto.fab_pb2 import GetFabRequest, GetFabResponse  # pylint: disable=E0611
 from flwr.proto.fleet_pb2 import (  # pylint: disable=E0611
     CreateNodeRequest,
@@ -32,10 +31,10 @@ from flwr.proto.fleet_pb2 import (  # pylint: disable=E0611
     DeleteNodeResponse,
     PingRequest,
     PingResponse,
-    PullTaskInsRequest,
-    PullTaskInsResponse,
-    PushTaskResRequest,
-    PushTaskResResponse,
+    PullMessagesRequest,
+    PullMessagesResponse,
+    PushMessagesRequest,
+    PushMessagesResponse,
 )
 from flwr.proto.run_pb2 import GetRunRequest, GetRunResponse  # pylint: disable=E0611
 from flwr.server.superlink.ffs.ffs import Ffs
@@ -51,7 +50,7 @@ try:
     from starlette.responses import Response
     from starlette.routing import Route
 except ModuleNotFoundError:
-    sys.exit(MISSING_EXTRA_REST)
+    flwr_exit(ExitCode.COMMON_MISSING_EXTRA_REST)
 
 
 GrpcRequest = TypeVar("GrpcRequest", bound=GrpcMessage)
@@ -107,25 +106,24 @@ async def delete_node(request: DeleteNodeRequest) -> DeleteNodeResponse:
     return message_handler.delete_node(request=request, state=state)
 
 
-@rest_request_response(PullTaskInsRequest)
-async def pull_task_ins(request: PullTaskInsRequest) -> PullTaskInsResponse:
-    """Pull TaskIns."""
+@rest_request_response(PullMessagesRequest)
+async def pull_message(request: PullMessagesRequest) -> PullMessagesResponse:
+    """Pull PullMessages."""
     # Get state from app
     state: LinkState = cast(LinkStateFactory, app.state.STATE_FACTORY).state()
 
     # Handle message
-    return message_handler.pull_task_ins(request=request, state=state)
+    return message_handler.pull_messages(request=request, state=state)
 
 
-# Check if token is needed here
-@rest_request_response(PushTaskResRequest)
-async def push_task_res(request: PushTaskResRequest) -> PushTaskResResponse:
-    """Push TaskRes."""
+@rest_request_response(PushMessagesRequest)
+async def push_message(request: PushMessagesRequest) -> PushMessagesResponse:
+    """Pull PushMessages."""
     # Get state from app
     state: LinkState = cast(LinkStateFactory, app.state.STATE_FACTORY).state()
 
     # Handle message
-    return message_handler.push_task_res(request=request, state=state)
+    return message_handler.push_messages(request=request, state=state)
 
 
 @rest_request_response(PingRequest)
@@ -154,15 +152,18 @@ async def get_fab(request: GetFabRequest) -> GetFabResponse:
     # Get ffs from app
     ffs: Ffs = cast(FfsFactory, app.state.FFS_FACTORY).ffs()
 
+    # Get state from app
+    state: LinkState = cast(LinkStateFactory, app.state.STATE_FACTORY).state()
+
     # Handle message
-    return message_handler.get_fab(request=request, ffs=ffs)
+    return message_handler.get_fab(request=request, ffs=ffs, state=state)
 
 
 routes = [
     Route("/api/v0/fleet/create-node", create_node, methods=["POST"]),
     Route("/api/v0/fleet/delete-node", delete_node, methods=["POST"]),
-    Route("/api/v0/fleet/pull-task-ins", pull_task_ins, methods=["POST"]),
-    Route("/api/v0/fleet/push-task-res", push_task_res, methods=["POST"]),
+    Route("/api/v0/fleet/pull-messages", pull_message, methods=["POST"]),
+    Route("/api/v0/fleet/push-messages", push_message, methods=["POST"]),
     Route("/api/v0/fleet/ping", ping, methods=["POST"]),
     Route("/api/v0/fleet/get-run", get_run, methods=["POST"]),
     Route("/api/v0/fleet/get-fab", get_fab, methods=["POST"]),

flwr/server/superlink/fleet/vce/__init__.py

@@ -14,6 +14,7 @@
 # ==============================================================================
 """Fleet Simulation Engine side."""
 
+
 from .vce_api import start_vce
 
 __all__ = [

flwr/server/superlink/fleet/vce/backend/__init__.py

@@ -14,6 +14,7 @@
 # ==============================================================================
 """Simulation Engine Backends."""
 
+
 import importlib
 
 from .backend import Backend, BackendConfig

flwr/server/superlink/fleet/vce/backend/raybackend.py

@@ -14,6 +14,7 @@
 # ==============================================================================
 """Ray backend for the Fleet API using the Simulation Engine."""
 
+
 import sys
 from logging import DEBUG, ERROR
 from typing import Callable, Optional, Union

flwr/server/superlink/fleet/vce/vce_api.py

@@ -182,8 +182,8 @@ def run_api(
     f_stop: threading.Event,
 ) -> None:
     """Run the VCE."""
-    taskins_queue: "Queue[TaskIns]" = Queue()
-    taskres_queue: "Queue[TaskRes]" = Queue()
+    taskins_queue: Queue[TaskIns] = Queue()
+    taskres_queue: Queue[TaskRes] = Queue()
 
     try:
 

flwr/server/superlink/linkstate/in_memory_linkstate.py

@@ -28,6 +28,7 @@ from flwr.common.constant import (
     MESSAGE_TTL_TOLERANCE,
     NODE_ID_NUM_BYTES,
     RUN_ID_NUM_BYTES,
+    SUPERLINK_NODE_ID,
     Status,
 )
 from flwr.common.record import ConfigsRecord
@@ -62,6 +63,7 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
         # Map node_id to (online_until, ping_interval)
         self.node_ids: dict[int, tuple[float, float]] = {}
         self.public_key_to_node_id: dict[bytes, int] = {}
+        self.node_id_to_public_key: dict[int, bytes] = {}
 
         # Map run_id to RunRecord
         self.run_ids: dict[int, RunRecord] = {}
@@ -72,8 +74,6 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
         self.task_ins_id_to_task_res_id: dict[UUID, UUID] = {}
 
         self.node_public_keys: set[bytes] = set()
-        self.server_public_key: Optional[bytes] = None
-        self.server_private_key: Optional[bytes] = None
 
         self.lock = threading.RLock()
 
@@ -89,7 +89,7 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
             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:
+        if task_ins.task.producer.node_id != SUPERLINK_NODE_ID:
             log(
                 ERROR,
                 "Invalid source node ID for TaskIns: %s",
@@ -97,14 +97,13 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
             )
             return None
         # Validate destination node ID
-        if not task_ins.task.consumer.anonymous:
-            if task_ins.task.consumer.node_id not in self.node_ids:
-                log(
-                    ERROR,
-                    "Invalid destination node ID for TaskIns: %s",
-                    task_ins.task.consumer.node_id,
-                )
-                return None
+        if task_ins.task.consumer.node_id not in self.node_ids:
+            log(
+                ERROR,
+                "Invalid destination node ID for TaskIns: %s",
+                task_ins.task.consumer.node_id,
+            )
+            return None
 
         # Create task_id
         task_id = uuid4()
@@ -117,9 +116,7 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
         # Return the new task_id
         return task_id
 
-    def get_task_ins(
-        self, node_id: Optional[int], limit: Optional[int]
-    ) -> list[TaskIns]:
+    def get_task_ins(self, node_id: int, limit: Optional[int]) -> list[TaskIns]:
         """Get all TaskIns that have not been delivered yet."""
         if limit is not None and limit < 1:
             raise AssertionError("`limit` must be >= 1")
@@ -129,17 +126,8 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
         current_time = time.time()
         with self.lock:
             for _, task_ins in self.task_ins_store.items():
-                # pylint: disable=too-many-boolean-expressions
                 if (
-                    node_id is not None  # Not anonymous
-                    and task_ins.task.consumer.anonymous is False
-                    and task_ins.task.consumer.node_id == node_id
-                    and task_ins.task.delivered_at == ""
-                    and task_ins.task.created_at + task_ins.task.ttl > current_time
-                ) or (
-                    node_id is None  # Anonymous
-                    and task_ins.task.consumer.anonymous is True
-                    and task_ins.task.consumer.node_id == 0
+                    task_ins.task.consumer.node_id == node_id
                     and task_ins.task.delivered_at == ""
                     and task_ins.task.created_at + task_ins.task.ttl > current_time
                 ):
@@ -173,9 +161,6 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
             if (
                 task_ins
                 and task_res
-                and not (
-                    task_ins.task.consumer.anonymous or task_res.task.producer.anonymous
-                )
                 and task_ins.task.consumer.node_id != task_res.task.producer.node_id
             ):
                 return None
@@ -265,41 +250,15 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
             for task_res in task_res_found:
                 task_res.task.delivered_at = delivered_at
 
-            # Cleanup
-            self._force_delete_tasks_by_ids(set(ret.keys()))
-
         return list(ret.values())
 
-    def delete_tasks(self, task_ids: set[UUID]) -> None:
-        """Delete all delivered TaskIns/TaskRes pairs."""
-        task_ins_to_be_deleted: set[UUID] = set()
-        task_res_to_be_deleted: set[UUID] = set()
-
-        with self.lock:
-            for task_ins_id in task_ids:
-                # Find the task_id of the matching task_res
-                for task_res_id, task_res in self.task_res_store.items():
-                    if UUID(task_res.task.ancestry[0]) != task_ins_id:
-                        continue
-                    if task_res.task.delivered_at == "":
-                        continue
-
-                    task_ins_to_be_deleted.add(task_ins_id)
-                    task_res_to_be_deleted.add(task_res_id)
-
-            for task_id in task_ins_to_be_deleted:
-                del self.task_ins_store[task_id]
-                del self.task_ins_id_to_task_res_id[task_id]
-            for task_id in task_res_to_be_deleted:
-                del self.task_res_store[task_id]
-
-    def _force_delete_tasks_by_ids(self, task_ids: set[UUID]) -> None:
-        """Delete tasks based on a set of TaskIns IDs."""
-        if not task_ids:
+    def delete_tasks(self, task_ins_ids: set[UUID]) -> None:
+        """Delete TaskIns/TaskRes pairs based on provided TaskIns IDs."""
+        if not task_ins_ids:
             return
 
         with self.lock:
-            for task_id in task_ids:
+            for task_id in task_ins_ids:
                 # Delete TaskIns
                 if task_id in self.task_ins_store:
                     del self.task_ins_store[task_id]
@@ -308,6 +267,16 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
                     task_res_id = self.task_ins_id_to_task_res_id.pop(task_id)
                     del self.task_res_store[task_res_id]
 
+    def get_task_ids_from_run_id(self, run_id: int) -> set[UUID]:
+        """Get all TaskIns IDs for the given run_id."""
+        task_id_list: set[UUID] = set()
+        with self.lock:
+            for task_id, task_ins in self.task_ins_store.items():
+                if task_ins.run_id == run_id:
+                    task_id_list.add(task_id)
+
+        return task_id_list
+
     def num_task_ins(self) -> int:
         """Calculate the number of task_ins in store.
 
@@ -322,45 +291,30 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
         """
         return len(self.task_res_store)
 
-    def create_node(
-        self, ping_interval: float, public_key: Optional[bytes] = None
-    ) -> int:
+    def create_node(self, ping_interval: float) -> int:
         """Create, store in the link state, and return `node_id`."""
         # Sample a random int64 as node_id
-        node_id = generate_rand_int_from_bytes(NODE_ID_NUM_BYTES)
+        node_id = generate_rand_int_from_bytes(
+            NODE_ID_NUM_BYTES, exclude=[SUPERLINK_NODE_ID, 0]
+        )
 
         with self.lock:
             if node_id in self.node_ids:
                 log(ERROR, "Unexpected node registration failure.")
                 return 0
 
-            if public_key is not None:
-                if (
-                    public_key in self.public_key_to_node_id
-                    or node_id in self.public_key_to_node_id.values()
-                ):
-                    log(ERROR, "Unexpected node registration failure.")
-                    return 0
-
-                self.public_key_to_node_id[public_key] = node_id
-
             self.node_ids[node_id] = (time.time() + ping_interval, ping_interval)
             return node_id
 
-    def delete_node(self, node_id: int, public_key: Optional[bytes] = None) -> None:
+    def delete_node(self, node_id: int) -> None:
         """Delete a node."""
         with self.lock:
             if node_id not in self.node_ids:
                 raise ValueError(f"Node {node_id} not found")
 
-            if public_key is not None:
-                if (
-                    public_key not in self.public_key_to_node_id
-                    or node_id not in self.public_key_to_node_id.values()
-                ):
-                    raise ValueError("Public key or node_id not found")
-
-                del self.public_key_to_node_id[public_key]
+            # Remove node ID <> public key mappings
+            if pk := self.node_id_to_public_key.pop(node_id, None):
+                del self.public_key_to_node_id[pk]
 
             del self.node_ids[node_id]
 
@@ -382,6 +336,26 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
                 if online_until > current_time
             }
 
+    def set_node_public_key(self, node_id: int, public_key: bytes) -> None:
+        """Set `public_key` for the specified `node_id`."""
+        with self.lock:
+            if node_id not in self.node_ids:
+                raise ValueError(f"Node {node_id} not found")
+
+            if public_key in self.public_key_to_node_id:
+                raise ValueError("Public key already in use")
+
+            self.public_key_to_node_id[public_key] = node_id
+            self.node_id_to_public_key[node_id] = public_key
+
+    def get_node_public_key(self, node_id: int) -> Optional[bytes]:
+        """Get `public_key` for the specified `node_id`."""
+        with self.lock:
+            if node_id not in self.node_ids:
+                raise ValueError(f"Node {node_id} not found")
+
+            return self.node_id_to_public_key.get(node_id)
+
     def get_node_id(self, node_public_key: bytes) -> Optional[int]:
         """Retrieve stored `node_id` filtered by `node_public_keys`."""
         return self.public_key_to_node_id.get(node_public_key)
@@ -427,29 +401,15 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
         log(ERROR, "Unexpected run creation failure.")
         return 0
 
-    def store_server_private_public_key(
-        self, private_key: bytes, public_key: bytes
-    ) -> None:
-        """Store `server_private_key` and `server_public_key` in the link state."""
+    def clear_supernode_auth_keys(self) -> None:
+        """Clear stored `node_public_keys` in the link state if any."""
         with self.lock:
-            if self.server_private_key is None and self.server_public_key is None:
-                self.server_private_key = private_key
-                self.server_public_key = public_key
-            else:
-                raise RuntimeError("Server private and public key already set")
-
-    def get_server_private_key(self) -> Optional[bytes]:
-        """Retrieve `server_private_key` in urlsafe bytes."""
-        return self.server_private_key
-
-    def get_server_public_key(self) -> Optional[bytes]:
-        """Retrieve `server_public_key` in urlsafe bytes."""
-        return self.server_public_key
+            self.node_public_keys.clear()
 
     def store_node_public_keys(self, public_keys: set[bytes]) -> None:
         """Store a set of `node_public_keys` in the link state."""
         with self.lock:
-            self.node_public_keys = public_keys
+            self.node_public_keys.update(public_keys)
 
     def store_node_public_key(self, public_key: bytes) -> None:
         """Store a `node_public_key` in the link state."""
@@ -458,7 +418,8 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
 
     def get_node_public_keys(self) -> set[bytes]:
         """Retrieve all currently stored `node_public_keys` as a set."""
-        return self.node_public_keys
+        with self.lock:
+            return self.node_public_keys.copy()
 
     def get_run_ids(self) -> set[int]:
         """Retrieve all run IDs."""

flwr/server/superlink/linkstate/linkstate.py

@@ -40,20 +40,14 @@ class LinkState(abc.ABC):  # pylint: disable=R0904
 
         Constraints
         -----------
-        If `task_ins.task.consumer.anonymous` is `True`, then
-        `task_ins.task.consumer.node_id` MUST NOT be set (equal 0).
-
-        If `task_ins.task.consumer.anonymous` is `False`, then
-        `task_ins.task.consumer.node_id` MUST be set (not 0)
+        `task_ins.task.consumer.node_id` MUST be set (not constant.DRIVER_NODE_ID)
 
         If `task_ins.run_id` is invalid, then
         storing the `task_ins` MUST fail.
         """
 
     @abc.abstractmethod
-    def get_task_ins(
-        self, node_id: Optional[int], limit: Optional[int]
-    ) -> list[TaskIns]:
+    def get_task_ins(self, node_id: int, limit: Optional[int]) -> list[TaskIns]:
         """Get TaskIns optionally filtered by node_id.
 
         Usually, the Fleet API calls this for Nodes planning to work on one or more
@@ -61,15 +55,11 @@ class LinkState(abc.ABC):  # pylint: disable=R0904
 
         Constraints
         -----------
-        If `node_id` is not `None`, retrieve all TaskIns where
+        Retrieve all TaskIns where
 
             1. the `task_ins.task.consumer.node_id` equals `node_id` AND
-            2. the `task_ins.task.consumer.anonymous` equals `False` AND
-            3. the `task_ins.task.delivered_at` equals `""`.
+            2. the `task_ins.task.delivered_at` equals `""`.
 
-        If `node_id` is `None`, retrieve all TaskIns where the
-        `task_ins.task.consumer.node_id` equals `0` and
-        `task_ins.task.consumer.anonymous` is set to `True`.
 
         If `delivered_at` MUST BE set (not `""`) otherwise the TaskIns MUST not be in
         the result.
@@ -89,11 +79,8 @@ class LinkState(abc.ABC):  # pylint: disable=R0904
 
         Constraints
         -----------
-        If `task_res.task.consumer.anonymous` is `True`, then
-        `task_res.task.consumer.node_id` MUST NOT be set (equal 0).
 
-        If `task_res.task.consumer.anonymous` is `False`, then
-        `task_res.task.consumer.node_id` MUST be set (not 0)
+        `task_res.task.consumer.node_id` MUST be set (not constant.DRIVER_NODE_ID)
 
         If `task_res.run_id` is invalid, then
         storing the `task_res` MUST fail.
@@ -139,17 +126,26 @@ class LinkState(abc.ABC):  # pylint: disable=R0904
         """
 
     @abc.abstractmethod
-    def delete_tasks(self, task_ids: set[UUID]) -> None:
-        """Delete all delivered TaskIns/TaskRes pairs."""
+    def delete_tasks(self, task_ins_ids: set[UUID]) -> None:
+        """Delete TaskIns/TaskRes pairs based on provided TaskIns IDs.
+
+        Parameters
+        ----------
+        task_ins_ids : set[UUID]
+            A set of TaskIns IDs. For each ID in the set, the corresponding
+            TaskIns and its associated TaskRes will be deleted.
+        """
+
+    @abc.abstractmethod
+    def get_task_ids_from_run_id(self, run_id: int) -> set[UUID]:
+        """Get all TaskIns IDs for the given run_id."""
 
     @abc.abstractmethod
-    def create_node(
-        self, ping_interval: float, public_key: Optional[bytes] = None
-    ) -> int:
+    def create_node(self, ping_interval: float) -> int:
         """Create, store in the link state, and return `node_id`."""
 
     @abc.abstractmethod
-    def delete_node(self, node_id: int, public_key: Optional[bytes] = None) -> None:
+    def delete_node(self, node_id: int) -> None:
         """Remove `node_id` from the link state."""
 
     @abc.abstractmethod
@@ -162,6 +158,14 @@ class LinkState(abc.ABC):  # pylint: disable=R0904
         an empty `Set` MUST be returned.
         """
 
+    @abc.abstractmethod
+    def set_node_public_key(self, node_id: int, public_key: bytes) -> None:
+        """Set `public_key` for the specified `node_id`."""
+
+    @abc.abstractmethod
+    def get_node_public_key(self, node_id: int) -> Optional[bytes]:
+        """Get `public_key` for the specified `node_id`."""
+
     @abc.abstractmethod
     def get_node_id(self, node_public_key: bytes) -> Optional[int]:
         """Retrieve stored `node_id` filtered by `node_public_keys`."""
@@ -260,18 +264,8 @@ class LinkState(abc.ABC):  # pylint: disable=R0904
         """
 
     @abc.abstractmethod
-    def store_server_private_public_key(
-        self, private_key: bytes, public_key: bytes
-    ) -> None:
-        """Store `server_private_key` and `server_public_key` in the link state."""
-
-    @abc.abstractmethod
-    def get_server_private_key(self) -> Optional[bytes]:
-        """Retrieve `server_private_key` in urlsafe bytes."""
-
-    @abc.abstractmethod
-    def get_server_public_key(self) -> Optional[bytes]:
-        """Retrieve `server_public_key` in urlsafe bytes."""
+    def clear_supernode_auth_keys(self) -> None:
+        """Clear stored `node_public_keys` in the link state if any."""
 
     @abc.abstractmethod
     def store_node_public_keys(self, public_keys: set[bytes]) -> None: