examples/Parallel Computing/Using MPI with IPython Parallel.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Simple usage of a set of MPI engines"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This example assumes you've started a cluster of N engines (4 in this example) as part\n",
    "of an MPI world.  \n",
    "\n",
    "Our documentation describes [how to create an MPI profile](http://ipython.org/ipython-doc/dev/parallel/parallel_process.html#using-ipcluster-in-mpiexec-mpirun-mode)\n",
    "and explains [basic MPI usage of the IPython cluster](http://ipython.org/ipython-doc/dev/parallel/parallel_mpi.html).\n",
    "\n",
    "\n",
    "For the simplest possible way to start 4 engines that belong to the same MPI world, \n",
    "you can run this in a terminal:\n",
    "\n",
    "<pre>\n",
    "ipcluster start --engines=MPI -n 4\n",
    "</pre>\n",
    "\n",
    "or start an MPI cluster from the cluster tab if you have one configured.\n",
    "\n",
    "Once the cluster is running, we can connect to it and open a view into it:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from IPython.parallel import Client\n",
    "c = Client()\n",
    "view = c[:]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's define a simple function that gets the MPI rank from each engine."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "@view.remote(block=True)\n",
    "def mpi_rank():\n",
    "    from mpi4py import MPI\n",
    "    comm = MPI.COMM_WORLD\n",
    "    return comm.Get_rank()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[2, 3, 1, 0]"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mpi_rank()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To get a mapping of IPython IDs and MPI rank (these do not always match),\n",
    "you can use the get_dict method on AsyncResults."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{0: 2, 1: 3, 2: 1, 3: 0}"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mpi_rank.block = False\n",
    "ar = mpi_rank()\n",
    "ar.get_dict()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With %%px cell magic, the next cell will actually execute *entirely on each engine*:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "%%px\n",
    "from mpi4py import MPI\n",
    "\n",
    "comm = MPI.COMM_WORLD\n",
    "size = comm.Get_size()\n",
    "rank = comm.Get_rank()\n",
    "\n",
    "if rank == 0:\n",
    "   data = [(i+1)**2 for i in range(size)]\n",
    "else:\n",
    "   data = None\n",
    "data = comm.scatter(data, root=0)\n",
    "\n",
    "assert data == (rank+1)**2, 'data=%s, rank=%s' % (data, rank)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[9, 16, 4, 1]"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "view['data']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {},
 "nbformat": 4,
 "nbformat_minor": 0
}