init

adap · Feb 24, 2025 · 329e30f · 329e30f
1 parent 72aab3f
commit 329e30f
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 2 deletions.
diff --git a/src/py/flwr/server/superlink/linkstate/in_memory_linkstate.py b/src/py/flwr/server/superlink/linkstate/in_memory_linkstate.py
@@ -23,7 +23,7 @@
 from typing import Optional
 from uuid import UUID, uuid4
 
-from flwr.common import Context, log, now
+from flwr.common import Context, Metadata, log, now
 from flwr.common.constant import (
     MESSAGE_TTL_TOLERANCE,
     NODE_ID_NUM_BYTES,
@@ -72,6 +72,7 @@ def __init__(self) -> None:
         self.task_ins_store: dict[UUID, TaskIns] = {}
         self.task_res_store: dict[UUID, TaskRes] = {}
         self.task_ins_id_to_task_res_id: dict[UUID, UUID] = {}
+        self.in_processing_messages: dict[UUID, Metadata] = {}
 
         self.node_public_keys: set[bytes] = set()
 

diff --git a/src/py/flwr/server/superlink/linkstate/linkstate_test.py b/src/py/flwr/server/superlink/linkstate/linkstate_test.py
@@ -1224,7 +1224,7 @@ def test_initialize(self) -> None:
         result = state.query("SELECT name FROM sqlite_schema;")
 
         # Assert
-        assert len(result) == 15
+        assert len(result) == 17
 
 
 class SqliteFileBasedTest(StateTest, unittest.TestCase):

diff --git a/src/py/flwr/server/superlink/linkstate/sqlite_linkstate.py b/src/py/flwr/server/superlink/linkstate/sqlite_linkstate.py
@@ -151,6 +151,20 @@
 );
 """
 
+SQL_CREATE_TABLE_IN_PROCESSING_MESSAGES = """
+CREATE TABLE IF NOT EXISTS in_processing_messages(
+    message_id              TEXT UNIQUE,
+    run_id                  INTEGER,
+    src_node_id             INTEGER,
+    dst_node_id             INTEGER,
+    reply_to_message        TEXT,
+    group_id                TEXT,
+    ttl                     REAL,
+    message_type            TEXT,
+    FOREIGN KEY(run_id) REFERENCES run(run_id)
+);
+"""
+
 DictOrTuple = Union[tuple[Any, ...], dict[str, Any]]
 
 
@@ -198,6 +212,7 @@ def initialize(self, log_queries: bool = False) -> list[tuple[str]]:
         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_IN_PROCESSING_MESSAGES)
         cur.execute(SQL_CREATE_TABLE_NODE)
         cur.execute(SQL_CREATE_TABLE_PUBLIC_KEY)
         cur.execute(SQL_CREATE_INDEX_ONLINE_UNTIL)