examples/mnist/spark_horovod.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# The MNIST Application"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This example demostrates integration of the Horovod MPI-based distributed deep learning framework and the Spark platform within the context of the MNIST application. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import time\n",
    "from datetime import timedelta, datetime, tzinfo\n",
    "\n",
    "from pyspark.sql import SparkSession\n",
    "spark = SparkSession.builder.appName(\"spark-horovod-mnist\").getOrCreate()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "## Initialize the Spark RDD collection associated with MPI workers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "partitions = 4\n",
    "\n",
    "# Read the PMIx environmental variables\n",
    "env = {}\n",
    "with open('pmixsrv.env', 'r') as f:\n",
    "    lines = f.read().splitlines() \n",
    "    for line in lines:\n",
    "        words = line.split(\"=\")\n",
    "        env[words[0]] = words[1]\n",
    "        \n",
    "env[\"PATH\"] = os.getenv(\"PATH\")\n",
    "env[\"LD_LIBRARY_PATH\"] = os.getenv(\"LD_LIBRARY_PATH\")\n",
    "env[\"http_proxy\"] = os.getenv(\"http_proxy\")\n",
    "env[\"https_proxy\"] = os.getenv(\"https_proxy\")\n",
    "\n",
    "arg = []\n",
    "for id in range(partitions):\n",
    "    arg.append(env)\n",
    "\n",
    "rdd = spark.sparkContext.parallelize(arg, partitions)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Download the MNIST dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[0, 1, 2, 3]"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def read_data_sets(pid, it):\n",
    "    \n",
    "    import tensorflow as tf\n",
    "    \n",
    "    learn = tf.contrib.learn\n",
    "    learn.datasets.mnist.read_data_sets('MNIST-data-%d' % pid)\n",
    "    \n",
    "    yield pid\n",
    "\n",
    "rdd.mapPartitionsWithIndex(read_data_sets).collect()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "## Train the Horovod MPI-based distributed engine on the Spark workers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": false,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# The train method is defined after horovod'd example\n",
    "# https://github.com/uber/horovod/blob/master/examples/tensorflow_mnist.py\n",
    "\n",
    "def train(pid, parts):\n",
    "        \n",
    "    import tensorflow as tf\n",
    "    import horovod.tensorflow as hvd\n",
    "    import mnist_app\n",
    "    \n",
    "    log_string = mnist_app.get_log_string(1024)\n",
    "    \n",
    "    # define the MPI enviromental variables     \n",
    "    os.environ[\"PMIX_RANK\"] = str(pid)\n",
    "    for env in parts:\n",
    "        for key in env:\n",
    "            os.environ[key] = env[key]\n",
    "       \n",
    "    # initialize Horovod   \n",
    "    hvd.init()\n",
    "    \n",
    "    # Extract the MNIST dataset\n",
    "    learn = tf.contrib.learn\n",
    "    mnist = learn.datasets.mnist.read_data_sets('MNIST-data-%d' % hvd.rank())\n",
    "    \n",
    "    # Build model...\n",
    "    with tf.name_scope('input'):\n",
    "        image = tf.placeholder(tf.float32, [None, 784], name='image')\n",
    "        label = tf.placeholder(tf.float32, [None], name='label')\n",
    "    predict, loss = mnist_app.conv_model(image, label, tf.contrib.learn.ModeKeys.TRAIN)\n",
    "    \n",
    "    global_step = tf.train.get_or_create_global_step()\n",
    "    \n",
    "    # Horovod: add Horovod Distributed Optimizer.\n",
    "    opt = tf.train.RMSPropOptimizer(0.001 * hvd.size())\n",
    "    opt = hvd.DistributedOptimizer(opt)\n",
    "    train_op = opt.minimize(loss, global_step=global_step)\n",
    "    \n",
    "    # The MonitoredTrainingSession takes care of session initialization,\n",
    "    # restoring from a checkpoint, saving to a checkpoint, and closing \n",
    "    # when done or an error occurs.\n",
    "    \n",
    "    # Horovod: save checkpoints only on worker 0 \n",
    "    checkpoint_dir = './checkpoints' if hvd.rank() == 0 else None\n",
    "    \n",
    "    # Create hooks\n",
    "    hooks = [\n",
    "        hvd.BroadcastGlobalVariablesHook(0),\n",
    "        tf.train.StopAtStepHook(last_step=101),\n",
    "        tf.train.LoggingTensorHook(tensors={'step': global_step, 'loss': loss},\n",
    "                                   every_n_iter=100),\n",
    "    ]\n",
    "    \n",
    "    # Horovod: pin GPU to be used to process local rank (one GPU per process)\n",
    "    config = tf.ConfigProto()\n",
    "    config.gpu_options.allow_growth = True\n",
    "    config.gpu_options.visible_device_list = str(hvd.local_rank())\n",
    "    \n",
    "    with tf.train.MonitoredTrainingSession(checkpoint_dir=checkpoint_dir,\n",
    "                                           hooks=hooks,\n",
    "                                           config=config) as mon_sess:\n",
    "        while not mon_sess.should_stop():\n",
    "            # Run a training step synchronously.\n",
    "            image_, label_ = mnist.train.next_batch(100)\n",
    "            mon_sess.run(train_op, feed_dict={image: image_, label: label_})\n",
    "            \n",
    "    log_contents = log_string.getvalue()\n",
    "    log_string.close()\n",
    "    \n",
    "    yield log_contents\n",
    " \n",
    "log_contents = rdd.mapPartitionsWithIndex(train).collect()"
   ]
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "kernelspec": {
   "display_name": "Python [default]",
   "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.5.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}