Virtual Client Simulation

.github/workflows/docs.yml

@@ -21,7 +21,7 @@ jobs:
           python -m pip install -U poetry==1.1.6
           poetry config virtualenvs.create false
       - name: Install dependencies (mandatory only)
-        run: python -m poetry install --extras "baseline examples-tensorflow examples-pytorch http-logger ops"
+        run: python -m poetry install --extras "baseline examples-tensorflow examples-pytorch http-logger ops simulation"
       - name: Build and deploy docs
         env:
           AWS_DEFAULT_REGION: ${{ secrets. AWS_DEFAULT_REGION }}

.github/workflows/flower.yml

@@ -28,7 +28,7 @@ jobs:
           python -m pip install -U poetry==1.1.6
           poetry config virtualenvs.create false
       - name: Install dependencies (mandatory only)
-        run: python -m poetry install
+        run: python -m poetry install --extras "simulation"
       - name: Check if protos need recompilation
         run: ./dev/check-protos.sh
       - name: Lint + Test (isort/black/docformatter/mypy/pylint/flake8/pytest)

dev/bootstrap.sh

@@ -16,4 +16,5 @@ python -m poetry install \
   --extras "examples-pytorch" \
   --extras "examples-tensorflow" \
   --extras "http-logger" \
-  --extras "ops"
+  --extras "ops" \
+  --extras "simulation"

pyproject.toml

@@ -64,8 +64,10 @@ paramiko = { version = "^2.7.1", optional = true }
 docker = { version = "^4.2.0", optional = true }
 matplotlib = { version = "^3.2.1", optional = true }
 tqdm = { version = "^4.48.2", optional = true }
+ray = { version = "^1.4.0", optional = true}
 
 [tool.poetry.extras]
+simulation = ["ray"]
 baseline = ["tensorflow-cpu", "boto3", "boto3_type_annotations", "paramiko", "docker", "matplotlib"]
 examples-pytorch = ["torch", "torchvision", "tqdm"]
 examples-tensorflow = ["tensorflow-cpu"]

src/py/flwr/__init__.py

@@ -15,9 +15,10 @@
 """Flower main package."""
 
 
-from . import client, server
+from . import client, server, simulation
 
 __all__ = [
     "client",
     "server",
+    "simulation",
 ]

src/py/flwr/client/__init__.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""Flower Client."""
+"""Flower client."""
 
 
 from .app import start_client as start_client

src/py/flwr/dataset/__init__.py

@@ -12,4 +12,4 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""Flower Datasets."""
+"""Flower dataset."""

src/py/flwr/server/__init__.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""Flower Server."""
+"""Flower server."""
 
 
 from .app import start_server as start_server

src/py/flwr/server/grpc_server/grpc_client_proxy.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""Networked Flower client implementation."""
+"""gRPC-based Flower ClientProxy implementation."""
 
 
 from flwr import common

src/py/flwr/server/ray_server/__init__.py

@@ -0,0 +1,15 @@
+# Copyright 2020 Adap 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.
+# ==============================================================================
+"""Server for Ray-based ClientProxy objects."""

src/py/flwr/simulation/__init__.py

@@ -0,0 +1,22 @@
+# Copyright 2020 Adap 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 simulation."""
+
+
+from .app import start_simulation as start_simulation
+
+__all__ = [
+    "start_simulation",
+]

src/py/flwr/simulation/app.py

@@ -0,0 +1,106 @@
+# Copyright 2020 Adap 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 Simulation app."""
+
+
+from logging import INFO
+from typing import Any, Callable, Dict, Optional
+
+try:
+    import ray
+except ImportError:
+    ray = None  # type: ignore
+
+from flwr.client.client import Client
+from flwr.common.logger import log
+from flwr.server.app import _fl, _init_defaults
+from flwr.server.strategy import Strategy
+from flwr.simulation.ray_simulation.ray_client_proxy import RayClientProxy
+
+RAY_IMPORT_ERROR: str = """Unable to import module `ray`.
+
+To install the necessary dependencies, install `flwr` with the `simulation` extra:
+
+    pip install -U flwr["simulation"]
+"""
+
+
+def start_simulation(  # pylint: disable=too-many-arguments
+    client_fn: Callable[[str], Client],
+    num_clients: int,
+    client_resources: Optional[Dict[str, int]] = None,
+    num_rounds: int = 1,
+    strategy: Optional[Strategy] = None,
+    ray_init_args: Optional[Dict[str, Any]] = None,
+) -> None:
+    """Start a Ray-based Flower simulation server.
+
+    Parameters
+    ----------
+    client_fn : Callable[[str], Client]
+        A function creating client instances. The function must take a single
+        str argument called cid. It should return a single client instance
+        containing the local dataset and model of the client with the matching
+        client id (i.e., cid).
+    num_clients : int
+        The total number of clients in this simulation.
+    client_resources : Optional[Dict[str, int]] (default: None)
+        CPU and GPU resources for a single client. Supported keys are
+        `num_cpus` and `num_gpus`. Example: `{"num_cpus": 4, "num_gpus": 1}`
+    num_rounds : int (default: 1)
+        The number of rounds to train.
+    strategy : Optional[flwr.server.Strategy] (default: None)
+        An implementation of the abstract base class `flwr.server.Strategy`. If
+        no strategy is provided, then `start_server` will use
+        `flwr.server.strategy.FedAvg`.
+    ray_init_args : Optional[Dict[str, Any]] (default: None)
+        Optional dictionary containing `ray.init` arguments.
+    """
+
+    # Ray cannot be assumed to be installed
+    if ray is None:
+        raise ImportError(RAY_IMPORT_ERROR)
+
+    # Initialize Ray
+    if not ray_init_args:
+        ray_init_args = {}
+    ray.init(**ray_init_args)
+    log(
+        INFO,
+        "Ray initialized with resources: %s",
+        ray.cluster_resources(),
+    )
+
+    # Initialize server and server config
+    config = {"num_rounds": num_rounds}
+    initialized_server, initialized_config = _init_defaults(None, config, strategy)
+    log(
+        INFO,
+        "Starting Flower simulation running: %s",
+        initialized_config,
+    )
+
+    # Register one RayClientProxy object for each client with the ClientManager
+    resources = client_resources if client_resources is not None else {}
+    for i in range(num_clients):
+        client_proxy = RayClientProxy(
+            client_fn=client_fn,
+            cid=str(i),
+            resources=resources,
+        )
+        initialized_server.client_manager().register(client=client_proxy)
+
+    # Start training
+    _fl(server=initialized_server, config=initialized_config)

