multi_gpu_show.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "79e4dac4",
   "metadata": {},
   "source": [
    "# Multi-GPU workflow\n",
    "**Author:** [Severin Dicks](https://github.com/Intron7)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "37588552-b9d5-4113-9f53-ecc63d3da815",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import dask\n",
    "import time\n",
    "\n",
    "\n",
    "from dask_cuda import LocalCUDACluster\n",
    "from dask.distributed import Client"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "2f159c97-e435-40c5-ac7b-aa37a6812ced",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import rmm\n",
    "import cupy as cp\n",
    "\n",
    "from rmm.allocators.cupy import rmm_cupy_allocator\n",
    "\n",
    "def set_mem():\n",
    "    rmm.reinitialize(managed_memory=True)\n",
    "    cp.cuda.set_allocator(rmm_cupy_allocator)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a81da4d7",
   "metadata": {},
   "source": [
    "First we have to setup our dask-cluster since we are using a DGX we are using all 8 GPUs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "49a27e3c-5473-42eb-a0e3-73a22a1a9e16",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "preprocessing_gpus=\"0,1,2,3,4,5,6,7\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "79faf22b-629e-43d6-a762-02dd1fa2e65c",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 527 ms, sys: 719 ms, total: 1.25 s\n",
      "Wall time: 12.1 s\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "    <div style=\"width: 24px; height: 24px; background-color: #e1e1e1; border: 3px solid #9D9D9D; border-radius: 5px; position: absolute;\"> </div>\n",
       "    <div style=\"margin-left: 48px;\">\n",
       "        <h3 style=\"margin-bottom: 0px;\">Client</h3>\n",
       "        <p style=\"color: #9D9D9D; margin-bottom: 0px;\">Client-0837fc3f-7105-11ef-ae36-00001029fe80</p>\n",
       "        <table style=\"width: 100%; text-align: left;\">\n",
       "\n",
       "        <tr>\n",
       "        \n",
       "            <td style=\"text-align: left;\"><strong>Connection method:</strong> Cluster object</td>\n",
       "            <td style=\"text-align: left;\"><strong>Cluster type:</strong> dask_cuda.LocalCUDACluster</td>\n",
       "        \n",
       "        </tr>\n",
       "\n",
       "        \n",
       "            <tr>\n",
       "                <td style=\"text-align: left;\">\n",
       "                    <strong>Dashboard: </strong> <a href=\"http://127.0.0.1:8787/status\" target=\"_blank\">http://127.0.0.1:8787/status</a>\n",
       "                </td>\n",
       "                <td style=\"text-align: left;\"></td>\n",
       "            </tr>\n",
       "        \n",
       "\n",
       "        </table>\n",
       "\n",
       "        \n",
       "\n",
       "        \n",
       "            <details>\n",
       "            <summary style=\"margin-bottom: 20px;\"><h3 style=\"display: inline;\">Cluster Info</h3></summary>\n",
       "            <div class=\"jp-RenderedHTMLCommon jp-RenderedHTML jp-mod-trusted jp-OutputArea-output\">\n",
       "    <div style=\"width: 24px; height: 24px; background-color: #e1e1e1; border: 3px solid #9D9D9D; border-radius: 5px; position: absolute;\">\n",
       "    </div>\n",
       "    <div style=\"margin-left: 48px;\">\n",
       "        <h3 style=\"margin-bottom: 0px; margin-top: 0px;\">LocalCUDACluster</h3>\n",
       "        <p style=\"color: #9D9D9D; margin-bottom: 0px;\">ecb693e0</p>\n",
       "        <table style=\"width: 100%; text-align: left;\">\n",
       "            <tr>\n",
       "                <td style=\"text-align: left;\">\n",
       "                    <strong>Dashboard:</strong> <a href=\"http://127.0.0.1:8787/status\" target=\"_blank\">http://127.0.0.1:8787/status</a>\n",
       "                </td>\n",
       "                <td style=\"text-align: left;\">\n",
       "                    <strong>Workers:</strong> 8\n",
       "                </td>\n",
       "            </tr>\n",
       "            <tr>\n",
       "                <td style=\"text-align: left;\">\n",
       "                    <strong>Total threads:</strong> 8\n",
       "                </td>\n",
       "                <td style=\"text-align: left;\">\n",
       "                    <strong>Total memory:</strong> 1.95 TiB\n",
       "                </td>\n",
       "            </tr>\n",
       "            \n",
       "            <tr>\n",
       "    <td style=\"text-align: left;\"><strong>Status:</strong> running</td>\n",
       "    <td style=\"text-align: left;\"><strong>Using processes:</strong> True</td>\n",
       "</tr>\n",
       "\n",
       "            \n",
       "        </table>\n",
       "\n",
       "        <details>\n",
       "            <summary style=\"margin-bottom: 20px;\">\n",
       "                <h3 style=\"display: inline;\">Scheduler Info</h3>\n",
       "            </summary>\n",
       "\n",
       "            <div style=\"\">\n",
       "    <div>\n",
       "        <div style=\"width: 24px; height: 24px; background-color: #FFF7E5; border: 3px solid #FF6132; border-radius: 5px; position: absolute;\"> </div>\n",
       "        <div style=\"margin-left: 48px;\">\n",
       "            <h3 style=\"margin-bottom: 0px;\">Scheduler</h3>\n",
       "            <p style=\"color: #9D9D9D; margin-bottom: 0px;\">Scheduler-ecd26c72-1b91-4791-a23e-af10125e35b0</p>\n",
       "            <table style=\"width: 100%; text-align: left;\">\n",
       "                <tr>\n",
       "                    <td style=\"text-align: left;\">\n",
       "                        <strong>Comm:</strong> tcp://127.0.0.1:8183\n",
       "                    </td>\n",
       "                    <td style=\"text-align: left;\">\n",
       "                        <strong>Workers:</strong> 8\n",
       "                    </td>\n",
       "                </tr>\n",
       "                <tr>\n",
       "                    <td style=\"text-align: left;\">\n",
       "                        <strong>Dashboard:</strong> <a href=\"http://127.0.0.1:8787/status\" target=\"_blank\">http://127.0.0.1:8787/status</a>\n",
       "                    </td>\n",
       "                    <td style=\"text-align: left;\">\n",
       "                        <strong>Total threads:</strong> 8\n",
       "                    </td>\n",
       "                </tr>\n",
       "                <tr>\n",
       "                    <td style=\"text-align: left;\">\n",
       "                        <strong>Started:</strong> Just now\n",
       "                    </td>\n",
       "                    <td style=\"text-align: left;\">\n",
       "                        <strong>Total memory:</strong> 1.95 TiB\n",
       "                    </td>\n",
       "                </tr>\n",
       "            </table>\n",
       "        </div>\n",
       "    </div>\n",
       "\n",
       "    <details style=\"margin-left: 48px;\">\n",
       "        <summary style=\"margin-bottom: 20px;\">\n",
       "            <h3 style=\"display: inline;\">Workers</h3>\n",
       "        </summary>\n",
       "\n",
       "        \n",
       "        <div style=\"margin-bottom: 20px;\">\n",
       "            <div style=\"width: 24px; height: 24px; background-color: #DBF5FF; border: 3px solid #4CC9FF; border-radius: 5px; position: absolute;\"> </div>\n",
       "            <div style=\"margin-left: 48px;\">\n",
       "            <details>\n",
       "                <summary>\n",
       "                    <h4 style=\"margin-bottom: 0px; display: inline;\">Worker: 0</h4>\n",
       "                </summary>\n",
       "                <table style=\"width: 100%; text-align: left;\">\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Comm: </strong> tcp://127.0.0.1:8467\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Total threads: </strong> 1\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Dashboard: </strong> <a href=\"http://127.0.0.1:8041/status\" target=\"_blank\">http://127.0.0.1:8041/status</a>\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Memory: </strong> 250.00 GiB\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Nanny: </strong> tcp://127.0.0.1:8299\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\"></td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td colspan=\"2\" style=\"text-align: left;\">\n",
       "                            <strong>Local directory: </strong> /tmp/dask-scratch-space/worker-qzl185wj\n",
       "                        </td>\n",
       "                    </tr>\n",
       "\n",
       "                    \n",
       "\n",
       "                    \n",
       "\n",
       "                </table>\n",
       "            </details>\n",
       "            </div>\n",
       "        </div>\n",
       "        \n",
       "        <div style=\"margin-bottom: 20px;\">\n",
       "            <div style=\"width: 24px; height: 24px; background-color: #DBF5FF; border: 3px solid #4CC9FF; border-radius: 5px; position: absolute;\"> </div>\n",
       "            <div style=\"margin-left: 48px;\">\n",
       "            <details>\n",
       "                <summary>\n",
       "                    <h4 style=\"margin-bottom: 0px; display: inline;\">Worker: 1</h4>\n",
       "                </summary>\n",
       "                <table style=\"width: 100%; text-align: left;\">\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Comm: </strong> tcp://127.0.0.1:8369\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Total threads: </strong> 1\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Dashboard: </strong> <a href=\"http://127.0.0.1:8385/status\" target=\"_blank\">http://127.0.0.1:8385/status</a>\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Memory: </strong> 250.00 GiB\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Nanny: </strong> tcp://127.0.0.1:8329\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\"></td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td colspan=\"2\" style=\"text-align: left;\">\n",
       "                            <strong>Local directory: </strong> /tmp/dask-scratch-space/worker-teu3i76b\n",
       "                        </td>\n",
       "                    </tr>\n",
       "\n",
       "                    \n",
       "\n",
       "                    \n",
       "\n",
       "                </table>\n",
       "            </details>\n",
       "            </div>\n",
       "        </div>\n",
       "        \n",
       "        <div style=\"margin-bottom: 20px;\">\n",
       "            <div style=\"width: 24px; height: 24px; background-color: #DBF5FF; border: 3px solid #4CC9FF; border-radius: 5px; position: absolute;\"> </div>\n",
       "            <div style=\"margin-left: 48px;\">\n",
       "            <details>\n",
       "                <summary>\n",
       "                    <h4 style=\"margin-bottom: 0px; display: inline;\">Worker: 2</h4>\n",
       "                </summary>\n",
       "                <table style=\"width: 100%; text-align: left;\">\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Comm: </strong> tcp://127.0.0.1:8373\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Total threads: </strong> 1\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Dashboard: </strong> <a href=\"http://127.0.0.1:8295/status\" target=\"_blank\">http://127.0.0.1:8295/status</a>\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Memory: </strong> 250.00 GiB\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Nanny: </strong> tcp://127.0.0.1:8445\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\"></td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td colspan=\"2\" style=\"text-align: left;\">\n",
       "                            <strong>Local directory: </strong> /tmp/dask-scratch-space/worker-gv3wqhd8\n",
       "                        </td>\n",
       "                    </tr>\n",
       "\n",
       "                    \n",
       "\n",
       "                    \n",
       "\n",
       "                </table>\n",
       "            </details>\n",
       "            </div>\n",
       "        </div>\n",
       "        \n",
       "        <div style=\"margin-bottom: 20px;\">\n",
       "            <div style=\"width: 24px; height: 24px; background-color: #DBF5FF; border: 3px solid #4CC9FF; border-radius: 5px; position: absolute;\"> </div>\n",
       "            <div style=\"margin-left: 48px;\">\n",
       "            <details>\n",
       "                <summary>\n",
       "                    <h4 style=\"margin-bottom: 0px; display: inline;\">Worker: 3</h4>\n",
       "                </summary>\n",
       "                <table style=\"width: 100%; text-align: left;\">\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Comm: </strong> tcp://127.0.0.1:8081\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Total threads: </strong> 1\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Dashboard: </strong> <a href=\"http://127.0.0.1:8039/status\" target=\"_blank\">http://127.0.0.1:8039/status</a>\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Memory: </strong> 250.00 GiB\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Nanny: </strong> tcp://127.0.0.1:8099\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\"></td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td colspan=\"2\" style=\"text-align: left;\">\n",
       "                            <strong>Local directory: </strong> /tmp/dask-scratch-space/worker-kfo_8zuo\n",
       "                        </td>\n",
       "                    </tr>\n",
       "\n",
       "                    \n",
       "\n",
       "                    \n",
       "\n",
       "                </table>\n",
       "            </details>\n",
       "            </div>\n",
       "        </div>\n",
       "        \n",
       "        <div style=\"margin-bottom: 20px;\">\n",
       "            <div style=\"width: 24px; height: 24px; background-color: #DBF5FF; border: 3px solid #4CC9FF; border-radius: 5px; position: absolute;\"> </div>\n",
       "            <div style=\"margin-left: 48px;\">\n",
       "            <details>\n",
       "                <summary>\n",
       "                    <h4 style=\"margin-bottom: 0px; display: inline;\">Worker: 4</h4>\n",
       "                </summary>\n",
       "                <table style=\"width: 100%; text-align: left;\">\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Comm: </strong> tcp://127.0.0.1:8433\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Total threads: </strong> 1\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Dashboard: </strong> <a href=\"http://127.0.0.1:8181/status\" target=\"_blank\">http://127.0.0.1:8181/status</a>\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Memory: </strong> 250.00 GiB\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Nanny: </strong> tcp://127.0.0.1:8251\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\"></td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td colspan=\"2\" style=\"text-align: left;\">\n",
       "                            <strong>Local directory: </strong> /tmp/dask-scratch-space/worker-1xo60jlx\n",
       "                        </td>\n",
       "                    </tr>\n",
       "\n",
       "                    \n",
       "\n",
       "                    \n",
       "\n",
       "                </table>\n",
       "            </details>\n",
       "            </div>\n",
       "        </div>\n",
       "        \n",
       "        <div style=\"margin-bottom: 20px;\">\n",
       "            <div style=\"width: 24px; height: 24px; background-color: #DBF5FF; border: 3px solid #4CC9FF; border-radius: 5px; position: absolute;\"> </div>\n",
       "            <div style=\"margin-left: 48px;\">\n",
       "            <details>\n",
       "                <summary>\n",
       "                    <h4 style=\"margin-bottom: 0px; display: inline;\">Worker: 5</h4>\n",
       "                </summary>\n",
       "                <table style=\"width: 100%; text-align: left;\">\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Comm: </strong> tcp://127.0.0.1:8285\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Total threads: </strong> 1\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Dashboard: </strong> <a href=\"http://127.0.0.1:8093/status\" target=\"_blank\">http://127.0.0.1:8093/status</a>\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Memory: </strong> 250.00 GiB\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Nanny: </strong> tcp://127.0.0.1:8283\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\"></td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td colspan=\"2\" style=\"text-align: left;\">\n",
       "                            <strong>Local directory: </strong> /tmp/dask-scratch-space/worker-g0owzccw\n",
       "                        </td>\n",
       "                    </tr>\n",
       "\n",
       "                    \n",
       "\n",
       "                    \n",
       "\n",
       "                </table>\n",
       "            </details>\n",
       "            </div>\n",
       "        </div>\n",
       "        \n",
       "        <div style=\"margin-bottom: 20px;\">\n",
       "            <div style=\"width: 24px; height: 24px; background-color: #DBF5FF; border: 3px solid #4CC9FF; border-radius: 5px; position: absolute;\"> </div>\n",
       "            <div style=\"margin-left: 48px;\">\n",
       "            <details>\n",
       "                <summary>\n",
       "                    <h4 style=\"margin-bottom: 0px; display: inline;\">Worker: 6</h4>\n",
       "                </summary>\n",
       "                <table style=\"width: 100%; text-align: left;\">\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Comm: </strong> tcp://127.0.0.1:8289\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Total threads: </strong> 1\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Dashboard: </strong> <a href=\"http://127.0.0.1:8059/status\" target=\"_blank\">http://127.0.0.1:8059/status</a>\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Memory: </strong> 250.00 GiB\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Nanny: </strong> tcp://127.0.0.1:8227\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\"></td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td colspan=\"2\" style=\"text-align: left;\">\n",
       "                            <strong>Local directory: </strong> /tmp/dask-scratch-space/worker-vl_9ezdy\n",
       "                        </td>\n",
       "                    </tr>\n",
       "\n",
       "                    \n",
       "\n",
       "                    \n",
       "\n",
       "                </table>\n",
       "            </details>\n",
       "            </div>\n",
       "        </div>\n",
       "        \n",
       "        <div style=\"margin-bottom: 20px;\">\n",
       "            <div style=\"width: 24px; height: 24px; background-color: #DBF5FF; border: 3px solid #4CC9FF; border-radius: 5px; position: absolute;\"> </div>\n",
       "            <div style=\"margin-left: 48px;\">\n",
       "            <details>\n",
       "                <summary>\n",
       "                    <h4 style=\"margin-bottom: 0px; display: inline;\">Worker: 7</h4>\n",
       "                </summary>\n",
       "                <table style=\"width: 100%; text-align: left;\">\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Comm: </strong> tcp://127.0.0.1:8083\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Total threads: </strong> 1\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Dashboard: </strong> <a href=\"http://127.0.0.1:8055/status\" target=\"_blank\">http://127.0.0.1:8055/status</a>\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Memory: </strong> 250.00 GiB\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Nanny: </strong> tcp://127.0.0.1:8331\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\"></td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td colspan=\"2\" style=\"text-align: left;\">\n",
       "                            <strong>Local directory: </strong> /tmp/dask-scratch-space/worker-cr59y1by\n",
       "                        </td>\n",
       "                    </tr>\n",
       "\n",
       "                    \n",
       "\n",
       "                    \n",
       "\n",
       "                </table>\n",
       "            </details>\n",
       "            </div>\n",
       "        </div>\n",
       "        \n",
       "\n",
       "    </details>\n",
       "</div>\n",
       "\n",
       "        </details>\n",
       "    </div>\n",
       "</div>\n",
       "            </details>\n",
       "        \n",
       "\n",
       "    </div>\n",
       "</div>"
      ],
      "text/plain": [
       "<Client: 'tcp://127.0.0.1:8183' processes=8 threads=8, memory=1.95 TiB>"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%time\n",
    "cluster = LocalCUDACluster(CUDA_VISIBLE_DEVICES=preprocessing_gpus)\n",
    "\n",
    "client = Client(cluster)\n",
    "\n",
    "client.run(set_mem)\n",
    "\n",
    "client"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "b4ac3268-f54b-40c6-9f56-6f0d942c6753",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import rapids_singlecell as rsc\n",
    "import anndata as ad"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "c8541cf5-9a40-4b1a-baea-309bfdc8d15a",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from cupyx.scipy import sparse"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "658cc4c7-b9c0-432e-a447-af470da52d37",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "dask.array.register_chunk_type(sparse.csr_matrix)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5b94f7c9",
   "metadata": {},
   "source": [
    "The data can be loaded directly from an `h5` or `zarr` file into an AnnData object as a `dask-array` using the `read_elem_as_dask` function, enabling lazy loading and parallel computation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "3fc9a948-80ea-457f-8ded-151505a23418",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/ictstr01/home/icb/severin.dicks/git/anndata/src/anndata/_core/anndata.py:1759: UserWarning: Observation names are not unique. To make them unique, call `.obs_names_make_unique`.\n",
      "  utils.warn_names_duplicates(\"obs\")\n"
     ]
    }
   ],
   "source": [
    "from anndata.experimental import read_elem_as_dask\n",
    "\n",
    "import h5py\n",
    "\n",
    "SPARSE_CHUNK_SIZE = 200_000\n",
    "data_pth = \"../2024_gpu_severin.dicks/notebooks/h5/6M_CxG.h5ad\"\n",
    "\n",
    "\n",
    "f = h5py.File(data_pth)\n",
    "X = f[\"X\"]\n",
    "shape = X.attrs[\"shape\"]\n",
    "adata = ad.AnnData(\n",
    "    X = read_elem_as_dask(X, (SPARSE_CHUNK_SIZE, shape[1])),\n",
    "    obs = ad.io.read_elem(f[\"obs\"]),\n",
    "    var = ad.io.read_elem(f[\"var\"]))\n",
    "f.close()\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "d1a4bed2-9e22-4c59-a1be-c851884c77f3",
   "metadata": {},
   "outputs": [],
   "source": [
    "t1 = time.time()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "114352ef-3262-4103-a1a5-f9b33b712656",
   "metadata": {},
   "outputs": [],
   "source": [
    "rsc.get.anndata_to_GPU(adata)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1c4640d3",
   "metadata": {},
   "source": [
    "Persisting the data with `adata.X.persist()` transfers it to GPU memory. \n",
    "Note: Once persisted, the data is no longer managed in an out-of-core manner.\n",
    "\n",
    "All functions that accept dask-arrays operate lazily, performing computations only when `compute()` is called.\n",
    "Use `compute()` or `persist()` to explicitly trigger execution and synchronize results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "2762631c-1b7d-49da-8bd4-3c9d5a0a0ef4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<table>\n",
       "    <tr>\n",
       "        <td>\n",
       "            <table style=\"border-collapse: collapse;\">\n",
       "                <thead>\n",
       "                    <tr>\n",
       "                        <td> </td>\n",
       "                        <th> Array </th>\n",
       "                        <th> Chunk </th>\n",
       "                    </tr>\n",
       "                </thead>\n",
       "                <tbody>\n",
       "                    \n",
       "                    <tr>\n",
       "                        <th> Shape </th>\n",
       "                        <td> (6948864, 19331) </td>\n",
       "                        <td> (200000, 19331) </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <th> Dask graph </th>\n",
       "                        <td colspan=\"2\"> 35 chunks in 1 graph layer </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <th> Data type </th>\n",
       "                        <td colspan=\"2\"> float32 cupyx.scipy.sparse._csr.csr_matrix </td>\n",
       "                    </tr>\n",
       "                </tbody>\n",
       "            </table>\n",
       "        </td>\n",
       "        <td>\n",
       "        <svg width=\"75\" height=\"170\" style=\"stroke:rgb(0,0,0);stroke-width:1\" >\n",
       "\n",
       "  <!-- Horizontal lines -->\n",
       "  <line x1=\"0\" y1=\"0\" x2=\"25\" y2=\"0\" style=\"stroke-width:2\" />\n",
       "  <line x1=\"0\" y1=\"3\" x2=\"25\" y2=\"3\" />\n",
       "  <line x1=\"0\" y1=\"10\" x2=\"25\" y2=\"10\" />\n",
       "  <line x1=\"0\" y1=\"17\" x2=\"25\" y2=\"17\" />\n",
       "  <line x1=\"0\" y1=\"24\" x2=\"25\" y2=\"24\" />\n",
       "  <line x1=\"0\" y1=\"31\" x2=\"25\" y2=\"31\" />\n",
       "  <line x1=\"0\" y1=\"37\" x2=\"25\" y2=\"37\" />\n",
       "  <line x1=\"0\" y1=\"41\" x2=\"25\" y2=\"41\" />\n",
       "  <line x1=\"0\" y1=\"48\" x2=\"25\" y2=\"48\" />\n",
       "  <line x1=\"0\" y1=\"55\" x2=\"25\" y2=\"55\" />\n",
       "  <line x1=\"0\" y1=\"62\" x2=\"25\" y2=\"62\" />\n",
       "  <line x1=\"0\" y1=\"69\" x2=\"25\" y2=\"69\" />\n",
       "  <line x1=\"0\" y1=\"75\" x2=\"25\" y2=\"75\" />\n",
       "  <line x1=\"0\" y1=\"79\" x2=\"25\" y2=\"79\" />\n",
       "  <line x1=\"0\" y1=\"86\" x2=\"25\" y2=\"86\" />\n",
       "  <line x1=\"0\" y1=\"93\" x2=\"25\" y2=\"93\" />\n",
       "  <line x1=\"0\" y1=\"100\" x2=\"25\" y2=\"100\" />\n",
       "  <line x1=\"0\" y1=\"107\" x2=\"25\" y2=\"107\" />\n",
       "  <line x1=\"0\" y1=\"113\" x2=\"25\" y2=\"113\" />\n",
       "  <line x1=\"0\" y1=\"120\" x2=\"25\" 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=\"25\" y1=\"0\" x2=\"25\" y2=\"120\" style=\"stroke-width:2\" />\n",
       "\n",
       "  <!-- Colored Rectangle -->\n",
       "  <polygon points=\"0.0,0.0 25.412616514582485,0.0 25.412616514582485,120.0 0.0,120.0\" style=\"fill:#8B4903A0;stroke-width:0\"/>\n",
       "\n",
       "  <!-- Text -->\n",
       "  <text x=\"12.706308\" y=\"140.000000\" font-size=\"1.0rem\" font-weight=\"100\" text-anchor=\"middle\" >19331</text>\n",
       "  <text x=\"45.412617\" y=\"60.000000\" font-size=\"1.0rem\" font-weight=\"100\" text-anchor=\"middle\" transform=\"rotate(-90,45.412617,60.000000)\">6948864</text>\n",
       "</svg>\n",
       "        </td>\n",
       "    </tr>\n",
       "</table>"
      ],
      "text/plain": [
       "dask.array<csr_matrix, shape=(6948864, 19331), dtype=float32, chunksize=(200000, 19331), chunktype=cupyx.csr_matrix>"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "adata.X = adata.X.persist()\n",
    "adata.X.compute_chunk_sizes()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "870b5367-05c7-4d9a-933e-6ddd2198af19",
   "metadata": {},
   "outputs": [],
   "source": [
    "t2 = time.time()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "4f4f2b47-d363-40c1-b168-cff95b084358",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 735 ms, sys: 178 ms, total: 914 ms\n",
      "Wall time: 10.2 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "rsc.pp.calculate_qc_metrics(adata)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "d7f90ad3-52ac-47da-8840-ccf9db43c6ad",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/ictstr01/home/icb/severin.dicks/git/anndata/src/anndata/_core/index.py:159: PerformanceWarning: Slicing is producing a large chunk. To accept the large\n",
      "chunk and silence this warning, set the option\n",
      "    >>> with dask.config.set(**{'array.slicing.split_large_chunks': False}):\n",
      "    ...     array[indexer]\n",
      "\n",
      "To avoid creating the large chunks, set the option\n",
      "    >>> with dask.config.set(**{'array.slicing.split_large_chunks': True}):\n",
      "    ...     array[indexer]\n",
      "  return a[subset_idx]\n",
      "/ictstr01/home/icb/severin.dicks/git/anndata/src/anndata/_core/anndata.py:1759: UserWarning: Observation names are not unique. To make them unique, call `.obs_names_make_unique`.\n",
      "  utils.warn_names_duplicates(\"obs\")\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "filtered out 132 genes based on n_cells_by_counts\n"
     ]
    }
   ],
   "source": [
    "adata = adata[(adata.obs[\"n_genes_by_counts\"]<=10000) \n",
    "            & (adata.obs[\"n_genes_by_counts\"]>=200)].copy()\n",
    "rsc.pp.filter_genes(adata, min_count=10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "299580e9-44a9-41c5-bfe2-570ddd61ce7d",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/icb/severin.dicks/miniconda3/envs/rapids-24.06/lib/python3.11/site-packages/distributed/client.py:3245: UserWarning: Sending large graph of size 53.08 MiB.\n",
      "This may cause some slowdown.\n",
      "Consider scattering data ahead of time and using futures.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<table>\n",
       "    <tr>\n",
       "        <td>\n",
       "            <table style=\"border-collapse: collapse;\">\n",
       "                <thead>\n",
       "                    <tr>\n",
       "                        <td> </td>\n",
       "                        <th> Array </th>\n",
       "                        <th> Chunk </th>\n",
       "                    </tr>\n",
       "                </thead>\n",
       "                <tbody>\n",
       "                    \n",
       "                    <tr>\n",
       "                        <th> Shape </th>\n",
       "                        <td> (6937990, 19199) </td>\n",
       "                        <td> (199767, 19199) </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <th> Dask graph </th>\n",
       "                        <td colspan=\"2\"> 35 chunks in 1 graph layer </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <th> Data type </th>\n",
       "                        <td colspan=\"2\"> float32 cupyx.scipy.sparse._csr.csr_matrix </td>\n",
       "                    </tr>\n",
       "                </tbody>\n",
       "            </table>\n",
       "        </td>\n",
       "        <td>\n",
       "        <svg width=\"75\" height=\"170\" style=\"stroke:rgb(0,0,0);stroke-width:1\" >\n",
       "\n",
       "  <!-- Horizontal lines -->\n",
       "  <line x1=\"0\" y1=\"0\" x2=\"25\" y2=\"0\" style=\"stroke-width:2\" />\n",
       "  <line x1=\"0\" y1=\"3\" x2=\"25\" y2=\"3\" />\n",
       "  <line x1=\"0\" y1=\"10\" x2=\"25\" y2=\"10\" />\n",
       "  <line x1=\"0\" y1=\"17\" x2=\"25\" y2=\"17\" />\n",
       "  <line x1=\"0\" y1=\"24\" x2=\"25\" y2=\"24\" />\n",
       "  <line x1=\"0\" y1=\"31\" x2=\"25\" y2=\"31\" />\n",
       "  <line x1=\"0\" y1=\"38\" x2=\"25\" y2=\"38\" />\n",
       "  <line x1=\"0\" y1=\"41\" x2=\"25\" y2=\"41\" />\n",
       "  <line x1=\"0\" y1=\"48\" x2=\"25\" y2=\"48\" />\n",
       "  <line x1=\"0\" y1=\"55\" x2=\"25\" y2=\"55\" />\n",
       "  <line x1=\"0\" y1=\"62\" x2=\"25\" y2=\"62\" />\n",
       "  <line x1=\"0\" y1=\"69\" x2=\"25\" y2=\"69\" />\n",
       "  <line x1=\"0\" y1=\"75\" x2=\"25\" y2=\"75\" />\n",
       "  <line x1=\"0\" y1=\"79\" x2=\"25\" y2=\"79\" />\n",
       "  <line x1=\"0\" y1=\"86\" x2=\"25\" y2=\"86\" />\n",
       "  <line x1=\"0\" y1=\"93\" x2=\"25\" y2=\"93\" />\n",
       "  <line x1=\"0\" y1=\"100\" x2=\"25\" y2=\"100\" />\n",
       "  <line x1=\"0\" y1=\"107\" x2=\"25\" y2=\"107\" />\n",
       "  <line x1=\"0\" y1=\"113\" x2=\"25\" y2=\"113\" />\n",
       "  <line x1=\"0\" y1=\"120\" x2=\"25\" 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=\"25\" y1=\"0\" x2=\"25\" y2=\"120\" style=\"stroke-width:2\" />\n",
       "\n",
       "  <!-- Colored Rectangle -->\n",
       "  <polygon points=\"0.0,0.0 25.412616514582485,0.0 25.412616514582485,120.0 0.0,120.0\" style=\"fill:#8B4903A0;stroke-width:0\"/>\n",
       "\n",
       "  <!-- Text -->\n",
       "  <text x=\"12.706308\" y=\"140.000000\" font-size=\"1.0rem\" font-weight=\"100\" text-anchor=\"middle\" >19199</text>\n",
       "  <text x=\"45.412617\" y=\"60.000000\" font-size=\"1.0rem\" font-weight=\"100\" text-anchor=\"middle\" transform=\"rotate(-90,45.412617,60.000000)\">6937990</text>\n",
       "</svg>\n",
       "        </td>\n",
       "    </tr>\n",
       "</table>"
      ],
      "text/plain": [
       "dask.array<getitem, shape=(6937990, 19199), dtype=float32, chunksize=(199767, 19199), chunktype=cupyx.csr_matrix>"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "adata.X = adata.X.persist()\n",
    "adata.X.compute_chunk_sizes()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "fdbe6f8f-9086-474b-9583-d85b9190593f",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 11.1 ms, sys: 3.65 ms, total: 14.8 ms\n",
      "Wall time: 137 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "rsc.pp.normalize_total(adata,target_sum = 10000)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "3f700e76-94ee-4f84-9db2-74cd9aaab999",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 651 μs, sys: 848 μs, total: 1.5 ms\n",
      "Wall time: 1.34 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "rsc.pp.log1p(adata)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "fcda2919-95da-4354-9386-9507347ccdc1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<table>\n",
       "    <tr>\n",
       "        <td>\n",
       "            <table style=\"border-collapse: collapse;\">\n",
       "                <thead>\n",
       "                    <tr>\n",
       "                        <td> </td>\n",
       "                        <th> Array </th>\n",
       "                        <th> Chunk </th>\n",
       "                    </tr>\n",
       "                </thead>\n",
       "                <tbody>\n",
       "                    \n",
       "                    <tr>\n",
       "                        <th> Shape </th>\n",
       "                        <td> (6937990, 19199) </td>\n",
       "                        <td> (199767, 19199) </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <th> Dask graph </th>\n",
       "                        <td colspan=\"2\"> 35 chunks in 1 graph layer </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <th> Data type </th>\n",
       "                        <td colspan=\"2\"> float32 cupyx.scipy.sparse._csr.csr_matrix </td>\n",
       "                    </tr>\n",
       "                </tbody>\n",
       "            </table>\n",
       "        </td>\n",
       "        <td>\n",
       "        <svg width=\"75\" height=\"170\" style=\"stroke:rgb(0,0,0);stroke-width:1\" >\n",
       "\n",
       "  <!-- Horizontal lines -->\n",
       "  <line x1=\"0\" y1=\"0\" x2=\"25\" y2=\"0\" style=\"stroke-width:2\" />\n",
       "  <line x1=\"0\" y1=\"3\" x2=\"25\" y2=\"3\" />\n",
       "  <line x1=\"0\" y1=\"10\" x2=\"25\" y2=\"10\" />\n",
       "  <line x1=\"0\" y1=\"17\" x2=\"25\" y2=\"17\" />\n",
       "  <line x1=\"0\" y1=\"24\" x2=\"25\" y2=\"24\" />\n",
       "  <line x1=\"0\" y1=\"31\" x2=\"25\" y2=\"31\" />\n",
       "  <line x1=\"0\" y1=\"38\" x2=\"25\" y2=\"38\" />\n",
       "  <line x1=\"0\" y1=\"41\" x2=\"25\" y2=\"41\" />\n",
       "  <line x1=\"0\" y1=\"48\" x2=\"25\" y2=\"48\" />\n",
       "  <line x1=\"0\" y1=\"55\" x2=\"25\" y2=\"55\" />\n",
       "  <line x1=\"0\" y1=\"62\" x2=\"25\" y2=\"62\" />\n",
       "  <line x1=\"0\" y1=\"69\" x2=\"25\" y2=\"69\" />\n",
       "  <line x1=\"0\" y1=\"75\" x2=\"25\" y2=\"75\" />\n",
       "  <line x1=\"0\" y1=\"79\" x2=\"25\" y2=\"79\" />\n",
       "  <line x1=\"0\" y1=\"86\" x2=\"25\" y2=\"86\" />\n",
       "  <line x1=\"0\" y1=\"93\" x2=\"25\" y2=\"93\" />\n",
       "  <line x1=\"0\" y1=\"100\" x2=\"25\" y2=\"100\" />\n",
       "  <line x1=\"0\" y1=\"107\" x2=\"25\" y2=\"107\" />\n",
       "  <line x1=\"0\" y1=\"113\" x2=\"25\" y2=\"113\" />\n",
       "  <line x1=\"0\" y1=\"120\" x2=\"25\" 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=\"25\" y1=\"0\" x2=\"25\" y2=\"120\" style=\"stroke-width:2\" />\n",
       "\n",
       "  <!-- Colored Rectangle -->\n",
       "  <polygon points=\"0.0,0.0 25.412616514582485,0.0 25.412616514582485,120.0 0.0,120.0\" style=\"fill:#8B4903A0;stroke-width:0\"/>\n",
       "\n",
       "  <!-- Text -->\n",
       "  <text x=\"12.706308\" y=\"140.000000\" font-size=\"1.0rem\" font-weight=\"100\" text-anchor=\"middle\" >19199</text>\n",
       "  <text x=\"45.412617\" y=\"60.000000\" font-size=\"1.0rem\" font-weight=\"100\" text-anchor=\"middle\" transform=\"rotate(-90,45.412617,60.000000)\">6937990</text>\n",
       "</svg>\n",
       "        </td>\n",
       "    </tr>\n",
       "</table>"
      ],
      "text/plain": [
       "dask.array<lambda, shape=(6937990, 19199), dtype=float32, chunksize=(199767, 19199), chunktype=cupyx.csr_matrix>"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "adata.X = adata.X.persist()\n",
    "adata.X.compute_chunk_sizes()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "bb84f792-2b06-4bfa-8cca-3b4eaec042ea",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 369 ms, sys: 112 ms, total: 482 ms\n",
      "Wall time: 4.75 s\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/ictstr01/home/icb/severin.dicks/git/anndata/src/anndata/_core/index.py:159: PerformanceWarning: Slicing is producing a large chunk. To accept the large\n",
      "chunk and silence this warning, set the option\n",
      "    >>> with dask.config.set(**{'array.slicing.split_large_chunks': False}):\n",
      "    ...     array[indexer]\n",
      "\n",
      "To avoid creating the large chunks, set the option\n",
      "    >>> with dask.config.set(**{'array.slicing.split_large_chunks': True}):\n",
      "    ...     array[indexer]\n",
      "  return a[subset_idx]\n",
      "/ictstr01/home/icb/severin.dicks/git/anndata/src/anndata/_core/anndata.py:1759: UserWarning: Observation names are not unique. To make them unique, call `.obs_names_make_unique`.\n",
      "  utils.warn_names_duplicates(\"obs\")\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "rsc.pp.highly_variable_genes(adata,n_top_genes=5000, flavor=\"cell_ranger\")\n",
    "adata = adata[:,adata.var.highly_variable].copy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "fa17722a-ffae-45cd-bd65-a29d4abf8bdc",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(6937990, 5000)"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "adata.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "73f66c1a-8892-44ba-9a15-c88d921548dd",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<table>\n",
       "    <tr>\n",
       "        <td>\n",
       "            <table style=\"border-collapse: collapse;\">\n",
       "                <thead>\n",
       "                    <tr>\n",
       "                        <td> </td>\n",
       "                        <th> Array </th>\n",
       "                        <th> Chunk </th>\n",
       "                    </tr>\n",
       "                </thead>\n",
       "                <tbody>\n",
       "                    \n",
       "                    <tr>\n",
       "                        <th> Shape </th>\n",
       "                        <td> (6937990, 5000) </td>\n",
       "                        <td> (867248, 5000) </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <th> Dask graph </th>\n",
       "                        <td colspan=\"2\"> 9 chunks in 1 graph layer </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <th> Data type </th>\n",
       "                        <td colspan=\"2\"> float32 cupyx.scipy.sparse._csr.csr_matrix </td>\n",
       "                    </tr>\n",
       "                </tbody>\n",
       "            </table>\n",
       "        </td>\n",
       "        <td>\n",
       "        <svg width=\"75\" height=\"170\" style=\"stroke:rgb(0,0,0);stroke-width:1\" >\n",
       "\n",
       "  <!-- Horizontal lines -->\n",
       "  <line x1=\"0\" y1=\"0\" x2=\"25\" y2=\"0\" style=\"stroke-width:2\" />\n",
       "  <line x1=\"0\" y1=\"14\" x2=\"25\" y2=\"14\" />\n",
       "  <line x1=\"0\" y1=\"29\" x2=\"25\" y2=\"29\" />\n",
       "  <line x1=\"0\" y1=\"44\" x2=\"25\" y2=\"44\" />\n",
       "  <line x1=\"0\" y1=\"59\" x2=\"25\" y2=\"59\" />\n",
       "  <line x1=\"0\" y1=\"74\" x2=\"25\" y2=\"74\" />\n",
       "  <line x1=\"0\" y1=\"89\" x2=\"25\" y2=\"89\" />\n",
       "  <line x1=\"0\" y1=\"104\" x2=\"25\" y2=\"104\" />\n",
       "  <line x1=\"0\" y1=\"119\" x2=\"25\" y2=\"119\" />\n",
       "  <line x1=\"0\" y1=\"120\" x2=\"25\" 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=\"25\" y1=\"0\" x2=\"25\" y2=\"120\" style=\"stroke-width:2\" />\n",
       "\n",
       "  <!-- Colored Rectangle -->\n",
       "  <polygon points=\"0.0,0.0 25.412616514582485,0.0 25.412616514582485,120.0 0.0,120.0\" style=\"fill:#ECB172A0;stroke-width:0\"/>\n",
       "\n",
       "  <!-- Text -->\n",
       "  <text x=\"12.706308\" y=\"140.000000\" font-size=\"1.0rem\" font-weight=\"100\" text-anchor=\"middle\" >5000</text>\n",
       "  <text x=\"45.412617\" y=\"60.000000\" font-size=\"1.0rem\" font-weight=\"100\" text-anchor=\"middle\" transform=\"rotate(-90,45.412617,60.000000)\">6937990</text>\n",
       "</svg>\n",
       "        </td>\n",
       "    </tr>\n",
       "</table>"
      ],
      "text/plain": [
       "dask.array<rechunk-merge, shape=(6937990, 5000), dtype=float32, chunksize=(867248, 5000), chunktype=cupyx.csr_matrix>"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "n_rows = adata.shape[0]\n",
    "n_cols = adata.shape[1]\n",
    "cols_per_worker = int(n_rows / 8)\n",
    "adata.X = adata.X.rechunk((cols_per_worker, n_cols)).persist()\n",
    "\n",
    "adata.X.compute_chunk_sizes()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "ab10f1e1-2fe5-4f5d-aa9d-1357d2d55a49",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/icb/severin.dicks/miniconda3/envs/rapids-24.06/lib/python3.11/site-packages/distributed/client.py:3245: UserWarning: Sending large graph of size 26.49 MiB.\n",
      "This may cause some slowdown.\n",
      "Consider scattering data ahead of time and using futures.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 494 ms, sys: 96.4 ms, total: 591 ms\n",
      "Wall time: 4.97 s\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<table>\n",
       "    <tr>\n",
       "        <td>\n",
       "            <table style=\"border-collapse: collapse;\">\n",
       "                <thead>\n",
       "                    <tr>\n",
       "                        <td> </td>\n",
       "                        <th> Array </th>\n",
       "                        <th> Chunk </th>\n",
       "                    </tr>\n",
       "                </thead>\n",
       "                <tbody>\n",
       "                    \n",
       "                    <tr>\n",
       "                        <th> Shape </th>\n",
       "                        <td> (6937990, 5000) </td>\n",
       "                        <td> (867248, 5000) </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <th> Dask graph </th>\n",
       "                        <td colspan=\"2\"> 9 chunks in 1 graph layer </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <th> Data type </th>\n",
       "                        <td colspan=\"2\"> float32 cupyx.scipy.sparse._csr.csr_matrix </td>\n",
       "                    </tr>\n",
       "                </tbody>\n",
       "            </table>\n",
       "        </td>\n",
       "        <td>\n",
       "        <svg width=\"75\" height=\"170\" style=\"stroke:rgb(0,0,0);stroke-width:1\" >\n",
       "\n",
       "  <!-- Horizontal lines -->\n",
       "  <line x1=\"0\" y1=\"0\" x2=\"25\" y2=\"0\" style=\"stroke-width:2\" />\n",
       "  <line x1=\"0\" y1=\"14\" x2=\"25\" y2=\"14\" />\n",
       "  <line x1=\"0\" y1=\"29\" x2=\"25\" y2=\"29\" />\n",
       "  <line x1=\"0\" y1=\"44\" x2=\"25\" y2=\"44\" />\n",
       "  <line x1=\"0\" y1=\"59\" x2=\"25\" y2=\"59\" />\n",
       "  <line x1=\"0\" y1=\"74\" x2=\"25\" y2=\"74\" />\n",
       "  <line x1=\"0\" y1=\"89\" x2=\"25\" y2=\"89\" />\n",
       "  <line x1=\"0\" y1=\"104\" x2=\"25\" y2=\"104\" />\n",
       "  <line x1=\"0\" y1=\"119\" x2=\"25\" y2=\"119\" />\n",
       "  <line x1=\"0\" y1=\"120\" x2=\"25\" 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=\"25\" y1=\"0\" x2=\"25\" y2=\"120\" style=\"stroke-width:2\" />\n",
       "\n",
       "  <!-- Colored Rectangle -->\n",
       "  <polygon points=\"0.0,0.0 25.412616514582485,0.0 25.412616514582485,120.0 0.0,120.0\" style=\"fill:#ECB172A0;stroke-width:0\"/>\n",
       "\n",
       "  <!-- Text -->\n",
       "  <text x=\"12.706308\" y=\"140.000000\" font-size=\"1.0rem\" font-weight=\"100\" text-anchor=\"middle\" >5000</text>\n",
       "  <text x=\"45.412617\" y=\"60.000000\" font-size=\"1.0rem\" font-weight=\"100\" text-anchor=\"middle\" transform=\"rotate(-90,45.412617,60.000000)\">6937990</text>\n",
       "</svg>\n",
       "        </td>\n",
       "    </tr>\n",
       "</table>"
      ],
      "text/plain": [
       "dask.array<scale_kernel_csr, shape=(6937990, 5000), dtype=float32, chunksize=(867248, 5000), chunktype=cupyx.csr_matrix>"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%time\n",
    "rsc.pp.scale(adata, zero_center= False)\n",
    "adata.X = adata.X.persist()\n",
    "adata.X.compute_chunk_sizes()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "b525a19c-6ed6-4e2c-a249-c4fc536862bb",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 1.36 s, sys: 349 ms, total: 1.71 s\n",
      "Wall time: 17.3 s\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<table>\n",
       "    <tr>\n",
       "        <td>\n",
       "            <table style=\"border-collapse: collapse;\">\n",
       "                <thead>\n",
       "                    <tr>\n",
       "                        <td> </td>\n",
       "                        <th> Array </th>\n",
       "                        <th> Chunk </th>\n",
       "                    </tr>\n",
       "                </thead>\n",
       "                <tbody>\n",
       "                    \n",
       "                    <tr>\n",
       "                        <th> Bytes </th>\n",
       "                        <td> 2.58 GiB </td>\n",
       "                        <td> 330.83 MiB </td>\n",
       "                    </tr>\n",
       "                    \n",
       "                    <tr>\n",
       "                        <th> Shape </th>\n",
       "                        <td> (6937990, 100) </td>\n",
       "                        <td> (867248, 100) </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <th> Dask graph </th>\n",
       "                        <td colspan=\"2\"> 9 chunks in 1 graph layer </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <th> Data type </th>\n",
       "                        <td colspan=\"2\"> float32 cupy.ndarray </td>\n",
       "                    </tr>\n",
       "                </tbody>\n",
       "            </table>\n",
       "        </td>\n",
       "        <td>\n",
       "        <svg width=\"75\" height=\"170\" style=\"stroke:rgb(0,0,0);stroke-width:1\" >\n",
       "\n",
       "  <!-- Horizontal lines -->\n",
       "  <line x1=\"0\" y1=\"0\" x2=\"25\" y2=\"0\" style=\"stroke-width:2\" />\n",
       "  <line x1=\"0\" y1=\"14\" x2=\"25\" y2=\"14\" />\n",
       "  <line x1=\"0\" y1=\"29\" x2=\"25\" y2=\"29\" />\n",
       "  <line x1=\"0\" y1=\"44\" x2=\"25\" y2=\"44\" />\n",
       "  <line x1=\"0\" y1=\"59\" x2=\"25\" y2=\"59\" />\n",
       "  <line x1=\"0\" y1=\"74\" x2=\"25\" y2=\"74\" />\n",
       "  <line x1=\"0\" y1=\"89\" x2=\"25\" y2=\"89\" />\n",
       "  <line x1=\"0\" y1=\"104\" x2=\"25\" y2=\"104\" />\n",
       "  <line x1=\"0\" y1=\"119\" x2=\"25\" y2=\"119\" />\n",
       "  <line x1=\"0\" y1=\"120\" x2=\"25\" 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=\"25\" y1=\"0\" x2=\"25\" y2=\"120\" style=\"stroke-width:2\" />\n",
       "\n",
       "  <!-- Colored Rectangle -->\n",
       "  <polygon points=\"0.0,0.0 25.41261651458249,0.0 25.41261651458249,120.0 0.0,120.0\" style=\"fill:#ECB172A0;stroke-width:0\"/>\n",
       "\n",
       "  <!-- Text -->\n",
       "  <text x=\"12.706308\" y=\"140.000000\" font-size=\"1.0rem\" font-weight=\"100\" text-anchor=\"middle\" >100</text>\n",
       "  <text x=\"45.412617\" y=\"60.000000\" font-size=\"1.0rem\" font-weight=\"100\" text-anchor=\"middle\" transform=\"rotate(-90,45.412617,60.000000)\">6937990</text>\n",
       "</svg>\n",
       "        </td>\n",
       "    </tr>\n",
       "</table>"
      ],
      "text/plain": [
       "dask.array<_transform, shape=(6937990, 100), dtype=float32, chunksize=(867248, 100), chunktype=cupy.ndarray>"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%time\n",
    "rsc.pp.pca(adata, n_comps = 100,mask_var=None)\n",
    "adata.obsm[\"X_pca\"]=adata.obsm[\"X_pca\"].persist()\n",
    "adata.obsm[\"X_pca\"].compute_chunk_sizes()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "c22c2937-9cc0-47e7-b02b-c0b7de6397ad",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "60.341403007507324\n",
      "6.557059117158254\n"
     ]
    }
   ],
   "source": [
    "print((time.time()-t2))\n",
    "print((time.time()-t1)/60)"
   ]
  }
 ],
 "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.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}