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

{
 "metadata": {
  "name": "parallel_mpi"
 },
 "nbformat": 3,
 "nbformat_minor": 0,
 "worksheets": [
  {
   "cells": [
    {
     "cell_type": "heading",
     "level": 1,
     "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",
     "collapsed": true,
     "input": [
      "from IPython.parallel import Client\n",
      "c = Client()\n",
      "view = c[:]"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 1
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "Let's define a simple function that gets the MPI rank from each engine."
     ]
    },
    {
     "cell_type": "code",
     "collapsed": true,
     "input": [
      "@view.remote(block=True)\n",
      "def mpi_rank():\n",
      "    from mpi4py import MPI\n",
      "    comm = MPI.COMM_WORLD\n",
      "    return comm.Get_rank()"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 2
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "mpi_rank()"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [
      {
       "output_type": "pyout",
       "prompt_number": 3,
       "text": [
        "[2, 3, 1, 0]"
       ]
      }
     ],
     "prompt_number": 3
    },
    {
     "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",
     "collapsed": false,
     "input": [
      "mpi_rank.block = False\n",
      "ar = mpi_rank()\n",
      "ar.get_dict()"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [
      {
       "output_type": "pyout",
       "prompt_number": 4,
       "text": [
        "{0: 2, 1: 3, 2: 1, 3: 0}"
       ]
      }
     ],
     "prompt_number": 4
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "With %%px cell magic, the next cell will actually execute *entirely on each engine*:"
     ]
    },
    {
     "cell_type": "code",
     "collapsed": true,
     "input": [
      "%%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)"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 5
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "view['data']"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [
      {
       "output_type": "pyout",
       "prompt_number": 6,
       "text": [
        "[9, 16, 4, 1]"
       ]
      }
     ],
     "prompt_number": 6
    },
    {
     "cell_type": "code",
     "collapsed": true,
     "input": [],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 6
    }
   ],
   "metadata": {}
  }
 ]
}