Middleware

doc/source/how-to-use-middleware.rst

@@ -0,0 +1,84 @@
+Use Built-in Middleware Layers
+==============================
+
+In this tutorial, we will learn how to utilize built-in middleware layers to augment the behavior of an application. Middleware allows us to perform operations before and after a task is processed in the application.
+
+What is Middleware?
+-------------------
+
+Middleware is a callable that wraps around an application. It can manipulate or inspect incoming tasks (``TaskIns``) in the ``Fwd`` and the resulting tasks (``TaskRes``) in the ``Bwd``. The signature for a middleware layer (``Layer``) is as follows:
+
+.. code-block:: python
+    APP = Callable[[Fwd], Bwd]
+    Layer = Callable[[Fwd, App], Bwd]
+
+A typical middleware function might look something like this:
+
+.. code-block:: python
+
+    def example_middleware(fwd: Fwd, app: App) -> Bwd:
+        # Do something with Fwd before passing to app.
+        bwd = app(fwd)
+        # Do something with Bwd before returning.
+        return bwd
+
+Using Middleware Layers
+-----------------------
+
+To use middleware layers in your application, you can follow these steps:
+
+1. Import the Required Middleware
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+First, import the built-in middleware layers you intend to use:
+
+.. code-block:: python
+
+    import flwr as fl
+    from flwr.client.middleware import example_middleware1, example_middleware2
+
+2. Define Your Client Function
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+Define your client function (``client_fn``) that will be wrapped by the middleware:
+
+.. code-block:: python
+
+    def client_fn():
+        # Your client code goes here.
+        pass
+
+3. Create the Application with Middleware
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+Create your application and pass the middleware layers as a list to the ``middleware`` argument. The order in which you provide the middleware layers matters:
+
+.. code-block:: python
+
+    app = fl.app.Flower(
+        client_fn=client_fn,
+        middleware=[
+            example_middleware1,  # Middleware layer 1
+            example_middleware2,  # Middleware layer 2
+        ]
+    )
+
+Order of Execution
+------------------
+
+When the application runs, the middleware layers are executed in the order they are provided in the list:
+
+1. ``example_middleware1`` (outermost layer)
+2. ``example_middleware2`` (next layer)
+3. Message handler (core app functionality)
+4. ``example_middleware2`` (on the way back)
+5. ``example_middleware1`` (outermost layer on the way back)
+
+Each middleware has a chance to inspect and modify the ``TaskIns`` in the ``Fwd`` before passing it to the next layer, and likewise with the ``TaskRes`` in the ``Bwd`` before returning it up the stack.
+
+Conclusion
+----------
+
+By following this guide, you have learned how to effectively use middleware layers to enhance your application's functionality. Remember that the order of middleware is crucial and affects how the input and output are processed.
+
+Enjoy building more robust and flexible applications with middleware layers!

doc/source/index.rst

@@ -91,6 +91,7 @@ Problem-oriented how-to guides show step-by-step how to achieve a specific goal.
    how-to-configure-logging
    how-to-enable-ssl-connections
    how-to-upgrade-to-flower-1.0
+   how-to-use-middleware
 
 .. toctree::
    :maxdepth: 1

examples/vertical-fl/simulation.py

@@ -22,4 +22,4 @@ def client_fn(cid):
 
 results_dir = Path("_static/results")
 results_dir.mkdir(exist_ok=True)
-np.save(str(results_dir/"hist.npy"), hist)
+np.save(str(results_dir / "hist.npy"), hist)

src/py/flwr/client/app.py

@@ -23,8 +23,8 @@
 from typing import Callable, ContextManager, Optional, Tuple, Union
 
 from flwr.client.client import Client
