dask/dask-cupy-multi-node-gpu.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "opposed-recycling",
   "metadata": {},
   "source": [
    "# Dask-jobqueue in Action"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "similar-writer",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<dask.config.set at 0x2b7e89a7a520>"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from dask_jobqueue import PBSCluster\n",
    "\n",
    "import dask\n",
    "\n",
    "dask.config.set(\n",
    "    {\n",
    "        'distributed.dashboard.link': 'https://jupyterhub.hpc.ucar.edu/stable/user/{USER}/dav-compute/proxy/{port}/status'\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "assisted-vocabulary",
   "metadata": {},
   "source": [
    "## Instatiate a cluster object"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "unavailable-comment",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8e4fb40b26b54a1ca44ee3366f472f33",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "VBox(children=(HTML(value='<h2>PBSCluster</h2>'), HBox(children=(HTML(value='\\n<div>\\n  <style scoped>\\n    .d…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "cluster = PBSCluster(\n",
    "    cores=1,\n",
    "    processes=1,\n",
    "    memory=\"10GB\",\n",
    "    queue=\"casper\",\n",
    "    project=\"NTDD0005\",\n",
    "    walltime=\"00:30:00\",\n",
    "    resource_spec=\"select=1:ncpus=1:mem=10GB:ngpus=1\",\n",
    "    extra=[\n",
    "        '--resources GPU=1'\n",
    "    ],  # tag dask-workers: # specify special hardware availability that the scheduler is not aware of\n",
    "    job_extra=['-l gpu_type=v100'],\n",
    "    env_extra=['module load cuda/11.0.3'],  # ensure cuda is loaded\n",
    ")\n",
    "\n",
    "cluster"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "naughty-fellowship",
   "metadata": {},
   "source": [
    "## Batch Job Script"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "foreign-implementation",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "#!/usr/bin/env bash\n",
      "\n",
      "#PBS -N dask-worker\n",
      "#PBS -q casper\n",
      "#PBS -A NTDD0005\n",
      "#PBS -l select=1:ncpus=1:mem=10GB:ngpus=1\n",
      "#PBS -l walltime=00:30:00\n",
      "#PBS -e /glade/scratch/abanihi/\n",
      "#PBS -o /glade/scratch/abanihi/\n",
      "#PBS -l gpu_type=v100\n",
      "module load cuda/11.0.3\n",
      "/glade/work/abanihi/opt/miniconda/envs/dask-gpu/bin/python -m distributed.cli.dask_worker tcp://10.12.205.17:41700 --nthreads 1 --memory-limit 9.31GiB --name dummy-name --nanny --death-timeout 60 --local-directory /glade/scratch/abanihi --resources GPU=1 --interface ib0 --protocol tcp://\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(cluster.job_script())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "voluntary-lover",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "                                                            Req'd  Req'd   Elap\n",
      "Job ID          Username Queue    Jobname    SessID NDS TSK Memory Time  S Time\n",
      "--------------- -------- -------- ---------- ------ --- --- ------ ----- - -----\n",
      "240546.casper-* abanihi  jhublog* STDIN      217212   1   1    4gb 720:0 R 226:0\n",
      "251037.casper-* abanihi  jhublog* STDIN      238092   1   1    4gb 720:0 R 149:3\n",
      "282298.casper-* abanihi  tdd      STDIN       10906   1   1   10gb 00:30 R 00:01\n",
      "282301.casper-* abanihi  tdd      dask-work*  31679   1   1   10gb 00:30 R 00:00\n",
      "282302.casper-* abanihi  tdd      dask-work*    --    1   1   10gb 00:30 R   -- \n"
     ]
    }
   ],
   "source": [
    "!qstat -u abanihi"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "tough-syntax",
   "metadata": {},
   "source": [
    "## Cluster Scaling APIs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "widespread-montgomery",
   "metadata": {},
   "outputs": [],
   "source": [
    "cluster.scale(2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "amino-poker",
   "metadata": {},
   "outputs": [],
   "source": [
    "cluster.scale(jobs=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "differential-ranch",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<distributed.deploy.adaptive.Adaptive at 0x2b7e8985eac0>"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cluster.adapt(minimum=2, maximum=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "hydraulic-greeting",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<distributed.deploy.adaptive.Adaptive at 0x2b7eacd3f400>"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cluster.adapt(minimum_jobs=2, maximum_jobs=4)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dressed-ethernet",
   "metadata": {},
   "source": [
    "## Connect created cluster to a client"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "sustained-simple",
   "metadata": {},
   "outputs": [],
   "source": [
    "from distributed import Client\n",
    "\n",
    "client = Client(cluster)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "documented-ecology",
   "metadata": {},
   "source": [
    "## Confirm that our dask workers have access to GPUs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "twelve-guatemala",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_nvidia_smi_info():\n",
    "    import subprocess\n",
    "\n",
    "    p = subprocess.check_output('nvidia-smi').strip().decode('utf-8')\n",
    "    return p\n",
    "\n",
    "\n",
    "def nvidia_smi(on='workers'):\n",
    "    if on == 'workers':\n",
    "        x = client.run(get_nvidia_smi_info)\n",
    "        print(\" ***** NVIDIA-SMI info on Workers *****\")\n",
    "        for key, value in x.items():\n",
    "            print(\"*\" * 80)\n",
    "            print(key)\n",
    "            print(value, end=\"\\n\\n\")\n",
    "\n",
    "    elif on == 'scheduler':\n",
    "        print(\"***** NVIDIA-SMI info on Scheduler *****\")\n",
    "        print(client.run_on_scheduler(get_nvidia_smi_info))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "least-colors",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " ***** NVIDIA-SMI info on Workers *****\n",
      "********************************************************************************\n",
      "tcp://10.12.205.38:40174\n",
      "Wed May 19 22:18:08 2021       \n",
      "+-----------------------------------------------------------------------------+\n",
      "| NVIDIA-SMI 450.51.06    Driver Version: 450.51.06    CUDA Version: 11.0     |\n",
      "|-------------------------------+----------------------+----------------------+\n",
      "| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |\n",
      "| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |\n",
      "|                               |                      |               MIG M. |\n",
      "|===============================+======================+======================|\n",
      "|   0  Tesla V100-SXM2...  On   | 00000000:B2:00.0 Off |                    0 |\n",
      "| N/A   31C    P0    40W / 300W |      0MiB / 32510MiB |      0%      Default |\n",
      "|                               |                      |                  N/A |\n",
      "+-------------------------------+----------------------+----------------------+\n",
      "                                                                               \n",
      "+-----------------------------------------------------------------------------+\n",
      "| Processes:                                                                  |\n",
      "|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |\n",
      "|        ID   ID                                                   Usage      |\n",
      "|=============================================================================|\n",
      "|  No running processes found                                                 |\n",
      "+-----------------------------------------------------------------------------+\n",
      "\n",
      "********************************************************************************\n",
      "tcp://10.12.205.38:41472\n",
      "Wed May 19 22:18:08 2021       \n",
      "+-----------------------------------------------------------------------------+\n",
      "| NVIDIA-SMI 450.51.06    Driver Version: 450.51.06    CUDA Version: 11.0     |\n",
      "|-------------------------------+----------------------+----------------------+\n",
      "| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |\n",
      "| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |\n",
      "|                               |                      |               MIG M. |\n",
      "|===============================+======================+======================|\n",
      "|   0  Tesla V100-SXM2...  On   | 00000000:1B:00.0 Off |                    0 |\n",
      "| N/A   31C    P0    40W / 300W |      0MiB / 32510MiB |      0%      Default |\n",
      "|                               |                      |                  N/A |\n",
      "+-------------------------------+----------------------+----------------------+\n",
      "                                                                               \n",
      "+-----------------------------------------------------------------------------+\n",
      "| Processes:                                                                  |\n",
      "|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |\n",
      "|        ID   ID                                                   Usage      |\n",
      "|=============================================================================|\n",
      "|  No running processes found                                                 |\n",
      "+-----------------------------------------------------------------------------+\n",
      "\n"
     ]
    }
   ],
   "source": [
    "nvidia_smi()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "global-hollywood",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "***** NVIDIA-SMI info on Scheduler *****\n",
      "Wed May 19 22:18:09 2021       \n",
      "+-----------------------------------------------------------------------------+\n",
      "| NVIDIA-SMI 450.51.06    Driver Version: 450.51.06    CUDA Version: 11.0     |\n",
      "|-------------------------------+----------------------+----------------------+\n",
      "| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |\n",
      "| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |\n",
      "|                               |                      |               MIG M. |\n",
      "|===============================+======================+======================|\n",
      "|   0  Tesla V100-SXM2...  On   | 00000000:B3:00.0 Off |                    0 |\n",
      "| N/A   28C    P0    40W / 300W |      0MiB / 32510MiB |      0%      Default |\n",
      "|                               |                      |                  N/A |\n",
      "+-------------------------------+----------------------+----------------------+\n",
      "                                                                               \n",
      "+-----------------------------------------------------------------------------+\n",
      "| Processes:                                                                  |\n",
      "|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |\n",
      "|        ID   ID                                                   Usage      |\n",
      "|=============================================================================|\n",
      "|  No running processes found                                                 |\n",
      "+-----------------------------------------------------------------------------+\n"
     ]
    }
   ],
   "source": [
    "nvidia_smi(on='scheduler')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "found-split",
   "metadata": {},
   "source": [
    "## Run some computation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "wooden-quarterly",
   "metadata": {},
   "outputs": [],
   "source": [
    "import cupy\n",
    "\n",
    "import dask.array as da"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "bulgarian-point",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<table>\n",
       "<tr>\n",
       "<td>\n",
       "<table>\n",
       "  <thead>\n",
       "    <tr><td> </td><th> Array </th><th> Chunk </th></tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr><th> Bytes </th><td> 149.01 GiB </td> <td> 152.59 MiB </td></tr>\n",
       "    <tr><th> Shape </th><td> (200000, 200000) </td> <td> (10000, 4000) </td></tr>\n",
       "    <tr><th> Count </th><td> 1000 Tasks </td><td> 1000 Chunks </td></tr>\n",
       "    <tr><th> Type </th><td> float32 </td><td> cupy.ndarray </td></tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</td>\n",
       "<td>\n",
       "<svg width=\"170\" height=\"170\" style=\"stroke:rgb(0,0,0);stroke-width:1\" >\n",
       "\n",
       "  <!-- Horizontal lines -->\n",
       "  <line x1=\"0\" y1=\"0\" x2=\"120\" y2=\"0\" style=\"stroke-width:2\" />\n",
       "  <line x1=\"0\" y1=\"6\" x2=\"120\" y2=\"6\" />\n",
       "  <line x1=\"0\" y1=\"12\" x2=\"120\" y2=\"12\" />\n",
       "  <line x1=\"0\" y1=\"18\" x2=\"120\" y2=\"18\" />\n",
       "  <line x1=\"0\" y1=\"24\" x2=\"120\" y2=\"24\" />\n",
       "  <line x1=\"0\" y1=\"30\" x2=\"120\" y2=\"30\" />\n",
       "  <line x1=\"0\" y1=\"36\" x2=\"120\" y2=\"36\" />\n",
       "  <line x1=\"0\" y1=\"42\" x2=\"120\" y2=\"42\" />\n",
       "  <line x1=\"0\" y1=\"48\" x2=\"120\" y2=\"48\" />\n",
       "  <line x1=\"0\" y1=\"54\" x2=\"120\" y2=\"54\" />\n",
       "  <line x1=\"0\" y1=\"60\" x2=\"120\" y2=\"60\" />\n",
       "  <line x1=\"0\" y1=\"66\" x2=\"120\" y2=\"66\" />\n",
       "  <line x1=\"0\" y1=\"72\" x2=\"120\" y2=\"72\" />\n",
       "  <line x1=\"0\" y1=\"78\" x2=\"120\" y2=\"78\" />\n",
       "  <line x1=\"0\" y1=\"84\" x2=\"120\" y2=\"84\" />\n",
       "  <line x1=\"0\" y1=\"90\" x2=\"120\" y2=\"90\" />\n",
       "  <line x1=\"0\" y1=\"96\" x2=\"120\" y2=\"96\" />\n",
       "  <line x1=\"0\" y1=\"102\" x2=\"120\" y2=\"102\" />\n",
       "  <line x1=\"0\" y1=\"108\" x2=\"120\" y2=\"108\" />\n",
       "  <line x1=\"0\" y1=\"120\" x2=\"120\" y2=\"120\" style=\"stroke-width:2\" />\n",
       "\n",
       "  <!-- Vertical lines -->\n",
       "  <line x1=\"0\" y1=\"0\" x2=\"0\" y2=\"120\" style=\"stroke-width:2\" />\n",
       "  <line x1=\"4\" y1=\"0\" x2=\"4\" y2=\"120\" />\n",
       "  <line x1=\"12\" y1=\"0\" x2=\"12\" y2=\"120\" />\n",
       "  <line x1=\"16\" y1=\"0\" x2=\"16\" y2=\"120\" />\n",
       "  <line x1=\"24\" y1=\"0\" x2=\"24\" y2=\"120\" />\n",
       "  <line x1=\"31\" y1=\"0\" x2=\"31\" y2=\"120\" />\n",
       "  <line x1=\"36\" y1=\"0\" x2=\"36\" y2=\"120\" />\n",
       "  <line x1=\"43\" y1=\"0\" x2=\"43\" y2=\"120\" />\n",
       "  <line x1=\"50\" y1=\"0\" x2=\"50\" y2=\"120\" />\n",
       "  <line x1=\"55\" y1=\"0\" x2=\"55\" y2=\"120\" />\n",
       "  <line x1=\"62\" y1=\"0\" x2=\"62\" y2=\"120\" />\n",
       "  <line x1=\"67\" y1=\"0\" x2=\"67\" y2=\"120\" />\n",
       "  <line x1=\"74\" y1=\"0\" x2=\"74\" y2=\"120\" />\n",
       "  <line x1=\"81\" y1=\"0\" x2=\"81\" y2=\"120\" />\n",
       "  <line x1=\"86\" y1=\"0\" x2=\"86\" y2=\"120\" />\n",
       "  <line x1=\"93\" y1=\"0\" x2=\"93\" y2=\"120\" />\n",
       "  <line x1=\"100\" y1=\"0\" x2=\"100\" y2=\"120\" />\n",
       "  <line x1=\"105\" y1=\"0\" x2=\"105\" y2=\"120\" />\n",
       "  <line x1=\"112\" y1=\"0\" x2=\"112\" y2=\"120\" />\n",
       "  <line x1=\"120\" y1=\"0\" x2=\"120\" y2=\"120\" style=\"stroke-width:2\" />\n",
       "\n",
       "  <!-- Colored Rectangle -->\n",
       "  <polygon points=\"0.0,0.0 120.0,0.0 120.0,120.0 0.0,120.0\" style=\"fill:#8B4903A0;stroke-width:0\"/>\n",
       "\n",
       "  <!-- Text -->\n",
       "  <text x=\"60.000000\" y=\"140.000000\" font-size=\"1.0rem\" font-weight=\"100\" text-anchor=\"middle\" >200000</text>\n",
       "  <text x=\"140.000000\" y=\"60.000000\" font-size=\"1.0rem\" font-weight=\"100\" text-anchor=\"middle\" transform=\"rotate(-90,140.000000,60.000000)\">200000</text>\n",
       "</svg>\n",
       "</td>\n",
       "</tr>\n",
       "</table>"
      ],
      "text/plain": [
       "dask.array<normal, shape=(200000, 200000), dtype=float32, chunksize=(10000, 4000), chunktype=cupy.ndarray>"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# generate chunked dask arrays of many cupy random arrays\n",
    "rs = da.random.RandomState(RandomState=cupy.random.RandomState)  # <-- we specify cupy here\n",
    "x = rs.normal(10, 1, size=(200000, 200000), chunks=(10000, 4000), dtype=cupy.float32)\n",
    "x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "historic-shipping",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<table>\n",
       "<tr>\n",
       "<td>\n",
       "<table>\n",
       "  <thead>\n",
       "    <tr><td> </td><th> Array </th><th> Chunk </th></tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr><th> Bytes </th><td> 390.62 kiB </td> <td> 7.81 kiB </td></tr>\n",
       "    <tr><th> Shape </th><td> (100000,) </td> <td> (2000,) </td></tr>\n",
       "    <tr><th> Count </th><td> 50 Tasks </td><td> 50 Chunks </td></tr>\n",
       "    <tr><th> Type </th><td> float32 </td><td> cupy.ndarray </td></tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</td>\n",
       "<td>\n",
       "<svg width=\"170\" height=\"75\" style=\"stroke:rgb(0,0,0);stroke-width:1\" >\n",
       "\n",
       "  <!-- Horizontal lines -->\n",
       "  <line x1=\"0\" y1=\"0\" x2=\"120\" y2=\"0\" style=\"stroke-width:2\" />\n",
       "  <line x1=\"0\" y1=\"25\" x2=\"120\" y2=\"25\" style=\"stroke-width:2\" />\n",
       "\n",
       "  <!-- Vertical lines -->\n",
       "  <line x1=\"0\" y1=\"0\" x2=\"0\" y2=\"25\" style=\"stroke-width:2\" />\n",
       "  <line x1=\"4\" y1=\"0\" x2=\"4\" y2=\"25\" />\n",
       "  <line x1=\"12\" y1=\"0\" x2=\"12\" y2=\"25\" />\n",
       "  <line x1=\"16\" y1=\"0\" x2=\"16\" y2=\"25\" />\n",
       "  <line x1=\"24\" y1=\"0\" x2=\"24\" y2=\"25\" />\n",
       "  <line x1=\"31\" y1=\"0\" x2=\"31\" y2=\"25\" />\n",
       "  <line x1=\"36\" y1=\"0\" x2=\"36\" y2=\"25\" />\n",
       "  <line x1=\"43\" y1=\"0\" x2=\"43\" y2=\"25\" />\n",
       "  <line x1=\"50\" y1=\"0\" x2=\"50\" y2=\"25\" />\n",
       "  <line x1=\"55\" y1=\"0\" x2=\"55\" y2=\"25\" />\n",
       "  <line x1=\"62\" y1=\"0\" x2=\"62\" y2=\"25\" />\n",
       "  <line x1=\"67\" y1=\"0\" x2=\"67\" y2=\"25\" />\n",
       "  <line x1=\"74\" y1=\"0\" x2=\"74\" y2=\"25\" />\n",
       "  <line x1=\"81\" y1=\"0\" x2=\"81\" y2=\"25\" />\n",
       "  <line x1=\"86\" y1=\"0\" x2=\"86\" y2=\"25\" />\n",
       "  <line x1=\"93\" y1=\"0\" x2=\"93\" y2=\"25\" />\n",
       "  <line x1=\"100\" y1=\"0\" x2=\"100\" y2=\"25\" />\n",
       "  <line x1=\"105\" y1=\"0\" x2=\"105\" y2=\"25\" />\n",
       "  <line x1=\"112\" y1=\"0\" x2=\"112\" y2=\"25\" />\n",
       "  <line x1=\"120\" y1=\"0\" x2=\"120\" y2=\"25\" style=\"stroke-width:2\" />\n",
       "\n",
       "  <!-- Colored Rectangle -->\n",
       "  <polygon points=\"0.0,0.0 120.0,0.0 120.0,25.412616514582485 0.0,25.412616514582485\" style=\"fill:#8B4903A0;stroke-width:0\"/>\n",
       "\n",
       "  <!-- Text -->\n",
       "  <text x=\"60.000000\" y=\"45.412617\" font-size=\"1.0rem\" font-weight=\"100\" text-anchor=\"middle\" >100000</text>\n",
       "  <text x=\"140.000000\" y=\"12.706308\" font-size=\"1.0rem\" font-weight=\"100\" text-anchor=\"middle\" transform=\"rotate(0,140.000000,12.706308)\">1</text>\n",
       "</svg>\n",
       "</td>\n",
       "</tr>\n",
       "</table>"
      ],
      "text/plain": [
       "dask.array<_sqrt, shape=(100000,), dtype=float32, chunksize=(2000,), chunktype=cupy.ndarray>"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "y = (x + 1)[::2, ::2].std(axis=0)\n",
    "y = y.persist()\n",
    "y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "presidential-patio",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 30 ms, sys: 2 ms, total: 32 ms\n",
      "Wall time: 36.7 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "result = y.compute()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "expired-weather",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([0.9968655 , 0.99782085, 1.0016662 , ..., 1.0015868 , 0.9996461 ,\n",
       "       0.9991475 ], dtype=float32)"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "matched-greene",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "cupy._core.core.ndarray"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "type(result)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "express-mongolia",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " ***** NVIDIA-SMI info on Workers *****\n",
      "********************************************************************************\n",
      "tcp://10.12.205.38:40174\n",
      "Wed May 19 22:23:17 2021       \n",
      "+-----------------------------------------------------------------------------+\n",
      "| NVIDIA-SMI 450.51.06    Driver Version: 450.51.06    CUDA Version: 11.0     |\n",
      "|-------------------------------+----------------------+----------------------+\n",
      "| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |\n",
      "| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |\n",
      "|                               |                      |               MIG M. |\n",
      "|===============================+======================+======================|\n",
      "|   0  Tesla V100-SXM2...  On   | 00000000:B2:00.0 Off |                    0 |\n",
      "| N/A   32C    P0    52W / 300W |    926MiB / 32510MiB |      0%      Default |\n",
      "|                               |                      |                  N/A |\n",
      "+-------------------------------+----------------------+----------------------+\n",
      "                                                                               \n",
      "+-----------------------------------------------------------------------------+\n",
      "| Processes:                                                                  |\n",
      "|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |\n",
      "|        ID   ID                                                   Usage      |\n",
      "|=============================================================================|\n",
      "|    0   N/A  N/A     31956      C   .../envs/dask-gpu/bin/python      923MiB |\n",
      "+-----------------------------------------------------------------------------+\n",
      "\n",
      "********************************************************************************\n",
      "tcp://10.12.205.38:41472\n",
      "Wed May 19 22:23:17 2021       \n",
      "+-----------------------------------------------------------------------------+\n",
      "| NVIDIA-SMI 450.51.06    Driver Version: 450.51.06    CUDA Version: 11.0     |\n",
      "|-------------------------------+----------------------+----------------------+\n",
      "| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |\n",
      "| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |\n",
      "|                               |                      |               MIG M. |\n",
      "|===============================+======================+======================|\n",
      "|   0  Tesla V100-SXM2...  On   | 00000000:1B:00.0 Off |                    0 |\n",
      "| N/A   33C    P0    54W / 300W |    926MiB / 32510MiB |      0%      Default |\n",
      "|                               |                      |                  N/A |\n",
      "+-------------------------------+----------------------+----------------------+\n",
      "                                                                               \n",
      "+-----------------------------------------------------------------------------+\n",
      "| Processes:                                                                  |\n",
      "|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |\n",
      "|        ID   ID                                                   Usage      |\n",
      "|=============================================================================|\n",
      "|    0   N/A  N/A     31847      C   .../envs/dask-gpu/bin/python      923MiB |\n",
      "+-----------------------------------------------------------------------------+\n",
      "\n"
     ]
    }
   ],
   "source": [
    "nvidia_smi()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "herbal-pendant",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:miniconda-dask-gpu]",
   "language": "python",
   "name": "conda-env-miniconda-dask-gpu-py"
  },
  "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.8.10"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "state": {
     "02cae7acb2234a5280667860807ab9ed": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "1.5.0",
      "model_name": "HBoxModel",
      "state": {
       "children": [
        "IPY_MODEL_4864192b774a452c9253e3d2c9120d3f",
        "IPY_MODEL_77e8a945c429446c9fdac519ccbf4cf8"
       ],
       "layout": "IPY_MODEL_78a5fb90edd74a5397bc0322c2602ee1"
      }
     },
     "02e45b1ddd064408a1928991815f76ea": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "1.5.0",
      "model_name": "HBoxModel",
      "state": {
       "children": [
        "IPY_MODEL_858712101e594f958e008a77b480870b",
        "IPY_MODEL_a4b3396c04be47eab437678bed1b745e",
        "IPY_MODEL_4b0bfbbf543944d581462371240c87bd"
       ],
       "layout": "IPY_MODEL_e1af4bddf7b24f679de43c3470e211ab"
      }
     },
     "0565deab87d14cf4b2f15ee5e463e63d": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "1.2.0",
      "model_name": "LayoutModel",
      "state": {}
     },
     "157d171d8139497cb553d7d0ee348cb9": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "1.5.0",
      "model_name": "HBoxModel",
      "state": {
       "children": [
        "IPY_MODEL_9f4cb6ddc4b945728be294aa2a6805ce",
        "IPY_MODEL_e5ae9eb9bdb944d7bdf51d873982b3ca"
       ],
       "layout": "IPY_MODEL_51c993229cb341b8a7a253ad5784bf3a"
      }
     },
     "3ca20f9462fa4d648fe5f2de8a0bac25": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "1.5.0",
      "model_name": "HTMLModel",
      "state": {
       "layout": "IPY_MODEL_58e66b04611c40b794e573b07d10f217",
       "style": "IPY_MODEL_72b2e0026713449cae6922e98982ba64",
       "value": "<h2>PBSCluster</h2>"
      }
     },
     "4864192b774a452c9253e3d2c9120d3f": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "1.5.0",
      "model_name": "IntTextModel",
      "state": {
       "description": "Workers",
       "layout": "IPY_MODEL_97ad47d5572c439db2dd578af929e01b",
       "step": 1,
       "style": "IPY_MODEL_9157410c11164bab8e740888e521e050"
      }
     },
     "4b0bfbbf543944d581462371240c87bd": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "1.5.0",
      "model_name": "ButtonModel",
      "state": {
       "description": "Adapt",
       "layout": "IPY_MODEL_97ad47d5572c439db2dd578af929e01b",
       "style": "IPY_MODEL_d4ac4c834c3e47248d77dc9b980d6e79"
      }
     },
     "51c993229cb341b8a7a253ad5784bf3a": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "1.2.0",
      "model_name": "LayoutModel",
      "state": {}
     },
     "58e66b04611c40b794e573b07d10f217": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "1.2.0",
      "model_name": "LayoutModel",
      "state": {}
     },
     "72b2e0026713449cae6922e98982ba64": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "1.5.0",
      "model_name": "DescriptionStyleModel",
      "state": {
       "description_width": ""
      }
     },
     "77e8a945c429446c9fdac519ccbf4cf8": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "1.5.0",
      "model_name": "ButtonModel",
      "state": {
       "description": "Scale",
       "layout": "IPY_MODEL_97ad47d5572c439db2dd578af929e01b",
       "style": "IPY_MODEL_a506f402d9ca490fa13330f287c6f157"
      }
     },
     "78a5fb90edd74a5397bc0322c2602ee1": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "1.2.0",
      "model_name": "LayoutModel",
      "state": {}
     },
     "7b8a38b6910d46ce83ceb577f5f77718": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "1.5.0",
      "model_name": "DescriptionStyleModel",
      "state": {
       "description_width": ""
      }
     },
     "858712101e594f958e008a77b480870b": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "1.5.0",
      "model_name": "IntTextModel",
      "state": {
       "description": "Minimum",
       "layout": "IPY_MODEL_97ad47d5572c439db2dd578af929e01b",
       "step": 1,
       "style": "IPY_MODEL_c79a59a89ac7437fbcea81217b22e3c9",
       "value": 2
      }
     },
     "88fa9d1682474b5f99b80e0a081def02": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "1.2.0",
      "model_name": "LayoutModel",
      "state": {
       "min_width": "500px"
      }
     },
     "88fac047c80c449d917588246c1b5621": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "1.5.0",
      "model_name": "DescriptionStyleModel",
      "state": {
       "description_width": ""
      }
     },
     "8e4fb40b26b54a1ca44ee3366f472f33": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "1.5.0",
      "model_name": "VBoxModel",
      "state": {
       "children": [
        "IPY_MODEL_3ca20f9462fa4d648fe5f2de8a0bac25",
        "IPY_MODEL_157d171d8139497cb553d7d0ee348cb9",
        "IPY_MODEL_ab96334b912f4dac9f2454311df4b9aa"
       ],
       "layout": "IPY_MODEL_fb941851b9774d3aa4082581b0f171f2"
      }
     },
     "9157410c11164bab8e740888e521e050": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "1.5.0",
      "model_name": "DescriptionStyleModel",
      "state": {
       "description_width": ""
      }
     },
     "97ad47d5572c439db2dd578af929e01b": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "1.2.0",
      "model_name": "LayoutModel",
      "state": {
       "width": "150px"
      }
     },
     "9ccffffe365642d0a4907bd95e954bb3": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "1.5.0",
      "model_name": "DescriptionStyleModel",
      "state": {
       "description_width": ""
      }
     },
     "9f4cb6ddc4b945728be294aa2a6805ce": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "1.5.0",
      "model_name": "HTMLModel",
      "state": {
       "layout": "IPY_MODEL_d50931b9950d41268f58203bee552fc9",
       "style": "IPY_MODEL_88fac047c80c449d917588246c1b5621",
       "value": "\n<div>\n  <style scoped>\n    .dataframe tbody tr th:only-of-type {\n        vertical-align: middle;\n    }\n\n    .dataframe tbody tr th {\n        vertical-align: top;\n    }\n\n    .dataframe thead th {\n        text-align: right;\n    }\n  </style>\n  <table style=\"text-align: right;\">\n    <tr> <th>Workers</th> <td>2</td></tr>\n    <tr> <th>Cores</th> <td>2</td></tr>\n    <tr> <th>Memory</th> <td>18.62 GiB</td></tr>\n  </table>\n</div>\n"
      }
     },
     "a4b3396c04be47eab437678bed1b745e": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "1.5.0",
      "model_name": "IntTextModel",
      "state": {
       "description": "Maximum",
       "layout": "IPY_MODEL_97ad47d5572c439db2dd578af929e01b",
       "step": 1,
       "style": "IPY_MODEL_9ccffffe365642d0a4907bd95e954bb3",
       "value": 4
      }
     },
     "a506f402d9ca490fa13330f287c6f157": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "1.5.0",
      "model_name": "ButtonStyleModel",
      "state": {}
     },
     "ab96334b912f4dac9f2454311df4b9aa": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "1.5.0",
      "model_name": "HTMLModel",
      "state": {
       "layout": "IPY_MODEL_0565deab87d14cf4b2f15ee5e463e63d",
       "style": "IPY_MODEL_7b8a38b6910d46ce83ceb577f5f77718",
       "value": "<p><b>Dashboard: </b><a href=\"https://jupyterhub.hpc.ucar.edu/stable/user/abanihi/dav-compute/proxy/8787/status\" target=\"_blank\">https://jupyterhub.hpc.ucar.edu/stable/user/abanihi/dav-compute/proxy/8787/status</a></p>\n"
      }
     },
     "c79a59a89ac7437fbcea81217b22e3c9": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "1.5.0",
      "model_name": "DescriptionStyleModel",
      "state": {
       "description_width": ""
      }
     },
     "d4ac4c834c3e47248d77dc9b980d6e79": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "1.5.0",
      "model_name": "ButtonStyleModel",
      "state": {}
     },
     "d50931b9950d41268f58203bee552fc9": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "1.2.0",
      "model_name": "LayoutModel",
      "state": {
       "min_width": "150px"
      }
     },
     "e1af4bddf7b24f679de43c3470e211ab": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "1.2.0",
      "model_name": "LayoutModel",
      "state": {}
     },
     "e5ae9eb9bdb944d7bdf51d873982b3ca": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "1.5.0",
      "model_name": "AccordionModel",
      "state": {
       "_titles": {
        "0": "Manual Scaling",
        "1": "Adaptive Scaling"
       },
       "children": [
        "IPY_MODEL_02cae7acb2234a5280667860807ab9ed",
        "IPY_MODEL_02e45b1ddd064408a1928991815f76ea"
       ],
       "layout": "IPY_MODEL_88fa9d1682474b5f99b80e0a081def02",
       "selected_index": 1
      }
     },
     "fb941851b9774d3aa4082581b0f171f2": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "1.2.0",
      "model_name": "LayoutModel",
      "state": {}
     }
    },
    "version_major": 2,
    "version_minor": 0
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}