04-RayServe.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "46ebd169-cbf4-44dc-b863-0aa147b8050e",
   "metadata": {},
   "source": [
    "# Scalable, Flexible Serving with Ray Serve\n",
    "\n",
    "<img src='images/servelogo.svg' width=400>\n",
    "\n",
    "Earlier, we saw a server that performed simple request/response operation.\n",
    "\n",
    "While it's nice to have that feature without deploying additional software, that pattern is fairly well understood and easily scaled with existing technology.\n",
    "\n",
    "Ray Serve provides more value as we move to more complex patterns such as\n",
    "- stateful services\n",
    "- batching\n",
    "- composition\n",
    "- integration to model registries\n",
    "\n",
    "Let's start with a simple service deployment that could work with just about any model or logic. This will also help \"take some of the magic away\" from that 3-line serve demo featuring a `XGBoostPredictor` and `PredictorDeployment`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a837f298-7987-480f-b9e1-5a4f43d86ef9",
   "metadata": {},
   "outputs": [],
   "source": [
    "import ray\n",
    "\n",
    "ray.init(num_cpus=8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "65198e11-3d2f-49db-a682-a205e7af8529",
   "metadata": {},
   "outputs": [],
   "source": [
    "from starlette.requests import Request\n",
    "from typing import Dict\n",
    "import json\n",
    "from ray import serve\n",
    "\n",
    "class MyModel:\n",
    "    def __init__(self, demo_param: int):\n",
    "        self._demo_param = demo_param\n",
    "        \n",
    "    def predict(self, data):\n",
    "        return data * self._demo_param\n",
    "\n",
    "@serve.deployment(route_prefix=\"/\", num_replicas=2)\n",
    "class GenericDeployment:\n",
    "    def __init__(self, demo_param:int):        \n",
    "        self._model = MyModel(demo_param)\n",
    "\n",
    "    async def __call__(self, request: Request) -> Dict:\n",
    "        data = await request.json()\n",
    "        data = json.loads(data)\n",
    "        return { \"result\" : self._model.predict(data['input']) }\n",
    "\n",
    "serve.run(GenericDeployment.bind(demo_param=42))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "55178764-dd56-42f7-8857-d215cee9a686",
   "metadata": {},
   "outputs": [],
   "source": [
    "sample_json = '{\"input\":7}'\n",
    "sample_json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d6ad7074-90a0-4c88-a888-fb2b2d8320a9",
   "metadata": {},
   "outputs": [],
   "source": [
    "import requests\n",
    "\n",
    "print(requests.post(\"http://localhost:8000/\", json = sample_json).json())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f5b73f78-0d06-45a2-abb3-5121819011ed",
   "metadata": {},
   "source": [
    "> __Loading models__\n",
    "> \n",
    "> Commonly, we want to load a model once per process, not on every request. This might be because the model is large/expensive to load, or we're retrieving it from a another system like a model registry or model database, and we want to minimize traffic against that other system while caching the model locally for performance.\n",
    ">\n",
    "> Note how this RayServe example shows a pattern where we load/create the model in the service constructor (`__init__`), assign it to an instance variable, and then use that instance variable as needed for prediction.\n",
    ">\n",
    "> If you're loading very large models and want to improve performance further, there a few cool Ray tricks explained in this article, where the authors speed up loading a ~420MB flavor of the BERT language model: https://medium.com/ibm-data-ai/how-to-load-pytorch-models-340-times-faster-with-ray-8be751a6944c"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "85a0acd8-3817-4ffb-b542-59cc7ee245b5",
   "metadata": {},
   "source": [
    "Next let's look at some more interesting features that can power more complex use cases."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "292127ed-d10c-43f6-837b-f1beafb30779",
   "metadata": {},
   "outputs": [],
   "source": [
    "@serve.deployment\n",
    "class Counter:\n",
    "    def __init__(self):\n",
    "        self.count = 0\n",
    "\n",
    "    def __call__(self, *args):\n",
    "        self.count += 1\n",
    "        return {\"count\": self.count}\n",
    "\n",
    "Counter.deploy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1a372ed0-0a11-4084-b02d-365d2f9896a3",
   "metadata": {},
   "outputs": [],
   "source": [
    "import requests\n",
    "\n",
    "requests.get(\"http://127.0.0.1:8000/Counter\").json()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6e397944-3926-4c62-95b1-a0de1546afd9",
   "metadata": {},
   "source": [
    "We can also invoke these services directly from other applications or services in our Ray application"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ccbde4da-b9bf-4f0a-b7c1-bcc80cccc145",
   "metadata": {},
   "outputs": [],
   "source": [
    "Counter.get_handle().remote()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a805de54-6397-4eaf-8655-c44090c90b73",
   "metadata": {},
   "outputs": [],
   "source": [
    "ray.get(Counter.get_handle().remote())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12224fbb-1cbb-4acb-b7da-a261ae7d8d87",
   "metadata": {},
   "source": [
    "Many models achieve much better per-record performance when evaluating batches of records.\n",
    "\n",
    "We can use Ray Serve to build that batching layer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eee28cbb-65ae-4ff4-ba05-61d49911ba42",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import time\n",
    "\n",
    "@serve.deployment(route_prefix=\"/adder\")\n",
    "class BatchAdder:\n",
    "    @serve.batch(max_batch_size=4)\n",
    "    async def handle_batch(self, numbers): # numbers is a list\n",
    "        input_array = np.array(numbers)\n",
    "        print(\"Our input array has shape:\", input_array.shape)\n",
    "        # Sleep for 200ms, this could be performing CPU intensive computation\n",
    "        # in real models\n",
    "        time.sleep(0.2)\n",
    "        output_array = input_array + 1\n",
    "        return output_array.astype(int).tolist()\n",
    "\n",
    "    async def __call__(self, request):\n",
    "        return await self.handle_batch(int(request.query_params[\"number\"]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11519947-390b-4512-842f-4678a4c9c584",
   "metadata": {},
   "outputs": [],
   "source": [
    "BatchAdder.deploy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "77fc5601-b4c9-4112-b4f9-919f2b045937",
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_request(i):\n",
    "    return requests.get(\"http://localhost:8000/adder?number={}\".format(i)).text"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9c658155-700b-4e53-8353-3b228b1d71c1",
   "metadata": {},
   "outputs": [],
   "source": [
    "make_request(17)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8a411ab2-4f1c-46e5-8572-d71948e4d8fe",
   "metadata": {},
   "outputs": [],
   "source": [
    "from concurrent.futures import ThreadPoolExecutor\n",
    "\n",
    "executor = ThreadPoolExecutor()\n",
    "\n",
    "results = executor.map(make_request, range(0, 20, 2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9d8dc01d-4986-4a53-921e-587a8db7cc03",
   "metadata": {},
   "outputs": [],
   "source": [
    "list(results)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "17de24af-9b7f-409c-bfe3-8271d551019b",
   "metadata": {},
   "source": [
    "__Featurization/Model Composition__"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0cef046d-c41d-4453-b7f6-c5b941c5cc4c",
   "metadata": {},
   "source": [
    "Our pipeline will be structured as follows:\n",
    "- Input comes in, the composed model sends it to `model_one`\n",
    "- `model_one` outputs a random number between 0 and 1, if the value is\n",
    "  greater than 0.5, then the data is sent to `model_two`\n",
    "- otherwise, the data is returned to the user."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a857dae9-3d3c-4028-88bf-ce578c885f11",
   "metadata": {},
   "outputs": [],
   "source": [
    "from random import random\n",
    "\n",
    "@serve.deployment(route_prefix='/one') # remove /refactor this (route) for stateless?\n",
    "def model_one(data):\n",
    "    print(\"Model 1 called with data \", data)\n",
    "    return random()\n",
    "\n",
    "model_one.deploy()\n",
    "\n",
    "@serve.deployment\n",
    "def model_two(data):\n",
    "    print(\"Model 2 called with data \", data)\n",
    "    return data\n",
    "\n",
    "model_two.deploy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "82229466-d5e2-4e79-a390-17250df47b12",
   "metadata": {},
   "outputs": [],
   "source": [
    "resp = requests.get(\"http://127.0.0.1:8000/one\", data=\"hey!\") # stateless demo only\n",
    "resp.json()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a7c49a93-b646-4663-ac69-b88120e56233",
   "metadata": {},
   "outputs": [],
   "source": [
    "# max_concurrent_queries is optional. By default, if you pass in an async\n",
    "# function, Ray Serve sets the limit to a high number.\n",
    "@serve.deployment(max_concurrent_queries=10, route_prefix=\"/composed\")\n",
    "class ComposedModel:\n",
    "    def __init__(self):\n",
    "        self.model_one = model_one.get_handle()\n",
    "        self.model_two = model_two.get_handle()\n",
    "\n",
    "    # This method can be called concurrently!\n",
    "    async def __call__(self, starlette_request):\n",
    "        data = await starlette_request.body()\n",
    "\n",
    "        score = await self.model_one.remote(data=data)\n",
    "        if score > 0.5:\n",
    "            result = await self.model_two.remote(data=data)\n",
    "            result = {\"model_used\": 2, \"score\": score}\n",
    "        else:\n",
    "            result = {\"model_used\": 1, \"score\": score}\n",
    "\n",
    "        return result\n",
    "\n",
    "ComposedModel.deploy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c184de3e-b1f7-408f-91cc-38684a63dc1c",
   "metadata": {},
   "outputs": [],
   "source": [
    "for _ in range(5):\n",
    "    resp = requests.get(\"http://127.0.0.1:8000/composed\", data=\"hey!\")\n",
    "    print(resp.json())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "50cf7da0-3467-4e7c-ae7e-d99bde3729b4",
   "metadata": {},
   "outputs": [],
   "source": [
    "ray.shutdown()"
   ]
  }
 ],
 "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.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}