flwr-nightly 1.9.0.dev20240517__tar.gz → 1.9.0.dev20240519__tar.gz

{flwr_nightly-1.9.0.dev20240517 → flwr_nightly-1.9.0.dev20240519}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: flwr-nightly
-Version: 1.9.0.dev20240517
+Version: 1.9.0.dev20240519
 Summary: Flower: A Friendly Federated Learning Framework
 Home-page: https://flower.ai
 License: Apache-2.0

{flwr_nightly-1.9.0.dev20240517 → flwr_nightly-1.9.0.dev20240519}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
 name = "flwr-nightly"
-version = "1.9.0.dev20240517"
+version = "1.9.0.dev20240519"
 description = "Flower: A Friendly Federated Learning Framework"
 license = "Apache-2.0"
 authors = ["The Flower Authors <hello@flower.ai>"]

{flwr_nightly-1.9.0.dev20240517 → flwr_nightly-1.9.0.dev20240519}/src/py/flwr/cli/new/new.py

@@ -37,6 +37,7 @@ class MlFramework(str, Enum):
     NUMPY = "NumPy"
     PYTORCH = "PyTorch"
     TENSORFLOW = "TensorFlow"
+    JAX = "JAX"
     HUGGINGFACE = "HF"
     MLX = "MLX"
     SKLEARN = "sklearn"
