making federation_options a compulsory arg to create_run()

adap · Nov 6, 2024 · 41c37fc · 41c37fc
1 parent 08bacdc
commit 41c37fc
Show file tree

Hide file tree

Showing 8 changed files with 77 additions and 64 deletions.
diff --git a/src/py/flwr/server/driver/inmemory_driver_test.py b/src/py/flwr/server/driver/inmemory_driver_test.py
@@ -21,7 +21,7 @@
 from unittest.mock import MagicMock, patch
 from uuid import uuid4
 
-from flwr.common import RecordSet
+from flwr.common import ConfigsRecord, RecordSet
 from flwr.common.constant import NODE_ID_NUM_BYTES, PING_MAX_INTERVAL
 from flwr.common.message import Error
 from flwr.common.serde import (
@@ -232,7 +232,7 @@ def test_task_store_consistency_after_push_pull_sqlitestate(self) -> None:
         """Test tasks are deleted in sqlite state once messages are pulled."""
         # Prepare
         state = LinkStateFactory("").state()
-        run_id = state.create_run("", "", "", {})
+        run_id = state.create_run("", "", "", {}, ConfigsRecord())
         self.driver = InMemoryDriver(MagicMock(state=lambda: state))
         self.driver.init_run(run_id=run_id)
         msg_ids, node_id = push_messages(self.driver, self.num_nodes)
@@ -259,7 +259,7 @@ def test_task_store_consistency_after_push_pull_inmemory_state(self) -> None:
         # Prepare
         state_factory = LinkStateFactory(":flwr-in-memory-state:")
         state = state_factory.state()
-        run_id = state.create_run("", "", "", {})
+        run_id = state.create_run("", "", "", {}, ConfigsRecord())
         self.driver = InMemoryDriver(state_factory)
         self.driver.init_run(run_id=run_id)
         msg_ids, node_id = push_messages(self.driver, self.num_nodes)

diff --git a/src/py/flwr/server/superlink/driver/serverappio_servicer.py b/src/py/flwr/server/superlink/driver/serverappio_servicer.py
@@ -23,6 +23,7 @@
 
 import grpc
 
+from flwr.common import ConfigsRecord
 from flwr.common.constant import Status
 from flwr.common.logger import log
 from flwr.common.serde import (
@@ -112,6 +113,7 @@ def CreateRun(
             request.fab_version,
             fab_hash,
             user_config_from_proto(request.override_config),
+            ConfigsRecord(),
         )
         return CreateRunResponse(run_id=run_id)
 

diff --git a/src/py/flwr/server/superlink/fleet/grpc_rere/server_interceptor_test.py b/src/py/flwr/server/superlink/fleet/grpc_rere/server_interceptor_test.py
@@ -20,6 +20,7 @@
 
 import grpc
 
+from flwr.common import ConfigsRecord
 from flwr.common.constant import FLEET_API_GRPC_RERE_DEFAULT_ADDRESS
 from flwr.common.secure_aggregation.crypto.symmetric_encryption import (
     compute_hmac,
@@ -334,7 +335,7 @@ def test_successful_get_run_with_metadata(self) -> None:
         self.state.create_node(
             ping_interval=30, public_key=public_key_to_bytes(self._node_public_key)
         )
-        run_id = self.state.create_run("", "", "", {})
+        run_id = self.state.create_run("", "", "", {}, ConfigsRecord())
         request = GetRunRequest(run_id=run_id)
         shared_secret = generate_shared_key(
             self._node_private_key, self._server_public_key
@@ -365,7 +366,7 @@ def test_unsuccessful_get_run_with_metadata(self) -> None:
         self.state.create_node(
             ping_interval=30, public_key=public_key_to_bytes(self._node_public_key)
         )
-        run_id = self.state.create_run("", "", "", {})
+        run_id = self.state.create_run("", "", "", {}, ConfigsRecord())
         request = GetRunRequest(run_id=run_id)
         node_private_key, _ = generate_key_pairs()
         shared_secret = generate_shared_key(node_private_key, self._server_public_key)

diff --git a/src/py/flwr/server/superlink/linkstate/in_memory_linkstate.py b/src/py/flwr/server/superlink/linkstate/in_memory_linkstate.py
@@ -387,7 +387,7 @@ def create_run(
         fab_version: Optional[str],
         fab_hash: Optional[str],
         override_config: UserConfig,
-        federation_options: Optional[ConfigsRecord] = None,
+        federation_options: ConfigsRecord,
     ) -> int:
         """Create a new run for the specified `fab_hash`."""
         # Sample a random int64 as run_id
@@ -413,7 +413,7 @@ def create_run(
                 self.run_ids[run_id] = run_record
 
                 # Record federation options. Leave empty if not passed
-                self.federation_options[run_id] = federation_options or ConfigsRecord()
+                self.federation_options[run_id] = federation_options
                 return run_id
         log(ERROR, "Unexpected run creation failure.")
         return 0

diff --git a/src/py/flwr/server/superlink/linkstate/linkstate.py b/src/py/flwr/server/superlink/linkstate/linkstate.py
@@ -159,7 +159,7 @@ def create_run(  # pylint: disable=too-many-arguments,too-many-positional-argume
         fab_version: Optional[str],
         fab_hash: Optional[str],
         override_config: UserConfig,
-        federation_options: Optional[ConfigsRecord] = None,
+        federation_options: ConfigsRecord,
     ) -> int:
         """Create a new run for the specified `fab_hash`."""