notebooks/ch_06_data_processing.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "372c8436",
   "metadata": {},
   "source": [
    "# Data Processing with Ray"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f7365d60",
   "metadata": {},
   "source": [
    "\n",
    "You can run this notebook directly in\n",
    "[Colab](https://colab.research.google.com/github/maxpumperla/learning_ray/blob/main/notebooks/ch_06_data_processing.ipynb).\n",
    "<a target=\"_blank\" href=\"https://colab.research.google.com/github/maxpumperla/learning_ray/blob/main/notebooks/ch_06_data_processing.ipynb\">\n",
    "<img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "650c6212",
   "metadata": {},
   "source": [
    "For this chapter you need to install the following dependencies:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "60a6140b",
   "metadata": {},
   "outputs": [],
   "source": [
    "! pip install \"ray[data]==2.2.0\"\n",
    "! pip install \"scikit-learn==1.0.2\"\n",
    "! pip install \"dask==2022.2.0\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "40becbf9",
   "metadata": {},
   "source": [
    "\n",
    "To import utility files for this chapter, on Colab you will also have to clone\n",
    "the repo and copy the code files to the base path of the runtime:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17aca7a4",
   "metadata": {},
   "outputs": [],
   "source": [
    "!git clone https://github.com/maxpumperla/learning_ray\n",
    "%cp -r learning_ray/notebooks/* ."
   ]
  },
  {
   "cell_type": "markdown",
   "source": [
    "![Simple Ray Data](https://raw.githubusercontent.com/maxpumperla/learning_ray/main/notebooks/images/chapter_06/AIR_data.png)\n"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "![Data Pipeline 1](https://raw.githubusercontent.com/maxpumperla/learning_ray/main/notebooks/images/chapter_06/data_pipeline_1.png)"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "![Data Pipeline 2](https://raw.githubusercontent.com/maxpumperla/learning_ray/main/notebooks/images/chapter_06/data_pipeline_2.png)"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "![Data Positioning 1](https://raw.githubusercontent.com/maxpumperla/learning_ray/main/notebooks/images/chapter_06/data_positioning_1.png)"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "![Data Positioning 2](https://raw.githubusercontent.com/maxpumperla/learning_ray/main/notebooks/images/chapter_06/data_positioning_2.png)"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "![Data Architecture](https://raw.githubusercontent.com/maxpumperla/learning_ray/main/notebooks/images/chapter_06/datasets_arch.png)"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "![Data ML Workflow](https://raw.githubusercontent.com/maxpumperla/learning_ray/main/notebooks/images/chapter_06/ml_workflow.png)"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20447273",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "import ray\n",
    "\n",
    "# Create a dataset containing integers in the range [0, 10000).\n",
    "ds = ray.data.range(10000)\n",
    "\n",
    "# Basic operations: show the size of the dataset, get a few samples, print the schema.\n",
    "print(ds.count())  # -> 10000\n",
    "print(ds.take(5))  # -> [0, 1, 2, 3, 4]\n",
    "print(ds.schema())  # -> <class 'int'>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "665a2ff7",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "# Save the dataset to a local file and load it back.\n",
    "ray.data.range(10000).write_csv(\"local_dir\")\n",
    "ds = ray.data.read_csv(\"local_dir\")\n",
    "print(ds.count())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5847b020",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "ds1 = ray.data.range(10000)\n",
    "ds2 = ray.data.range(10000)\n",
    "ds3 = ds1.union(ds2)\n",
    "print(ds3.count())  # -> 20000\n",
    "\n",
    "# Filter the combined dataset to only the even elements.\n",
    "ds3 = ds3.filter(lambda x: x % 2 == 0)\n",
    "print(ds3.count())  # -> 10000\n",
    "print(ds3.take(5))  # -> [0, 2, 4, 6, 8]\n",
    "\n",
    "# Sort the filtered dataset.\n",
    "ds3 = ds3.sort()  # <3>\n",
    "print(ds3.take(5))  # -> [0, 0, 2, 2, 4]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "09c5ecc7",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "ds1 = ray.data.range(10000)\n",
    "print(ds1.num_blocks())  # -> 200\n",
    "ds2 = ray.data.range(10000)\n",
    "print(ds2.num_blocks())  # -> 200\n",
    "ds3 = ds1.union(ds2)\n",
    "print(ds3.num_blocks())  # -> 400\n",
    "\n",
    "print(ds3.repartition(200).num_blocks())  # -> 200"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "08f6cbde",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "ds = ray.data.from_items([{\"id\": \"abc\", \"value\": 1}, {\"id\": \"def\", \"value\": 2}])\n",
    "print(ds.schema())  # -> id: string, value: int64"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8c4dcf23",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "pandas_df = ds.to_pandas()  # pandas_df will inherit the schema from our Dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bdb1d60e",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "ds = ray.data.range(10000).map(lambda x: x ** 2)\n",
    "ds.take(5)  # -> [0, 1, 4, 9, 16]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "51bd5b81",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "\n",
    "ds = ray.data.range(10000).map_batches(lambda batch: np.square(batch).tolist())\n",
    "ds.take(5)  # -> [0, 1, 4, 9, 16]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d5553d8b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_model():\n",
    "    # Returns a dummy model for this example.\n",
    "    # In reality, this would likely load some model weights onto a GPU.\n",
    "    class DummyModel:\n",
    "        def __call__(self, batch):\n",
    "            return batch\n",
    "\n",
    "    return DummyModel()\n",
    "\n",
    "\n",
    "class MLModel:\n",
    "    def __init__(self):\n",
    "        # load_model() will only run once per actor that's started.\n",
    "        self._model = load_model()\n",
    "\n",
    "    def __call__(self, batch):\n",
    "        return self._model(batch)\n",
    "\n",
    "\n",
    "ds.map_batches(MLModel, compute=\"actors\")\n",
    "\n",
    "\n",
    "cpu_intensive_preprocessing = lambda batch: batch\n",
    "gpu_intensive_inference = lambda batch: batch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1e2350d3",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "# NOTE: this only works if you create an S3 bucket and upload the data there.\n",
    "ds = (ray.data.read_parquet(\"s3://my_bucket/input_data\")\n",
    "      .map(cpu_intensive_preprocessing)\n",
    "      .map_batches(gpu_intensive_inference, compute=\"actors\", num_gpus=1)\n",
    "      .repartition(10))\n",
    "\n",
    "ds.write_parquet(\"s3://my_bucket/output_predictions\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9f3c1e05",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "# NOTE: this only works if you create an S3 bucket and upload the data there.\n",
    "ds = (ray.data.read_parquet(\"s3://my_bucket/input_data\")\n",
    "      .window(blocks_per_window=5)\n",
    "      .map(cpu_intensive_preprocessing)\n",
    "      .map_batches(gpu_intensive_inference, compute=\"actors\", num_gpus=1)\n",
    "      .repartition(10))\n",
    "ds.write_parquet(\"s3://my_bucket/output_predictions\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6cc33042",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn import datasets\n",
    "from sklearn.linear_model import SGDClassifier\n",
    "from sklearn.model_selection import train_test_split\n",
    "\n",
    "\n",
    "@ray.remote\n",
    "class TrainingWorker:\n",
    "    def __init__(self, alpha: float):\n",
    "        self._model = SGDClassifier(alpha=alpha)\n",
    "\n",
    "    def train(self, train_shard: ray.data.Dataset):\n",
    "        for i, epoch in enumerate(train_shard.iter_epochs()):\n",
    "            X, Y = zip(*list(epoch.iter_rows()))\n",
    "            self._model.partial_fit(X, Y, classes=[0, 1])\n",
    "\n",
    "        return self._model\n",
    "\n",
    "    def test(self, X_test: np.ndarray, Y_test: np.ndarray):\n",
    "        return self._model.score(X_test, Y_test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7437935a",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "ALPHA_VALS = [0.00008, 0.00009, 0.0001, 0.00011, 0.00012]\n",
    "\n",
    "print(f\"Starting {len(ALPHA_VALS)} training workers.\")\n",
    "workers = [TrainingWorker.remote(alpha) for alpha in ALPHA_VALS]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "46198d41",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "X_train, X_test, Y_train, Y_test = train_test_split(\n",
    "    *datasets.make_classification()\n",
    ")\n",
    "\n",
    "train_ds = ray.data.from_items(list(zip(X_train, Y_train)))\n",
    "shards = (train_ds.repeat(10)\n",
    "          .random_shuffle_each_window()\n",
    "          .split(len(workers), locality_hints=workers))\n",
    "\n",
    "ray.get([worker.train.remote(shard) for worker, shard in zip(workers, shards)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9969a4ef",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get validation results from each worker.\n",
    "print(ray.get([worker.test.remote(X_test, Y_test) for worker in workers]))\n",
    "\n",
    "ray.shutdown()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "adf0b4b7",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "import ray\n",
    "from ray.util.dask import enable_dask_on_ray\n",
    "\n",
    "ray.init()  # Start or connect to Ray.\n",
    "enable_dask_on_ray()  # Enable the Ray scheduler backend for Dask."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ef41df5b",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "import dask\n",
    "\n",
    "df = dask.datasets.timeseries()\n",
    "df = df[df.y > 0].groupby(\"name\").x.std()\n",
    "df.compute()  # Trigger the task graph to be evaluated."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17bb6db1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import ray\n",
    "ds = ray.data.range(10000)\n",
    "\n",
    "# Convert the Dataset to a Dask DataFrame.\n",
    "df = ds.to_dask()\n",
    "print(df.std().compute())  # -> 2886.89568\n",
    "\n",
    "# Convert the Dask DataFrame back to a Dataset.\n",
    "ds = ray.data.from_dask(df)\n",
    "print(ds.std())  # -> 2886.89568"
   ]
  }
 ],
 "metadata": {
  "jupytext": {
   "cell_metadata_filter": "-all",
   "main_language": "python",
   "notebook_metadata_filter": "-all"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}