@@ -155,6 +156,7 @@ def new(
     # Depending on the framework, generate task.py file
     frameworks_with_tasks = [
         MlFramework.PYTORCH.value.lower(),
+        MlFramework.JAX.value.lower(),
         MlFramework.HUGGINGFACE.value.lower(),
         MlFramework.MLX.value.lower(),
         MlFramework.TENSORFLOW.value.lower(),

flwr_nightly-1.9.0.dev20240519/src/py/flwr/cli/new/templates/app/code/client.jax.py.tpl

@@ -0,0 +1,55 @@
+"""$project_name: A Flower / JAX app."""
+
+import jax
+from flwr.client import NumPyClient, ClientApp
+
+from $import_name.task import (
+    evaluation,
+    get_params,
+    load_data,
+    load_model,
+    loss_fn,
+    set_params,
+    train,
+)
+
+
+# Define Flower Client and client_fn
+class FlowerClient(NumPyClient):
+    def __init__(self):
+        self.train_x, self.train_y, self.test_x, self.test_y = load_data()
+        self.grad_fn = jax.grad(loss_fn)
+        model_shape = self.train_x.shape[1:]
+
+        self.params = load_model(model_shape)
+
+    def get_parameters(self, config):
+        return get_params(self.params)
+
+    def set_parameters(self, parameters):
+        set_params(self.params, parameters)
+
+    def fit(self, parameters, config):
+        self.set_parameters(parameters)
+        self.params, loss, num_examples = train(
+            self.params, self.grad_fn, self.train_x, self.train_y
+        )
+        parameters = self.get_parameters(config={})
+        return parameters, num_examples, {"loss": float(loss)}
+
+    def evaluate(self, parameters, config):
+        self.set_parameters(parameters)
+        loss, num_examples = evaluation(
+            self.params, self.grad_fn, self.test_x, self.test_y
+        )
+        return float(loss), num_examples, {"loss": float(loss)}
+
+def client_fn(cid):
+    # Return Client instance
+    return FlowerClient().to_client()
+
+
+# Flower ClientApp
+app = ClientApp(
+    client_fn,
+)

flwr_nightly-1.9.0.dev20240519/src/py/flwr/cli/new/templates/app/code/server.jax.py.tpl

@@ -0,0 +1,12 @@
+"""$project_name: A Flower / JAX app."""
+
+import flwr as fl
+
+# Configure the strategy
+strategy = fl.server.strategy.FedAvg()
+
+# Flower ServerApp
+app = fl.server.ServerApp(
+    config=fl.server.ServerConfig(num_rounds=3),
+    strategy=strategy,
+)

flwr_nightly-1.9.0.dev20240519/src/py/flwr/cli/new/templates/app/code/task.jax.py.tpl

@@ -0,0 +1,57 @@
+"""$project_name: A Flower / JAX app."""
+
+import jax
+import jax.numpy as jnp
+from sklearn.datasets import make_regression
+from sklearn.model_selection import train_test_split
+import numpy as np
+
+key = jax.random.PRNGKey(0)
+
+
+def load_data():
+    # Load dataset
+    X, y = make_regression(n_features=3, random_state=0)
+    X, X_test, y, y_test = train_test_split(X, y)
+    return X, y, X_test, y_test
+
+
+def load_model(model_shape):
+    # Extract model parameters
+    params = {"b": jax.random.uniform(key), "w": jax.random.uniform(key, model_shape)}
+    return params
+
+
+def loss_fn(params, X, y):
+    # Return MSE as loss
+    err = jnp.dot(X, params["w"]) + params["b"] - y
+    return jnp.mean(jnp.square(err))
+
+
+def train(params, grad_fn, X, y):
+    loss = 1_000_000
+    num_examples = X.shape[0]
+    for epochs in range(50):
+        grads = grad_fn(params, X, y)
+        params = jax.tree.map(lambda p, g: p - 0.05 * g, params, grads)
+        loss = loss_fn(params, X, y)
+    return params, loss, num_examples
+
+
+def evaluation(params, grad_fn, X_test, y_test):
+    num_examples = X_test.shape[0]
+    err_test = loss_fn(params, X_test, y_test)
+    loss_test = jnp.mean(jnp.square(err_test))
+    return loss_test, num_examples
+
+
+def get_params(params):
+    parameters = []
+    for _, val in params.items():
+        parameters.append(np.array(val))
+    return parameters
+
+
+def set_params(local_params, global_params):
+    for key, value in list(zip(local_params.keys(), global_params)):
+        local_params[key] = value

flwr_nightly-1.9.0.dev20240519/src/py/flwr/cli/new/templates/app/pyproject.jax.toml.tpl

@@ -0,0 +1,28 @@
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[project]
+name = "$package_name"
+version = "1.0.0"
+description = ""
+authors = [
+    { name = "The Flower Authors", email = "hello@flower.ai" },
+]
+license = {text = "Apache License (2.0)"}
+dependencies = [
+    "flwr[simulation]>=1.8.0,<2.0",
+    "jax==0.4.26",
+    "jaxlib==0.4.26",
+    "scikit-learn==1.4.2",
+]
+
+[tool.hatch.build.targets.wheel]
+packages = ["."]
+
+[flower]
+publisher = "$username"
+
+[flower.components]
+serverapp = "$import_name.server:app"
+clientapp = "$import_name.client:app"

{flwr_nightly-1.9.0.dev20240517 → flwr_nightly-1.9.0.dev20240519}/src/py/flwr/server/driver/__init__.py

@@ -16,10 +16,11 @@
 
 
 from .driver import Driver
-from .grpc_driver import GrpcDriver, GrpcDriverHelper
+from .grpc_driver import GrpcDriver
+from .inmemory_driver import InMemoryDriver
 
 __all__ = [
     "Driver",
     "GrpcDriver",
-    "GrpcDriverHelper",
+    "InMemoryDriver",
 ]

flwr_nightly-1.9.0.dev20240519/src/py/flwr/server/driver/inmemory_driver.py

@@ -0,0 +1,181 @@
+# 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 in-memory Driver."""
+
+
+import time
+import warnings
+from typing import Iterable, List, Optional
+from uuid import UUID
+
+from flwr.common import DEFAULT_TTL, Message, Metadata, RecordSet
+from flwr.common.serde import message_from_taskres, message_to_taskins
+from flwr.proto.node_pb2 import Node  # pylint: disable=E0611
+from flwr.server.superlink.state import StateFactory
+
+from .driver import Driver
+
+
+class InMemoryDriver(Driver):
+    """`InMemoryDriver` class provides an interface to the Driver API.
+
+    Parameters
+    ----------
+    state_factory : StateFactory
+        A StateFactory embedding a state that this driver can interface with.
+    fab_id : str (default: None)
+        The identifier of the FAB used in the run.
+    fab_version : str (default: None)
+        The version of the FAB used in the run.
+    """
+
+    def __init__(
+        self,
+        state_factory: StateFactory,
+        fab_id: Optional[str] = None,
+        fab_version: Optional[str] = None,
+    ) -> None:
+        self.run_id: Optional[int] = None
+        self.fab_id = fab_id if fab_id is not None else ""
+        self.fab_version = fab_version if fab_version is not None else ""
+        self.node = Node(node_id=0, anonymous=True)
+        self.state = state_factory.state()
+
+    def _check_message(self, message: Message) -> None:
+        # Check if the message is valid
+        if not (
+            message.metadata.run_id == self.run_id
+            and message.metadata.src_node_id == self.node.node_id
+            and message.metadata.message_id == ""
+            and message.metadata.reply_to_message == ""
+            and message.metadata.ttl > 0
+        ):
+            raise ValueError(f"Invalid message: {message}")
+
+    def _get_run_id(self) -> int:
+        """Return run_id.
+
+        If unset, create a new run.
+        """
+        if self.run_id is None:
+            self.run_id = self.state.create_run(
+                fab_id=self.fab_id, fab_version=self.fab_version
+            )
+        return self.run_id
+
+    def create_message(  # pylint: disable=too-many-arguments
+        self,
+        content: RecordSet,
+        message_type: str,
+        dst_node_id: int,
+        group_id: str,
+        ttl: Optional[float] = None,
+    ) -> Message:
+        """Create a new message with specified parameters.
+
+        This method constructs a new `Message` with given content and metadata.
+        The `run_id` and `src_node_id` will be set automatically.
+        """
+        run_id = self._get_run_id()
+        if ttl:
+            warnings.warn(
+                "A custom TTL was set, but note that the SuperLink does not enforce "
+                "the TTL yet. The SuperLink will start enforcing the TTL in a future "
+                "version of Flower.",
+                stacklevel=2,
+            )
+        ttl_ = DEFAULT_TTL if ttl is None else ttl
+
+        metadata = Metadata(
+            run_id=run_id,
+            message_id="",  # Will be set by the server
+            src_node_id=self.node.node_id,
+            dst_node_id=dst_node_id,
+            reply_to_message="",
+            group_id=group_id,
+            ttl=ttl_,
+            message_type=message_type,
+        )
+        return Message(metadata=metadata, content=content)
+
+    def get_node_ids(self) -> List[int]:
+        """Get node IDs."""
+        run_id = self._get_run_id()
+        return list(self.state.get_nodes(run_id))
+
+    def push_messages(self, messages: Iterable[Message]) -> Iterable[str]:
+        """Push messages to specified node IDs.
+
+        This method takes an iterable of messages and sends each message
+        to the node specified in `dst_node_id`.
+        """
+        task_ids: List[str] = []
+        for msg in messages:
+            # Check message
+            self._check_message(msg)
+            # Convert Message to TaskIns
+            taskins = message_to_taskins(msg)
+            # Store in state
+            taskins.task.pushed_at = time.time()
+            task_id = self.state.store_task_ins(taskins)
+            if task_id:
+                task_ids.append(str(task_id))
+
+        return task_ids
+
+    def pull_messages(self, message_ids: Iterable[str]) -> Iterable[Message]:
+        """Pull messages based on message IDs.
+
+        This method is used to collect messages from the SuperLink that correspond to a
+        set of given message IDs.
+        """
+        msg_ids = {UUID(msg_id) for msg_id in message_ids}
+        # Pull TaskRes
+        task_res_list = self.state.get_task_res(task_ids=msg_ids, limit=len(msg_ids))
+        # Delete tasks in state
+        self.state.delete_tasks(msg_ids)
+        # Convert TaskRes to Message
+        msgs = [message_from_taskres(taskres) for taskres in task_res_list]
+        return msgs
+
+    def send_and_receive(
+        self,
+        messages: Iterable[Message],
+        *,
+        timeout: Optional[float] = None,
+    ) -> Iterable[Message]:
+        """Push messages to specified node IDs and pull the reply messages.
+
+        This method sends a list of messages to their destination node IDs and then
+        waits for the replies. It continues to pull replies until either all replies are
+        received or the specified timeout duration is exceeded.
+        """
+        # Push messages
+        msg_ids = set(self.push_messages(messages))
+
+        # Pull messages
+        end_time = time.time() + (timeout if timeout is not None else 0.0)
+        ret: List[Message] = []
+        while timeout is None or time.time() < end_time:
+            res_msgs = self.pull_messages(msg_ids)
+            ret.extend(res_msgs)
+            msg_ids.difference_update(
+                {msg.metadata.reply_to_message for msg in res_msgs}
+            )
+            if len(msg_ids) == 0:
+                break
+            # Sleep
+            time.sleep(3)
+        return ret

{flwr_nightly-1.9.0.dev20240517 → flwr_nightly-1.9.0.dev20240519}/src/py/flwr/server/server.py

@@ -282,7 +282,14 @@ class Server:
         get_parameters_res = random_client.get_parameters(
             ins=ins, timeout=timeout, group_id=server_round
         )
-        log(INFO, "Received initial parameters from one random client")
+        if get_parameters_res.status.code == Code.OK:
+            log(INFO, "Received initial parameters from one random client")
+        else:
+            log(
+                WARN,
+                "Failed to receive initial parameters from the client."
+                " Empty initial parameters will be used.",
+            )
         return get_parameters_res.parameters
 
 

{flwr_nightly-1.9.0.dev20240517 → flwr_nightly-1.9.0.dev20240519}/src/py/flwr/server/superlink/fleet/vce/backend/raybackend.py

@@ -55,7 +55,12 @@ class RayBackend(Backend):
         runtime_env = (
             self._configure_runtime_env(work_dir=work_dir) if work_dir else None
         )
-        if backend_config.get("silent", False):
+
+        if backend_config.get("mute_logging", False):
+            init_ray(
+                logging_level=WARNING, log_to_driver=False, runtime_env=runtime_env
+            )
+        elif backend_config.get("silent", False):
             init_ray(logging_level=WARNING, log_to_driver=True, runtime_env=runtime_env)
         else:
             init_ray(runtime_env=runtime_env)

{flwr_nightly-1.9.0.dev20240517 → flwr_nightly-1.9.0.dev20240519}/src/py/flwr/server/workflow/default_workflows.py

@@ -18,12 +18,22 @@
 import io
 import timeit
 from logging import INFO, WARN
-from typing import Optional, cast
+from typing import List, Optional, Tuple, Union, cast
 
 import flwr.common.recordset_compat as compat
-from flwr.common import ConfigsRecord, Context, GetParametersIns, ParametersRecord, log
+from flwr.common import (
+    Code,
+    ConfigsRecord,
+    Context,
+    EvaluateRes,
+    FitRes,
+    GetParametersIns,
+    ParametersRecord,
+    log,
+)
 from flwr.common.constant import MessageType, MessageTypeLegacy
 
+from ..client_proxy import ClientProxy
 from ..compat.app_utils import start_update_client_manager_thread
 from ..compat.legacy_context import LegacyContext
 from ..driver import Driver
@@ -136,7 +146,14 @@ def default_init_params_workflow(driver: Driver, context: Context) -> None:
             ]
         )
         msg = list(messages)[0]