src/py/flwr/simulation/ray_simulation/__init__.py

@@ -0,0 +1,15 @@
+# Copyright 2020 Adap 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.
+# ==============================================================================
+"""Simulation for Ray-based ClientProxy objects."""

src/py/flwr/simulation/ray_simulation/ray_client_proxy.py

@@ -0,0 +1,109 @@
+# Copyright 2020 Adap 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.
+# ==============================================================================
+"""Ray-based Flower ClientProxy implementation."""
+
+
+from typing import Callable, Dict, Union, cast
+
+try:
+    import ray
+except ImportError:
+    ray = None  # type: ignore
+
+from flwr import common
+from flwr.client import Client, NumPyClient
+from flwr.client.numpy_client import NumPyClientWrapper
+from flwr.server.client_proxy import ClientProxy
+
+ClientFn = Callable[[str], Client]
+
+
+class RayClientProxy(ClientProxy):
+    """Flower client proxy which delegates work using Ray."""
+
+    def __init__(self, client_fn: ClientFn, cid: str, resources: Dict[str, int]):
+        super().__init__(cid)
+        self.client_fn = client_fn
+        self.resources = resources
+
+    def get_parameters(self) -> common.ParametersRes:
+        """Return the current local model parameters."""
+        future_paramseters_res = launch_and_get_parameters.options(
+            **self.resources
+        ).remote(self.client_fn, self.cid)
+        res = ray.get(future_paramseters_res)
+        return cast(
+            common.ParametersRes,
+            res,
+        )
+
+    def fit(self, ins: common.FitIns) -> common.FitRes:
+        """Train model parameters on the locally held dataset."""
+        future_fit_res = launch_and_fit.options(**self.resources).remote(
+            self.client_fn, self.cid, ins
+        )
+        res = ray.get(future_fit_res)
+        return cast(
+            common.FitRes,
+            res,
+        )
+
+    def evaluate(self, ins: common.EvaluateIns) -> common.EvaluateRes:
+        """Evaluate model parameters on the locally held dataset."""
+        future_evaluate_res = launch_and_evaluate.options(**self.resources).remote(
+            self.client_fn, self.cid, ins
+        )
+        res = ray.get(future_evaluate_res)
+        return cast(
+            common.EvaluateRes,
+            res,
+        )
+
+    def reconnect(self, reconnect: common.Reconnect) -> common.Disconnect:
+        """Disconnect and (optionally) reconnect later."""
+        return common.Disconnect(reason="")  # Nothing to do here (yet)
+
+
+@ray.remote  # type: ignore
+def launch_and_get_parameters(client_fn: ClientFn, cid: str) -> common.ParametersRes:
+    """Exectue get_parameters remotely."""
+    client: Client = _create_client(client_fn, cid)
+    return client.get_parameters()
+
+
+@ray.remote  # type: ignore
+def launch_and_fit(
+    client_fn: ClientFn, cid: str, fit_ins: common.FitIns
+) -> common.FitRes:
+    """Exectue fit remotely."""
+    client: Client = _create_client(client_fn, cid)
+    return client.fit(fit_ins)
+
+
+@ray.remote  # type: ignore
+def launch_and_evaluate(
+    client_fn: ClientFn, cid: str, evaluate_ins: common.EvaluateIns
+) -> common.EvaluateRes:
+    """Exectue evaluate remotely."""
+    client: Client = _create_client(client_fn, cid)
+    return client.evaluate(evaluate_ins)
+
+
+def _create_client(client_fn: ClientFn, cid: str) -> Client:
+    """Create a client instance."""
+    client: Union[Client, NumPyClient] = client_fn(cid)
+    if isinstance(client, NumPyClient):
+        client = NumPyClientWrapper(numpy_client=client)
+    return client