ray-crash-course/02-Ray-Actors.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# A Guided Tour of Ray Core: Remote Stateful Classes\n",
    "\n",
    "![Anyscale Academy](../images/AnyscaleAcademyLogo.png)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[*Remote Classes*](https://docs.ray.io/en/latest/walkthrough.html#remote-classes-actors)\n",
    "involve using a `@ray.remote` decorator on a class. \n",
    "\n",
    "This implements an [*actor*](https://patterns.eecs.berkeley.edu/?page_id=258) pattern, with properties: *stateful*, *message-passing semantics*\n",
    "\n",
    "Actors are extremely powerful. They allow you to take a Python class and instantiate it as a stateful microservice that can be queried from other actors and tasks and even other Python applications.\n",
    "\n",
    "When you instantiate a remote Actor, a separate worker process is created as a worker process and becomes an Actor process on the workder node, for the purpose of running methods called on the actor. Other Ray tasks and actors can invoke its methods on that process, mutating its internal state. Actors can also be terminated manually if needed. The examples code below show all these cases.\n",
    "\n",
    "<img src=\"../images/ray_worker_actor_1.png\" height=\"30%\" width=\"60%\">\n",
    "<img src=\"../images/ray_worker_actor_2.png\" height=\"30%\" width=\"60%\">"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, let's start Ray…"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import logging\n",
    "import time\n",
    "import ray\n",
    "import random\n",
    "from random import randint\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'node_ip_address': '127.0.0.1',\n",
       " 'raylet_ip_address': '127.0.0.1',\n",
       " 'redis_address': None,\n",
       " 'object_store_address': '/tmp/ray/session_2022-03-16_15-50-05_688620_57443/sockets/plasma_store',\n",
       " 'raylet_socket_name': '/tmp/ray/session_2022-03-16_15-50-05_688620_57443/sockets/raylet',\n",
       " 'webui_url': '127.0.0.1:8265',\n",
       " 'session_dir': '/tmp/ray/session_2022-03-16_15-50-05_688620_57443',\n",
       " 'metrics_export_port': 65431,\n",
       " 'gcs_address': '127.0.0.1:60006',\n",
       " 'address': '127.0.0.1:60006',\n",
       " 'node_id': '99a61e2e13f4abf096d4bbee6af4971171e6c5375ade0a13d448e4bc'}"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ray.init(\n",
    "    ignore_reinit_error=True,\n",
    "    logging_level=logging.ERROR,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Remote Class as a Stateful Actor Pattern"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To start, we'll define a class and use the decorator: `@ray.remote`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's use Python class and convert that to a remote Actor class actor service as a Parameter Server.\n",
    "This is a common example in machine learning where you may have a central Parameter server updating gradients\n",
    "from other worker processes computing individual gradients. \n",
    "\n",
    "<img src=\"https://terrytangyuan.github.io/img/inblog/mpi-operator-1.png\" width=\"40%\" height=\"20%\">"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "@ray.remote\n",
    "class ParameterSever:\n",
    "    def __init__(self):\n",
    "        # Initialized our gradients to zero\n",
    "        self.params = np.zeros(10)\n",
    "\n",
    "    def get_params(self):\n",
    "        # Return current gradients\n",
    "        return self.params\n",
    "\n",
    "    def update_params(self, grad):\n",
    "        # Update the gradients \n",
    "        self.params -= grad"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define work or task as a function for a remote Worker process. This could be a machine learning task that\n",
    "computes gradients and sends them to the parameter server"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "@ray.remote\n",
    "def worker(ps):\n",
    "    # Iterate over some epoch\n",
    "    for i in range(100):\n",
    "        time.sleep(1.5)  # this could be your task computing gradients\n",
    "        grad = np.ones(10)\n",
    "        # update the gradients in the parameter server\n",
    "        ps.update_params.remote(grad)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Start our Parameter Server actor. This will be scheduled as a process on a remote Ray Worker"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Actor(ParameterSever, a1d3d52739b22c3eecc5588101000000)"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "param_server = ParameterSever.remote()\n",
    "param_server"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's get the initial values of the parameter server"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Initial params: [0. 0. 0. 0. 0. 0. 0. 0. 0. 0.]\n"
     ]
    }
   ],
   "source": [
    "print(f\"Initial params: {ray.get(param_server.get_params.remote())}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create Workers Nodes Computing Gradients\n",
    "Let's create three separate workers as our machine learning tasks that compute gradients.\n",
    "These will be scheduled as tasks on a Ray cluster.\n",
    "\n",
    "You can use list comprehension. Quite Pythonic!\n",
    "\n",
    "If we need more workers to scale, we can always bump them up."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[ObjectRef(c2668a65bda616c1ffffffffffffffffffffffff0100000001000000),\n",
       " ObjectRef(32d950ec0ccf9d2affffffffffffffffffffffff0100000001000000),\n",
       " ObjectRef(e0dc174c83599034ffffffffffffffffffffffff0100000001000000)]"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "[worker.remote(param_server) for _ in range(3)]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, let's iterate over a loop and query the Parameter Server \n",
    "as the workers are running independently and updating the gradients"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Updated params: [0. 0. 0. 0. 0. 0. 0. 0. 0. 0.]\n",
      "Updated params: [-3. -3. -3. -3. -3. -3. -3. -3. -3. -3.]\n",
      "Updated params: [-6. -6. -6. -6. -6. -6. -6. -6. -6. -6.]\n",
      "Updated params: [-6. -6. -6. -6. -6. -6. -6. -6. -6. -6.]\n",
      "Updated params: [-9. -9. -9. -9. -9. -9. -9. -9. -9. -9.]\n",
      "Updated params: [-12. -12. -12. -12. -12. -12. -12. -12. -12. -12.]\n",
      "Updated params: [-12. -12. -12. -12. -12. -12. -12. -12. -12. -12.]\n",
      "Updated params: [-15. -15. -15. -15. -15. -15. -15. -15. -15. -15.]\n",
      "Updated params: [-18. -18. -18. -18. -18. -18. -18. -18. -18. -18.]\n",
      "Updated params: [-18. -18. -18. -18. -18. -18. -18. -18. -18. -18.]\n",
      "Updated params: [-21. -21. -21. -21. -21. -21. -21. -21. -21. -21.]\n",
      "Updated params: [-24. -24. -24. -24. -24. -24. -24. -24. -24. -24.]\n",
      "Updated params: [-24. -24. -24. -24. -24. -24. -24. -24. -24. -24.]\n",
      "Updated params: [-27. -27. -27. -27. -27. -27. -27. -27. -27. -27.]\n",
      "Updated params: [-30. -30. -30. -30. -30. -30. -30. -30. -30. -30.]\n",
      "Updated params: [-30. -30. -30. -30. -30. -30. -30. -30. -30. -30.]\n",
      "Updated params: [-33. -33. -33. -33. -33. -33. -33. -33. -33. -33.]\n",
      "Updated params: [-36. -36. -36. -36. -36. -36. -36. -36. -36. -36.]\n",
      "Updated params: [-36. -36. -36. -36. -36. -36. -36. -36. -36. -36.]\n",
      "Updated params: [-39. -39. -39. -39. -39. -39. -39. -39. -39. -39.]\n"
     ]
    }
   ],
   "source": [
    "for _i in range(20):\n",
    "    print(f\"Updated params: {ray.get(param_server.get_params.remote())}\")\n",
    "    time.sleep(1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Tree of Actors Pattern\n",
    "\n",
    "A common pattern 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 this common pattern, tree of actors, a collection of workers as actors, are managed by a supervisor. For example, you want to train multiple models at the same time, while being able to checkpoint/inspect its state.\n",
    "\n",
    "<img src=\"https://docs.ray.io/en/latest/_images/tree-of-actors.svg\" width=\"40%\" height=\"20%\">\n",
    "\n",
    "Let's implement a simple example to illustrate this pattern."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "STATES = [\"RUNNING\", \"DONE\"]\n",
    "\n",
    "class Model:\n",
    "\n",
    "    def __init__(self, m:str):\n",
    "        self._model = m\n",
    "\n",
    "    def train(self):\n",
    "        # do some training work here\n",
    "        time.sleep(1)\n",
    "\n",
    "# Factory function to return an instance of a model type\n",
    "def model_factory(m: str):\n",
    "    return Model(m)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create a Worker Actor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "@ray.remote\n",
    "class Worker(object):\n",
    "    def __init__(self, m:str):\n",
    "        # type of a model: lr, cl, or nn\n",
    "        self._model = m                  \n",
    "        \n",
    "    def state(self) -> str:\n",
    "        return random.choice(STATES)\n",
    "    # Do the work for this model\n",
    "    def work(self) -> None:\n",
    "        model_factory(self._model).train()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create Supervisor Actor "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "@ray.remote\n",
    "class Supervisor:\n",
    "    def __init__(self):\n",
    "        # Create three Actor Workers, each by its unique model type\n",
    "        self.workers = [Worker.remote(name) for name in [\"lr\", \"cl\", \"nn\"]]\n",
    "                        \n",
    "    def work(self):\n",
    "        # do the work \n",
    "        [w.work.remote() for w in self.workers]\n",
    "        \n",
    "    def terminate(self):\n",
    "        [ray.kill(w) for w in self.workers]\n",
    "        \n",
    "    def state(self):\n",
    "        return ray.get([w.state.remote() for w in self.workers])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create a Actor instance for supervisor and launch its workers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ObjectRef(bcb4fef46b376cafe9fcb45911f1848cba1000850100000001000000)"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sup = Supervisor.remote()\n",
    "\n",
    "# Launch remote actors as workers\n",
    "sup.work.remote()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Look at the Ray Dashboard\n",
    "http://127.0.0.1:8265"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['RUNNING', 'DONE', 'RUNNING']\n",
      "['DONE', 'RUNNING', 'RUNNING']\n",
      "['RUNNING', 'RUNNING', 'RUNNING']\n",
      "['DONE', 'DONE', 'RUNNING']\n",
      "['RUNNING', 'DONE', 'RUNNING']\n",
      "['RUNNING', 'RUNNING', 'DONE']\n",
      "['DONE', 'RUNNING', 'DONE']\n",
      "['RUNNING', 'RUNNING', 'RUNNING']\n",
      "['DONE', 'DONE', 'DONE']\n"
     ]
    }
   ],
   "source": [
    "# check their status\n",
    "while True:\n",
    "    # Fetch the states of all its workers\n",
    "    states = ray.get(sup.state.remote())\n",
    "    print(states)\n",
    "    # check if all are DONE\n",
    "    result = all('DONE' == e for e in states)\n",
    "    if result:\n",
    "        # Note: Actor processes will be terminated automatically when the initial actor handle goes out of scope in Python. \n",
    "        # If we create an actor with actor_handle = ActorClass.remote(), then when actor_handle goes out of scope and is destructed, \n",
    "        # the actor process will be terminated. Note that this only applies to the original actor handle created for the actor \n",
    "        # and not to subsequent actor handles created by passing the actor handle to other tasks.\n",
    "        \n",
    "        # kill supervisors all worker manually, only for illustrtation and demo\n",
    "        sup.terminate.remote()\n",
    "\n",
    "        # kill the supervisor manually, only for illustration and demo\n",
    "        ray.kill(sup)\n",
    "        break"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Passing Actor handles to Ray Tasks\n",
    "\n",
    "You can pass actor handle instances to remote Ray tasks, which can change its \n",
    "state. The `MessageActor` keeps or clears messages, depending on the its method\n",
    "invoked."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "@ray.remote\n",
    "class MessageActor(object):\n",
    "    def __init__(self):\n",
    "        # Keep the state of the messages\n",
    "        self.messages = []\n",
    "    \n",
    "    def add_message(self, message):\n",
    "        self.messages.append(message)\n",
    "    \n",
    "    # reset and clear all messages\n",
    "    def get_and_clear_messages(self):\n",
    "        messages = self.messages\n",
    "        self.messages = []\n",
    "        return messages"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define a remote function which loops around and pushes messages to the actor, having access to a handle instance as an argument. That is, we are sending it a `MessageActor` instance handle ref as an argument to it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "@ray.remote\n",
    "def worker(message_actor, j):\n",
    "    for i in range(10):\n",
    "        time.sleep(1)\n",
    "        message_actor.add_message.remote(\n",
    "            f\"Message {i} from worker {j}.\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create a message actor."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "message_actor = MessageActor.remote()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Start 3 tasks that push messages to the actor."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[ObjectRef(c7528efcb2fd36edffffffffffffffffffffffff0100000001000000),\n",
       " ObjectRef(6efb86ef2d286c40ffffffffffffffffffffffff0100000001000000),\n",
       " ObjectRef(89af82725933373effffffffffffffffffffffff0100000001000000)]"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "[worker.remote(message_actor, j) for j in range(3)]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Periodically get the messages and print them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "New messages\n",
      ": ['Message 0 from worker 2.', 'Message 0 from worker 0.', 'Message 0 from worker 1.']\n",
      "New messages\n",
      ": ['Message 1 from worker 2.', 'Message 1 from worker 0.', 'Message 1 from worker 1.']\n",
      "New messages\n",
      ": ['Message 2 from worker 2.', 'Message 2 from worker 0.', 'Message 2 from worker 1.']\n",
      "New messages\n",
      ": ['Message 3 from worker 2.', 'Message 3 from worker 0.', 'Message 3 from worker 1.']\n",
      "New messages\n",
      ": ['Message 4 from worker 2.', 'Message 4 from worker 0.', 'Message 4 from worker 1.']\n",
      "New messages\n",
      ": ['Message 5 from worker 2.', 'Message 5 from worker 0.', 'Message 5 from worker 1.']\n",
      "New messages\n",
      ": ['Message 6 from worker 2.', 'Message 6 from worker 1.', 'Message 6 from worker 0.']\n",
      "New messages\n",
      ": ['Message 7 from worker 2.', 'Message 7 from worker 1.', 'Message 7 from worker 0.']\n",
      "New messages\n",
      ": ['Message 8 from worker 1.', 'Message 8 from worker 2.', 'Message 8 from worker 0.']\n",
      "New messages\n",
      ": ['Message 9 from worker 0.', 'Message 9 from worker 1.', 'Message 9 from worker 2.']\n"
     ]
    }
   ],
   "source": [
    "for _ in range(10):\n",
    "    new_messages = ray.get(message_actor.get_and_clear_messages.remote())\n",
    "    print(\"New messages\\n:\", new_messages)\n",
    "    time.sleep(1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, shutdown Ray"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "ray.shutdown()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Excercises\n",
    "\n",
    "1. Add a remote class, such as a logging actgor, that keeps states by logging info (may be only in memory)\n",
    "2. Implement methods that alters the state\n",
    "3. Instantiate it and call its methods"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Solution hints\n",
    "\n",
    "This solution is just a structural hint. There are few missing bits:\n",
    " * instantiation of `LoggingActor`\n",
    " * Need to use `ray.get()` to fetch the values from the object store"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import defaultdict\n",
    "@ray.remote\n",
    "class LoggingActor(object):\n",
    "    def __init__(self):\n",
    "        self.logs = defaultdict(list)\n",
    "    \n",
    "    def log(self, index, message):\n",
    "        self.logs[index].append(message)\n",
    "    \n",
    "    def get_logs(self):\n",
    "        return dict(self.logs)\n",
    "    \n",
    "@ray.remote\n",
    "def run_experiment(experiment_index, logging_actor):\n",
    "    for i in range(60):\n",
    "        time.sleep(1)\n",
    "        # Push a logging message to the actor.\n",
    "        logging_actor.log.remote(experiment_index, 'On iteration {}'.format(i))    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# logging_actor = # TODO Instantiate Actor here\n",
    "experiment_ids = []\n",
    "for i in range(3):\n",
    "    experiment_ids.append(run_experiment.remote(i, logging_actor))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "logs = logging_actor.get_logs.remote()\n",
    "# TODO use ray.get() to fetch the logs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "## 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",
    " * [Advanced Patterns and Anti-Patterns in Ray](https://docs.ray.io/en/latest/ray-design-patterns/index.htmlhttps://docs.ray.io/en/latest/ray-design-patterns/index.html)"
   ]
  }
 ],
 "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.8.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}