docs/source/examples/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](https://ipyparallel.readthedocs.io/en/stable/tutorial/process.html#using-ipython-parallel-with-mpi)\n",
    "and explains [basic MPI usage of the IPython cluster](https://ipyparallel.readthedocs.io/en/stable/reference/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,
    "jupyter": {
     "outputs_hidden": true
    }
   },
   "outputs": [],
   "source": [
    "import ipyparallel as ipp\n",
    "\n",
    "rc = ipp.Client()\n",
    "view = rc[:]"
   ]
  },
  {
   "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,
    "jupyter": {
     "outputs_hidden": true
    }
   },
   "outputs": [],
   "source": [
    "@view.remote(block=True)\n",
    "def mpi_rank():\n",
    "    from mpi4py import MPI\n",
    "\n",
    "    comm = MPI.COMM_WORLD\n",
    "    return comm.Get_rank()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[0, 1, 3, 2]"
      ]
     },
     "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": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{0: 0, 1: 1, 2: 3, 3: 2}"
      ]
     },
     "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": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\u001b[0;31mOut[0:8]: \u001b[0m{'data': 1, 'rank': 0}"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "\u001b[0;31mOut[1:8]: \u001b[0m{'data': 4, 'rank': 1}"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "\u001b[0;31mOut[2:8]: \u001b[0m{'data': 16, 'rank': 3}"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "\u001b[0;31mOut[3:8]: \u001b[0m{'data': 9, 'rank': 2}"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "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, f'data={data}, rank={rank}'\n",
    "{\n",
    "    'data': data,\n",
    "    'rank': rank,\n",
    "}"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.6"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "state": {},
    "version_major": 2,
    "version_minor": 0
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}