06_Core_Tasks_Image_Processing.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "76ff7ebf-95b8-49a1-b78f-ca16d5315cd3",
   "metadata": {},
   "source": [
    "# Ray Core: Quick Intro and Image Processing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "79c8feae-3ec6-42f1-a284-47b9d8c5214d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import ray\n",
    "\n",
    "if ray.is_initialized:\n",
    "    ray.shutdown()\n",
    "\n",
    "ray.init(runtime_env={\"working_dir\": \".\"})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5e2bcec0-7015-4f79-83c8-c9bb0cb72621",
   "metadata": {},
   "source": [
    "### Ray core: primitves for fast, simple remote tasks, actors, and shared storage"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "730c01e5-3061-43e4-868b-12c55646082a",
   "metadata": {},
   "outputs": [],
   "source": [
    "@ray.remote\n",
    "def sum(arr):\n",
    "    return arr.sum()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eb971fd1-0b6f-4578-b6a5-fdaefaf1179c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "i = np.eye(10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10c515b2-fc89-4389-a59f-5c033f016a59",
   "metadata": {},
   "outputs": [],
   "source": [
    "result = sum.remote(i)\n",
    "result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "82376063-4835-40eb-b048-697db728c293",
   "metadata": {},
   "outputs": [],
   "source": [
    "ray.get(result)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ac338ba2-9d6f-4ce7-b51c-147e13771a2f",
   "metadata": {},
   "source": [
    "Tasks can include arbitrary dependencies, whether or not those are finished"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ce354ddd-bd31-4c37-a33f-3fbfee56a4e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "@ray.remote\n",
    "def double(arr):\n",
    "    return 2*arr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3e131068-645b-4b20-837a-83cd1d30fc66",
   "metadata": {},
   "outputs": [],
   "source": [
    "i = np.eye(10)\n",
    "j = double.remote(i)\n",
    "k = sum.remote(j)\n",
    "k"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "62722369-9336-4648-9d59-4945d0d19868",
   "metadata": {},
   "source": [
    "We can block to get any values we need ... although it's best to do so only when necessary"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "02daff9f-1c9b-4229-9219-b6fef9f6b177",
   "metadata": {},
   "outputs": [],
   "source": [
    "ray.get(k)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4c759281-5dc7-4eb3-a749-0c288405cffe",
   "metadata": {},
   "source": [
    "Where are the objects before we \"get\" them? In the object store. We can directly place data there as well -- e.g., if we would like to share it across nodes and workers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4a3c3ffd-c40c-4ec0-9951-90b09cf000ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "handle = ray.put(np.eye(5))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b08830c8-341a-45f4-b1f1-0a708448ff06",
   "metadata": {},
   "outputs": [],
   "source": [
    "handle"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "25f6e163-1850-462f-b757-cb55831c0287",
   "metadata": {},
   "source": [
    "Like the results of remote calls, these handles can be transparently passed to remote functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b7acd1ef-fea8-4dd2-921e-e01b276b3312",
   "metadata": {},
   "outputs": [],
   "source": [
    "another_sum = sum.remote(handle)\n",
    "another_sum"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "91d6201d-0aee-41b1-a288-706077ebf6b7",
   "metadata": {},
   "outputs": [],
   "source": [
    "ray.get(another_sum)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fb6794ed-164e-425f-b30a-041f1ccb4114",
   "metadata": {},
   "source": [
    "## How to use Ray distributed tasks for image transformation and computation\n",
    "For this example, we will simulate a compute-intensive task by transforming and computing some operations on large high-resolution images. These tasks are not uncommon in image classification in a DNN for training and transposing\n",
    "images. \n",
    "\n",
    "PyTorch `torchvision.transforms` API provides many transformation APIs. We will use a couple here, along with some `numpy` and `torch.tensor` operations. Our tasks will perform the following compute-intensive transformations:\n",
    "\n",
    " 1. Use PIL APIs to [blur the image](https://pillow.readthedocs.io/en/stable/reference/ImageFilter.html) with a filter intensity\n",
    " 2. Use Torchvision random [trivial wide augmentation](https://pytorch.org/vision/stable/generated/torchvision.transforms.TrivialAugmentWide.html#torchvision.transforms.TrivialAugmentWide)\n",
    " 3. Convert images into numpy array and tensors and do numpy and torch tensor operations such as [transpose](https://pytorch.org/docs/stable/generated/torch.transpose.html), element-wise [multiplication](https://pytorch.org/docs/stable/generated/torch.mul.html) with a random integers\n",
    " 4. Do more exponential [tensor power](https://pytorch.org/docs/stable/generated/torch.pow.html) and [multiplication with tensors](https://pytorch.org/docs/stable/generated/torch.mul.html)\n",
    "\n",
    "The goal is to compare execution times running these task serially vs. distributed as a Ray Task.\n",
    "\n",
    "|<img src=\"https://technical-training-assets.s3.us-west-2.amazonaws.com/Ray_Core/images_for_transformation.png\" width=\"80%\" height=\"30%\"> |\n",
    "|:--|\n",
    "|High resolution images for transformation and computation.|"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bee29917",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import time\n",
    "import logging\n",
    "import math\n",
    "import random\n",
    "\n",
    "from pathlib import Path\n",
    "from typing import Tuple, List\n",
    "from PIL import Image\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import pyarrow.parquet as pq\n",
    "from tqdm.notebook import tqdm\n",
    "import tasks_helper_utils as t_utils"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "394af3c2-f68e-4696-84ce-f9d716c63c30",
   "metadata": {},
   "source": [
    "Define some constants that can be tweaked for experimentation with different batch sizes as part of your exercsie."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "07a19d90-9278-4d8d-82b2-a6c3ff900480",
   "metadata": {},
   "outputs": [],
   "source": [
    "DATA_DIR = Path(\"/tmp/task_images\")\n",
    "BATCHES = [5, 10, 20] # , 30, 40, 50]\n",
    "SERIAL_BATCH_TIMES = []\n",
    "DISTRIBUTED_BATCH_TIMES = []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1f499cb1-dbf8-43d7-ae6c-c6a584018562",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "torch.set_num_threads(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "161098df-f1c1-498e-8256-60efce8f11d6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define a function to run these transformation tasks serially, on a single node, single core\n",
    "def run_serially(img_list: List) -> List[Tuple[int, float]]:\n",
    "    return [t_utils.transform_image(Image.open(image)) for image in tqdm(img_list)]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0760d0a7-210a-48cc-b4f5-294dc8251267",
   "metadata": {},
   "source": [
    "Let's download 100 large images, each betwen 5-20 MB+ with high-resolution greater (4000, 3500) pixels. It will only download once."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8c37d5a7-bf6d-4b0d-b0a4-63f602beb135",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check if dir exists. If so ignore download.\n",
    "# Just assume we have done from a prior run\n",
    "if not os.path.exists(DATA_DIR):\n",
    "    os.mkdir(DATA_DIR)\n",
    "    print(f\"downloading images ...\")\n",
    "    for url in tqdm(t_utils.URLS):\n",
    "        t_utils.download_images(url, DATA_DIR)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9828ce0a-37f6-4bda-8043-ce45c758b695",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Fetch the the entire image list\n",
    "image_list = list(DATA_DIR.glob(\"*.jpg\"))\n",
    "image_list[:2]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "90336873-bfdc-4627-8d9e-01e61bdc83be",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Let's look at some of random images, five for now, we are working with. Nice to be one with the data.\n",
    "t_utils.display_random_images(image_list, n=5)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "93659bf2-ab75-4a97-94e2-bf16e0420d87",
   "metadata": {},
   "source": [
    "### Run serially: each image transformation with a Python function\n",
    "\n",
    "We will iterate through the images with batches of 10 (this can be changed 20 or 25, etc) and process them. To simulate a computer-intensive operation on images, we are doing the tensor transformation and computations described above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "37c09e04-fa04-4c45-b9db-302a0dde3922",
   "metadata": {},
   "outputs": [],
   "source": [
    "for idx in BATCHES:\n",
    "    # Use the index to get N number of URLs to images\n",
    "    image_batch_list = image_list[:idx]\n",
    "    print(f\"\\nRunning {len(image_batch_list)} tasks serially....\")\n",
    "    \n",
    "    # Run each one serially\n",
    "    start = time.perf_counter()\n",
    "    serial_results = run_serially(image_batch_list)\n",
    "    end = time.perf_counter()\n",
    "    elapsed = end - start\n",
    "    \n",
    "    # Keep track of batches, execution times as a Tuple\n",
    "    SERIAL_BATCH_TIMES.append((idx, round(elapsed, 2)))\n",
    "    print(f\"Serial transformations/computations of {len(image_batch_list)} images: {elapsed:.2f} sec\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a0a9e510-0152-41ae-a2b4-b08e00a5bcac",
   "metadata": {},
   "source": [
    "### Run distributed: each image transformation with a Ray task\n",
    "\n",
    "Let's create a Ray task for an image within each batch and process them. Since \n",
    "our images are large, let's put them in the [Ray Distributed object store](https://docs.ray.io/en/latest/ray-core/key-concepts.html#objects).\n",
    "\n",
    "|<img src=\"https://technical-training-assets.s3.us-west-2.amazonaws.com/Overview_of_Ray/object_store.png\" width=\"70%\" loading=\"lazy\">|\n",
    "|:--|\n",
    "|Diagram of workers in worker nodes using `ray.put()` to store values and using `ray.get()` to retrieve them from each node's object store.|"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fd62e454-f2a5-49f7-9783-f10b619343d8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Put images in object store\n",
    "object_refs_list = [ray.put(Image.open(img)) for img in tqdm(image_list)]\n",
    "\n",
    "object_refs_list[:2], len(object_refs_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5e2f9aa8-0899-4e3f-ac28-2146f31339c9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define a Ray task to transform, augment and do some compute intensive tasks on an image\n",
    "@ray.remote\n",
    "def augment_image_distributed(img: Image) -> List[object]:\n",
    "    return t_utils.transform_image(img)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "beb18a7d-a09f-42fe-a960-0655df05dc75",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define function to run these transformation tasks distributed\n",
    "def run_distributed(obj_ref_list:List) ->  List[Tuple[int, float]]:\n",
    "    return ray.get([augment_image_distributed.remote(img) for img in tqdm(obj_ref_list)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1cb12234-e7a1-4172-b896-c7268dd2ef73",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Iterate over batches of 10, launching Ray task for each image within the processing\n",
    "# batch\n",
    "for idx in BATCHES:\n",
    "    image_obj_ref_batch_list = object_refs_list[:idx]\n",
    "    print(f\"\\nRunning {len(image_obj_ref_batch_list)} tasks distributed....\")\n",
    "    \n",
    "    # Run each one serially\n",
    "    start = time.perf_counter()\n",
    "    distributed_results = run_distributed(image_obj_ref_batch_list)\n",
    "    end = time.perf_counter()\n",
    "    elapsed = end - start\n",
    "    \n",
    "     # Keep track of batchs, execution times as a Tuple\n",
    "    DISTRIBUTED_BATCH_TIMES.append((idx, round(elapsed, 2)))\n",
    "    print(f\"Distributed transformations/computations of {len(image_obj_ref_batch_list)} images: {elapsed:.2f} sec\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2e0b1e0a-64e6-4e82-906e-5a2fe7649cdb",
   "metadata": {},
   "source": [
    "### Compare and plot the serial vs. distributed computational times"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e7e74c33-c85f-4916-99f8-9de21aebf7f8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Print times for each and plot them for comparison\n",
    "print(f\"Serial times & batches     : {SERIAL_BATCH_TIMES}\")\n",
    "print(f\"Distributed times & batches: {DISTRIBUTED_BATCH_TIMES}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17db3ef2-c1d1-4341-8a12-04b229dae6f0",
   "metadata": {},
   "outputs": [],
   "source": [
    "t_utils.plot_times(BATCHES, SERIAL_BATCH_TIMES, DISTRIBUTED_BATCH_TIMES)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a31bb3a5-57c8-4396-a75c-896c274e1127",
   "metadata": {},
   "source": [
    "### Recap\n",
    "\n",
    "We can clearly observe that the overall execution times by Ray tasks is in order of **3-4x** faster 🚅 than serial. Converting an existing serial compute-intensive Python function is as simple as adding the `ray.remote(...)` operator to your Python function. And Ray will handle all the hard bits: scheduling, execution, scaling, memory management, etc.\n",
    "\n",
    "As you can see the benefits are tangible in execution times with Ray tasks."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "08f363a1",
   "metadata": {},
   "source": [
    "### References\n",
    "\n",
    "1. [Modern Parallel and Distributed Python: A Quick Tutorial on Ray](https://towardsdatascience.com/modern-parallel-and-distributed-python-a-quick-tutorial-on-ray-99f8d70369b8) by Robert Nishihara, co-creator of Ray and co-founder Anyscale\n",
    "2. [Ray Core Introduction](https://www.anyscale.com/events/2022/02/03/introduction-to-ray-core-and-its-ecosystem) by Jules S. Damji"
   ]
  }
 ],
 "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.9.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}