examples/hello-world/hello_world.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "d686a43c-7f91-465c-b09e-6d41ae92b215",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Hello World Examples\n",
    "\n",
    "In this notebook, we will walk you through some Hello World examples in `NVFlare/examples/hello-world` to get familiar with basic workflow of NVIDIA FLARE.\n",
    "\n",
    "We will run the examples in FLARE with the POC mode using the [FLARE API](../tutorials/flare_api.ipynb). You can also run these examples in the [FLARE simulator](../tutorials/flare_simulator.ipynb).\n",
    "\n",
    "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.\n",
    "\n"
   ]
  },
  {
   "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 Scatter and Gather\n",
    "\n",
    "The example job in `hello-world/hello-numpy-sag/jobs/hello-numpy-sag` showcases the scatter and gather workflow. See [this](https://nvflare.readthedocs.io/en/main/examples/hello_scatter_and_gather.html#hello-scatter-and-gather) for details on the example.\n",
    "\n",
    "### 1. Submit job using FLARE API\n",
    "\n",
    "Start a FLARE API session and submit the `hello-numpy-sag` job:"
   ]
  },
  {
   "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(), \"hello-numpy-sag/jobs/hello-numpy-sag\")\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 example job in `hello-world/hello-numpy-cross/jobs/hello-numpy-cross` 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-numpy-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(), \"hello-numpy-cross-val/jobs/hello-numpy-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",
    "This 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. The job is `hello-world/hello-cyclic/jobs/hello-cyclic`.\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(), \"hello-cyclic/jobs/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": [
    "from nvflare.fuel.utils import fobs\n",
    "from nvflare.app_common.decomposers import common_decomposers\n",
    "import pprint\n",
    "\n",
    "# This example stores numpy arrays in FOBS format. Decomposers for Numpy is not registered automatically.\n",
    "common_decomposers.register()\n",
    "\n",
    "result = sess.download_job_result(job_id)\n",
    "with open(result + \"/workspace/app_server/tf2weights.fobs\", \"rb\") as f:\n",
    "    bytes = f.read()\n",
    "\n",
    "weights = fobs.loads(bytes)\n",
    "\n",
    "pp = pprint.PrettyPrinter(indent=4)\n",
    "pp.pprint(weights)"
   ]
  },
  {
   "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",
    "This example demonstrates how to use NVFlare with the popular deep learning framework PyTorch. The job is `hello-world/hello-pt/jobs/hello-pt`.\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(), \"hello-pt/jobs/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 2\n",
    "\n",
    "This example demonstrates how to use NVFlare with the popular deep learning framework TensorFlow 2. The job is `examples/hello-world/hello-tf2/jobs/hello-tf2`.\n",
    "\n",
    "Refer to the [documentation](https://nvflare.readthedocs.io/en/main/examples/hello_tf2.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-tf2/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-tf2/requirements.txt"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "425c7f1d-7cb6-4602-bb88-4b87ff517529",
   "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. \n",
    "\n",
    "The First method is set the environmental variable TF_FORCE_GPU_ALLOW_GROWTH to true. This configuration is platform specific. \n",
    "The 2nd method is using the piece of code below"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "026ce33a-90b1-4ef7-8c7c-f25722f2d2ae",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%env TF_FORCE_GPU_ALLOW_GROWTH=true"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "91cf84ec-57f1-439f-b72b-b33fea1a7f7c",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "gpus = tf.config.list_physical_devices('GPU')\n",
    "if gpus:\n",
    "  # Restrict TensorFlow to only allocate 1GB of memory on the first GPU\n",
    "  try:\n",
    "    tf.config.set_logical_device_configuration(\n",
    "        gpus[0],\n",
    "        [tf.config.LogicalDeviceConfiguration(memory_limit=1024)])\n",
    "    logical_gpus = tf.config.list_logical_devices('GPU')\n",
    "    print(len(gpus), \"Physical GPUs,\", len(logical_gpus), \"Logical GPUs\")\n",
    "  except RuntimeError as e:\n",
    "    # Virtual devices must be set before GPUs have been initialized\n",
    "    print(e)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0c3f1149-d54d-4f3b-b497-c06deba22fef",
   "metadata": {},
   "source": [
    "### 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(), \"hello-tf2/jobs/hello-tf2\")\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/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": [
    "from nvflare.fuel.utils import fobs\n",
    "from nvflare.app_common.decomposers import common_decomposers\n",
    "import pprint\n",
    "\n",
    "common_decomposers.register()\n",
    "result = sess.download_job_result(job_id)\n",
    "with open(result + \"/workspace/app_server/tf2weights.fobs\", \"rb\") as f:\n",
    "    bytes = f.read()\n",
    "\n",
    "weights = fobs.loads(bytes)\n",
    "\n",
    "pp = pprint.PrettyPrinter(indent=4)\n",
    "pp.pprint(weights)"
   ]
  },
  {
   "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.17"
  },
  "vscode": {
   "interpreter": {
    "hash": "31f2aee4e71d21fbe5cf8b01ff0e069b9275f58929596ceb00d14d90e3e16cd6"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}