{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "source": [
    "# Using cloudknot to perform matrix-vector multiplication of random matrices"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "source": [
    "This example uses cloudknot to perform matrix-vector multiplication of some random matrices with varying standard deviations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "outputs": [],
   "source": [
    "import cloudknot as ck"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "source": [
    "First, we write the python script that we want to run on AWS batch. Note that we import the necessary python packages within the function `random_mv_prod`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "outputs": [],
   "source": [
    "def random_mv_prod(b):    \n",
    "    import numpy as np\n",
    "    import pandas as pd\n",
    "    import s3fs\n",
    "    import json\n",
    "    import logging\n",
    "    import os.path as op\n",
    "    import nibabel as nib\n",
    "    import dipy.data as dpd\n",
    "    import dipy.tracking.utils as dtu\n",
    "    import dipy.tracking.streamline as dts\n",
    "    from dipy.io.streamline import save_tractogram, load_tractogram\n",
    "    from dipy.stats.analysis import afq_profile, gaussian_weights\n",
    "    from dipy.io.stateful_tractogram import StatefulTractogram\n",
    "    from dipy.io.stateful_tractogram import Space\n",
    "    import dipy.core.gradients as dpg\n",
    "    from dipy.segment.mask import median_otsu\n",
    "    \n",
    "    x = np.random.normal(0, b, 1024)\n",
    "    A = np.random.normal(0, b, (1024, 1024))\n",
    "    \n",
    "    return np.dot(A, x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "source": [
    "Create a knot using the `random_mv_prod` function and a job definition memory of 128 MiB."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "outputs": [],
   "source": [
    "knot = ck.Knot(name='random-mv-prod', base_image=\"python:3.7\", func=random_mv_prod, memory=128, retries=3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "source": [
    "Submit 20 batch jobs to the knot. The `map()` method returns a list of futures for the results of each batch job. You can optionally supply a list of environment variables to each job."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": true,
    "deletable": true,
    "editable": true,
    "jupyter": {
     "outputs_hidden": true
    }
   },
   "outputs": [],
   "source": [
    "# import numpy since it was only imported in the `random_mv_prod` function above\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "collapsed": false,
    "deletable": true,
    "editable": true,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "# Submit the jobs\n",
    "result_future = knot.map(np.linspace(0.1, 100, 17), env_vars=[{'name': 'MY_ENV_VAR', 'value': 'foo'}])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "source": [
    "We can query the jobs associated with this knot by calling `knot.view_jobs()`, prints a bunch of job info and provides a consice summary of job statuses."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "collapsed": false,
    "deletable": true,
    "editable": true,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Job ID              Name                        Status   \n",
      "---------------------------------------------------------\n",
      "565605cc-6c10-45dc-9634-430c92a04d36        random_mv_prod-0            SUBMITTED\n"
     ]
    }
   ],
   "source": [
    "# Rerun this cell as often as you like to update your job status info\n",
    "knot.view_jobs()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "source": [
    "We can also inspect each BatchJob instance by looking at `knot.jobs` which returns a list of BatchJob instances for each submitted job, e.g.:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "collapsed": true,
    "deletable": true,
    "editable": true,
    "jupyter": {
     "outputs_hidden": true
    }
   },
   "outputs": [],
   "source": [
    "last_job = knot.jobs[-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "collapsed": false,
    "deletable": true,
    "editable": true,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "False\n"
     ]
    },
    {
     "ename": "CKTimeoutError",
     "evalue": "The job with job-id 565605cc-6c10-45dc-9634-430c92a04d36 did not finish within the requested timeout period",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mCKTimeoutError\u001b[0m                            Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-8-8f32aad30717>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m()\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mlast_job\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdone\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 2\u001b[0;31m \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mlast_job\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mresult\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtimeout\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m5\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;32m~/code/projects/cloudknot/cloudknot/aws/batch.py\u001b[0m in \u001b[0;36mresult\u001b[0;34m(self, timeout)\u001b[0m\n\u001b[1;32m   1898\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1899\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1900\u001b[0;31m             \u001b[0;32mraise\u001b[0m \u001b[0mCKTimeoutError\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mjob_id\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1901\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1902\u001b[0m         \u001b[0mstatus\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstatus\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mCKTimeoutError\u001b[0m: The job with job-id 565605cc-6c10-45dc-9634-430c92a04d36 did not finish within the requested timeout period"
     ]
    }
   ],
   "source": [
    "print(last_job.done)\n",
    "print(last_job.result(timeout=5))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'arrayProperties': {'size': 17,\n",
       "  'statusSummary': {'FAILED': 0,\n",
       "   'PENDING': 0,\n",
       "   'RUNNABLE': 17,\n",
       "   'RUNNING': 0,\n",
       "   'STARTING': 0,\n",
       "   'SUBMITTED': 0,\n",
       "   'SUCCEEDED': 0}},\n",
       " 'attempts': [],\n",
       " 'status': 'PENDING',\n",
       " 'statusReason': None}"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "last_job.status"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "source": [
    "`Knot.map()` returns a list of futures so you can use any of the futures methods to query the results, e.g. `done()` or `result()`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True\n"
     ]
    }
   ],
   "source": [
    "print(result_future.done())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "collapsed": false,
    "deletable": true,
    "editable": true,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[array([ 0.33573444, -0.17816871,  0.04182915, ...,  0.18977458,\n",
      "        0.18525011,  0.28909461]), array([-2170.50641905,  -914.96192525, -2935.42394086, ...,  -845.61602121,\n",
      "       -1066.55543101,  3786.07847293]), array([ 5518.26250287,  2629.82191285,  3973.26642432, ...,   245.35879457,\n",
      "        2378.91610231,  1440.85016483]), array([ 12145.92971008,  -1752.67075638,  -2236.9687673 , ...,\n",
      "       -10473.04861508,   4948.93624631,  -5164.02257178]), array([ 19220.37103598, -11636.39028164,   4435.09553021, ...,\n",
      "          275.34159951, -40660.23318601, -31436.39319062]), array([-12139.67549791, -37191.77517368,  17676.0058413 , ...,\n",
      "         3871.21504177, -15745.62622439, -15256.27259278]), array([-33710.44216824,  26970.72741577, -25699.97672714, ...,\n",
      "          907.34711287, -12405.55383236,  25991.10145291]), array([ 20642.78548984, -51662.87680929, -13446.32639517, ...,\n",
      "       -13242.89581277,   8647.19798733,  69037.63436123]), array([ -11077.38603201,   25959.65404178, -123167.99186893, ...,\n",
      "        -25486.75276358,   33565.61649554,   68675.4970642 ]), array([-139402.98727787,   35577.84187373,    8167.51240975, ...,\n",
      "       -115266.37136274,   53858.43230435,  -22372.13590293]), array([ -85660.27928164,  -87524.78978449,    5849.47007401, ...,\n",
      "        108583.87339172,  -38679.96416028,   -6415.63103812]), array([  67318.53630424,    1914.09187751,  -20859.14110294, ...,\n",
      "       -156192.64555252, -293051.5218775 , -180415.32923358]), array([ -26414.10480353, -145369.12520887,  -82828.38001282, ...,\n",
      "       -273897.86414521, -161427.48206156,  -88802.77429876]), array([ 107773.50860438,  315763.8425277 ,  -64905.07963653, ...,\n",
      "         32352.21473818,  191698.54867767,  215704.20427246]), array([ 308656.37474034,  147695.23439019,  -26775.28966502, ...,\n",
      "       -159280.69577612,  -88390.58938526,  290458.81708465]), array([ 108520.36949807, -232753.97901184,  162913.86153282, ...,\n",
      "         70634.99906695,   40860.97626671, -597361.98929967]), array([ 453458.8684153 ,  148091.52259826, -430783.77627923, ...,\n",
      "          6989.64826569,  299962.44403145,  370271.46710545])]\n"
     ]
    }
   ],
   "source": [
    "print(result_future.result())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "source": [
    "Once you're all done, clobber the knot, including the underlying PARS and the remote repo."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "collapsed": false,
    "deletable": true,
    "editable": true,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "knot.clobber(clobber_pars=True, clobber_repo=True, clobber_image=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "deletable": true,
    "editable": true,
    "jupyter": {
     "outputs_hidden": true
    }
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}