flwr-nightly 1.9.0.dev20240516__py3-none-any.whl → 1.9.0.dev20240519__py3-none-any.whl

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.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/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/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/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/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/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/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",

{flwr_nightly-1.9.0.dev20240516.dist-info → flwr_nightly-1.9.0.dev20240519.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: flwr-nightly
-Version: 1.9.0.dev20240516
+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.dev20240516.dist-info → flwr_nightly-1.9.0.dev20240519.dist-info}/RECORD

@@ -5,7 +5,7 @@ flwr/cli/build.py,sha256=W30wnPSgFuHRnGB9G_vKO14rsaibWk7m-jv9r8rDqo4,5106
 flwr/cli/config_utils.py,sha256=Hql5A5hbSpJ51hgpwaTkKqfPoaZN4Zq7FZfBuQYLMcQ,4899
 flwr/cli/example.py,sha256=1bGDYll3BXQY2kRqSN-oICqS5n1b9m0g0RvXTopXHl4,2215
 flwr/cli/new/__init__.py,sha256=cQzK1WH4JP2awef1t2UQ2xjl1agVEz9rwutV18SWV1k,789
-flwr/cli/new/new.py,sha256=whQvNN-r_opeAEpB8i7X21u53FMUKOKWbdY8gJVY-L8,6168
+flwr/cli/new/new.py,sha256=7BWziuEOE15MXX4xNLH-w0-x0ytOEfYn_AUrbaDp13Y,6223
 flwr/cli/new/templates/__init__.py,sha256=4luU8RL-CK8JJCstQ_ON809W9bNTkY1l9zSaPKBkgwY,725
 flwr/cli/new/templates/app/.gitignore.tpl,sha256=XixnHdyeMB2vwkGtGnwHqoWpH-9WChdyG0GXe57duhc,3078
 flwr/cli/new/templates/app/README.md.tpl,sha256=_qGtgpKYKoCJVjQnvlBMKvFs_1gzTcL908I3KJg0oAM,668
@@ -13,22 +13,26 @@ flwr/cli/new/templates/app/__init__.py,sha256=DU7QMY7IhMQyuwm_tja66xU0KXTWQFqzfT
 flwr/cli/new/templates/app/code/__init__.py,sha256=EM6vfvgAILKPaPn7H1wMV1Wi01WyZCP_Eg6NxD6oWg8,736
 flwr/cli/new/templates/app/code/__init__.py.tpl,sha256=olwrBeJemHNBWvjc6gJURloFRqW40dAy7FRQA5pDqHU,21
 flwr/cli/new/templates/app/code/client.hf.py.tpl,sha256=RaN89A8HgKp6kjhzH8tgtDSWW8BwwcvJdqRLcvG04zw,1450
+flwr/cli/new/templates/app/code/client.jax.py.tpl,sha256=MtxhwQzxAWNvlz8B-L-2a8LXcgaLPW8dp5K0vBZHR_o,1434
 flwr/cli/new/templates/app/code/client.mlx.py.tpl,sha256=53wJy6s3zk4CZwob_qPmMoOqJ-LZNKbdDe_hw5LwOXE,2113
 flwr/cli/new/templates/app/code/client.numpy.py.tpl,sha256=mTh7Y_jOJrPUvDYHVJy4wJCnjXZV_q-jlDkB07U5GSk,521
 flwr/cli/new/templates/app/code/client.pytorch.py.tpl,sha256=MgCtMSv1Th16Faod11HubVaARkLYt7vS9RYH962-2pk,1172
 flwr/cli/new/templates/app/code/client.sklearn.py.tpl,sha256=S71SZiHaRXtKqUk3m5Elc_c6HhKAIKLalrKOQ3p20No,2801
 flwr/cli/new/templates/app/code/client.tensorflow.py.tpl,sha256=dxrTO9JwYrDBjLsmCiRLetN9KxbnWRTeGA0BQbnOu_A,1280
 flwr/cli/new/templates/app/code/server.hf.py.tpl,sha256=Mld452y3SUkejlFzac5hpCjT7_mbA0ZEEMJIUyHtSTI,338
