examples/hello-world/hello-fedavg-numpy/hello-fedavg-numpy_getting_started.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "7a5c3d67-a6ea-4f59-84d2-effc3ef016e1",
   "metadata": {},
   "source": [
    "# Getting Started with NVFlare (Numpy)\n",
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/NVIDIA/NVFlare/blob/main/examples/getting_started/pt/nvflare_pt_getting_started.ipynb)\n",
    "\n",
    "NVFlare is an open-source framework that allows researchers and\n",
    "data scientists to seamlessly move their machine learning and deep\n",
    "learning workflows into a federated paradigm."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fcf2b4a8-ed42-421d-8898-c0c93f9d8a09",
   "metadata": {},
   "source": [
    "## Basic Concepts\n",
    "At the heart of NVFlare lies the concept of collaboration through\n",
    "\"tasks.\" An FL controller assigns tasks (e.g., training on local data) to one or more FL clients, processes returned\n",
    "results (e.g., model weight updates), and may assign additional\n",
    "tasks based on these results and other factors (e.g., a pre-configured\n",
    "number of training rounds). The clients run executors which can listen for tasks and perform the necessary computations locally, such as model training. This task-based interaction repeats\n",
    "until the experiment’s objectives are met. \n",
    "\n",
    "We can also add data filters (for example, for [homomorphic encryption](https://www.usenix.org/conference/atc20/presentation/zhang-chengliang) or [differential privacy filters](https://arxiv.org/abs/1910.00962)) to the task data or results received or produced by the server or clients.\n",
    "\n",
    "<img src=\"../../../docs/resources/nvflare_overview.svg\" alt=\"NVIDIA FLARE Overview\" width=75% height=75% />"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "907933a8-20fd-4aa7-a3bf-3f5b5829a544",
   "metadata": {},
   "source": [
    "## Setup environment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33bba668-72ac-4e69-aaed-8d4254f547c0",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -q -r nvflare~=2.5.0rc"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b68cb248-dc6a-48d1-880d-33c4324d9723",
   "metadata": {},
   "source": [
    "## Federated Averaging with NVFlare\n",
    "Given the flexible controller and executor concepts, it is easy to implement different computing & communication patterns with NVFlare, such as [FedAvg](https://proceedings.mlr.press/v54/mcmahan17a?ref=https://githubhelp.com). \n",
    "\n",
    "The controller's `run()` routine is responsible for assigning tasks and processing task results from the Executors. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b2f84fb1-9dd3-4c72-a727-c4614260f02f",
   "metadata": {},
   "source": [
    "### Server Code\n",
    "First, we use a simple implementation of the [FedAvg](https://proceedings.mlr.press/v54/mcmahan17a?ref=https://githubhelp.com) algorithm with NVFlare. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d62a13d5-1130-44e6-8818-70e30de401e6",
   "metadata": {},
   "source": [
    "```python\n",
    "class FedAvg(BaseFedAvg):\n",
    "    def run(self) -> None:\n",
    "        self.info(\"Start FedAvg.\")\n",
    "\n",
    "        model = self.load_model()\n",
    "        model.start_round = self.start_round\n",
    "        model.total_rounds = self.num_rounds\n",
    "\n",
    "        for self.current_round in range(self.start_round, self.start_round + self.num_rounds):\n",
    "            self.info(f\"Round {self.current_round} started.\")\n",
    "            model.current_round = self.current_round\n",
    "\n",
    "            clients = self.sample_clients(self.num_clients)\n",
    "\n",
    "            results = self.send_model_and_wait(targets=clients, data=model)\n",
    "\n",
    "            aggregate_results = self.aggregate(results)\n",
    "\n",
    "            model = self.update_model(model, aggregate_results)\n",
    "\n",
    "            self.save_model(model)\n",
    "\n",
    "        self.info(\"Finished FedAvg.\")\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d24b6476-089a-4e9d-825b-07107bd5d84a",
   "metadata": {},
   "source": [
    "### Client Code \n",
    "For this numpy example, we just have mock training by adding one to the data, so there are no deep learning concepts and we can focus on the flow of NVFlare."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1c551053-5460-4d83-8578-796074170342",
   "metadata": {},
   "source": [
    "On the client side, the training workflow is as follows:\n",
    "1. Receive the model from the FL server (for this example we initialize the model in the client code to the numpy array [[1, 2, 3], [4, 5, 6], [7, 8, 9]] if the model params are empty).\n",
    "2. Perform local training on the received global model\n",
    "and/or evaluate the received global model for model\n",
    "selection.\n",
    "3. Send the new model back to the FL server."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c02bfc2a-783c-494f-9427-c38f40a2e870",
   "metadata": {},
   "source": [
    "Using NVFlare's client API, we can easily adapt machine learning code that was written for centralized training and apply it in a federated scenario.\n",
    "For a general use case, there are three essential methods to achieve this using the Client API :\n",
    "- `init()`: Initializes NVFlare Client API environment.\n",
    "- `receive()`: Receives model from the FL server.\n",
    "- `send()`: Sends the model to the FL server."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9115ee07-d848-4a7c-99ad-64e20ab7093c",
   "metadata": {},
   "source": [
    "With these simple methods, the developers can use the Client API\n",
    "to change their centralized training code to an FL scenario with\n",
    "five lines of code changes as shown below.\n",
    "\n",
    "```python\n",
    "    import nvflare.client as flare\n",
    "    \n",
    "    flare.init() # 1. Initializes NVFlare Client API environment.\n",
    "    input_model = flare.receive() # 2. Receives model from the FL server.\n",
    "    params = input_model.params # 3. Obtain the required information from the received model.\n",
    "    \n",
    "    # original local training code\n",
    "    new_params = local_train(params)\n",
    "    \n",
    "    output_model = flare.FLModel(params=new_params) # 4. Put the results in a new `FLModel`\n",
    "    flare.send(output_model) # 5. Sends the model to the FL server.  \n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "67432f44-4144-4347-8d74-e7f57e065a14",
   "metadata": {},
   "source": [
    "The full client training script is saved in a separate file, e.g. [./src/hello-numpy_fl.py](./src/hello-numpy_fl.py)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5da34414-bac4-4352-8077-ab7ade998eec",
   "metadata": {},
   "source": [
    "## Run an NVFlare Job\n",
    "Now that we have defined the FedAvg controller to run our workflow on the FL server and our client training script to receive the global models, run local training, and send the results back to the FL server, we can put everything together using NVFlare's Job API."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1b70da5d-ba8b-4e65-b47f-44bb9bddae4d",
   "metadata": {},
   "source": [
    "#### 1. Define a FedJob\n",
    "The `FedJob` is used to define how controllers and executors are placed within a federated job using the `to(object, target)` routine."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13771bfb-901f-485a-9a23-84db1ccd5fe4",
   "metadata": {},
   "outputs": [],
   "source": [
    "from nvflare import FedJob\n",
    "from nvflare.app_common.workflows.fedavg import FedAvg\n",
    "from nvflare.job_config.script_runner import FrameworkType, ScriptRunner\n",
    "\n",
    "\n",
    "job = FedJob(name=\"hello-fedavg-numpy\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9361d9f8-54f3-4363-b3ba-706a7ae3a8e9",
   "metadata": {},
   "source": [
    "#### 2. Define the Controller Workflow\n",
    "Define the controller workflow with the persistor and send to server."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6962e6cc-995e-4356-8156-3ceba2c7a249",
   "metadata": {},
   "outputs": [],
   "source": [
    "from nvflare.app_common.np.np_model_persistor import NPModelPersistor\n",
    "\n",
    "n_clients = 2\n",
    "num_rounds = 3\n",
    "\n",
    "persistor_id = job.to_server(NPModelPersistor(), \"persistor\")\n",
    "\n",
    "controller = FedAvg(\n",
    "    num_clients=n_clients,\n",
    "    num_rounds=num_rounds,\n",
    "    persistor_id=persistor_id,\n",
    ")\n",
    "job.to_server(controller)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4171fe9e",
   "metadata": {},
   "source": [
    "#### 3. Add ModelSelector\n",
    "Add IntimeModelSelector for global best model selection."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fb7b93a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "from nvflare.app_common.widgets.intime_model_selector import IntimeModelSelector\n",
    "\n",
    "job.to(IntimeModelSelector(key_metric=\"accuracy\"), \"server\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "548966c2-90bf-47ad-91d2-5c6c22c3c4f0",
   "metadata": {},
   "source": [
    "#### 3. Add Clients\n",
    "Next, we can use the `ScriptRunner` and send it to each of the clients to run our training script.\n",
    "\n",
    "Note that our script could have additional input arguments, such as batch size or data path, but we don't use them here for simplicity.\n",
    "We can also specify, which GPU should be used to run this client, which is helpful for simulated environments."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ad5d36fe-9ae5-43c3-80bc-2cdc66bf7a7e",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_script = \"src/hello-numpy_fl.py\"\n",
    "\n",
    "for i in range(n_clients):\n",
    "    executor = ScriptRunner(\n",
    "        script=train_script, script_args=\"\", framework=FrameworkType.NUMPY\n",
    "    )\n",
    "    job.to(executor, f\"site-{i}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "113fd6af-85be-4f75-8a8e-4666771252b3",
   "metadata": {},
   "source": [
    "That's it!\n",
    "\n",
    "#### 4. Optionally export the job\n",
    "Now, we could export the job and submit it to a real NVFlare deployment using the [Admin client](https://nvflare.readthedocs.io/en/main/real_world_fl/operation.html) or [FLARE API](https://nvflare.readthedocs.io/en/main/real_world_fl/flare_api.html). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "99a270bf-c906-425b-b999-2306cb76eb62",
   "metadata": {},
   "outputs": [],
   "source": [
    "job.export_job(\"/tmp/nvflare/jobs/job_config\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9ac3f0a8-06bb-4bea-89d3-4a5fc5b76c63",
   "metadata": {},
   "source": [
    "#### 5. Run FL Simulation\n",
    "Finally, we can run our FedJob in simulation using NVFlare's [simulator](https://nvflare.readthedocs.io/en/main/user_guide/nvflare_cli/fl_simulator.html) under the hood. The results will be saved in the specified `workdir`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13068ab7-35cf-49e7-91ed-10993049ef0d",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "job.simulator_run(\"/tmp/nvflare/jobs/workdir\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}