dask/kubernetes/Dask_cuML_Exploration.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "gothic-metro",
   "metadata": {},
   "outputs": [],
   "source": [
    "import certifi\n",
    "import cudf\n",
    "import cuml\n",
    "import cupy as cp\n",
    "import gcsfs\n",
    "import json\n",
    "import numpy as np\n",
    "import os\n",
    "import pandas as pd\n",
    "import random\n",
    "import seaborn as sns\n",
    "import time\n",
    "import uuid\n",
    "import yaml\n",
    "\n",
    "from collections import OrderedDict\n",
    "from functools import partial\n",
    "from math import cos, sin, asin, sqrt, pi\n",
    "from tqdm import tqdm\n",
    "from typing import Optional\n",
    "\n",
    "import dask\n",
    "import dask.array as da\n",
    "import dask.dataframe as dd\n",
    "import dask_cudf\n",
    "\n",
    "from dask_kubernetes import KubeCluster, make_pod_from_dict\n",
    "from dask.distributed import Client, WorkerPlugin, wait, progress, get_worker\n",
    "\n",
    "from cuml import ForestInference\n",
    "\n",
    " \n",
    "def create_pod_from_yaml(yaml_file):\n",
    "    with open(yaml_file, 'r') as reader:\n",
    "        d = yaml.safe_load(reader)\n",
    "        d = dask.config.expand_environment_variables(d)\n",
    "    return make_pod_from_dict(d)\n",
    "\n",
    "\n",
    "def build_worker_and_scheduler_pods(sched_spec, worker_spec):\n",
    "    assert os.path.isfile(sched_spec)\n",
    "    assert os.path.isfile(worker_spec)\n",
    "\n",
    "    sched_pod = create_pod_from_yaml(sched_spec)\n",
    "    worker_pod = create_pod_from_yaml(worker_spec)\n",
    "\n",
    "    return sched_pod, worker_pod\n",
    "\n",
    "\n",
    "dask.config.set({\"logging.kubernetes\": \"info\",\n",
    "                 \"logging.distributed\": \"info\",\n",
    "                 \"kubernetes.scheduler-service-type\": \"LoadBalancer\",\n",
    "                 \"kubernetes.idle-timeout\": None,\n",
    "                 \"kubernetes.scheduler-service-wait-timeout\": 3600,\n",
    "                 \"kubernetes.deploy-mode\": \"remote\",\n",
    "                 \"kubernetes.logging\": \"info\",\n",
    "                 \"distributed.logging\": \"info\",\n",
    "                 \"distributed.scheduler.idle-timeout\": None,\n",
    "                 \"distributed.scheduler.locks.lease-timeout\": None,\n",
    "                 \"distributed.comm.timeouts.connect\": 3600,\n",
    "                 \"distributed.comm.tls.ca-file\": certifi.where()})\n",
    "\n",
    "sched_spec_path = \"./specs/sched-spec.yaml\"\n",
    "worker_spec_path = \"./specs/worker-spec.yaml\"\n",
    "\n",
    "sched_pod, worker_pod = build_worker_and_scheduler_pods(sched_spec=sched_spec_path,\n",
    "                                                        worker_spec=worker_spec_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "440f3cb8-9f83-4f1e-b9f9-26ee57d52dd9",
   "metadata": {},
   "outputs": [],
   "source": [
    "cluster = KubeCluster(pod_template=worker_pod,\n",
    "                      scheduler_pod_template=sched_pod)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "34d175df-617f-4150-90ea-24adbb01d9fa",
   "metadata": {},
   "outputs": [],
   "source": [
    "client = Client(cluster)\n",
    "scheduler_address = cluster.scheduler_address"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "awful-fourth",
   "metadata": {},
   "outputs": [],
   "source": [
    "def scale_workers(client, n_workers, timeout=300):\n",
    "    client.cluster.scale(n_workers)\n",
    "    \n",
    "    m = len(client.has_what().keys())    \n",
    "    start = end = time.perf_counter_ns()\n",
    "    while ((m != n_workers) and (((end - start) / 1e9) < timeout) ):\n",
    "        time.sleep(5)\n",
    "        m = len(client.has_what().keys())\n",
    "        \n",
    "        end = time.perf_counter_ns()\n",
    "        \n",
    "    if (((end - start) / 1e9) >= timeout):\n",
    "        raise RuntimeError(f\"Failed to rescale cluster in {timeout} sec.\"\n",
    "              \"Try increasing timeout for very large containers, and verify available compute resources.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "parental-nickel",
   "metadata": {},
   "outputs": [],
   "source": [
    "scale_workers(client, n_workers=8 timeout=800)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "mobile-european",
   "metadata": {},
   "source": [
    "### Helper Tools"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "synthetic-yesterday",
   "metadata": {},
   "outputs": [],
   "source": [
    "class SimpleTimer:\n",
    "    def __init__(self):\n",
    "        self.start = None\n",
    "        self.end = None\n",
    "        self.elapsed = None\n",
    "\n",
    "    def __enter__(self):\n",
    "        self.start = time.perf_counter_ns()\n",
    "        return self\n",
    "\n",
    "    def __exit__(self, exc_type, exc_val, exc_tb):\n",
    "        self.end = time.perf_counter_ns()\n",
    "        self.elapsed = self.end - self.start\n",
    "\n",
    "def construct_worker_pool(client, n_workers, auto_scale=False, timeout=300):\n",
    "    workers = [w for w in client.has_what().keys()]\n",
    "    if (len(workers) < n_workers):\n",
    "        if (auto_scale):\n",
    "            scale_workers(client=client, n_workers=n_workers, timeout=timeout)\n",
    "            workers = [w for w in client.has_what().keys()]\n",
    "        else:\n",
    "            print(\"Attempt to construct worker pool larger than available worker set, and auto_scale is False.\"\n",
    "                  \" Returning entire pool.\")\n",
    "    else:\n",
    "        workers = random.sample(population=workers, k=n_workers)\n",
    "        \n",
    "    return workers\n",
    "\n",
    "\n",
    "def estimate_df_rows(client, files, storage_opts={}, testpct=0.01):\n",
    "    workers = client.has_what().keys()\n",
    "    \n",
    "    est_size = 0\n",
    "    for file in files:\n",
    "        if (file.endswith('.csv')):\n",
    "            df = dask_cudf.read_csv(file, npartitions=len(workers), storage_options=storage_opts)\n",
    "        elif (file.endswith('.parquet')):\n",
    "            df = dask_cudf.read_parquet(file, npartitions=len(workers))           \n",
    "        \n",
    "        # Select only the index column from our subsample\n",
    "        est_size += (df.sample(frac=testpct).iloc[:,0].shape[0] / testpct).compute()\n",
    "        del df\n",
    "    \n",
    "    return est_size"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "photographic-voice",
   "metadata": {},
   "source": [
    "### Taxi Data Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "exotic-knitting",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def clean(df_part, remap, must_haves):\n",
    "    \"\"\"\n",
    "    This function performs the various clean up tasks for the data\n",
    "    and returns the cleaned dataframe.\n",
    "    \"\"\"\n",
    "    tmp = {col:col.strip().lower() for col in list(df_part.columns)}\n",
    "    df_part = df_part.rename(columns=tmp)\n",
    "    \n",
    "    # rename using the supplied mapping\n",
    "    df_part = df_part.rename(columns=remap)\n",
    "    \n",
    "    # iterate through columns in this df partition\n",
    "    for col in df_part.columns:\n",
    "        # drop anything not in our expected list\n",
    "        if col not in must_haves:\n",
    "            df_part = df_part.drop(col, axis=1)\n",
    "            continue\n",
    "\n",
    "        # fixes datetime error found by Ty Mckercher and fixed by Paul Mahler\n",
    "        if df_part[col].dtype == 'object' and col in ['pickup_datetime', 'dropoff_datetime']:\n",
    "            df_part[col] = df_part[col].astype('datetime64[ms]')\n",
    "            continue\n",
    "\n",
    "        # if column was read as a string, recast as float\n",
    "        if df_part[col].dtype == 'object':\n",
    "            df_part[col] = df_part[col].astype('float32')\n",
    "        else:\n",
    "            # downcast from 64bit to 32bit types\n",
    "            # Tesla T4 are faster on 32bit ops\n",
    "            if 'int' in str(df_part[col].dtype):\n",
    "                df_part[col] = df_part[col].astype('int32')\n",
    "            if 'float' in str(df_part[col].dtype):\n",
    "                df_part[col] = df_part[col].astype('float32')\n",
    "            df_part[col] = df_part[col].fillna(-1)\n",
    "            \n",
    "    return df_part\n",
    "\n",
    "\n",
    "def coalesce_taxi_data(fraction, random_state):\n",
    "    base_path = 'gcs://anaconda-public-data/nyc-taxi/csv'\n",
    "\n",
    "    # list of column names that need to be re-mapped\n",
    "    remap = {}\n",
    "    remap['tpep_pickup_datetime'] = 'pickup_datetime'\n",
    "    remap['tpep_dropoff_datetime'] = 'dropoff_datetime'\n",
    "    remap['ratecodeid'] = 'rate_code'\n",
    "\n",
    "    #create a list of columns & dtypes the df must have\n",
    "    must_haves = {\n",
    "     'pickup_datetime': 'datetime64[ms]',\n",
    "     'dropoff_datetime': 'datetime64[ms]',\n",
    "     'passenger_count': 'int32',\n",
    "     'trip_distance': 'float32',\n",
    "     'pickup_longitude': 'float32',\n",
    "     'pickup_latitude': 'float32',\n",
    "     'rate_code': 'int32',\n",
    "     'dropoff_longitude': 'float32',\n",
    "     'dropoff_latitude': 'float32',\n",
    "     'fare_amount': 'float32'\n",
    "    }\n",
    "    \n",
    "    # apply a list of filter conditions to throw out records with missing or outlier values\n",
    "    query_fragments = [\n",
    "        'fare_amount > 0 and fare_amount < 500',\n",
    "        'passenger_count > 0 and passenger_count < 6',\n",
    "        'pickup_longitude > -75 and pickup_longitude < -73',\n",
    "        'dropoff_longitude > -75 and dropoff_longitude < -73',\n",
    "        'pickup_latitude > 40 and pickup_latitude < 42',\n",
    "        'dropoff_latitude > 40 and dropoff_latitude < 42'\n",
    "    ]\n",
    "    \n",
    "    \n",
    "    valid_months_2016 = [str(x).rjust(2, '0') for x in range(1, 7)]\n",
    "    valid_files_2016 = [f'{base_path}/2016/yellow_tripdata_2016-{month}.csv' for month in valid_months_2016]\n",
    "    \n",
    "    df_2014_fractional = dask_cudf.read_csv(f'{base_path}/2014/yellow_*.csv', chunksize=25e6).sample(\n",
    "        frac=fraction, random_state=random_state)\n",
    "    df_2014_fractional = clean(df_2014_fractional, remap, must_haves)\n",
    "    \n",
    "    df_2015_fractional = dask_cudf.read_csv(f'{base_path}/2015/yellow_*.csv', chunksize=25e6).sample(\n",
    "        frac=fraction, random_state=random_state)\n",
    "    df_2015_fractional = clean(df_2015_fractional, remap, must_haves)\n",
    "    \n",
    "    df_2016_fractional = dask_cudf.read_csv(valid_files_2016, chunksize=25e6).sample(\n",
    "        frac=fraction, random_state=random_state)\n",
    "    df_2016_fractional = clean(df_2016_fractional, remap, must_haves)\n",
    "    \n",
    "    df_taxi = dask.dataframe.multi.concat([df_2014_fractional, df_2015_fractional, df_2016_fractional])\n",
    "    df_taxi = df_taxi.query(' and '.join(query_fragments))\n",
    "    \n",
    "    return df_taxi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "changing-traffic",
   "metadata": {},
   "outputs": [],
   "source": [
    "def persist_train_infer_split(client, df, response_dtype, response_id, infer_frac=1.0):\n",
    "    workers = client.has_what().keys()\n",
    "    infer_frac = max(0, min(infer_frac, 1.0))\n",
    "    \n",
    "    df_train = df\n",
    "    \n",
    "    X_train = df_train[df.columns.difference([response_id])].astype(np.float32)\n",
    "    y_train = df_train[response_id].astype(response_dtype)\n",
    "    \n",
    "    with dask.annotate(workers=set(workers)):\n",
    "        X_train, y_train = client.persist(\n",
    "            collections=[X_train, y_train]) \n",
    "    \n",
    "    if (infer_frac != 1.0):\n",
    "        df_split = df.sample(frac=infer_frac)\n",
    "        X_infer = df_split[df.columns.difference([response_id])].astype(np.float32)\n",
    "        y_infer = df_split[response_id].astype(response_dtype)\n",
    "        \n",
    "        with dask.annotate(workers=set(workers)):\n",
    "            X_infer, y_infer = client.persist(\n",
    "                collections=[X_infer, y_infer])\n",
    "\n",
    "        wait([X_train, y_train, X_infer, y_infer])\n",
    "    else:\n",
    "        X_infer = X_train\n",
    "        y_infer = y_train\n",
    "\n",
    "        wait([X_train, y_train])\n",
    "    \n",
    "    return X_train, y_train, X_train, y_train\n",
    "\n",
    "def mortgage_parquet_loader(client, response_dtype=np.float32, fraction=1.0, infer_frac=1.0, random_state=0):\n",
    "    response_id = 'foreclosure_costs'\n",
    "    km_fields = [\n",
    "        'loan_id', 'interest_rate', 'current_actual_upb', 'loan_age', 'remaining_months_to_legal_maturity',\n",
    "        'msa', 'current_loan_delinquency_status',\n",
    "        'prop_preservation_and_repair_costs', 'asset_recovery_costs', 'misc_holding_expenses', 'holding_taxes',\n",
    "        'net_sale_proceeds', 'credit_enhancement_proceeds', 'repurchase_make_whole_proceeds',\n",
    "        'other_foreclosure_proceeds', 'non_interest_bearing_upb', 'principal_forgiveness_upb',\n",
    "        'foreclosure_principal_write_off_amount', 'foreclosure_costs'\n",
    "    ]\n",
    "\n",
    "    mort_df = dask_cudf.read_parquet(YOUR_MORTGAGE_DATA_PATH, storage_options=YOUR_STORAGE_OPTS).sample(\n",
    "        frac=fraction, random_state=random_state).fillna(value=0.0)\n",
    "    \n",
    "    mort_df = mort_df[km_fields]\n",
    "\n",
    "    \n",
    "    return persist_train_infer_split(client, mort_df, response_dtype, response_id, infer_frac)\n",
    "\n",
    "\n",
    "def taxi_csv_data_loader(client, response_dtype=np.float32, fraction=1.0, random_state=0):\n",
    "    response_id = 'fare_amount'\n",
    "    workers = client.has_what().keys()\n",
    "    km_fields = ['passenger_count', 'trip_distance', 'pickup_longitude',\n",
    "                 'pickup_latitude', 'rate_code', 'dropoff_longitude',\n",
    "                 'dropoff_latitude', 'fare_amount']\n",
    "    \n",
    "    taxi_df = coalesce_taxi_data(fraction=fraction, random_state=random_state)\n",
    "    \n",
    "    taxi_df = taxi_df[km_fields]\n",
    "    X = taxi_df[taxi_df.columns.difference([response_id])].astype(np.float32)\n",
    "    y = taxi_df[response_id].astype(response_dtype)\n",
    "    \n",
    "    with dask.annotate(workers=set(workers)):\n",
    "        taxi_df = client.persist(collections=taxi_df)\n",
    "        \n",
    "    with dask.annotate(workers=set(workers)):\n",
    "        X = client.persist(collections=X)\n",
    "    \n",
    "    with dask.annotate(workers=set(workers)):\n",
    "        y = client.persist(collections=y)\n",
    "    \n",
    "    wait([taxi_df, X, y])\n",
    "    \n",
    "    return taxi_df, X, y\n",
    "\n",
    "\n",
    "def taxi_parquet_data_loader(client, response_dtype=np.float32, fraction=1.0, infer_frac=1.0, random_state=0):\n",
    "    # list of column names that need to be re-mapped\n",
    "    remap = {}\n",
    "    remap['tpep_pickup_datetime'] = 'pickup_datetime'\n",
    "    remap['tpep_dropoff_datetime'] = 'dropoff_datetime'\n",
    "    remap['ratecodeid'] = 'rate_code'\n",
    "\n",
    "    #create a list of columns & dtypes the df must have\n",
    "    must_haves = {\n",
    "     'pickup_datetime': 'datetime64[ms]',\n",
    "     'dropoff_datetime': 'datetime64[ms]',\n",
    "     'passenger_count': 'int32',\n",
    "     'trip_distance': 'float32',\n",
    "     'pickup_longitude': 'float32',\n",
    "     'pickup_latitude': 'float32',\n",
    "     'rate_code': 'int32',\n",
    "     'dropoff_longitude': 'float32',\n",
    "     'dropoff_latitude': 'float32',\n",
    "     'fare_amount': 'float32'\n",
    "    }\n",
    "\n",
    "    # apply a list of filter conditions to throw out records with missing or outlier values\n",
    "    query_fragments = [\n",
    "        'fare_amount > 0 and fare_amount < 500',\n",
    "        'passenger_count > 0 and passenger_count < 6',\n",
    "        'pickup_longitude > -75 and pickup_longitude < -73',\n",
    "        'dropoff_longitude > -75 and dropoff_longitude < -73',\n",
    "        'pickup_latitude > 40 and pickup_latitude < 42',\n",
    "        'dropoff_latitude > 40 and dropoff_latitude < 42'\n",
    "    ]\n",
    "\n",
    "    workers = client.has_what().keys()\n",
    "    taxi_parquet_path = \"gs://anaconda-public-data/nyc-taxi/nyc.parquet\"\n",
    "    response_id = 'fare_amount'\n",
    "    fields = ['passenger_count', 'trip_distance', 'pickup_longitude', 'pickup_latitude', 'rate_code',\n",
    "                 'dropoff_longitude', 'dropoff_latitude', 'fare_amount']\n",
    "    \n",
    "    taxi_df = dask_cudf.read_parquet(taxi_parquet_path, npartitions=len(workers))\n",
    "    taxi_df = clean(taxi_df, remap, must_haves)\n",
    "    taxi_df = taxi_df.query(' and '.join(query_fragments))\n",
    "    taxi_df = taxi_df[fields]\n",
    "    \n",
    "    return persist_train_infer_split(client, taxi_df, response_dtype, response_id, infer_frac)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "searching-christmas",
   "metadata": {},
   "source": [
    "## Performance Validation Code"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "pacific-telephone",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "def record_elapsed_timings_to_df(df, timings, record_template, type, columns, write_to=None):\n",
    "    records = [dict(record_template, **{\"sample_index\": i,\n",
    "                                        \"elapsed\": elapsed,\n",
    "                                        \"type\": type})\n",
    "                  for i, elapsed in enumerate(timings)]\n",
    "\n",
    "    df = df.append(other=records, ignore_index=True)\n",
    "    \n",
    "    if (write_to):\n",
    "        df.to_csv(write_to, columns=columns) \n",
    "            \n",
    "    return df\n",
    "\n",
    "\n",
    "def collect_load_time_samples(load_func, count, return_final_sample=True, verbose=False):\n",
    "    timings = []\n",
    "    for m in tqdm(range(count)):\n",
    "        with SimpleTimer() as timer:\n",
    "            X_train, y_train, X_infer, y_infer = load_func()\n",
    "        timings.append(timer.elapsed)\n",
    "    \n",
    "    if (return_final_sample):\n",
    "        return X_train, y_train, X_infer, y_infer, timings\n",
    "    \n",
    "    return None, None, None, timings\n",
    "\n",
    "\n",
    "def collect_func_time_samples(func, count, verbose=False):\n",
    "    timings = []\n",
    "    for k in tqdm(range(count)):\n",
    "        with SimpleTimer() as timer:\n",
    "            func()\n",
    "        timings.append(timer.elapsed)\n",
    "        \n",
    "    return timings\n",
    "\n",
    "\n",
    "def sweep_fit_func(model, func_id, require_compute, X, y, xy_fit, count):\n",
    "    _fit_func_attr = getattr(model, func_id)\n",
    "    if (require_compute):\n",
    "        if (xy_fit):\n",
    "            fit_func = partial(lambda X, y: _fit_func_attr(X, y).compute(), X, y)\n",
    "        else:\n",
    "            fit_func = partial(lambda X: _fit_func_attr(X).compute(), X)\n",
    "    else:\n",
    "        if (xy_fit):\n",
    "            fit_func = partial(_fit_func_attr, X, y)\n",
    "        else:\n",
    "            fit_func = partial(_fit_func_attr, X)                \n",
    "\n",
    "    return collect_func_time_samples(func=fit_func, count=count)\n",
    "\n",
    "\n",
    "def sweep_predict_func(model, func_id, require_compute, X, count):\n",
    "    _predict_func_attr = getattr(model, func_id)\n",
    "    predict_func = partial(lambda X: _predict_func_attr(X).compute(), X)\n",
    "    \n",
    "    return collect_func_time_samples(func=predict_func, count=count)\n",
    "    \n",
    "\n",
    "def performance_sweep(client, model, data_loader, hardware_type, worker_counts=[1], samples=1, load_samples=1, max_data_frac=1.0,\n",
    "                    predict_frac=1.0, scaling_type='weak', xy_fit=True, fit_requires_compute=False, update_workers_in_kwargs=True,\n",
    "                    response_dtype=np.float32, out_path='./perf_sweep.csv', append_to_existing=False, model_name=None, infer_with_fil=False,\n",
    "                    fit_func_id=\"fit\", predict_func_id=\"predict\", scaling_denom=None, post_fit_handler=None, model_args={}, model_kwargs={}):\n",
    "    \"\"\"\n",
    "    Primary performance sweep entrypoint.\n",
    "    \n",
    "    \n",
    "    Parameters\n",
    "    ------------\n",
    "    client: DASK client associated with the cluster we're interesting in collecting performance data for.\n",
    "    \n",
    "    model: Model object on which to gather performance data. This will be created and destroyed,\n",
    "        once for each element of 'worker_counts'\n",
    "    \n",
    "    data_loader: arbitrary data loading function that will be called to load the appropriate testing data.\n",
    "        Function that is responsible for loading and returning the data to be used for a given performance run. Function\n",
    "        signature must accept (client, fraction, and random_state). Client should be used to distribute data, and loaders\n",
    "        should utilize fraction and random_state with dask's dataframe.sample method to allow for control of how much data\n",
    "        is loaded.\n",
    "        \n",
    "        When called, its return value should be of the form: df, X, y, where df is the full dask_cudf dataframe, X is a\n",
    "        dask_cudf dataframe which contains all explanatory variables that will be passed to the 'fit' function, and y is a\n",
    "        dask_cudf series or dataframe that contains response variables which should be passed to fit/predict as fit(X, y)\n",
    "    \n",
    "    hardware_type: indicates the core hardware the current sweep is running on. ex. 'T4', 'V100', 'A100'\n",
    "    \n",
    "    worker_counts: List indicating the number of workers that should be swept. Ex [1, 2, 4]\n",
    "        worker counts must fit within the cluster associated with 'client', if the current DASK worker count is different\n",
    "        from what is requested on a given sweep, attempt to automatically scale the worker count. NOTE: this does not \n",
    "        mean we will scale the available cluster nodes, just the number of deployed worker pods.\n",
    "    \n",
    "    samples: number of fit/predict samples to record per worker count\n",
    "    \n",
    "    load_samples: number of times to sample data loads. This effectively times how long 'data_loader' runs.\n",
    "    \n",
    "    max_data_frac: maximum fraction of data to return.\n",
    "        Strong scaling: each run will utilize max_data_frac data.\n",
    "        Weak scaling: each run will utilize (current worker count) / (max worker count) * max_data_frac data.\n",
    "        \n",
    "    predict_frac: fraction of training data used to test inference\n",
    "    \n",
    "    scaling_type: values can be 'weak' or 'strong' indicating the type of scaling sweep to perform.\n",
    "    \n",
    "    xy_fit: indicates whether or not the model's 'fit' function is of the form (X, y), when xy_fit is False, we assume that\n",
    "        fit is of the form (X), as is the case with various unsupervised methods ex. KNN.\n",
    "    \n",
    "    fit_requires_compute: False generally, set this to True if the model's 'fit' function requires a corresponding '.compute()'\n",
    "        call to execute the required work.\n",
    "    \n",
    "    update_workers_in_kwargs: Some algorithms accept a 'workers' list, much like DASK, and will require their kwargs to have\n",
    "        workers populated. Setting this flag handles this automatically.\n",
    "        \n",
    "    response_dtype: defaults to np.float32, some algorithms require another dtype, such as int32\n",
    "    \n",
    "    out_path: path where performance data csv should be saved\n",
    "    \n",
    "    append_to_existing: When true, append results to an existing csv, otherwise overwrite.\n",
    "    \n",
    "    model_name: Override what we output as the model name\n",
    "    \n",
    "    fit_func_id: Defaults to 'fit', only set this if the model has a non-standard naming.\n",
    "    \n",
    "    predict_func_id: Defaults to 'predict', only set this if the model has a on-standard predict naming.\n",
    "    \n",
    "    scaling_denom: (weak scaling) defaults to max(workers) if unset. Specifies the maximum worker count that weak scaling\n",
    "        should scale against. For example, when using 1 worker in a weak scaling sweep, the worker will attempt to\n",
    "        process a fraction of the total data equal to 1/scaling_denom\n",
    "    \n",
    "    model_args: args that will be passed to the model's constructor\n",
    "    \n",
    "    model_kwargs: keyword args that will be passed to the model's constructor\n",
    "\n",
    "    Returns\n",
    "    --------\n",
    "    \n",
    "    \"\"\"\n",
    "    \n",
    "    cols = ['n_workers', 'sample_index', 'elapsed', 'type', 'algorithm', 'scaling_type', 'data_fraction', 'hardware', 'trial_id']\n",
    "    perf_df = cudf.DataFrame(columns=cols)\n",
    "    if (append_to_existing):\n",
    "        try:\n",
    "            perf_df = cudf.read_csv(out_path)\n",
    "        except:\n",
    "            pass\n",
    "    \n",
    "    model_name = model_name if model_name else str(model)\n",
    "    scaling_denom = scaling_denom if (scaling_denom is not None) else max(worker_counts)\n",
    "    max_data_frac = min(1.0, max_data_frac)\n",
    "\n",
    "    start_msg = f\"Starting {scaling_type}-scaling performance sweep for:\\n\"\n",
    "    start_msg += f\" model      : {model_name}\\n\"\n",
    "    start_msg += f\" data loader: {data_loader}.\\n\"\n",
    "    start_msg += f\"Configuration\\n\"\n",
    "    start_msg += \"==========================\\n\"\n",
    "    start_msg += f\"{'Worker counts':<25} : {worker_counts}\\n\"\n",
    "    start_msg += f\"{'Fit/Predict samples':<25} : {samples}\\n\"\n",
    "    start_msg += f\"{'Data load samples':<25} : {load_samples}\\n\"\n",
    "    start_msg += f\"- {'Max data fraction':<23} : {max_data_frac:0.2f}\\n\"\n",
    "    start_msg += f\" - {'Train':<22} : {max_data_frac:0.2f}\\n\"\n",
    "    start_msg += f\" - {'Infer':<22} : {predict_frac*max_data_frac:0.2f}\\n\"\n",
    "    start_msg += f\"{'Model fit':<25} : {'X ~ y' if xy_fit else 'X'}\\n\"\n",
    "    start_msg += f\"- {'Response DType':<23} : {response_dtype}\\n\"\n",
    "    start_msg += f\"{'Writing results to':<25} : {out_path}\\n\"\n",
    "    start_msg += f\"- {'Method':<23} : {'overwrite' if not append_to_existing else 'append'}\\n\"\n",
    "    print(start_msg, flush=True)\n",
    "    \n",
    "    for n in worker_counts:\n",
    "        fraction = (n / scaling_denom) * max_data_frac if scaling_type == 'weak' else max_data_frac\n",
    "        fraction = min(1.0, fraction)\n",
    "        record_template = {\n",
    "                            \"n_workers\": n,\n",
    "                            \"type\": \"predict\",\n",
    "                            \"algorithm\": model_name,\n",
    "                            \"scaling_type\": scaling_type,\n",
    "                            \"data_fraction\": fraction,\n",
    "                            \"hardware\": hardware_type,\n",
    "                            \"trial_id\": str(uuid.uuid4()),\n",
    "                          }\n",
    "        scale_workers(client, n)\n",
    "\n",
    "        print(f\"Sampling <{load_samples}> load times with {n} workers.\",\n",
    "              f\" With {fraction*100:0.1f} percent of total data\", flush=True)\n",
    "        \n",
    "        # Todo: add fractional selection passthrough\n",
    "        load_func = partial(data_loader, client=client, response_dtype=response_dtype, fraction=fraction, random_state=0)\n",
    "        X_train, y_train, X_infer, y_infer, load_timings = collect_load_time_samples(load_func=load_func, count=load_samples)\n",
    "        \n",
    "        perf_df = record_elapsed_timings_to_df(df=perf_df, timings=load_timings, type='load',\n",
    "                                                    record_template=record_template, columns=cols,write_to=out_path)\n",
    "\n",
    "        print(f\"Finished loading <{load_samples}>, samples, to <{n}>\",\n",
    "              f\"workers with a mean time of {np.mean(load_timings)/1e9:0.4f} sec.\", flush=True)\n",
    "        print(f\"Sweeping {model_name} '{fit_func_id}' with <{n}> workers. Sampling\",\n",
    "              f\" <{samples}> times with {fraction*100:0.1f} percent of total data.\", flush=True)\n",
    "\n",
    "        if (update_workers_in_kwargs and 'workers' in model_kwargs):\n",
    "            model_kwargs['workers'] = workers = list(client.has_what().keys())\n",
    "    \n",
    "        m = model(*model_args, **model_kwargs)\n",
    "        if (fit_func_id):\n",
    "            fit_timings = sweep_fit_func(model=m, func_id=fit_func_id,\n",
    "                                             require_compute=fit_requires_compute,\n",
    "                                             X=X_train, y=y_train, xy_fit=xy_fit, count=samples)\n",
    "\n",
    "            perf_df = record_elapsed_timings_to_df(df=perf_df, timings=fit_timings, type='fit',\n",
    "                                                        record_template=record_template, columns=cols, write_to=out_path)\n",
    "\n",
    "            print(f\"Finished gathering <{samples}>, 'fit' samples using <{n}>\",\n",
    "                  f\" workers, with a mean time of {np.mean(fit_timings)/1e9:0.4f} sec.\", flush=True)\n",
    "        else:\n",
    "            print(f\"Skipping fit sweep, fit_func_id is None\")\n",
    "            \n",
    "        if (post_fit_handler):\n",
    "            post_fit_handler(m)\n",
    "\n",
    "        if (predict_func_id):\n",
    "            print(f\"Sweeping {model_name} '{predict_func_id}' with <{n}> workers.\" \n",
    "                  f\" Sampling <{samples}> times with {fraction*predict_frac*100:0.1f} percent of total data.\", flush=True)\n",
    "            predict_timings = sweep_predict_func(model=m, func_id=predict_func_id,\n",
    "                                                     require_compute=True, X=X_infer, count=samples)\n",
    "\n",
    "            perf_df = record_elapsed_timings_to_df(df=perf_df, timings=predict_timings, type='predict',\n",
    "                                                        record_template=record_template, columns=cols, write_to=out_path)\n",
    "            \n",
    "            print(f\"Finished gathering <{samples}>, 'predict' samples using <{n}>\",\n",
    "                  f\" workers, with a mean time of {np.mean(predict_timings)/1e9:0.4f} sec.\",\n",
    "                  flush=True)\n",
    "        else:\n",
    "            print(f\"Skipping inference sweep. predict_func_id is None\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "genuine-evening",
   "metadata": {},
   "source": [
    "### Vis and Analysis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "communist-complaint",
   "metadata": {},
   "outputs": [],
   "source": [
    "def simple_ci(df, fields, groupby):\n",
    "    gbdf = df[fields].groupby(groupby).agg(['mean', 'std', 'count'])   \n",
    "    \n",
    "    ci = (1.96 + gbdf['elapsed']['std'] / np.sqrt(gbdf['elapsed']['count']))\n",
    "    \n",
    "    ci_df = ci.reset_index()\n",
    "    ci_df['ci.low'] = gbdf['elapsed'].reset_index()['mean'] - ci_df[0]\n",
    "    ci_df['ci.high'] = gbdf['elapsed'].reset_index()['mean'] + ci_df[0]\n",
    "    \n",
    "    return ci_df\n",
    "\n",
    "\n",
    "def visualize_csv_data(csv_path, filter_query=None):\n",
    "    import pandas as pd\n",
    "    df = cudf.read_csv(csv_path)\n",
    "    \n",
    "    fields = ['elapsed', 'elapsed_sec', 'type', 'n_workers', 'hardware', 'scaling_type']\n",
    "    groupby = ['n_workers', 'type', 'hardware', 'scaling_type']\n",
    "    df['elapsed_sec'] = df['elapsed']/1e9\n",
    "\n",
    "    ci_df = simple_ci(df, fields, groupby=groupby)\n",
    "\n",
    "    # Rescale to seconds\n",
    "    ci_df[['ci.low', 'ci.high']] = ci_df[['ci.low', 'ci.high']]/1e9\n",
    "\n",
    "    # Print confidence intervals\n",
    "    print(ci_df[['hardware', 'n_workers', 'type', 'ci.low', 'ci.high']][ci_df['type'] != 'load'])\n",
    "\n",
    "    sns.set_theme(style=\"whitegrid\")\n",
    "    sns.set(rc={'figure.figsize':(20, 10)}, font_scale=2)\n",
    "\n",
    "    # Boxplots for elapsed time at each worker count.\n",
    "    # [df[fields].type != 'load']\n",
    "    plot_df = df[fields].to_pandas()\n",
    "    plot_df = plot_df.query(\"type != 'load'\")\n",
    "    if (filter_query):\n",
    "        plot_df = plot_df.query(filter_query)\n",
    "        \n",
    "    ax = sns.catplot(data=plot_df, x=\"n_workers\", y=\"elapsed_sec\",\n",
    "                     col=\"type\", row=\"scaling_type\", hue=\"hardware\", kind=\"box\",\n",
    "                     height=8, order=None)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "basic-charm",
   "metadata": {},
   "source": [
    "### Taxi Data Configuration (Medium)\n",
    "We can use the parquet data from the anaconda public repo here. Which will illustrate how much faster it is to read parquet, and gives us around 150 million rows of data to work with."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "neutral-mortality",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Uncomment to test with Taxi Dataset\n",
    "append_to_existing = True\n",
    "samples = 5\n",
    "load_samples = 1\n",
    "worker_counts = [8]\n",
    "scaling_denom = 8\n",
    "hardware_type = None\n",
    "max_data_frac = 0.75\n",
    "scale_type = 'weak' # weak | strong\n",
    "out_prefix = 'taxi_medium'\n",
    "data_loader = taxi_parquet_data_loader\n",
    "\n",
    "    \n",
    "if (not hardware_type):\n",
    "    raise RuntimeError(\"Please specify the hardware type for this run! ex. (T4, V100, A100)\")\n",
    "\n",
    "sweep_kwargs = {\n",
    "    'append_to_existing': append_to_existing,\n",
    "    'samples': samples,\n",
    "    'load_samples': load_samples,\n",
    "    'worker_counts': worker_counts,\n",
    "    'scaling_denom': scaling_denom,\n",
    "    'hardware_type': hardware_type,\n",
    "    'data_loader': data_loader,\n",
    "    'max_data_frac': max_data_frac,\n",
    "    'scaling_type': scale_type\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "authentic-mount",
   "metadata": {},
   "outputs": [],
   "source": [
    "taxi_parquet_path = [\"gs://anaconda-public-data/nyc-taxi/nyc.parquet\"]\n",
    "estimated_rows = estimate_df_rows(client, files=taxi_parquet_path, testpct=0.0001)\n",
    "print(estimated_rows)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "acoustic-terry",
   "metadata": {},
   "source": [
    "## Taxi Data Configuration (Large)\n",
    "The largest dataset we'll work with, contains up to 450 million rows of taxi data, stored as CSV files."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "committed-physiology",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Uncomment to sweep with the large Taxi Dataset\n",
    "append_to_existing = True\n",
    "samples = 5\n",
    "load_samples = 1\n",
    "worker_counts = [8,4,2]\n",
    "scaling_denom = 8\n",
    "hardware_type = None\n",
    "max_data_frac = 0.75\n",
    "scale_type = 'weak'\n",
    "out_prefix = 'taxi_large'\n",
    "data_loader = taxi_csv_data_loader\n",
    "\n",
    "\n",
    "if (not hardware_type):\n",
    "    raise RuntimeError(\"Please specify the hardware type for this run! ex. (T4, V100, A100)\")\n",
    "    \n",
    "\n",
    "sweep_kwargs = {\n",
    "    'append_to_existing': append_to_existing,\n",
    "    'samples': samples,\n",
    "    'load_samples': load_samples,\n",
    "    'worker_counts': worker_counts,\n",
    "    'scaling_denom': scaling_denom,\n",
    "    'hardware_type': hardware_type,\n",
    "    'data_loader': data_loader,\n",
    "    'max_data_frac': max_data_frac,\n",
    "    'scaling_type': scale_type\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "clear-sword",
   "metadata": {},
   "source": [
    "# Mortgage Data Configuration (Very Large)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "identical-value",
   "metadata": {},
   "outputs": [],
   "source": [
    "append_to_existing = True\n",
    "samples = 5\n",
    "load_samples = 1\n",
    "worker_counts = [4]\n",
    "scaling_denom = 8\n",
    "hardware_type = None\n",
    "max_data_frac = 0.15\n",
    "scale_type = 'weak'\n",
    "out_prefix = 'mortgage_large'\n",
    "data_loader = mortgage_parquet_loader\n",
    "\n",
    "if (not hardware_type):\n",
    "    raise RuntimeError(\"Please specify the hardware type for this run! ex. (T4, V100, A100)\")\n",
    "\n",
    "sweep_kwargs = {\n",
    "    'append_to_existing': append_to_existing,\n",
    "    'samples': samples,\n",
    "    'load_samples': load_samples,\n",
    "    'worker_counts': worker_counts,\n",
    "    'scaling_denom': scaling_denom,\n",
    "    'hardware_type': hardware_type,\n",
    "    'data_loader': data_loader,\n",
    "    'max_data_frac': max_data_frac,\n",
    "    'scaling_type': scale_type\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "elect-replication",
   "metadata": {},
   "source": [
    "# ETL Exploration CSV vs Parquet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "desperate-sarah",
   "metadata": {},
   "outputs": [],
   "source": [
    "remap = {}\n",
    "remap['tpep_pickup_datetime'] = 'pickup_datetime'\n",
    "remap['tpep_dropoff_datetime'] = 'dropoff_datetime'\n",
    "remap['ratecodeid'] = 'rate_code'\n",
    "\n",
    "#create a list of columns & dtypes the df must have\n",
    "must_haves = {\n",
    " 'pickup_datetime': 'datetime64[ms]',\n",
    " 'dropoff_datetime': 'datetime64[ms]',\n",
    " 'passenger_count': 'int32',\n",
    " 'trip_distance': 'float32',\n",
    " 'pickup_longitude': 'float32',\n",
    " 'pickup_latitude': 'float32',\n",
    " 'rate_code': 'int32',\n",
    " 'dropoff_longitude': 'float32',\n",
    " 'dropoff_latitude': 'float32',\n",
    " 'fare_amount': 'float32'\n",
    "}\n",
    "\n",
    "# apply a list of filter conditions to throw out records with missing or outlier values\n",
    "query_fragments = [\n",
    "    'fare_amount > 0 and fare_amount < 500',\n",
    "    'passenger_count > 0 and passenger_count < 6',\n",
    "    'pickup_longitude > -75 and pickup_longitude < -73',\n",
    "    'dropoff_longitude > -75 and dropoff_longitude < -73',\n",
    "    'pickup_latitude > 40 and pickup_latitude < 42',\n",
    "    'dropoff_latitude > 40 and dropoff_latitude < 42'\n",
    "]\n",
    "\n",
    "workers = client.has_what().keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "imported-norman",
   "metadata": {},
   "outputs": [],
   "source": [
    "base_path = 'gcs://anaconda-public-data/nyc-taxi/csv'\n",
    "\n",
    "with SimpleTimer() as timer_csv:\n",
    "    df_csv_2014 = dask_cudf.read_csv(f'{base_path}/2014/yellow_*.csv', chunksize=25e6, dtype={' tolls_amount': 'float64'})\n",
    "    df_csv_2014 = clean(df_csv_2014, remap, must_haves)\n",
    "    df_csv_2014 = df_csv_2014.query(' and '.join(query_fragments))\n",
    "    \n",
    "    with dask.annotate(workers=set(workers)):\n",
    "        df_csv_2014 = client.persist(collections=df_csv_2014)\n",
    "        \n",
    "    wait(df_csv_2014)\n",
    "\n",
    "print(df_csv_2014.columns)\n",
    "rows_csv = df_csv_2014.iloc[:,0].shape[0].compute()\n",
    "print(f\"CSV load took {timer_csv.elapsed/1e9} sec. For {rows_csv} rows of data => {rows_csv/(timer_csv.elapsed/1e9)} rows/sec\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "patent-albert",
   "metadata": {},
   "outputs": [],
   "source": [
    "client.cancel(df_csv_2014)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "increased-resolution",
   "metadata": {},
   "outputs": [],
   "source": [
    "with SimpleTimer() as timer_parquet:\n",
    "    df_parquet = dask_cudf.read_parquet(f'gs://anaconda-public-data/nyc-taxi/nyc.parquet', chunksize=25e6)\n",
    "    df_parquet = clean(df_parquet, remap, must_haves)\n",
    "    df_parquet = df_parquet.query(' and '.join(query_fragments))\n",
    "    \n",
    "    with dask.annotate(workers=set(workers)):\n",
    "        df_parquet = client.persist(collections=df_parquet)\n",
    "    \n",
    "    wait(df_parquet)\n",
    "\n",
    "print(df_parquet.columns)\n",
    "rows_parquet = df_parquet.iloc[:,0].shape[0].compute()\n",
    "print(f\"Parquet load took {timer_parquet.elapsed/1e9} sec. For {rows_parquet} rows of data => {rows_parquet/(timer_parquet.elapsed/1e9)} rows/sec\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "conscious-render",
   "metadata": {},
   "outputs": [],
   "source": [
    "client.cancel(df_parquet)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "weird-triumph",
   "metadata": {},
   "source": [
    "Speedup with 8 T4 nodes should be approximately 3-5x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "burning-business",
   "metadata": {},
   "outputs": [],
   "source": [
    "speedup = (rows_parquet/(timer_parquet.elapsed/1e9))/(rows_csv/(timer_csv.elapsed/1e9))\n",
    "print(speedup)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "logical-conducting",
   "metadata": {},
   "source": [
    "# cuML Algorithms -- Performance Sweeps"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "continued-plane",
   "metadata": {},
   "source": [
    "### RandomForestRegressor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "collect-broadcasting",
   "metadata": {},
   "outputs": [],
   "source": [
    "from cuml.dask.ensemble import RandomForestRegressor\n",
    "\n",
    "rf_kwargs = {\n",
    "    \"workers\": client.has_what().keys(),\n",
    "    \"n_estimators\": 200,\n",
    "    \"max_depth\": 6 # match xgboost's default\n",
    "}\n",
    "rf_csv_path = f\"./{out_prefix}_random_forest_regression.csv\"\n",
    "\n",
    "performance_sweep(client=client, model=RandomForestRegressor,\n",
    "                **sweep_kwargs,\n",
    "                out_path=rf_csv_path,\n",
    "                response_dtype=np.int32,\n",
    "                model_kwargs=rf_kwargs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "structured-friendship",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "rf_csv_path = f\"./{out_prefix}_random_forest_regression.csv\"\n",
    "visualize_csv_data(rf_csv_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "packed-adoption",
   "metadata": {},
   "source": [
    "##### XGBoost"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "swedish-burlington",
   "metadata": {},
   "outputs": [],
   "source": [
    "import xgboost as xgb\n",
    "\n",
    "xg_args = [client]\n",
    "xg_kwargs = {\n",
    "    'params': {\n",
    "        'tree_method': 'gpu_hist',\n",
    "    },\n",
    "    'num_boost_round': 100\n",
    "}\n",
    "\n",
    "xgb_csv_path = f'./{out_prefix}_xgb.csv'\n",
    "\n",
    "\n",
    "def xgb_post_fit(model, gcs_model_path, project, token):\n",
    "    model.trained_model.save_model(\"./xgb_perf_sweep_model.model\")\n",
    "    gcs_fs = gcsfs.core.GCSFileSystem(project=project, token=token)\n",
    "    \n",
    "    # Push the model to central storage\n",
    "    gcs_fs.put(\"./xgb_perf_sweep_model.model\", gcs_model_path)\n",
    "\n",
    "    client.run(attach_fil_to_worker, gcs_model_path, project, token,\n",
    "               wait=True)\n",
    "\n",
    "\n",
    "def attach_fil_to_worker(gcs_model_path, project, token):\n",
    "    worker = get_worker()\n",
    "\n",
    "    gcs_fs = gcsfs.core.GCSFileSystem(project=project, token=token)\n",
    "    with gcs_fs.open(gcs_model_path, 'rb') as fmod:\n",
    "        data = fmod.read()\n",
    "        \n",
    "    with open(\"./xgb_perf_sweep_model.model\", 'wb') as wmod:\n",
    "        wmod.write(data)\n",
    "    \n",
    "    worker.data[\"fil_model\"] = ForestInference.load(\"./xgb_perf_sweep_model.model\",\n",
    "                               algo=\"BATCH_TREE_REORG\",\n",
    "                               output_class=False,\n",
    "                               model_type='xgboost')\n",
    "\n",
    "\n",
    "def predict_on_worker(X):\n",
    "    worker = get_worker()\n",
    "    worker.data['fil_model'].predict(X)\n",
    "    \n",
    "    return\n",
    "    \n",
    "    \n",
    "class XGBProxy():\n",
    "    \"\"\"\n",
    "    Create a simple API wrapper around XGBoost so that it supports our fit/predict workflow.\n",
    "    \n",
    "    Parameters\n",
    "    -------------\n",
    "    data_loader: data loader object intended to be used by performance sweep.\n",
    "    \"\"\"\n",
    "    def __init__(self, data_loader):\n",
    "        self.args = []\n",
    "        self.kwargs = {}\n",
    "        self.data_loader = data_loader\n",
    "        self.trained_model = None\n",
    "        self.saved = False\n",
    "        self.fm = None\n",
    "        self.X_train = None\n",
    "        self.y_train = None\n",
    "        self.X_infer = None\n",
    "        self.y_infer = None\n",
    "        \n",
    "    def loader(self, client, response_dtype, fraction, random_state):\n",
    "        \"\"\"\n",
    "        Wrap the data loader method so that it creates a DMatrix from the returned data.\n",
    "        \"\"\"\n",
    "        for df in [self.X_train, self.X_infer, self.y_train, self.y_infer]:\n",
    "            if df is not None:\n",
    "                del df\n",
    "        \n",
    "        X_train, y_train, X_infer, y_infer = self.data_loader(client, response_dtype, fraction, random_state)\n",
    "        self.X_train = X_train\n",
    "        self.y_train = y_train\n",
    "        self.X_infer = X_infer\n",
    "        self.y_infer = y_infer\n",
    "        \n",
    "        self.dmatrix = xgb.dask.DaskDMatrix(client, X_train, y_train)\n",
    "        \n",
    "        return X_train, y_train, X_infer, y_infer\n",
    "\n",
    "    def __call__(self, *args, **kwargs):\n",
    "        \"\"\"\n",
    "        Acts as a pseudo init function which initializes our model args.\n",
    "        \"\"\"\n",
    "        self.args = args\n",
    "        self.kwargs = kwargs\n",
    "        \n",
    "        return self\n",
    "        \n",
    "    def fit(self, X):\n",
    "        \"\"\"\n",
    "        Wrap dask.train, and store the model on our proxy object.\n",
    "        \"\"\"\n",
    "        if (self.trained_model):\n",
    "            del self.trained_model\n",
    "            \n",
    "        self.trained_model = xgb.dask.train(*self.args,\n",
    "                              dtrain=self.dmatrix,\n",
    "                              evals=[(self.dmatrix, 'train')],\n",
    "                              **self.kwargs)['booster']\n",
    "        \n",
    "        return self\n",
    "    \n",
    "    def predict(self, X):\n",
    "        part_map = self.X_infer.map_partitions(predict_on_worker,\n",
    "                                               meta=pd.Series(dtype='float32'))\n",
    "        \n",
    "        return part_map\n",
    "    \n",
    "\n",
    "_xgb_post_fit = partial(xgb_post_fit,\n",
    "                gcs_model_path=YOUR_GCS_MODEL_PATH,\n",
    "                project=YOUR_GCP_PROJECT,\n",
    "                token=YOUR_GCP_TOKEN)\n",
    "xgb_proxy = XGBProxy(data_loader)\n",
    "performance_sweep(client=client, model=xgb_proxy, data_loader=xgb_proxy.loader, hardware_type=hardware_type,\n",
    "                worker_counts=worker_counts, \n",
    "                samples=samples,\n",
    "                load_samples=load_samples,\n",
    "                max_data_frac=max_data_frac, \n",
    "                scaling_type=scale_type,\n",
    "                out_path=xgb_csv_path,\n",
    "                append_to_existing=append_to_existing,\n",
    "                update_workers_in_kwargs=False,\n",
    "                xy_fit=False,\n",
    "                scaling_denom = scaling_denom,\n",
    "                post_fit_handler=_xgb_post_fit,\n",
    "                model_args=xg_args,\n",
    "                model_kwargs=xg_kwargs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "rural-token",
   "metadata": {},
   "outputs": [],
   "source": [
    "xgb_csv_path = f'./{out_prefix}_xgb.csv'\n",
    "visualize_csv_data(xgb_csv_path)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}