+flwr/cli/new/templates/app/code/server.jax.py.tpl,sha256=YTi-wroUpjRDY_AZqnoN5X-n3U5V7laL6UJgqFLEbKE,246
 flwr/cli/new/templates/app/code/server.mlx.py.tpl,sha256=Cqk3PvM0e7hzohXPqD5hG_cthXoxCfc30bpEThqMy7M,272
 flwr/cli/new/templates/app/code/server.numpy.py.tpl,sha256=fRxrDXV7pB1aDhQUXMBmrCsC1zp0uKwsBxZBx1JzbHA,248
 flwr/cli/new/templates/app/code/server.pytorch.py.tpl,sha256=ltdsnFSvFGPcycVmRL4ITlr-TV0CmmXcperZe7Vamow,593
 flwr/cli/new/templates/app/code/server.sklearn.py.tpl,sha256=cLzOpQzGIUzEazuFsjBpXAQUNPy6in6zR33SCqhix6o,341
 flwr/cli/new/templates/app/code/server.tensorflow.py.tpl,sha256=gsNrWCKTU77_65_gw9nlp1LSQojgP5QQIWILvqdjx2s,579
 flwr/cli/new/templates/app/code/task.hf.py.tpl,sha256=Rw8cnds4Ym8o8TOq6kMkwlBJfIfvsfnb02jwyulOgF8,2857
+flwr/cli/new/templates/app/code/task.jax.py.tpl,sha256=u4o3V019EH79szOw2xzVeC5r9xgQiayPi9ZTIopV2TA,1519
 flwr/cli/new/templates/app/code/task.mlx.py.tpl,sha256=y7aVj3F_98-wBnDcbPsCNnFs9BOHTn0y6XIYkByzv7Y,2598
 flwr/cli/new/templates/app/code/task.pytorch.py.tpl,sha256=NvajdZN-eTyfdqKK0v2MrvWITXw9BjJ3Ri5c1haPJDs,3684
 flwr/cli/new/templates/app/code/task.tensorflow.py.tpl,sha256=cPOUUS07QbblT9PGFucwu9lY1clRA4-W4DQGA7cpcao,1044
 flwr/cli/new/templates/app/pyproject.hf.toml.tpl,sha256=PNGBNTfWmNJ23aVnW5f1TMMJ0uEwIljevpOsI-mqX08,676
+flwr/cli/new/templates/app/pyproject.jax.toml.tpl,sha256=o34H5MvQeu4H2nRolbIas9G63mR7nDDL4rqQMlJW6LA,568
 flwr/cli/new/templates/app/pyproject.mlx.toml.tpl,sha256=JCEsuHZffO1KKkN65rSp6N-A9-OW8-kl6EQp5Z2H3uE,585
 flwr/cli/new/templates/app/pyproject.numpy.toml.tpl,sha256=m276SKsjOZ4awGdXasUKvLim66agrpAsPNP9-PN6q4I,523
 flwr/cli/new/templates/app/pyproject.pytorch.toml.tpl,sha256=QikP3u5ht6qr2BkgcnvB3rCYK7jt1cS0nAm7V8g_zFc,592
@@ -143,12 +147,13 @@ flwr/server/compat/app_utils.py,sha256=06NHrPRPrjMjz5FglSPicJ9lAWZ-rIZ1cKQFs4nD6
 flwr/server/compat/driver_client_proxy.py,sha256=Wc6jyyHY4OrJzeiy8tdXtkF8IdGREdxUPnom7VvvWPI,5444
 flwr/server/compat/legacy_context.py,sha256=D2s7PvQoDnTexuRmf1uG9Von7GUj4Qqyr7qLklSlKAM,1766
 flwr/server/criterion.py,sha256=ypbAexbztzGUxNen9RCHF91QeqiEQix4t4Ih3E-42MM,1061
