flwr-nightly 1.8.0.dev20240314__py3-none-any.whl → 1.11.0.dev20240813__py3-none-any.whl

flwr/server/superlink/state/in_memory_state.py

@@ -1,4 +1,4 @@
-# Copyright 2023 Flower Labs GmbH. All Rights Reserved.
+# Copyright 2024 Flower Labs GmbH. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -15,27 +15,40 @@
 """In-memory State implementation."""
 
 
-import os
 import threading
-from datetime import datetime, timedelta
+import time
 from logging import ERROR
-from typing import Dict, List, Optional, Set
+from typing import Dict, List, Optional, Set, Tuple
 from uuid import UUID, uuid4
 
 from flwr.common import log, now
+from flwr.common.constant import NODE_ID_NUM_BYTES, RUN_ID_NUM_BYTES
+from flwr.common.typing import Run, UserConfig
 from flwr.proto.task_pb2 import TaskIns, TaskRes  # pylint: disable=E0611
 from flwr.server.superlink.state.state import State
 from flwr.server.utils import validate_task_ins_or_res
 
+from .utils import generate_rand_int_from_bytes, make_node_unavailable_taskres
 
-class InMemoryState(State):
+
+class InMemoryState(State):  # pylint: disable=R0902,R0904
     """In-memory State implementation."""
 
     def __init__(self) -> None:
-        self.node_ids: Set[int] = set()
-        self.run_ids: Set[int] = set()
+
+        # 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] = {}
+
+        # Map run_id to (fab_id, fab_version)
+        self.run_ids: Dict[int, Run] = {}
         self.task_ins_store: Dict[UUID, TaskIns] = {}
         self.task_res_store: Dict[UUID, TaskRes] = {}
+
+        self.client_public_keys: Set[bytes] = set()
+        self.server_public_key: Optional[bytes] = None
+        self.server_private_key: Optional[bytes] = None
+
         self.lock = threading.Lock()
 
     def store_task_ins(self, task_ins: TaskIns) -> Optional[UUID]:
@@ -50,15 +63,11 @@ class InMemoryState(State):
             log(ERROR, "`run_id` is invalid")
             return None
 
-        # Create task_id, created_at and ttl
+        # Create task_id
         task_id = uuid4()
-        created_at: datetime = now()
-        ttl: datetime = created_at + timedelta(hours=24)
 
         # Store TaskIns
         task_ins.task_id = str(task_id)
-        task_ins.task.created_at = created_at.isoformat()
-        task_ins.task.ttl = ttl.isoformat()
         with self.lock:
             self.task_ins_store[task_id] = task_ins
 
@@ -113,15 +122,11 @@ class InMemoryState(State):
             log(ERROR, "`run_id` is invalid")
             return None
 
-        # Create task_id, created_at and ttl
+        # Create task_id
         task_id = uuid4()
-        created_at: datetime = now()
-        ttl: datetime = created_at + timedelta(hours=24)
 
         # Store TaskRes
         task_res.task_id = str(task_id)
-        task_res.task.created_at = created_at.isoformat()
-        task_res.task.ttl = ttl.isoformat()
         with self.lock:
             self.task_res_store[task_id] = task_res
 
@@ -136,14 +141,31 @@ class InMemoryState(State):
         with self.lock:
             # Find TaskRes that were not delivered yet
             task_res_list: List[TaskRes] = []
+            replied_task_ids: Set[UUID] = set()
             for _, task_res in self.task_res_store.items():
-                if (
-                    UUID(task_res.task.ancestry[0]) in task_ids
-                    and task_res.task.delivered_at == ""
-                ):
+                reply_to = UUID(task_res.task.ancestry[0])
+                if reply_to in task_ids and task_res.task.delivered_at == "":
                     task_res_list.append(task_res)
+                    replied_task_ids.add(reply_to)
+                if limit and len(task_res_list) == limit:
+                    break
+
+            # Check if the node is offline
+            for task_id in task_ids - replied_task_ids:
                 if limit and len(task_res_list) == limit:
                     break
