docs(framework) Update the How to Use Built-in Mods guide (#4962)

Co-authored-by: Javier <[email protected]>
Co-authored-by: Daniel J. Beutel <[email protected]>

framework/docs/source/how-to-use-built-in-mods.rst

@@ -5,8 +5,10 @@
 Use Built-in Mods
 =================
 
-**Note: This tutorial covers experimental features. The functionality and interfaces may
-change in future versions.**
+.. note::
+
+    This tutorial covers preview features. The functionality and interfaces may change
+    in future versions.
 
 In this tutorial, we will learn how to utilize built-in mods to augment the behavior of
 a ``ClientApp``. Mods (sometimes also called Modifiers) allow us to perform operations
@@ -21,80 +23,153 @@ is as follows:
 
 .. code-block:: python
 
-    ClientApp = Callable[[Message, Context], Message]
-    Mod = Callable[[Message, Context, ClientApp], Message]
+    ClientAppCallable = Callable[[Message, Context], Message]
+    Mod = Callable[[Message, Context, ClientAppCallable], Message]
 
 A typical mod function might look something like this:
 
 .. code-block:: python
 
-    def example_mod(msg: Message, ctx: Context, nxt: ClientApp) -> Message:
+    from flwr.client.typing import ClientAppCallable
+    from flwr.common import Context, Message
+
+
+    def example_mod(msg: Message, ctx: Context, call_next: ClientAppCallable) -> Message:
         # Do something with incoming Message (or Context)
-        # before passing to the inner ``ClientApp``
-        msg = nxt(msg, ctx)
+        # before passing it to the next layer in the chain.
+        # This could be another Mod or, if this is the last Mod, the ClientApp itself.
+        msg = call_next(msg, ctx)
         # Do something with outgoing Message (or Context)
         # before returning
         return msg
 
 Using Mods
 ----------
 
-To use mods in your ``ClientApp``, you can follow these steps:
+Mods can be registered in two ways: **Application-wide mods** and **Function-specific
+mods**.
+
+1. **Application-wide mods**: These mods apply to all functions within the
+   ``ClientApp``.
+2. **Function-specific mods**: These mods apply only to a specific function (e.g, the
+   function decorated by ``@app.train()``)
+
+1. Registering Application-wide Mods
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
-1. Import the required mods
-~~~~~~~~~~~~~~~~~~~~~~~~~~~
+To use application-wide mods in your ``ClientApp``, follow these steps:
 
-First, import the built-in mod you intend to use:
+Import the required mods
+++++++++++++++++++++++++
 
 .. code-block:: python
 
     import flwr as fl
     from flwr.client.mod import example_mod_1, example_mod_2
 
-2. Define your client function
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-Define your client function (``client_fn``) that will be wrapped by the mod(s):
-
-.. code-block:: python
-
-    def client_fn(cid):
-        # Your client code goes here.
-        return  # your client
-
-3. Create the ``ClientApp`` with mods
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+Create the ``ClientApp`` with application-wide mods
++++++++++++++++++++++++++++++++++++++++++++++++++++
 
 Create your ``ClientApp`` and pass the mods as a list to the ``mods`` argument. The
 order in which you provide the mods matters:
 
 .. code-block:: python
 
     app = fl.client.ClientApp(
-        client_fn=client_fn,
+        client_fn=client_fn,  # Not needed if using decorators
         mods=[
-            example_mod_1,  # Mod 1
-            example_mod_2,  # Mod 2
+            example_mod_1,  # Application-wide Mod 1
+            example_mod_2,  # Application-wide Mod 2
         ],
     )
 
-Order of execution
+If you define functions to handle messages using decorators instead of ``client_fn``,
+e.g., ``@app.train()``, you do not need to pass the ``client_fn`` argument.
+
+2. Registering Function-specific Mods
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+Instead of applying mods to the entire ``ClientApp``, you can specify them for a
+particular function:
+
+.. code-block:: python
+
+    import flwr as fl
+    from flwr.client.mod import example_mod_3, example_mod_4
+
+    app = fl.client.ClientApp()
+
+
+    @app.train(mods=[example_mod_3, example_mod_4])
+    def train(msg, ctx):
+        # Training logic here
+        return reply_msg
+
+
+    @app.evaluate()
+    def evaluate(msg, ctx):
+        # Evaluation logic here
+        return reply_msg
+
+In this case, ``example_mod_3`` and ``example_mod_4`` are only applied to the ``train``
+function.
+
+Order of Execution
 ------------------
 
-When the ``ClientApp`` runs, the mods are executed in the order they are provided in the
-list:
+When the ``ClientApp`` runs, the mods execute in the following order:
 
-1. ``example_mod_1`` (outermost mod)
-2. ``example_mod_2`` (next mod)
-3. Message handler (core function that handles the incoming ``Message`` and returns the
+1. **Application-wide mods** (executed first, in the order they are provided)
+2. **Function-specific mods** (executed after application-wide mods, in the order they
+   are provided)
+3. **ClientApp** (core function that handles the incoming ``Message`` and returns the
    outgoing ``Message``)
-4. ``example_mod_2`` (on the way back)
-5. ``example_mod_1`` (outermost mod on the way back)
+4. **Function-specific mods** (on the way back, in reverse order)
+5. **Application-wide mods** (on the way back, in reverse order)
 
 Each mod has a chance to inspect and modify the incoming ``Message`` before passing it
 to the next mod, and likewise with the outgoing ``Message`` before returning it up the
 stack.
 
+Example Execution Flow
+~~~~~~~~~~~~~~~~~~~~~~
+
+Assuming the following registration:
+
+.. code-block:: python
+
+    app = fl.client.ClientApp(mods=[example_mod_1, example_mod_2])
+
+
+    @app.train(mods=[example_mod_3, example_mod_4])
+    def train(msg, ctx):
+        return msg.create_reply(fl.common.RecordSet())
+
+
+    @app.evaluate()
+    def evaluate(msg, ctx):
+        return msg.create_reply(fl.common.RecordSet())
+
+The execution order for an incoming **train** message is as follows:
+
+1. ``example_mod_1`` (before handling)
+2. ``example_mod_2`` (before handling)
+3. ``example_mod_3`` (before handling)
+4. ``example_mod_4`` (before handling)
+5. ``train`` (handling message)
+6. ``example_mod_4`` (after handling)
+7. ``example_mod_3`` (after handling)
+8. ``example_mod_2`` (after handling)
+9. ``example_mod_1`` (after handling)
+
+The execution order for an incoming **evaluate** message is as follows:
+
+1. ``example_mod_1`` (before handling)
+2. ``example_mod_2`` (before handling)
+3. ``evaluate`` (handling message)
+4. ``example_mod_2`` (after handling)
+5. ``example_mod_1`` (after handling)
+
 Conclusion
 ----------