mortgage_e2e.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Mortgage Workflow\n",
    "\n",
    "## The Dataset\n",
    "The dataset used with this workflow is derived from [Fannie Mae’s Single-Family Loan Performance Data](http://www.fanniemae.com/portal/funding-the-market/data/loan-performance-data.html) with all rights reserved by Fannie Mae. This processed dataset is redistributed with permission and consent from Fannie Mae.\n",
    "\n",
    "To acquire this dataset, please visit [RAPIDS Datasets Homepage](https://docs.rapids.ai/datasets/mortgage-data)\n",
    "\n",
    "## Introduction\n",
    "The Mortgage workflow is composed of three core phases:\n",
    "\n",
    "1. ETL - Extract, Transform, Load\n",
    "2. Data Conversion\n",
    "3. ML - Training\n",
    "\n",
    "### ETL\n",
    "Data is \n",
    "1. Read in from storage\n",
    "2. Transformed to emphasize key features\n",
    "3. Loaded into volatile memory for conversion\n",
    "\n",
    "### Data Conversion\n",
    "Features are\n",
    "1. Broken into (labels, data) pairs\n",
    "2. Distributed across many workers\n",
    "3. Converted into compressed sparse row (CSR) matrix format for XGBoost\n",
    "\n",
    "### Machine Learning\n",
    "The CSR data is fed into a distributed training session with `xgboost.dask`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "If required, the notebook can be converted to a python script for execution using tools like `nbconvert`\n",
    "\n",
    "```sh\n",
    "$ jupyter nbconvert --to python mortgage_e2e.ipynb\n",
    "$ python mortgage_e2e.py\n",
    "```\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Imports statements"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from utils.utils import (\n",
    "    determine_dataset,\n",
    "    get_data,\n",
    "    memory_info,\n",
    ")\n",
    "\n",
    "from dask_cuda import LocalCUDACluster\n",
    "from dask.delayed import delayed\n",
    "from dask.distributed import Client, wait\n",
    "import rmm\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "from collections import OrderedDict\n",
    "import argparse\n",
    "import gc\n",
    "from glob import glob\n",
    "import os\n",
    "import subprocess\n",
    "import time"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Define functions to encapsulate the workflow into a single call"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_dask_task(func, **kwargs):\n",
    "    task = func(**kwargs)\n",
    "    return task\n",
    "\n",
    "\n",
    "def process_quarter_gpu(\n",
    "    year=2000, quarter=1, perf_file=\"\", data_dir=\"\", client=None, **kwargs\n",
    "):\n",
    "    ml_arrays = run_dask_task(\n",
    "        delayed(run_gpu_workflow), quarter=quarter, year=year, perf_file=perf_file\n",
    "    )\n",
    "    return client.compute(ml_arrays, optimize_graph=False, fifo_timeout=\"0ms\")\n",
    "\n",
    "\n",
    "def run_gpu_workflow(\n",
    "    quarter=1, year=2000, perf_file=\"\", acq_file=\"\", names_file=\"\", **kwargs\n",
    "):\n",
    "    names = gpu_load_names(col_names_path=data_dir + \"names.csv\")\n",
    "    names = hash_df_string_columns(names)\n",
    "    acq_gdf = gpu_load_acquisition_csv(\n",
    "        acquisition_path=data_dir\n",
    "        + \"acq\"\n",
    "        + \"/Acquisition_\"\n",
    "        + str(year)\n",
    "        + \"Q\"\n",
    "        + str(quarter)\n",
    "        + \".txt\"\n",
    "    )\n",
    "    acq_gdf = hash_df_string_columns(acq_gdf)\n",
    "    acq_gdf = acq_gdf.merge(names, how=\"left\", on=[\"seller_name\"])\n",
    "    acq_gdf[\"seller_name\"] = acq_gdf[\"new\"]\n",
    "    acq_gdf.drop(columns=\"new\", inplace=True)\n",
    "    perf_df_tmp = gpu_load_performance_csv(perf_file)\n",
    "    perf_df_tmp = hash_df_string_columns(perf_df_tmp)\n",
    "    gdf = perf_df_tmp\n",
    "    everdf = create_ever_features(gdf)\n",
    "    delinq_merge = create_delinq_features(gdf)\n",
    "    everdf = join_ever_delinq_features(everdf, delinq_merge)\n",
    "    del delinq_merge\n",
    "    joined_df = create_joined_df(gdf, everdf)\n",
    "    testdf = create_12_mon_features(joined_df)\n",
    "    joined_df = combine_joined_12_mon(joined_df, testdf)\n",
    "    del testdf\n",
    "    perf_df = final_performance_delinquency(gdf, joined_df)\n",
    "    del (gdf, joined_df)\n",
    "    final_gdf = join_perf_acq_gdfs(perf_df, acq_gdf)\n",
    "    del perf_df\n",
    "    del acq_gdf\n",
    "    final_gdf = last_mile_cleaning(final_gdf)\n",
    "    return final_gdf\n",
    "\n",
    "\n",
    "def gpu_load_performance_csv(performance_path, **kwargs):\n",
    "    \"\"\" \n",
    "    Loads performance data\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    GPU DataFrame\n",
    "    \"\"\"\n",
    "\n",
    "    cols = [\n",
    "        \"loan_id\",\n",
    "        \"monthly_reporting_period\",\n",
    "        \"servicer\",\n",
    "        \"interest_rate\",\n",
    "        \"current_actual_upb\",\n",
    "        \"loan_age\",\n",
    "        \"remaining_months_to_legal_maturity\",\n",
    "        \"adj_remaining_months_to_maturity\",\n",
    "        \"maturity_date\",\n",
    "        \"msa\",\n",
    "        \"current_loan_delinquency_status\",\n",
    "        \"mod_flag\",\n",
    "        \"zero_balance_code\",\n",
    "        \"zero_balance_effective_date\",\n",
    "        \"last_paid_installment_date\",\n",
    "        \"foreclosed_after\",\n",
    "        \"disposition_date\",\n",
    "        \"foreclosure_costs\",\n",
    "        \"prop_preservation_and_repair_costs\",\n",
    "        \"asset_recovery_costs\",\n",
    "        \"misc_holding_expenses\",\n",
    "        \"holding_taxes\",\n",
    "        \"net_sale_proceeds\",\n",
    "        \"credit_enhancement_proceeds\",\n",
    "        \"repurchase_make_whole_proceeds\",\n",
    "        \"other_foreclosure_proceeds\",\n",
    "        \"non_interest_bearing_upb\",\n",
    "        \"principal_forgiveness_upb\",\n",
    "        \"repurchase_make_whole_proceeds_flag\",\n",
    "        \"foreclosure_principal_write_off_amount\",\n",
    "        \"servicing_activity_indicator\",\n",
    "    ]\n",
    "\n",
    "    dtypes = OrderedDict(\n",
    "        [\n",
    "            (\"loan_id\", \"int64\"),\n",
    "            (\"monthly_reporting_period\", \"date\"),\n",
    "            (\"servicer\", \"str\"),\n",
    "            (\"interest_rate\", \"float64\"),\n",
    "            (\"current_actual_upb\", \"float64\"),\n",
    "            (\"loan_age\", \"float64\"),\n",
    "            (\"remaining_months_to_legal_maturity\", \"float64\"),\n",
    "            (\"adj_remaining_months_to_maturity\", \"float64\"),\n",
    "            (\"maturity_date\", \"date\"),\n",
    "            (\"msa\", \"float64\"),\n",
    "            (\"current_loan_delinquency_status\", \"int32\"),\n",
    "            (\"mod_flag\", \"str\"),\n",
    "            (\"zero_balance_code\", \"str\"),\n",
    "            (\"zero_balance_effective_date\", \"date\"),\n",
    "            (\"last_paid_installment_date\", \"date\"),\n",
    "            (\"foreclosed_after\", \"date\"),\n",
    "            (\"disposition_date\", \"date\"),\n",
    "            (\"foreclosure_costs\", \"float64\"),\n",
    "            (\"prop_preservation_and_repair_costs\", \"float64\"),\n",
    "            (\"asset_recovery_costs\", \"float64\"),\n",
    "            (\"misc_holding_expenses\", \"float64\"),\n",
    "            (\"holding_taxes\", \"float64\"),\n",
    "            (\"net_sale_proceeds\", \"float64\"),\n",
    "            (\"credit_enhancement_proceeds\", \"float64\"),\n",
    "            (\"repurchase_make_whole_proceeds\", \"float64\"),\n",
    "            (\"other_foreclosure_proceeds\", \"float64\"),\n",
    "            (\"non_interest_bearing_upb\", \"float64\"),\n",
    "            (\"principal_forgiveness_upb\", \"float64\"),\n",
    "            (\"repurchase_make_whole_proceeds_flag\", \"str\"),\n",
    "            (\"foreclosure_principal_write_off_amount\", \"float64\"),\n",
    "            (\"servicing_activity_indicator\", \"str\"),\n",
    "        ]\n",
    "    )\n",
    "\n",
    "    return cudf.read_csv(\n",
    "        performance_path,\n",
    "        names=cols,\n",
    "        delimiter=\"|\",\n",
    "        dtype=list(dtypes.values()),\n",
    "        skiprows=1,\n",
    "    )\n",
    "\n",
    "\n",
    "def gpu_load_acquisition_csv(acquisition_path, **kwargs):\n",
    "    \"\"\" \n",
    "    Loads acquisition data\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    GPU DataFrame\n",
    "    \"\"\"\n",
    "\n",
    "    cols = [\n",
    "        \"loan_id\",\n",
    "        \"orig_channel\",\n",
    "        \"seller_name\",\n",
    "        \"orig_interest_rate\",\n",
    "        \"orig_upb\",\n",
    "        \"orig_loan_term\",\n",
    "        \"orig_date\",\n",
    "        \"first_pay_date\",\n",
    "        \"orig_ltv\",\n",
    "        \"orig_cltv\",\n",
    "        \"num_borrowers\",\n",
    "        \"dti\",\n",
    "        \"borrower_credit_score\",\n",
    "        \"first_home_buyer\",\n",
    "        \"loan_purpose\",\n",
    "        \"property_type\",\n",
    "        \"num_units\",\n",
    "        \"occupancy_status\",\n",
    "        \"property_state\",\n",
    "        \"zip\",\n",
    "        \"mortgage_insurance_percent\",\n",
    "        \"product_type\",\n",
    "        \"coborrow_credit_score\",\n",
    "        \"mortgage_insurance_type\",\n",
    "        \"relocation_mortgage_indicator\",\n",
    "    ]\n",
    "\n",
    "    dtypes = OrderedDict(\n",
    "        [\n",
    "            (\"loan_id\", \"int64\"),\n",
    "            (\"orig_channel\", \"str\"),\n",
    "            (\"seller_name\", \"str\"),\n",
    "            (\"orig_interest_rate\", \"float64\"),\n",
    "            (\"orig_upb\", \"int64\"),\n",
    "            (\"orig_loan_term\", \"int64\"),\n",
    "            (\"orig_date\", \"date\"),\n",
    "            (\"first_pay_date\", \"date\"),\n",
    "            (\"orig_ltv\", \"float64\"),\n",
    "            (\"orig_cltv\", \"float64\"),\n",
    "            (\"num_borrowers\", \"float64\"),\n",
    "            (\"dti\", \"float64\"),\n",
    "            (\"borrower_credit_score\", \"float64\"),\n",
    "            (\"first_home_buyer\", \"str\"),\n",
    "            (\"loan_purpose\", \"str\"),\n",
    "            (\"property_type\", \"str\"),\n",
    "            (\"num_units\", \"int64\"),\n",
    "            (\"occupancy_status\", \"str\"),\n",
    "            (\"property_state\", \"str\"),\n",
    "            (\"zip\", \"int64\"),\n",
    "            (\"mortgage_insurance_percent\", \"float64\"),\n",
    "            (\"product_type\", \"str\"),\n",
    "            (\"coborrow_credit_score\", \"float64\"),\n",
    "            (\"mortgage_insurance_type\", \"float64\"),\n",
    "            (\"relocation_mortgage_indicator\", \"str\"),\n",
    "        ]\n",
    "    )\n",
    "\n",
    "    return cudf.read_csv(\n",
    "        acquisition_path,\n",
    "        names=cols,\n",
    "        delimiter=\"|\",\n",
    "        dtype=list(dtypes.values()),\n",
    "        skiprows=1,\n",
    "    )\n",
    "\n",
    "\n",
    "def gpu_load_names(col_names_path=\"\", **kwargs):\n",
    "    \"\"\" \n",
    "    Loads names used for renaming the banks\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    GPU DataFrame\n",
    "    \"\"\"\n",
    "\n",
    "    cols = [\"seller_name\", \"new\"]\n",
    "\n",
    "    dtypes = OrderedDict([(\"seller_name\", \"str\"), (\"new\", \"str\"),])\n",
    "\n",
    "    return cudf.read_csv(\n",
    "        col_names_path,\n",
    "        names=cols,\n",
    "        delimiter=\"|\",\n",
    "        dtype=list(dtypes.values()),\n",
    "        skiprows=1,\n",
    "    )\n",
    "\n",
    "\n",
    "def hash_df_string_columns(gdf):\n",
    "    \"\"\"\n",
    "    Hash all string columns in a cudf dataframe\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    Dataframe with all string columns replaced by hashed values for the strings\n",
    "    \"\"\"\n",
    "    for col in gdf.columns:\n",
    "        if cudf.utils.dtypes.is_string_dtype(gdf[col]):\n",
    "            gdf[col] = gdf[col].hash_values()\n",
    "    return gdf\n",
    "\n",
    "\n",
    "def create_ever_features(gdf, **kwargs):\n",
    "    \"\"\"\n",
    "    Creates features denoting whether a loan_id has ever been delinquent\n",
    "    for over 30, 90 and 180 days.\n",
    "    \"\"\"\n",
    "    everdf = gdf[[\"loan_id\", \"current_loan_delinquency_status\"]]\n",
    "    everdf = everdf.groupby(\"loan_id\", as_index=False).max()\n",
    "    del gdf\n",
    "    everdf[\"ever_30\"] = (everdf[\"current_loan_delinquency_status\"] >= 1).astype(\"int8\")\n",
    "    everdf[\"ever_90\"] = (everdf[\"current_loan_delinquency_status\"] >= 3).astype(\"int8\")\n",
    "    everdf[\"ever_180\"] = (everdf[\"current_loan_delinquency_status\"] >= 6).astype(\"int8\")\n",
    "    everdf.drop(columns=\"current_loan_delinquency_status\", inplace=True)\n",
    "    return everdf\n",
    "\n",
    "\n",
    "def create_delinq_features(gdf, **kwargs):\n",
    "    \"\"\"\n",
    "    Computes features denoting the earliest reported date when a loan_id\n",
    "    became delinquent for more than 30, 90 and 180 days.\n",
    "    \"\"\"\n",
    "    delinq_gdf = gdf[\n",
    "        [\"loan_id\", \"monthly_reporting_period\", \"current_loan_delinquency_status\",]\n",
    "    ]\n",
    "    del gdf\n",
    "    delinq_30 = (\n",
    "        delinq_gdf.query(\"current_loan_delinquency_status >= 1\")[\n",
    "            [\"loan_id\", \"monthly_reporting_period\"]\n",
    "        ]\n",
    "        .groupby(\"loan_id\", as_index=False)\n",
    "        .min()\n",
    "    )\n",
    "    delinq_30[\"delinquency_30\"] = delinq_30[\"monthly_reporting_period\"]\n",
    "    delinq_30.drop(columns=\"monthly_reporting_period\", inplace=True)\n",
    "    delinq_90 = (\n",
    "        delinq_gdf.query(\"current_loan_delinquency_status >= 3\")[\n",
    "            [\"loan_id\", \"monthly_reporting_period\"]\n",
    "        ]\n",
    "        .groupby(\"loan_id\", as_index=False)\n",
    "        .min()\n",
    "    )\n",
    "    delinq_90[\"delinquency_90\"] = delinq_90[\"monthly_reporting_period\"]\n",
    "    delinq_90.drop(columns=\"monthly_reporting_period\", inplace=True)\n",
    "    delinq_180 = (\n",
    "        delinq_gdf.query(\"current_loan_delinquency_status >= 6\")[\n",
    "            [\"loan_id\", \"monthly_reporting_period\"]\n",
    "        ]\n",
    "        .groupby(\"loan_id\", as_index=False)\n",
    "        .min()\n",
    "    )\n",
    "    delinq_180[\"delinquency_180\"] = delinq_180[\"monthly_reporting_period\"]\n",
    "    delinq_180.drop(columns=\"monthly_reporting_period\", inplace=True)\n",
    "    del delinq_gdf\n",
    "    delinq_merge = delinq_30.merge(delinq_90, how=\"left\", on=[\"loan_id\"])\n",
    "    delinq_merge = delinq_merge.merge(\n",
    "        delinq_180, how=\"left\", on=[\"loan_id\"]\n",
    "    )\n",
    "    del delinq_30\n",
    "    del delinq_90\n",
    "    del delinq_180\n",
    "    return delinq_merge\n",
    "\n",
    "\n",
    "def join_ever_delinq_features(everdf_tmp, delinq_merge, **kwargs):\n",
    "    \"\"\"\n",
    "    Merges the ever and delinq features table on loan_id\n",
    "    \"\"\"\n",
    "    everdf = everdf_tmp.merge(delinq_merge, on=[\"loan_id\"], how=\"left\")\n",
    "    del everdf_tmp\n",
    "    del delinq_merge\n",
    "    return everdf\n",
    "\n",
    "\n",
    "def create_joined_df(gdf, everdf, **kwargs):\n",
    "    \"\"\"\n",
    "    Join the performance table with the features table. (delinq and ever features)\n",
    "    \"\"\"\n",
    "    test = gdf[\n",
    "        [\n",
    "            \"loan_id\",\n",
    "            \"monthly_reporting_period\",\n",
    "            \"current_loan_delinquency_status\",\n",
    "            \"current_actual_upb\",\n",
    "        ]\n",
    "    ]\n",
    "    del gdf\n",
    "    test[\"timestamp\"] = test[\"monthly_reporting_period\"]\n",
    "    test.drop(columns=\"monthly_reporting_period\", inplace=True)\n",
    "    test[\"timestamp_month\"] = test[\"timestamp\"].dt.month\n",
    "    test[\"timestamp_year\"] = test[\"timestamp\"].dt.year\n",
    "    test[\"delinquency_12\"] = test[\"current_loan_delinquency_status\"]\n",
    "    test.drop(columns=\"current_loan_delinquency_status\", inplace=True)\n",
    "    test[\"upb_12\"] = test[\"current_actual_upb\"]\n",
    "    test.drop(columns=\"current_actual_upb\", inplace=True)\n",
    "\n",
    "    joined_df = test.merge(everdf, how=\"left\", on=[\"loan_id\"])\n",
    "    del everdf\n",
    "    del test\n",
    "\n",
    "    joined_df[\"timestamp_year\"] = joined_df[\"timestamp_year\"].astype(\"int32\")\n",
    "    joined_df[\"timestamp_month\"] = joined_df[\"timestamp_month\"].astype(\"int32\")\n",
    "\n",
    "    return joined_df\n",
    "\n",
    "\n",
    "def create_12_mon_features(joined_df, **kwargs):\n",
    "    \"\"\"\n",
    "    For every loan_id in a 12 month window compute a feature denoting\n",
    "    whether it has been delinquent for over 3 months or had an unpaid principal balance.\n",
    "    The 12 month window moves by a month to span across all months of the year.\n",
    "    \n",
    "    The computations windows for each loan_id follows the pattern below\n",
    "    Window 1: Jan 2000 - Jan 2001, Jan 2001 - Jan 2002\n",
    "    Window 2: Feb 2000- Feb 2001, Feb 2001 - Feb 2002\n",
    "    \"\"\"\n",
    "    testdfs = []\n",
    "    n_months = 12\n",
    "    for y in range(1, n_months + 1):\n",
    "        tmpdf = joined_df[\n",
    "            [\"loan_id\", \"timestamp_year\", \"timestamp_month\", \"delinquency_12\", \"upb_12\"]\n",
    "        ]\n",
    "        tmpdf[\"josh_months\"] = tmpdf[\"timestamp_year\"] * 12 + tmpdf[\"timestamp_month\"]\n",
    "        tmpdf[\"josh_mody_n\"] = (\n",
    "            (tmpdf[\"josh_months\"].astype(\"float64\") - 24000 - y) / 12\n",
    "        ).floor()\n",
    "        tmpdf = tmpdf.groupby(\n",
    "            [\"loan_id\", \"josh_mody_n\"], as_index=False\n",
    "        ).agg({\"delinquency_12\": \"max\", \"upb_12\": \"min\"})\n",
    "        tmpdf[\"delinquency_12\"] = (tmpdf[\"delinquency_12\"] > 3).astype(\"int32\")\n",
    "        tmpdf[\"delinquency_12\"] += (tmpdf[\"upb_12\"] == 0).astype(\"int32\")\n",
    "        tmpdf[\"timestamp_year\"] = (\n",
    "            (((tmpdf[\"josh_mody_n\"] * n_months) + 24000 + (y - 1)) / 12)\n",
    "            .floor()\n",
    "            .astype(\"int16\")\n",
    "        )\n",
    "        tmpdf[\"timestamp_month\"] = np.int8(y)\n",
    "        tmpdf.drop(columns=\"josh_mody_n\", inplace=True)\n",
    "        testdfs.append(tmpdf)\n",
    "        del tmpdf\n",
    "    del joined_df\n",
    "\n",
    "    return cudf.concat(testdfs)\n",
    "\n",
    "\n",
    "def combine_joined_12_mon(joined_df, testdf, **kwargs):\n",
    "    \"\"\"\n",
    "    Combines the 12_mon features table with the ever_delinq features tables\n",
    "    \"\"\"\n",
    "    joined_df.drop(columns=[\"delinquency_12\", \"upb_12\"], inplace=True)\n",
    "    joined_df[\"timestamp_year\"] = joined_df[\"timestamp_year\"].astype(\"int16\")\n",
    "    joined_df[\"timestamp_month\"] = joined_df[\"timestamp_month\"].astype(\"int8\")\n",
    "    return joined_df.merge(\n",
    "        testdf,\n",
    "        how=\"left\",\n",
    "        on=[\"loan_id\", \"timestamp_year\", \"timestamp_month\"]\n",
    "    )\n",
    "\n",
    "\n",
    "def final_performance_delinquency(gdf, joined_df, **kwargs):\n",
    "    \"\"\"\n",
    "    Combines the grouped table with all features with the original Performance table\n",
    "    \"\"\"\n",
    "    merged = gdf\n",
    "    joined_df[\"timestamp_month\"] = joined_df[\"timestamp_month\"].astype(\"int8\")\n",
    "    joined_df[\"timestamp_year\"] = joined_df[\"timestamp_year\"].astype(\"int16\")\n",
    "    merged[\"timestamp_month\"] = merged[\"monthly_reporting_period\"].dt.month\n",
    "    merged[\"timestamp_month\"] = merged[\"timestamp_month\"].astype(\"int8\")\n",
    "    merged[\"timestamp_year\"] = merged[\"monthly_reporting_period\"].dt.year\n",
    "    merged[\"timestamp_year\"] = merged[\"timestamp_year\"].astype(\"int16\")\n",
    "    merged = merged.merge(\n",
    "        joined_df,\n",
    "        how=\"left\",\n",
    "        on=[\"loan_id\", \"timestamp_year\", \"timestamp_month\"]\n",
    "    )\n",
    "    merged.drop(columns=[\"timestamp_year\",\"timestamp_month\"], inplace=True)\n",
    "    return merged\n",
    "\n",
    "\n",
    "def join_perf_acq_gdfs(perf, acq, **kwargs):\n",
    "    \"\"\"\n",
    "    Combines the Acquisition and Performance tables on loan_id\n",
    "    \"\"\"\n",
    "    return perf.merge(acq, how=\"left\", on=[\"loan_id\"])\n",
    "\n",
    "\n",
    "def last_mile_cleaning(df, **kwargs):\n",
    "    \"\"\"\n",
    "    Final cleanup to drop columns not passed to the XGBoost model for training.\n",
    "    Convert all string/categorical features to numeric features.\n",
    "\n",
    "    Returns\n",
    "    ------\n",
    "    Arrow Table (Host memory)\n",
    "    \"\"\"\n",
    "    drop_list = [\n",
    "        \"loan_id\",\n",
    "        \"orig_date\",\n",
    "        \"first_pay_date\",\n",
    "        \"seller_name\",\n",
    "        \"monthly_reporting_period\",\n",
    "        \"last_paid_installment_date\",\n",
    "        \"maturity_date\",\n",
    "        \"ever_30\",\n",
    "        \"ever_90\",\n",
    "        \"ever_180\",\n",
    "        \"delinquency_30\",\n",
    "        \"delinquency_90\",\n",
    "        \"delinquency_180\",\n",
    "        \"upb_12\",\n",
    "        \"zero_balance_effective_date\",\n",
    "        \"foreclosed_after\",\n",
    "        \"disposition_date\",\n",
    "        \"timestamp\",\n",
    "    ]\n",
    "    df.drop(columns=drop_list, inplace=True)\n",
    "    for col, dtype in df.dtypes.iteritems():\n",
    "        if str(dtype) == \"category\":\n",
    "            df[col] = df[col].cat.codes\n",
    "        df[col] = df[col].astype(\"float32\")\n",
    "    df[\"delinquency_12\"] = df[\"delinquency_12\"] > 0\n",
    "    df[\"delinquency_12\"] = df[\"delinquency_12\"].fillna(False).astype(\"int32\")\n",
    "    for column in df.columns:\n",
    "        df[column] = df[column].fillna(np.dtype(str(df[column].dtype)).type(-1))\n",
    "    return df.to_arrow(preserve_index=False)\n",
    "\n",
    "\n",
    "def prepare_data(arrow_input):\n",
    "    \"\"\"\n",
    "    Convert a list of arrow tables to a single GPU dataframe\n",
    "    \n",
    "    Returns\n",
    "    -------\n",
    "    GPU Dataframe\n",
    "    \"\"\"\n",
    "    gpu_dataframes = []\n",
    "    for arrow_df in arrow_input:\n",
    "        gpu_dataframes.append(cudf.DataFrame.from_arrow(arrow_df))\n",
    "\n",
    "    concat_df = cudf.concat(gpu_dataframes)\n",
    "    del gpu_dataframes\n",
    "    return concat_df\n",
    "\n",
    "\n",
    "def xgb_training(arrow_dfs, client=None):\n",
    "    \"\"\"\n",
    "    Convert the post ETL data to Dmatrix format for XGBoost training input.\n",
    "    Train the XGBoost model.\n",
    "    \n",
    "    Returns\n",
    "    -------\n",
    "    The trained model and time taken for preparing, training data.\n",
    "    \"\"\"\n",
    "    dxgb_gpu_params = {\n",
    "        \"max_depth\": 8,\n",
    "        \"max_leaves\": 2 ** 8,\n",
    "        \"alpha\": 0.9,\n",
    "        \"eta\": 0.1,\n",
    "        \"gamma\": 0.1,\n",
    "        \"learning_rate\": 0.1,\n",
    "        \"subsample\": 1,\n",
    "        \"reg_lambda\": 1,\n",
    "        \"scale_pos_weight\": 2,\n",
    "        \"min_child_weight\": 30,\n",
    "        \"tree_method\": \"gpu_hist\",\n",
    "        \"objective\": \"binary:logistic\",\n",
    "        \"grow_policy\": \"lossguide\",\n",
    "    }\n",
    "    NUM_BOOST_ROUND = 100\n",
    "\n",
    "    part_count = len(arrow_dfs)\n",
    "    print(f\"Preparing data for training with part count: {part_count}\")\n",
    "    t1 = time.time()\n",
    "    tmp_map = [\n",
    "        (arrow_df, list(client.who_has(arrow_df).values())[0][0])\n",
    "        for arrow_df in arrow_dfs\n",
    "    ]\n",
    "    new_map = OrderedDict()\n",
    "    for key, value in tmp_map:\n",
    "        if value not in new_map:\n",
    "            new_map[value] = [key]\n",
    "        else:\n",
    "            new_map[value].append(key)\n",
    "\n",
    "    del (tmp_map, key, value)\n",
    "\n",
    "    train_x_y = []\n",
    "    for list_delayed in new_map.values():\n",
    "        train_x_y.append(delayed(prepare_data)(list_delayed))\n",
    "\n",
    "    del (new_map, list_delayed)\n",
    "\n",
    "    worker_list = OrderedDict()\n",
    "    for task in train_x_y:\n",
    "        worker_list[task] = list(client.who_has(task).values())[0][0]\n",
    "\n",
    "    del task\n",
    "\n",
    "    persisted_train_x_y = []\n",
    "    for task in train_x_y:\n",
    "        persisted_train_x_y.append(\n",
    "            client.persist(\n",
    "                collections=task,\n",
    "                workers=worker_list[task],\n",
    "                optimize_graph=False,\n",
    "                fifo_timeout=\"0ms\",\n",
    "            )\n",
    "        )\n",
    "\n",
    "    del (arrow_dfs, train_x_y, worker_list, task)\n",
    "\n",
    "    wait(persisted_train_x_y)\n",
    "    persisted_train_x_y = dask_cudf.from_delayed(persisted_train_x_y)\n",
    "\n",
    "    dmat = xgb.dask.DaskDeviceQuantileDMatrix(\n",
    "        client=client,\n",
    "        data=persisted_train_x_y[\n",
    "            persisted_train_x_y.columns.difference([\"delinquency_12\"])\n",
    "        ],\n",
    "        label=persisted_train_x_y[[\"delinquency_12\"]],\n",
    "        missing=-1,\n",
    "    )\n",
    "\n",
    "    del persisted_train_x_y\n",
    "    gc.collect()\n",
    "\n",
    "    dmat_time = time.time() - t1\n",
    "    print(\"Prepared data for XGB training\")\n",
    "\n",
    "    print(\"Training model\")\n",
    "    t1 = time.time()\n",
    "\n",
    "    print(\"XGB training for part_count:{}\".format(part_count))\n",
    "    bst = xgb.dask.train(\n",
    "        client, dxgb_gpu_params, dmat, num_boost_round=NUM_BOOST_ROUND,\n",
    "    )\n",
    "\n",
    "    train_time = time.time() - t1\n",
    "    print(\"Training complete\")\n",
    "    return (bst, dmat_time, train_time)\n",
    "\n",
    "\n",
    "def run_etl(start_year, end_year, data_dir, client):\n",
    "    \"\"\"\n",
    "    Driver function for the ETL step\n",
    "    \n",
    "    Iterates through all files in `data_dir` between `start_year` \n",
    "    and `end_year` and calls the ETL function for each file.\n",
    "    \n",
    "    Returns\n",
    "    -------\n",
    "    Dask futures to arrow tables containing post ETL data for all processed files.\n",
    "    \"\"\"\n",
    "    print(\"Starting ETL\")\n",
    "    t1 = time.time()\n",
    "\n",
    "    perf_data_path = data_dir + \"perf/\"\n",
    "\n",
    "    gpu_dfs = []\n",
    "    quarter = 1\n",
    "    year = start_year\n",
    "    count = 0\n",
    "    while year <= end_year:\n",
    "        for file in glob(\n",
    "            os.path.join(\n",
    "                perf_data_path + \"/Performance_\" + str(year) + \"Q\" + str(quarter) + \"*\"\n",
    "            )\n",
    "        ):\n",
    "            gpu_dfs.append(\n",
    "                process_quarter_gpu(\n",
    "                    year=year, quarter=quarter, perf_file=file, client=client\n",
    "                )\n",
    "            )\n",
    "            count += 1\n",
    "        quarter += 1\n",
    "        if quarter == 5:\n",
    "            year += 1\n",
    "            quarter = 1\n",
    "    print(\"ETL for start_year:{} and end_year:{}\".format(start_year, end_year))\n",
    "    wait(gpu_dfs)\n",
    "\n",
    "    etl_time = time.time() - t1\n",
    "\n",
    "    print(\"ETL done!\")\n",
    "    return (gpu_dfs, etl_time)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### The cell below runs the workflow end to end including the ETL and XGBoost model training step\n",
    "\n",
    "**Notes** \n",
    "\n",
    "The mortgage dataset for years 2000-2016 is about 200GB of data. There are two key factors that determine the `start_year`, `end_year`, `part_count` and `use_1GB_splits` params used in the notebook for processing this data. \n",
    "\n",
    "_Total GPU memory_: Determines the amount of data that can be trained using XGBoost (`part_count`). The ETL is performed on one part file at a time (per GPU) whereas XGBoost training requires all the training data to be loaded in GPU memory.\n",
    "\n",
    "_Memory per GPU_: Determines the variation of the dataset to use (1GB vs 2GB splits). The 2GB splits version of the data results in larger partitions being processed per task resulting in better utilization of the GPU, with the tradeoff of increased memory usage that can be handled by GPUs cards with greater than `32GB` of memory.\n",
    "\n",
    "The `determine_dataset` utility used below automatically queries these two parameters based on the machine and decides suitable values for `part_count` and consequently `start_year`, `end_year`(to ensure ETL is performed on enough parts for training), as well as the variation of the dataset (1GB split part files vs 2GB split part files) that should work on such systems.\n",
    "\n",
    "If you'd like to use existing data that has already been downloaded to your own location, or manually adjust these parameters based on the amount of data needed for processing, you can change these parameters provided in the notebook, by assigning new values to the variables or setting enivronment variables for `MORTGAGE_DATA_DIR` and `part_count`. You can visit the [RAPIDS Datasets Homepage](https://docs.rapids.ai/datasets/mortgage-data) for more information on downloading the data manually."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Downloading data for year 2000\n",
      "Download complete\n",
      "Decompressing and extracting data\n",
      "Done extracting year 2000\n",
      "Downloading data for year 2001\n",
      "Download complete\n",
      "Decompressing and extracting data\n",
      "Done extracting year 2001\n",
      "Downloading data for year 2002\n",
      "Download complete\n",
      "Decompressing and extracting data\n",
      "Done extracting year 2002\n",
      "Downloading data for year 2003\n",
      "Download complete\n",
      "Decompressing and extracting data\n",
      "Done extracting year 2003\n",
      "Downloading data for year 2004\n",
      "Download complete\n",
      "Decompressing and extracting data\n",
      "Done extracting year 2004\n",
      "Starting ETL\n",
      "ETL for start_year:2000 and end_year:2004\n",
      "ETL done!\n",
      "Preparing data for training with part count: 12\n",
      "Prepared data for XGB training\n",
      "Training model\n",
      "XGB training for part_count:12\n",
      "Training complete\n",
      "\n",
      "Time taken to run ETL from 2000 to 2004 (108 parts) was 68.7227 s\n",
      "Time taken to prepare 12 parts for XGB training 3.3915 s\n",
      "Time taken to train XGB model 87.521 s\n",
      "Total E2E time: 159.6352 s\n"
     ]
    }
   ],
   "source": [
    "if __name__ == \"__main__\":\n",
    "\n",
    "    import cudf\n",
    "    import xgboost as xgb\n",
    "    import dask_cudf\n",
    "\n",
    "    cmd = \"hostname --all-ip-addresses\"\n",
    "    process = subprocess.Popen(cmd.split(), stdout=subprocess.PIPE)\n",
    "    output, error = process.communicate()\n",
    "    IPADDR = str(output.decode()).split()[0]\n",
    "\n",
    "    cluster = LocalCUDACluster(ip=IPADDR)\n",
    "    client = Client(cluster)\n",
    "\n",
    "    data_dir = os.environ.get(\"MORTGAGE_DATA_DIR\", \"\") # Default to current working directory\n",
    "    res = client.run(memory_info)\n",
    "    # Total GPU memory on the system\n",
    "    total_mem = sum(res.values()) \n",
    "    # Memory of a single GPU on the machine\n",
    "    # If the machine has multiple GPUs of different sizes, this is the size of the smallest GPU\n",
    "    min_mem = min(res.values()) \n",
    " \n",
    "    # Start year for processing mortgage data\n",
    "    start_year = None\n",
    "    # End year for processing mortgage data\n",
    "    end_year = None\n",
    "    # The number of part files to train against. \n",
    "    # If not provided, default to auto selection based on GPU memory available on the system\n",
    "    part_count = os.environ.get(\"part_count\")\n",
    "\n",
    "    start_year, end_year, part_count, use_1GB_splits = determine_dataset(\n",
    "        total_mem=total_mem, min_mem=min_mem, part_count=part_count\n",
    "    )\n",
    "\n",
    "    # Download data based on these parameters\n",
    "    # The 2GB split mortgage performance files are used if the system has 32GB GPUs.\n",
    "    # On machines with GPUs less than 32GB we use the 1GB split files (to help reduce memory load)\n",
    "    get_data(data_dir, start_year, end_year, use_1GB_splits)\n",
    "\n",
    "    # Initialize a GPU pool allocating 95% of GPU memory for each worker\n",
    "    mem_fraction = int(0.95 * min_mem)\n",
    "    initial_pool_size = mem_fraction - (mem_fraction % 256) #Rounding to the nearest value divisible by 256\n",
    "    client.run(rmm.reinitialize, pool_allocator=True, initial_pool_size=initial_pool_size)\n",
    "    etl_result, etl_time = run_etl(start_year, end_year, data_dir, client)\n",
    "\n",
    "    # Clear the existing RMM pool post-ETL to make space for GPU accelerated XGBoost\n",
    "    # This makes space for XGBoost to operate since it doesn't have visibility into the cuDF memory pool\n",
    "    client.run(rmm.reinitialize, pool_allocator=False)\n",
    "\n",
    "    total_file_count = len(etl_result)\n",
    "    etl_result = etl_result[:part_count]  # Select subset for training\n",
    "    model, dmat_time, train_time = xgb_training(etl_result, client)\n",
    "\n",
    "    print(\n",
    "        f\"\\nTime taken to run ETL from {start_year} to {end_year}\"\n",
    "        f\" ({total_file_count} parts) was {round(etl_time,4)} s\"\n",
    "    )\n",
    "    print(\n",
    "        f\"Time taken to prepare {len(etl_result)} parts\"\n",
    "        f\" for XGB training {round(dmat_time,4)} s\"\n",
    "    )\n",
    "    print(f\"Time taken to train XGB model {round(train_time, 4)} s\")\n",
    "    print(f\"Total E2E time: {round(etl_time+dmat_time+train_time, 4)} s\")\n",
    "    client.close()\n",
    "    cluster.close()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.8.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}