+                task_ins = self.task_ins_store.get(task_id)
+                if task_ins is None:
+                    continue
+                node_id = task_ins.task.consumer.node_id
+                online_until, _ = self.node_ids[node_id]
+                # Generate a TaskRes containing an error reply if the node is offline.
+                if online_until < time.time():
+                    err_taskres = make_node_unavailable_taskres(
+                        ref_taskins=task_ins,
+                    )
+                    self.task_res_store[UUID(err_taskres.task_id)] = err_taskres
+                    task_res_list.append(err_taskres)
 
             # Mark all of them as delivered
             delivered_at = now().isoformat()
@@ -189,22 +211,47 @@ class InMemoryState(State):
         """
         return len(self.task_res_store)
 
-    def create_node(self) -> int:
+    def create_node(
+        self, ping_interval: float, public_key: Optional[bytes] = None
+    ) -> int:
         """Create, store in state, and return `node_id`."""
         # Sample a random int64 as node_id
-        node_id: int = int.from_bytes(os.urandom(8), "little", signed=True)
+        node_id = generate_rand_int_from_bytes(NODE_ID_NUM_BYTES)
 
-        if node_id not in self.node_ids:
-            self.node_ids.add(node_id)
+        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
-        log(ERROR, "Unexpected node registration failure.")
-        return 0
 
-    def delete_node(self, node_id: int) -> None:
+    def delete_node(self, node_id: int, public_key: Optional[bytes] = None) -> None:
         """Delete a client node."""
-        if node_id not in self.node_ids:
-            raise ValueError(f"Node {node_id} not found")
-        self.node_ids.remove(node_id)
+        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]
+
+            del self.node_ids[node_id]
 
     def get_nodes(self, run_id: int) -> Set[int]:
         """Return all available client nodes.
@@ -214,17 +261,87 @@ class InMemoryState(State):
         If the provided `run_id` does not exist or has no matching nodes,
         an empty `Set` MUST be returned.
         """
-        if run_id not in self.run_ids:
-            return set()
-        return self.node_ids
-
-    def create_run(self) -> int:
-        """Create one run."""
+        with self.lock:
+            if run_id not in self.run_ids:
+                return set()
+            current_time = time.time()
+            return {
+                node_id
+                for node_id, (online_until, _) in self.node_ids.items()
+                if online_until > current_time
+            }
+
+    def get_node_id(self, client_public_key: bytes) -> Optional[int]:
+        """Retrieve stored `node_id` filtered by `client_public_keys`."""
+        return self.public_key_to_node_id.get(client_public_key)
+
+    def create_run(
+        self,
+        fab_id: str,
+        fab_version: str,
+        override_config: UserConfig,
+    ) -> int:
+        """Create a new run for the specified `fab_id` and `fab_version`."""
         # Sample a random int64 as run_id
-        run_id: int = int.from_bytes(os.urandom(8), "little", signed=True)
-
-        if run_id not in self.run_ids:
-            self.run_ids.add(run_id)
-            return run_id
+        with self.lock:
+            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,
+                    fab_version=fab_version,
+                    override_config=override_config,
+                )
+                return run_id
         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 state."""
+        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
+
+    def store_client_public_keys(self, public_keys: Set[bytes]) -> None:
+        """Store a set of `client_public_keys` in state."""
+        with self.lock:
+            self.client_public_keys = public_keys
+
+    def store_client_public_key(self, public_key: bytes) -> None:
+        """Store a `client_public_key` in state."""
+        with self.lock:
+            self.client_public_keys.add(public_key)
+
+    def get_client_public_keys(self) -> Set[bytes]:
+        """Retrieve all currently stored `client_public_keys` as a set."""
+        return self.client_public_keys
+
+    def get_run(self, run_id: int) -> Optional[Run]:
+        """Retrieve information about the run with the specified `run_id`."""
+        with self.lock:
+            if run_id not in self.run_ids:
+                log(ERROR, "`run_id` is invalid")
+                return None
+            return self.run_ids[run_id]
+
+    def acknowledge_ping(self, node_id: int, ping_interval: float) -> bool:
+        """Acknowledge a ping received from a node, serving as a heartbeat."""
+        with self.lock:
+            if node_id in self.node_ids:
+                self.node_ids[node_id] = (time.time() + ping_interval, ping_interval)
+                return True
+        return False