-        if msg.has_content():
+
+        if (
+            msg.has_content()
+            and compat._extract_status_from_recordset(  # pylint: disable=W0212
+                "getparametersres", msg.content
+            ).code
+            == Code.OK
+        ):
             log(INFO, "Received initial parameters from one random client")
             paramsrecord = next(iter(msg.content.parameters_records.values()))
         else:
@@ -257,18 +274,20 @@ def default_fit_workflow(  # pylint: disable=R0914
     )
 
     # Aggregate training results
-    results = [
-        (
-            node_id_to_proxy[msg.metadata.src_node_id],
-            compat.recordset_to_fitres(msg.content, False),
-        )
-        for msg in messages
-        if msg.has_content()
-    ]
-    failures = [Exception(msg.error) for msg in messages if msg.has_error()]
-    aggregated_result = context.strategy.aggregate_fit(
-        current_round, results, failures  # type: ignore
-    )
+    results: List[Tuple[ClientProxy, FitRes]] = []
+    failures: List[Union[Tuple[ClientProxy, FitRes], BaseException]] = []
+    for msg in messages:
+        if msg.has_content():
+            proxy = node_id_to_proxy[msg.metadata.src_node_id]
+            fitres = compat.recordset_to_fitres(msg.content, False)
+            if fitres.status.code == Code.OK:
+                results.append((proxy, fitres))
+            else:
+                failures.append((proxy, fitres))
+        else:
+            failures.append(Exception(msg.error))
+
+    aggregated_result = context.strategy.aggregate_fit(current_round, results, failures)
     parameters_aggregated, metrics_aggregated = aggregated_result
 
     # Update the parameters and write history
