DaskExample.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "92162538-672b-48c2-ae61-653c2a5733c8",
   "metadata": {},
   "source": [
    "# How to easily parallelize your code in python with Dask"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e6cd0140-1b81-4c94-803b-c6b692c9417c",
   "metadata": {},
   "source": [
    "## Importing essential libraries for tutorial"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5c3e25b3-0a42-4e2d-b3c0-b13569a98879",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import pickle\n",
    "import gzip\n",
    "import shutil\n",
    "import psutil\n",
    "from collections import Counter\n",
    "from Bio import SeqIO\n",
    "from dask.distributed import Client, LocalCluster, progress\n",
    "from fuzzysearch import find_near_matches"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5d751410-40af-4f6b-a755-61e0785b861e",
   "metadata": {},
   "source": [
    "## Creating workspace and downloading dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "736d32ae-9f5d-4922-9d20-9f1a55591c0b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "os.makedirs(\"DaskWorkingSpace\", exist_ok = True)\n",
    "os.chdir(\"DaskWorkingSpace\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "77bc48a1-8edb-4c77-ac68-a7f25160ed10",
   "metadata": {},
   "source": [
    "We will use curl, to be OS agnostic (works on Windows, Linux and Mac)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5849b5c9-79f1-44d0-9c75-2de586d0b4d1",
   "metadata": {},
   "outputs": [],
   "source": [
    "!curl http://ftp.ebi.ac.uk/ensemblgenomes/pub/release-51/plants/fasta/arabidopsis_thaliana/dna/Arabidopsis_thaliana.TAIR10.dna.toplevel.fa.gz --output athaliana.fa.gz"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9abd52df-2afb-4f72-86f5-e424db633a87",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Ungziping the dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dff17adb-bd14-41d8-ba8e-55807bcdc3ec",
   "metadata": {},
   "source": [
    "It's really easy with Python standard library"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a2b4d006-eea7-4050-8d46-df175c2daab9",
   "metadata": {},
   "outputs": [],
   "source": [
    "with gzip.open('athaliana.fa.gz', 'rb') as f_in:\n",
    "    with open('athaliana.fa', 'wb') as f_out:\n",
    "        shutil.copyfileobj(f_in, f_out)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5c2eef87-2127-443d-8ead-8408029bb043",
   "metadata": {},
   "source": [
    "## Parsing dataset with Biopython"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c13d0687-a271-4cb0-9e8f-0ddd4c25410c",
   "metadata": {},
   "outputs": [],
   "source": [
    "sequences = list(SeqIO.parse(\"athaliana.fa\", \"fasta\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5c7796fc-8157-4640-b222-f34a6b3343c2",
   "metadata": {},
   "source": [
    "Dataset contains 7 records, each is a set of DNA alphabet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3b4ebd96-2bf0-428c-8499-b19075b21b0a",
   "metadata": {},
   "outputs": [],
   "source": [
    "len(sequences)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6b4c253e-d970-434b-a1d4-5c9776a2c6c6",
   "metadata": {},
   "outputs": [],
   "source": [
    "sequences[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8f7cf943-9e24-43d8-8367-bf98d7a16466",
   "metadata": {},
   "source": [
    "Our task is to count <b> how many unique nucleotides are within each record </b>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "67473d3c-0715-40de-965f-96a62c01d8d6",
   "metadata": {},
   "source": [
    "## Defining function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e5c8142e-a1eb-4496-9585-5046a5b8f6bc",
   "metadata": {},
   "outputs": [],
   "source": [
    "def CountBases(record):\n",
    "    header = record.description\n",
    "    sequence = record.seq.__str__()\n",
    "    c = Counter(sequence)\n",
    "    result = (header,c)\n",
    "    return result"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8488583f-961e-46da-aa49-f98d17c35705",
   "metadata": {},
   "source": [
    "## Single process approach"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a19e75cd-2ecc-4e36-ad66-497fb9f5da7e",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "results_count = []\n",
    "for record in sequences:\n",
    "    result = CountBases(record)\n",
    "    results_count.append(result)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "540c2e2d-a741-4489-958b-58c5d7193258",
   "metadata": {},
   "source": [
    "## Multi process approach"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "878f2550-b089-4030-9dcb-614bc69cb718",
   "metadata": {},
   "source": [
    "### First we need to set up dask distributed LocalCluster and connect Client to it (for workers to communicate)\n",
    "<b> We want to use all of our systems power </b>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b91f5ab8-d36c-4559-8442-a6e0e843236c",
   "metadata": {},
   "outputs": [],
   "source": [
    "nproc = psutil.cpu_count()\n",
    "print(nproc)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f69b68ff-f44c-48dd-b41e-fb422ee98cfa",
   "metadata": {},
   "source": [
    "<b> If your CPU count is exceeding number of records within `sequences` list, set the nproc to the length of it (to be more kosher) </b>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "89407a5f-5067-455a-86b2-c3c420206880",
   "metadata": {},
   "outputs": [],
   "source": [
    "if nproc > len(sequences):\n",
    "    nproc = len(sequences)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c3ba4c97-3f9a-4efe-9dcf-11b9504462db",
   "metadata": {},
   "source": [
    "<b> Finally we define the cluster object </b>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0121e68e-f665-400b-b115-7c015fdb15ad",
   "metadata": {},
   "outputs": [],
   "source": [
    "cluster = LocalCluster(n_workers = nproc)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "71725d07-6ab3-4110-81e9-8b3b84fab7c8",
   "metadata": {},
   "source": [
    "<b> Now we can connect our client to the cluster </b>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7d21e4ae-5946-4642-9b77-5228650522be",
   "metadata": {},
   "outputs": [],
   "source": [
    "client = Client(cluster)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c3b9dfe1-18a3-4b7e-8c6b-35a408b18083",
   "metadata": {},
   "outputs": [],
   "source": [
    "client"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "30a3a3d2-712d-4f66-b43b-db289b981e62",
   "metadata": {},
   "source": [
    "<b> Scheduler automatically chooses its port, but the best thing is that at 127.0.0.1:8787 you'll find dashboard on which you can see how well your computation is going (if you are using defaults) </b>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2c46eac0-cf53-4b8b-ba41-5c9b8f553bb3",
   "metadata": {},
   "source": [
    "### And run parallelized computation with the same function as above"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "55927fbb-c8bb-46af-9864-c6bc70c06c77",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "futures = client.map(CountBases, sequences)\n",
    "progress(futures) #the progress widget, more informative within pure python scripts\n",
    "results_count_parallel = client.gather(futures)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d18eaaaf-7035-4a6b-b9bd-63e4eb9430b6",
   "metadata": {},
   "outputs": [],
   "source": [
    "results_count == results_count_parallel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dcc9c2c1-30a2-4aa4-9d8a-9ed8595e8fc9",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(results_count_parallel)\n",
    "del results_count\n",
    "del results_count_parallel"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c26e8860-20fd-4921-a645-a15e9a0f4174",
   "metadata": {},
   "source": [
    "<b> It is very good practice to delete ram residuals by method `cancel` of `Client` </b>\n",
    "\n",
    "```\n",
    "    client.cancel(futures)\n",
    "```\n",
    "\n",
    "You can also do it by restarting whole Client (which is not recommended), but for the educational purpose we will do that (to have clean state of scheduler that does not have a clue about data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "edcf05bd-3780-43bd-be5f-966303bc5678",
   "metadata": {},
   "outputs": [],
   "source": [
    "client.restart()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "15d1046d-66f1-4be3-b006-4a861fbbcd67",
   "metadata": {},
   "source": [
    "## Now let's do it again on less trivial task of pattern searching\n",
    "We do not want to exceed the memory on Your machines, so we will use only 1/8 of each record"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f8dd5706-5994-426b-aebf-26c3eb9db15b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def FindPattern(record, pattern, subs, ins, dels):\n",
    "    header = record.description\n",
    "    sequence = record.seq.__str__()\n",
    "    matches = find_near_matches(pattern, sequence, max_substitutions = subs, max_insertions = ins, max_deletions = dels)\n",
    "    return (header,matches)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ac4bd65a-77d8-4d55-8827-5ddb9dd8bcf1",
   "metadata": {},
   "source": [
    "### Let's define a preset of function arguments"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "04c32a9e-91c9-4008-8536-38358d1d33ad",
   "metadata": {},
   "outputs": [],
   "source": [
    "pattern = \"TGATTTGGATGATTCAAGACTTCTCGGTACTGCA\"\n",
    "subs = 1\n",
    "ins = 1\n",
    "dels = 1"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9038d1d1-d612-4322-b7b6-cbea3af2816c",
   "metadata": {},
   "source": [
    "### Single process"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "57569c5f-d6c4-4694-b72c-9a5c497b663d",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "results_fuzz = []\n",
    "for record in sequences:\n",
    "    result = FindPattern(record, pattern, subs, ins, dels)\n",
    "    results_fuzz.append(result)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3de1e274-17f5-4247-9162-f80bb3dd63c7",
   "metadata": {},
   "source": [
    "### Multi process"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0dec573e-8183-41e8-8416-a4c85bfe5581",
   "metadata": {},
   "source": [
    "#### This time function takes multiple arguments, so we need to define helper function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f893a68b-da2b-4893-a33a-637e783f20d5",
   "metadata": {},
   "outputs": [],
   "source": [
    "def FuzzHelper(x):\n",
    "    return FindPattern(x, pattern, subs, ins, dels)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "35cbdff9-43f9-43b1-9996-3eb0821b696f",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "futures = client.map(FuzzHelper, sequences)\n",
    "progress(futures)\n",
    "results_fuzz_parallel = client.gather(futures)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "88768433-ff90-4bfa-9ea9-9a2e36fba56b",
   "metadata": {},
   "outputs": [],
   "source": [
    "results_fuzz == results_fuzz_parallel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f2d6456c-b9e4-4f22-8b1c-c78258d92ae3",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(results_fuzz_parallel)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6549c5c4-4f6d-4e5b-8972-88d9d23893d4",
   "metadata": {},
   "outputs": [],
   "source": [
    "client.restart()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5699d8a2-f6c1-40bd-b606-3de298d42a83",
   "metadata": {},
   "source": [
    "## Working with data chunks and I/O streaming"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "22529d19-c448-4002-95ac-f60bd6f22e37",
   "metadata": {},
   "outputs": [],
   "source": [
    "def SplitByteChunks(records):\n",
    "    os.makedirs(\"tmp\", exist_ok = True)\n",
    "    chunkpaths = []\n",
    "    for i in range(len(records)):\n",
    "        header = records[i].description\n",
    "        sequence = records[i].seq.__str__()\n",
    "        chunk = (header, sequence)\n",
    "        fpath = os.path.join(\"tmp\", f\"chunk_{i}.pickle\")\n",
    "        with open(fpath, \"wb\") as handle:\n",
    "            pickle.dump(chunk, handle)\n",
    "        chunkpaths.append(fpath)\n",
    "    return chunkpaths"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7afeeb68-03a2-43db-a6ae-19b9e68542c2",
   "metadata": {},
   "outputs": [],
   "source": [
    "pickles = SplitByteChunks(sequences)\n",
    "print(pickles)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5895bc59-945b-4968-97cb-e08656a9ecca",
   "metadata": {},
   "outputs": [],
   "source": [
    "def FindPatternPickle(fpath, pattern, subs, ins, dels):\n",
    "    with open(fpath, \"rb\") as handle:\n",
    "        record = pickle.load(handle)\n",
    "    header, sequence = record\n",
    "    matches = find_near_matches(pattern, sequence, max_substitutions = subs, max_insertions = ins, max_deletions = dels)\n",
    "    return (header,matches)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7d3b4fa5-dd9a-40fb-8d30-8d9358493275",
   "metadata": {},
   "outputs": [],
   "source": [
    "def PickleHelper(x):\n",
    "    return FindPatternPickle(x, pattern, subs, ins, dels)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cad8797b-d967-4722-b985-85292bb0baeb",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "futures = client.map(PickleHelper, pickles)\n",
    "progress(futures)\n",
    "results_pickle = client.gather(futures)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ab95ab9c-e05a-4bd8-9abb-a1ce3021fd48",
   "metadata": {},
   "outputs": [],
   "source": [
    "results_fuzz == results_fuzz_parallel == results_pickle"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e0532484-7273-4430-bb71-1a7404602c67",
   "metadata": {},
   "source": [
    "# Dask docs\n",
    "- https://dask.org/ (Main page)\n",
    "- http://distributed.dask.org/en/latest/ (Dask Distributed)\n",
    "- https://docs.dask.org/en/latest/dataframe.html (Dask DataFrame)\n",
    "- https://docs.dask.org/en/latest/bag.html (Dask Bag)\n",
    "- https://ml.dask.org/ (Dask Machine Learning)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}