Cpu vs GPU nsmallest perf

nsmallest_perf.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "fb477301-c2a1-4de7-ac3f-e61bed019788",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<Client: 'tcp://127.0.0.1:44137' processes=8 threads=64, memory=376.55 GiB>\n"
     ]
    }
   ],
   "source": [
    "gpu = False\n",
    "\n",
    "from distributed import Client, wait\n",
    "from dask_cuda import LocalCUDACluster\n",
    "\n",
    "try:\n",
    "    cluster.close()\n",
    "    client.close()\n",
    "except NameError:\n",
    "    client = None\n",
    "\n",
    "if gpu:\n",
    "    cluster = LocalCUDACluster(protocol=\"tcp\", rmm_pool_size=\"12GB\")\n",
    "    client = Client(cluster)\n",
    "else:\n",
    "    client = Client()\n",
    "\n",
    "print(client)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "85dbeb68-4bba-4dcc-9289-a651d0816469",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "14400000\n"
     ]
    }
   ],
   "source": [
    "from dask.datasets import timeseries\n",
    "ddf = timeseries(freq=\"180ms\", partition_freq=\"3D\", seed=42)\n",
    "pdf = ddf.compute()\n",
    "if gpu:\n",
    "    import cudf\n",
    "    ddf = ddf.map_partitions(cudf.from_pandas)\n",
    "ddf = ddf.persist()\n",
    "wait(ddf)\n",
    "print(len(ddf))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "2c160f0e-7420-48e6-b813-10e251cab0a1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "8.29 s ± 235 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit \n",
    "wait(ddf.nsmallest(len(pdf)//2, ['id']).persist())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "e29c306c-d8ed-4797-a223-afcb8d1d439d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "3.2 s ± 83.7 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit \n",
    "wait(ddf.sort_values(by=['id']).persist())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f987b13e-b2db-4538-ab68-559600ff722d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from dask.distributed import performance_report\n",
    "fname = (\"gpu\" if gpu else \"cpu\")+ \"-nostring.html\"\n",
    "with performance_report(filename=fname):\n",
    "    wait(ddf.nsmallest(len(pdf)//2, ['id']).persist())"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "dask-sql",
   "language": "python",
   "name": "dask-sql"
  },
  "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.9.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}