examples/hello-world/step-by-step/higgs/xgboost/xgboost_horizontal.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0ccbfaee-9efe-4fe1-bab6-60462769fede",
   "metadata": {},
   "source": [
    "# Federated Horizontal XGBoost with Tree-based Collaboration "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1a5ea300-8fd1-4568-8121-9bb9d3d303b7",
   "metadata": {
    "tags": []
   },
   "source": [
    "This tutorial illustrates a federated horizontal xgboost learning on tabular data with bagging collaboration. \n",
    "\n",
    "Before do the training, we need to setup NVFLARE\n",
    "\n",
    "## Setup NVFLARE\n",
    "\n",
    "Follow [Getting Started](https://nvflare.readthedocs.io/en/main/getting_started.html) to set up a virtual environment and install NVFLARE.\n",
    "\n",
    "You can also follow this [notebook](https://github.com/NVIDIA/NVFlare/blob/main/examples/nvflare_setup.ipynb) to get set up.\n",
    "\n",
    "> Make sure you have installed nvflare from **terminal** \n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "162c1eb8-3070-45b8-a83c-f16d060d0385",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Install requirements\n",
    "assuming the current directory is '/examples/hello-world/step-by-step/higgs/xgboost'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b3f5a8f-cd09-441e-a4c3-f1a121fd8244",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!pwd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a7bf5e9-1f21-430a-9227-3e912eba5f4c",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%pip install -r requirements.txt"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5f87c3da-59dc-4551-9448-b20b64a57137",
   "metadata": {},
   "source": [
    "\n",
    "## Prepare data\n",
    "Please reference [prepare_higgs_data](../prepare_data.ipynb) notebooks. Pay attention to the current location. You need to switch \"higgs\" directory to run the data split.\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bfe5bc9a-d589-4d48-89b8-73fc1d8fe44b",
   "metadata": {
    "tags": []
   },
   "source": [
    "Now we have our data prepared. we are ready to do the training"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8ddc62f2-5571-44d0-bb8b-b6a660a500c5",
   "metadata": {},
   "source": [
    "### Data Cleaning \n",
    "\n",
    "We noticed from time-to-time the Higgs dataset is making small changes which causing job to fail. so we need to do some clean up or skip certain rows. \n",
    "For example: certain floating number mistakenly add an alphabetical letter at some point of time. This may have already fixed by UCI. \n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "64dbec4e-b4a9-44fe-954e-62f1d0f99c40",
   "metadata": {},
   "source": [
    "## XGBoost\n",
    "This tutorial uses [XGBoost](https://github.com/dmlc/xgboost), which is an optimized distributed gradient boosting library."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c80eb797-f1aa-4537-ab50-4c27fb9cff60",
   "metadata": {},
   "source": [
    "### Federated XGBoost Model\n",
    "Here we use tree-based collaboration for horizontal federated XGBoost.\n",
    "\n",
    "Under this setting, individual trees are independently trained on each client's local data without aggregating the global sample gradient histogram information.\n",
    "Trained trees are collected and passed to the server / other clients for bagging aggregation and further boosting rounds.\n",
    "\n",
    "The XGBoost Booster api is leveraged to create in-memory Booster objects that persist across rounds to cache predictions from trees added in previous rounds and retain other data structures needed for training.\n",
    "\n",
    "Let's look at the code see how we convert the local training script to the federated training script."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "70b5d82b-3c5a-492a-bc92-86ca7367ff92",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!pwd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1fdbcc96-3ce0-4e9e-a986-f877e082c0ee",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!cat code/xgboost_fl.py"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1d932e8c-7c7c-470b-8448-43288feba261",
   "metadata": {
    "tags": []
   },
   "source": [
    "The code is pretty much like the local xgboost training script of `code/xgboost_local_oneshot.py`, and its derivative `code/xgboost_local_iter.py`.\n",
    "\n",
    "#### load data\n",
    "\n",
    "We first load the features from the header file: \n",
    "    \n",
    "```\n",
    "    site_name = flare.get_site_name()\n",
    "    feature_data_path = f\"{data_root_dir}/{site_name}_header.csv\"\n",
    "    features = load_features(feature_data_path)\n",
    "    n_features = len(features) -1\n",
    "\n",
    "    data_path = f\"{data_root_dir}/{site_name}.csv\"\n",
    "    data = load_data(data_path=data_path, data_features=features, test_size=test_size, skip_rows=skip_rows)\n",
    "\n",
    "```\n",
    "\n",
    "then load the data from the main csv file, then transform the data and split the training and test data based on the test_size provided.  \n",
    "\n",
    "```\n",
    "    data = to_dataset_tuple(data)\n",
    "    dataset = transform_data(data)\n",
    "    x_train, y_train, train_size = dataset[\"train\"]\n",
    "    x_test, y_test, test_size = dataset[\"test\"]\n",
    "\n",
    "```\n",
    "\n",
    "The part that's specific to Federated Learning is in the following codes\n",
    "\n",
    "```\n",
    "# (1) import nvflare client API\n",
    "from nvflare import client as flare\n",
    "\n",
    "```\n",
    "```\n",
    "# (2) initializes NVFlare client API\n",
    "    flare.init()\n",
    "\n",
    "    site_name = flare.get_site_name()\n",
    "    \n",
    "```\n",
    "    \n",
    "These few lines, import NVFLARE Client API and initialize it, then use the API to find the site_name (such as site-1, site-2 etc.). With the site-name, we can construct the site-specific \n",
    "data path such as\n",
    "\n",
    "```\n",
    "    feature_data_path = f\"{data_root_dir}/{site_name}_header.csv\"\n",
    "\n",
    "    data_path = f\"{data_root_dir}/{site_name}.csv\"\n",
    "```\n",
    "\n",
    "#### Training \n",
    "\n",
    "In the standard traditional xgboost, we would train the model such as\n",
    "```\n",
    "  model = xgb.train(...) \n",
    "```\n",
    "\n",
    "with federated learning, using FLARE Client API, we need to make a few changes\n",
    "* 1) we are not only training in local iterations, but also global rounds, we need to keep the program running until we reached to the totoal number of rounds \n",
    "  \n",
    "  ```\n",
    "      while flare.is_running():\n",
    "          ... rest of code\n",
    "  \n",
    "  ```\n",
    "  \n",
    "* 2) Unlike local learning, we have now have more than one clients/sites participating the training. To ensure every site starts with the same model parameters, we use server to broadcast the initial model parameters to every sites at the first round ( current_round = 0). \n",
    "\n",
    "* 3) We will need to use FLARE client API to receive global model updates \n",
    "\n",
    "```\n",
    "        # (3) receives FLModel from NVFlare\n",
    "        input_model = flare.receive()\n",
    "        global_params = input_model.params\n",
    "        curr_round = input_model.current_round\n",
    "```\n",
    "\n",
    "```\n",
    "        if curr_round == 0:\n",
    "            # (4) first round, no global model\n",
    "            model = xgb.train(\n",
    "                xgb_params,\n",
    "                dmat_train,\n",
    "                num_boost_round=1,\n",
    "                evals=[(dmat_train, \"train\"), (dmat_test, \"test\")],\n",
    "            )\n",
    "            config = model.save_config()\n",
    "        ....\n",
    "```\n",
    "* 4) if it is not the first round, we need to use the global model to update the local model before training the next round. \n",
    "\n",
    "```\n",
    "            # (5) update model based on global updates\n",
    "            model_updates = global_params[\"model_data\"]\n",
    "            for update in model_updates:\n",
    "                global_model_as_dict = update_model(global_model_as_dict, json.loads(update))\n",
    "            loadable_model = bytearray(json.dumps(global_model_as_dict), \"utf-8\")\n",
    "            # load model\n",
    "            model.load_model(loadable_model)\n",
    "            model.load_config(config)\n",
    "```\n",
    "\n",
    "* 5) we then evaluate the global model using the local data\n",
    "\n",
    "```\n",
    "            # (6) evaluate model\n",
    "            auc = evaluate_model(x_test, model, y_test)\n",
    "```\n",
    "* 6) finally we do the training \n",
    "\n",
    "```\n",
    "            # train model in two steps\n",
    "            # first, eval on train and test\n",
    "            eval_results = model.eval_set(\n",
    "                evals=[(dmat_train, \"train\"), (dmat_test, \"test\")], iteration=model.num_boosted_rounds() - 1\n",
    "            )\n",
    "            print(eval_results)\n",
    "            # second, train for one round\n",
    "            model.update(dmat_train, model.num_boosted_rounds())\n",
    "        \n",
    "```\n",
    "\n",
    "* 7) we need the new training result (new tree) back to server for aggregation, to do that, we have the following code\n",
    "\n",
    "```\n",
    "        # (7) construct trained FL model\n",
    "        # Extract newly added tree using xgboost_bagging slicing api\n",
    "        bst_new = model[model.num_boosted_rounds() - 1 : model.num_boosted_rounds()]\n",
    "        local_model_update = bst_new.save_raw(\"json\")\n",
    "        params = {\"model_data\": local_model_update}\n",
    "        metrics = {\"accuracy\": auc}\n",
    "        output_model = flare.FLModel(params=params, metrics=metrics)\n",
    "\n",
    "        # (8) send model back to NVFlare\n",
    "        flare.send(output_model)\n",
    "```\n",
    "\n",
    "## Prepare Job  \n",
    "\n",
    "Now, we have the code, we need to prepare job folder with configurations to run in NVFLARE. To do this, we can leveage the job template for scikit learn. First look at the the available job templates"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "01fdd962-e63c-4b58-81e7-beeedd05509b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!nvflare config -jt ../../../../../job_templates/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a5ed150d-4692-49fe-87a6-1779ec64d9df",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!nvflare job list_templates"
   ]
  },
  {
   "cell_type": "raw",
   "id": "2b702033-6084-4002-8306-2d9aeab901b9",
   "metadata": {},
   "source": [
    "the `xgboost_tree` is the one we need. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9899a684-444f-4c04-a388-4bdd9b7dc649",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!nvflare job create -j /tmp/nvflare/jobs/xgboost -force -w xgboost_tree \\\n",
    "-sd code \\\n",
    "-f config_fed_client.conf app_script=\"xgboost_fl.py\" app_config=\"--data_root_dir /tmp/nvflare/dataset/output\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17dc5dcc-573e-479e-965d-0cca8b58995a",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!cat /tmp/nvflare/jobs/xgboost/app/config/config_fed_client.conf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7b3086fd-ecf3-4584-a643-5706ae2069b5",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!tree /tmp/nvflare/jobs/xgboost"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6111b784-1722-4aff-a48b-fc50c6c67281",
   "metadata": {},
   "source": [
    ">Note \n",
    " For potential on-the-fly data cleaning, we use skip_rows = 0 to skip 1st row. We could skip_rows = [0, 3] to skip first and 4th rows.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "488ff6b6-5e39-4636-8a48-44d2b503dcdb",
   "metadata": {},
   "source": [
    "\n",
    "## Run job in simulator\n",
    "\n",
    "We use the simulator to run this job"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "09800691-46e0-4a95-bea6-d2a4a425052b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!nvflare simulator /tmp/nvflare/jobs/xgboost -w /tmp/nvflare/xgboost -n 3 -t 3"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ebe301fa-8f51-4d3f-b61e-fd851b9f92d2",
   "metadata": {
    "tags": []
   },
   "source": [
    "Let's examine the results.\n",
    "\n",
    "We can notice from the FL training log, at the last round (`current_round=99`) of local training, site-1 reports `site-1: global model AUC: 0.82085`\n",
    "Under tree-based collaboration, the newly boosted trees from all clients will be bagged together to be added to the local models, hence, at the end of our 100-round training with 3 clients, the global model should have 300 trees in total. Let's double check if this holds:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fa6f85dd-ee68-41e3-975e-70dea2103b00",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "# Specify the global model file path\n",
    "file_path =  '/tmp/nvflare/xgboost/simulate_job/app_server/xgboost_model.json'\n",
    "\n",
    "with open(file_path) as json_data:\n",
    "    model = json.load(json_data)\n",
    "    print(f\"num_trees: {model['learner']['gradient_booster']['model']['gbtree_model_param']['num_trees']}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4d1fd077-a5bb-4135-a2c9-fd2c60db544b",
   "metadata": {
    "tags": []
   },
   "source": [
    "And yes, we do have 300 trees for the global model as expected.\n",
    "\n",
    "Now let's run some local trainings to verify if this number makes sense."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7f2fac8a-f3a6-467b-a6d9-b93f3c2ad4c4",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!python3 ./code/xgboost_local_oneshot.py --data_root_dir /tmp/nvflare/dataset/output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8009591f-83c4-4c84-be79-176573296a6a",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!python3 ./code/xgboost_local_iter.py --data_root_dir /tmp/nvflare/dataset/output"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ea7bbacc-b059-4f82-9785-2b22bf840ef9",
   "metadata": {},
   "source": [
    "Both oneshot and iterative training schemes yield idential results of `local model AUC: 0.81928`. As compared with FL's `global model AUC: 0.82085`, we can notice FL gives some benefits, even under homogeneous data distribution across clients.\n",
    "\n",
    "## We are done !\n",
    "Congratulations! you have just completed the federated xgboost model for tabular data. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b5329a14-1a81-488a-8edc-2b7c85cffd8c",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}