Distributed TensorFlow - A Gentle Introduction.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Distributed TensorFlow: A Gentle Introduction"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Distributed TensorFlow allows us to **share parts of a TensorFlow graph between multiple processes**, possibly each on a different machine.\n",
    "\n",
    "<img src=\"images/distributed_tensorflow.png\" style=\"max-width: 300px\" />\n",
    "\n",
    "Why might we want to do this? The classic use case is to harness the power of multiple machines for training, with shared parameters between all machines. Even if we're just running on a single machine, though, I've come across two examples from reinforcement learning where sharing between processes has been necessary:\n",
    "* In [A3C](https://arxiv.org/pdf/1602.01783.pdf), multiple agents run in parallel in multiple processes, exploring different copies of the environment at the same time. Each agent generates parameter updates, which must also be sent to the other agents in different processes.\n",
    "* In [Deep Reinforcement Learning from Human Preferences](https://arxiv.org/pdf/1706.03741.pdf), one process runs an agent exploring the environment, with rewards calculated by a network trained from human preferences about agent behaviour. This network is trained asynchronously from the agent, in a separate process, so that the agent doesn't have to wait for each training cycle to complete before it can continue exploring.\n",
    "\n",
    "Unfortunately, the [official documentation](https://www.tensorflow.org/deploy/distributed) on Distributed TensorFlow rather jumps in at the deep end. **For a slightly more gentle introduction - let's run through some really basic examples.**\n",
    "\n",
    "(Note: some of the explanations in this notebook are my own interpretation of empirical results/TensorFlow documentation. If you see anything that's wrong, give me a shout!)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Running this Notebook"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you're running this notebook interactively and get a message saying\n",
    "\n",
    "```UnknownError: Could not start gRPC server```\n",
    "\n",
    "when running certain cells, it's probably because something is already listening on the specified port.\n",
    "\n",
    "If you've run cells more than once, this could be a previous instance of a TensorFlow server. In this case, try restarting the notebook kernel."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Introduction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import tensorflow as tf"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's say we want multiple processes to have shared access to some common parameters. For simplicity, suppose this is just a single variable:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "var = tf.Variable(initial_value=0.0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As a first step, we can imagine that each process would need its own session. (Pretend session 1 is created in one process, and session 2 in another.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "sess1 = tf.Session()\n",
    "sess2 = tf.Session()\n",
    "\n",
    "sess1.run(tf.global_variables_initializer())\n",
    "sess2.run(tf.global_variables_initializer())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Each call to `tf.Session()` creates a separate **execution engine**, then connects the session handle to the execution engine. The execution engine is what actually stores variable values and runs operations.\n",
    "\n",
    "Normally, execution engines in different processes are unlinked. Changing `var` in one session (on one execution engine) won't affect `var` in the other session."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Initial value of var in session 1: 0.0\n",
      "Initial value of var in session 2: 0.0\n",
      "Incremented var in session 1\n",
      "Value of var in session 1: 1.0\n",
      "Value of var in session 2: 0.0\n"
     ]
    }
   ],
   "source": [
    "print(\"Initial value of var in session 1:\", sess1.run(var))\n",
    "print(\"Initial value of var in session 2:\", sess2.run(var))\n",
    "\n",
    "sess1.run(var.assign_add(1.0))\n",
    "print(\"Incremented var in session 1\")\n",
    "\n",
    "print(\"Value of var in session 1:\", sess1.run(var))\n",
    "print(\"Value of var in session 2:\", sess2.run(var))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Distributed TensorFlow"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In order to share variables between processes, we need to link the different execution engines together. Enter Distributed TensorFlow.\n",
    "\n",
    "With Distributed TensorFlow, each process runs a special execution engine, a TensorFlow **server**. Servers are linked together as part of a **cluster**. (Each server in the cluster is also known as a **task**.)\n",
    "\n",
    "The first step is to define what the cluster looks like. We start off with the simplest possible cluster: two servers (two tasks), both on the same machine; one that will listen on port 2222, one on port 2223."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "tasks = [\"localhost:2222\", \"localhost:2223\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Each task is associated with a **job**, which is a collection of related tasks. We associate both tasks with a job called \"local\"."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "jobs = {\"local\": tasks}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This completes the definition of the cluster."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "cluster = tf.train.ClusterSpec(jobs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now launch the servers, specifying which server in the cluster definition each server corresponds to. Each server starts immediately, listening on the port specified in the cluster definition."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# \"This server corresponds to the the first task (task_index=0)\n",
    "# of the tasks associated with the 'local' job.\"\n",
    "server1 = tf.train.Server(cluster, job_name=\"local\", task_index=0)\n",
    "\n",
    "server2 = tf.train.Server(cluster, job_name=\"local\", task_index=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**With the servers linked together in the same cluster**, we can now experience the main magic of Distributed TensorFlow: **any variable with the same name will be shared between all servers**.\n",
    "\n",
    "The simplest example is to run the same graph on all servers, each graph with just one variable, as before:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "tf.reset_default_graph()\n",
    "var = tf.Variable(initial_value=0.0, name='var')\n",
    "sess1 = tf.Session(server1.target)\n",
    "sess2 = tf.Session(server2.target)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Modifications made to the variable on one server will now be mirrored on the second server."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Initial value of var in session 1: 0.0\n",
      "Initial value of var in session 2: 0.0\n",
      "Incremented var in session 1\n",
      "Value of var in session 1: 1.0\n",
      "Value of var in session 2: 1.0\n"
     ]
    }
   ],
   "source": [
    "sess1.run(tf.global_variables_initializer())\n",
    "sess2.run(tf.global_variables_initializer())\n",
    "\n",
    "print(\"Initial value of var in session 1:\", sess1.run(var))\n",
    "print(\"Initial value of var in session 2:\", sess2.run(var))\n",
    "\n",
    "sess1.run(var.assign_add(1.0))\n",
    "print(\"Incremented var in session 1\")\n",
    "\n",
    "print(\"Value of var in session 1:\", sess1.run(var))\n",
    "print(\"Value of var in session 2:\", sess2.run(var))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "(Note that because we only have one variable, and that variable is shared between both sessions, the second run of `global_variables_initializer` here is redundant.)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Placement"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A question that might be in our minds at this point is: which server does the variable actually get stored on? And for operations, which server actually runs them?\n",
    "\n",
    "Empirically, it seems that by default, variables and operations get placed on the first task in the cluster."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_with_location_trace(sess, op):\n",
    "    # From https://stackoverflow.com/a/41525764/7832197\n",
    "    run_options = tf.RunOptions(trace_level=tf.RunOptions.FULL_TRACE)\n",
    "    run_metadata = tf.RunMetadata()\n",
    "    sess.run(op, options=run_options, run_metadata=run_metadata)\n",
    "    for device in run_metadata.step_stats.dev_stats:\n",
    "      print(device.device)\n",
    "      for node in device.node_stats:\n",
    "        print(\"  \", node.node_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For example, if we do something to `var` using the session connected to the first task, everything happens on that task:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/job:local/replica:0/task:0/device:CPU:0\n",
      "   _SOURCE\n",
      "   var\n"
     ]
    }
   ],
   "source": [
    "run_with_location_trace(sess1, var)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/job:local/replica:0/task:0/device:CPU:0\n",
      "   _SOURCE\n",
      "   AssignAdd_1/value\n",
      "   var\n",
      "   AssignAdd_1\n"
     ]
    }
   ],
   "source": [
    "run_with_location_trace(sess1, var.assign_add(1.0))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "But if we try and try do something to `var` using the session connected to the second task, the graph nodes still get run on the first task."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/job:local/replica:0/task:1/device:CPU:0\n",
      "   _SOURCE\n",
      "/job:local/replica:0/task:0/device:CPU:0\n",
      "   _SOURCE\n",
      "   var\n"
     ]
    }
   ],
   "source": [
    "run_with_location_trace(sess2, var)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**To fix a variable or an operation to a specific task, we can use `tf.device`:**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "with tf.device(\"/job:local/task:0\"):\n",
    "    var1 = tf.Variable(0.0, name='var1')\n",
    "with tf.device(\"/job:local/task:1\"):\n",
    "    var2 = tf.Variable(0.0, name='var2')\n",
    "    \n",
    "# (This will initialize both variables)\n",
    "sess1.run(tf.global_variables_initializer())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now `var1` runs on the first task, as before."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/job:local/replica:0/task:0/device:CPU:0\n",
      "   _SOURCE\n",
      "   var1\n"
     ]
    }
   ],
   "source": [
    "run_with_location_trace(sess1, var1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "But `var2` runs on the second task. Even if we try to evaluate it using the session connected to the first task, it still runs on the second task."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/job:local/replica:0/task:0/device:CPU:0\n",
      "   _SOURCE\n",
      "/job:local/replica:0/task:1/device:CPU:0\n",
      "   _SOURCE\n",
      "   var2\n"
     ]
    }
   ],
   "source": [
    "run_with_location_trace(sess1, var2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And vice-versa with `var2`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/job:local/replica:0/task:1/device:CPU:0\n",
      "   _SOURCE\n",
      "   var2\n"
     ]
    }
   ],
   "source": [
    "run_with_location_trace(sess2, var2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/job:local/replica:0/task:1/device:CPU:0\n",
      "   _SOURCE\n",
      "/job:local/replica:0/task:0/device:CPU:0\n",
      "   _SOURCE\n",
      "   var1\n"
     ]
    }
   ],
   "source": [
    "run_with_location_trace(sess2, var1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Graphs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "There are a couple of things to note about how graphs work with Distributed TensorFlow."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Who builds the graph?\n",
    "\n",
    "First, although variable _values_ are shared throughout the cluster, the _graph_ is _not_ automatically shared.\n",
    "\n",
    "Let's create a fresh cluster with two servers, and set up the first server with an explicitly-created graph."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "cluster = tf.train.ClusterSpec({\"local\": [\"localhost:2224\", \"localhost:2225\"]})\n",
    "server1 = tf.train.Server(cluster, job_name=\"local\", task_index=0)\n",
    "server2 = tf.train.Server(cluster, job_name=\"local\", task_index=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[<tf.Operation 'var/initial_value' type=Const>, <tf.Operation 'var' type=VariableV2>, <tf.Operation 'var/Assign' type=Assign>, <tf.Operation 'var/read' type=Identity>]\n"
     ]
    }
   ],
   "source": [
    "graph1 = tf.Graph()\n",
    "with graph1.as_default():\n",
    "    var1 = tf.Variable(0.0, name='var')\n",
    "sess1 = tf.Session(target=server1.target, graph=graph1)\n",
    "print(graph1.get_operations())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If we then create a session connected to the second server, note that the graph does not automatically get mirrored."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[]\n"
     ]
    }
   ],
   "source": [
    "graph2 = tf.Graph()\n",
    "sess2 = tf.Session(target=server2.target, graph=graph2)\n",
    "print(graph2.get_operations())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To access the shared variable, we must manually add a variable with the same name to the second graph."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "with graph2.as_default():\n",
    "    var2 = tf.Variable(0.0, name='var')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Only then can we access it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1.0"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sess1.run(var1.assign(1.0))\n",
    "sess2.run(var2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The takeaway is: **each server is responsible for building its own graph.**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Does the graph have to be the same on all servers?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So far, all our examples have run the same graph structure on both servers. This is known as **in-graph replication**.\n",
    "\n",
    "For example, let's say we have a cluster containing three servers. Server 1 holds shared parameters, while server 2 and server 3 are worker nodes, each with local variables. With in-graph replication, each server's graphs would look like:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<img src=\"images/in_graph_replication.png\" style=\"max-width: 600px\" />"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The issue with in-graph replication is that every server has to have a copy of the entire graph, including the parts of the graph that might only be relevant for other servers. This can lead to graphs growing very large.\n",
    "\n",
    "The alternative is **between-graph replication**. Here, each server runs a graph containing only the shared parameters, and whatever variables and operations are relevant to that individual server."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<img src=\"images/between_graph_replication.png\" style=\"max-width: 600px\" />"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Because it keeps graph sizes smaller, **between-graph replication is the recommended approach.**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Practical Details"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Before we finish with a full example, there are a few practical details to discuss.\n",
    "\n",
    "(If you're following along interactively, please restart the notebook kernel at this point. TensorFlow servers started in the main process interfere with servers started in forked processes.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import tensorflow as tf"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## What happens if we try to run something on the cluster before all servers have connected?\n",
    "\n",
    "Let's create another two-task cluster."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "cluster = tf.train.ClusterSpec({\n",
    "    \"local\": [\"localhost:2226\", \"localhost:2227\"]\n",
    "})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This time, let's start each server in a separate process. (This allows us to kill the servers, so that we can start them again for later experiments. There's currently no way of killing servers other than killing the process which started them.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from multiprocessing import Process\n",
    "from time import sleep\n",
    "\n",
    "def s1():\n",
    "    server1 = tf.train.Server(cluster,\n",
    "                              job_name=\"local\",\n",
    "                              task_index=0)\n",
    "    sess1 = tf.Session(server1.target)\n",
    "    print(\"server 1: running no-op...\")\n",
    "    sess1.run(tf.no_op())\n",
    "    print(\"server 1: no-op run!\")\n",
    "    server1.join() # Block\n",
    "\n",
    "def s2():\n",
    "    for i in range(3):\n",
    "        print(\"server 2: %d seconds left before connecting...\"\n",
    "              % (3 - i))\n",
    "        sleep(1.0)\n",
    "    server2 = tf.train.Server(cluster,\n",
    "                              job_name=\"local\",\n",
    "                              task_index=1)\n",
    "    print(\"server 2: connected!\")\n",
    "    server2.join() # Block\n",
    "\n",
    "# daemon=True so that these processes will definitely be killed\n",
    "# when the notebook restarts\n",
    "p1 = Process(target=s1, daemon=True)\n",
    "p2 = Process(target=s2, daemon=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Server 1 joins the cluster immediately, but server 2 waits a little while before connecting. The result is shown below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "server 2: 3 seconds left before connecting...\n",
      "server 1: running no-op...\n",
      "server 2: 2 seconds left before connecting...\n",
      "server 2: 1 seconds left before connecting...\n",
      "server 2: connected!\n",
      "server 1: no-op run!\n"
     ]
    }
   ],
   "source": [
    "p1.start()\n",
    "p2.start()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As can be seen, **any attempt to run an operation on the cluster blocks until all servers have joined.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "p1.terminate()\n",
    "p2.terminate()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## What happens if a server leaves the cluster?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's set up a cluster with two servers. Server 1 will just repeatedly try and run a no-op located on server 1. Server 2 will die after two seconds."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Server 1: about to run no-op...success!\n",
      "Server 1: about to run no-op...success!\n",
      "Server 2 dieing...\n",
      "Server 1: about to run no-op...success!\n",
      "Server 1: about to run no-op...success!\n",
      "Server 1: about to run no-op...success!\n",
      "Server 1: about to run no-op...success!\n"
     ]
    }
   ],
   "source": [
    "def s1():\n",
    "    server1 = tf.train.Server(cluster,\n",
    "                              job_name=\"local\",\n",
    "                              task_index=0)\n",
    "    \n",
    "    with tf.device(\"/job:local/task:0\"):\n",
    "        no_op = tf.no_op()\n",
    "        \n",
    "    sess1 = tf.Session(server1.target)\n",
    "    for _ in range(6):\n",
    "        print(\"Server 1: about to run no-op...\", end=\"\")\n",
    "        sess1.run(no_op)\n",
    "        print(\"success!\")\n",
    "        sleep(1.0)\n",
    "\n",
    "def s2():\n",
    "    server2 = tf.train.Server(cluster,\n",
    "                              job_name=\"local\",\n",
    "                              task_index=1)\n",
    "    sleep(2.0)\n",
    "    print(\"Server 2 dieing...\")\n",
    "    \n",
    "p1 = Process(target=s1, daemon=True)\n",
    "p2 = Process(target=s2, daemon=True)\n",
    "\n",
    "p1.start()\n",
    "p2.start()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the short term, it seems there's no problem, as long as the operation we're trying to run isn't on the server that's left. (I haven't tested what happens long-term.) (Update: see comment from Yaroslav Bulatov about this at <http://amid.fish/distributed-tensorflow-a-gentle-introduction>.)\n",
    "\n",
    "If the operation _is_ on the server that leaves..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Server 1: about to run no-op...success!\n",
      "Server 1: about to run no-op...success!\n",
      "Server 2 dieing...\n"
     ]
    }
   ],
   "source": [
    "def s1():\n",
    "    server1 = tf.train.Server(cluster,\n",
    "                              job_name=\"local\",\n",
    "                              task_index=0)\n",
    "    \n",
    "    # This time, we place the no-op on server 2,\n",
    "    # which is going to leave\n",
    "    with tf.device(\"/job:local/task:1\"):\n",
    "        no_op = tf.no_op()\n",
    "        \n",
    "    sess1 = tf.Session(server1.target)\n",
    "    for _ in range(5):\n",
    "        print(\"Server 1: about to run no-op...\", end=\"\")\n",
    "        sess1.run(no_op)\n",
    "        print(\"success!\")\n",
    "        sleep(1.0)\n",
    "    \n",
    "p1 = Process(target=s1, daemon=True)\n",
    "p2 = Process(target=s2, daemon=True)\n",
    "\n",
    "p1.start()\n",
    "p2.start()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "...then the attempt to run the operation blocks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "p1.terminate()\n",
    "p2.terminate()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "What happens if the server then comes back?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Server 1: about to run no-op...success!\n",
      "Server 1: about to run no-op...success!\n",
      "Server 2 dieing...\n",
      "Restarting server 2...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Process Process-7:\n",
      "Traceback (most recent call last):\n",
      "  File \"/Users/matthew/tensorflow/lib/python3.6/site-packages/tensorflow/python/client/session.py\", line 1323, in _do_call\n",
      "    return fn(*args)\n",
      "  File \"/Users/matthew/tensorflow/lib/python3.6/site-packages/tensorflow/python/client/session.py\", line 1302, in _run_fn\n",
      "    status, run_metadata)\n",
      "  File \"/Users/matthew/tensorflow/lib/python3.6/site-packages/tensorflow/python/framework/errors_impl.py\", line 473, in __exit__\n",
      "    c_api.TF_GetCode(self.status.status))\n",
      "tensorflow.python.framework.errors_impl.AbortedError: Graph handle is not found: 0000000000000001\n",
      "\n",
      "During handling of the above exception, another exception occurred:\n",
      "\n",
      "Traceback (most recent call last):\n",
      "  File \"/usr/local/Cellar/python3/3.6.3/Frameworks/Python.framework/Versions/3.6/lib/python3.6/multiprocessing/process.py\", line 258, in _bootstrap\n",
      "    self.run()\n",
      "  File \"/usr/local/Cellar/python3/3.6.3/Frameworks/Python.framework/Versions/3.6/lib/python3.6/multiprocessing/process.py\", line 93, in run\n",
      "    self._target(*self._args, **self._kwargs)\n",
      "  File \"<ipython-input-7-b02182aeb5bc>\", line 14, in s1\n",
      "    sess1.run(no_op)\n",
      "  File \"/Users/matthew/tensorflow/lib/python3.6/site-packages/tensorflow/python/client/session.py\", line 889, in run\n",
      "    run_metadata_ptr)\n",
      "  File \"/Users/matthew/tensorflow/lib/python3.6/site-packages/tensorflow/python/client/session.py\", line 1120, in _run\n",
      "    feed_dict_tensor, options, run_metadata)\n",
      "  File \"/Users/matthew/tensorflow/lib/python3.6/site-packages/tensorflow/python/client/session.py\", line 1317, in _do_run\n",
      "    options, run_metadata)\n",
      "  File \"/Users/matthew/tensorflow/lib/python3.6/site-packages/tensorflow/python/client/session.py\", line 1336, in _do_call\n",
      "    raise type(e)(node_def, op, message)\n",
      "tensorflow.python.framework.errors_impl.AbortedError: Graph handle is not found: 0000000000000001\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Server 1: about to run no-op...Server 2 dieing...\n"
     ]
    }
   ],
   "source": [
    "p1 = Process(target=s1, daemon=True)\n",
    "p2 = Process(target=s2, daemon=True)\n",
    "p1.start()\n",
    "p2.start()\n",
    "sleep(3.0)\n",
    "# At this point, server 1 is blocked, and server 2 is dead.\n",
    "print(\"Restarting server 2...\")\n",
    "p2 = Process(target=s2, daemon=True)\n",
    "p2.start()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We get a `Graph handle is not found` error.\n",
    "\n",
    "So the takeaway is that **Distributed TensorFlow isn't automatically resilient to server failures**. (If you _are_ interested in fault tolerance, check out the [TensorFlow Dev Summit video on Distributed TensorFlow](https://www.youtube.com/watch?v=la_M6bCV91M).) (Update: also, see comments from Yaroslav Bulatov at <http://amid.fish/distributed-tensorflow-a-gentle-introduction>.)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Who's responsible for initializing shared variables?\n",
    "\n",
    "One approach is just to have all workers run `tf.global_variables_initializer()`.\n",
    "\n",
    "But if we want to keep it clean and have only one server do initialization, we could run into a problem if other servers try to use the variables before they've been initialized. One option for avoiding this problem is to **have the other workers wait until initialization has taken place using `tf.report_uninitialized_variables`.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Server 1: waiting for connection...\n",
      "Server 2: waiting 3 seconds before initializing...\n",
      "Server 1: waiting for initialization...\n",
      "Server 2: waiting 2 seconds before initializing...\n",
      "Server 1: waiting for initialization...\n",
      "Server 2: waiting 1 seconds before initializing...\n",
      "Server 1: waiting for initialization...\n",
      "Server 1: variables initialized!\n"
     ]
    }
   ],
   "source": [
    "def s1():\n",
    "    server1 = tf.train.Server(cluster,\n",
    "                              job_name=\"local\",\n",
    "                              task_index=0)\n",
    "    var = tf.Variable(0.0, name='var')\n",
    "    sess1 = tf.Session(server1.target)\n",
    "    \n",
    "    print(\"Server 1: waiting for connection...\")\n",
    "    sess1.run(tf.report_uninitialized_variables())\n",
    "    while len(sess1.run(tf.report_uninitialized_variables())) > 0:\n",
    "        print(\"Server 1: waiting for initialization...\")\n",
    "        sleep(1.0)\n",
    "    print(\"Server 1: variables initialized!\")\n",
    "\n",
    "def s2():\n",
    "    server2 = tf.train.Server(cluster,\n",
    "                              job_name=\"local\",\n",
    "                              task_index=1)\n",
    "    var = tf.Variable(0.0, name='var')\n",
    "    sess2 = tf.Session(server2.target)\n",
    "    \n",
    "    for i in range(3):\n",
    "        print(\"Server 2: waiting %d seconds before initializing...\"\n",
    "              % (3 - i))\n",
    "        sleep(1.0)\n",
    "    sess2.run(tf.global_variables_initializer())\n",
    "    \n",
    "p1 = Process(target=s1, daemon=True)\n",
    "p2 = Process(target=s2, daemon=True)\n",
    "p1.start()\n",
    "p2.start()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "p1.terminate()\n",
    "p2.terminate()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Example"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's put everything we've learned into one final example using multiple processes.\n",
    "\n",
    "We'll create:\n",
    "* One parameter server, which will store a single variable, `var`.\n",
    "* Two worker tasks. Each worker will increment `var` a few times.\n",
    "\n",
    "We'll have the parameter server print out the value of `var` a few times too so we can really see it changing."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "from multiprocessing import Process\n",
    "from time import sleep\n",
    "\n",
    "cluster = tf.train.ClusterSpec({\n",
    "    \"worker\": [\n",
    "        \"localhost:3333\",\n",
    "        \"localhost:3334\",\n",
    "    ],\n",
    "    \"ps\": [\n",
    "        \"localhost:3335\"\n",
    "    ]\n",
    "})\n",
    "\n",
    "def parameter_server():\n",
    "    with tf.device(\"/job:ps/task:0\"):\n",
    "        var = tf.Variable(0.0, name='var')\n",
    "\n",
    "    server = tf.train.Server(cluster,\n",
    "                             job_name=\"ps\",\n",
    "                             task_index=0)\n",
    "    sess = tf.Session(target=server.target)\n",
    "    \n",
    "    print(\"Parameter server: waiting for cluster connection...\")\n",
    "    sess.run(tf.report_uninitialized_variables())\n",
    "    print(\"Parameter server: cluster ready!\")\n",
    "    \n",
    "    print(\"Parameter server: initializing variables...\")\n",
    "    sess.run(tf.global_variables_initializer())\n",
    "    print(\"Parameter server: variables initialized\")\n",
    "    \n",
    "    for i in range(5):\n",
    "        val = sess.run(var)\n",
    "        print(\"Parameter server: var has value %.1f\" % val)\n",
    "        sleep(1.0)\n",
    "\n",
    "    print(\"Parameter server: blocking...\")\n",
    "    server.join()\n",
    "    \n",
    "\n",
    "def worker(worker_n):\n",
    "    with tf.device(\"/job:ps/task:0\"):\n",
    "        var = tf.Variable(0.0, name='var')\n",
    "        \n",
    "    server = tf.train.Server(cluster,\n",
    "                             job_name=\"worker\",\n",
    "                             task_index=worker_n)\n",
    "    sess = tf.Session(target=server.target)\n",
    "    \n",
    "    print(\"Worker %d: waiting for cluster connection...\" % worker_n)\n",
    "    sess.run(tf.report_uninitialized_variables())\n",
    "    print(\"Worker %d: cluster ready!\" % worker_n)\n",
    "    \n",
    "    while sess.run(tf.report_uninitialized_variables()):\n",
    "        print(\"Worker %d: waiting for variable initialization...\" % worker_n)\n",
    "        sleep(1.0)\n",
    "    print(\"Worker %d: variables initialized\" % worker_n)\n",
    "    \n",
    "    for i in range(5):\n",
    "        print(\"Worker %d: incrementing var\" % worker_n)\n",
    "        sess.run(var.assign_add(1.0))\n",
    "        sleep(1.0)\n",
    "    \n",
    "    print(\"Worker %d: blocking...\" % worker_n)\n",
    "    server.join()\n",
    "\n",
    "ps_proc = Process(target=parameter_server, daemon=True)\n",
    "w1_proc = Process(target=worker, args=(0, ), daemon=True)\n",
    "w2_proc = Process(target=worker, args=(1, ), daemon=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Parameter server: waiting for cluster connection...\n",
      "Parameter server: cluster ready!\n",
      "Parameter server: initializing variables...\n",
      "Parameter server: variables initialized\n",
      "Parameter server: var has value 0.0\n",
      "Parameter server: var has value 2.0\n",
      "Parameter server: var has value 4.0\n",
      "Parameter server: var has value 5.0\n",
      "Parameter server: var has value 7.0\n",
      "Parameter server: blocking...\n"
     ]
    }
   ],
   "source": [
    "ps_proc.start()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Worker 0: waiting for cluster connection...\n",
      "Worker 0: cluster ready!\n",
      "Worker 0: waiting for variable initialization...\n",
      "Worker 0: variables initialized\n",
      "Worker 0: incrementing var\n",
      "Worker 0: incrementing var\n",
      "Worker 0: incrementing var\n",
      "Worker 0: incrementing var\n",
      "Worker 0: incrementing var\n",
      "Worker 0: blocking...\n"
     ]
    }
   ],
   "source": [
    "w1_proc.start()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Worker 1: waiting for cluster connection...\n",
      "Worker 1: cluster ready!\n",
      "Worker 1: waiting for variable initialization...\n",
      "Worker 1: variables initialized\n",
      "Worker 1: incrementing var\n",
      "Worker 1: incrementing var\n",
      "Worker 1: incrementing var\n",
      "Worker 1: incrementing var\n",
      "Worker 1: incrementing var\n",
      "Worker 1: blocking...\n"
     ]
    }
   ],
   "source": [
    "w2_proc.start()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "for proc in [w1_proc, w2_proc, ps_proc]:\n",
    "    proc.terminate()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Final words"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We've looked at:\n",
    "* How to join together multiple TensorFlow execution engines (running on different processes or different machines) into a cluster, so that they can share variables.\n",
    "* How to place variables or operations on a specific server.\n",
    "* In-graph and between-graph replication.\n",
    "* What happens if we try to run operations on the cluster before all servers have connected, or after a server has left.\n",
    "* How to wait until variables have been initialized by another task in the cluster.\n",
    "\n",
    "For more information and some more realistic examples, check out the official documentation at [Distributed TensorFlow](https://www.tensorflow.org/deploy/distributed)."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}