@@ -341,17 +360,21 @@ def default_evaluate_workflow(driver: Driver, context: Context) -> None:
     )
 
     # Aggregate the evaluation results
-    results = [
-        (
-            node_id_to_proxy[msg.metadata.src_node_id],
-            compat.recordset_to_evaluateres(msg.content),
-        )
-        for msg in messages
-        if msg.has_content()
-    ]
-    failures = [Exception(msg.error) for msg in messages if msg.has_error()]
+    results: List[Tuple[ClientProxy, EvaluateRes]] = []
+    failures: List[Union[Tuple[ClientProxy, EvaluateRes], BaseException]] = []
+    for msg in messages:
+        if msg.has_content():
+            proxy = node_id_to_proxy[msg.metadata.src_node_id]
+            evalres = compat.recordset_to_evaluateres(msg.content)
+            if evalres.status.code == Code.OK:
+                results.append((proxy, evalres))
+            else:
+                failures.append((proxy, evalres))
+        else:
+            failures.append(Exception(msg.error))
+
     aggregated_result = context.strategy.aggregate_evaluate(
-        current_round, results, failures  # type: ignore
+        current_round, results, failures
     )
 
     loss_aggregated, metrics_aggregated = aggregated_result

{flwr_nightly-1.9.0.dev20240517 → flwr_nightly-1.9.0.dev20240519}/src/py/flwr/simulation/run_simulation.py

