Launch a hybrid CPU and GPU Dask cluster and use annotations to pin work to specific hardware

notebook.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "596b6423-6978-4f8c-9181-6264e5326a61",
   "metadata": {},
   "outputs": [],
   "source": [
    "import dask\n",
    "from dask.system import CPU_COUNT\n",
    "from dask.distributed import Client, Nanny, wait\n",
    "from dask_cuda import LocalCUDACluster"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7a942263-9233-4fb1-811a-d924f564ef98",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Create a GPU cluster\n",
    "cluster = await LocalCUDACluster(asynchronous=True, resources={\"GPU\": 1})\n",
    "\n",
    "# Launch a CPU worker to grab the rest of the cores\n",
    "GPU_COUNT = len(cluster.workers)\n",
    "UNUSED_CORES = CPU_COUNT - GPU_COUNT\n",
    "nanny = await Nanny(cluster.scheduler_address, nthreads=UNUSED_CORES, resources={\"CPU\": UNUSED_CORES})\n",
    "\n",
    "# Connect a  client\n",
    "client = await Client(cluster, asynchronous=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dcbe123e-aaf1-4fff-beac-f5090707ed1b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import dask.array as da\n",
    "import cupy as cp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ad8d5a61-27b5-4010-a4e3-fdbf4b0f34cb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create some data on the CPU\n",
    "with dask.annotate(resources={'CPU': 1}):\n",
    "    arr = da.random.random((10_000, 10_000), chunks=(1000, 1000)).persist()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e8f31474-6c40-474d-9a46-c06fba27f5b4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Wait for CPU tasks (workaround for https://github.com/dask/distributed/issues/5937)\n",
    "_ = await wait(arr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "57edeb67-15fc-40fd-b374-a0ec5efb0606",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Use the data on the GPU\n",
    "with dask.annotate(resources={'GPU': 1}):\n",
    "    arr = arr.map_blocks(cp.asarray)\n",
    "    result = arr.mean().persist()"
   ]
  }
 ],
 "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.8.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}

screenshot.png