-flwr/server/driver/__init__.py,sha256=bbVL5pyA0Y2HcUK4s5U0B4epI-BuUFyEJbchew_8tJY,862
+flwr/server/driver/__init__.py,sha256=bikRv6CjTwSvYh7tf10gziU5o2YotOWhhftz2tr3KDc,886
 flwr/server/driver/driver.py,sha256=t9SSSDlo9wT_y2Nl7waGYMTm2VlkvK3_bOb7ggPPlho,5090
 flwr/server/driver/grpc_driver.py,sha256=rdjkcAmtRWKeqJw4xDFqULuwVf0G2nLhfbOTrNUvPeY,11832
+flwr/server/driver/inmemory_driver.py,sha256=XfdLV3mVorTWBfthBkErJDLm8jXZ834IHF3139lTS5o,6490
 flwr/server/history.py,sha256=bBOHKyX1eQONIsUx4EUU-UnAk1i0EbEl8ioyMq_UWQ8,5063
 flwr/server/run_serverapp.py,sha256=avLi_yRNE5jD2ql95gzh04BTUbHvzH-N848_mdnnkVk,5972
-flwr/server/server.py,sha256=QySLaE-1z3JCwYIEHUK37a6dFmboSYjXMLTqYaM3sBk,17565
+flwr/server/server.py,sha256=wsXsxMZ9SQ0B42nBnUlcV83NJPycgrgg5bFwcQ4BYBE,17821
 flwr/server/server_app.py,sha256=KgAT_HqsfseTLNnfX2ph42PBbVqQ0lFzvYrT90V34y0,4402
 flwr/server/server_config.py,sha256=CZaHVAsMvGLjpWVcLPkiYxgJN4xfIyAiUrCI3fETKY4,1349
 flwr/server/strategy/__init__.py,sha256=7eVZ3hQEg2BgA_usAeL6tsLp9T6XI1VYYoFy08Xn-ew,2836
@@ -195,7 +200,7 @@ flwr/server/superlink/fleet/rest_rere/rest_api.py,sha256=8gNziOjBA8ygTzfVPYiNkg_
 flwr/server/superlink/fleet/vce/__init__.py,sha256=36MHKiefnJeyjwMQzVUK4m06Ojon3WDcwZGQsAcyVhQ,783
 flwr/server/superlink/fleet/vce/backend/__init__.py,sha256=oBIzmnrSSRvH_H0vRGEGWhWzQQwqe3zn6e13RsNwlIY,1466
 flwr/server/superlink/fleet/vce/backend/backend.py,sha256=LJsKl7oixVvptcG98Rd9ejJycNWcEVB0ODvSreLGp-A,2260
-flwr/server/superlink/fleet/vce/backend/raybackend.py,sha256=bGoqfaxsFaxmNLwZMiCCvzo-SnHVAV7niQXOIfcTmFU,6545
+flwr/server/superlink/fleet/vce/backend/raybackend.py,sha256=KCzV-n-czXxIKPwNfuD-JEVCl4-xAJaHe4taGmw9cTQ,6722
 flwr/server/superlink/fleet/vce/vce_api.py,sha256=aH-1h1EhTPCxdiqgH0_t8oDPiXX8VNNLV_BiDvu6kRk,12456
 flwr/server/superlink/state/__init__.py,sha256=ij-7Ms-hyordQdRmGQxY1-nVa4OhixJ0jr7_YDkys0s,1003
 flwr/server/superlink/state/in_memory_state.py,sha256=WoIOwgayuCu1DLRkkV6KgBsc28SKzSDxtXwO2a9Phuw,12750