@@ -24,16 +24,13 @@ from logging import DEBUG, ERROR, INFO, WARNING
 from time import sleep
 from typing import Dict, Optional
 
-import grpc
-
 from flwr.client import ClientApp
 from flwr.common import EventType, event, log
 from flwr.common.logger import set_logger_propagation, update_console_handler
 from flwr.common.typing import ConfigsRecordValues
-from flwr.server.driver import Driver, GrpcDriver
+from flwr.server.driver import Driver, InMemoryDriver
 from flwr.server.run_serverapp import run
 from flwr.server.server_app import ServerApp
-from flwr.server.superlink.driver.driver_grpc import run_driver_api_grpc
 from flwr.server.superlink.fleet import vce
 from flwr.server.superlink.state import StateFactory
 from flwr.simulation.ray_transport.utils import (
@@ -56,7 +53,6 @@ def run_simulation_from_cli() -> None:
         backend_name=args.backend,
         backend_config=backend_config_dict,
         app_dir=args.app_dir,
-        driver_api_address=args.driver_api_address,
         enable_tf_gpu_growth=args.enable_tf_gpu_growth,
         verbose_logging=args.verbose,
     )
@@ -177,7 +173,6 @@ def _main_loop(
     num_supernodes: int,
     backend_name: str,
     backend_config_stream: str,
-    driver_api_address: str,
     app_dir: str,
     enable_tf_gpu_growth: bool,
     client_app: Optional[ClientApp] = None,
@@ -194,21 +189,11 @@ def _main_loop(
     # Initialize StateFactory
     state_factory = StateFactory(":flwr-in-memory-state:")
 
-    # Start Driver API
-    driver_server: grpc.Server = run_driver_api_grpc(
-        address=driver_api_address,
-        state_factory=state_factory,
-        certificates=None,
-    )
-
     f_stop = asyncio.Event()
     serverapp_th = None
     try:
         # Initialize Driver
-        driver = GrpcDriver(
-            driver_service_address=driver_api_address,
-            root_certificates=None,
-        )
+        driver = InMemoryDriver(state_factory)
 
         # Get and run ServerApp thread
         serverapp_th = run_serverapp_th(
@@ -239,9 +224,6 @@ def _main_loop(
         raise RuntimeError("An error was encountered. Ending simulation.") from ex
 
     finally:
-        # Stop Driver
-        driver_server.stop(grace=0)
-        driver.close()
         # Trigger stop event
         f_stop.set()
 
@@ -262,7 +244,6 @@ def _run_simulation(
     client_app_attr: Optional[str] = None,
     server_app_attr: Optional[str] = None,
     app_dir: str = "",
-    driver_api_address: str = "0.0.0.0:9091",
     enable_tf_gpu_growth: bool = False,
     verbose_logging: bool = False,
 ) -> None:
@@ -302,9 +283,6 @@ def _run_simulation(
         Add specified directory to the PYTHONPATH and load `ClientApp` from there.
         (Default: current working directory.)
 
-    driver_api_address : str (default: "0.0.0.0:9091")
-        Driver API (gRPC) server address (IPv4, IPv6, or a domain name)
-
     enable_tf_gpu_growth : bool (default: False)
         A boolean to indicate whether to enable GPU growth on the main thread. This is
         desirable if you make use of a TensorFlow model on your `ServerApp` while
@@ -342,7 +320,6 @@ def _run_simulation(
         num_supernodes,
         backend_name,
         backend_config_stream,
-        driver_api_address,
         app_dir,
         enable_tf_gpu_growth,
         client_app,
@@ -399,12 +376,6 @@ def _parse_args_run_simulation() -> argparse.ArgumentParser:
         required=True,
         help="Number of simulated SuperNodes.",
     )
-    parser.add_argument(
-        "--driver-api-address",
-        default="0.0.0.0:9091",
-        type=str,
-        help="For example: `server:app` or `project.package.module:wrapper.app`",
-    )
     parser.add_argument(
         "--backend",
         default="ray",