notebooks/experiment_spark_lr.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "# Criteo 1 TiB benchmark - Spark.ML logistic regression\n",
    "\n",
    "Specialization of the experiment notebook for Spark.ML logistic regression."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Table of contents\n",
    "\n",
    "* [Configuration](#Configuration)\n",
    "* [Distributed training](#Distributed-training)\n",
    "* [End](#End)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "%load_ext autotime\n",
    "%matplotlib inline\n",
    "\n",
    "from __future__ import print_function"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Configuration\n",
    "[_(back to toc)_](#Table-of-contents)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Paths:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "libsvm_data_remote_path = 'criteo/libsvm'\n",
    "local_runtime_path = 'criteo/runtime'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "\n",
    "libsvm_train_template = os.path.join(libsvm_data_remote_path, 'train', '{}')\n",
    "libsvm_test_template = os.path.join(libsvm_data_remote_path, 'test', '{}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def ensure_directory_exists(path):\n",
    "    if not os.path.exists(path):\n",
    "        os.makedirs(path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Samples to take:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "train_samples = [\n",
    "    10000, 30000,  # tens of thousands\n",
    "    100000, 300000,  # hundreds of thousands\n",
    "    1000000, 3000000,  # millions\n",
    "    10000000, 30000000,  # tens of millions\n",
    "    100000000, 300000000,  # hundreds of millions\n",
    "    1000000000, 3000000000,  # billions\n",
    "]\n",
    "\n",
    "test_samples = [1000000]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Spark configuration and initialization:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "total_cores = 256"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "executor_cores = 4\n",
    "executor_instances = total_cores / executor_cores\n",
    "memory_per_core = 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "app_name = 'Criteo experiment - LR on 128 cores'\n",
    "\n",
    "master = 'yarn'\n",
    "\n",
    "settings = {\n",
    "    'spark.network.timeout': '600',\n",
    "    \n",
    "    'spark.driver.cores': '16',\n",
    "    'spark.driver.maxResultSize': '16G',\n",
    "    'spark.driver.memory': '32G',\n",
    "    \n",
    "    'spark.executor.cores': str(executor_cores),\n",
    "    'spark.executor.instances': str(executor_instances),\n",
    "    'spark.executor.memory': str(memory_per_core * executor_cores) + 'G',\n",
    "    \n",
    "    'spark.speculation': 'true',\n",
    "    \n",
    "    'spark.yarn.queue': 'root.HungerGames',\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "from pyspark.sql import SparkSession\n",
    "\n",
    "\n",
    "builder = SparkSession.builder\n",
    "\n",
    "builder.appName(app_name)\n",
    "builder.master(master)\n",
    "for k, v in settings.items():\n",
    "    builder.config(k, v)\n",
    "\n",
    "spark = builder.getOrCreate()\n",
    "sc = spark.sparkContext\n",
    "\n",
    "sc.setLogLevel('ERROR')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Logging:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import sys\n",
    "import logging\n",
    "reload(logging)\n",
    "\n",
    "\n",
    "handler = logging.StreamHandler(stream=sys.stdout)\n",
    "formatter = logging.Formatter('[%(asctime)s] %(message)s')\n",
    "handler.setFormatter(formatter)\n",
    "\n",
    "ensure_directory_exists(local_runtime_path)\n",
    "file_handler = logging.FileHandler(filename=os.path.join(local_runtime_path, 'mylog.log'), mode='a')\n",
    "file_handler.setFormatter(formatter)\n",
    "\n",
    "logger = logging.getLogger()\n",
    "logger.addHandler(handler)\n",
    "logger.addHandler(file_handler)\n",
    "logger.setLevel(logging.DEBUG)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "logger.info('Spark version: %s.', spark.version)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot measurements:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "import pandas\n",
    "\n",
    "\n",
    "def extract_data_for_plotting(df, what):\n",
    "    return reduce(\n",
    "        lambda left, right: pandas.merge(left, right, how='outer', on='Train size'),\n",
    "        map(\n",
    "            lambda name: df[df.Engine == name][['Train size', what]].rename(columns={what: name}),\n",
    "            df.Engine.unique(),\n",
    "        ),\n",
    "    )   \n",
    "\n",
    "def plot_stuff(df, what, ylabel=None, **kwargs):\n",
    "    data = extract_data_for_plotting(df, what).set_index('Train size')\n",
    "    ax = data.plot(marker='o', figsize=(6, 6), title=what, grid=True, linewidth=2.0, **kwargs)  # xlim=(1e4, 4e9)\n",
    "    if ylabel is not None:\n",
    "        ax.set_ylabel(ylabel)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's name samples as their shortened \"engineering\" notation - 1e5 is 100k etc.:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def sample_name(sample):\n",
    "    return str(sample)[::-1].replace('000', 'k')[::-1]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Distributed training\n",
    "[_(back to toc)_](#Table-of-contents)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Loading of LibSVM data as Spark.ML dataset:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import hashlib\n",
    "import math\n",
    "import struct\n",
    "from pyspark.ml.linalg import SparseVector\n",
    "\n",
    "\n",
    "total_features = 100000\n",
    "\n",
    "\n",
    "def hash_fun(x):\n",
    "    return int(struct.unpack('L', hashlib.md5(x).digest()[:8])[0] % total_features)\n",
    "\n",
    "def parse_kv_pair(kv):\n",
    "    k, _, v = kv.partition(':')\n",
    "    return int(k), int(v)\n",
    "\n",
    "def parse_libsvm_line(line):\n",
    "    parts = line.split(' ')\n",
    "    \n",
    "    label = int(parts[0])\n",
    "    pairs = map(parse_kv_pair, parts[1:])\n",
    "    \n",
    "    for i in range(len(pairs)):\n",
    "        k, v = pairs[i]\n",
    "        if k < 14:\n",
    "            if v > 2:\n",
    "                pairs[i] = (k, int(math.log(v) ** 2))\n",
    "        else:\n",
    "            pairs[i] = (k, '{:08x}'.format(v))\n",
    "    \n",
    "    indices = sorted({hash_fun('{}_{}'.format(k, v)) for k, v in pairs})\n",
    "    values = [1.0] * len(indices)\n",
    "    \n",
    "    return (label, SparseVector(total_features, indices, values))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "task_splitting = 1  # tasks per core\n",
    "\n",
    "def load_ml_data(template, sample):\n",
    "    path = template.format(sample_name(sample))\n",
    "    return sc.textFile(path).map(parse_libsvm_line).toDF(['label', 'features']).repartition(executor_cores * executor_instances * task_splitting)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Evaluating a model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from matplotlib import pyplot\n",
    "from sklearn.metrics import auc, log_loss, roc_curve\n",
    "\n",
    "\n",
    "def calculate_roc(predictions):\n",
    "    labels, scores = zip(*predictions.rdd.map(lambda row: (row.label, row.probability[1])).collect())\n",
    "    fpr, tpr, _ = roc_curve(labels, scores)\n",
    "    roc_auc = auc(fpr, tpr)\n",
    "    ll = log_loss(labels, scores)\n",
    "    return fpr, tpr, roc_auc, ll\n",
    "\n",
    "def evaluate_model(name, model, test, train=None):\n",
    "    metrics = dict()\n",
    "    \n",
    "    figure = pyplot.figure(figsize=(6, 6))\n",
    "    ax = figure.gca()\n",
    "    ax.set_title('ROC - ' + name)\n",
    "    \n",
    "    if train is not None:\n",
    "        train_predictions = model.transform(train)\n",
    "        train_fpr, train_tpr, train_roc_auc, train_log_loss = calculate_roc(train_predictions)\n",
    "        \n",
    "        metrics['train_roc_auc'] = train_roc_auc\n",
    "        metrics['train_log_loss'] = train_log_loss\n",
    "        \n",
    "        ax.plot(train_fpr, train_tpr, linewidth=2.0, label='train (auc = {:.3f})'.format(train_roc_auc))\n",
    "    \n",
    "    test_predictions = model.transform(test)\n",
    "    test_fpr, test_tpr, test_roc_auc, test_log_loss = calculate_roc(test_predictions)\n",
    "    \n",
    "    metrics['test_roc_auc'] = test_roc_auc\n",
    "    metrics['test_log_loss'] = test_log_loss\n",
    "    \n",
    "    ax.plot(test_fpr, test_tpr, linewidth=2.0, label='test (auc = {:.3f})'.format(test_roc_auc))\n",
    "    \n",
    "    ax.plot([0.0, 1.0], [0.0, 1.0], linestyle='--', c='gray')\n",
    "    ax.legend()\n",
    "    pyplot.show()\n",
    "    \n",
    "    return metrics"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Models to work on:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from pyspark.ml.classification import (\n",
    "    LogisticRegression,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "classifiers = {\n",
    "    'lr': LogisticRegression(regParam=0.03),\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Monkey-patch RDDs and DataFrames for context persistence:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import pyspark\n",
    "\n",
    "\n",
    "def enter_method(self):\n",
    "    self.persist()\n",
    "\n",
    "def exit_method(self,exc_type, exc, traceback):\n",
    "    self.unpersist()\n",
    "\n",
    "\n",
    "pyspark.sql.dataframe.DataFrame.__enter__ = enter_method\n",
    "pyspark.sql.dataframe.DataFrame.__exit__ = exit_method"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Do distributed training:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "import time\n",
    "\n",
    "from pyspark.ml.evaluation import BinaryClassificationEvaluator\n",
    "\n",
    "\n",
    "test_sample = test_samples[-1]\n",
    "\n",
    "evaluator = BinaryClassificationEvaluator(labelCol='label', rawPredictionCol='probability', metricName='areaUnderROC')\n",
    "\n",
    "new_quality_data = []\n",
    "new_data_rows = []\n",
    "\n",
    "test_dfs = dict()\n",
    "\n",
    "logger.info('Loading \"%s\" test samples.', test_sample)\n",
    "test_df = load_ml_data(libsvm_test_template, test_sample)\n",
    "with test_df:\n",
    "    logger.info('Loaded \"%s\" lines.', test_df.count())\n",
    "\n",
    "    for train_sample in train_samples:\n",
    "\n",
    "        logger.info('Working on \"%s\" train sample.', train_sample)\n",
    "\n",
    "        for classifier_name, classifier in classifiers.items():\n",
    "\n",
    "            logger.info('Training a model \"%s\" on sample \"%s\".', classifier_name, train_sample)\n",
    "\n",
    "            logger.info('Loading \"%s\" train samples.', train_sample)\n",
    "            train_df = load_ml_data(libsvm_train_template, train_sample)\n",
    "            with train_df:\n",
    "                logger.info('Loaded \"%s\" lines.', train_df.count())\n",
    "                \n",
    "                logger.info('Training a model \"%s\" on sample \"%s\".', classifier_name, train_sample)\n",
    "\n",
    "                start = time.time()\n",
    "                model = classifier.fit(train_df)\n",
    "                duration = time.time() - start\n",
    "\n",
    "                logger.info('Training a model \"%s\" on sample \"%s\" took \"%g\" seconds.', classifier_name, train_sample, duration)\n",
    "\n",
    "                logger.info('Evaluating the model \"%s\" trained on sample \"%s\".', classifier_name, train_sample)\n",
    "                metrics = evaluate_model(classifier_name + ' - ' + sample_name(train_sample), model, test_df, train=(train_df if train_sample <= 1000000 else None))\n",
    "\n",
    "                test_predictions = model.transform(test_df)\n",
    "                ml_metric_value = evaluator.evaluate(test_predictions)\n",
    "\n",
    "                logger.info(\n",
    "                    'For the model \"%s\" trained on sample \"%s\" metrics are: \"%s\"; ROC AUC calculated by Spark is \"%s\".',\n",
    "                    classifier_name,\n",
    "                    train_sample,\n",
    "                    metrics,\n",
    "                    ml_metric_value,\n",
    "                )\n",
    "\n",
    "                data_row = {\n",
    "                    'Train size': train_sample,\n",
    "                    'ROC AUC': metrics['test_roc_auc'],\n",
    "                    'Log loss': metrics['test_log_loss'],\n",
    "                    'Duration': duration,\n",
    "                    'Engine': classifier_name,\n",
    "                }\n",
    "                new_quality_data.append(data_row)\n",
    "                data_row_string = '\\t'.join(str(data_row[field]) for field in ['Engine', 'Train size', 'ROC AUC', 'Log loss', 'Duration'])\n",
    "                new_data_rows.append(data_row_string)\n",
    "                logger.info('Data row: \"%s\".', data_row_string)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "measurements_df = pandas.DataFrame(new_quality_data).sort_values(by=['Train size'])\n",
    "plot_stuff(measurements_df, 'ROC AUC', logx=True)\n",
    "plot_stuff(measurements_df, 'Log loss', logx=True, ylim=(0.13, 0.18))\n",
    "plot_stuff(measurements_df, 'Duration', loglog=True, ylabel='s')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "for row in new_data_rows:\n",
    "    print(row)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## End\n",
    "[_(back to toc)_](#Table-of-contents)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Work done, stop Spark:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "spark.stop()"
   ]
  }
 ],
 "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.6.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}