{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!--- Licensed to the Apache Software Foundation (ASF) under one -->\n",
    "<!--- or more contributor license agreements.  See the NOTICE file -->\n",
    "<!--- distributed with this work for additional information -->\n",
    "<!--- regarding copyright ownership.  The ASF licenses this file -->\n",
    "<!--- to you under the Apache License, Version 2.0 (the -->\n",
    "<!--- \"License\"); you may not use this file except in compliance -->\n",
    "<!--- with the License.  You may obtain a copy of the License at -->\n",
    "\n",
    "<!---   http://www.apache.org/licenses/LICENSE-2.0 -->\n",
    "\n",
    "<!--- Unless required by applicable law or agreed to in writing, -->\n",
    "<!--- software distributed under the License is distributed on an -->\n",
    "<!--- \"AS IS\" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -->\n",
    "<!--- KIND, either express or implied.  See the License for the -->\n",
    "<!--- specific language governing permissions and limitations -->\n",
    "<!--- under the License. -->\n",
    "\n",
    "# Distributed Key-Value Store\n",
    "\n",
    "KVStore is a place for data sharing. Think of it as a single object shared\n",
    "across different devices (GPUs and computers), where each device can push data in\n",
    "and pull data out.\n",
    "\n",
    "## Initialization\n",
    "\n",
    "Let's consider a simple example: initializing\n",
    "a (`int`, `NDArray`) pair into the store, and then pulling the value out:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```python\n",
    "import mxnet as mx\n",
    "\n",
    "kv = mx.kv.create('local') # create a local kv store.\n",
    "shape = (2,3)\n",
    "kv.init(3, mx.nd.ones(shape)*2)\n",
    "a = mx.nd.zeros(shape)\n",
    "kv.pull(3, out = a)\n",
    "print(a.asnumpy())\n",
    "```\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`[[ 2.  2.  2.],[ 2.  2.  2.]]`<!--notebook-skip-line-->\n",
    "\n",
    "## Push, Aggregate, and Update\n",
    "\n",
    "For any key that has been initialized, you can push a new value with the same shape to the key:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```python\n",
    "kv.push(3, mx.nd.ones(shape)*8)\n",
    "kv.pull(3, out = a) # pull out the value\n",
    "print(a.asnumpy())\n",
    "```\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`[[ 8.  8.  8.],[ 8.  8.  8.]]`<!--notebook-skip-line-->\n",
    "\n",
    "The data for pushing can be stored on any device. Furthermore, you can push multiple\n",
    "values into the same key, where KVStore will first sum all of these\n",
    "values and then push the aggregated value. Here we will just demonstrate pushing a list of values on CPU.\n",
    "Please note summation only happens if the value list is longer than one"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```python\n",
    "contexts = [mx.cpu(i) for i in range(4)]\n",
    "b = [mx.nd.ones(shape, ctx) for ctx in contexts]\n",
    "kv.push(3, b)\n",
    "kv.pull(3, out = a)\n",
    "print(a.asnumpy())\n",
    "```\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`[[ 4.  4.  4.],[ 4.  4.  4.]]`<!--notebook-skip-line-->\n",
    "\n",
    "For each push, KVStore combines the pushed value with the value stored using an\n",
    "`updater`. The default updater is `ASSIGN`. You can replace the default to\n",
    "control how data is merged:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```python\n",
    "def update(key, input, stored):\n",
    "    print(\"update on key: %d\" % key)\n",
    "    stored += input * 2\n",
    "kv._set_updater(update)\n",
    "kv.pull(3, out=a)\n",
    "print(a.asnumpy())\n",
    "```\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`[[ 4.  4.  4.],[ 4.  4.  4.]]`<!--notebook-skip-line-->"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```python\n",
    "kv.push(3, mx.nd.ones(shape))\n",
    "kv.pull(3, out=a)\n",
    "print(a.asnumpy())\n",
    "```\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`update on key: 3`<!--notebook-skip-line-->\n",
    "\n",
    "`[[ 6.  6.  6.],[ 6.  6.  6.]]`<!--notebook-skip-line-->\n",
    "\n",
    "\n",
    "## Pull\n",
    "\n",
    "You've already seen how to pull a single key-value pair. Similarly, to push, you can\n",
    "pull the value onto several devices with a single call:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```python\n",
    "b = [mx.nd.ones(shape, ctx) for ctx in contexts]\n",
    "kv.pull(3, out = b)\n",
    "print(b[1].asnumpy())\n",
    "```\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`[ 6.  6.  6.]],[[ 6.  6.  6.]`<!--notebook-skip-line-->\n",
    "\n",
    "## Handle a List of Key-Value Pairs\n",
    "\n",
    "All operations introduced so far involve a single key. KVStore also provides\n",
    "an interface for a list of key-value pairs.\n",
    "\n",
    "For a single device:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```python\n",
    "keys = [5, 7, 9]\n",
    "kv.init(keys, [mx.nd.ones(shape)]*len(keys))\n",
    "kv.push(keys, [mx.nd.ones(shape)]*len(keys))\n",
    "b = [mx.nd.zeros(shape)]*len(keys)\n",
    "kv.pull(keys, out = b)\n",
    "print(b[1].asnumpy())\n",
    "```\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`update on key: 5`<!--notebook-skip-line-->\n",
    "\n",
    "`update on key: 7`<!--notebook-skip-line-->\n",
    "\n",
    "`update on key: 9`<!--notebook-skip-line-->\n",
    "\n",
    "`[[ 3.  3.  3.],[ 3.  3.  3.]]`<!--notebook-skip-line-->\n",
    "\n",
    "For multiple devices:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```python\n",
    "b = [[mx.nd.ones(shape, ctx) for ctx in contexts]] * len(keys)\n",
    "kv.push(keys, b)\n",
    "kv.pull(keys, out = b)\n",
    "print(b[1][1].asnumpy())\n",
    "```\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`update on key: 5`<!--notebook-skip-line-->\n",
    "\n",
    "`update on key: 7`<!--notebook-skip-line-->\n",
    "\n",
    "`update on key: 9`<!--notebook-skip-line-->\n",
    "\n",
    "`[[ 11.  11.  11.],[ 11.  11.  11.]]`<!--notebook-skip-line-->\n",
    "\n",
    "## Run on Multiple Machines\n",
    "Based on parameter server, the `updater` runs on the server nodes.\n",
    "When the distributed version is ready, we will update this section.\n",
    "\n",
    "\n",
    "<!-- ## How to Choose Between APIs -->\n",
    "\n",
    "<!-- You can mix APIs as much as you like. Here are some guidelines -->\n",
    "<!-- * Use the Symbolic API and a coarse-grained operator to create  an established structure. -->\n",
    "<!-- * Use a fine-grained operator to extend parts of a more flexible symbolic graph. -->\n",
    "<!-- * Do some dynamic NDArray tricks, which are even more flexible, between the calls of forward and backward executors. -->\n",
    "\n",
    "<!-- Different approaches offer you different levels of flexibility and -->\n",
    "<!-- efficiency. Normally, you do not need to be flexible in all parts of the -->\n",
    "<!-- network, so use the parts optimized for speed, and compose it -->\n",
    "<!-- flexibly with a fine-grained operator or a dynamic NDArray. Such a -->\n",
    "<!-- mixture allows you to build the deep learning architecture both efficiently and -->\n",
    "<!-- flexibly as your choice.  -->\n",
    "\n",
    "## Next Steps\n",
    "* [MXNet tutorials index](/api/python/docs/tutorials/)\n",
    "\n",
    "<!-- INSERT SOURCE DOWNLOAD BUTTONS -->"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}