Ray_Core/Ray_Core_4_Remote_Classes_part_2.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Ray Core: Remote Classes as Actors, part 2: Tree of Actors\n",
    "\n",
    "© 2019-2023, Anyscale. All Rights Reserved\n",
    "\n",
    "## Overview\n",
    "\n",
    "Ray has a myriad of design patterns for [tasks](https://docs.ray.io/en/latest/ray-core/tasks/patterns/index.html#task-patterns) and [actors](https://docs.ray.io/en/latest/ray-core/actors/patterns/index.html). These patterns allows you to write distributed applications. In this tutorial we'll explore one of the design pattern, commonly used in Ray libraries to scale workloads.\n",
    "\n",
    "### Tree of Actors Pattern\n",
    "\n",
    "This pattern is primarily used in Ray libraries [Ray Tune](https://docs.ray.io/en/latest/tune/index.html), [Ray Train](https://docs.ray.io/en/latest/train/train.html), and [RLlib](https://docs.ray.io/en/latest/rllib/index.html) to train models in a parallel or conduct distributed HPO.\n",
    "\n",
    "In a tree of actors pattern, a collection of workers as Ray actors (or in some cases Ray tasks too), are managed by a supervisor actor. For example, you want to train multiple models, each of a different ML algorithm, at the same time, while being able to inspect its state during its training. As methology to examine simple baseline models rapidly, this pattern helps machine learning engineers to quickly build a set of baseline models for comparison.\n",
    "\n",
    "|<img src=\"https://technical-training-assets.s3.us-west-2.amazonaws.com/Ray_Core/tree_of_actors.png\" width=\"50%\" height=\"45%\">|\n",
    "|:--|\n",
    "| Common tree actor pattern with a supervisor launching a supervisor actor|\n",
    "\n",
    "This pattern facilitates **Same Data Different Function/Model (SDDF)** scaling pattern. Popular in [AutoML](https://en.wikipedia.org/wiki/Automated_machine_learning) scenarios, where you may want to train different models, with their respective algorithms, at the same time using the same dataset.\n",
    "\n",
    "|<img src=\"https://technical-training-assets.s3.us-west-2.amazonaws.com/Ray_Core/same_data_different_function.png\" width=\"85%\" height=\"75%\">|\n",
    "|:--|\n",
    "| Same data different function/model is a common scaling pattern|\n",
    "\n",
    "## Learning objectives\n",
    "\n",
    "In this this tutorial, we revisit Ray Actors, diving deeper its use for scaling patterns, and learn more about:\n",
    " * Common Ray Actors patterns used in Ray native libraries for writing distributed Actors\n",
    "   * Tree of Actors \n",
    "   * Same Data Different Functions (SDDF)\n",
    " * How to use Actors and ActorPool for Batch Inference\n",
    "\n",
    "Let's implement a simple example to illustrate this pattern."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Some preliminary imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import logging\n",
    "import time\n",
    "import ray\n",
    "import random\n",
    "from random import randint\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import pyarrow.parquet as pq\n",
    "from pprint import pprint\n",
    "\n",
    "import ray\n",
    "from ray.util.actor_pool import ActorPool"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, let's start Ray…"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if ray.is_initialized:\n",
    "    ray.shutdown()\n",
    "ray.init(logging_level=logging.ERROR)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Example 1: Supervisor and worker actor pattern\n",
    "\n",
    "**Problem**: We want to train multiple small models, say of type linear regression, each with their respective machine learning algorithm, on the same dataset, to create a set of baseline models for comparison. In our case, we use the scikit-learn [California house price](https://scikit-learn.org/stable/modules/generated/sklearn.datasets.fetch_california_housing.html) dataset and use models of type linear regression. We'll train different linear regression models: Decision tree, random forest, and xgboost regressor.  Each model is trained and evaluated based on [mean square error](https://en.wikipedia.org/wiki/Mean_squared_error)(MSE).\n",
    "\n",
    "To see this pattern used in production at scale, read [Training One Million Machine Learning Models in Record Time with Ray](https://www.anyscale.com/blog/training-one-million-machine-learning-models-in-record-time-with-ray).\n",
    "Another blog that employs this pattern is training many models per a particular feature, such as a zip code or a product SKU. [Many Models Batch Training at Scale with Ray Core](https://www.anyscale.com/blog/many-models-batch-training-at-scale-with-ray-core)\n",
    "\n",
    "|<img src=\"https://technical-training-assets.s3.us-west-2.amazonaws.com/Overview_of_Ray/California_dataset.png\" width=\"75%\" loading=\"lazy\">|\n",
    "|:--|\n",
    "|`n_samples = 20640`, target is numeric and corresponds to the average house value in units of 100k.|"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Generic model factory utility  \n",
    "\n",
    "This factory generates three different algorithms for linear regression–random forest, decsion tree, and xgboost–and trains on the same Scikit learn dataset [California housing price](https://scikit-learn.org/stable/modules/generated/sklearn.datasets.fetch_california_housing.html)\n",
    "Each training model returns its MSE score, along with time to train and relevant parameters. \n",
    "\n",
    "Each model can be in a particular state during training. The final state is `DONE`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from model_helper_utils import RFRActor\n",
    "from model_helper_utils import DTActor\n",
    "from model_helper_utils import XGBoostActor\n",
    "from model_helper_utils import RANDOM_FOREST_CONFIGS, DECISION_TREE_CONFIGS, XGBOOST_CONFIGS\n",
    "\n",
    "class ModelFactory:\n",
    "    \"\"\"\n",
    "    Mode factory to create different ML models\n",
    "    \"\"\"\n",
    "    MODEL_TYPES = [\"random_forest\", \"decision_tree\", \"xgboost\"]\n",
    "    \n",
    "    @staticmethod\n",
    "    def create_model(model_name: str) -> ray.actor.ActorHandle:\n",
    "        if model_name not in ModelFactory.MODEL_TYPES:\n",
    "            raise Exception(f\"{model_name} not supported\")\n",
    "        if model_name == \"random_forest\":\n",
    "            configs = RANDOM_FOREST_CONFIGS\n",
    "            return RFRActor.remote(configs)\n",
    "        elif model_name == \"decision_tree\":\n",
    "            configs = DECISION_TREE_CONFIGS\n",
    "            return DTActor.remote(configs)\n",
    "        else: \n",
    "            configs = XGBOOST_CONFIGS\n",
    "            return XGBoostActor.remote(configs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Create an Actor instance as a supervisor\n",
    "The supervisor creates three worker actors, each with its own respective training model ML algorithm and its training function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@ray.remote\n",
    "class Supervisor:\n",
    "    def __init__(self):\n",
    "        # Create three Actor Workers, each by its unique model type and \n",
    "        # their respective training function\n",
    "        self.worker_models =  [ModelFactory.create_model(name) for name in ModelFactory.MODEL_TYPES]\n",
    "\n",
    "    def work(self):\n",
    "        # do the train work for each Actor model\n",
    "        results = [worker_model.train_and_evaluate_model.remote() for worker_model in self.worker_models]\n",
    "        \n",
    "        # Return the final results\n",
    "        return ray.get(results)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Launch Supervisor's workers\n",
    " * create the Supervisor actor\n",
    " * launch its workers\n",
    " * fetch the final results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "supervisor = Supervisor.remote()\n",
    "results = supervisor.work.remote()\n",
    "values = ray.get(results)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let for the workers to finish by returning their `DONE` state"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "states = []\n",
    "# Wait for all models to finish\n",
    "while True:\n",
    "    for value in values:\n",
    "        states.append(value[\"state\"])\n",
    "    result = all('DONE' == e for e in states)\n",
    "    if result:\n",
    "        break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from operator import itemgetter\n",
    "sorted_by_mse = sorted(values, key=itemgetter('mse'))\n",
    "print(f\"\\nResults from three training models sorted by MSE ascending order:\")\n",
    "pprint(sorted_by_mse)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Look at the Ray Dashboard\n",
    "\n",
    "You should see Actors running as process on the workers nodes\n",
    " * Supervisor\n",
    " * Eeach model actor: RFRActor, DTActor, XGBoostActor\n",
    " \n",
    "Also, click on the `Actors` to view more metrics and data on individual Ray Actors"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise\n",
    "Change the estimators or max_depth in [model_helper_utils.py](model_helper_utils.py) "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Solution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# DECISION_TREE_CONFIGS = {\"max_depth\": 15,\n",
    "#                          \"name\": \"decision_tree\"}\n",
    "\n",
    "# RANDOM_FOREST_CONFIGS = {\"n_estimators\": 150,\n",
    "#                         \"name\": \"random_forest\"}\n",
    "\n",
    "# XGBOOST_CONFIGS = {\"max_depth\": 10,\n",
    "#                    \"n_estimators\": 150,\n",
    "#                    \"lr\": 0.1,\n",
    "#                    \"eta\": 0.3,\n",
    "#                    \"colsample_bytree\": 1,\n",
    "#                    \"name\": \"xgboost\"}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Recap\n",
    "To sum up, we went through implementing a tree actor design pattern, in which a Supervisor\n",
    "launched worker models. Each model, in our case, is a separate linear regression model training on the same data. This tree actor pattern facilitates the **Same Data Different Functions (SDDF)** scaling pattern, a common machine learning workload.\n",
    "\n",
    "For modularity, we used `model_helper_utils`, where all the model-related utility code is implemented. From all three linear regression models trained in parallel, with the same dataset, XGBoost seems to have faired well with the best MSE result. This gaves us an initial baseline of set of models to further experiment with different [hyperparameter optimization](https://en.wikipedia.org/wiki/Hyperparameter_optimization). For that I would turn to [Ray Train](https://docs.ray.io/en/latest/train/train.html) and [Ray Tune](https://docs.ray.io/en/latest/tune/index.html), whhich are part of [Ray AIR](https://docs.ray.io/en/latest/ray-air/getting-started.html)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Homework Exercise (Optional):\n",
    "\n",
    "Add another linear regression model and train four models. \n",
    "\n",
    "**HINTS**:\n",
    " * modify `model_helper_utils` to create a new super class of `ActorCls`\n",
    " * add new model configs for the new model\n",
    " * modify the `ModelFactory` class instantiate this new model\n",
    " * update the `Supervisor` to include the new model in its training\n",
    " \n",
    " You can use one of the linear models from this [array of linear models](https://www.kaggle.com/code/faressayah/practical-introduction-to-10-regression-algorithm).\n",
    " \n",
    "Have fun 😜!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Example 2: Actor-based batch inference\n",
    "\n",
    "In our first tutorial on [Ray Tasks](./ex_01_remote_funcs.ipynb) we covered a use case to parallelize batch inference. In short, we\n",
    "used the **Different Data Same Function(DDSF)** pattern. Here we will use the same pattern but with Ray Actors and [ActorPool](https://docs.ray.io/en/latest/ray-core/actors/actor-utils.html?highlight=ActorPool), which are state-based and preferred method to do batch inference. Primarily, the elements remain the same except for few modifications.\n",
    "\n",
    "**Input dataset**: This is a large collection of observations to generate predictions for. The data is usually stored in an external storage system like S3, HDFS or database, across\n",
    "many, files.\n",
    "\n",
    "**ML model**: This is a trained ML model that is usually also stored in an external storage system or in a model store.\n",
    "\n",
    "**Predictions**: These are the outputs when applying the ML model on observations. Normally, predictions are usually written back to the storage system. Unlike tasks\n",
    "doing the predictions, we employ a pool of Actors.\n",
    "\n",
    "For purpose of this tutorial, we make the following provisions:\n",
    " * create a dummy model that returns some fake prediction\n",
    " * use real-world NYC taxi data to provide large data set for batch inference\n",
    " * create a pool of actors and submit each shard to the pool.\n",
    " * return the predictions instead of writing it back to the disk\n",
    "\n",
    "As an example of scaling pattern called **Different Data Same Function (DDSF)**, also known as **Distributed Data Parallel** (DDP) paradigm, our function in this digaram is the \n",
    "pretrained **model** and the data is split and disributed as **shards**.\n",
    "\n",
    "|<img src=\"https://technical-training-assets.s3.us-west-2.amazonaws.com/Ray_Core/different_data_same_function.png\" width=\"75%\" height=\"35%\">|\n",
    "|:--|\n",
    "|Different data same function is another scaling pattern|"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "NUM_ACTORS = 6             # You can always increase the number of actors to scale\n",
    "NUM_SHARD_FILES = 12       # number of shard files you want each each actor to work on"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Our load model closure remains the same\n",
    "def load_trained_model():\n",
    "    # A fake model that predicts whether tips were given based on number of passengers in the taxi cab.\n",
    "    def model(batch: pd.DataFrame) -> pd.DataFrame:\n",
    "        # Some model payload so Ray copies the model in the shared plasma store to tasks scheduled across nodes.\n",
    "        model.payload = np.arange(10, 10_000, dtype=float)\n",
    "        #model.payload = np.arange(100, 100_000_000, dtype=float)\n",
    "        model.cls = \"regression\"\n",
    "        \n",
    "        # give a tip if 2 or more passengers\n",
    "        predict = batch[\"passenger_count\"] >= 2 \n",
    "        return pd.DataFrame({\"score\": predict})\n",
    "    \n",
    "    return model    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create a Ray Actor that stores a model reference and does the prediction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@ray.remote\n",
    "class NYCBatchPredictor:\n",
    "    def __init__(self, model):\n",
    "        self.model = model\n",
    "\n",
    "    def predict(self, shard_path):\n",
    "        # read each shard and convert to pandas\n",
    "        df = pq.read_table(shard_path).to_pandas()\n",
    "        \n",
    "        # do the inference with our model and return the result\n",
    "        result = self.model(df)\n",
    "        return result"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "1. Get our trained model instance\n",
    "2. Store it into the plasma object store"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = load_trained_model()\n",
    "model_ref = ray.put(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Fetch our NYC taxi shard files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Iterate thorough our NYC files ~ 2GB\n",
    "input_shard_files = [\n",
    "    f\"s3://anonymous@air-example-data/ursa-labs-taxi-data/downsampled_2009_full_year_data.parquet\"\n",
    "    f\"/fe41422b01c04169af2a65a83b753e0f_{i:06d}.parquet\"\n",
    "    for i in range(NUM_SHARD_FILES) ]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_shard_files"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "1. Create five Actor instances, each initialized with the same model reference\n",
    "2. Create a pool of five actors\n",
    "\n",
    "We use the Ray actor pool utility [ActorPool](https://docs.ray.io/en/latest/ray-core/actors/actor-utils.html?highlight=ActorPool#actor-pool).\n",
    "\n",
    "[Actool Pool API](https://docs.ray.io/en/latest/ray-core/package-ref.html?highlight=ActorPool#ray-util-actorpool) reference package."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "actors = [NYCBatchPredictor.remote(model_ref) for _ in range(NUM_ACTORS)]\n",
    "actors_pool = ActorPool(actors)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Submit each shard to the pool of actors for batch reference\n",
    "# The API syntax is not dissimilar to Python or Ray Multiprocessor pool APIs\n",
    "\n",
    "for shard_path in input_shard_files:\n",
    "    # Submit file shard for prediction to the pool\n",
    "    actors_pool.submit(lambda actor, shard: actor.predict.remote(shard), shard_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Iterate over finised actor's predictions\n",
    "while actors_pool.has_next():\n",
    "    r =  actors_pool.get_next()\n",
    "    print(f\"Predictions dataframe size: {len(r)} | Total score for tips: {r['score'].sum()}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Recap\n",
    "\n",
    "What we have demonstrated above is an Actor tree design pattern, commonly used in Ray for writing distributed applications. In particular, Ray's native libraries such as Train, Tune, Serve, and RLib and [Ray AIR's](https://docs.ray.io/en/latest/ray-air/getting-started.html) components use it for distributed training and tuning trials. \n",
    "\n",
    "Additionally, we implemented a DDSF scaling design pattern using an Actor-based predictor model function, using an `ActorPool` utility class instead of task.  \n",
    "\n",
    "Task-based batch inferene has an overhead cost that can be significant if the model size is large, since it has to fetch the model from the driver's plasma store. We can optimize it by using Ray actors, \n",
    "which will fetch the model just once and reuse it for all predictions assigned to the same actor in the pool."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ray.shutdown()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Homework\n",
    "\n",
    "1. Read references below."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## References"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " * [Writing your First Distributed Python Application with Ray](https://www.anyscale.com/blog/writing-your-first-distributed-python-application-with-ray)\n",
    " * [Using and Programming with Actors](https://docs.ray.io/en/latest/actors.html)\n",
    " * [Ray Asynchronous and Threaded Actors: A way to achieve concurrency](https://medium.com/@2twitme/ray-asynchronous-and-threaded-actors-a-way-to-achieve-concurrency-ad9f86145f72)\n",
    " * [Training One Million Machine Learning Models in Record Time with Ray](https://www.anyscale.com/blog/training-one-million-machine-learning-models-in-record-time-with-ray)\n",
    " * [Many Models Batch Training at Scale with Ray Core](https://www.anyscale.com/blog/many-models-batch-training-at-scale-with-ray-core)"
   ]
  }
 ],
 "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.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}