Draft of gpu setup in multinode env

gpu-setup.ipynb

{
 "metadata": {
  "name": "",
  "signature": "sha256:ce19f2fd638f156fe7be711c682492a6cb58ac36de2960a646cd18e8d91b6f78"
 },
 "nbformat": 3,
 "nbformat_minor": 0,
 "worksheets": [
  {
   "cells": [
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "# Setting up a cluster with gpu nodes"
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "from IPython import parallel\n",
      "rc = parallel.Client()\n",
      "eall = rc[:]"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 24
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "def num_gpus():\n",
      "    import random\n",
      "    return random.randint(4, 12)"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 26
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "class EngineGroup(object):\n",
      "    def __init__(self, n=1):\n",
      "        import subprocess as sp\n",
      "        self.engines = [sp.Popen(['ipengine']) for i in range(n)]\n",
      "        self._stopped = False\n",
      "        \n",
      "    def stop(self):\n",
      "        import signal, time\n",
      "        \n",
      "        if self._stopped: return\n",
      "        \n",
      "        for e in self.engines:\n",
      "            e.send_signal(signal.SIGINT)\n",
      "            \n",
      "        time.sleep(1)  # wait 1s for all engines to terminate\n",
      "        \n",
      "        for e in self.engines:\n",
      "            e.terminate()\n",
      "            \n",
      "        self._stopped = True"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 27
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "def start_gpu():\n",
      "    pass\n",
      "\n",
      "def init_gpus():\n",
      "    ng = num_gpus()\n",
      "    gpu_engines = EngineGroup(1, ng)\n",
      "    "
     ],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 28
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "eall.apply_sync( num_gpus)"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [
      {
       "metadata": {},
       "output_type": "pyout",
       "prompt_number": 29,
       "text": [
        "[6, 10, 9, 12]"
       ]
      }
     ],
     "prompt_number": 29
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [],
     "language": "python",
     "metadata": {},
     "outputs": []
    }
   ],
   "metadata": {}
  }
 ]
}