docs/examples/parallel/parallel_mpi.ipynb

{
 "metadata": {
  "name": "parallel_mpi"
 },
 "nbformat": 3,
 "worksheets": [
  {
   "cells": [
    {
     "cell_type": "markdown",
     "source": [
      "# Simple usage of a set of MPI engines",
      "",
      "This example assumes you've started a cluster of N engines (4 in this example) as part",
      "of an MPI world.  ",
      "",
      "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)",
      "and explains [basic MPI usage of the IPython cluster](http://ipython.org/ipython-doc/dev/parallel/parallel_mpi.html).",
      "",
      "",
      "For the simplest possible way to start 4 engines that belong to the same MPI world, ",
      "you can run this in a terminal or antoher notebook:",
      "",
      "<pre>",
      "ipcluster start --engines=MPI -n 4",
      "</pre>",
      "",
      "Note: to run the above in a notebook, use a *new* notebook and prepend the command with `!`, but do not run",
      "it in *this* notebook, as this command will block until you shut down the cluster.  To stop the cluster, use ",
      "the 'Interrupt' button on the left, which is the equivalent of sending `Ctrl-C` to the kernel.",
      "",
      "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",
      "c = Client()",
      "view = c[:]"
     ],
     "language": "python",
     "outputs": [],
     "prompt_number": 21
    },
    {
     "cell_type": "markdown",
     "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)",
      "def mpi_rank():",
      "    from mpi4py import MPI",
      "    comm = MPI.COMM_WORLD",
      "    return comm.Get_rank()"
     ],
     "language": "python",
     "outputs": [],
     "prompt_number": 22
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "mpi_rank()"
     ],
     "language": "python",
     "outputs": [
      {
       "output_type": "pyout",
       "prompt_number": 23,
       "text": [
        "[3, 0, 2, 1]"
       ]
      }
     ],
     "prompt_number": 23
    },
    {
     "cell_type": "markdown",
     "source": [
      "For interactive convenience, we load the parallel magic extensions and make this view",
      "the active one for the automatic parallelism magics.",
      "",
      "This is not necessary and in production codes likely won't be used, as the engines will ",
      "load their own MPI codes separately.  But it makes it easy to illustrate everything from",
      "within a single notebook here."
     ]
    },
    {
     "cell_type": "code",
     "collapsed": true,
     "input": [
      "%load_ext parallelmagic",
      "view.activate()"
     ],
     "language": "python",
     "outputs": [],
     "prompt_number": 4
    },
    {
     "cell_type": "markdown",
     "source": [
      "Use the autopx magic to make the rest of this cell execute on the engines instead",
      "of locally"
     ]
    },
    {
     "cell_type": "code",
     "collapsed": true,
     "input": [
      "view.block = True"
     ],
     "language": "python",
     "outputs": [],
     "prompt_number": 24
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "%autopx"
     ],
     "language": "python",
     "outputs": [
      {
       "output_type": "stream",
       "stream": "stdout",
       "text": [
        "%autopx enabled"
       ]
      }
     ],
     "prompt_number": 32
    },
    {
     "cell_type": "markdown",
     "source": [
      "With autopx enabled, the next cell will actually execute *entirely on each engine*:"
     ]
    },
    {
     "cell_type": "code",
     "collapsed": true,
     "input": [
      "from mpi4py import MPI",
      "",
      "comm = MPI.COMM_WORLD",
      "size = comm.Get_size()",
      "rank = comm.Get_rank()",
      "",
      "if rank == 0:",
      "   data = [(i+1)**2 for i in range(size)]",
      "else:",
      "   data = None",
      "data = comm.scatter(data, root=0)",
      "",
      "assert data == (rank+1)**2, 'data=%s, rank=%s' % (data, rank)"
     ],
     "language": "python",
     "outputs": [],
     "prompt_number": 29
    },
    {
     "cell_type": "markdown",
     "source": [
      "Though the assertion at the end of the previous block validated the code, we can now ",
      "pull the 'data' variable from all the nodes for local inspection.",
      "First, don't forget to toggle off `autopx` mode so code runs again in the notebook:"
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "%autopx"
     ],
     "language": "python",
     "outputs": [
      {
       "output_type": "stream",
       "stream": "stdout",
       "text": [
        "%autopx disabled"
       ]
      }
     ],
     "prompt_number": 33
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "view['data']"
     ],
     "language": "python",
     "outputs": [
      {
       "output_type": "pyout",
       "prompt_number": 34,
       "text": [
        "[16, 1, 9, 4]"
       ]
      }
     ],
     "prompt_number": 34
    },
    {
     "cell_type": "code",
     "collapsed": true,
     "input": [],
     "language": "python",
     "outputs": []
    }
   ]
  }
 ]
}