flwr-nightly 1.9.0.dev20240509__py3-none-any.whl → 1.9.0.dev20240531__py3-none-any.whl

flwr/cli/build.py

@@ -24,7 +24,7 @@ import pathspec
 import typer
 from typing_extensions import Annotated
 
-from .config_utils import load_and_validate_with_defaults
+from .config_utils import load_and_validate
 from .utils import is_valid_project_name
 
 
@@ -67,9 +67,7 @@ def build(
         )
         raise typer.Exit(code=1)
 
-    conf, errors, warnings = load_and_validate_with_defaults(
-        directory / "pyproject.toml"
-    )
+    conf, errors, warnings = load_and_validate(directory / "pyproject.toml")
     if conf is None:
         typer.secho(
             "Project configuration could not be loaded.\npyproject.toml is invalid:\n"

flwr/cli/config_utils.py

@@ -22,7 +22,7 @@ import tomli
 from flwr.common import object_ref
 
 
-def load_and_validate_with_defaults(
+def load_and_validate(
     path: Optional[Path] = None,
 ) -> Tuple[Optional[Dict[str, Any]], List[str], List[str]]:
     """Load and validate pyproject.toml as dict.
@@ -47,14 +47,6 @@ def load_and_validate_with_defaults(
     if not is_valid:
         return (None, errors, warnings)
 
-    # Apply defaults
-    defaults = {
-        "flower": {
-            "engine": {"name": "simulation", "simulation": {"supernode": {"num": 2}}}
-        }
-    }
-    config = apply_defaults(config, defaults)
-
     return (config, errors, warnings)
 
 
@@ -129,17 +121,3 @@ def validate(config: Dict[str, Any]) -> Tuple[bool, List[str], List[str]]:
         return False, [reason], []
 
     return True, [], []
-
-
-def apply_defaults(
-    config: Dict[str, Any],
-    defaults: Dict[str, Any],
-) -> Dict[str, Any]:
-    """Apply defaults to config."""
-    for key in defaults:
-        if key in config:
-            if isinstance(config[key], dict) and isinstance(defaults[key], dict):
-                apply_defaults(config[key], defaults[key])
-        else:
-            config[key] = defaults[key]
-    return config

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/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/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/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/cli/new/templates/app/pyproject.hf.toml.tpl

@@ -29,3 +29,9 @@ publisher = "$username"
 [flower.components]
 serverapp = "$import_name.server:app"
 clientapp = "$import_name.client:app"
+
+[flower.engine]
+name = "simulation"
+
+[flower.engine.simulation.supernode]
+num = 2

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/cli/new/templates/app/pyproject.mlx.toml.tpl

@@ -26,3 +26,9 @@ publisher = "$username"
 [flower.components]
 serverapp = "$import_name.server:app"
 clientapp = "$import_name.client:app"
+
+[flower.engine]
+name = "simulation"
+
+[flower.engine.simulation.supernode]
+num = 2

flwr/cli/new/templates/app/pyproject.numpy.toml.tpl

@@ -24,3 +24,9 @@ publisher = "$username"
 [flower.components]
 serverapp = "$import_name.server:app"
 clientapp = "$import_name.client:app"
+
+[flower.engine]
+name = "simulation"
+
+[flower.engine.simulation.supernode]
+num = 2

flwr/cli/new/templates/app/pyproject.pytorch.toml.tpl

@@ -26,3 +26,9 @@ publisher = "$username"
 [flower.components]
 serverapp = "$import_name.server:app"
 clientapp = "$import_name.client:app"
+
+[flower.engine]
+name = "simulation"
+
+[flower.engine.simulation.supernode]
+num = 2

flwr/cli/new/templates/app/pyproject.sklearn.toml.tpl

@@ -25,3 +25,9 @@ publisher = "$username"
 [flower.components]
 serverapp = "$import_name.server:app"
 clientapp = "$import_name.client:app"
+
+[flower.engine]
+name = "simulation"
+
+[flower.engine.simulation.supernode]
+num = 2

flwr/cli/new/templates/app/pyproject.tensorflow.toml.tpl

@@ -25,3 +25,9 @@ publisher = "$username"
 [flower.components]
 serverapp = "$import_name.server:app"
 clientapp = "$import_name.client:app"
+
+[flower.engine]
+name = "simulation"
+
+[flower.engine.simulation.supernode]
+num = 2

flwr/cli/run/run.py

@@ -15,18 +15,32 @@
 """Flower command line interface `run` command."""
 
 import sys
+from enum import Enum
+from typing import Optional
 
 import typer
+from typing_extensions import Annotated
 
 from flwr.cli import config_utils
 from flwr.simulation.run_simulation import _run_simulation
 
 
-def run() -> None:
+class Engine(str, Enum):
+    """Enum defining the engine to run on."""
+
+    SIMULATION = "simulation"
+
+
+def run(
+    engine: Annotated[
+        Optional[Engine],
+        typer.Option(case_sensitive=False, help="The ML framework to use"),
+    ] = None,
+) -> None:
     """Run Flower project."""
     typer.secho("Loading project configuration... ", fg=typer.colors.BLUE)
 
-    config, errors, warnings = config_utils.load_and_validate_with_defaults()
+    config, errors, warnings = config_utils.load_and_validate()
 
     if config is None:
         typer.secho(
@@ -49,9 +63,11 @@ def run() -> None:
 
     server_app_ref = config["flower"]["components"]["serverapp"]
     client_app_ref = config["flower"]["components"]["clientapp"]
-    engine = config["flower"]["engine"]["name"]
 
-    if engine == "simulation":
+    if engine is None:
+        engine = config["flower"]["engine"]["name"]
+
+    if engine == Engine.SIMULATION:
         num_supernodes = config["flower"]["engine"]["simulation"]["supernode"]["num"]
 
         typer.secho("Starting run... ", fg=typer.colors.BLUE)

flwr/client/grpc_rere_client/connection.py

@@ -21,7 +21,7 @@ from contextlib import contextmanager
 from copy import copy
 from logging import DEBUG, ERROR
 from pathlib import Path
-from typing import Callable, Iterator, Optional, Sequence, Tuple, Union, cast
+from typing import Callable, Iterator, Optional, Sequence, Tuple, Type, Union, cast
 
 import grpc
 from cryptography.hazmat.primitives.asymmetric import ec
@@ -73,6 +73,7 @@ def grpc_request_response(  # pylint: disable=R0913, R0914, R0915
     authentication_keys: Optional[
         Tuple[ec.EllipticCurvePrivateKey, ec.EllipticCurvePublicKey]
     ] = None,
+    adapter_cls: Optional[Type[FleetStub]] = None,
 ) -> Iterator[
     Tuple[
         Callable[[], Optional[Message]],
@@ -133,7 +134,9 @@ def grpc_request_response(  # pylint: disable=R0913, R0914, R0915
     channel.subscribe(on_channel_state_change)
 
     # Shared variables for inner functions
-    stub = FleetStub(channel)
+    if adapter_cls is None:
+        adapter_cls = FleetStub
+    stub = adapter_cls(channel)
     metadata: Optional[Metadata] = None
     node: Optional[Node] = None
     ping_thread: Optional[threading.Thread] = None

flwr/client/mod/comms_mods.py

@@ -29,7 +29,7 @@ def message_size_mod(
 ) -> Message:
     """Message size mod.
 
-    This mod logs the size in Bytes of the message being transmited.
+    This mod logs the size in bytes of the message being transmited.
     """
     message_size_in_bytes = 0
 
@@ -42,7 +42,7 @@ def message_size_mod(
     for m_record in msg.content.metrics_records.values():
         message_size_in_bytes += m_record.count_bytes()
 
-    log(INFO, "Message size: %i Bytes", message_size_in_bytes)
+    log(INFO, "Message size: %i bytes", message_size_in_bytes)
 
     return call_next(msg, ctxt)
 
@@ -53,7 +53,7 @@ def parameters_size_mod(
     """Parameters size mod.
 
     This mod logs the number of parameters transmitted in the message as well as their
-    size in Bytes.
+    size in bytes.
     """
     model_size_stats = {}
     parameters_size_in_bytes = 0
@@ -74,6 +74,6 @@ def parameters_size_mod(
     if model_size_stats:
         log(INFO, model_size_stats)
 
-    log(INFO, "Total parameters transmited: %i Bytes", parameters_size_in_bytes)
+    log(INFO, "Total parameters transmitted: %i bytes", parameters_size_in_bytes)
 
     return call_next(msg, ctxt)

flwr/client/mod/localdp_mod.py

@@ -145,8 +145,7 @@ class LocalDpMod:
         )
         log(
             INFO,
-            "LocalDpMod: local DP noise with "
-            "standard deviation: %.4f added to parameters.",
+            "LocalDpMod: local DP noise with %.4f stedv added to parameters",
             noise_value_sd,
         )
 

flwr/client/supernode/app.py

@@ -20,6 +20,7 @@ from logging import DEBUG, INFO, WARN
 from pathlib import Path
 from typing import Callable, Optional, Tuple
 
+from cryptography.exceptions import UnsupportedAlgorithm
 from cryptography.hazmat.primitives.asymmetric import ec
 from cryptography.hazmat.primitives.serialization import (
     load_ssh_private_key,
@@ -31,9 +32,6 @@ from flwr.common import EventType, event
 from flwr.common.exit_handlers import register_exit_handlers
 from flwr.common.logger import log
 from flwr.common.object_ref import load_app, validate
-from flwr.common.secure_aggregation.crypto.symmetric_encryption import (
-    ssh_types_to_elliptic_curve,
-)
 
 from ..app import _start_client_internal
 
@@ -242,40 +240,60 @@ def _parse_args_common(parser: argparse.ArgumentParser) -> None:
         " Default: current working directory.",
     )
     parser.add_argument(
-        "--authentication-keys",
-        nargs=2,
-        metavar=("CLIENT_PRIVATE_KEY", "CLIENT_PUBLIC_KEY"),
+        "--auth-supernode-private-key",
+        type=str,
+        help="The SuperNode's private key (as a path str) to enable authentication.",
+    )
+    parser.add_argument(
+        "--auth-supernode-public-key",
         type=str,
-        help="Provide two file paths: (1) the client's private "
-        "key file, and (2) the client's public key file.",
+        help="The SuperNode's public key (as a path str) to enable authentication.",
     )
 
 
 def _try_setup_client_authentication(
     args: argparse.Namespace,
 ) -> Optional[Tuple[ec.EllipticCurvePrivateKey, ec.EllipticCurvePublicKey]]:
-    if not args.authentication_keys:
+    if not args.auth_supernode_private_key and not args.auth_supernode_public_key:
         return None
 
-    ssh_private_key = load_ssh_private_key(
-        Path(args.authentication_keys[0]).read_bytes(),
-        None,
-    )
-    ssh_public_key = load_ssh_public_key(Path(args.authentication_keys[1]).read_bytes())
+    if not args.auth_supernode_private_key or not args.auth_supernode_public_key:
+        sys.exit(
+            "Authentication requires file paths to both "
+            "'--auth-supernode-private-key' and '--auth-supernode-public-key'"
+            "to be provided (providing only one of them is not sufficient)."
+        )
+
+    try:
+        ssh_private_key = load_ssh_private_key(
+            Path(args.auth_supernode_private_key).read_bytes(),
+            None,
+        )
+        if not isinstance(ssh_private_key, ec.EllipticCurvePrivateKey):
+            raise ValueError()
+    except (ValueError, UnsupportedAlgorithm):
+        sys.exit(
+            "Error: Unable to parse the private key file in "
+            "'--auth-supernode-private-key'. Authentication requires elliptic "
+            "curve private and public key pair. Please ensure that the file "
+            "path points to a valid private key file and try again."
+        )
 
     try:
-        client_private_key, client_public_key = ssh_types_to_elliptic_curve(
-            ssh_private_key, ssh_public_key
+        ssh_public_key = load_ssh_public_key(
+            Path(args.auth_supernode_public_key).read_bytes()
         )
-    except TypeError:
+        if not isinstance(ssh_public_key, ec.EllipticCurvePublicKey):
+            raise ValueError()
+    except (ValueError, UnsupportedAlgorithm):
         sys.exit(
-            "The file paths provided could not be read as a private and public "
-            "key pair. Client authentication requires an elliptic curve public and "
-            "private key pair. Please provide the file paths containing elliptic "
-            "curve private and public keys to '--authentication-keys'."
+            "Error: Unable to parse the public key file in "
+            "'--auth-supernode-public-key'. Authentication requires elliptic "
+            "curve private and public key pair. Please ensure that the file "
+            "path points to a valid public key file and try again."
         )
 
     return (
-        client_private_key,
-        client_public_key,
+        ssh_private_key,
+        ssh_public_key,
     )

flwr/common/recordset_compat.py

@@ -35,6 +35,8 @@ from .typing import (
     Status,
 )
 
+EMPTY_TENSOR_KEY = "_empty"
+
 
 def parametersrecord_to_parameters(
     record: ParametersRecord, keep_input: bool
@@ -59,7 +61,8 @@ def parametersrecord_to_parameters(
     parameters = Parameters(tensors=[], tensor_type="")
 
     for key in list(record.keys()):
-        parameters.tensors.append(record[key].data)
+        if key != EMPTY_TENSOR_KEY:
+            parameters.tensors.append(record[key].data)
 
         if not parameters.tensor_type:
             # Setting from first array in record. Recall the warning in the docstrings
@@ -103,6 +106,10 @@ def parameters_to_parametersrecord(
             data=tensor, dtype="", stype=tensor_type, shape=[]
         )
 
+    if num_arrays == 0:
+        ordered_dict[EMPTY_TENSOR_KEY] = Array(
+            data=b"", dtype="", stype=tensor_type, shape=[]
+        )
     return ParametersRecord(ordered_dict, keep_input=keep_input)
 
 

flwr/common/secure_aggregation/crypto/symmetric_encryption.py

@@ -117,18 +117,3 @@ def verify_hmac(key: bytes, message: bytes, hmac_value: bytes) -> bool:
         return True
     except InvalidSignature:
         return False
-
-
-def ssh_types_to_elliptic_curve(
-    private_key: serialization.SSHPrivateKeyTypes,
-    public_key: serialization.SSHPublicKeyTypes,
-) -> Tuple[ec.EllipticCurvePrivateKey, ec.EllipticCurvePublicKey]:
-    """Cast SSH key types to elliptic curve."""
-    if isinstance(private_key, ec.EllipticCurvePrivateKey) and isinstance(
-        public_key, ec.EllipticCurvePublicKey
-    ):
-        return (private_key, public_key)
-
-    raise TypeError(
-        "The provided key is not an EllipticCurvePrivateKey or EllipticCurvePublicKey"
-    )

flwr/proto/grpcadapter_pb2.py

@@ -0,0 +1,32 @@
+# -*- coding: utf-8 -*-
+# Generated by the protocol buffer compiler.  DO NOT EDIT!
+# source: flwr/proto/grpcadapter.proto
+# Protobuf Python Version: 4.25.0
+"""Generated protocol buffer code."""
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import descriptor_pool as _descriptor_pool
+from google.protobuf import symbol_database as _symbol_database
+from google.protobuf.internal import builder as _builder
+# @@protoc_insertion_point(imports)
+
+_sym_db = _symbol_database.Default()
+
+
+
+
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x1c\x66lwr/proto/grpcadapter.proto\x12\nflwr.proto\"\xba\x01\n\x10MessageContainer\x12<\n\x08metadata\x18\x01 \x03(\x0b\x32*.flwr.proto.MessageContainer.MetadataEntry\x12\x19\n\x11grpc_message_name\x18\x02 \x01(\t\x12\x1c\n\x14grpc_message_content\x18\x03 \x01(\x0c\x1a/\n\rMetadataEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\t:\x02\x38\x01\x32Z\n\x0bGrpcAdapter\x12K\n\x0bSendReceive\x12\x1c.flwr.proto.MessageContainer\x1a\x1c.flwr.proto.MessageContainer\"\x00\x62\x06proto3')
+
+_globals = globals()
+_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
+_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'flwr.proto.grpcadapter_pb2', _globals)
+if _descriptor._USE_C_DESCRIPTORS == False:
+  DESCRIPTOR._options = None
+  _globals['_MESSAGECONTAINER_METADATAENTRY']._options = None
+  _globals['_MESSAGECONTAINER_METADATAENTRY']._serialized_options = b'8\001'
+  _globals['_MESSAGECONTAINER']._serialized_start=45
+  _globals['_MESSAGECONTAINER']._serialized_end=231
+  _globals['_MESSAGECONTAINER_METADATAENTRY']._serialized_start=184
+  _globals['_MESSAGECONTAINER_METADATAENTRY']._serialized_end=231
+  _globals['_GRPCADAPTER']._serialized_start=233
+  _globals['_GRPCADAPTER']._serialized_end=323
+# @@protoc_insertion_point(module_scope)

flwr/proto/grpcadapter_pb2.pyi

@@ -0,0 +1,43 @@
+"""
+@generated by mypy-protobuf.  Do not edit manually!
+isort:skip_file
+"""
+import builtins
+import google.protobuf.descriptor
+import google.protobuf.internal.containers
+import google.protobuf.message
+import typing
+import typing_extensions
+
+DESCRIPTOR: google.protobuf.descriptor.FileDescriptor
+
+class MessageContainer(google.protobuf.message.Message):
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+    class MetadataEntry(google.protobuf.message.Message):
+        DESCRIPTOR: google.protobuf.descriptor.Descriptor
+        KEY_FIELD_NUMBER: builtins.int
+        VALUE_FIELD_NUMBER: builtins.int
+        key: typing.Text
+        value: typing.Text
+        def __init__(self,
+            *,
+            key: typing.Text = ...,
+            value: typing.Text = ...,
+            ) -> None: ...
+        def ClearField(self, field_name: typing_extensions.Literal["key",b"key","value",b"value"]) -> None: ...
+
+    METADATA_FIELD_NUMBER: builtins.int
+    GRPC_MESSAGE_NAME_FIELD_NUMBER: builtins.int
+    GRPC_MESSAGE_CONTENT_FIELD_NUMBER: builtins.int
+    @property
+    def metadata(self) -> google.protobuf.internal.containers.ScalarMap[typing.Text, typing.Text]: ...
+    grpc_message_name: typing.Text
+    grpc_message_content: builtins.bytes
+    def __init__(self,
+        *,
+        metadata: typing.Optional[typing.Mapping[typing.Text, typing.Text]] = ...,
+        grpc_message_name: typing.Text = ...,
+        grpc_message_content: builtins.bytes = ...,
+        ) -> None: ...
+    def ClearField(self, field_name: typing_extensions.Literal["grpc_message_content",b"grpc_message_content","grpc_message_name",b"grpc_message_name","metadata",b"metadata"]) -> None: ...
+global___MessageContainer = MessageContainer