feat(framework) Add override config to Run (#3730)

src/py/flwr/client/app.py

@@ -371,7 +371,7 @@ def _on_backoff(retry_state: RetryState) -> None:
                             run_info[run_id] = get_run(run_id)
                         # If get_run is None, i.e., in grpc-bidi mode
                         else:
-                            run_info[run_id] = Run(run_id, "", "")
+                            run_info[run_id] = Run(run_id, "", "", {})
 
                     # Register context for this run
                     node_state.register_context(run_id=run_id)

src/py/flwr/client/grpc_rere_client/connection.py

@@ -280,6 +280,7 @@ def get_run(run_id: int) -> Run:
             run_id,
             get_run_response.run.fab_id,
             get_run_response.run.fab_version,
+            dict(get_run_response.run.override_config.items()),
         )
 
     try:

src/py/flwr/client/rest_client/connection.py

@@ -352,12 +352,13 @@ def get_run(run_id: int) -> Run:
         # Send the request
         res = _request(req, GetRunResponse, PATH_GET_RUN)
         if res is None:
-            return Run(run_id, "", "")
+            return Run(run_id, "", "", {})
 
         return Run(
             run_id,
             res.run.fab_id,
             res.run.fab_version,
+            dict(res.run.override_config.items()),
         )
 
     try:

src/py/flwr/common/typing.py

@@ -194,3 +194,4 @@ class Run:
     run_id: int
     fab_id: str
     fab_version: str
+    override_config: Dict[str, str]

src/py/flwr/server/driver/grpc_driver.py

@@ -127,6 +127,7 @@ def _init_run(self) -> None:
             run_id=res.run.run_id,
             fab_id=res.run.fab_id,
             fab_version=res.run.fab_version,
+            override_config=dict(res.run.override_config.items()),
         )
 
     @property

src/py/flwr/server/driver/inmemory_driver_test.py

@@ -86,7 +86,10 @@ def setUp(self) -> None:
             for _ in range(self.num_nodes)
         ]
         self.state.get_run.return_value = Run(
-            run_id=61016, fab_id="mock/mock", fab_version="v1.0.0"
+            run_id=61016,
+            fab_id="mock/mock",
+            fab_version="v1.0.0",
+            override_config={"test_key": "test_value"},
         )
         state_factory = MagicMock(state=lambda: self.state)
         self.driver = InMemoryDriver(run_id=61016, state_factory=state_factory)
@@ -98,6 +101,7 @@ def test_get_run(self) -> None:
         self.assertEqual(self.driver.run.run_id, 61016)
         self.assertEqual(self.driver.run.fab_id, "mock/mock")
         self.assertEqual(self.driver.run.fab_version, "v1.0.0")
+        self.assertEqual(self.driver.run.override_config["test_key"], "test_value")
 
     def test_get_nodes(self) -> None:
         """Test retrieval of nodes."""
@@ -223,7 +227,7 @@ def test_task_store_consistency_after_push_pull_sqlitestate(self) -> None:
         # Prepare
         state = StateFactory("").state()
         self.driver = InMemoryDriver(
-            state.create_run("", ""), MagicMock(state=lambda: state)
+            state.create_run("", "", {}), MagicMock(state=lambda: state)
         )
         msg_ids, node_id = push_messages(self.driver, self.num_nodes)
         assert isinstance(state, SqliteState)
@@ -249,7 +253,7 @@ def test_task_store_consistency_after_push_pull_inmemory_state(self) -> None:
         # Prepare
         state_factory = StateFactory(":flwr-in-memory-state:")
         state = state_factory.state()
-        self.driver = InMemoryDriver(state.create_run("", ""), state_factory)
+        self.driver = InMemoryDriver(state.create_run("", "", {}), state_factory)
         msg_ids, node_id = push_messages(self.driver, self.num_nodes)
         assert isinstance(state, InMemoryState)
 

src/py/flwr/server/superlink/driver/driver_servicer.py

