8.0 Multi-GPU with Dask.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "4dae51fa-32b1-499a-a21d-639d1da8840b",
   "metadata": {},
   "source": [
    "# Multi-GPU with Dask\n",
    "\n",
    "Leveraging GPUs to accelerate your workloads can result in orders of magnitude increase in performance, but once your workload is fully utilizing the device you will start to reach a new ceiling of performance.\n",
    "\n",
    "This is where multi-GPU and multi-node workloads come in. It is possible to use many GPUs together and see another step change in performance.\n",
    "\n",
    "Before we dive into multi-GPU workloads I do want to caution that distributed computation can increase the complexity of your code. The tools discussed in this chapter do everything they can to ease the burden of distributed computing but we should be sure to check that we have squeezed out every last drop of performance on a single GPU before we start scaling out."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!git clone https://github.com/rapidsai/rapidsai-csp-utils.git\n",
    "!python rapidsai-csp-utils/colab/pip-install.py\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8f39055d-e211-4c7f-8b8b-75e291880946",
   "metadata": {},
   "source": [
    "## Dask\n",
    "\n",
    "[Dask](https://dask.org) is a Python library for scaling out your Python code. At it's core Dask takes your Python code and converts it into a computation graph of function calls, inputs and outputs. It then has a selection of schedulers that it can use to execute through this graph in parallel. Here we are going to focus on Dask's distributed scheduler."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1cd0c93e-e58a-4d15-919a-88e62b5be814",
   "metadata": {},
   "outputs": [],
   "source": [
    "from dask.distributed import Client\n",
    "\n",
    "client = Client()\n",
    "client\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "28ced429-50cb-487d-b0d0-78ad503f2d9a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Submit a function to be executed on the Dask cluster\n",
    "f = client.submit(lambda: 10 + 1)\n",
    "f.result()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "57b604b2-8d8b-4a29-9674-d291b5ead139",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Use a high level collection API to distribute familiar work on the cluster\n",
    "import dask.array as da\n",
    "arr = da.random.random((1000, 1000), chunks=(100, 100))\n",
    "arr.mean().compute()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5cb28fc9-1e26-4743-8a38-0adf6a045b8e",
   "metadata": {},
   "source": [
    "Dask doesn't hugely care about what your code it doing, it just attempts to run through the graph as quickly as possible with its pool of workers. Because we've done all of our GPU computation in Python Dask can also distribute our GPU code too."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ac465c9c-13ee-4ff9-8aad-0d35592c4f32",
   "metadata": {},
   "outputs": [],
   "source": [
    "client.close()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f7ff1b78-225f-4277-af69-ffd65d9093c6",
   "metadata": {},
   "source": [
    "### Distributed clusters\n",
    "\n",
    "In order for Dask to distribute a graph onto many machines it needs a scheduler process and a number of worker processes. We can start these manually, either via the CLI commands `dask-scheduler` and `dask-worker` or using any number of Dask's cluster managers."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aeeb108b-5e07-4786-a78e-97f51dc4685b",
   "metadata": {},
   "source": [
    "#### Cluster managers\n",
    "\n",
    "Dask can be run in any number of compute environments as long as you have a Python environment, network connectivity and can start the scheduler and worker processes.\n",
    "\n",
    "To make creating a Dask cluster a consistent experience there are a number of cluster manager classes that you can import and instantiate which will construct a cluster for you.\n",
    "\n",
    "The first that most folks interact with is `LocalCluster`. When you create an instance of this class it will inspect the CPU and memory resources available on the local computer and create subprocesses for the scheduler and an appropriate number of workers automatically."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0de43e9b-3931-41f2-b216-1c9b6502f5b2",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from dask.distributed import LocalCluster\n",
    "\n",
    "cluster = LocalCluster()\n",
    "cluster\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fabd6943-16f1-4031-9d2c-cb69835d48a1",
   "metadata": {},
   "source": [
    "This is great for trying Dask out and using it to leverage all of the CPU cores available on your local machine.\n",
    "\n",
    "Once you are ready to break beyond the confines of your computer there are cluster managers for HPC platforms like SLURM, PBS and SGE. Cluster managers for Kubernetes, Hadoop and public cloud providers including Amazon Web Services, Microsoft Azure and Google Cloud Platform.\n",
    "\n",
    "The most popular cluster manager we see folks using is the `SSHCluster` class which opens secure shell connections to other machines on your network and starts Dask processes on them. This can be great for remotely leveraging servers or even other desktops that aren't being used.\n",
    "\n",
    "```python\n",
    "from dask.distributed import SSHCluster\n",
    "\n",
    "cluster = SSHCluster([\n",
    "    \"localhost\",  # Hostname to start the scheduler on\n",
    "    \"10.0.0.2\",   # Hostname to start the first worker on\n",
    "    \"10.0.0.3\",   # Hostname to start the second worker on\n",
    "    ...           # etc\n",
    "])\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0e0234e4-0e2f-4efc-9446-242b98697cb9",
   "metadata": {},
   "outputs": [],
   "source": [
    "cluster.close()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "45e048a5-dba2-45a4-822a-316c60f6f5fb",
   "metadata": {},
   "source": [
    "#### Dask CUDA\n",
    "\n",
    "When it comes to using GPUs with Dask there are a few things we need to bear in mind. Each Dask worker needs to have exactly one GPU, so if your machine has multiple GPUs you'll need one worker per device. There are also a few other things that need to be done in order for a Dask worker to successfully be able to leverage a GPU. To simplify this for the user you can use the tools found in the Python package `dask-cuda`.\n",
    "\n",
    "The Dask CUDA package has a cluster manager called `LocalCUDACluster` and an alternative worker CLI command called `dask-cuda-worker`. Both of these inspect your hardware and start one worker per GPU and correctly configure each workers to only use their allocated device."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8611f07b-ae1d-4257-8946-ab9d6802ee09",
   "metadata": {},
   "outputs": [],
   "source": [
    "from dask_cuda import LocalCUDACluster\n",
    "\n",
    "cluster = LocalCUDACluster()\n",
    "cluster\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0bde8f10-3045-47c6-a7ba-aaecfe0800e0",
   "metadata": {},
   "outputs": [],
   "source": [
    "client = Client(cluster)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "03b71eae-af05-4fcc-92d4-eae3966ef263",
   "metadata": {},
   "source": [
    "It is also possible to configure the other cluster managers that leverage HPC and Cloud capabilities to use the Dask CUDA worker instead of the regular worker.\n",
    "\n",
    "Once we have a Dask cluster with GPU workers we could manually submit some CUDA kernels written with Numba to be executed on those GPUS."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9fae98b4-4aaf-4a99-853e-755a64c21850",
   "metadata": {},
   "outputs": [],
   "source": [
    "from numba import cuda\n",
    "\n",
    "@cuda.jit\n",
    "def some_kernel():\n",
    "    i = 0\n",
    "    while i < 1_000_000:\n",
    "        i += 1\n",
    "\n",
    "f = client.submit(some_kernel[1024*1024, 1024])\n",
    "f\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "16435a6b-9cbc-4ee4-b398-88c88122509e",
   "metadata": {},
   "source": [
    "### High level collections\n",
    "\n",
    "Thankfully we don't have to do everything manually with Dask. Dask has a concept of high-level collections which implement the API of a popular Python package but chunk/partition up data structures and tasks so that they can be run on a Dask cluster. Commonly folks use `dask.array` which follows the NumPy API, `dask.dataframe` which follows the Pandas API and `dask.ml` which follows the Scikit-Learn API.\n",
    "\n",
    "This approach may sound familiar, we've already seen RAPIDS libraries that mimic the APIs of these libraries to provide accelerated computation. Dask does the same but for distributed computation. One of the benefits of this approach is that we can combine them to get distributed and accelerated computation of the tools we already know and love.\n",
    "\n",
    "When `dask.dataframe` create a DataFrame it constructs a task graph that consists of many smaller Pandas DataFrames. Then operations like taking the mean of a series will first be performed on each Pandas DataFrame before the results are aggregated to get the overall mean. But Dask is not limited to using Pandas in its DataFrame collection, it can also leverage other libraries that follow the Pandas API like cuDF.\n",
    "\n",
    "cuDF comes with a useful helper library for constructing Dask DataFrames made up of cuDF DataFrames and we can load our data and perform operations just as we have seen before."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8ac962fb-8a46-4d79-b4cd-81c1cae3ada8",
   "metadata": {},
   "outputs": [],
   "source": [
    "import dask\n",
    "import cudf\n",
    "import dask_cudf\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8d46ac77-3e1a-473f-81c2-fa0aa54f09ed",
   "metadata": {},
   "outputs": [],
   "source": [
    "@dask.delayed\n",
    "def gen_partition():\n",
    "    return cudf.datasets.timeseries()\n",
    "\n",
    "gddf = dask_cudf.from_delayed([gen_partition() for i in range(30)])\n",
    "gddf\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "477b02f9-add8-441d-a59a-d4cb59114413",
   "metadata": {},
   "outputs": [],
   "source": [
    "gddf.head()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4fd9670c-40b3-4e02-98b0-bb8b1bb55cd1",
   "metadata": {},
   "outputs": [],
   "source": [
    "len(gddf)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2d253f08-2144-47e8-9863-e3193ad92c10",
   "metadata": {},
   "outputs": [],
   "source": [
    "gddf.groupby(\"name\").x.mean().compute()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "800c19bc-662d-4336-82c7-7f59225ac992",
   "metadata": {},
   "source": [
    "But now our DataFrame is distributed across all of our GPUs and computations can leverage the performance of all of our hardware."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "96a35fe7-4356-4eb9-9eb4-b4e328bcfb2b",
   "metadata": {},
   "source": [
    "### Communication\n",
    "\n",
    "In Chapter 1 when we were exploring Numba CUDA we saw that there were penalties when it comes to moving data from CPU memory into the GPU memory. The same applies with moving data between GPU memories and between GPUs on different machines.\n",
    "\n",
    "By default Dask uses a custom TCP protocol for communication between workers. This means that any memory transfer from one GPU to another has to make its way back up the PCI-E lanes to the CPU, into the operating system's network stack to be routed to it's destination. If the GPU it is sending it to is in the same machine it will head back down the PCI-E lanes and into the GPU. If it is located on another machine it will make its way our via the IP network, most likely via an ethernet connection.\n",
    "\n",
    "In the case where our two GPUs are sat next to each other on the motherboard this is very wasteful. They could even be connected to each other directly via NVLINK, or at least connected to the same PCI-E switch on the motherboard. Routing every transfer via the CPU is wasteful, and that's where UCX comes in."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d189139b-bed9-4842-b293-5eb5c0a5b905",
   "metadata": {},
   "source": [
    "#### UCX\n",
    "\n",
    "[UCX](https://openucx.org/) is a network protocol which can inspect the topology of the systems and find an optimal route via accelerated hardware. If two GPUs are connected via NVLINK then UCX will use that to transfer data, if they are connected to the same PCI-E switch that is the next best bet. If the GPUs are on two separate machines but those machines have Infiniband network cards then UCX can leverage RDMA over Infiniband to also transfer data directly between the GPUs.\n",
    "\n",
    "UCX will do everything in its power to transfer data as directly and performantly as possible between two locations before ultimately falling back to TCP."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "db5e93cf-4d7c-4ee3-8569-d669e2e4356e",
   "metadata": {},
   "source": [
    "#### Dask communication protocols\n",
    "\n",
    "Dask supports alternative communication protocols which can be configured by the user. This includes UCX which we can leverage for more performance but also other protocols like websockets which may be more flexible in modern system architectures due to being easier to proxy.\n",
    "\n",
    "If we use UCX with our GPU workers and have accelerated networking hardware like NVLINK or Infiniband then we can see much reduced memory transfer times between our GPU workers."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d088b413-88cf-4c16-af56-a9fd0bf9e7c4",
   "metadata": {},
   "source": [
    "### Resource annotations\n",
    "\n",
    "The last topic I wanted to cover with Dask and GPUs is annotations. Dask has a capability where each task in the task graph can be annotated with requirements that a worker needs to have in order to be able to run it.\n",
    "\n",
    "When we start up a worker we can also add our resource labels so that the scheduler can place appropriate tasks on appropriate workers. This feature is most powerful when our workers are a mixed bag of configurations.\n",
    "\n",
    "\n",
    "```console\n",
    "$ dask-cuda-worker scheduler:8786 --resources \"GPU=2\"\n",
    "```\n",
    "\n",
    "There may be steps in your task graph where memory usage increases heavily during an intermediate calculation. It can be helpful to steer these tasks towards workers that have more memory than the rest.\n",
    "\n",
    "We can also use this for GPU work if not all of our workers have GPUs. It would be reasonable to have a number of regular Dask workers that take on most of your tasks but also have a couple of GPU workers to run steps that have been optimised to run on the GPU.\n",
    "\n",
    "This would likely be most useful if you have an existing workload that leverages Dask and you want to experiment with GPUs. You could add another worker that has a GPU, choose some tasks in your workflow to optimize with Numba and annotate those tasks to only run on your GPU worker.\n",
    "\n",
    "```python\n",
    "foo = client.submit(some_non_gpu_function)\n",
    "\n",
    "with dask.annotate(resources={'GPU': 1}):\n",
    "    bar = client.submit(a_gpu_function, foo)\n",
    "    \n",
    "baz = client.submit(another_non_gpu_function, bar)\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c07a58b1-48c3-4ae7-b31c-35b40cdd41bc",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}