-from flwr.client.flower import Bwd, Flower, Fwd
-from flwr.client.typing import ClientFn
+from flwr.client.flower import Flower
+from flwr.client.typing import Bwd, ClientFn, Fwd
 from flwr.common import GRPC_MAX_MESSAGE_LENGTH, EventType, event
 from flwr.common.address import parse_address
 from flwr.common.constant import (

src/py/flwr/client/flower.py

@@ -16,32 +16,12 @@
 
 
 import importlib
-from dataclasses import dataclass
-from typing import Callable, cast
+from typing import List, Optional, cast
 
 from flwr.client.message_handler.message_handler import handle
-from flwr.client.typing import ClientFn
-from flwr.client.workload_state import WorkloadState
-from flwr.proto.task_pb2 import TaskIns, TaskRes
-
-
-@dataclass
-class Fwd:
-    """."""
-
-    task_ins: TaskIns
-    state: WorkloadState
-
-
-@dataclass
-class Bwd:
-    """."""
-
-    task_res: TaskRes
-    state: WorkloadState
-
-
-FlowerCallable = Callable[[Fwd], Bwd]
+from flwr.client.middleware.typing import Layer
+from flwr.client.middleware.utils import make_app
+from flwr.client.typing import Bwd, ClientFn, Fwd
 
 
 class Flower:
@@ -73,21 +53,28 @@ class Flower:
     def __init__(
         self,
         client_fn: ClientFn,  # Only for backward compatibility
+        middleware: Optional[List[Layer]] = None,
     ) -> None:
         self.client_fn = client_fn
+        self.mw_list = middleware if middleware is not None else []
 
     def __call__(self, fwd: Fwd) -> Bwd:
         """."""
+
+        # Create wrapper function for `handle`
+        def handle_app(_fwd: Fwd) -> Bwd:
+            task_res, state_updated = handle(
+                client_fn=self.client_fn,
+                state=_fwd.state,
+                task_ins=_fwd.task_ins,
+            )
+            return Bwd(task_res=task_res, state=state_updated)
+
+        # Wrap middleware layers around handle_app
+        app = make_app(handle_app, self.mw_list)
+
         # Execute the task
-        task_res, state_updated = handle(
-            client_fn=self.client_fn,
-            state=fwd.state,
-            task_ins=fwd.task_ins,
-        )
-        return Bwd(
-            task_res=task_res,
-            state=state_updated,
-        )
+        return app(fwd)
 
 
 class LoadCallableError(Exception):

src/py/flwr/client/middleware/__init__.py

@@ -0,0 +1,25 @@
+# Copyright 2023 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.
+# ==============================================================================
+"""Middleware layers."""
+
+
+from .typing import App, Layer
+from .utils import make_app
+
+__all__ = [
+    "App",
+    "Layer",
+    "make_app",
+]

src/py/flwr/client/middleware/typing.py

@@ -0,0 +1,22 @@
+# Copyright 2023 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.
+# ==============================================================================
+"""Types for middleware layers."""
+
+from typing import Callable
+
+from flwr.client.typing import Bwd, Fwd
+
+App = Callable[[Fwd], Bwd]
+Layer = Callable[[Fwd, App], Bwd]

src/py/flwr/client/middleware/utils.py

@@ -0,0 +1,36 @@
+# Copyright 2023 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.
+# ==============================================================================
+"""Utility functions for middleware layers."""
+
+from typing import List
+
+from flwr.client.typing import Bwd, Fwd
+
+from .typing import App, Layer
+
+
+def make_app(app: App, middleware_layers: List[Layer]) -> App:
+    """."""
+
+    def wrap_app(_app: App, _layer: Layer) -> App:
+        def new_app(fwd: Fwd) -> Bwd:
+            return _layer(fwd, _app)
+
+        return new_app
+
+    for layer in reversed(middleware_layers):
+        app = wrap_app(app, layer)
+
+    return app

src/py/flwr/client/middleware/utils_test.py

@@ -0,0 +1,100 @@
+# Copyright 2023 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 for the utility functions."""
+
+
+import unittest
+from typing import List
+
+from flwr.client.typing import Bwd, Fwd
+from flwr.client.workload_state import WorkloadState
+from flwr.proto.task_pb2 import TaskIns, TaskRes
+
+from .typing import App, Layer
+from .utils import make_app
+
+
+def make_mock_middleware(name: str, footprint: List[str]) -> Layer:
+    """Make a mock middleware layer."""
+
+    def middleware(fwd: Fwd, app: App) -> Bwd:
+        footprint.append(name)
+        fwd.task_ins.task_id += f"{name}"
+        bwd = app(fwd)
+        footprint.append(name)
+        bwd.task_res.task_id += f"{name}"
+        return bwd
+
+    return middleware
+
+
+def make_mock_app(name: str, footprint: List[str]) -> App:
+    """Make a mock app."""
+
+    def app(fwd: Fwd) -> Bwd:
+        footprint.append(name)
+        fwd.task_ins.task_id += f"{name}"
+        return Bwd(task_res=TaskRes(task_id=name), state=WorkloadState({}))
+
+    return app
+
+
+class TestMakeApp(unittest.TestCase):
+    """Tests for the `make_app` function."""
+
+    def test_multiple_middlewares(self) -> None:
+        """Test if multiple middlewares are called in the correct order."""
+        # Prepare
+        footprint: List[str] = []
+        mock_app = make_mock_app("app", footprint)
+        mock_middleware_names = [f"middleware{i}" for i in range(1, 15)]
+        mock_middleware_layers = [
+            make_mock_middleware(name, footprint) for name in mock_middleware_names
+        ]
+        task_ins = TaskIns()
+
+        # Execute
+        wrapped_app = make_app(mock_app, mock_middleware_layers)
+        task_res = wrapped_app(Fwd(task_ins=task_ins, state=WorkloadState({}))).task_res
+
+        # Assert
+        trace = mock_middleware_names + ["app"]
+        self.assertEqual(footprint, trace + list(reversed(mock_middleware_names)))
+        # pylint: disable-next=no-member
+        self.assertEqual(task_ins.task_id, "".join(trace))
+        self.assertEqual(task_res.task_id, "".join(reversed(trace)))
+
+    def test_filter(self) -> None:
+        """Test if a middleware can filter incoming TaskIns."""
+        # Prepare
+        footprint: List[str] = []
+        mock_app = make_mock_app("app", footprint)
+        task_ins = TaskIns()
+
+        def filter_layer(fwd: Fwd, _: App) -> Bwd:
+            footprint.append("filter")
+            fwd.task_ins.task_id += "filter"
+            # Skip calling app
+            return Bwd(task_res=TaskRes(task_id="filter"), state=WorkloadState({}))
+
+        # Execute
+        wrapped_app = make_app(mock_app, [filter_layer])
+        task_res = wrapped_app(Fwd(task_ins=task_ins, state=WorkloadState({}))).task_res
+
+        # Assert
+        self.assertEqual(footprint, ["filter"])
+        # pylint: disable-next=no-member
+        self.assertEqual(task_ins.task_id, "filter")
+        self.assertEqual(task_res.task_id, "filter")