@@ -209,7 +214,7 @@ flwr/server/utils/tensorboard.py,sha256=k0G6bqsLx7wfYbH2KtXsDYcOCfyIeE12-hefXA7l
 flwr/server/utils/validator.py,sha256=pzyXoOEEPSoYC2UEzened8IKSFRI-kIqqI0QlwRK9jk,5301
 flwr/server/workflow/__init__.py,sha256=SXY0XkwbkezFBxxrFB5hKUtmtAgnYISBkPouR1V71ss,902
 flwr/server/workflow/constant.py,sha256=q4DLdR8Krlxuewq2AQjwTL75hphxE5ODNz4AhViHMXk,1082
-flwr/server/workflow/default_workflows.py,sha256=p9pt8qMeX0T4uN4riT7zdaS_COuUl2YxkvRW_tzXXx0,13227
+flwr/server/workflow/default_workflows.py,sha256=_GqFCaxtiq3_UVCvZWgJ200QroGSI9qibeVcT2R71ao,14003
 flwr/server/workflow/secure_aggregation/__init__.py,sha256=3XlgDOjD_hcukTGl6Bc1B-8M_dPlVSJuTbvXIbiO-Ic,880
 flwr/server/workflow/secure_aggregation/secagg_workflow.py,sha256=wpAkYPId0nfK6SgpUAtsCni4_MQLd-uqJ81tUKu3xlI,5838
 flwr/server/workflow/secure_aggregation/secaggplus_workflow.py,sha256=BRqhlnVe8CYNoUvb_KCfRXay02NTT6a-pCrMaOqAxGc,29038
@@ -219,9 +224,9 @@ flwr/simulation/ray_transport/__init__.py,sha256=FsaAnzC4cw4DqoouBCix6496k29jACk
 flwr/simulation/ray_transport/ray_actor.py,sha256=_wv2eP7qxkCZ-6rMyYWnjLrGPBZRxjvTPjaVk8zIaQ4,19367
 flwr/simulation/ray_transport/ray_client_proxy.py,sha256=oDu4sEPIOu39vrNi-fqDAe10xtNUXMO49bM2RWfRcyw,6738
 flwr/simulation/ray_transport/utils.py,sha256=TYdtfg1P9VfTdLMOJlifInGpxWHYs9UfUqIv2wfkRLA,2392
-flwr/simulation/run_simulation.py,sha256=NdqplXCBRd9_VSqQuFWU1qG6r-KPDIpeHgQbaOCeutQ,16006
-flwr_nightly-1.9.0.dev20240516.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-flwr_nightly-1.9.0.dev20240516.dist-info/METADATA,sha256=UG54WsVtFORPeLJ_Yotx6yh5jcP91HJl6Zf8tKKN0iI,15302
-flwr_nightly-1.9.0.dev20240516.dist-info/WHEEL,sha256=FMvqSimYX_P7y0a7UY-_Mc83r5zkBZsCYPm7Lr0Bsq4,88
-flwr_nightly-1.9.0.dev20240516.dist-info/entry_points.txt,sha256=8JJPfpqMnXz9c5V_FSt07Xwd-wCWbAO3MFUDXQ5ZGsI,378
-flwr_nightly-1.9.0.dev20240516.dist-info/RECORD,,
+flwr/simulation/run_simulation.py,sha256=Jmc6DyN5UCY1U1PcDvL04NgYmEQ6ufJ1JisjG5yqfY8,15098
+flwr_nightly-1.9.0.dev20240519.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+flwr_nightly-1.9.0.dev20240519.dist-info/METADATA,sha256=Awsqxt2JHQrRwiGgDZhkqzGIoV-oco2gU58BVxb5YW8,15302
+flwr_nightly-1.9.0.dev20240519.dist-info/WHEEL,sha256=FMvqSimYX_P7y0a7UY-_Mc83r5zkBZsCYPm7Lr0Bsq4,88
+flwr_nightly-1.9.0.dev20240519.dist-info/entry_points.txt,sha256=8JJPfpqMnXz9c5V_FSt07Xwd-wCWbAO3MFUDXQ5ZGsI,378
+flwr_nightly-1.9.0.dev20240519.dist-info/RECORD,,