@@ -69,7 +69,11 @@ def CreateRun(
         """Create run ID."""
         log(DEBUG, "DriverServicer.CreateRun")
         state: State = self.state_factory.state()
-        run_id = state.create_run(request.fab_id, request.fab_version)
+        run_id = state.create_run(
+            request.fab_id,
+            request.fab_version,
+            dict(request.override_config.items()),
+        )
         return CreateRunResponse(run_id=run_id)
 
     def PushTaskIns(

src/py/flwr/server/superlink/fleet/grpc_rere/server_interceptor_test.py

@@ -328,7 +328,7 @@ def test_successful_get_run_with_metadata(self) -> None:
         self.state.create_node(
             ping_interval=30, public_key=public_key_to_bytes(self._client_public_key)
         )
-        run_id = self.state.create_run("", "")
+        run_id = self.state.create_run("", "", {})
         request = GetRunRequest(run_id=run_id)
         shared_secret = generate_shared_key(
             self._client_private_key, self._server_public_key
@@ -359,7 +359,7 @@ def test_unsuccessful_get_run_with_metadata(self) -> None:
         self.state.create_node(
             ping_interval=30, public_key=public_key_to_bytes(self._client_public_key)
         )
-        run_id = self.state.create_run("", "")
+        run_id = self.state.create_run("", "", {})
         request = GetRunRequest(run_id=run_id)
         client_private_key, _ = generate_key_pairs()
         shared_secret = generate_shared_key(client_private_key, self._server_public_key)

src/py/flwr/server/superlink/fleet/vce/vce_api_test.py

@@ -82,7 +82,9 @@ def register_messages_into_state(
 ) -> Dict[UUID, float]:
     """Register `num_messages` into the state factory."""
     state: InMemoryState = state_factory.state()  # type: ignore
-    state.run_ids[run_id] = Run(run_id=run_id, fab_id="Mock/mock", fab_version="v1.0.0")
+    state.run_ids[run_id] = Run(
+        run_id=run_id, fab_id="Mock/mock", fab_version="v1.0.0", override_config={}
+    )
     # Artificially add TaskIns to state so they can be processed
     # by the Simulation Engine logic
     nodes_cycle = cycle(nodes_mapping.keys())  # we have more messages than supernodes

src/py/flwr/server/superlink/state/in_memory_state.py

@@ -275,15 +275,23 @@ 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) -> int:
+    def create_run(
+        self,
+        fab_id: str,
+        fab_version: str,
+        override_config: Dict[str, str],
+    ) -> int:
         """Create a new run for the specified `fab_id` and `fab_version`."""
         # Sample a random int64 as 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
+                    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.")

src/py/flwr/server/superlink/state/sqlite_state.py

@@ -15,6 +15,7 @@
 """SQLite based implemenation of server state."""
 
 
+import json
 import re
 import sqlite3
 import time
@@ -61,9 +62,10 @@
 
 SQL_CREATE_TABLE_RUN = """
 CREATE TABLE IF NOT EXISTS run(
-    run_id          INTEGER UNIQUE,
-    fab_id          TEXT,
-    fab_version     TEXT
+    run_id                INTEGER UNIQUE,
+    fab_id                TEXT,
+    fab_version           TEXT,
+    override_config       TEXT
 );
 """
 
@@ -613,7 +615,12 @@ def get_node_id(self, client_public_key: bytes) -> Optional[int]:
             return node_id
         return None
 
-    def create_run(self, fab_id: str, fab_version: str) -> int:
+    def create_run(
+        self,
+        fab_id: str,
+        fab_version: str,
+        override_config: Dict[str, str],
+    ) -> int:
         """Create a new run for the specified `fab_id` and `fab_version`."""
         # Sample a random int64 as run_id
         run_id = generate_rand_int_from_bytes(RUN_ID_NUM_BYTES)
@@ -622,8 +629,13 @@ def create_run(self, fab_id: str, fab_version: str) -> int:
         query = "SELECT COUNT(*) FROM run WHERE run_id = ?;"
         # If run_id does not exist
         if self.query(query, (run_id,))[0]["COUNT(*)"] == 0:
-            query = "INSERT INTO run (run_id, fab_id, fab_version) VALUES (?, ?, ?);"
-            self.query(query, (run_id, fab_id, fab_version))
+            query = (
+                "INSERT INTO run (run_id, fab_id, fab_version, override_config)"
+                "VALUES (?, ?, ?, ?);"
+            )
+            self.query(
+                query, (run_id, fab_id, fab_version, json.dumps(override_config))
+            )
             return run_id
         log(ERROR, "Unexpected run creation failure.")
         return 0
@@ -687,7 +699,10 @@ def get_run(self, run_id: int) -> Optional[Run]:
         try:
             row = self.query(query, (run_id,))[0]
             return Run(
-                run_id=run_id, fab_id=row["fab_id"], fab_version=row["fab_version"]
+                run_id=run_id,
+                fab_id=row["fab_id"],
+                fab_version=row["fab_version"],
+                override_config=json.loads(row["override_config"]),
             )
         except sqlite3.IntegrityError:
             log(ERROR, "`run_id` does not exist.")

src/py/flwr/server/superlink/state/state.py

@@ -16,7 +16,7 @@
 
 
 import abc
-from typing import List, Optional, Set
+from typing import Dict, List, Optional, Set
 from uuid import UUID
 
 from flwr.common.typing import Run
@@ -157,7 +157,12 @@ def get_node_id(self, client_public_key: bytes) -> Optional[int]:
         """Retrieve stored `node_id` filtered by `client_public_keys`."""
 
     @abc.abstractmethod
-    def create_run(self, fab_id: str, fab_version: str) -> int:
+    def create_run(
+        self,
+        fab_id: str,
+        fab_version: str,
+        override_config: Dict[str, str],
+    ) -> int:
         """Create a new run for the specified `fab_id` and `fab_version`."""
 
     @abc.abstractmethod