Computer_vision_workloads/Semantic_segmentation/Scaling_batch_inference.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "fcb9a269-765a-46ef-abb6-d60682d38b0d",
   "metadata": {},
   "source": [
    "# Scalable Batch Inference with Ray"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1b114120-a325-4597-bf7b-8310c9c8d776",
   "metadata": {},
   "source": [
    "<img src=\"../../_static/assets/Generic/ray_logo.png\" width=\"20%\" loading=\"lazy\">\n",
    "\n",
    "## About this notebook\n",
    "\n",
    "### Is it right for you?\n",
    "\n",
    "This module focuses on batch inference task in the computer vision (CV) context, and presents several approaches for scaling it on Ray. It is right for you if:\n",
    "\n",
    "* You observe performance bottlenecks when working on batch inference problems in your CV projects.\n",
    "* You want to scale or increase throughput of your existing batch inference pipelines.\n",
    "* You wish to explore different architectures for scaling batch inference with Ray Core and Ray AIR.\n",
    "\n",
    "### Prerequisites\n",
    "\n",
    "For this notebook you should have:\n",
    "\n",
    "* Practical Python and machine learning experience.\n",
    "* Familiarity with batch inference task in ML.\n",
    "* Familiarity with Ray and Ray AIR equivalent to completing these training modules:\n",
    "  * [Overview of Ray](https://github.com/ray-project/ray-educational-materials/blob/main/Introductory_modules/Overview_of_Ray.ipynb)\n",
    "  * [Introduction to Ray AIR](https://github.com/ray-project/ray-educational-materials/blob/main/Introductory_modules/Introduction_to_Ray_AIR.ipynb)\n",
    "  * [Ray Core](https://github.com/ray-project/ray-educational-materials/tree/main/Ray_Core)\n",
    "\n",
    "### Learning objectives\n",
    "\n",
    "Upon completion of this notebook, you will be able to:\n",
    "\n",
    "* Evaluate common design patterns for distributed batch inference.\n",
    "* Decide about which design pattern to use for your batch inference problem in the CV context.\n",
    "* Implement scalable batch inference with Ray and tune its performance.\n",
    "\n",
    "### What will you do?\n",
    "\n",
    "* Learn about and evaluate several distributed batch inference design patterns with Ray Core and Ray AI Runtime.\n",
    "* Implement distributed batch inference through hands-on coding exercises."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "35a02709-0ed5-489a-b344-27005261eb72",
   "metadata": {},
   "source": [
    "## Part 1: Scalable batch inference design patterns with Ray\n",
    "\n",
    "One of the end goals for machine learning models is to generate predictions over a set of unseen data. In this notebook, you will look closely at the inference stage of the ML workflow and explore ways to scale it.\n",
    "\n",
    "Ray Core and Ray AIR APIs allow you to scale batch inference to millions of examples and offer various performance tuning opportunities.\n",
    "\n",
    "|<img src=\"../../_static/assets/Scaling_inference/example_ml_workflow.png\" width=\"70%\" loading=\"lazy\">|\n",
    "|:--|\n",
    "|Example machine learning workflow. It starts with reading raw data and preprocessing it. These stages are followed by training and tuning jobs that eventually produce trained model with desired quality and performance metrics. Trained models are used for inference, often on the large scale data sets.|"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2450181f",
   "metadata": {},
   "source": [
    "### What is (batch) inference?\n",
    "\n",
    "<div class=\"alert alert-info\">\n",
    "  <strong>Batch inference</strong> (also known as offline inference): is the process of generating predictions on the batch of data.\n",
    "</div>\n",
    "\n",
    "Unlike *online inference* where predictions are returned as soon as possible after an observation is produced, batch inference generates predictions over a large number of input data when immediate response is not required or feasible. For example, batch inference is relevant when generating product recommendations with historical customer data or forecasting using time-aggregated observations.\n",
    "\n",
    "|<img src=\"../../_static/assets/Scaling_inference/batch_inference.png\" width=\"70%\" loading=\"lazy\">|\n",
    "|:--|\n",
    "|Batch inference takes in data batch into trained model and outputs predictions.|\n",
    "\n",
    "Below, you will conceptually encounter five architectures for performing batch inference on Ray. Each one offers scalability and performance customization."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d8becdd0-3f1e-4503-92b6-0e275e03ad75",
   "metadata": {},
   "source": [
    "### Stateless inference using Ray Tasks\n",
    "\n",
    "In the most basic approach, inference could be performed sequentially where the model scores incoming batches of data one after another. However, performance of this approach is limited by the single machine or GPU and does not scale.\n",
    "\n",
    "Straightforward and easy-to-implement way to scale out batch inference is with Ray tasks. Task is an arbitrary Python function that will be evaluated remotely in the compute cluster. This is an example of the [Remote Procedure Call](https://en.wikipedia.org/wiki/Remote_procedure_call) or RPC. Ray is an example of the RPC system.\n",
    "\n",
    "A Ray task is *stateless* because it computes an output (e.g. predictions) determined purely by its input data without storing or modifying internal information. In other words, tasks are *stateless* because they do not have any internal state. Typical example of the stateless function in deep learning is [SGD optimizer](https://en.wikipedia.org/wiki/Stochastic_gradient_descent).\n",
    "\n",
    "|<img src=\"../../_static/assets/Scaling_inference/task_inference.png\" width=\"70%\" loading=\"lazy\">|\n",
    "|:--|\n",
    "|Stateless inference using Ray Tasks. Ray assigns batches to workers via tasks as soon as a worker becomes available. Each task loads the trained model and outputs predictions independent of the other concurrent inference jobs. This approach scales with the number of available CPUs or GPUs.|\n",
    "\n",
    "<img src=\"../../_static/assets/Scaling_inference/code_task.png\" width=\"70%\" loading=\"lazy\">"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5687ea84-b693-4a97-8de4-26954a4c97c4",
   "metadata": {},
   "source": [
    "### Stateful inference using Ray Actors\n",
    "\n",
    "Loading large, complex models into memory can be computationally expensive. In addition, you may want the flexibility to capture some persistent internal state. This second approach avoids loading and discarding the model after each batch. The overall recipe is like this:\n",
    "\n",
    "1. Creating a number of replicas (i.e. Ray Actors) of the trained model.\n",
    "2. Feeding data into these model replicas in parallel and retrieving inference results.\n",
    "\n",
    "Ray actors are *stateful* because they keep an internal state, just like Python classes. [Adam optimizer](https://arxiv.org/abs/1412.6980), commonly used in deep learning is an example of the stateful object.\n",
    "\n",
    "|<img src=\"../../_static/assets/Scaling_inference/actor_inference.png\" width=\"70%\" loading=\"lazy\">|\n",
    "|:--|\n",
    "|Stateful inference using Ray Actors. Ray Actors hold model replicas which generate predictions from batches of data. You can scale out this approach with the number of actors. Actors can be reused multiple times and run inference for many batches of data.|\n",
    "\n",
    "<img src=\"../../_static/assets/Scaling_inference/code_actor.png\" width=\"70%\" loading=\"lazy\">"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5fc92ba9-95c2-48b2-a115-064be421995d",
   "metadata": {},
   "source": [
    "### Stateful inference using Ray ActorPool utility\n",
    "\n",
    "When using Ray Actors, you need to implement *load balancing* in order to keep good level of resources utilization. You need to track when Actors become available to assign new work until the entire process completes. \n",
    "\n",
    "This design pattern introduces Ray [ActorPool](https://docs.ray.io/en/latest/ray-core/package-ref.html#ray-util-actorpool) utility that wraps a list of actors and automatically handles load balancing (futures management). This is convenient abstraction that let you focus on the inference logic and leave resources utilization to be managed by Ray.\n",
    "\n",
    "|<img src=\"../../_static/assets/Scaling_inference/actor_pool.png\" width=\"70%\" loading=\"lazy\">|\n",
    "|:--|\n",
    "|Using ActorPools for Batch Inference.|\n",
    "\n",
    "In essence, the ActorPool wraps around the `n` actors so you do not have to manage idle actors and manually distributing workloads.\n",
    "\n",
    "<img src=\"../../_static/assets/Scaling_inference/code_actorpool.png\" width=\"70%\" loading=\"lazy\">"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8c500d89-bdae-41fd-9b12-264c0839af90",
   "metadata": {},
   "source": [
    "### Batch inference using Ray AIR Datasets\n",
    "\n",
    "In the previous few approaches, there exist some unoptimized aspects to discuss:\n",
    "* dispatching file splits one at a time may be inefficient for small batches or cause OutOfMemory errors if batches are too large (e.g. on GPUs)\n",
    "* you may want to have multiple tasks sent to an Actor at once (i.e. pipelining task submission)\n",
    "* data fetching and batch preprocessing could be parallelized as well\n",
    "\n",
    "While you could control how Ray executes by implementing performance optimizations through Ray Core primitives, [Ray AIR](https://docs.ray.io/en/latest/ray-air/getting-started.html) offers high-level composable APIs that have these optimizations built-in.\n",
    "\n",
    "[Ray Datasets](https://docs.ray.io/en/latest/data/dataset.html) allows for:\n",
    "1. parallel reading and preprocessing of source data\n",
    "2. dynamic autoscaling of the actor pool\n",
    "3. automatic batching and pipelining of data\n",
    "\n",
    "|<img src=\"../../_static/assets/Scaling_inference/ray_datasets.png\" width=\"70%\" loading=\"lazy\">|\n",
    "|:--|\n",
    "|Ray Datasets replace the 'Batch preprocessing' stage.|\n",
    "\n",
    "In Ray AIR, a trained model is loaded into a `Checkpoint` object (could be from training or tuning). An AIR `Predictor` loads model from the `Checkpoint` to perform inference. Then, using the preprocessed batches provided by Ray Datasets, you extract predictions off of the testing data.\n",
    "\n",
    "<img src=\"../../_static/assets/Scaling_inference/code_dataset.png\" width=\"70%\" loading=\"lazy\">"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "02a4c211-95fe-4218-a10a-08c0cfeafc9c",
   "metadata": {},
   "source": [
    "### Batch inference using Ray AIR BatchPredictor\n",
    "\n",
    "Finally, Ray AIR's [`BatchPredictor`](https://docs.ray.io/en/latest/ray-air/package-ref.html#batch-predictor) takes in a [`Checkpoint`](https://docs.ray.io/en/latest/ray-air/package-ref.html#checkpoint) which represents the saved model. This high-level abstraction offers simple and composable APIs that enable preprocessing data in batches with [BatchMapper](https://docs.ray.io/en/latest/ray-air/package-ref.html#generic-preprocessors) and instantiate a distributed predictor given checkpoint data.\n",
    "\n",
    "As a part of Ray AIR, you specify what you want done through a set of declarative key-value arguments rather than concerning yourself with how to instruct Ray to scale.\n",
    "\n",
    "|<img src=\"../../_static/assets/Scaling_inference/air_batchpredictor.png\" width=\"70%\" loading=\"lazy\">|\n",
    "|:--|\n",
    "|Using Ray AIR's `BatchPredictor` for Batch Inference.|\n",
    "\n",
    "The AIR `BatchPredictor` takes both the `Checkpoint` and `Predictor` to replace the process of manually performing inference on a large dataset.\n",
    "\n",
    "<img src=\"../../_static/assets/Scaling_inference/code_batchpredictor.png\" width=\"70%\" loading=\"lazy\">"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0ee761b0-7549-46ca-99bc-e3588a7a1deb",
   "metadata": {},
   "source": [
    "## Part 2: Data and model - computer vision transformers for semantic segmentation\n",
    "\n",
    "To demonstrate each architecture, you will implement each approach by running inference on a variation on an object detection task: semantic segmentation.\n",
    "\n",
    "### MIT ADE20K - scene parsing benchmark\n",
    "\n",
    "Semantic, or image, segmentation takes a scene and classifies image objects into semantic [categories](https://docs.google.com/spreadsheets/d/1se8YEtb2detS7OuPE86fXGyD269pMycAWe2mtKUj2W8/edit?usp=sharing) pixel-by-pixel. Often used as a standard for assessing segmentation model quality, the [MIT ADE20K Dataset](http://sceneparsing.csail.mit.edu/) (also known as \"SceneParse150\") provides the largest open source data set for scene parsing.\n",
    "\n",
    "|<img src=\"../../_static/assets/Scaling_inference/scene.png\" width=\"70%\" loading=\"lazy\">|\n",
    "|:--|\n",
    "|Test image on the left vs. predicted result on the right.[Source](https://github.com/CSAILVision/semantic-segmentation-pytorch) *Date accessed: November 10, 2022*|\n",
    "\n",
    "Data set highlights:\n",
    "\n",
    "* 20k annotated, scene-centric training images\n",
    "* 3.3k test images\n",
    "* 150 total categories such as person, car, bed, sky, and more\n",
    "\n",
    "### SegFormer - transformer-based framework for semantic segmentation\n",
    "\n",
    "[SegFormer](https://arxiv.org/pdf/2105.15203.pdf) is a simple and powerful semantic segmentation method whose architecture consists of a hierarchical Transformer encoder and a lightweight All-MLP decoder. What sets SegFormer apart from previous approaches boils down to two key features:\n",
    "\n",
    "1. a novel hierarchically structured Transformer encoder which does not depend on positional encoding, avoiding interpolation when test resolution differs from training\n",
    "2. avoids complex decoders using a lightweight MLP layer\n",
    "\n",
    "With demonstrated success on benchmarks such as Cityscapes and [MIT ADE20K Dataset](http://sceneparsing.csail.mit.edu/), you will use a pre-trained version to perform inference on test images from MITADE20K/SceneParse150.\n",
    "\n",
    "|<img src=\"../../_static/assets/Scaling_inference/segformer_architecture.png\" width=\"70%\" loading=\"lazy\">|\n",
    "|:--|\n",
    "|Segformer architecture taken from [original paper](https://arxiv.org/pdf/2105.15203.pdf). *Date accessed: November 10, 2022*|\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c018d564-0e4a-453f-9f7a-74466c0273e8",
   "metadata": {},
   "source": [
    "## Part 3: Sequential batch inference"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1a07b49f-ee7b-48c7-a4d0-bc4043ff24c9",
   "metadata": {},
   "source": [
    "To begin, you will build a basic version of batch inference that is sequential.\n",
    "\n",
    "|<img src=\"../../_static/assets/Scaling_inference/single_sequential_timeline.png\" width=\"90%\" loading=\"lazy\">|\n",
    "|:--|\n",
    "|Sequential inference on the single worker. Performance is limited to the single machine performance.|"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6eb5419c-26af-4b09-8fd2-e885aaa60021",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from PIL import Image\n",
    "\n",
    "# set the seed to a fixed value for reproducibility\n",
    "torch.manual_seed(201)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e9a7f1cb-4009-4848-b589-974f82285963",
   "metadata": {},
   "source": [
    "### Load pre-trained model from the HuggingFace Hub"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5e32feab-170b-4220-a263-bb678592d106",
   "metadata": {},
   "outputs": [],
   "source": [
    "from utils import get_labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4346c5a0-36b4-46ce-bdc0-3102543d2dff",
   "metadata": {},
   "outputs": [],
   "source": [
    "id2label, label2id = get_labels()\n",
    "\n",
    "print(f\"Total number of labels: {len(id2label)}\")\n",
    "print(f\"Example labels: {list(id2label.values())[:5]}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d9354716-8ac2-4e4f-97e1-e48ae174f0d9",
   "metadata": {},
   "source": [
    "`get_labels`, a utility function, provides two dictionary mappings from [HuggingFace](https://huggingface.co/datasets/huggingface/label-files/blob/main/ade20k-id2label.json):\n",
    "* `id2label`\n",
    "* `label2id`\n",
    "\n",
    "which allows you to convert between ids (int) and labels (str) for the 150 available categories of objects in images."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eab7f82a-93e7-44c3-ac0e-322ebcdfd443",
   "metadata": {},
   "source": [
    "#### Load SegFormer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "046f3990-3c6f-4ae4-9186-fa9fb1800f0b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import SegformerForSemanticSegmentation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "68aec6f8-22bf-41f9-bd7b-fe54889d762f",
   "metadata": {},
   "outputs": [],
   "source": [
    "MODEL_NAME = \"nvidia/segformer-b0-finetuned-ade-512-512\"\n",
    "\n",
    "segformer = SegformerForSemanticSegmentation.from_pretrained(\n",
    "    MODEL_NAME, id2label=id2label, label2id=label2id\n",
    ")\n",
    "\n",
    "print(f\"Number of model parameters: {segformer.num_parameters()/(10**6):.2f} M\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "24f3ef57",
   "metadata": {},
   "source": [
    "From [HuggingFace](https://huggingface.co/nvidia/segformer-b0-finetuned-ade-512-512), you specify the b0-sized (the smallest, ranging up to b5) SegFormer model.\n",
    "\n",
    "This pre-trained model contains 3.75 million parameters and is fine-tuned on the MITADE20K (SceneParse150) dataset on images with a 512 x 512 resolution. Keep this in mind when comparing strengths and weaknesses of various batch inference approaches."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "41546abd-2f3d-4128-8176-19cd0604d6f4",
   "metadata": {},
   "source": [
    "#### Create feature extractor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b51be71f-1c3f-45fb-b9b2-af3393bfadc4",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import SegformerFeatureExtractor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f888ed08-cde2-4f4a-b05f-e467de4b7ccf",
   "metadata": {},
   "outputs": [],
   "source": [
    "segformer_feature_extractor = SegformerFeatureExtractor.from_pretrained(\n",
    "    MODEL_NAME, reduce_labels=True\n",
    ")\n",
    "segformer_feature_extractor"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4a9a23f3",
   "metadata": {},
   "source": [
    "[Feature extractors](https://huggingface.co/docs/transformers/main_classes/feature_extractor) preprocess input features by normalizing, resizing, padding, and converting raw images into the desired cleaned shape.\n",
    "\n",
    "The [`reduce_labels`](https://huggingface.co/docs/transformers/model_doc/segformer#segformer) flag ensures that the \"background\" of an image isn't counted as its own separate category when computing loss. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3109152b-3572-40f9-8fc6-b04ecce59896",
   "metadata": {},
   "source": [
    "### Prepare SceneParse150 dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f3d58d4b-d13c-4c34-8e80-194c8f6bdae9",
   "metadata": {},
   "source": [
    "#### Load dataset from the HuggingFace Hub"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "39e93db0-6218-4460-8f8c-14703a57f29f",
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import load_dataset\n",
    "from utils import convert_image_to_rgb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b24ed45d-8700-40de-99ac-2bd4c5999974",
   "metadata": {},
   "outputs": [],
   "source": [
    "SMALL_DATA = True"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b25fc198-d3f4-47a0-9c16-9054a342c4f7",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-warning\">\n",
    "  <strong>SMALL_DATA</strong>: default `True` - set to download only 160 images from the data set. Set to `False` (recommended) to work with full testing dataset (3352 images).\n",
    "</div>\n",
    "\n",
    "If you set `SMALL_DATA` to `False` it will take some time (depending on your connection download speed), because you download over 20k images to the local machine or cluster."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d70f31fe-e8a9-41ea-857c-e2e2b36503e4",
   "metadata": {},
   "outputs": [],
   "source": [
    "DATASET_NAME = \"scene_parse_150\"\n",
    "\n",
    "if SMALL_DATA:\n",
    "    train_dataset = load_dataset(DATASET_NAME, split=\"train[:10]\")\n",
    "    test_dataset = load_dataset(DATASET_NAME, split=\"test[:160]\")\n",
    "else:\n",
    "    train_dataset = load_dataset(DATASET_NAME, split=\"train[:10]\")\n",
    "    test_dataset = load_dataset(DATASET_NAME, split=\"test\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8fc70f94",
   "metadata": {},
   "source": [
    "Download MITADE20K (SceneParse150) dataset from HuggingFace's datasets repository using the `load_dataset` utility."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "caf19f5d-52cf-4d15-a958-e143bcf711c2",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dadba2c1-eb5e-42f9-8de7-bb3311a04786",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_dataset = test_dataset.map(convert_image_to_rgb)\n",
    "test_dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0c6e66d8-2587-43c9-91bb-cdfc6f22b33d",
   "metadata": {},
   "source": [
    "Each example in the data set includes:\n",
    "* `image` - the PIL image\n",
    "* `annotation` - human annotations of image regions (annotation mask is `None` in testing set)\n",
    "* `category` - category of the scene generally (e.g. driveway, voting booth, dairy_outdoor). \n",
    "\n",
    "For the datasets:\n",
    "* `train_dataset` - retrieve a small sample of images for visualization purposes. Full training dataset set is 20210 images.\n",
    "* `test_dataset` - used for batch inference purposes. Full test data set is 3352 images."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6e2a63ad-3bfd-453c-827e-6a0f664ac310",
   "metadata": {
    "tags": []
   },
   "source": [
    "#### Display example images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12c22f28-f63c-42e8-ab43-1684aa692d5f",
   "metadata": {},
   "outputs": [],
   "source": [
    "from utils import display_example_images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f1054dfa-90a6-4569-bc0f-cb059f4b4893",
   "metadata": {},
   "outputs": [],
   "source": [
    "# try running multiple times!\n",
    "display_example_images(train_dataset)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bf7243ec-d207-4b33-89ca-b80309c19c06",
   "metadata": {},
   "source": [
    "### Run sequential batch inference on the single batch and visualize predictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2ba82c68-f3cd-46ae-8b23-6fee43ccb3a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "def predict(model, feature_extractor, images):\n",
    "    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "    model.to(device)\n",
    "    model.eval()\n",
    "    inputs = feature_extractor(images=images, return_tensors=\"pt\")\n",
    "    with torch.no_grad():\n",
    "        outputs = model(pixel_values=inputs.pixel_values.to(device))\n",
    "\n",
    "    image_sizes = [image.size[::-1] for image in images]\n",
    "    segmentation_maps_postprocessed = (\n",
    "        feature_extractor.post_process_semantic_segmentation(\n",
    "            outputs=outputs, target_sizes=image_sizes\n",
    "        )\n",
    "    )\n",
    "\n",
    "    return [j.detach().cpu().numpy() for j in segmentation_maps_postprocessed]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "810e036a",
   "metadata": {},
   "source": [
    "The `predict` function forms the basis for inferencing, and you will reencounter it multiple times throughout this notebook's exploration of approaches.\n",
    "\n",
    "Inputs\n",
    "\n",
    "* `model` - which model to use; in this case, SegFormer b0 finetuned on 512x512 ADE20K\n",
    "* `feature_extractor` - the preprocessing mechanism associated with the model\n",
    "* `image` - preprocessed image\n",
    "* `labels` - labels of 150 possible categories\n",
    "* `device` - type of device responsible for loading into memory\n",
    "\n",
    "Output\n",
    "* list of segmentation maps\n",
    "\n",
    "Core Logic\n",
    "\n",
    "1. `inputs` - the images are converted using `feature_extractor` into `3x512x512` (`CxHxW`) arrays. This is independent of the original size of the input image.\n",
    "2. `outputs` (the inference step) - the 512x512 images are then passed to the model, which then produces 150 128x128 predicted images, one image for each available category. Each image is a mask representing the part of the image that belongs to a category.\n",
    "3. `image_sizes` - flip the dimensions of the image to expected order from PIL (`WxH`) to HuggingFace (`HxW`)\n",
    "4. `feature_extractor.postprocess_semantic_segmentation` - HuggingFace utility that generates segmentation maps from raw outputs\n",
    "5. returns a list of segmentation maps, detached from the computation to move from GPU to CPU"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "336d6a94-8399-4616-8221-e1f5c357baf2",
   "metadata": {},
   "source": [
    "#### Run batch prediction on 16 images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "70115952-fdb0-40d3-87d0-a7903502a289",
   "metadata": {},
   "outputs": [],
   "source": [
    "from utils import get_image_indices"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "930f4caf-8964-4170-8bde-b126a5a0ae57",
   "metadata": {},
   "outputs": [],
   "source": [
    "BATCH_SIZE = 16\n",
    "\n",
    "image_indices = get_image_indices(dataset=test_dataset, n=BATCH_SIZE)\n",
    "image_indices"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a5b62f05-26bc-46bf-941d-945bcdf0068e",
   "metadata": {},
   "source": [
    "Get `BATCH_SIZE` random image IDs from the test data set to run inference on. Each time you run this code, different images from the test set are selected."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a3302db4-6b93-4ef6-96b2-2bd316023044",
   "metadata": {},
   "outputs": [],
   "source": [
    "batch = [test_dataset[i][\"image\"] for i in image_indices]\n",
    "batch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76b9f0f1-be22-4910-a5a1-e6d861cf3a19",
   "metadata": {},
   "outputs": [],
   "source": [
    "segmentation_maps = predict(\n",
    "    model=segformer,\n",
    "    feature_extractor=segformer_feature_extractor,\n",
    "    images=batch,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a67a2244-089f-4af0-9d0f-7bd1c53820e6",
   "metadata": {},
   "source": [
    "Time how long it takes to run `predict` on a batch of images."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "72b7394d-eef4-498e-bee4-38214df85406",
   "metadata": {},
   "outputs": [],
   "source": [
    "segmentation_maps[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ee752e3a-d615-4773-a932-83018aee119e",
   "metadata": {},
   "source": [
    "#### Visualize example predictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b7d6a72-ed3e-4548-9b04-8490ee25f322",
   "metadata": {},
   "outputs": [],
   "source": [
    "from utils import visualize_predictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d99ef21c-879e-46b3-b97f-1c00efcae05e",
   "metadata": {},
   "outputs": [],
   "source": [
    "visualize_predictions(image=batch[0], segmentation_maps=segmentation_maps[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "52be8531-ad94-4fde-b13d-a864f3020a37",
   "metadata": {},
   "source": [
    "### Run sequential batch inference on 10 batches"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "93faafea-ddec-4148-a60a-a89d22b12aa8",
   "metadata": {
    "tags": []
   },
   "source": [
    "#### Prepare batches"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c3e17bf3-9c04-47eb-a7c5-604a0ffbb3b8",
   "metadata": {},
   "outputs": [],
   "source": [
    "BATCH_SIZE = 16\n",
    "N_BATCHES = 10\n",
    "\n",
    "image_indices = get_image_indices(dataset=test_dataset, n=BATCH_SIZE * N_BATCHES)\n",
    "image_indices_grouped = np.split(np.asarray(image_indices), N_BATCHES)\n",
    "image_indices_grouped"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d5f6004f-1cd5-46c9-89ad-708bbfe67312",
   "metadata": {},
   "outputs": [],
   "source": [
    "batches = []\n",
    "\n",
    "for image_idx in image_indices_grouped:\n",
    "    batch = [test_dataset[int(i)][\"image\"] for i in image_idx]\n",
    "    batches.append(batch)\n",
    "\n",
    "batches[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8c45ae67",
   "metadata": {},
   "source": [
    "To prepare batches, retrieve `BATCH_SIZE` number of images per `N_BATCHES` from the test set. The above code first fetches the indices of shuffled images, then prepares a list of images associated with the indices."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7907202d-bc99-40d0-aa15-e8200a206426",
   "metadata": {},
   "source": [
    "#### Run batch prediction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "45ca1763-8d9d-46aa-90c0-29b4b2be3c8a",
   "metadata": {},
   "outputs": [],
   "source": [
    "predictions = []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9154c7c3-8fae-44f9-9e4c-bca3b62b7328",
   "metadata": {},
   "outputs": [],
   "source": [
    "for batch in batches:\n",
    "    segmentation_maps = predict(\n",
    "        model=segformer,\n",
    "        feature_extractor=segformer_feature_extractor,\n",
    "        images=batch,\n",
    "    )\n",
    "    predictions.append(segmentation_maps)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d24d6603-50ee-4ed7-b17e-411fb5ca92b1",
   "metadata": {},
   "source": [
    "Notice that increasing the number of batches by 10 increases the runtime by ~10 which is the kind of linear scaling you expect in a sequential approach."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "89d425f7-d23c-41d1-adee-3cc9509e190f",
   "metadata": {},
   "outputs": [],
   "source": [
    "predictions[0][0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ab7f7f0a",
   "metadata": {},
   "source": [
    "Inspect the resulting `predictions` array to see predicted segmentation maps."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "21e69cad-d947-42d0-9cad-8f0a2d4941c1",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Summary: sequential batch inference implementation\n",
    "\n",
    "|<img src=\"../../_static/assets/Scaling_inference/single_sequential_timeline.png\" width=\"90%\" loading=\"lazy\">|\n",
    "|:--|\n",
    "|Sequential inference on the single worker. Performance is limited to the single machine performance.|\n",
    "\n",
    "#### Key concepts\n",
    "\n",
    "<div class=\"alert alert-info\">\n",
    "  <strong>Batch inference</strong> (also known as offline inference): is the process of generating predictions on the batch of data.\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d33c74b3-9d0f-469f-83ec-7e5557104489",
   "metadata": {},
   "source": [
    "## Part 4: Distributed, stateless batch inference with Ray Tasks"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3cba2f68-c762-439e-850b-3af83ba1eaae",
   "metadata": {},
   "source": [
    "In the first approach using Ray, this implementation transitions from sequential to parallel inferencing by loading the model across stateless functions to generate predictions.\n",
    "\n",
    "|<img src=\"../../_static/assets/Scaling_inference/task_inference.png\" width=\"70%\" loading=\"lazy\">|\n",
    "|:--|\n",
    "|Stateless inference using Ray Tasks.|"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "665bc8b6-878e-4fde-9054-2891acf83c1b",
   "metadata": {},
   "source": [
    "### Initialize Ray runtime"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ce928af7-2037-4b65-9dd1-6d67b5335e0f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import ray"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a2c97168",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "if ray.is_initialized:\n",
    "    ray.shutdown()\n",
    "\n",
    "ray.init()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fc24d09d-3b63-4307-a47e-1a72b2e33c9a",
   "metadata": {},
   "source": [
    "### Put the model and feature extractor in the object store\n",
    "\n",
    "|<img src=\"../../_static/assets/Overview_of_Ray/object_store.png\" width=\"70%\" loading=\"lazy\">|\n",
    "|:--|\n",
    "|Diagram of workers in worker nodes using `ray.put()` to place values and using `ray.get()` to retrieve them from each node's object store.|"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2abbe8d8-f00d-4ba8-a81f-aa0a5f480854",
   "metadata": {},
   "outputs": [],
   "source": [
    "segformer_ref = ray.put(segformer)\n",
    "segformer_feature_extractor_ref = ray.put(segformer_feature_extractor)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6ffc8a08-5629-42ae-94a0-1bdafad47c74",
   "metadata": {},
   "source": [
    "When passing an object as an argument to the remote function, Ray calls [ray.put()](https://docs.ray.io/en/latest/ray-core/package-ref.html#ray-put) implicitly to store that object in the local object store (local worker node), making it available to all local tasks. However, when that object is large, you want to avoid re-copying it every time the object is passed to a remote function or method.\n",
    "\n",
    "By explicitly storing both the model and feature extractor into the object store, you avoid having multiple copies which improves performance.\n",
    "\n",
    "If there are multiple worker nodes in your cluster, the object needs to be copied in memory when it is used on a worker node different from where it is stored. *Zero copy* is not guaranteed.\n",
    "\n",
    "<div class=\"alert alert-warning\">\n",
    "  <strong>Pro Tip</strong>\n",
    "\n",
    "  Passing the same large argument (model), by value repeatedly <a href=\"https://docs.ray.io/en/latest/ray-core/patterns/pass-large-arg-by-value.html\">harms performance</a>.\n",
    "  \n",
    "  Use **object store** and **ray.put()** to pass by reference instead (for example: model_ref, instead of model).\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "946db436-ac8b-4ac9-8dca-e557dbd126fb",
   "metadata": {},
   "source": [
    "### Implement remote function for inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "797b9e38",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "@ray.remote\n",
    "def inference_task(model, feature_extractor, images):\n",
    "    return predict(\n",
    "        model=model,\n",
    "        feature_extractor=feature_extractor,\n",
    "        images=images,\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e695c89f",
   "metadata": {},
   "source": [
    "Notice here that `inference_task` wraps the `predict()` function from before, and `@ray.remote` specifies this as the remote function.\n",
    "\n",
    "Stateless (lambda style) way of parallelising prediction is to create Ray tasks that load the trained model internally when called. This way we can make the prediction task \"stateless\", but at the cost of incurring the overhead of loading the model every single time.\n",
    "\n",
    "When called, each Ray task loads the trained model from the local object store to perform inference. This pattern works well for small models which do not encounter the same level of bottleneck issues upon model loading.\n",
    "\n",
    "<div class=\"alert alert-warning\">\n",
    "  <strong>Pro Tip</strong>: Batches should be large enough to avoid <a href=\"https://docs.ray.io/en/latest/ray-core/patterns/too-fine-grained-tasks.html\">too fine grained tasks</a> anti-pattern.\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0c403178-5b95-4849-bb9d-fe453c88f0f5",
   "metadata": {},
   "source": [
    "### Run parallel batch inference on 10 batches"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "132c9494-40f9-4ddc-904e-f1e12a5e0647",
   "metadata": {},
   "outputs": [],
   "source": [
    "prediction_refs = []\n",
    "predictions = []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bfc90802-72b2-4ac1-b61d-60b32cdbad39",
   "metadata": {},
   "outputs": [],
   "source": [
    "for batch in batches:\n",
    "    task_ref = inference_task.remote(\n",
    "        model=segformer_ref,\n",
    "        feature_extractor=segformer_feature_extractor_ref,\n",
    "        images=batch,\n",
    "    )\n",
    "    prediction_refs.append(task_ref)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5366e165-3aec-4e91-9693-53bbef29c2aa",
   "metadata": {},
   "outputs": [],
   "source": [
    "predictions = ray.get(prediction_refs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dbc0dbf2",
   "metadata": {},
   "source": [
    "Ray schedules each task to execute in parallel using the available resources. For each image:\n",
    "* call `inference_task.remote` to assign a task (returns immediately)\n",
    "* store the Object Reference `task_ref` to a list `prediction_refs`\n",
    "\n",
    "Lastly, you use `ray.get()` on the list of prediction references to retrieve the final results, and this step takes the longest to execute because it waits on all processes to complete in order to access predictions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "97d4d685-315f-4bc9-8d25-e3422e12054c",
   "metadata": {},
   "outputs": [],
   "source": [
    "predictions[0][0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0fadeb4b",
   "metadata": {},
   "source": [
    "**Coding Exercise**\n",
    "\n",
    "You have seen how the sequential version and stateless inference using Ray Tasks performs on 10 batches of 16 images each. Try scaling the number of batches as well as the number of images per batch to see the effect on performance.\n",
    "\n",
    "Hint: `BATCH_SIZE` and `N_BATCHES` is set in the Part 3 under \"Prepare batches\"\n",
    "\n",
    "Note: In order to perform inference on more than 160 total images, you need to set the `SMALL_DATA` flag to `False` to download the complete testing set. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "711cf20f-4cb2-4e16-ad66-bd273352c9d9",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Summary: stateless inference - Ray Tasks\n",
    "\n",
    "#### Key concepts\n",
    "\n",
    "<div class=\"alert alert-info\">\n",
    "  <strong>Object store</strong>: Ray's distributed shared-memory store that makes remote objects available anywhere in a Ray cluster.\n",
    "</div>\n",
    "\n",
    "<div class=\"alert alert-info\">\n",
    "  <strong>Stateless inference</strong>: inference that depends only on an inputted trained model and does not preserve state once predictions are generated.\n",
    "</div>\n",
    "\n",
    "#### Key API elements\n",
    "\n",
    "* `ray.init()` - start Ray runtime and connect to the Ray cluster\n",
    "* `@ray.remote` - functions and classes decorator specifying that it will be executed as a task (remote function) or actor (remote class) in a different process\n",
    "* `.remote` - postfix to the remote functions and classes. Remote operations are asynchronous\n",
    "* `ray.put()` - put an object in the in-memory object store and return its ID. Use this ID to pass object to any remote function or method call\n",
    "* `ray.get()` - get a remote object or a list of remote objects from the object store"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "44e2f9e2-945b-4b08-95ec-74b09b8b07d5",
   "metadata": {},
   "source": [
    "## Part 5: Distributed, stateful batch inference with Ray Actors"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9ff1bb7d-b8f9-4a2d-89d2-01446e114525",
   "metadata": {},
   "source": [
    "Moving from stateless to stateful, Ray Actors offer the advantage of holding some mutable internal state as well as avoiding reloading large models for each inference job.\n",
    "\n",
    "|<img src=\"../../_static/assets/Scaling_inference/actor_inference.png\" width=\"70%\" loading=\"lazy\">|\n",
    "|:--|\n",
    "|Stateful batch inference using Ray Actors.|"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "70563f89-c95c-4ee8-8374-0288d918d5fe",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Implement remote class for inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4030c8a1",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "@ray.remote\n",
    "class PredictionActor:\n",
    "    def __init__(self, model, feature_extractor):\n",
    "        self.model = model\n",
    "        self.feature_extractor = feature_extractor\n",
    "\n",
    "    def predict(self, images):\n",
    "        device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "        self.model.to(device)\n",
    "        self.model.eval()\n",
    "\n",
    "        inputs = self.feature_extractor(images=images, return_tensors=\"pt\")\n",
    "        with torch.no_grad():\n",
    "            outputs = self.model(pixel_values=inputs.pixel_values.to(device))\n",
    "\n",
    "        image_sizes = [image.size[::-1] for image in images]\n",
    "        segmentation_maps_postprocessed = (\n",
    "            self.feature_extractor.post_process_semantic_segmentation(\n",
    "                outputs=outputs, target_sizes=image_sizes\n",
    "            )\n",
    "        )\n",
    "\n",
    "        return [j.detach().cpu().numpy() for j in segmentation_maps_postprocessed]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "28da8dd3-08ff-47b5-8d29-be9e27a99f6b",
   "metadata": {},
   "source": [
    "Once again, `@ray.remote` declares which class will be a Ray Actor. This actor can then execute remote method calls and maintain its own internal state.\n",
    "\n",
    "Each instance of `PredictionActor`, will hold its own replica of the model, feature extractor, and device to avoid loading these every time a call to `predict` is made.\n",
    "\n",
    "Note: the `predict` function contains the same core logic as the ones you have encountered previously, with minor tweaks to fit this pattern."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cd911b90-2006-4fc9-bbbb-90e92ef79ee3",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Create list of Ray Actors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f910cd41",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "N_ACTORS = 2\n",
    "\n",
    "idle_actors = []\n",
    "for i in range(N_ACTORS):\n",
    "    idle_actors.append(\n",
    "        PredictionActor.remote(\n",
    "            model=segformer_ref, feature_extractor=segformer_feature_extractor_ref\n",
    "        )\n",
    "    )\n",
    "\n",
    "idle_actors"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b915bfba-7baf-4761-a99d-e6e36dfbfa0f",
   "metadata": {},
   "source": [
    "Create a list of `idle_actors` filled with each instance of `PredictionActor` to maintain a revolving record of which actors are available for assignment.\n",
    "\n",
    "Note: `N_ACTORS` is initally set to 2 here, which hinders performance. Ideally, you want to set the number of actors to be proportional to the amount of resources you have available, such as number of CPUs and/or GPUs."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "93d39dc3-a2ac-45a6-aab6-6a26147dd9b3",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Run parallel batch inference on 10 batches and assess scalability"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "29a80f96-3561-4435-8661-1b989c22692e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def prediction_results_postprocessing(predictions, segmentation_maps):\n",
    "    predictions.append(segmentation_maps)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c26d89d1",
   "metadata": {},
   "source": [
    "`prediction_results_postprocessing` is simple function in this tutorial and exists to abstract away the final processing step. In practice it will likely be much more complex."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aa802a4b-24cc-4716-bb94-a033b0394b5a",
   "metadata": {},
   "outputs": [],
   "source": [
    "predictions = []\n",
    "future_to_actor_mapping = {}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eeb849ae-a56b-43a9-8af5-57bd776faac5",
   "metadata": {},
   "source": [
    "To set up batch inference, create:\n",
    "* `predictions` - list of final predictions\n",
    "* `future_to_actor_mapping` - a dictionary that maps ObejctReferences to the actor that promised them"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cbce1966-c4be-406e-8df1-4ff1650b531b",
   "metadata": {},
   "outputs": [],
   "source": [
    "while batches:\n",
    "    if idle_actors:\n",
    "        actor = idle_actors.pop()\n",
    "        batch = batches.pop()\n",
    "        future = actor.predict.remote(images=batch)\n",
    "        future_to_actor_mapping[future] = actor\n",
    "    else:\n",
    "        [ready], _ = ray.wait(list(future_to_actor_mapping.keys()), num_returns=1)\n",
    "        actor = future_to_actor_mapping.pop(ready)\n",
    "        idle_actors.append(actor)\n",
    "        prediction_results_postprocessing(\n",
    "            predictions=predictions, segmentation_maps=ray.get(ready)\n",
    "        )\n",
    "\n",
    "# Process any leftover results at the end.\n",
    "for future in future_to_actor_mapping.keys():\n",
    "    prediction_results_postprocessing(\n",
    "        predictions=predictions, segmentation_maps=ray.get(future)\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "462f5b7b",
   "metadata": {},
   "source": [
    "While there remain in-flight tasks:\n",
    "* if any actors are idle\n",
    "    * take the first actor and assign it an image\n",
    "    * store the ObjectReference as a key in `future_to_actor_mapping` with the actor as a value\n",
    "* else\n",
    "    * use [ray.wait()](https://docs.ray.io/en/latest/ray-core/package-ref.html#ray-wait) to retrieve the first future to return and [limit pending tasks](https://docs.ray.io/en/master/ray-core/patterns/limit-pending-tasks.html)\n",
    "    * pop the actor that computed on the result object and add to the list of `idle_actors`\n",
    "    * send the prediction via `ray.get(ready)` to the postprocessing function\n",
    "\n",
    "Finally, to ensure that all objects have been retrieved, call `ray.get()` on any remaining futures left in the `future_to_actor_mapping` dictionary.\n",
    "\n",
    "|<img src=\"../../_static/assets/Scaling_inference/distributed_timeline.png\" width=\"70%\" loading=\"lazy\">|\n",
    "|:--|\n",
    "|Timeline of distributed bath inference where a scheduler orchestrates batch assignment as soon as a worker is available.|"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20a7b0fd-b23e-4212-bfcd-483c6a604376",
   "metadata": {},
   "outputs": [],
   "source": [
    "predictions[0][0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "67a5df5a",
   "metadata": {},
   "source": [
    "Print out the first prediction to verify that `predictions` contains results."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ec9f2c2f-f4b5-4129-a154-af23d65a9909",
   "metadata": {},
   "source": [
    "#### Optional: terminate actors after the prediction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dac36aef-83b5-4b7d-af7e-666ab359c5a5",
   "metadata": {},
   "outputs": [],
   "source": [
    "[actor.__ray_terminate__.remote() for actor in idle_actors]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c60cb444",
   "metadata": {},
   "source": [
    "**Coding Exercise**\n",
    "\n",
    "In this tutorial, the default setting for `N_ACTORS` is 2. Try setting the number of actors to the number of CPUs/GPUs you have available minus 1. How does this affect runtime performance?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b6e0209e-ca44-41d9-9c61-18150dc78285",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Summary: stateful inference with Ray Actors\n",
    "\n",
    "#### Key concepts\n",
    "\n",
    "<div class=\"alert alert-info\">\n",
    "  <strong>Stateful inference</strong>: inference carried out over stateful processes where Ray actors hold model replicas and can mutate and persist state\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7036efa3-2dd1-4b55-82e3-da81a0a33f13",
   "metadata": {},
   "source": [
    "## Part 6: Distributed batch inference with Ray ActorPool utility"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e914515e-6987-42d9-a29a-d9233db7971f",
   "metadata": {},
   "source": [
    "Building off of the previous approach, the [ActorPool](https://docs.ray.io/en/latest/ray-core/package-ref.html#ray-util-actorpool) utility wraps the list of actors to automatically handle futures management with the trade-off of giving up more granular control.\n",
    "\n",
    "|<img src=\"../../_static/assets/Scaling_inference/actor_pool.png\" width=\"70%\" loading=\"lazy\">|\n",
    "|:--|\n",
    "|Using Actor Pools for Batch Inference.|"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b35ca662-e250-449d-9e20-cd233d61452c",
   "metadata": {},
   "source": [
    "### Prepare batches"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "44e7ee7e-d79b-48b3-b2a3-b387e738a745",
   "metadata": {},
   "outputs": [],
   "source": [
    "BATCH_SIZE = 16\n",
    "N_BATCHES = 10\n",
    "\n",
    "image_indices = get_image_indices(dataset=test_dataset, n=BATCH_SIZE * N_BATCHES)\n",
    "image_indices_grouped = np.split(np.asarray(image_indices), N_BATCHES)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "596c3661-3970-49f0-8cac-4f0c41164d37",
   "metadata": {},
   "outputs": [],
   "source": [
    "batches = []\n",
    "\n",
    "for image_idx in image_indices_grouped:\n",
    "    batch = [test_dataset[int(i)][\"image\"] for i in image_idx]\n",
    "    batches.append(batch)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2fe849c0",
   "metadata": {},
   "source": [
    "Recreate the batches for inference since the last batches were popped from the existing list."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2f6a75b0-648e-4806-9ae6-f49f9ce10ba0",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Create ActorPool"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0e2baad5-ea29-44d3-88d8-28f81190caf8",
   "metadata": {},
   "outputs": [],
   "source": [
    "from ray.util.actor_pool import ActorPool"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "756bce4f-a8ed-45c4-8f28-a44d5e4c16db",
   "metadata": {},
   "outputs": [],
   "source": [
    "N_ACTORS = 2\n",
    "\n",
    "actors = [\n",
    "    PredictionActor.remote(\n",
    "        model=segformer_ref, feature_extractor=segformer_feature_extractor_ref\n",
    "    )\n",
    "    for _ in range(N_ACTORS)\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3770add5",
   "metadata": {},
   "source": [
    "Just as before, you instantiate the `N_ACTORS` of the `PredictionActor` class with model and feature extractor replicas."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cb97c5e3-0cec-41b2-ad44-0a5e4dda0b5d",
   "metadata": {},
   "outputs": [],
   "source": [
    "actor_pool = ActorPool(actors)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8b0c7732",
   "metadata": {},
   "source": [
    "Then, wrap the actors in an [ActorPool](https://docs.ray.io/en/latest/ray-core/package-ref.html#ray-util-actorpool) utility to automatically handle futures management."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "755b1725-82ce-409f-8251-6eb24964fabc",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Run parallel batch inference on 10 batches and assess scalability"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d50e7b08",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "def actor_call(actor, batch_of_images):\n",
    "    return actor.predict.remote(images=batch_of_images)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bf1e7ea3",
   "metadata": {},
   "source": [
    "`actor_call` takes in an `actor` and image and returns an ObjectRef that computes the image segmentation prediction."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "058dbfae-10b2-452e-9054-e01062ca4321",
   "metadata": {},
   "outputs": [],
   "source": [
    "predictions = []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aaacec5d-1e99-40bc-b160-642c2a32b554",
   "metadata": {},
   "outputs": [],
   "source": [
    "for segmentation_maps in actor_pool.map_unordered(actor_call, batches):\n",
    "    prediction_results_postprocessing(\n",
    "        predictions=predictions, segmentation_maps=segmentation_maps\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "029149ee",
   "metadata": {},
   "source": [
    "`map_unordered` takes in:\n",
    "- `actor_call`: a function that takes `(actor, data_item)` as argument and returns an ObjectRef computing the result over the value. The actor will be considered busy until the ObjectRef completes.\n",
    "- `data`: a list of values that `actor_call(actor, data_item)` should be applied to\n",
    "\n",
    "Note: `map_unordered` has slightly better efficiency that a similar method `actor_pool.map` since we don't care about the order of the results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5730b12c-c5e2-4935-8b06-d0d01b36e3b3",
   "metadata": {},
   "outputs": [],
   "source": [
    "predictions[0][0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e215c15a-83ee-45d0-a8f7-53ef97cd0d28",
   "metadata": {},
   "source": [
    "#### Optional: terminate actors after the prediction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1827dc61-4aed-49f5-80d4-926805cbe444",
   "metadata": {},
   "outputs": [],
   "source": [
    "if actor_pool.has_next() == False:\n",
    "    while actor_pool.has_free():\n",
    "        actor = actor_pool.pop_idle()\n",
    "        actor.__ray_terminate__.remote()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "476402b4",
   "metadata": {},
   "source": [
    "**Coding Exercise**\n",
    "\n",
    "While the `ActorPool` utility offers a good level of abstraction above orchestrating actors directly, there are [methods](https://docs.ray.io/en/latest/ray-core/package-ref.html?highlight=actorpool#ray-util-actorpool) available to you to schedule tasks, inspect in-flight jobs, and retrieve idle actors.\n",
    "\n",
    "Before terminating the actors after the prediction, try look into the actor pool by printing out which actors are idle and which tasks remain during the inferencing step."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "941f3d6d-42c8-4726-ab46-21527fd737e8",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Summary: stateful inference with Ray ActorPool utility\n",
    "\n",
    "#### Key API elements\n",
    "\n",
    "* `ActorPool()` - wraps the list of actors that run inference\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7289f04c-337c-47e1-a415-a5c5d80b61cc",
   "metadata": {},
   "source": [
    "## Part 7: Distributed batch inference with Ray Datasets\n",
    "\n",
    "Moving towards the higher-level APIs offered by Ray AIR, you'll see how to use [Ray Datasets](https://docs.ray.io/en/latest/data/dataset.html) to perform pre-processing and batching of data in parallel."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "68550c41-7755-4863-900a-d8220d9d9b55",
   "metadata": {},
   "source": [
    "|<img src=\"../../_static/assets/Scaling_inference/ray_datasets.png\" width=\"70%\" loading=\"lazy\">|\n",
    "|:--|\n",
    "|Ray Datasets replace the 'Batch preprocessing' stage.|"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d050e876-1b42-41ee-b3f2-d0caa8a31fbc",
   "metadata": {},
   "source": [
    "### Create Ray dataset with 160 images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e6ccdeb8-4c31-4752-8f16-9f13e3cf3366",
   "metadata": {},
   "outputs": [],
   "source": [
    "BATCH_SIZE = 16\n",
    "N_BATCHES = 10\n",
    "\n",
    "image_indices = get_image_indices(dataset=test_dataset, n=BATCH_SIZE * N_BATCHES)\n",
    "data = [test_dataset[i][\"image\"] for i in image_indices]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "968ed9fd",
   "metadata": {},
   "source": [
    "Once again, you prepare data batches by retrieving a random `BATCH_SIZE` number of images for every `N_BATCHES` and store them in the `data` list."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2d1f4066-7d32-4cb6-813f-9a880263b2c5",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = ray.data.from_items(data)\n",
    "dataset.show(limit=3)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7d2cf451",
   "metadata": {},
   "source": [
    "Then, you create a Ray Dataset from the list of data, and you can inspect that each item is a PIL image."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6ca4ba67-72f8-405c-b9c9-faccbe7900d1",
   "metadata": {},
   "source": [
    "### Implement class that computes predictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "54efdaa4-d0d6-473d-b647-0d7ee9a3e6be",
   "metadata": {},
   "outputs": [],
   "source": [
    "class PredictionClass:\n",
    "    def __init__(self, model, feature_extractor):\n",
    "        self.model = model\n",
    "        self.feature_extractor = feature_extractor\n",
    "\n",
    "    def __call__(self, batch):\n",
    "        device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "        self.model.to(device)\n",
    "        self.model.eval()\n",
    "\n",
    "        inputs = self.feature_extractor(images=batch, return_tensors=\"pt\")\n",
    "        with torch.no_grad():\n",
    "            outputs = self.model(pixel_values=inputs.pixel_values.to(device))\n",
    "\n",
    "        image_sizes = [image.size[::-1] for image in batch]\n",
    "        segmentation_maps_postprocessed = (\n",
    "            self.feature_extractor.post_process_semantic_segmentation(\n",
    "                outputs=outputs, target_sizes=image_sizes\n",
    "            )\n",
    "        )\n",
    "\n",
    "        return [j.detach().cpu().numpy() for j in segmentation_maps_postprocessed]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6147a110-10f8-4ee1-ab83-48116ac0e4e3",
   "metadata": {},
   "source": [
    "Each [instance](https://docs.ray.io/en/latest/data/transforming-datasets.html#transform-datasets-writing-udfs) of `PredictionClass` contains a replica of the model and feature extractor.\n",
    "\n",
    "Define the `__call__` method of the class to make it a callable class and specify the target method. The core logic of `__call__` remains the same as previous `predict()` functions.\n",
    "\n",
    "Given a `batch` (list), the [return type](https://docs.ray.io/en/latest/data/transforming-datasets.html#batch-udf-output-types) must be one of:\n",
    "\n",
    "* `pandas.DataFrame`\n",
    "* `pyarrow.Table`\n",
    "* `numpy.ndarray`\n",
    "* `Dict[str, numpy.ndarray]`\n",
    "* `list`"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "916d526b-8d86-4316-9279-9f8ada88e863",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Run parallel batch inference on 160 images and assess scalability"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "58424a37-642a-4906-86d4-f1e7b5a31058",
   "metadata": {},
   "outputs": [],
   "source": [
    "from ray.data import ActorPoolStrategy"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "744df868",
   "metadata": {},
   "source": [
    "In Ray Datasets, transformations can either be carried out by Ray Tasks or Actors. With `ActorPoolStrategy`, you can specify an [autoscaling](https://docs.ray.io/en/latest/data/transforming-datasets.html#compute-strategy) pool of `min` to `max` actors to carry out the transforms."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4daec098",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "predictions_dataset = dataset.map_batches(\n",
    "    PredictionClass,\n",
    "    batch_size=1,\n",
    "    num_gpus=0,\n",
    "    num_cpus=1,\n",
    "    compute=ActorPoolStrategy(min_size=1, max_size=2),\n",
    "    fn_constructor_args=(segformer, segformer_feature_extractor),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2824b4f4-94f3-4ac5-a06a-9bff6fed4086",
   "metadata": {},
   "source": [
    "Use the Dataset `map_batches()` [function](https://docs.ray.io/en/latest/data/api/dataset.html#ray.data.Dataset.map_batches) to apply the model to the Dataset in parallel. You can specify the batch size, any resources, as well as any autoscaling options for the actor pool.\n",
    "\n",
    "Note: don't forget to pass `fn_constructor_args` to construct `PredictionClass`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d3ac047c-8885-4c09-abae-99d2023c78a7",
   "metadata": {},
   "outputs": [],
   "source": [
    "predictions_dataset.take(limit=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dc0c1abb",
   "metadata": {},
   "source": [
    "**Coding Exercise**\n",
    "\n",
    "In this approach, you are able to control the actors using an `ActorPoolStrategy` which sets an upper and lower limit on the dynamic autoscaling of the pool. Try toggling the `min_size` and `max_size` of the actor pool in the inference step to see the effect on runtime performance."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f181c6e2",
   "metadata": {},
   "source": [
    "After running inference, you can inspect predictions to probe the resulting image array."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a34c05d8-7b3b-4524-9f03-7ec64166bdd4",
   "metadata": {},
   "source": [
    "### Summary: batch inference with Ray AIR Datasets\n",
    "\n",
    "#### Key concepts\n",
    "* parallel reading and [preprocessing](https://docs.ray.io/en/master/data/transforming-datasets.html) of the source data\n",
    "* managing the autoscaling of the ActorPool using `ActorPoolStrategy`\n",
    "* declarative key-value arguments over fine-grain control over Ray\n",
    "\n",
    "#### Key API elements\n",
    "* `map_batches` - a function to apply a transformation and/or model class to all batches"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "460ac28a-3739-4456-99a3-1579c1081db3",
   "metadata": {},
   "source": [
    "## Part 8: Distributed batch inference with Ray AIR BatchPredictor\n",
    "\n",
    "With the last approach, you will use Ray's highest-level API for batch inference: BatchPredictor."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9302f394-a33b-45e2-b527-6c74aea58f90",
   "metadata": {},
   "source": [
    "|<img src=\"../../_static/assets/Scaling_inference/air_batchpredictor.png\" width=\"70%\" loading=\"lazy\">|\n",
    "|:--|\n",
    "|Using Ray AIR's `BatchPredictor` for Batch Inference.|"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "49ef6e96-f29e-4db9-be57-0aee104a34d9",
   "metadata": {},
   "source": [
    "### Implement Predictor for image data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "62ae3c1e-31f6-4388-86a8-10ba84860f54",
   "metadata": {},
   "outputs": [],
   "source": [
    "from ray.train.predictor import Predictor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3e3024c0-f651-44f4-93f6-8964918b3bff",
   "metadata": {},
   "outputs": [],
   "source": [
    "class SemanticSegmentationPredictor(Predictor):\n",
    "    def __init__(self, model, feature_extractor):\n",
    "        super().__init__()\n",
    "        self.model = model\n",
    "        self.feature_extractor = feature_extractor\n",
    "\n",
    "    def _predict_pandas(self, batch):\n",
    "        device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "        self.model.to(device)\n",
    "        self.model.eval()\n",
    "\n",
    "        batch = [batch[\"value\"][0]]\n",
    "        inputs = self.feature_extractor(images=batch, return_tensors=\"pt\")\n",
    "\n",
    "        with torch.no_grad():\n",
    "            outputs = self.model(pixel_values=inputs.pixel_values.to(device))\n",
    "\n",
    "        image_sizes = [image.size[::-1] for image in batch]\n",
    "        segmentation_maps_postprocessed = (\n",
    "            self.feature_extractor.post_process_semantic_segmentation(\n",
    "                outputs=outputs, target_sizes=image_sizes\n",
    "            )\n",
    "        )\n",
    "\n",
    "        df = pd.DataFrame(columns=[\"segmentation_maps\"])\n",
    "        df.loc[0, \"segmentation_maps\"] = segmentation_maps_postprocessed\n",
    "\n",
    "        return df\n",
    "\n",
    "    @classmethod\n",
    "    def from_checkpoint(self, checkpoint, **kwargs):\n",
    "        checkpoint_data = checkpoint.to_dict()\n",
    "        return SemanticSegmentationPredictor(\n",
    "            model=checkpoint_data[\"model\"],\n",
    "            feature_extractor=checkpoint_data[\"feature_extractor\"],\n",
    "        )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c09cd7a2-3825-49c4-a4f6-74ba71859ad1",
   "metadata": {},
   "source": [
    "Before you run inference, define a custom [Predictor](https://docs.ray.io/en/latest/ray-air/package-ref.html#predictor), `SemanticSegmentationPredictor`, with the same replicas and core `predict()` logic as before, with a few tweaks to fit this pattern.\n",
    "\n",
    "BatchPredictor also supports multiple framework specific predictors such as TorchPredictor and TensorflowPredictor along with providing support for framework native batch conversions, the ability to resume from an AIR checkpoint, keeping columns, and aggregating batch metrics.\n",
    "\n",
    "Note: batch in `_predict_pandas` is DataFrame rather than a list."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b0790eeb-3207-4f8b-b520-e3751f93cf2b",
   "metadata": {},
   "source": [
    "### Implement BatchPredictor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e5b73b9a-b04b-4ded-8322-fbc4fa2620a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "from ray.air import Checkpoint\n",
    "from ray.train.batch_predictor import BatchPredictor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d3118f64",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "batch_predictor = BatchPredictor(\n",
    "    checkpoint=Checkpoint.from_dict(\n",
    "        {\"model\": segformer, \"feature_extractor\": segformer_feature_extractor}\n",
    "    ),\n",
    "    predictor_cls=SemanticSegmentationPredictor,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dce3a7c8",
   "metadata": {},
   "source": [
    "[`BatchPredictor`](https://docs.ray.io/en/latest/ray-air/predictors.html#batch-prediction) takes a [`Checkpoint`](https://docs.ray.io/en/latest/ray-air/package-ref.html#checkpoint) representing the saved model, and allows you to perform inference on an input dataset."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b0509bf8-d124-4a6c-9720-82bf01b03faa",
   "metadata": {},
   "source": [
    "### Run parallel batch inference on 160 images and assess scalability"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "74717154-8664-46b6-b54e-ec6d687d644f",
   "metadata": {},
   "outputs": [],
   "source": [
    "predictions_dataset = batch_predictor.predict(data=dataset, batch_size=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a13caadb",
   "metadata": {},
   "source": [
    "Perform batch inference by using the simple API `batch_predictor.predict()` without specifying *how* execution should occur."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "01247c69-c2f4-4563-a8d2-71b4e5b7b364",
   "metadata": {},
   "outputs": [],
   "source": [
    "predictions_dataset.take(limit=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9e61745f",
   "metadata": {},
   "source": [
    "Once again, you can inspect the predictions to look at the resulting segmentation maps in this Pandas Dataframe."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "425505c4",
   "metadata": {},
   "source": [
    "**Coding Exercise**\n",
    "\n",
    "In our example, we used a custom `Predictor`, but Ray AIR's BatchPredictor offers support for a number of framework specific predictors. Referring to this [user guide] for assistance, try to implement the same inferencing logic, but this time, use a [HuggingFacePredictor](https://docs.ray.io/en/master/train/api.html?highlight=huggingfacepredictor#ray.train.huggingface.HuggingFacePredictor.predict) instead."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cbfc8a4e-a38e-4b7e-8c7c-d3fd6d2fc2f8",
   "metadata": {},
   "outputs": [],
   "source": [
    "### YOUR CODE HERE ###"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "37914995",
   "metadata": {},
   "source": [
    "**Solution**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1930578f",
   "metadata": {},
   "outputs": [],
   "source": [
    "### SAMPLE IMPLEMENTATION ###\n",
    "\n",
    "import tempfile\n",
    "from ray.train.huggingface import HuggingFaceCheckpoint, HuggingFacePredictor\n",
    "\n",
    "with tempfile.TemporaryDirectory() as tmpdir:\n",
    "    huggingface_checkpoint = HuggingFaceCheckpoint.from_model(\n",
    "        model=segformer, path=tmpdir\n",
    "    )\n",
    "    predictor = BatchPredictor.from_checkpoint(\n",
    "        checkpoint=huggingface_checkpoint,\n",
    "        predictor_cls=HuggingFacePredictor,\n",
    "        feature_extractor=segformer_feature_extractor,  # passed to HF pipeline\n",
    "        task=\"image-segmentation\",  # passed to HF pipeline\n",
    "        device=-1,\n",
    "    )\n",
    "\n",
    "predictions_dataset = predictor.predict(data=dataset, batch_size=1)\n",
    "predictions_dataset.take(1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0d8a0b7a-3cff-4aea-9b92-c28200560728",
   "metadata": {},
   "source": [
    "### Shutdown Ray runtime"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fc3aae2c-491f-41a9-bd67-09395f398890",
   "metadata": {},
   "outputs": [],
   "source": [
    "ray.shutdown()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "72cdc4c3-5c2e-489a-9382-d8ca46a26a1f",
   "metadata": {},
   "source": [
    "Disconnect the worker, and terminate processes started by `ray.init()`."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fbf1a41b-0e05-427f-8cf2-ec393484bf16",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Summary: BatchPredictor\n",
    "\n",
    "#### Key API elements\n",
    "\n",
    "* `BatchPredictor` - takes a predictor class and checkpoint and provides an interface to run batch scoring on Ray datasets; this batch predictor wraps around a predictor class and executes it in a distributed way when calling `predict()`\n",
    "* `Checkpoint` - a common interface for accessing models across different AIR components and libraries\n",
    "* `predict()` - run batch scoring on Dataset\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d86ade2c-5e33-46af-804a-a10dd0e2bfc8",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Part 9: Architectures for scalable batch inference with Ray - final summary\n",
    "\n",
    "You learned about scalable batch inference patterns. Which one to choose depends on your use cases and needs. Here, you can find analysis that will help you make decisions.\n",
    "\n",
    "Overall, all patterns fall into intro 2 categories:\n",
    "\n",
    "1. **Ray core-based solutions** with task and actors (optionally with [ActorPool](https://docs.ray.io/en/latest/ray-core/package-ref.html#ray-util-actorpool) utility) are suitable for users who want to specify how Ray should execute batch inference and have fine-grained control over application behavior.\n",
    "1. **Ray AIR-based solutions** with [Datasets](https://docs.ray.io/en/latest/data/api/dataset.html) and [BatchPredictor](https://docs.ray.io/en/latest/ray-air/package-ref.html#batch-predictor) allow you to define \"what\" instead of \"how\" and delegate responsibility for parallel execution to Ray.\n",
    "\n",
    "**Key differences**\n",
    "\n",
    "| |Ray Core|Ray AI Runtime|\n",
    "|:-:|:-:|:-:|\n",
    "|**Expose parallelism**|Yes|No|\n",
    "|**Scalable from workstation to large cluster**|Yes|Yes|\n",
    "|**Integrations**|Build yourself|Out-of-the-box (PyTorch, TF and [more](https://docs.ray.io/en/latest/ray-air/package-ref.html#trainer-and-predictor-integrations))|\n",
    "|**Data pre-processing**|Build yourself|Out-of-the-box support|\n",
    "\n",
    "**Developer experience**\n",
    "\n",
    "When working with [Ray core](https://docs.ray.io/en/latest/ray-core/walkthrough.html) and [Ray AIR](https://docs.ray.io/en/latest/ray-air/getting-started.html), you will notice that they have different focus. Ray core enables you to build and scale distributed Python applications. Ray AIR is a toolkit for simplifying ML compute.\n",
    "\n",
    "| |Ray Core|Ray AI Runtime|\n",
    "|:-:|:-:|:-:|\n",
    "|**Level of control**|High|Medium|\n",
    "|**Directly program distributed apps**|Yes|No|\n",
    "|**Flexibility**|High|Medium|\n",
    "|**Ease of use**|Medium|High|\n",
    "|**Entry barrier**|Depend on use case|Low|\n",
    "\n",
    "Ray core-based batch inference gives you exposed parallelism and high level of control, because you write Python code and scale it using Ray primitives. You control compute resources assigned to individual Ray tasks. In case of actors you also have access to the internal state of the distributed objects.\n",
    "\n",
    "With Ray AIR, parallel computing primitives (tasks, actors, object store) are abstracted away and handled internally. Because of that you do not need to worry about complex parallel computing logic and focus on your use case.\n",
    "\n",
    "**Recommendations***\n",
    "\n",
    "* Use Ray core-based solutions if you have batch inference pipeline with stages that are \"embarrassingly parallel\".\n",
    "* For most batch prediction problems **BatchPredictor is recommended**. Here is why:\n",
    "    * Ray Datasets handle distributed processing, creating batches of data and ensuring compute resources utilization.\n",
    "    * On top of that [BatchPredictor](https://docs.ray.io/en/latest/ray-air/package-ref.html#batch-predictor) gives you out-of-the-box predictors, handles framework native batch conversions and gives you an option to resume from AIR checkpoint.\n",
    "    * Ray AIR libraries are connected, giving you an option to scale other parts of your pipeline in the future.\n",
    "\n",
    "*Above-mentioned recommendations are meant to provide initial advice. They are not definitive guide, because final recommendations depend on your use case. Reach out on our Slack to discuss it deeper.*\n",
    "\n",
    "|<img src=\"../../_static/assets/Scaling_inference/actor_inference.png\" width=\"90%\" loading=\"lazy\">|<img src=\"../../_static/assets/Scaling_inference/air_batchpredictor.png\" width=\"90%\" loading=\"lazy\">|\n",
    "|:-:|:-:|\n",
    "|Ray core: actors|Ray AIR: batch prediction|"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fe90b042-e31b-483e-8368-5d7b7aa1c761",
   "metadata": {},
   "source": [
    "# Connect with the Ray community\n",
    "\n",
    "You can learn and get more involved with the Ray community of developers and researchers:\n",
    "\n",
    "* [Ray documentation](https://docs.ray.io/en/latest)\n",
    "* [Official Ray Website](https://www.ray.io/): Browse the ecosystem and use this site as a hub to get the information that you need to get going and building with Ray.\n",
    "* [Join the Community on Slack](https://forms.gle/9TSdDYUgxYs8SA9e8): Find friends to discuss your new learnings in our Slack space.\n",
    "* [Use the Discussion Board](https://discuss.ray.io/): Ask questions, follow topics, and view announcements on this community forum.\n",
    "* [Join a Meetup Group](https://www.meetup.com/Bay-Area-Ray-Meetup/): Tune in on meet-ups to listen to compelling talks, get to know other users, and meet the team behind Ray.\n",
    "* [Open an Issue](https://github.com/ray-project/ray/issues/new/choose): Ray is constantly evolving to improve developer experience. Submit feature requests, bug-reports, and get help via GitHub issues."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5891f0e7-56f7-41c5-b130-6bbee6e878f8",
   "metadata": {},
   "source": [
    "<img src=\"../../_static/assets/Generic/ray_logo.png\" width=\"20%\" loading=\"lazy\">"
   ]
  }
 ],
 "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.13"
  },
  "toc-autonumbering": false,
  "toc-showmarkdowntxt": false,
  "toc-showtags": false,
  "vscode": {
   "interpreter": {
    "hash": "567405a8058597909526349386224fe35dd047505a91307e44ed44be00113429"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}