Merge branch 'main' into refactor-grpc-driver

.github/workflows/e2e.yml

@@ -67,7 +67,7 @@ jobs:
           - connection: insecure
             authentication: client-auth
     name: |
-      SuperExec / 
+      Exec API / 
       Python ${{ matrix.python-version }} /
       ${{ matrix.connection }} / 
       ${{ matrix.authentication }} / 

e2e/test_exec_api.sh

@@ -81,17 +81,19 @@ timeout 2m flower-superlink $combined_args --executor-config "$executor_config"
 sl_pid=$(pgrep -f "flower-superlink")
 sleep 2
 
-timeout 2m flower-supernode ./ $client_arg \
-    --superlink $server_address $client_auth_1 \
-    --node-config "partition-id=0 num-partitions=2" --max-retries 0 &
-cl1_pid=$!
-sleep 2
+if [ "$3" = "deployment-engine" ]; then
+  timeout 2m flower-supernode ./ $client_arg \
+      --superlink $server_address $client_auth_1 \
+      --node-config "partition-id=0 num-partitions=2" --max-retries 0 &
+  cl1_pid=$!
+  sleep 2
 
-timeout 2m flower-supernode ./ $client_arg \
-    --superlink $server_address $client_auth_2 \
-    --node-config "partition-id=1 num-partitions=2" --max-retries 0 &
-cl2_pid=$!
-sleep 2
+  timeout 2m flower-supernode ./ $client_arg \
+      --superlink $server_address $client_auth_2 \
+      --node-config "partition-id=1 num-partitions=2" --max-retries 0 &
+  cl2_pid=$!
+  sleep 2
+fi
 
 timeout 1m flwr run --run-config num-server-rounds=1 ../e2e-tmp-test e2e
 
@@ -105,7 +107,10 @@ while [ "$found_success" = false ] && [ $elapsed -lt $timeout ]; do
     if grep -q "Run finished" flwr_output.log; then
         echo "Training worked correctly!"
         found_success=true
-        kill $cl1_pid; kill $cl2_pid; sleep 1; kill $sl_pid;
+        if $3 = "deployment-engine"; then
+          kill $cl1_pid; kill $cl2_pid;
+        fi
+        sleep 1; kill $sl_pid;
     else
         echo "Waiting for training ... ($elapsed seconds elapsed)"
     fi
@@ -116,5 +121,8 @@ done
 
 if [ "$found_success" = false ]; then
     echo "Training had an issue and timed out."
-    kill $cl1_pid; kill $cl2_pid; kill $sl_pid;
+    if $3 = "deployment-engine"; then
+      kill $cl1_pid; kill $cl2_pid;
+    fi
+    kill $sl_pid;
 fi

src/proto/flwr/proto/exec.proto

@@ -19,6 +19,7 @@ package flwr.proto;
 
 import "flwr/proto/fab.proto";
 import "flwr/proto/transport.proto";
