examples/hello-world/hello_world.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "d686a43c-7f91-465c-b09e-6d41ae92b215",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Hello World Examples\n",
    "\n",
    "Before NVIDIA FLARE 2.5, the Hello World Examples were presented in a notebook, with instructions on setting up a POC environment and then using the FLARE API to submit and follow along each job. While this can still be done, the new Job API simplifies the user experience to get started.\n",
    "\n",
    "The Job API allows for the easy creation of jobs programatically, and even allows for the ability to easily run the job with the FLARE Simulator.\n",
    "\n",
    "Previously, the jobs were in a jobs folder in each example, but since each job is now created by the Job API, you just have to follow along the Python script that creates the job in each example.  At the end of each script is a line that runs the job with a specified workspace directory:\n",
    "\n",
    "```\n",
    "job.simulator_run(\"/tmp/nvflare/jobs/workdir\")\n",
    "```\n",
    "\n",
    "Simply executing the script (for example `python fedavg_script_runner_hello-numpy.py`) will create and run the job with the FLARE Simulator.\n",
    "\n",
    "## Running with POC mode and FLARE API\n",
    "\n",
    "If you would like to run the examples in FLARE with the POC mode using the [FLARE API](../tutorials/flare_api.ipynb) like before, you can export the job to a specified location and then you can submit the job from there.\n",
    "\n",
    "```\n",
    "job.export_job(\"/tmp/nvflare/exported/job/hello-fedavg-numpy\")\n",
    "```\n",
    "\n",
    "The following instuctions in this notebook guide you through the steps to set up FLARE with POC mode and then submit jobs with the FLARE API, assuming that you have run the `export_job()` command after creating the job with the script in the folder with each example. In each `submit_job(job_folder)` command below, make sure that the job folder you are submitting matches where you exported each job to."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a9b26543",
   "metadata": {},
   "source": [
    "Each example below is self-contained. You can start from any example, but you must run through the 3 steps of each example in sequence.\n",
    "\n",
    "## Prerequisites\n",
    "Before you can run the examples here, the following preparation work must be done:\n",
    "\n",
    "1. Install a virturalenv following the instructions in [README.md](https://github.com/NVIDIA/NVFlare/tree/main/examples)\n",
    "2. Install Jupyter Lab and install a new kernel for the virtualenv called `nvflare_example`\n",
    "3. Install NVFlare following this [notebook](../nvflare_setup.ipynb)\n",
    "4. Start NVFlare in POC mode following this [notebook](../tutorials/setup_poc.ipynb). All the examples in this notebook require 2 clients to run."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f67a779d-3b54-41ab-bf19-544cdbdb79d2",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Hello World start scripts\n",
    "\n",
    "To help you quickly get started, we prepared a set of start/stop NVFLARE in POC mode scripts which capture the steps described above. These scripts must be run **from a terminal**.\n",
    "\n",
    "In the terminal, you can use the following convenience script to create and enter a ``nvflare_example`` venv:\n",
    "\n",
    "```../set_env.sh```\n",
    "\n",
    "Then making sure you are in the ``nvflare_example`` venv, you can run the ``./hw_pre_start.sh`` script to install NVFlare, provision, and start the FL system in POC mode.\n",
    "\n",
    "If the you getting errors, **do not repeatedly run ./hw_pre_start.sh**. First, you need to try to shut down NVFLARE system, using:\n",
    "\n",
    "```\n",
    "./hw_post_cleanup.sh \n",
    "```\n",
    "\n",
    "You can check if the system is shut down cleanly with:\n",
    "\n",
    "```\n",
    "ps -eaf | grep nvflare\n",
    "```\n",
    "\n",
    "If you see output like the following, then the nvflare systems are still running\n",
    "\n",
    "```\n",
    "510535    1932  1 18:54 pts/1    00:00:03 python3 -u -m nvflare.private.fed.app.client.client_train -m /tmp/workspace/example_project/prod_00/site-1/startup/.. -s fed_client.json --set secure_train=true uid=site-1 org=nvidia config_folder=config\n",
    "510539    1932  1 18:54 pts/1    00:00:03 python3 -u -m nvflare.private.fed.app.client.client_train -m /tmp/workspace/example_project/prod_00/site-2/startup/.. -s fed_client.json --set secure_train=true uid=site-2 org=nvidia config_folder=config\n",
    "510543    1932  1 18:54 pts/1    00:00:04 python3 -u -m nvflare.private.fed.app.server.server_train -m /tmp/workspace/example_project/prod_00/localhost/startup/.. -s fed_server.json --set secure_train=true org=nvidia config_folder=config\n",
    "```\n",
    "\n",
    "Make sure they are cleanly exited before you try to start the nvflare again. You can use ``nvflare poc stop`` and kill the processes if needed.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "798ee5cd-9e3b-4c09-a69c-77017942725b",
   "metadata": {},
   "source": [
    "\n",
    "## Check FL System Status\n",
    "\n",
    "If the NVFLARE system is up and running, then we are ready to check FL system status and run jobs.  \n",
    "\n",
    "**Warning**:  this step will fail if FL system is not running\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aab48af3-bd22-44ba-b3a2-61809414b252",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import time\n",
    "\n",
    "from nvflare.fuel.flare_api.flare_api import new_secure_session\n",
    "from nvflare.fuel.flare_api.flare_api import NoConnection \n",
    "\n",
    "workspace = \"/tmp/nvflare/poc\"\n",
    "default_poc_prepared_dir = os.path.join(workspace, \"example_project/prod_00\")\n",
    "admin_dir = os.path.join(default_poc_prepared_dir, \"admin@nvidia.com\")\n",
    "\n",
    "# the following try/except is usually not needed, we need it here to handle the case when you \"Run all cells\" or use notebook automation. \n",
    "# in the \"Run all cells\" case, JupyterLab seems to try to connect to the server before it starts (even though the execution is supposed to be sequential),\n",
    "# which will result in a connection timeout. We use try/except to capture the scenario since extra sleep time doesn't seem to help.\n",
    "\n",
    "try: \n",
    "   sess = new_secure_session(\"admin@nvidia.com\", admin_dir, timeout=5)\n",
    "except NoConnection:\n",
    "    time.sleep(10)\n",
    "    \n",
    "    \n",
    "flare_not_ready = True\n",
    "while flare_not_ready: \n",
    "    print(\"trying to connect to server\")\n",
    "    try:\n",
    "        sess = new_secure_session(\"admin@nvidia.com\", admin_dir)\n",
    "    except NoConnection:\n",
    "        print(\"CANNOT CONNECT AFTER 10 SECONDS\")\n",
    "        continue\n",
    "\n",
    "    sys_info = sess.get_system_info()\n",
    "\n",
    "    print(f\"Server info:\\n{sys_info.server_info}\")\n",
    "    print(\"\\nClient info\")\n",
    "    for client in sys_info.client_info:\n",
    "        print(client)\n",
    "    flare_not_ready = len( sys_info.client_info) < 2\n",
    "        \n",
    "    time.sleep(2)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a5ac45cc-ee2f-4703-9024-10af9d68ad67",
   "metadata": {},
   "source": [
    "### Utilities\n",
    "\n",
    "**Monitoring Job**\n",
    "\n",
    "You can use a custom function to format the output for the `monitor_job()` command. Here is one example function to display the job information:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "39bae42c-f054-4058-b722-388ee5059c7d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import json\n",
    "from nvflare.fuel.flare_api.flare_api import Session\n",
    "\n",
    "def status_monitor_cb(\n",
    "        session: Session, job_id: str, job_meta, *cb_args, **cb_kwargs\n",
    "    ) -> bool:\n",
    "    if job_meta[\"status\"] == \"RUNNING\":\n",
    "        if cb_kwargs[\"cb_run_counter\"][\"count\"] < 3 or cb_kwargs[\"cb_run_counter\"][\"count\"]%15 == 0:\n",
    "            print(job_meta)            \n",
    "        else:\n",
    "            # avoid printing job_meta repeatedly to save space on the screen and not overwhelm the user\n",
    "            print(\".\", end=\"\")\n",
    "    else:\n",
    "        print(\"\\n\" + str(job_meta))\n",
    "    \n",
    "    cb_kwargs[\"cb_run_counter\"][\"count\"] += 1\n",
    "    return True\n",
    "\n",
    "\n",
    "def format_json( data: dict):\n",
    "    # Helper function to format output of list_jobs()\n",
    "    print(json.dumps(data, sort_keys=True, indent=4,separators=(',', ': ')))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "86b7ca4b-d4f1-49c4-88e6-60099ec14bdd",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Hello FedAvg NumPy\n",
    "\n",
    "The `hello-world/hello-fedavg-numpy` example showcases the FedAvg workflow. See [this](https://nvflare.readthedocs.io/en/main/examples/hello_fedavg_numpy.html) for details on the example.\n",
    "\n",
    "### 1. Submit job using FLARE API\n",
    "\n",
    "Start a FLARE API session and submit the `hello-fedavg-numpy` job that you exported:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c3dbde69",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import os\n",
    "from nvflare.fuel.flare_api.flare_api import new_secure_session\n",
    "\n",
    "poc_workspace = \"/tmp/nvflare/poc\"\n",
    "poc_prepared = os.path.join(poc_workspace, \"example_project/prod_00\")\n",
    "admin_dir = os.path.join(poc_prepared, \"admin@nvidia.com\")\n",
    "sess = new_secure_session(\"admin@nvidia.com\", startup_kit_location=admin_dir)\n",
    "\n",
    "job_folder = os.path.join(os.getcwd(), \"/tmp/nvflare/exported/job/hello-fedavg-numpy\")\n",
    "job_id = sess.submit_job(job_folder)\n",
    "\n",
    "print(f\"Job is running with ID {job_id}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "42317cf3",
   "metadata": {},
   "source": [
    "### 2. Wait for the job\n",
    "\n",
    "The command `monitor_job()` will wait for the job until it is done."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ead94f38",
   "metadata": {},
   "outputs": [],
   "source": [
    "sess.monitor_job(job_id, cb=status_monitor_cb, cb_run_counter={\"count\":0})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2cf217f1",
   "metadata": {},
   "source": [
    "You can get additional information about the jobs on the server with `list_jobs()`, and you can get the information for a specific job with the `get_job_meta()` command with the ``job_id``."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c453f1e5-cf6f-4164-8d3b-b21f696ca153",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "list_jobs_output_detailed = sess.list_jobs(detailed=True)\n",
    "print(format_json(list_jobs_output_detailed))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "023584f0-9e5c-4483-98b2-fe839ac4998f",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "sess.get_job_meta(job_id)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "19a1fdfc",
   "metadata": {},
   "source": [
    "### 3. Get the result\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a9b0b73",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "result = sess.download_job_result(job_id)\n",
    "array = np.load(result + \"/workspace/models/server.npy\")\n",
    "print(array)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "93e3dbf5-5175-4a7c-80f0-04756af4a423",
   "metadata": {},
   "source": [
    "#### Clean up result directory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9e106af0-b61a-49e2-ab67-f8deba4d35d3",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "rm -r {result}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "43ade87b",
   "metadata": {},
   "source": [
    "## Hello Cross-Site Validation\n",
    "\n",
    "The `hello-world/hello-cross-val` example demonstrates how to perform cross site validation after training.\n",
    "\n",
    "Please refer to the [documentation](https://nvflare.readthedocs.io/en/main/examples/hello_cross_val.html) for the details.\n",
    "\n",
    "### 1. Submit job using FLARE API\n",
    "\n",
    "Starting a FLARE API session and submit the `hello-cross-val` job"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c2412eee-b19f-4756-b046-b70ca8250977",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import os\n",
    "from nvflare.fuel.flare_api.flare_api import new_secure_session\n",
    "\n",
    "poc_workspace = \"/tmp/nvflare/poc\"\n",
    "poc_prepared = os.path.join(poc_workspace, \"example_project/prod_00\")\n",
    "admin_dir = os.path.join(poc_prepared, \"admin@nvidia.com\")\n",
    "sess = new_secure_session(\"admin@nvidia.com\", admin_dir)\n",
    "\n",
    "job_folder = os.path.join(os.getcwd(), \"/tmp/nvflare/exported/job/hello-cross-val\")\n",
    "job_id = sess.submit_job(job_folder)\n",
    "\n",
    "print(f\"Job is running with ID {job_id}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b7dced59",
   "metadata": {},
   "source": [
    "### 2. Wait for the job"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d69769e2-a0be-48b6-a58e-aba10b382c56",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "sess.get_job_meta(job_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8bcde35c-80cd-4016-8df2-1ca9d48a2286",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "sess.monitor_job(job_id, cb=status_monitor_cb, cb_run_counter={\"count\":0})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4bfbcf71",
   "metadata": {},
   "source": [
    "### 3. Get the result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "187cb413",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import json\n",
    "import pprint\n",
    "\n",
    "result = sess.download_job_result(job_id)\n",
    "with open(result + \"/workspace/cross_site_val/cross_val_results.json\", \"r\") as f:\n",
    "  cross_val_result = json.load(f)\n",
    "\n",
    "pp = pprint.PrettyPrinter(indent=2)\n",
    "pp.pprint(cross_val_result)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3abb6a89-dc30-469a-9f4b-7b4fa287f021",
   "metadata": {},
   "source": [
    "#### Clean up result directory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ddf2cfd9-d09d-4278-9473-88477cd42c34",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "rm -r {result}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9ff5919d-7e4d-40b7-ac28-0e581fb61090",
   "metadata": {},
   "source": [
    "## Hello Cyclic Weight Transfer\n",
    "\n",
    "The `hello-world/hello-cyclic` example uses the CyclicController workflow to implement [Cyclic Weight Transfer](https://pubmed.ncbi.nlm.nih.gov/29617797/) with TensorFlow as the deep learning training framework.\n",
    "\n",
    "To use this example, tensorflow must be installed using the `requirements.txt`,\n",
    "\n",
    "    pip install -r hello-world/hello-cyclic/requirements.txt\n",
    "    \n",
    "This examples needs access to [MNIST dataset](http://yann.lecun.com/exdb/mnist/)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8d78c258-1239-4410-9200-71fabc2266f7",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "! pwd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c966687f-11f8-4b11-991e-1e7bcfbd9f55",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%pip install -r ../hello-world/hello-cyclic/requirements.txt    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d7450046-d64f-4af2-8820-e7410fbbc4dc",
   "metadata": {},
   "source": [
    "\n",
    "### 1. Submit job using FLARE API\n",
    "\n",
    "Starting a FLARE API session and submit the hello-cyclic job"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a0a01d9a-de8b-4db1-bf01-b104c9712a1c",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import os\n",
    "from nvflare.fuel.flare_api.flare_api import new_secure_session\n",
    "\n",
    "poc_workspace = \"/tmp/nvflare/poc\"\n",
    "poc_prepared = os.path.join(poc_workspace, \"example_project/prod_00\")\n",
    "admin_dir = os.path.join(poc_prepared, \"admin@nvidia.com\")\n",
    "sess = new_secure_session(\"admin@nvidia.com\", admin_dir)\n",
    "\n",
    "job_folder = os.path.join(os.getcwd(), \"/tmp/nvflare/exported/job/location/hello-cyclic\")\n",
    "job_id = sess.submit_job(job_folder)\n",
    "print(f\"Job is running with ID {job_id}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d54632ec-553a-4009-bc56-42e5c9aeb69c",
   "metadata": {},
   "source": [
    "### 2. Wait for the job"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c7146328-7136-4cdb-8dd3-7bad32d33b1e",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "sess.monitor_job(job_id)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fc6ded7e-1f28-47e2-89ca-c48de20d5162",
   "metadata": {},
   "source": [
    "### 3. Get the result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cc4bd485-fc25-4c63-ada4-3a4cf9ea3cc1",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import pprint\n",
    "from nvflare.app_opt.tf.utils import flat_layer_weights_dict\n",
    "from tensorflow.keras import layers, models\n",
    "\n",
    "result = sess.download_job_result(job_id)\n",
    "\n",
    "class Net(models.Sequential):\n",
    "    def __init__(self, input_shape=(None, 28, 28)):\n",
    "        super().__init__()\n",
    "        self._input_shape = input_shape\n",
    "        self.add(layers.Flatten())\n",
    "        self.add(layers.Dense(128, activation=\"relu\"))\n",
    "        self.add(layers.Dropout(0.2))\n",
    "        self.add(layers.Dense(10))\n",
    "\n",
    "model = Net()\n",
    "model.build(input_shape=(None, 28, 28))\n",
    "model.load_weights(result + \"/workspace/app_server/tf_model.weights.h5\")\n",
    "model.summary()\n",
    "\n",
    "layer_weights_dict = flat_layer_weights_dict({layer.name: layer.get_weights() for layer in model.layers})\n",
    "\n",
    "pp = pprint.PrettyPrinter(indent=4)\n",
    "pp.pprint(layer_weights_dict)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6046cc7f-9066-4295-bfd8-8177486274bc",
   "metadata": {},
   "source": [
    "#### Clean up result directory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c3b8710c-31e5-4b7d-ac33-ecd26fe81b6c",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "rm -r {result}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c602c834-ca96-46ed-81b8-ff33eaa24840",
   "metadata": {},
   "source": [
    "## Hello PyTorch\n",
    "\n",
    "The `hello-world/hello-pt` example demonstrates how to use NVFlare with the popular deep learning framework PyTorch.\n",
    "\n",
    "Refer to the [documentation](https://nvflare.readthedocs.io/en/main/examples/hello_pt.html) for details.\n",
    "\n",
    "To use this example, PyTorch must be installed using the `requirements.txt`,\n",
    "\n",
    "    pip install -r hello-world/hello-pt/requirements.txt\n",
    "    \n",
    "This examples also needs access to CIFAR10 dataset.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d754824a-e84f-424a-9348-9900a81e96b2",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%pip install -r ../hello-world/hello-pt/requirements.txt    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "390201f7-11ee-472d-a4f3-5499d64b89a9",
   "metadata": {},
   "source": [
    "### 1. Submit job using FLARE API\n",
    "\n",
    "Starting a FLARE API session and submit the hello-pt job"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "48a98fd8-7974-4cce-be32-3317379fe96e",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import os\n",
    "from nvflare.fuel.flare_api.flare_api import new_secure_session\n",
    "\n",
    "poc_workspace = \"/tmp/nvflare/poc\"\n",
    "poc_prepared = os.path.join(poc_workspace, \"example_project/prod_00\")\n",
    "admin_dir = os.path.join(poc_prepared, \"admin@nvidia.com\")\n",
    "sess = new_secure_session(\"admin@nvidia.com\", admin_dir)\n",
    "\n",
    "job_folder = os.path.join(os.getcwd(), \"/tmp/nvflare/exported/job/hello-pt\")\n",
    "job_id = sess.submit_job(job_folder)\n",
    "\n",
    "print(f\"Job is running with ID {job_id}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b9109755-f221-4ec6-a7f4-3fd06b88bdde",
   "metadata": {},
   "source": [
    "### 2. Wait for the job"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "128d2297-eb87-415f-ad04-ff4b51d1bff1",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "sess.monitor_job(job_id)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ccbdf205-ff82-42ae-8c89-9b7ac91e2981",
   "metadata": {},
   "source": [
    "### 3. Get the result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cf6a2808-ff53-484c-854d-16d7c00f57cc",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import pprint\n",
    "import torch\n",
    "\n",
    "print(\"this will take a bit of time\")\n",
    "result = sess.download_job_result(job_id)\n",
    "model_path = os.path.join(result, \"workspace/app_server/FL_global_model.pt\")\n",
    "\n",
    "model = torch.load(model_path)\n",
    "\n",
    "pp = pprint.PrettyPrinter(indent=4)\n",
    "pp.pprint(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "52f5b93b-e4a6-4d56-a905-b46cbaecdb20",
   "metadata": {},
   "source": [
    "#### Clean up result directory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0c18cee8-c6b3-421a-b518-7679d06e1b8a",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "rm -r {result}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c15407d8-7b6e-41df-8b03-e5caf15f1bce",
   "metadata": {},
   "source": [
    "## Hello TensorFlow\n",
    "\n",
    "The `examples/hello-world/hello-tf` example demonstrates how to use NVFlare with the popular deep learning framework TensorFlow.\n",
    "\n",
    "Refer to the [documentation](https://nvflare.readthedocs.io/en/main/examples/hello_tf.html) for details.\n",
    "\n",
    "To use this example, PyTorch must be installed using the `requirements.txt`,\n",
    "\n",
    "    python3 -m pip install -r hello-tf/requirements.txt\n",
    "    \n",
    "This examples also needs access to [MNIST dataset](http://yann.lecun.com/exdb/mnist/)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fecf1220-3cec-4d2f-88b3-0e28d6ef4ad0",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%pip install -r hello-tf/requirements.txt"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "696a384e-f6da-4044-a7b9-db4c464f52ac",
   "metadata": {},
   "source": [
    "#### Running Tensorflow on local host with GPU \n",
    "\n",
    "Before we start to run the tensorflow job, we must aware the way we are running this job. \n",
    "We are running with 1 server, 2 sites in a local machine, which means three process involved for this federated training. \n",
    "If the local host has GPU, you might enter OOM error, due to the way Tensorflow consumes GPU memory. By default, TensorFlow maps nearly all of the GPU memory of all GPUs (subject to CUDA_VISIBLE_DEVICES) visible to the process. If one has multiple process, some of the process will be OOM. To avoid multiple processes grabbing all GPU memory in TF, use the options described in [Limiting GPU memory growth]( https://www.tensorflow.org/guide/gpu#limiting_gpu_memory_growth). \n",
    "\n",
    "In our cases,  we prefer that the process only allocates a subset of the available memory, or to only grow the memory usage as is needed by the process. TensorFlow provides two methods to control this, as described in the above link.\n",
    "\n",
    "In this example, we explictly set the environment varialble `TF_FORCE_GPU_ALLOW_GROWTH` to `true` at the very beginning of the trainer.py file, which runs in the clients and will allocate GPU memory for training.  With the env var been set, TF will not grab the entire GPU memory and will not cause GPU OOM error when running POC on local host.\n",
    "\n",
    "Note that setting the env var `TF_FORCE_GPU_ALLOW_GROWTH` inside this notebook takes no effect because the clients of POC have already started and their env vars are set at the starting time.\n",
    "\n",
    "\n",
    "### 1. Submit job using FLARE API\n",
    "\n",
    "Starting a FLARE API session and submit the hello-tf2 job\n",
    "\n",
    "This time, we tail the server log"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bdf1ee4e-af12-455f-a143-082b2297c73c",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import os\n",
    "from nvflare.fuel.flare_api.flare_api import new_secure_session\n",
    "\n",
    "poc_workspace = \"/tmp/nvflare/poc\"\n",
    "poc_prepared = os.path.join(poc_workspace, \"example_project/prod_00\")\n",
    "admin_dir = os.path.join(poc_prepared, \"admin@nvidia.com\")\n",
    "sess = new_secure_session(\"admin@nvidia.com\", admin_dir)\n",
    "\n",
    "job_folder = os.path.join(os.getcwd(), \"/tmp/nvflare/exported/job/hello-tf\")\n",
    "job_id = sess.submit_job(job_folder)                          \n",
    "print(f\"Job is running with ID {job_id}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8b6b7d10-b6d2-4b2b-bc89-d36322d9cd45",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "! tail -100 /tmp/nvflare/poc/example_project/prod_00/server/log.txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2e63065c-b883-4edd-b8df-10d3813c3adb",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "list_jobs_output = sess.list_jobs()\n",
    "print(format_json(list_jobs_output))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fbc9f03e-c5b8-4530-9e77-13bbf3d5201e",
   "metadata": {},
   "source": [
    "### 2. Wait for the job"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b7c08c5f-6f49-4903-aff3-4fe908131585",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "sess.monitor_job(job_id)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7e175a52-4424-4714-9abb-b25aaf476e10",
   "metadata": {},
   "source": [
    "### 3. Get the result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "63c28eff-2d3e-4d29-958b-d82ac9622bf9",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import pprint\n",
    "from nvflare.app_opt.tf.utils import flat_layer_weights_dict\n",
    "from tensorflow.keras import layers, models\n",
    "\n",
    "result = sess.download_job_result(job_id)\n",
    "\n",
    "class Net(models.Sequential):\n",
    "    def __init__(self, input_shape=(None, 28, 28)):\n",
    "        super().__init__()\n",
    "        self._input_shape = input_shape\n",
    "        self.add(layers.Flatten())\n",
    "        self.add(layers.Dense(128, activation=\"relu\"))\n",
    "        self.add(layers.Dropout(0.2))\n",
    "        self.add(layers.Dense(10))\n",
    "\n",
    "model = Net()\n",
    "model.build(input_shape=(None, 28, 28))\n",
    "model.load_weights(result + \"/workspace/app_server/tf_model.weights.h5\")\n",
    "model.summary()\n",
    "\n",
    "layer_weights_dict = flat_layer_weights_dict({layer.name: layer.get_weights() for layer in model.layers})\n",
    "\n",
    "pp = pprint.PrettyPrinter(indent=4)\n",
    "pp.pprint(layer_weights_dict)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dc111a72-d7ba-4f3f-af23-9632c7be38a3",
   "metadata": {},
   "source": [
    "#### Clean up result directory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "58b215d5-d888-4941-9a90-7cdbff6f1a34",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "rm -r {result}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b3ce5164-b32a-450a-867a-ffc36a2efa7f",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Cleanup\n",
    "We need to shutdown NVFLARE system and clean up POC workspace. This can be down in the following steps. \n",
    "you can change the cell into the code cell from markdown cell.\n",
    "or you can simply execute from a **terminal**\n",
    "\n",
    "```hw_post_cleanup.sh```\n",
    "\n",
    "bash shudown script basically does the followings\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "97748fb1-b465-4bfc-aefb-000f37f69141",
   "metadata": {
    "tags": []
   },
   "source": [
    "```! nvflare poc stop```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c1270bd8-f432-4be6-8b2f-a1ef50b0f1f0",
   "metadata": {
    "tags": []
   },
   "source": [
    "```! nvflare poc clean```\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "nvflare_example",
   "language": "python",
   "name": "nvflare_example"
  },
  "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.8.15"
  },
  "vscode": {
   "interpreter": {
    "hash": "31f2aee4e71d21fbe5cf8b01ff0e069b9275f58929596ceb00d14d90e3e16cd6"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}