+import "flwr/proto/recordset.proto";
 
 service Exec {
   // Start run upon request
@@ -31,7 +32,7 @@ service Exec {
 message StartRunRequest {
   Fab fab = 1;
   map<string, Scalar> override_config = 2;
-  map<string, Scalar> federation_config = 3;
+  ConfigsRecord federation_options = 3;
 }
 message StartRunResponse { uint64 run_id = 1; }
 message StreamLogsRequest {

src/py/flwr/cli/run/run.py

@@ -29,10 +29,18 @@
     validate_federation_in_project_config,
     validate_project_config,
 )
-from flwr.common.config import flatten_dict, parse_config_args
+from flwr.common.config import (
+    flatten_dict,
+    parse_config_args,
+    user_config_to_configsrecord,
+)
 from flwr.common.grpc import GRPC_MAX_MESSAGE_LENGTH, create_channel
 from flwr.common.logger import log
-from flwr.common.serde import fab_to_proto, user_config_to_proto
+from flwr.common.serde import (
+    configs_record_to_proto,
+    fab_to_proto,
+    user_config_to_proto,
+)
 from flwr.common.typing import Fab
 from flwr.proto.exec_pb2 import StartRunRequest  # pylint: disable=E0611
 from flwr.proto.exec_pb2_grpc import ExecStub
@@ -94,6 +102,7 @@ def run(
         _run_without_exec_api(app, federation_config, config_overrides, federation)
 
 
+# pylint: disable-next=too-many-locals
 def _run_with_exec_api(
     app: Path,
     federation_config: dict[str, Any],
@@ -118,12 +127,14 @@ def _run_with_exec_api(
     content = Path(fab_path).read_bytes()
     fab = Fab(fab_hash, content)
 
+    # Construct a `ConfigsRecord` out of a flattened `UserConfig`
+    fed_conf = flatten_dict(federation_config.get("options", {}))
+    c_record = user_config_to_configsrecord(fed_conf)
+
     req = StartRunRequest(
         fab=fab_to_proto(fab),
         override_config=user_config_to_proto(parse_config_args(config_overrides)),
-        federation_config=user_config_to_proto(
-            flatten_dict(federation_config.get("options"))
-        ),
+        federation_options=configs_record_to_proto(c_record),
     )
     res = stub.StartRun(req)
 

src/py/flwr/client/app.py

@@ -32,6 +32,7 @@
 from flwr.cli.install import install_from_fab
 from flwr.client.client import Client
 from flwr.client.client_app import ClientApp, LoadClientAppError
+from flwr.client.nodestate.nodestate_factory import NodeStateFactory
 from flwr.client.typing import ClientFnExt
 from flwr.common import GRPC_MAX_MESSAGE_LENGTH, Context, EventType, Message, event
 from flwr.common.address import parse_address
@@ -365,6 +366,8 @@ def _on_backoff(retry_state: RetryState) -> None:
 
     # DeprecatedRunInfoStore gets initialized when the first connection is established
     run_info_store: Optional[DeprecatedRunInfoStore] = None
+    state_factory = NodeStateFactory()
+    state = state_factory.state()
 
     runs: dict[int, Run] = {}
 
@@ -396,13 +399,14 @@ def _on_backoff(retry_state: RetryState) -> None:
                     )
                 else:
                     # Call create_node fn to register node
-                    node_id: Optional[int] = (  # pylint: disable=assignment-from-none
-                        create_node()
-                    )  # pylint: disable=not-callable
-                    if node_id is None:
-                        raise ValueError("Node registration failed")
+                    # and store node_id in state
+                    if (node_id := create_node()) is None:
+                        raise ValueError(
+                            "Failed to register SuperNode with the SuperLink"
+                        )
+                    state.set_node_id(node_id)
                     run_info_store = DeprecatedRunInfoStore(
-                        node_id=node_id,
+                        node_id=state.get_node_id(),
                         node_config=node_config,
                     )
 

src/py/flwr/client/nodestate/__init__.py

@@ -0,0 +1,25 @@
+# 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.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Flower NodeState."""
+
+from .in_memory_nodestate import InMemoryNodeState as InMemoryNodeState
+from .nodestate import NodeState as NodeState
+from .nodestate_factory import NodeStateFactory as NodeStateFactory
+
+__all__ = [
+    "InMemoryNodeState",
+    "NodeState",
+    "NodeStateFactory",
+]

src/py/flwr/client/nodestate/in_memory_nodestate.py

@@ -0,0 +1,38 @@
+# 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.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""In-memory NodeState implementation."""
+
+
+from typing import Optional
+
+from flwr.client.nodestate.nodestate import NodeState
+
+
+class InMemoryNodeState(NodeState):
+    """In-memory NodeState implementation."""
+
+    def __init__(self) -> None:
+        # Store node_id
+        self.node_id: Optional[int] = None
+
+    def set_node_id(self, node_id: Optional[int]) -> None:
+        """Set the node ID."""
+        self.node_id = node_id
+
+    def get_node_id(self) -> int:
+        """Get the node ID."""
+        if self.node_id is None:
+            raise ValueError("Node ID not set")
+        return self.node_id

src/py/flwr/client/nodestate/nodestate.py

@@ -0,0 +1,30 @@
+# 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.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Abstract base class NodeState."""
+
+import abc
+from typing import Optional
+
+
+class NodeState(abc.ABC):
+    """Abstract NodeState."""
+
+    @abc.abstractmethod
+    def set_node_id(self, node_id: Optional[int]) -> None:
+        """Set the node ID."""
+
+    @abc.abstractmethod
+    def get_node_id(self) -> int:
+        """Get the node ID."""

src/py/flwr/client/nodestate/nodestate_factory.py

@@ -0,0 +1,37 @@
+# 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.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Factory class that creates NodeState instances."""
+
+import threading
+from typing import Optional
+
+from .in_memory_nodestate import InMemoryNodeState
+from .nodestate import NodeState
+
+
+class NodeStateFactory:
+    """Factory class that creates NodeState instances."""
+
+    def __init__(self) -> None:
+        self.state_instance: Optional[NodeState] = None
+        self.lock = threading.RLock()
+
+    def state(self) -> NodeState:
+        """Return a State instance and create it, if necessary."""
+        # Lock access to NodeStateFactory to prevent returning different instances
+        with self.lock:
+            if self.state_instance is None:
+                self.state_instance = InMemoryNodeState()
+            return self.state_instance

src/py/flwr/client/nodestate/nodestate_test.py

@@ -0,0 +1,69 @@
+# 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.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests all NodeState implementations have to conform to."""
+
+import unittest
+from abc import abstractmethod
+
+from flwr.client.nodestate import InMemoryNodeState, NodeState
+
+
+class StateTest(unittest.TestCase):
+    """Test all state implementations."""
+
+    # This is to True in each child class
+    __test__ = False
+
+    @abstractmethod
+    def state_factory(self) -> NodeState:
+        """Provide state implementation to test."""
+        raise NotImplementedError()
+
+    def test_get_set_node_id(self) -> None:
+        """Test set_node_id."""
+        # Prepare
+        state: NodeState = self.state_factory()
+        node_id = 123
+
+        # Execute
+        state.set_node_id(node_id)
+
+        retrieved_node_id = state.get_node_id()
+
+        # Assert
+        assert node_id == retrieved_node_id
+
+    def test_get_node_id_fails(self) -> None:
+        """Test get_node_id fails correctly if node_id is not set."""
+        # Prepare
+        state: NodeState = self.state_factory()
+
+        # Execute and assert
+        with self.assertRaises(ValueError):
+            state.get_node_id()
+
+
+class InMemoryStateTest(StateTest):
+    """Test InMemoryState implementation."""
+
+    __test__ = True
+
+    def state_factory(self) -> NodeState:
+        """Return InMemoryState."""
+        return InMemoryNodeState()
+
+
+if __name__ == "__main__":
+    unittest.main(verbosity=2)

src/py/flwr/common/config.py

@@ -22,6 +22,7 @@
 import tomli
 
 from flwr.cli.config_utils import get_fab_config, validate_fields
+from flwr.common import ConfigsRecord
 from flwr.common.constant import (
     APP_DIR,
     FAB_CONFIG_FILE,
@@ -229,3 +230,12 @@ def get_metadata_from_config(config: dict[str, Any]) -> tuple[str, str]:
         config["project"]["version"],
         f"{config['tool']['flwr']['app']['publisher']}/{config['project']['name']}",
     )
+
+
+def user_config_to_configsrecord(config: UserConfig) -> ConfigsRecord:
+    """Construct a `ConfigsRecord` out of a `UserConfig`."""
+    c_record = ConfigsRecord()
+    for k, v in config.items():
+        c_record[k] = v
+
+    return c_record