chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,40 @@
|
||||
load("//bazel:python.bzl", "py_test_run_all_notebooks")
|
||||
|
||||
filegroup(
|
||||
name = "core_examples",
|
||||
srcs = glob(["*.ipynb"]),
|
||||
visibility = ["//doc:__subpackages__"]
|
||||
)
|
||||
|
||||
# --------------------------------------------------------------------
|
||||
# Test the lightweight, CPU-only Ray Core example notebooks. The
|
||||
# excluded notebooks need GPU, long-running training, or live network
|
||||
# access and aren't runnable in the standard CPU CI job.
|
||||
# --------------------------------------------------------------------
|
||||
|
||||
py_test_run_all_notebooks(
|
||||
size = "large",
|
||||
include = ["*.ipynb"],
|
||||
exclude = [
|
||||
"batch_prediction.ipynb", # Requires GPU (num_gpus=1) and torch.
|
||||
"highly_parallel.ipynb", # Needs an existing multi-node cluster and belongs in a release test.
|
||||
"plot_hyperparameter.ipynb", # Model-training loop.
|
||||
"plot_parameter_server.ipynb", # Torch training loop; slow/non-deterministic.
|
||||
"plot_pong_example.ipynb", # Long-running RL (Pong) training.
|
||||
"web_crawler.ipynb", # Makes live network requests.
|
||||
],
|
||||
data = ["//doc/source/ray-core/examples:core_examples"],
|
||||
tags = [
|
||||
"exclusive",
|
||||
"team:core",
|
||||
],
|
||||
)
|
||||
|
||||
filegroup(
|
||||
name = "core_examples_ci_configs",
|
||||
srcs = glob([
|
||||
"**/ci/aws.yaml",
|
||||
"**/ci/gce.yaml",
|
||||
]),
|
||||
visibility = ["//doc:__pkg__"],
|
||||
)
|
||||
@@ -0,0 +1,345 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"attachments": {},
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# Batch Prediction with Ray Core\n",
|
||||
"\n",
|
||||
"<a id=\"try-anyscale-quickstart-batch_prediction\" href=\"https://console.anyscale.com/register/ha?render_flow=ray&utm_source=ray_docs&utm_medium=docs&utm_campaign=batch_prediction\">\n",
|
||||
" <img src=\"../../_static/img/run-on-anyscale.svg\" alt=\"try-anyscale-quickstart\">\n",
|
||||
"</a>\n",
|
||||
"<br></br>\n",
|
||||
"\n",
|
||||
"```{note}\n",
|
||||
"For a higher level API for batch inference on large datasets, see [batch inference with Ray Data](batch_inference_home). This example is for users who want more control over data sharding and execution.\n",
|
||||
"```\n",
|
||||
"\n",
|
||||
"The batch prediction is the process of using a trained model to generate predictions for a collection of observations. It has the following elements:\n",
|
||||
"* Input dataset: this is a collection of observations to generate predictions for. The data is usually stored in an external storage system like S3, HDFS or database, and can be large.\n",
|
||||
"* ML model: this is a trained ML model which is usually also stored in an external storage system.\n",
|
||||
"* Predictions: these are the outputs when applying the ML model on observations. The predictions are usually written back to the storage system.\n",
|
||||
"\n",
|
||||
"With Ray, you can build scalable batch prediction for large datasets at high prediction throughput. Ray Data provides a [higher-level API for offline batch inference](batch_inference_home), with built-in optimizations. However, for more control, you can use the lower-level Ray Core APIs. This example demonstrates batch inference with Ray Core by splitting the dataset into disjoint shards and executing them in parallel, with either Ray Tasks or Ray Actors across a Ray Cluster."
|
||||
]
|
||||
},
|
||||
{
|
||||
"attachments": {},
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Task-based batch prediction\n",
|
||||
"\n",
|
||||
"With Ray tasks, you can build a batch prediction program in this way:\n",
|
||||
"1. Loads the model\n",
|
||||
"2. Launches Ray tasks, with each taking in the model and a shard of input dataset\n",
|
||||
"3. Each worker executes predictions on the assigned shard, and writes out results\n",
|
||||
"\n",
|
||||
"Let’s take NYC taxi data in 2009 for example. Suppose we have this simple model:"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"import pandas as pd\n",
|
||||
"import numpy as np\n",
|
||||
"\n",
|
||||
"def load_model():\n",
|
||||
" # A dummy model.\n",
|
||||
" def model(batch: pd.DataFrame) -> pd.DataFrame:\n",
|
||||
" # Dummy payload so copying the model will actually copy some data\n",
|
||||
" # across nodes.\n",
|
||||
" model.payload = np.zeros(100_000_000)\n",
|
||||
" return pd.DataFrame({\"score\": batch[\"passenger_count\"] % 2 == 0})\n",
|
||||
" \n",
|
||||
" return model"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"The dataset has 12 files (one for each month) so we can naturally have each Ray task to take one file. By taking the model and a shard of input dataset (i.e. a single file), we can define a Ray remote task for prediction:"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"import pyarrow.parquet as pq\n",
|
||||
"import ray\n",
|
||||
"\n",
|
||||
"@ray.remote\n",
|
||||
"def make_prediction(model, shard_path):\n",
|
||||
" df = pq.read_table(shard_path).to_pandas()\n",
|
||||
" result = model(df)\n",
|
||||
"\n",
|
||||
" # Write out the prediction result.\n",
|
||||
" # NOTE: unless the driver will have to further process the\n",
|
||||
" # result (other than simply writing out to storage system),\n",
|
||||
" # writing out at remote task is recommended, as it can avoid\n",
|
||||
" # congesting or overloading the driver.\n",
|
||||
" # ...\n",
|
||||
"\n",
|
||||
" # Here we just return the size about the result in this example.\n",
|
||||
" return len(result)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"The driver launches all tasks for the entire input dataset. "
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"# 12 files, one for each remote task.\n",
|
||||
"input_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(12)\n",
|
||||
"]\n",
|
||||
"\n",
|
||||
"# ray.put() the model just once to local object store, and then pass the\n",
|
||||
"# reference to the remote tasks.\n",
|
||||
"model = load_model()\n",
|
||||
"model_ref = ray.put(model)\n",
|
||||
"\n",
|
||||
"result_refs = []\n",
|
||||
"\n",
|
||||
"# Launch all prediction tasks.\n",
|
||||
"for file in input_files:\n",
|
||||
" # Launch a prediction task by passing model reference and shard file to it.\n",
|
||||
" # NOTE: it would be highly inefficient if you are passing the model itself\n",
|
||||
" # like make_prediction.remote(model, file), which in order to pass the model\n",
|
||||
" # to remote node will ray.put(model) for each task, potentially overwhelming\n",
|
||||
" # the local object store and causing out-of-disk error.\n",
|
||||
" result_refs.append(make_prediction.remote(model_ref, file))\n",
|
||||
"\n",
|
||||
"results = ray.get(result_refs)\n",
|
||||
"\n",
|
||||
"# Let's check prediction output size.\n",
|
||||
"for r in results:\n",
|
||||
" print(\"Prediction output size:\", r)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"In order to not overload the cluster and cause OOM, we can control the parallelism by setting the proper resource requirement for tasks, see details about this design pattern in {doc}`/ray-core/patterns/limit-running-tasks`.\n",
|
||||
"For example, if it's easy for your to get a good estimate of the in-memory size for data loaded from external storage, you can control the parallelism by specifying the amount of memory needed for each task, e.g. launching tasks with ``make_prediction.options(memory=100*1023*1025).remote(model_ref, file)``. Ray will then do the right thing and make sure tasks scheduled to a node will not exceed its total memory."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"```{tip}\n",
|
||||
"To avoid repeatedly storing the same model into object store (this can cause Out-of-disk for driver node), use ray.put() to store the model once, and then pass the reference around.\n",
|
||||
"```\n",
|
||||
"```{tip}\n",
|
||||
"To avoid congest or overload the driver node, it’s preferable to have each task to write out the predictions (instead of returning results back to driver which actualy does nothing but write out to storage system).\n",
|
||||
"```\n",
|
||||
"\n",
|
||||
"## Actor-based batch prediction\n",
|
||||
"In the above solution, each Ray task will have to fetch the model from the driver node before it can start performing the prediction. This is an overhead cost that can be significant if the model size is large. We can optimize it by using Ray actors, which will fetch the model just once and reuse it for all tasks assigned to the actor.\n",
|
||||
"\n",
|
||||
"First, we define a callable class that’s structured with an interface (i.e. constructor) to load/cache the model, and the other to take in a file and perform prediction."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"import pandas as pd\n",
|
||||
"import pyarrow.parquet as pq\n",
|
||||
"import ray\n",
|
||||
"\n",
|
||||
"@ray.remote\n",
|
||||
"class BatchPredictor:\n",
|
||||
" def __init__(self, model):\n",
|
||||
" self.model = model\n",
|
||||
" \n",
|
||||
" def predict(self, shard_path):\n",
|
||||
" df = pq.read_table(shard_path).to_pandas()\n",
|
||||
" result =self.model(df)\n",
|
||||
"\n",
|
||||
" # Write out the prediction result.\n",
|
||||
" # NOTE: unless the driver will have to further process the\n",
|
||||
" # result (other than simply writing out to storage system),\n",
|
||||
" # writing out at remote task is recommended, as it can avoid\n",
|
||||
" # congesting or overloading the driver.\n",
|
||||
" # ...\n",
|
||||
"\n",
|
||||
" # Here we just return the size about the result in this example.\n",
|
||||
" return len(result)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"The constructor is called only once per actor worker. We use ActorPool to manage a set of actors that can receive prediction requests."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"from ray.util.actor_pool import ActorPool\n",
|
||||
"\n",
|
||||
"model = load_model()\n",
|
||||
"model_ref = ray.put(model)\n",
|
||||
"num_actors = 4\n",
|
||||
"actors = [BatchPredictor.remote(model_ref) for _ in range(num_actors)]\n",
|
||||
"pool = ActorPool(actors)\n",
|
||||
"input_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(12)\n",
|
||||
"]\n",
|
||||
"for file in input_files:\n",
|
||||
" pool.submit(lambda a, v: a.predict.remote(v), file)\n",
|
||||
"while pool.has_next():\n",
|
||||
" print(\"Prediction output size:\", pool.get_next())"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"Note that the ActorPool is fixed in size, unlike task-based approach where the number of parallel tasks can be dynamic (as long as it's not exceeding max_in_flight_tasks). To have autoscaling actor pool, you will need to use the [Ray Data batch prediction](batch_inference_home)."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Batch prediction with GPUs\n",
|
||||
"\n",
|
||||
"If your cluster has GPU nodes and your predictor can utilize the GPUs, you can direct the tasks or actors to those GPU nodes by specifying num_gpus. Ray will schedule them onto GPU nodes accordingly. On the node, you will need to move the model to GPU. The following is an example for Torch model."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"import torch\n",
|
||||
"\n",
|
||||
"@ray.remote(num_gpus=1)\n",
|
||||
"def make_torch_prediction(model: torch.nn.Module, shard_path):\n",
|
||||
" # Move model to GPU.\n",
|
||||
" model.to(torch.device(\"cuda\"))\n",
|
||||
" inputs = pq.read_table(shard_path).to_pandas().to_numpy()\n",
|
||||
"\n",
|
||||
" results = []\n",
|
||||
" # for each tensor in inputs:\n",
|
||||
" # results.append(model(tensor))\n",
|
||||
" #\n",
|
||||
" # Write out the results right in task instead of returning back\n",
|
||||
" # to the driver node (unless you have to), to avoid congest/overload\n",
|
||||
" # driver node.\n",
|
||||
" # ...\n",
|
||||
"\n",
|
||||
" # Here we just return simple/light meta information.\n",
|
||||
" return len(results)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## FAQs\n",
|
||||
"\n",
|
||||
"### How to load and pass model efficiently in Ray cluster if the model is large?\n",
|
||||
"The recommended way is to (taking task-based batch prediction for example, the actor-based is the same):\n",
|
||||
"1. let the driver load the model (e.g. from storage system)\n",
|
||||
"2. let the driver ray.put(model) to store the model into object store; and\n",
|
||||
"3. pass the same reference of the model to each remote tasks when launching them.\n",
|
||||
"The remote task will fetch the model (from driver's object store) to its local object store before start performing prediction.\n",
|
||||
"\n",
|
||||
"Note it's highly inefficient if you skip the step 2 and pass the model (instead of reference) to remote tasks. If the model is large and there are many tasks, it'll likely cause out-of-disk crash for the driver node."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"# GOOD: the model will be stored to driver's object store only once\n",
|
||||
"model = load_model()\n",
|
||||
"model_ref = ray.put(model)\n",
|
||||
"for file in input_files:\n",
|
||||
" make_prediction.remote(model_ref, file)\n",
|
||||
"\n",
|
||||
"# BAD: the same model will be stored to driver's object store repeatedly for each task\n",
|
||||
"model = load_model()\n",
|
||||
"for file in input_files:\n",
|
||||
" make_prediction.remote(model, file)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"For more details, check out {doc}`/ray-core/patterns/pass-large-arg-by-value`."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"### How to improve the GPU utilization rate?\n",
|
||||
"To keep GPUs busy, there are following things to look at:\n",
|
||||
"- **Schedule multiple tasks on the same GPU node if it has multiple GPUs**: If there are multiple GPUs on same node and a single task cannot use them all, you can direct multiple tasks to the node. This is automatically handled by Ray, e.g. if you specify num_gpus=1 and there are 4 GPUs, Ray will schedule 4 tasks to the node, provided there are enough tasks and no other resource constraints.\n",
|
||||
"- **Use actor-based approach**: as mentioned above, actor-based approach is more efficient because it reuses model initialization for many tasks, so the node will spend more time on the actual workload."
|
||||
]
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"kernelspec": {
|
||||
"display_name": "Python 3.8.10 ('venv': venv)",
|
||||
"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.10"
|
||||
},
|
||||
"orig_nbformat": 4,
|
||||
"vscode": {
|
||||
"interpreter": {
|
||||
"hash": "3c0d54d489a08ae47a06eae2fd00ff032d6cddb527c382959b7b2575f6a8167f"
|
||||
}
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 2
|
||||
}
|
||||
@@ -0,0 +1,781 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "9cdddbae",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# A Gentle Introduction to Ray Core by Example\n",
|
||||
"\n",
|
||||
"<a id=\"try-anyscale-quickstart-gentle_walkthrough\" href=\"https://console.anyscale.com/register/ha?render_flow=ray&utm_source=ray_docs&utm_medium=docs&utm_campaign=gentle_walkthrough\">\n",
|
||||
" <img src=\"../../_static/img/run-on-anyscale.svg\" alt=\"try-anyscale-quickstart\">\n",
|
||||
"</a>\n",
|
||||
"<br></br>"
|
||||
]
|
||||
},
|
||||
{
|
||||
"attachments": {},
|
||||
"cell_type": "markdown",
|
||||
"id": "9d4d0ecd",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"Implement a function in Ray Core to understand how Ray works and its basic concepts.\n",
|
||||
"Python programmers from those with less experience to those who are interested in advanced tasks,\n",
|
||||
"can start working with distributed computing using Python by learning the Ray Core API.\n",
|
||||
"\n",
|
||||
"## Install Ray\n",
|
||||
"Install Ray with the following command:"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "6115afbb",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"! pip install ray"
|
||||
]
|
||||
},
|
||||
{
|
||||
"attachments": {},
|
||||
"cell_type": "markdown",
|
||||
"id": "be9d3c98",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Ray Core\n",
|
||||
"\n",
|
||||
"Start a local cluster by running the following commands:"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"collapsed": false
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"import ray\n",
|
||||
"ray.init()"
|
||||
]
|
||||
},
|
||||
{
|
||||
"attachments": {},
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"collapsed": false
|
||||
},
|
||||
"source": [
|
||||
"Note the following lines in the output:"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "e75826f4",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"```\n",
|
||||
"... INFO services.py:1263 -- View the Ray dashboard at http://127.0.0.1:8265\n",
|
||||
"{'node_ip_address': '192.168.1.41',\n",
|
||||
"...\n",
|
||||
"'node_id': '...'}\n",
|
||||
"```"
|
||||
]
|
||||
},
|
||||
{
|
||||
"attachments": {},
|
||||
"cell_type": "markdown",
|
||||
"id": "4c837608",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"These messages indicate that the Ray cluster is working as expected. In this example output, the address of the Ray dashboard is `http://127.0.0.1:8265`. Access the Ray dashboard at the address on the first line of your output. The Ray dashboard displays information such as the number of CPU cores available and the total utilization of the current Ray application.\n",
|
||||
"This is a typical output for a laptop:\n",
|
||||
"\n",
|
||||
"```\n",
|
||||
"{'CPU': 12.0,\n",
|
||||
"'memory': 14203886388.0,\n",
|
||||
"'node:127.0.0.1': 1.0,\n",
|
||||
"'object_store_memory': 2147483648.0}\n",
|
||||
"```"
|
||||
]
|
||||
},
|
||||
{
|
||||
"attachments": {},
|
||||
"cell_type": "markdown",
|
||||
"id": "cf542ed9",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"Next, is a brief introduction to the Ray Core API, which we refer to as the Ray API.\n",
|
||||
"The Ray API builds on concepts such as decorators, functions, and classes, that are familiar to Python programmers.\n",
|
||||
"It is a universal programming interface for distributed computing. \n",
|
||||
"The engine handles the complicated work, allowing developers to use Ray with existing Python libraries and systems."
|
||||
]
|
||||
},
|
||||
{
|
||||
"attachments": {},
|
||||
"cell_type": "markdown",
|
||||
"id": "83d1a3bf",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Your First Ray API Example\n",
|
||||
"\n",
|
||||
"The following function retrieves and processes\n",
|
||||
"data from a database. The dummy `database` is a plain Python list containing the\n",
|
||||
"words of the title of the [\"Learning Ray\" book](https://www.amazon.com/Learning-Ray-Flexible-Distributed-Machine/dp/1098117220/).\n",
|
||||
"The `sleep` function pauses for a certain amount of time to simulate the cost of accessing and processing data from the database. "
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 2,
|
||||
"id": "e053331e",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"import time\n",
|
||||
"\n",
|
||||
"database = [\n",
|
||||
" \"Learning\", \"Ray\",\n",
|
||||
" \"Flexible\", \"Distributed\", \"Python\", \"for\", \"Machine\", \"Learning\"\n",
|
||||
"]\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"def retrieve(item):\n",
|
||||
" time.sleep(item / 10.)\n",
|
||||
" return item, database[item]"
|
||||
]
|
||||
},
|
||||
{
|
||||
"attachments": {},
|
||||
"cell_type": "markdown",
|
||||
"id": "2b518f4f",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"If the item with index 5 takes half a second `(5 / 10.)`, an estimate of the total runtime to retrieve all eight items sequentially is `(0+1+2+3+4+5+6+7)/10. = 2.8` seconds.\n",
|
||||
"Run the following code to get the actual time:"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 3,
|
||||
"id": "b0091149",
|
||||
"metadata": {
|
||||
"lines_to_next_cell": 2
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"Runtime: 2.82 seconds, data:\n",
|
||||
"(0, 'Learning')\n",
|
||||
"(1, 'Ray')\n",
|
||||
"(2, 'Flexible')\n",
|
||||
"(3, 'Distributed')\n",
|
||||
"(4, 'Python')\n",
|
||||
"(5, 'for')\n",
|
||||
"(6, 'Machine')\n",
|
||||
"(7, 'Learning')\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"def print_runtime(input_data, start_time):\n",
|
||||
" print(f'Runtime: {time.time() - start_time:.2f} seconds, data:')\n",
|
||||
" print(*input_data, sep=\"\\n\")\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"start = time.time()\n",
|
||||
"data = [retrieve(item) for item in range(8)]\n",
|
||||
"print_runtime(data, start)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"attachments": {},
|
||||
"cell_type": "markdown",
|
||||
"id": "97aa047d",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"The total time to run the function is 2.82 seconds in this example, but time may be different for your computer.\n",
|
||||
"Note that this basic Python version cannot run the function simultaneously."
|
||||
]
|
||||
},
|
||||
{
|
||||
"attachments": {},
|
||||
"cell_type": "markdown",
|
||||
"id": "30291db3",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"You may expect that Python list comprehensions are more efficient. The measured runtime of 2.8 seconds is actually the worst case scenario.\n",
|
||||
"Although this program \"sleeps\" for most of its runtime, it is slow because of the Global Interpreter Lock (GIL)."
|
||||
]
|
||||
},
|
||||
{
|
||||
"attachments": {},
|
||||
"cell_type": "markdown",
|
||||
"id": "2ebb32d4",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"### Ray Tasks\n",
|
||||
"\n",
|
||||
"This task can benefit from parallelization. If it is perfectly distributed, the runtime should not take much longer than the slowest subtask,\n",
|
||||
"that is, `7/10. = 0.7` seconds.\n",
|
||||
"To extend this example to run in parallel on Ray, start by using the @ray.remote decorator:"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 5,
|
||||
"id": "1e21e7c7",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"import ray \n",
|
||||
"\n",
|
||||
"\n",
|
||||
"@ray.remote\n",
|
||||
"def retrieve_task(item):\n",
|
||||
" return retrieve(item)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"attachments": {},
|
||||
"cell_type": "markdown",
|
||||
"id": "935a4062",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"With the decorator, the function `retrieve_task` becomes a :ref:`ray-remote-functions<Ray task>`_.\n",
|
||||
"A Ray task is a function that Ray executes on a different process from where\n",
|
||||
"it was called, and possibly on a different machine."
|
||||
]
|
||||
},
|
||||
{
|
||||
"attachments": {},
|
||||
"cell_type": "markdown",
|
||||
"id": "78afc80a",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"Ray is convenient to use because you can continue writing Python code,\n",
|
||||
"without having to significantly change your approach or programming style.\n",
|
||||
"Using the :func:`ray.remote()<@ray.remote>` decorator on the retrieve function is the intended use of decorators,\n",
|
||||
"and you did not modify the original code in this example."
|
||||
]
|
||||
},
|
||||
{
|
||||
"attachments": {},
|
||||
"cell_type": "markdown",
|
||||
"id": "c5dc2e18",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"To retrieve database entries and measure performance, you do not need to make many changes to the code. Here's an overview of the process:"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 6,
|
||||
"id": "a34697da",
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stderr",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"2022-12-20 13:52:34,632\tINFO worker.py:1529 -- Started a local Ray instance. View the dashboard at \u001b[1m\u001b[32m127.0.0.1:8265 \u001b[39m\u001b[22m\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"Runtime: 0.71 seconds, data:\n",
|
||||
"(0, 'Learning')\n",
|
||||
"(1, 'Ray')\n",
|
||||
"(2, 'Flexible')\n",
|
||||
"(3, 'Distributed')\n",
|
||||
"(4, 'Python')\n",
|
||||
"(5, 'for')\n",
|
||||
"(6, 'Machine')\n",
|
||||
"(7, 'Learning')\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"start = time.time()\n",
|
||||
"object_references = [\n",
|
||||
" retrieve_task.remote(item) for item in range(8)\n",
|
||||
"]\n",
|
||||
"data = ray.get(object_references)\n",
|
||||
"print_runtime(data, start)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"attachments": {},
|
||||
"cell_type": "markdown",
|
||||
"id": "8df4087c",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"Running the task in parallel requires two minor code modifications.\n",
|
||||
"To execute your Ray task remotely, you must use a `.remote()` call.\n",
|
||||
"Ray executes remote tasks asynchronously, even on a local cluster.\n",
|
||||
"The items in the `object_references` list in the code snippet do not directly contain the results.\n",
|
||||
"If you check the Python type of the first item using `type(object_references[0])`,\n",
|
||||
"you see that it is actually an `ObjectRef`.\n",
|
||||
"These object references correspond to _futures_ for which you need to request the result.\n",
|
||||
"The call :func:`ray.get()<ray.get(...)>` is for requesting the result. Whenever you call remote on a Ray task,\n",
|
||||
"it immediately returns one or more object references.\n",
|
||||
"Consider Ray tasks as the primary way of creating objects.\n",
|
||||
"The following section is an example that links multiple tasks together and allows\n",
|
||||
"Ray to pass and resolve the objects between them."
|
||||
]
|
||||
},
|
||||
{
|
||||
"attachments": {},
|
||||
"cell_type": "markdown",
|
||||
"id": "2373ddd9",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"Let's review the previous steps.\n",
|
||||
"You started with a Python function, then decorated it with `@ray.remote`, making the function a Ray task.\n",
|
||||
"Instead of directly calling the original function in the code, you called `.remote(...)` on the Ray task.\n",
|
||||
"Finally, you retrieved the results from the Ray cluster using `.get(...)`.\n",
|
||||
"Consider creating a Ray task from one of your own functions as an additional exercise."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"collapsed": false
|
||||
},
|
||||
"outputs": [],
|
||||
"source": []
|
||||
},
|
||||
{
|
||||
"attachments": {},
|
||||
"cell_type": "markdown",
|
||||
"id": "e008a500",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"Let's review the performance gain from using Ray tasks.\n",
|
||||
"On most laptops the runtime is around 0.71 seconds,\n",
|
||||
"which is slightly more than the slowest subtask, which is 0.7 seconds.\n",
|
||||
"You can further improve the program by leveraging more of Ray’s API."
|
||||
]
|
||||
},
|
||||
{
|
||||
"attachments": {},
|
||||
"cell_type": "markdown",
|
||||
"id": "54f53644",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"### Object Stores\n",
|
||||
"\n",
|
||||
"The retrieve definition directly accesses items from the `database`. While this works well on a local Ray cluster, consider how it functions on an actual cluster with multiple computers.\n",
|
||||
"A Ray cluster has a head node with a driver process and multiple worker nodes with worker processes executing tasks.\n",
|
||||
"In this scenario the database is only defined on the driver, but the worker processes need access to it to run the retrieve task.\n",
|
||||
"Ray's solution for sharing objects between the driver and workers or between workers is to use\n",
|
||||
"the `ray.put` function to place the data into Ray's distributed object store.\n",
|
||||
"In the `retrieve_task` definition, you can add a `db` argument to pass later as the `db_object_ref` object."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 7,
|
||||
"id": "da66a836",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"db_object_ref = ray.put(database)\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"@ray.remote\n",
|
||||
"def retrieve_task(item, db):\n",
|
||||
" time.sleep(item / 10.)\n",
|
||||
" return item, db[item]"
|
||||
]
|
||||
},
|
||||
{
|
||||
"attachments": {},
|
||||
"cell_type": "markdown",
|
||||
"id": "72f37eb4",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"By using the object store, you allow Ray to manage data access throughout the entire cluster.\n",
|
||||
"Although the object store involves some overhead, it improves performance for larger datasets.\n",
|
||||
"This step is crucial for a truly distributed environment.\n",
|
||||
"Rerun the example with the `retrieve_task` function to confirm that it executes as you expect."
|
||||
]
|
||||
},
|
||||
{
|
||||
"attachments": {},
|
||||
"cell_type": "markdown",
|
||||
"id": "453e312f",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"### Non-blocking calls\n",
|
||||
"\n",
|
||||
"In the previous section, you used `ray.get(object_references)` to retrieve results.\n",
|
||||
"This call blocks the driver process until all results are available.\n",
|
||||
"This dependency can cause problems if each database item takes several minutes to process.\n",
|
||||
"More efficiency gains are possible if you allow the driver process to perform other tasks while waiting for results,\n",
|
||||
"and to process results as they are completed rather than waiting for all items to finish.\n",
|
||||
"Additionally, if one of the database items cannot be retrieved due to an issue like a deadlock in the database connection,\n",
|
||||
"the driver hangs indefinitely.\n",
|
||||
"To prevent indefinite hangs, set reasonable `timeout` values when using the `wait` function.\n",
|
||||
"For example, if you want to wait less than ten times the time of the slowest data retrieval task,\n",
|
||||
"use the `wait` function to stop the task after that time has passed."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 8,
|
||||
"id": "75da06ec",
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"Runtime: 0.11 seconds, data:\n",
|
||||
"(0, 'Learning')\n",
|
||||
"(1, 'Ray')\n",
|
||||
"Runtime: 0.31 seconds, data:\n",
|
||||
"(2, 'Flexible')\n",
|
||||
"(3, 'Distributed')\n",
|
||||
"Runtime: 0.51 seconds, data:\n",
|
||||
"(4, 'Python')\n",
|
||||
"(5, 'for')\n",
|
||||
"Runtime: 0.71 seconds, data:\n",
|
||||
"(6, 'Machine')\n",
|
||||
"(7, 'Learning')\n",
|
||||
"Runtime: 0.71 seconds, data:\n",
|
||||
"(0, 'Learning')\n",
|
||||
"(1, 'Ray')\n",
|
||||
"(2, 'Flexible')\n",
|
||||
"(3, 'Distributed')\n",
|
||||
"(4, 'Python')\n",
|
||||
"(5, 'for')\n",
|
||||
"(6, 'Machine')\n",
|
||||
"(7, 'Learning')\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"start = time.time()\n",
|
||||
"object_references = [\n",
|
||||
" retrieve_task.remote(item, db_object_ref) for item in range(8)\n",
|
||||
"]\n",
|
||||
"all_data = []\n",
|
||||
"\n",
|
||||
"while len(object_references) > 0:\n",
|
||||
" finished, object_references = ray.wait(\n",
|
||||
" object_references, timeout=7.0\n",
|
||||
" )\n",
|
||||
" data = ray.get(finished)\n",
|
||||
" print_runtime(data, start)\n",
|
||||
" all_data.extend(data)\n",
|
||||
"\n",
|
||||
"print_runtime(all_data, start)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"attachments": {},
|
||||
"cell_type": "markdown",
|
||||
"id": "cf6f00c3",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"Instead of printing the results, you can use the retrieved values\n",
|
||||
"within the `while` loop to initiate new tasks on other workers."
|
||||
]
|
||||
},
|
||||
{
|
||||
"attachments": {},
|
||||
"cell_type": "markdown",
|
||||
"id": "1a9f6be5",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"### Task dependencies\n",
|
||||
"\n",
|
||||
"You may want to perform an additional processing task on the retrieved data. For example, \n",
|
||||
"use the results from the first retrieval task to query other related data from the same database (perhaps from a different table).\n",
|
||||
"The code below sets up this follow-up task and executes both the `retrieve_task` and `follow_up_task` in sequence."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 9,
|
||||
"id": "f5734bb1",
|
||||
"metadata": {
|
||||
"lines_to_next_cell": 2
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"((0, 'Learning'), (1, 'Ray'))\n",
|
||||
"((2, 'Flexible'), (3, 'Distributed'))\n",
|
||||
"((4, 'Python'), (5, 'for'))\n",
|
||||
"((6, 'Machine'), (7, 'Learning'))\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"@ray.remote\n",
|
||||
"def follow_up_task(retrieve_result):\n",
|
||||
" original_item, _ = retrieve_result\n",
|
||||
" follow_up_result = retrieve(original_item + 1)\n",
|
||||
" return retrieve_result, follow_up_result\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"retrieve_refs = [retrieve_task.remote(item, db_object_ref) for item in [0, 2, 4, 6]]\n",
|
||||
"follow_up_refs = [follow_up_task.remote(ref) for ref in retrieve_refs]\n",
|
||||
"\n",
|
||||
"result = [print(data) for data in ray.get(follow_up_refs)]"
|
||||
]
|
||||
},
|
||||
{
|
||||
"attachments": {},
|
||||
"cell_type": "markdown",
|
||||
"id": "6e4a6945",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"If you're unfamiliar with asynchronous programming, this example may not be particularly impressive.\n",
|
||||
"However, at second glance it might be surprising that the code runs at all.\n",
|
||||
"The code appears to be a regular Python function with a few list comprehensions.\n",
|
||||
"\n",
|
||||
"The function body of `follow_up_task` expects a Python tuple for its input argument `retrieve_result`.\n",
|
||||
"However, when you use the `[follow_up_task.remote(ref) for ref in retrieve_refs]` command,\n",
|
||||
"you are not passing tuples to the follow-up task.\n",
|
||||
"Instead, you are using the `retrieve_refs` to pass in Ray object references.\n",
|
||||
"\n",
|
||||
"Behind the scenes, Ray recognizes that the `follow_up_task` needs actual values,\n",
|
||||
"so it _automatically_ uses the `ray.get` function to resolve these futures.\n",
|
||||
"Additionally, Ray creates a dependency graph for all the tasks and executes them in a way that respects their dependencies.\n",
|
||||
"You don't have to explicitly tell Ray when to wait for a previous task to be completed––it infers the order of execution.\n",
|
||||
"This feature of the Ray object store is useful because you avoid copying large intermediate values\n",
|
||||
"back to the driver by passing the object references to the next task and letting Ray handle the rest."
|
||||
]
|
||||
},
|
||||
{
|
||||
"attachments": {},
|
||||
"cell_type": "markdown",
|
||||
"id": "f8722673",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"The next steps in the process are only scheduled once the tasks specifically designed to retrieve information are completed.\n",
|
||||
"In fact, if `retrieve_refs` was called `retrieve_result`, you might not have noticed this crucial and intentional naming nuance. Ray allows you to concentrate on your work rather than the technicalities of cluster computing.\n",
|
||||
"The dependency graph for the two tasks looks like this:"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "d03b8e46",
|
||||
"metadata": {
|
||||
"lines_to_next_cell": 2
|
||||
},
|
||||
"source": [
|
||||
""
|
||||
]
|
||||
},
|
||||
{
|
||||
"attachments": {},
|
||||
"cell_type": "markdown",
|
||||
"id": "e4001edb",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"### Ray Actors\n",
|
||||
"\n",
|
||||
"This example covers one more significant aspect of Ray Core.\n",
|
||||
"Up until this step, everything is essentially a function.\n",
|
||||
"You used the `@ray.remote` decorator to make certain functions remote, but aside from that, you only used standard Python."
|
||||
]
|
||||
},
|
||||
{
|
||||
"attachments": {},
|
||||
"cell_type": "markdown",
|
||||
"id": "2a7a6e69",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"If you want to keep track of how often the database is being queried, you could count the results of the retrieve tasks.\n",
|
||||
"However, is there a more efficient way to do this? Ideally, you want to track this in a distributed manner that can handle a large amount of data.\n",
|
||||
"Ray provides a solution with actors, which run stateful computations on a cluster and can also communicate with each other.\n",
|
||||
"Similar to how you create Ray tasks using decorated functions, create Ray actors using decorated Python classes.\n",
|
||||
"Therefore, you can create a simple counter using a Ray actor to track the number of database calls."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 10,
|
||||
"id": "717df7d0",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"@ray.remote\n",
|
||||
"class DataTracker:\n",
|
||||
" def __init__(self):\n",
|
||||
" self._counts = 0\n",
|
||||
"\n",
|
||||
" def increment(self):\n",
|
||||
" self._counts += 1\n",
|
||||
"\n",
|
||||
" def counts(self):\n",
|
||||
" return self._counts"
|
||||
]
|
||||
},
|
||||
{
|
||||
"attachments": {},
|
||||
"cell_type": "markdown",
|
||||
"id": "2e003cc6",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"The DataTracker class becomes an actor when you give it the `ray.remote` decorator. This actor is capable of tracking state,\n",
|
||||
"such as a count, and its methods are Ray actor tasks that you can invoke in the same way as functions using `.remote()`.\n",
|
||||
"Modify the retrieve_task to incorporate this actor."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 11,
|
||||
"id": "6843b8d9",
|
||||
"metadata": {
|
||||
"lines_to_next_cell": 2
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"[(0, 'Learning'), (1, 'Ray'), (2, 'Flexible'), (3, 'Distributed'), (4, 'Python'), (5, 'for'), (6, 'Machine'), (7, 'Learning')]\n",
|
||||
"8\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"@ray.remote\n",
|
||||
"def retrieve_tracker_task(item, tracker, db):\n",
|
||||
" time.sleep(item / 10.)\n",
|
||||
" tracker.increment.remote()\n",
|
||||
" return item, db[item]\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"tracker = DataTracker.remote()\n",
|
||||
"\n",
|
||||
"object_references = [\n",
|
||||
" retrieve_tracker_task.remote(item, tracker, db_object_ref) for item in range(8)\n",
|
||||
"]\n",
|
||||
"data = ray.get(object_references)\n",
|
||||
"\n",
|
||||
"print(data)\n",
|
||||
"print(ray.get(tracker.counts.remote()))"
|
||||
]
|
||||
},
|
||||
{
|
||||
"attachments": {},
|
||||
"cell_type": "markdown",
|
||||
"id": "ae886162",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"As expected, the outcome of this calculation is 8.\n",
|
||||
"Although you don't need actors to perform this calculation, this demonstrates a way to maintain state across the cluster, possibly involving multiple tasks.\n",
|
||||
"In fact, you could pass the actor into any related task or even into the constructor of a different actor.\n",
|
||||
"The Ray API is flexible, allowing for limitless possibilities.\n",
|
||||
"It's rare for distributed Python tools to allow for stateful computations,\n",
|
||||
"which is especially useful for running complex distributed algorithms such as reinforcement learning."
|
||||
]
|
||||
},
|
||||
{
|
||||
"attachments": {},
|
||||
"cell_type": "markdown",
|
||||
"id": "fb8bd0f5",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Summary\n",
|
||||
"\n",
|
||||
"In this example, you only used six API methods.\n",
|
||||
"These included `ray.init()` to initiate the cluster, `@ray.remote` to transform functions and classes into tasks and actors,\n",
|
||||
"`ray.put()` to transfer values into Ray's object store, and `ray.get()` to retrieve objects from the cluster.\n",
|
||||
"Additionally, you used `.remote()` on actor methods or tasks to execute code on the cluster, and `ray.wait` to prevent blocking calls."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "0d936caa",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": []
|
||||
},
|
||||
{
|
||||
"attachments": {},
|
||||
"cell_type": "markdown",
|
||||
"id": "0d8693e9",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"The Ray API consists of more than these six calls, but these six are powerful, if you're just starting out.\n",
|
||||
"To summarize more generally, the methods are as follows:"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "f2fb14bd",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"- `ray.init()`: Initializes your Ray cluster. Pass in an address to connect to an existing cluster.\n",
|
||||
"- `@ray.remote`: Turns functions into tasks and classes into actors.\n",
|
||||
"- `ray.put()`: Puts values into Ray’s object store.\n",
|
||||
"- `ray.get()`: Gets values from the object store. Returns the values you’ve put there or that were computed by a task or actor.\n",
|
||||
"- `.remote()`: Runs actor methods or tasks on your Ray cluster and is used to instantiate actors.\n",
|
||||
"- `ray.wait()`: Returns two lists of object references, one with finished tasks we’re waiting for and one with unfinished tasks."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"collapsed": false
|
||||
},
|
||||
"source": [
|
||||
"## Want to learn more?\n",
|
||||
"\n",
|
||||
"This example is a simplified version of the Ray Core walkthrough of [our \"Learning Ray\" book](https://maxpumperla.com/learning_ray/).\n",
|
||||
"If you liked it, check out the [Ray Core Examples Gallery](./overview.rst) or some of the ML workloads in our [Use Case Gallery](../../ray-overview/use-cases.rst)."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"collapsed": false
|
||||
},
|
||||
"outputs": [],
|
||||
"source": []
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"jupytext": {
|
||||
"cell_metadata_filter": "-all",
|
||||
"main_language": "python",
|
||||
"notebook_metadata_filter": "-all"
|
||||
},
|
||||
"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"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 5
|
||||
}
|
||||
@@ -0,0 +1,338 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "515dffba",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# Using Ray for Highly Parallelizable Tasks\n",
|
||||
"\n",
|
||||
"<a id=\"try-anyscale-quickstart-highly_parallel\" href=\"https://console.anyscale.com/register/ha?render_flow=ray&utm_source=ray_docs&utm_medium=docs&utm_campaign=highly_parallel\">\n",
|
||||
" <img src=\"../../_static/img/run-on-anyscale.svg\" alt=\"try-anyscale-quickstart\">\n",
|
||||
"</a>\n",
|
||||
"<br></br>\n",
|
||||
"\n",
|
||||
"While Ray can be used for very complex parallelization tasks,\n",
|
||||
"often we just want to do something simple in parallel.\n",
|
||||
"For example, we may have 100,000 time series to process with exactly the same algorithm,\n",
|
||||
"and each one takes a minute of processing.\n",
|
||||
"\n",
|
||||
"Clearly running it on a single processor is prohibitive: this would take 70 days.\n",
|
||||
"Even if we managed to use 8 processors on a single machine,\n",
|
||||
"that would bring it down to 9 days. But if we can use 8 machines, each with 16 cores,\n",
|
||||
"it can be done in about 12 hours.\n",
|
||||
"\n",
|
||||
"How can we use Ray for these types of task? \n",
|
||||
"\n",
|
||||
"We take the simple example of computing the digits of pi.\n",
|
||||
"The algorithm is simple: generate random x and y, and if ``x^2 + y^2 < 1``, it's\n",
|
||||
"inside the circle, we count as in. This actually turns out to be pi/4\n",
|
||||
"(remembering your high school math).\n",
|
||||
"\n",
|
||||
"The following code (and this notebook) assumes you have already set up your Ray cluster and that you are running on the head node. For more details on how to set up a Ray cluster please see [Ray Clusters Getting Started](https://docs.ray.io/en/master/cluster/getting-started.html). \n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 2,
|
||||
"id": "8e3e7c4f",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"import ray\n",
|
||||
"import random\n",
|
||||
"import time\n",
|
||||
"import math\n",
|
||||
"from fractions import Fraction"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "92d2461b",
|
||||
"metadata": {
|
||||
"scrolled": true,
|
||||
"tags": [
|
||||
"remove-output"
|
||||
]
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"# Let's start Ray\n",
|
||||
"ray.init(address='auto')"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "b96f2eb9",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"We use the ``@ray.remote`` decorator to create a Ray task.\n",
|
||||
"A task is like a function, except the result is returned asynchronously.\n",
|
||||
"\n",
|
||||
"It also may not run on the local machine, it may run elsewhere in the cluster.\n",
|
||||
"This way you can run multiple tasks in parallel,\n",
|
||||
"beyond the limit of the number of processors you can have in a single machine."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 4,
|
||||
"id": "ece9887c",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"@ray.remote\n",
|
||||
"def pi4_sample(sample_count):\n",
|
||||
" \"\"\"pi4_sample runs sample_count experiments, and returns the \n",
|
||||
" fraction of time it was inside the circle. \n",
|
||||
" \"\"\"\n",
|
||||
" in_count = 0\n",
|
||||
" for i in range(sample_count):\n",
|
||||
" x = random.random()\n",
|
||||
" y = random.random()\n",
|
||||
" if x*x + y*y <= 1:\n",
|
||||
" in_count += 1\n",
|
||||
" return Fraction(in_count, sample_count)\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "05bf8675",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"To get the result of a future, we use ray.get() which \n",
|
||||
"blocks until the result is complete. "
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 5,
|
||||
"id": "9d9a3509",
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"Running 1000000 tests took 1.4935967922210693 seconds\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"SAMPLE_COUNT = 1000 * 1000\n",
|
||||
"start = time.time() \n",
|
||||
"future = pi4_sample.remote(sample_count = SAMPLE_COUNT)\n",
|
||||
"pi4 = ray.get(future)\n",
|
||||
"end = time.time()\n",
|
||||
"dur = end - start\n",
|
||||
"print(f'Running {SAMPLE_COUNT} tests took {dur} seconds')"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "cc17429d",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"Now let's see how good our approximation is."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 7,
|
||||
"id": "42d4c464",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"pi = pi4 * 4"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 8,
|
||||
"id": "4009bee0",
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"3.143024"
|
||||
]
|
||||
},
|
||||
"execution_count": 8,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"float(pi)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 9,
|
||||
"id": "d19155d6",
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"0.0004554042254233261"
|
||||
]
|
||||
},
|
||||
"execution_count": 9,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"abs(pi-math.pi)/pi"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "ddb3b095",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"Meh. A little off -- that's barely 4 decimal places.\n",
|
||||
"Why don't we do it a 100,000 times as much? Let's do 100 billion!"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 10,
|
||||
"id": "b7b9cff9",
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"Doing 100000 batches\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"FULL_SAMPLE_COUNT = 100 * 1000 * 1000 * 1000 # 100 billion samples! \n",
|
||||
"BATCHES = int(FULL_SAMPLE_COUNT / SAMPLE_COUNT)\n",
|
||||
"print(f'Doing {BATCHES} batches')\n",
|
||||
"results = []\n",
|
||||
"for _ in range(BATCHES):\n",
|
||||
" results.append(pi4_sample.remote(sample_count = SAMPLE_COUNT))\n",
|
||||
"output = ray.get(results)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "94264de4",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"Notice that in the above, we generated a list with 100,000 futures.\n",
|
||||
"Now all we do is have to do is wait for the result.\n",
|
||||
"\n",
|
||||
"Depending on your ray cluster's size, this might take a few minutes.\n",
|
||||
"But to give you some idea, if we were to do it on a single machine,\n",
|
||||
"when I ran this it took 0.4 seconds.\n",
|
||||
"\n",
|
||||
"On a single core, that means we're looking at 0.4 * 100000 = about 11 hours. \n",
|
||||
"\n",
|
||||
"Here's what the Dashboard looks like: \n",
|
||||
"\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"So now, rather than just a single core working on this,\n",
|
||||
"I have 168 working on the task together. And its ~80% efficient."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 12,
|
||||
"id": "76eba02d",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"pi = sum(output)*4/len(output)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 13,
|
||||
"id": "ede2bd8c",
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"3.14159518188"
|
||||
]
|
||||
},
|
||||
"execution_count": 13,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"float(pi)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 14,
|
||||
"id": "bb62cb27",
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"8.047791203506436e-07"
|
||||
]
|
||||
},
|
||||
"execution_count": 14,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"abs(pi-math.pi)/pi"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "30d12e50",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"Not bad at all -- we're off by a millionth. "
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "1b36747b",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": []
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"celltoolbar": "Tags",
|
||||
"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.13"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 5
|
||||
}
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 24 KiB |
File diff suppressed because one or more lines are too long
|
After Width: | Height: | Size: 9.7 KiB |
File diff suppressed because one or more lines are too long
|
After Width: | Height: | Size: 6.6 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 24 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 59 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 15 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 36 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 11 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 37 KiB |
@@ -0,0 +1,113 @@
|
||||
# An unique identifier for the head node and workers of this cluster.
|
||||
cluster_name: lm-cluster
|
||||
|
||||
# The minimum number of workers nodes to launch in addition to the head
|
||||
# node. This number should be >= 0.
|
||||
min_workers: 1
|
||||
|
||||
# The maximum number of workers nodes to launch in addition to the head
|
||||
# node. This takes precedence over min_workers.
|
||||
max_workers: 2
|
||||
|
||||
|
||||
# If a node is idle for this many minutes, it will be removed.
|
||||
idle_timeout_minutes: 5
|
||||
|
||||
# Cloud-provider specific configuration.
|
||||
provider:
|
||||
type: aws
|
||||
region: us-west-2
|
||||
# Availability zone(s), comma-separated, that nodes may be launched in.
|
||||
# Nodes will be launched in the first listed availability zone and will
|
||||
# be tried in the subsequent availability zones if launching fails.
|
||||
availability_zone: us-west-2a,us-west-2b
|
||||
|
||||
# How Ray will authenticate with newly launched nodes.
|
||||
auth:
|
||||
ssh_user: ubuntu
|
||||
# By default Ray creates a new private keypair, but you can also use your own.
|
||||
# If you do so, make sure to also set "KeyName" in the head and worker node
|
||||
# configurations below.
|
||||
# ssh_private_key: /path/to/your/key.pem
|
||||
|
||||
# Provider-specific config for the head node, e.g. instance type. By default
|
||||
# Ray will auto-configure unspecified fields such as SubnetId and KeyName.
|
||||
# For more documentation on available fields, see:
|
||||
# http://boto3.readthedocs.io/en/latest/reference/services/ec2.html#EC2.ServiceResource.create_instances
|
||||
head_node:
|
||||
InstanceType: m5.xlarge
|
||||
ImageId: ami-0b294f219d14e6a82 # Deep Learning AMI (Ubuntu) Version 21.0
|
||||
SecurityGroupIds:
|
||||
- "{{SecurityGroupId}}"
|
||||
# You can provision additional disk space with a conf as follows
|
||||
BlockDeviceMappings:
|
||||
- DeviceName: /dev/sda1
|
||||
Ebs:
|
||||
VolumeSize: 140
|
||||
|
||||
# Additional options in the boto docs.
|
||||
|
||||
# Provider-specific config for worker nodes, e.g. instance type. By default
|
||||
# Ray will auto-configure unspecified fields such as SubnetId and KeyName.
|
||||
# For more documentation on available fields, see:
|
||||
# http://boto3.readthedocs.io/en/latest/reference/services/ec2.html#EC2.ServiceResource.create_instances
|
||||
worker_nodes:
|
||||
InstanceType: p3.2xlarge
|
||||
ImageId: ami-0b294f219d14e6a82 # Deep Learning AMI (Ubuntu) Version 21.0
|
||||
SecurityGroupIds:
|
||||
- "{{SecurityGroupId}}"
|
||||
# Run workers on spot by default. Comment this out to use on-demand.
|
||||
InstanceMarketOptions:
|
||||
MarketType: spot
|
||||
# Additional options can be found in the boto docs, e.g.
|
||||
# SpotOptions:
|
||||
# MaxPrice: MAX_HOURLY_PRICE
|
||||
|
||||
# Additional options in the boto docs.
|
||||
|
||||
# List of shell commands to run to set up nodes.
|
||||
setup_commands:
|
||||
# Note: if you're developing Ray, you probably want to create an AMI that
|
||||
# has your Ray repo pre-cloned. Then, you can replace the pip installs
|
||||
# below with a git checkout <your_sha> (and possibly a recompile).
|
||||
- echo 'export PATH="$HOME/anaconda3/envs/pytorch_p36/bin:$PATH"' >> ~/.bashrc;
|
||||
source ~/.bashrc;
|
||||
pip install -U ray;
|
||||
pip install -U fairseq==0.8.0;
|
||||
- sudo kill -9 `sudo lsof /var/lib/dpkg/lock-frontend | awk '{print $2}' | tail -n 1`;
|
||||
sudo pkill -9 apt-get;
|
||||
sudo pkill -9 dpkg;
|
||||
sudo dpkg --configure -a;
|
||||
sudo apt-get -y install binutils;
|
||||
cd $HOME;
|
||||
git clone https://github.com/aws/efs-utils;
|
||||
cd $HOME/efs-utils;
|
||||
./build-deb.sh;
|
||||
sudo apt-get -y install ./build/amazon-efs-utils*deb;
|
||||
cd $HOME;
|
||||
mkdir efs;
|
||||
sudo mount -t efs {{FileSystemId}}:/ efs;
|
||||
sudo chmod 777 efs;
|
||||
|
||||
# Custom commands that will be run on the head node after common setup.
|
||||
head_setup_commands:
|
||||
- pip install boto3>=1.4.8 # 1.4.8 adds InstanceMarketOptions
|
||||
|
||||
# Custom commands that will be run on worker nodes after common setup.
|
||||
worker_setup_commands: []
|
||||
|
||||
# Command to start ray on the head node. You don't need to change this.
|
||||
head_start_ray_commands:
|
||||
- ray stop
|
||||
- ulimit -n 65536;
|
||||
ray start --head --port=6379
|
||||
--object-manager-port=8076
|
||||
--autoscaling-config=~/ray_bootstrap_config.yaml
|
||||
|
||||
# Command to start ray on worker nodes. You don't need to change this.
|
||||
worker_start_ray_commands:
|
||||
- ray stop
|
||||
- ulimit -n 65536;
|
||||
ray start
|
||||
--address=$RAY_HEAD_IP:6379
|
||||
--object-manager-port=8076
|
||||
@@ -0,0 +1,33 @@
|
||||
#!/usr/bin/env bash
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
cd ~/efs/lm
|
||||
|
||||
# download the dataset
|
||||
wget https://s3.amazonaws.com/research.metamind.io/wikitext/wikitext-103-raw-v1.zip
|
||||
unzip wikitext-103-raw-v1.zip
|
||||
# encode it with the GPT-2 BPE
|
||||
mkdir -p gpt2_bpe
|
||||
wget -O gpt2_bpe/encoder.json https://dl.fbaipublicfiles.com/fairseq/gpt2_bpe/encoder.json
|
||||
wget -O gpt2_bpe/vocab.bpe https://dl.fbaipublicfiles.com/fairseq/gpt2_bpe/vocab.bpe
|
||||
wget https://raw.githubusercontent.com/facebookresearch/fairseq/master/examples/roberta/multiprocessing_bpe_encoder.py
|
||||
for SPLIT in train valid test; do \
|
||||
python multiprocessing_bpe_encoder.py \
|
||||
--encoder-json gpt2_bpe/encoder.json \
|
||||
--vocab-bpe gpt2_bpe/vocab.bpe \
|
||||
--inputs wikitext-103-raw/wiki.${SPLIT}.raw \
|
||||
--outputs wikitext-103-raw/wiki.${SPLIT}.bpe \
|
||||
--keep-empty \
|
||||
--workers 60; \
|
||||
done
|
||||
# preprocess/binarize the data using the GPT-2 fairseq dictionary
|
||||
wget -O gpt2_bpe/dict.txt https://dl.fbaipublicfiles.com/fairseq/gpt2_bpe/dict.txt
|
||||
fairseq-preprocess \
|
||||
--only-source \
|
||||
--srcdict gpt2_bpe/dict.txt \
|
||||
--trainpref wikitext-103-raw/wiki.train.bpe \
|
||||
--validpref wikitext-103-raw/wiki.valid.bpe \
|
||||
--testpref wikitext-103-raw/wiki.test.bpe \
|
||||
--destdir data-bin/wikitext-103 \
|
||||
--workers 60
|
||||
@@ -0,0 +1,179 @@
|
||||
#!/usr/bin/env python3 -u
|
||||
|
||||
import copy
|
||||
import math
|
||||
import socket
|
||||
import time
|
||||
from contextlib import closing
|
||||
|
||||
import fairseq
|
||||
from fairseq import options
|
||||
from fairseq_cli.train import main
|
||||
|
||||
import ray
|
||||
|
||||
_original_save_checkpoint = fairseq.checkpoint_utils.save_checkpoint
|
||||
|
||||
|
||||
class RayDistributedActor:
|
||||
"""Actor to perform distributed training."""
|
||||
|
||||
def run(self, url, world_rank, args):
|
||||
"""Runs the fairseq training.
|
||||
|
||||
We set args for different ray actors for communication,
|
||||
add a checkpoint hook, and call the main function of fairseq.
|
||||
"""
|
||||
|
||||
# Set the init_method and rank of the process for distributed training.
|
||||
print("Ray worker at {url} rank {rank}".format(url=url, rank=world_rank))
|
||||
self.url = url
|
||||
self.world_rank = world_rank
|
||||
args.distributed_rank = world_rank
|
||||
args.distributed_init_method = url
|
||||
|
||||
# Add a checkpoint hook to make use of new resources.
|
||||
self.add_checkpoint_hook(args)
|
||||
|
||||
# Call the original main function of fairseq.
|
||||
main(args, init_distributed=(args.distributed_world_size > 1))
|
||||
|
||||
def add_checkpoint_hook(self, args):
|
||||
"""Add a hook to the original save_checkpoint function.
|
||||
|
||||
This checks if there are new computational resources available.
|
||||
If so, raise exception to restart the training process and
|
||||
make use of the new resources.
|
||||
"""
|
||||
|
||||
if args.cpu:
|
||||
original_n_cpus = args.distributed_world_size
|
||||
|
||||
def _new_save_checkpoint(*args, **kwargs):
|
||||
_original_save_checkpoint(*args, **kwargs)
|
||||
n_cpus = int(ray.cluster_resources()["CPU"])
|
||||
if n_cpus > original_n_cpus:
|
||||
raise Exception(
|
||||
"New CPUs find (original %d CPUs, now %d CPUs)"
|
||||
% (original_n_cpus, n_cpus)
|
||||
)
|
||||
|
||||
else:
|
||||
original_n_gpus = args.distributed_world_size
|
||||
|
||||
def _new_save_checkpoint(*args, **kwargs):
|
||||
_original_save_checkpoint(*args, **kwargs)
|
||||
n_gpus = int(ray.cluster_resources().get("GPU", 0))
|
||||
if n_gpus > original_n_gpus:
|
||||
raise Exception(
|
||||
"New GPUs find (original %d GPUs, now %d GPUs)"
|
||||
% (original_n_gpus, n_gpus)
|
||||
)
|
||||
|
||||
fairseq.checkpoint_utils.save_checkpoint = _new_save_checkpoint
|
||||
|
||||
def get_node_ip(self):
|
||||
"""Returns the IP address of the current node."""
|
||||
return ray.util.get_node_ip_address()
|
||||
|
||||
def find_free_port(self):
|
||||
"""Finds a free port on the current node."""
|
||||
with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as s:
|
||||
s.bind(("", 0))
|
||||
return s.getsockname()[1]
|
||||
|
||||
|
||||
def run_fault_tolerant_loop():
|
||||
"""Entrance function to the fairseq library, providing fault-tolerance."""
|
||||
|
||||
# Parse the command line arguments.
|
||||
parser = options.get_training_parser()
|
||||
add_ray_args(parser)
|
||||
args = options.parse_args_and_arch(parser)
|
||||
original_args = copy.deepcopy(args)
|
||||
|
||||
# Main loop for fault-tolerant training.
|
||||
retry = True
|
||||
while retry:
|
||||
args = copy.deepcopy(original_args)
|
||||
|
||||
# Initialize Ray.
|
||||
ray.init(address=args.ray_address)
|
||||
|
||||
set_num_resources(args)
|
||||
set_batch_size(args)
|
||||
|
||||
# Set up Ray distributed actors.
|
||||
Actor = ray.remote(num_cpus=1, num_gpus=int(not args.cpu))(RayDistributedActor)
|
||||
workers = [Actor.remote() for i in range(args.distributed_world_size)]
|
||||
|
||||
# Get the IP address and a free port of actor 0, which is used for
|
||||
# fairseq distributed training.
|
||||
ip = ray.get(workers[0].get_node_ip.remote())
|
||||
port = ray.get(workers[0].find_free_port.remote())
|
||||
address = f"tcp://{ip}:{port}"
|
||||
|
||||
# Start the remote processes, and check whether their are any process
|
||||
# fails. If so, restart all the processes.
|
||||
unfinished = [
|
||||
worker.run.remote(address, i, args) for i, worker in enumerate(workers)
|
||||
]
|
||||
try:
|
||||
while len(unfinished) > 0:
|
||||
finished, unfinished = ray.wait(unfinished)
|
||||
finished = ray.get(finished)
|
||||
retry = False
|
||||
except Exception as inst:
|
||||
print("Ray restart because following error occurs:")
|
||||
print(inst)
|
||||
retry = True
|
||||
ray.shutdown()
|
||||
|
||||
|
||||
def add_ray_args(parser):
|
||||
"""Add ray and fault-tolerance related parser arguments to the parser."""
|
||||
group = parser.add_argument_group("Ray related arguments")
|
||||
group.add_argument(
|
||||
"--ray-address", default="auto", type=str, help="address for ray initialization"
|
||||
)
|
||||
group.add_argument(
|
||||
"--fix-batch-size",
|
||||
default=None,
|
||||
metavar="B1,B2,...,B_N",
|
||||
type=lambda uf: options.eval_str_list(uf, type=int),
|
||||
help="fix the actual batch size (max_sentences * update_freq "
|
||||
"* n_GPUs) to be the fixed input values by adjusting update_freq "
|
||||
"according to actual n_GPUs; the batch size is fixed to B_i for "
|
||||
"epoch i; all epochs >N are fixed to B_N",
|
||||
)
|
||||
return group
|
||||
|
||||
|
||||
def set_num_resources(args):
|
||||
"""Get the number of resources and set the corresponding fields."""
|
||||
if args.cpu:
|
||||
args.distributed_world_size = int(ray.cluster_resources()["CPU"])
|
||||
else:
|
||||
n_gpus = int(ray.cluster_resources().get("GPU", 0))
|
||||
while n_gpus == 0:
|
||||
print("No GPUs available, wait 10 seconds")
|
||||
time.sleep(10)
|
||||
n_gpus = int(ray.cluster_resources().get("GPU", 0))
|
||||
args.distributed_world_size = n_gpus
|
||||
|
||||
|
||||
def set_batch_size(args):
|
||||
"""Fixes the total batch_size to be agnostic to the GPU count."""
|
||||
if args.fix_batch_size is not None:
|
||||
args.update_freq = [
|
||||
math.ceil(batch_size / (args.max_sentences * args.distributed_world_size))
|
||||
for batch_size in args.fix_batch_size
|
||||
]
|
||||
print(
|
||||
"Training on %d GPUs, max_sentences=%d, update_freq=%s"
|
||||
% (args.distributed_world_size, args.max_sentences, repr(args.update_freq))
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
run_fault_tolerant_loop()
|
||||
@@ -0,0 +1,26 @@
|
||||
#!/bin/bash
|
||||
|
||||
TOTAL_UPDATES=125000 # Total number of training steps
|
||||
WARMUP_UPDATES=10000 # Warmup the learning rate over this many updates
|
||||
PEAK_LR=0.0005 # Peak learning rate, adjust as needed
|
||||
TOKENS_PER_SAMPLE=512 # Max sequence length
|
||||
#MAX_POSITIONS=512 # Num. positional embeddings (usually same as above)
|
||||
MAX_SENTENCES=8 # Number of sequences per batch on one GPU (batch size)
|
||||
FIX_BATCH_SIZE=2048 # Number of batch size in total (max_sentences * update_freq * n_gpus)
|
||||
SAVE_INTERVAL_UPDATES=1000 # save a checkpoint every N updates
|
||||
|
||||
LOG_DIR=$HOME/efs/lm/log/
|
||||
DATA_DIR=$HOME/efs/lm/data-bin/wikitext-103/
|
||||
mkdir -p "$LOG_DIR"
|
||||
|
||||
python "$HOME"/efs/lm/ray_train.py --fp16 "$DATA_DIR" \
|
||||
--task masked_lm --criterion masked_lm \
|
||||
--arch roberta_base --sample-break-mode complete --tokens-per-sample $TOKENS_PER_SAMPLE \
|
||||
--optimizer adam --adam-betas '(0.9, 0.98)' --adam-eps 1e-6 --clip-norm 0.0 \
|
||||
--lr-scheduler polynomial_decay --lr $PEAK_LR --warmup-updates $WARMUP_UPDATES --total-num-update $TOTAL_UPDATES \
|
||||
--dropout 0.1 --attention-dropout 0.1 --weight-decay 0.01 \
|
||||
--max-sentences $MAX_SENTENCES \
|
||||
--fix-batch-size $FIX_BATCH_SIZE \
|
||||
--max-update $TOTAL_UPDATES --log-format simple --log-interval 1 \
|
||||
--save-interval-updates $SAVE_INTERVAL_UPDATES \
|
||||
--save-dir "$LOG_DIR" --ddp-backend=no_c10d
|
||||
@@ -0,0 +1,435 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"attachments": {},
|
||||
"cell_type": "markdown",
|
||||
"id": "999dbcc2",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# A Simple MapReduce Example with Ray Core\n",
|
||||
"\n",
|
||||
"<a id=\"try-anyscale-quickstart-map_reduce\" href=\"https://console.anyscale.com/register/ha?render_flow=ray&utm_source=ray_docs&utm_medium=docs&utm_campaign=map_reduce\">\n",
|
||||
" <img src=\"../../_static/img/run-on-anyscale.svg\" alt=\"try-anyscale-quickstart\">\n",
|
||||
"</a>\n",
|
||||
"<br></br>\n",
|
||||
"\n",
|
||||
" This example demonstrates how to use Ray for a common distributed computing example––counting word occurrences across multiple documents. The complexity lies in the handling of a large corpus, requiring multiple compute nodes to process the data. \n",
|
||||
" The simplicity of implementing MapReduce with Ray is a significant milestone in distributed computing. \n",
|
||||
" Many popular big data technologies, such as Hadoop, are built on this programming model, underscoring the impact \n",
|
||||
" of using Ray Core.\n",
|
||||
"\n",
|
||||
"The MapReduce approach has three phases:"
|
||||
]
|
||||
},
|
||||
{
|
||||
"attachments": {},
|
||||
"cell_type": "markdown",
|
||||
"id": "aae37e0a",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"1. Map phase\n",
|
||||
" The map phase applies a specified function to transform or _map_ elements within a set of data. It produces key-value pairs: the key represents an element and the value is a metric calculated for that element.\n",
|
||||
" To count the number of times each word appears in a document,\n",
|
||||
" the map function outputs the pair `(word, 1)` every time a word appears, to indicate that it has been found once.\n",
|
||||
"2. Shuffle phase\n",
|
||||
" The shuffle phase collects all the outputs from the map phase and organizes them by key. When the same key is found on multiple compute nodes, this phase includes transferring or _shuffling_ data between different nodes.\n",
|
||||
" If the map phase produces four occurrences of the pair `(word, 1)`, the shuffle phase puts all occurrences of the word on the same node.\n",
|
||||
"3. Reduce phase\n",
|
||||
" The reduce phase aggregates the elements from the shuffle phase.\n",
|
||||
" The total count of each word's occurrences is the sum of occurrences on each node.\n",
|
||||
" For example, four instances of `(word, 1)` combine for a final count of `word: 4`."
|
||||
]
|
||||
},
|
||||
{
|
||||
"attachments": {},
|
||||
"cell_type": "markdown",
|
||||
"id": "87bd5e5e",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"The first and last phases are in the MapReduce name, but the middle phase is equally crucial.\n",
|
||||
"These phases appear straightforward, but their power is in running them concurrently on multiple machines.\n",
|
||||
"This figure illustrates the three MapReduce phases on a set of documents:"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "0d0b1af6",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"\n",
|
||||
""
|
||||
]
|
||||
},
|
||||
{
|
||||
"attachments": {},
|
||||
"cell_type": "markdown",
|
||||
"id": "2225ae60",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Loading Data\n",
|
||||
"\n",
|
||||
"We use Python to implement the MapReduce algorithm for the word count and Ray to parallelize the computation.\n",
|
||||
"We start by loading some sample data from the Zen of Python, a collection of coding guidelines for the Python community. Access to the Zen of Python, according to Easter egg tradition, is by typing `import this` in a Python session. \n",
|
||||
"We divide the Zen of Python into three separate \"documents\" by treating each line as a separate entity\n",
|
||||
"and then splitting the lines into three partitions."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 13,
|
||||
"id": "91c6ddc0",
|
||||
"metadata": {
|
||||
"lines_to_next_cell": 2
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"import subprocess\n",
|
||||
"zen_of_python = subprocess.check_output([\"python\", \"-c\", \"import this\"])\n",
|
||||
"corpus = zen_of_python.split()\n",
|
||||
"\n",
|
||||
"num_partitions = 3\n",
|
||||
"chunk = len(corpus) // num_partitions\n",
|
||||
"partitions = [\n",
|
||||
" corpus[i * chunk: (i + 1) * chunk] for i in range(num_partitions)\n",
|
||||
"]"
|
||||
]
|
||||
},
|
||||
{
|
||||
"attachments": {},
|
||||
"cell_type": "markdown",
|
||||
"id": "1357e924",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Mapping Data\n",
|
||||
"\n",
|
||||
"To determine the map phase, we require a map function to use on each document.\n",
|
||||
"The output is the pair `(word, 1)` for every word found in a document.\n",
|
||||
"For basic text documents we load as Python strings, the process is as follows:"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 14,
|
||||
"id": "742193e2",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"def map_function(document):\n",
|
||||
" for word in document.lower().split():\n",
|
||||
" yield word, 1"
|
||||
]
|
||||
},
|
||||
{
|
||||
"attachments": {},
|
||||
"cell_type": "markdown",
|
||||
"id": "52009879",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"We use the `apply_map` function on a large collection of documents by marking it as a task in Ray using the [`@ray.remote`](https://docs.ray.io/en/latest/ray-core/api/doc/ray.remote.html) decorator.\n",
|
||||
"When we call `apply_map`, we apply it to three sets of document data (`num_partitions=3`).\n",
|
||||
"The `apply_map` function returns three lists, one for each partition so that Ray can rearrange the results of the map phase and distribute them to the appropriate nodes."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 15,
|
||||
"id": "a2fed469",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"import ray\n",
|
||||
"\n",
|
||||
"@ray.remote\n",
|
||||
"def apply_map(corpus, num_partitions=3):\n",
|
||||
" map_results = [list() for _ in range(num_partitions)]\n",
|
||||
" for document in corpus:\n",
|
||||
" for result in map_function(document):\n",
|
||||
" first_letter = result[0].decode(\"utf-8\")[0]\n",
|
||||
" word_index = ord(first_letter) % num_partitions\n",
|
||||
" map_results[word_index].append(result)\n",
|
||||
" return map_results"
|
||||
]
|
||||
},
|
||||
{
|
||||
"attachments": {},
|
||||
"cell_type": "markdown",
|
||||
"id": "a1ba13f8",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"For text corpora that can be stored on a single machine, the map phase is not necessasry.\n",
|
||||
"However, when the data needs to be divided across multiple nodes, the map phase is useful.\n",
|
||||
"To apply the map phase to the corpus in parallel, we use a remote call on `apply_map`, similar to the previous examples.\n",
|
||||
"The main difference is that we want three results returned (one for each partition) using the `num_returns` argument."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 16,
|
||||
"id": "360b19b8",
|
||||
"metadata": {
|
||||
"lines_to_next_cell": 2
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"Mapper 0, return value 0: [(b'of', 1), (b'is', 1)]\n",
|
||||
"Mapper 0, return value 1: [(b'python,', 1), (b'peters', 1)]\n",
|
||||
"Mapper 0, return value 2: [(b'the', 1), (b'zen', 1)]\n",
|
||||
"Mapper 1, return value 0: [(b'unless', 1), (b'in', 1)]\n",
|
||||
"Mapper 1, return value 1: [(b'although', 1), (b'practicality', 1)]\n",
|
||||
"Mapper 1, return value 2: [(b'beats', 1), (b'errors', 1)]\n",
|
||||
"Mapper 2, return value 0: [(b'is', 1), (b'is', 1)]\n",
|
||||
"Mapper 2, return value 1: [(b'although', 1), (b'a', 1)]\n",
|
||||
"Mapper 2, return value 2: [(b'better', 1), (b'than', 1)]\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"map_results = [\n",
|
||||
" apply_map.options(num_returns=num_partitions)\n",
|
||||
" .remote(data, num_partitions)\n",
|
||||
" for data in partitions\n",
|
||||
"]\n",
|
||||
"\n",
|
||||
"for i in range(num_partitions):\n",
|
||||
" mapper_results = ray.get(map_results[i])\n",
|
||||
" for j, result in enumerate(mapper_results):\n",
|
||||
" print(f\"Mapper {i}, return value {j}: {result[:2]}\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"attachments": {},
|
||||
"cell_type": "markdown",
|
||||
"id": "ce222cd3",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"This example demonstrates how to collect data on the driver with `ray.get`. To continue with another task after the mapping phase, you wouldn't do this. The following section shows how to run all phases together efficiently."
|
||||
]
|
||||
},
|
||||
{
|
||||
"attachments": {},
|
||||
"cell_type": "markdown",
|
||||
"id": "171744b1",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Shuffling and Reducing Data\n",
|
||||
"\n",
|
||||
"The objective for the reduce phase is to transfer all pairs from the `j`-th return value to the same node.\n",
|
||||
"In the reduce phase we create a dictionary that adds up all word occurrences on each partition:"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 17,
|
||||
"id": "5891b2c3",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"@ray.remote\n",
|
||||
"def apply_reduce(*results):\n",
|
||||
" reduce_results = dict()\n",
|
||||
" for res in results:\n",
|
||||
" for key, value in res:\n",
|
||||
" if key not in reduce_results:\n",
|
||||
" reduce_results[key] = 0\n",
|
||||
" reduce_results[key] += value\n",
|
||||
"\n",
|
||||
" return reduce_results"
|
||||
]
|
||||
},
|
||||
{
|
||||
"attachments": {},
|
||||
"cell_type": "markdown",
|
||||
"id": "21ee3c55",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"We can take the j-th return value from each mapper and send it to the j-th reducer using the following method.\n",
|
||||
"Note that this code works for large datasets that don't fit on one machine because we are passing references\n",
|
||||
"to the data using Ray objects rather than the actual data itself.\n",
|
||||
"Both the map and reduce phases can run on any Ray cluster and Ray handles the data shuffling."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 18,
|
||||
"id": "a395a7f9",
|
||||
"metadata": {
|
||||
"collapsed": true
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"is: 10\n",
|
||||
"better: 8\n",
|
||||
"than: 8\n",
|
||||
"the: 6\n",
|
||||
"to: 5\n",
|
||||
"of: 3\n",
|
||||
"although: 3\n",
|
||||
"be: 3\n",
|
||||
"unless: 2\n",
|
||||
"one: 2\n",
|
||||
"if: 2\n",
|
||||
"implementation: 2\n",
|
||||
"idea.: 2\n",
|
||||
"special: 2\n",
|
||||
"should: 2\n",
|
||||
"do: 2\n",
|
||||
"may: 2\n",
|
||||
"a: 2\n",
|
||||
"never: 2\n",
|
||||
"way: 2\n",
|
||||
"explain,: 2\n",
|
||||
"ugly.: 1\n",
|
||||
"implicit.: 1\n",
|
||||
"complex.: 1\n",
|
||||
"complex: 1\n",
|
||||
"complicated.: 1\n",
|
||||
"flat: 1\n",
|
||||
"readability: 1\n",
|
||||
"counts.: 1\n",
|
||||
"cases: 1\n",
|
||||
"rules.: 1\n",
|
||||
"in: 1\n",
|
||||
"face: 1\n",
|
||||
"refuse: 1\n",
|
||||
"one--: 1\n",
|
||||
"only: 1\n",
|
||||
"--obvious: 1\n",
|
||||
"it.: 1\n",
|
||||
"obvious: 1\n",
|
||||
"first: 1\n",
|
||||
"often: 1\n",
|
||||
"*right*: 1\n",
|
||||
"it's: 1\n",
|
||||
"it: 1\n",
|
||||
"idea: 1\n",
|
||||
"--: 1\n",
|
||||
"let's: 1\n",
|
||||
"python,: 1\n",
|
||||
"peters: 1\n",
|
||||
"simple: 1\n",
|
||||
"sparse: 1\n",
|
||||
"dense.: 1\n",
|
||||
"aren't: 1\n",
|
||||
"practicality: 1\n",
|
||||
"purity.: 1\n",
|
||||
"pass: 1\n",
|
||||
"silently.: 1\n",
|
||||
"silenced.: 1\n",
|
||||
"ambiguity,: 1\n",
|
||||
"guess.: 1\n",
|
||||
"and: 1\n",
|
||||
"preferably: 1\n",
|
||||
"at: 1\n",
|
||||
"you're: 1\n",
|
||||
"dutch.: 1\n",
|
||||
"good: 1\n",
|
||||
"are: 1\n",
|
||||
"great: 1\n",
|
||||
"more: 1\n",
|
||||
"zen: 1\n",
|
||||
"by: 1\n",
|
||||
"tim: 1\n",
|
||||
"beautiful: 1\n",
|
||||
"explicit: 1\n",
|
||||
"nested.: 1\n",
|
||||
"enough: 1\n",
|
||||
"break: 1\n",
|
||||
"beats: 1\n",
|
||||
"errors: 1\n",
|
||||
"explicitly: 1\n",
|
||||
"temptation: 1\n",
|
||||
"there: 1\n",
|
||||
"that: 1\n",
|
||||
"not: 1\n",
|
||||
"now: 1\n",
|
||||
"never.: 1\n",
|
||||
"now.: 1\n",
|
||||
"hard: 1\n",
|
||||
"bad: 1\n",
|
||||
"easy: 1\n",
|
||||
"namespaces: 1\n",
|
||||
"honking: 1\n",
|
||||
"those!: 1\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"outputs = []\n",
|
||||
"for i in range(num_partitions):\n",
|
||||
" outputs.append(\n",
|
||||
" apply_reduce.remote(*[partition[i] for partition in map_results])\n",
|
||||
" )\n",
|
||||
"\n",
|
||||
"counts = {k: v for output in ray.get(outputs) for k, v in output.items()}\n",
|
||||
"\n",
|
||||
"sorted_counts = sorted(counts.items(), key=lambda item: item[1], reverse=True)\n",
|
||||
"for count in sorted_counts:\n",
|
||||
" print(f\"{count[0].decode('utf-8')}: {count[1]}\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"attachments": {},
|
||||
"cell_type": "markdown",
|
||||
"id": "c57bd93b",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"For a thorough understanding of scaling MapReduce tasks across multiple nodes using Ray,\n",
|
||||
"including memory management, read the [blog post on the topic](https://medium.com/distributed-computing-with-ray/executing-adistributed-shuffle-without-a-mapreduce-system-d5856379426c).\n",
|
||||
"\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"attachments": {},
|
||||
"cell_type": "markdown",
|
||||
"id": "20346f17",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Wrapping up\n",
|
||||
"\n",
|
||||
"This MapReduce example demonstrates how flexible Ray’s programming model is.\n",
|
||||
"A production-grade MapReduce implementation requires more effort but being able to reproduce common algorithms like this one _quickly_ goes a long way.\n",
|
||||
"In the earlier years of MapReduce, around 2010, this paradigm was often the only model available for\n",
|
||||
"expressing workloads.\n",
|
||||
"With Ray, an entire range of interesting distributed computing patterns\n",
|
||||
"are accessible to any intermediate Python programmer.\n",
|
||||
"\n",
|
||||
"To learn more about Ray, and Ray Core and particular, see the [Ray Core Examples Gallery](./overview.rst),\n",
|
||||
"or the ML workloads in our [Use Case Gallery](../../ray-overview/use-cases.rst).\n",
|
||||
"This MapReduce example can be found in [\"Learning Ray\"](https://maxpumperla.com/learning_ray/),\n",
|
||||
"which contains more examples similar to this one."
|
||||
]
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"jupytext": {
|
||||
"cell_metadata_filter": "-all",
|
||||
"main_language": "python",
|
||||
"notebook_metadata_filter": "-all"
|
||||
},
|
||||
"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"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 5
|
||||
}
|
||||
@@ -0,0 +1,131 @@
|
||||
.. _monte-carlo-pi:
|
||||
|
||||
Monte Carlo Estimation of π
|
||||
===========================
|
||||
|
||||
.. raw:: html
|
||||
|
||||
<a id="try-anyscale-quickstart-monte_carlo_pi" target="_blank" href="https://console.anyscale.com/register/ha?render_flow=ray&utm_source=ray_docs&utm_medium=docs&utm_campaign=monte_carlo_pi">
|
||||
<img src="../../_static/img/run-on-anyscale.svg" alt="Run on Anyscale" />
|
||||
<br/><br/>
|
||||
</a>
|
||||
|
||||
This tutorial shows you how to estimate the value of π using a `Monte Carlo method <https://en.wikipedia.org/wiki/Monte_Carlo_method>`_
|
||||
that works by randomly sampling points within a 2x2 square.
|
||||
We can use the proportion of the points that are contained within the unit circle centered at the origin
|
||||
to estimate the ratio of the area of the circle to the area of the square.
|
||||
Given that we know the true ratio to be π/4, we can multiply our estimated ratio by 4 to approximate the value of π.
|
||||
The more points that we sample to calculate this approximation, the closer the value should be to the true value of π.
|
||||
|
||||
.. image:: ../images/monte_carlo_pi.png
|
||||
|
||||
We use Ray :ref:`tasks <ray-remote-functions>` to distribute the work of sampling and Ray :ref:`actors <ray-remote-classes>` to track the progress of these distributed sampling tasks.
|
||||
The code can run on your laptop and can be easily scaled to large :ref:`clusters <cluster-index>` to increase the accuracy of the estimate.
|
||||
|
||||
To get started, install Ray via ``pip install -U ray``. See :ref:`Installing Ray <installation>` for more installation options.
|
||||
|
||||
Starting Ray
|
||||
------------
|
||||
First, let's include all modules needed for this tutorial and start a local Ray cluster with :func:`ray.init() <ray.init>`:
|
||||
|
||||
.. literalinclude:: ../doc_code/monte_carlo_pi.py
|
||||
:language: python
|
||||
:start-after: __starting_ray_start__
|
||||
:end-before: __starting_ray_end__
|
||||
|
||||
|
||||
|
||||
Defining the Progress Actor
|
||||
---------------------------
|
||||
Next, we define a Ray actor that can be called by sampling tasks to update progress.
|
||||
Ray actors are essentially stateful services that anyone with an instance (a handle) of the actor can call its methods.
|
||||
|
||||
.. literalinclude:: ../doc_code/monte_carlo_pi.py
|
||||
:language: python
|
||||
:start-after: __defining_actor_start__
|
||||
:end-before: __defining_actor_end__
|
||||
|
||||
We define a Ray actor by decorating a normal Python class with :func:`ray.remote <ray.remote>`.
|
||||
The progress actor has ``report_progress()`` method that will be called by sampling tasks to update their progress individually
|
||||
and ``get_progress()`` method to get the overall progress.
|
||||
|
||||
Defining the Sampling Task
|
||||
--------------------------
|
||||
After our actor is defined, we now define a Ray task that does the sampling up to ``num_samples`` and returns the number of samples that are inside the circle.
|
||||
Ray tasks are stateless functions. They execute asynchronously, and run in parallel.
|
||||
|
||||
.. literalinclude:: ../doc_code/monte_carlo_pi.py
|
||||
:language: python
|
||||
:start-after: __defining_task_start__
|
||||
:end-before: __defining_task_end__
|
||||
|
||||
To convert a normal Python function as a Ray task, we decorate the function with :func:`ray.remote <ray.remote>`.
|
||||
The sampling task takes a progress actor handle as an input and reports progress to it.
|
||||
The above code shows an example of calling actor methods from tasks.
|
||||
|
||||
Creating a Progress Actor
|
||||
-------------------------
|
||||
Once the actor is defined, we can create an instance of it.
|
||||
|
||||
.. literalinclude:: ../doc_code/monte_carlo_pi.py
|
||||
:language: python
|
||||
:start-after: __creating_actor_start__
|
||||
:end-before: __creating_actor_end__
|
||||
|
||||
To create an instance of the progress actor, simply call ``ActorClass.remote()`` method with arguments to the constructor.
|
||||
This creates and runs the actor on a remote worker process.
|
||||
The return value of ``ActorClass.remote(...)`` is an actor handle that can be used to call its methods.
|
||||
|
||||
Executing Sampling Tasks
|
||||
------------------------
|
||||
Now the task is defined, we can execute it asynchronously.
|
||||
|
||||
.. literalinclude:: ../doc_code/monte_carlo_pi.py
|
||||
:language: python
|
||||
:start-after: __executing_task_start__
|
||||
:end-before: __executing_task_end__
|
||||
|
||||
We execute the sampling task by calling ``remote()`` method with arguments to the function.
|
||||
This immediately returns an ``ObjectRef`` as a future
|
||||
and then executes the function asynchronously on a remote worker process.
|
||||
|
||||
Calling the Progress Actor
|
||||
--------------------------
|
||||
While sampling tasks are running, we can periodically query the progress by calling the actor ``get_progress()`` method.
|
||||
|
||||
.. literalinclude:: ../doc_code/monte_carlo_pi.py
|
||||
:language: python
|
||||
:start-after: __calling_actor_start__
|
||||
:end-before: __calling_actor_end__
|
||||
|
||||
To call an actor method, use ``actor_handle.method.remote()``.
|
||||
This invocation immediately returns an ``ObjectRef`` as a future
|
||||
and then executes the method asynchronously on the remote actor process.
|
||||
To fetch the actual returned value of ``ObjectRef``, we use the blocking :func:`ray.get() <ray.get>`.
|
||||
|
||||
Calculating π
|
||||
-------------
|
||||
Finally, we get number of samples inside the circle from the remote sampling tasks and calculate π.
|
||||
|
||||
.. literalinclude:: ../doc_code/monte_carlo_pi.py
|
||||
:language: python
|
||||
:start-after: __calculating_pi_start__
|
||||
:end-before: __calculating_pi_end__
|
||||
|
||||
As we can see from the above code, besides a single ``ObjectRef``, :func:`ray.get() <ray.get>` can also take a list of ``ObjectRef`` and return a list of results.
|
||||
|
||||
If you run this tutorial, you will see output like:
|
||||
|
||||
.. code-block:: text
|
||||
|
||||
Progress: 0%
|
||||
Progress: 15%
|
||||
Progress: 28%
|
||||
Progress: 40%
|
||||
Progress: 50%
|
||||
Progress: 60%
|
||||
Progress: 70%
|
||||
Progress: 80%
|
||||
Progress: 90%
|
||||
Progress: 100%
|
||||
Estimated value of π is: 3.1412202
|
||||
@@ -0,0 +1,44 @@
|
||||
.. _ray-core-examples-tutorial:
|
||||
|
||||
Ray Core Examples
|
||||
=================
|
||||
|
||||
.. toctree::
|
||||
:hidden:
|
||||
:glob:
|
||||
|
||||
*
|
||||
|
||||
.. Organize example .rst files in the same manner as the
|
||||
.py files in ray/python/ray/train/examples.
|
||||
|
||||
Below are examples for using Ray Core for a variety use cases.
|
||||
|
||||
Beginner
|
||||
--------
|
||||
|
||||
.. list-table::
|
||||
|
||||
* - :doc:`A Gentle Introduction to Ray Core by Example <gentle_walkthrough>`
|
||||
* - :doc:`Using Ray for Highly Parallelizable Tasks <highly_parallel>`
|
||||
* - :doc:`Monte Carlo Estimation of π <monte_carlo_pi>`
|
||||
|
||||
|
||||
Intermediate
|
||||
------------
|
||||
|
||||
.. list-table::
|
||||
|
||||
* - :doc:`Running a Simple MapReduce Example with Ray Core <map_reduce>`
|
||||
* - :doc:`Speed Up Your Web Crawler by Parallelizing it with Ray <web_crawler>`
|
||||
|
||||
|
||||
Advanced
|
||||
--------
|
||||
|
||||
.. list-table::
|
||||
|
||||
* - :doc:`Build Batch Prediction Using Ray <batch_prediction>`
|
||||
* - :doc:`Build a Simple Parameter Server Using Ray <plot_parameter_server>`
|
||||
* - :doc:`Simple Parallel Model Selection <plot_hyperparameter>`
|
||||
* - :doc:`Learning to Play Pong <plot_pong_example>`
|
||||
@@ -0,0 +1,322 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "436ead19",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# Simple Parallel Model Selection\n",
|
||||
"\n",
|
||||
"<a id=\"try-anyscale-quickstart-plot_hyperparameter\" href=\"https://console.anyscale.com/register/ha?render_flow=ray&utm_source=ray_docs&utm_medium=docs&utm_campaign=plot_hyperparameter\">\n",
|
||||
" <img src=\"../../_static/img/run-on-anyscale.svg\" alt=\"try-anyscale-quickstart\">\n",
|
||||
"</a>\n",
|
||||
"<br></br>\n",
|
||||
"\n",
|
||||
"```{tip}\n",
|
||||
"For a production-grade implementation of distributed\n",
|
||||
"hyperparameter tuning, use [Ray Tune](https://docs.ray.io/en/master/tune.html), a scalable hyperparameter\n",
|
||||
"tuning library built using Ray's Actor API.\n",
|
||||
"```\n",
|
||||
"\n",
|
||||
"In this example, we'll demonstrate how to quickly write a hyperparameter\n",
|
||||
"tuning script that evaluates a set of hyperparameters in parallel.\n",
|
||||
"\n",
|
||||
"This script will demonstrate how to use two important parts of the Ray API:\n",
|
||||
"using ``ray.remote`` to define remote functions and ``ray.wait`` to wait for\n",
|
||||
"their results to be ready.\n",
|
||||
"\n",
|
||||
"```{image} /ray-core/images/hyperparameter.png\n",
|
||||
":align: center\n",
|
||||
"```\n",
|
||||
"\n",
|
||||
"## Setup: Dependencies\n",
|
||||
"\n",
|
||||
"First, import some dependencies and define functions to generate\n",
|
||||
"random hyperparameters and retrieve data."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "8e992dc3",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"import os\n",
|
||||
"import numpy as np\n",
|
||||
"from filelock import FileLock\n",
|
||||
"\n",
|
||||
"import torch\n",
|
||||
"import torch.nn as nn\n",
|
||||
"import torch.nn.functional as F\n",
|
||||
"import torch.optim as optim\n",
|
||||
"from torchvision import datasets, transforms\n",
|
||||
"\n",
|
||||
"import ray\n",
|
||||
"\n",
|
||||
"ray.init()\n",
|
||||
"\n",
|
||||
"# The number of sets of random hyperparameters to try.\n",
|
||||
"num_evaluations = 10\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"# A function for generating random hyperparameters.\n",
|
||||
"def generate_hyperparameters():\n",
|
||||
" return {\n",
|
||||
" \"learning_rate\": 10 ** np.random.uniform(-5, 1),\n",
|
||||
" \"batch_size\": np.random.randint(1, 100),\n",
|
||||
" \"momentum\": np.random.uniform(0, 1),\n",
|
||||
" }\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"def get_data_loaders(batch_size):\n",
|
||||
" mnist_transforms = transforms.Compose(\n",
|
||||
" [transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))]\n",
|
||||
" )\n",
|
||||
"\n",
|
||||
" # We add FileLock here because multiple workers will want to\n",
|
||||
" # download data, and this may cause overwrites since\n",
|
||||
" # DataLoader is not threadsafe.\n",
|
||||
" with FileLock(os.path.expanduser(\"~/data.lock\")):\n",
|
||||
" train_loader = torch.utils.data.DataLoader(\n",
|
||||
" datasets.MNIST(\n",
|
||||
" \"~/data\", train=True, download=True, transform=mnist_transforms\n",
|
||||
" ),\n",
|
||||
" batch_size=batch_size,\n",
|
||||
" shuffle=True,\n",
|
||||
" )\n",
|
||||
" test_loader = torch.utils.data.DataLoader(\n",
|
||||
" datasets.MNIST(\"~/data\", train=False, transform=mnist_transforms),\n",
|
||||
" batch_size=batch_size,\n",
|
||||
" shuffle=True,\n",
|
||||
" )\n",
|
||||
" return train_loader, test_loader"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "a3f0d421",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Setup: Defining the Neural Network\n",
|
||||
"\n",
|
||||
"We define a small neural network to use in training. In addition,\n",
|
||||
"we created methods to train and test this neural network."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "c02ed1db",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"class ConvNet(nn.Module):\n",
|
||||
" \"\"\"Simple two layer Convolutional Neural Network.\"\"\"\n",
|
||||
"\n",
|
||||
" def __init__(self):\n",
|
||||
" super(ConvNet, self).__init__()\n",
|
||||
" self.conv1 = nn.Conv2d(1, 3, kernel_size=3)\n",
|
||||
" self.fc = nn.Linear(192, 10)\n",
|
||||
"\n",
|
||||
" def forward(self, x):\n",
|
||||
" x = F.relu(F.max_pool2d(self.conv1(x), 3))\n",
|
||||
" x = x.view(-1, 192)\n",
|
||||
" x = self.fc(x)\n",
|
||||
" return F.log_softmax(x, dim=1)\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"def train(model, optimizer, train_loader, device=torch.device(\"cpu\")):\n",
|
||||
" \"\"\"Optimize the model with one pass over the data.\n",
|
||||
"\n",
|
||||
" Cuts off at 1024 samples to simplify training.\n",
|
||||
" \"\"\"\n",
|
||||
" model.train()\n",
|
||||
" for batch_idx, (data, target) in enumerate(train_loader):\n",
|
||||
" if batch_idx * len(data) > 1024:\n",
|
||||
" return\n",
|
||||
" data, target = data.to(device), target.to(device)\n",
|
||||
" optimizer.zero_grad()\n",
|
||||
" output = model(data)\n",
|
||||
" loss = F.nll_loss(output, target)\n",
|
||||
" loss.backward()\n",
|
||||
" optimizer.step()\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"def test(model, test_loader, device=torch.device(\"cpu\")):\n",
|
||||
" \"\"\"Checks the validation accuracy of the model.\n",
|
||||
"\n",
|
||||
" Cuts off at 512 samples for simplicity.\n",
|
||||
" \"\"\"\n",
|
||||
" model.eval()\n",
|
||||
" correct = 0\n",
|
||||
" total = 0\n",
|
||||
" with torch.no_grad():\n",
|
||||
" for batch_idx, (data, target) in enumerate(test_loader):\n",
|
||||
" if batch_idx * len(data) > 512:\n",
|
||||
" break\n",
|
||||
" data, target = data.to(device), target.to(device)\n",
|
||||
" outputs = model(data)\n",
|
||||
" _, predicted = torch.max(outputs.data, 1)\n",
|
||||
" total += target.size(0)\n",
|
||||
" correct += (predicted == target).sum().item()\n",
|
||||
"\n",
|
||||
" return correct / total"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "f3a9ed6f",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Evaluating the Hyperparameters\n",
|
||||
"\n",
|
||||
"For a given configuration, the neural network created previously\n",
|
||||
"will be trained and return the accuracy of the model. These trained\n",
|
||||
"networks will then be tested for accuracy to find the best set of\n",
|
||||
"hyperparameters.\n",
|
||||
"\n",
|
||||
"The ``@ray.remote`` decorator defines a remote process."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "2471f2db",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"@ray.remote\n",
|
||||
"def evaluate_hyperparameters(config):\n",
|
||||
" model = ConvNet()\n",
|
||||
" train_loader, test_loader = get_data_loaders(config[\"batch_size\"])\n",
|
||||
" optimizer = optim.SGD(\n",
|
||||
" model.parameters(), lr=config[\"learning_rate\"], momentum=config[\"momentum\"]\n",
|
||||
" )\n",
|
||||
" train(model, optimizer, train_loader)\n",
|
||||
" return test(model, test_loader)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "62f180e6",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Synchronous Evaluation of Randomly Generated Hyperparameters\n",
|
||||
"\n",
|
||||
"We will create multiple sets of random hyperparameters for our neural\n",
|
||||
"network that will be evaluated in parallel."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "c4075165",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"# Keep track of the best hyperparameters and the best accuracy.\n",
|
||||
"best_hyperparameters = None\n",
|
||||
"best_accuracy = 0\n",
|
||||
"# A list holding the object refs for all of the experiments that we have\n",
|
||||
"# launched but have not yet been processed.\n",
|
||||
"remaining_ids = []\n",
|
||||
"# A dictionary mapping an experiment's object ref to its hyperparameters.\n",
|
||||
"# hyerparameters used for that experiment.\n",
|
||||
"hyperparameters_mapping = {}"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "8dd73456",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"Launch asynchronous parallel tasks for evaluating different\n",
|
||||
"hyperparameters. ``accuracy_id`` is an ObjectRef that acts as a handle to\n",
|
||||
"the remote task. It is used later to fetch the result of the task\n",
|
||||
"when the task finishes."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "eda2226a",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"# Randomly generate sets of hyperparameters and launch a task to evaluate it.\n",
|
||||
"for i in range(num_evaluations):\n",
|
||||
" hyperparameters = generate_hyperparameters()\n",
|
||||
" accuracy_id = evaluate_hyperparameters.remote(hyperparameters)\n",
|
||||
" remaining_ids.append(accuracy_id)\n",
|
||||
" hyperparameters_mapping[accuracy_id] = hyperparameters"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "bd0e53ec",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"Process each hyperparameter and corresponding accuracy in the order that\n",
|
||||
"they finish to store the hyperparameters with the best accuracy."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "d95ca22b",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"# Fetch and print the results of the tasks in the order that they complete.\n",
|
||||
"while remaining_ids:\n",
|
||||
" # Use ray.wait to get the object ref of the first task that completes.\n",
|
||||
" done_ids, remaining_ids = ray.wait(remaining_ids)\n",
|
||||
" # There is only one return result by default.\n",
|
||||
" result_id = done_ids[0]\n",
|
||||
"\n",
|
||||
" hyperparameters = hyperparameters_mapping[result_id]\n",
|
||||
" accuracy = ray.get(result_id)\n",
|
||||
" print(\n",
|
||||
" \"\"\"We achieve accuracy {:.3}% with\n",
|
||||
" learning_rate: {:.2}\n",
|
||||
" batch_size: {}\n",
|
||||
" momentum: {:.2}\n",
|
||||
" \"\"\".format(\n",
|
||||
" 100 * accuracy,\n",
|
||||
" hyperparameters[\"learning_rate\"],\n",
|
||||
" hyperparameters[\"batch_size\"],\n",
|
||||
" hyperparameters[\"momentum\"],\n",
|
||||
" )\n",
|
||||
" )\n",
|
||||
" if accuracy > best_accuracy:\n",
|
||||
" best_hyperparameters = hyperparameters\n",
|
||||
" best_accuracy = accuracy\n",
|
||||
"\n",
|
||||
"# Record the best performing set of hyperparameters.\n",
|
||||
"print(\n",
|
||||
" \"\"\"Best accuracy over {} trials was {:.3} with\n",
|
||||
" learning_rate: {:.2}\n",
|
||||
" batch_size: {}\n",
|
||||
" momentum: {:.2}\n",
|
||||
" \"\"\".format(\n",
|
||||
" num_evaluations,\n",
|
||||
" 100 * best_accuracy,\n",
|
||||
" best_hyperparameters[\"learning_rate\"],\n",
|
||||
" best_hyperparameters[\"batch_size\"],\n",
|
||||
" best_hyperparameters[\"momentum\"],\n",
|
||||
" )\n",
|
||||
")"
|
||||
]
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"kernelspec": {
|
||||
"display_name": "Python 3",
|
||||
"language": "python",
|
||||
"name": "python3"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 5
|
||||
}
|
||||
@@ -0,0 +1,422 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "ac6a47d6",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# Parameter Server\n",
|
||||
"\n",
|
||||
"<a id=\"try-anyscale-quickstart-plot_parameter_server\" href=\"https://console.anyscale.com/register/ha?render_flow=ray&utm_source=ray_docs&utm_medium=docs&utm_campaign=plot_parameter_server\">\n",
|
||||
" <img src=\"../../_static/img/run-on-anyscale.svg\" alt=\"try-anyscale-quickstart\">\n",
|
||||
"</a>\n",
|
||||
"<br></br>\n",
|
||||
"\n",
|
||||
"```{tip}\n",
|
||||
"For a production-grade implementation of distributed\n",
|
||||
"training, use [Ray Train](https://docs.ray.io/en/master/train/train.html).\n",
|
||||
"```\n",
|
||||
"\n",
|
||||
"The parameter server is a framework for distributed machine learning training.\n",
|
||||
"\n",
|
||||
"In the parameter server framework, a centralized server (or group of server\n",
|
||||
"nodes) maintains global shared parameters of a machine-learning model\n",
|
||||
"(e.g., a neural network) while the data and computation of calculating\n",
|
||||
"updates (i.e., gradient descent updates) are distributed over worker nodes.\n",
|
||||
"\n",
|
||||
"```{image} /ray-core/images/param_actor.png\n",
|
||||
":align: center\n",
|
||||
"```\n",
|
||||
"\n",
|
||||
"Parameter servers are a core part of many machine learning applications. This\n",
|
||||
"document walks through how to implement simple synchronous and asynchronous\n",
|
||||
"parameter servers using Ray actors.\n",
|
||||
"\n",
|
||||
"To run the application, first install some dependencies.\n",
|
||||
"\n",
|
||||
"```bash\n",
|
||||
"pip install torch torchvision filelock\n",
|
||||
"```\n",
|
||||
"\n",
|
||||
"Let's first define some helper functions and import some dependencies."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "84677721",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"import os\n",
|
||||
"import torch\n",
|
||||
"import torch.nn as nn\n",
|
||||
"import torch.nn.functional as F\n",
|
||||
"from torchvision import datasets, transforms\n",
|
||||
"from filelock import FileLock\n",
|
||||
"import numpy as np\n",
|
||||
"\n",
|
||||
"import ray\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"def get_data_loader():\n",
|
||||
" \"\"\"Safely downloads data. Returns training/validation set dataloader.\"\"\"\n",
|
||||
" mnist_transforms = transforms.Compose(\n",
|
||||
" [transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))]\n",
|
||||
" )\n",
|
||||
"\n",
|
||||
" # We add FileLock here because multiple workers will want to\n",
|
||||
" # download data, and this may cause overwrites since\n",
|
||||
" # DataLoader is not threadsafe.\n",
|
||||
" with FileLock(os.path.expanduser(\"~/data.lock\")):\n",
|
||||
" train_loader = torch.utils.data.DataLoader(\n",
|
||||
" datasets.MNIST(\n",
|
||||
" \"~/data\", train=True, download=True, transform=mnist_transforms\n",
|
||||
" ),\n",
|
||||
" batch_size=128,\n",
|
||||
" shuffle=True,\n",
|
||||
" )\n",
|
||||
" test_loader = torch.utils.data.DataLoader(\n",
|
||||
" datasets.MNIST(\"~/data\", train=False, transform=mnist_transforms),\n",
|
||||
" batch_size=128,\n",
|
||||
" shuffle=True,\n",
|
||||
" )\n",
|
||||
" return train_loader, test_loader\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"def evaluate(model, test_loader):\n",
|
||||
" \"\"\"Evaluates the accuracy of the model on a validation dataset.\"\"\"\n",
|
||||
" model.eval()\n",
|
||||
" correct = 0\n",
|
||||
" total = 0\n",
|
||||
" with torch.no_grad():\n",
|
||||
" for batch_idx, (data, target) in enumerate(test_loader):\n",
|
||||
" # This is only set to finish evaluation faster.\n",
|
||||
" if batch_idx * len(data) > 1024:\n",
|
||||
" break\n",
|
||||
" outputs = model(data)\n",
|
||||
" _, predicted = torch.max(outputs.data, 1)\n",
|
||||
" total += target.size(0)\n",
|
||||
" correct += (predicted == target).sum().item()\n",
|
||||
" return 100.0 * correct / total"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "fa5c9aea",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Setup: Defining the Neural Network\n",
|
||||
"\n",
|
||||
"We define a small neural network to use in training. We provide\n",
|
||||
"some helper functions for obtaining data, including getter/setter\n",
|
||||
"methods for gradients and weights."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "9b064f87",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"class ConvNet(nn.Module):\n",
|
||||
" \"\"\"Small ConvNet for MNIST.\"\"\"\n",
|
||||
"\n",
|
||||
" def __init__(self):\n",
|
||||
" super(ConvNet, self).__init__()\n",
|
||||
" self.conv1 = nn.Conv2d(1, 3, kernel_size=3)\n",
|
||||
" self.fc = nn.Linear(192, 10)\n",
|
||||
"\n",
|
||||
" def forward(self, x):\n",
|
||||
" x = F.relu(F.max_pool2d(self.conv1(x), 3))\n",
|
||||
" x = x.view(-1, 192)\n",
|
||||
" x = self.fc(x)\n",
|
||||
" return F.log_softmax(x, dim=1)\n",
|
||||
"\n",
|
||||
" def get_weights(self):\n",
|
||||
" return {k: v.cpu() for k, v in self.state_dict().items()}\n",
|
||||
"\n",
|
||||
" def set_weights(self, weights):\n",
|
||||
" self.load_state_dict(weights)\n",
|
||||
"\n",
|
||||
" def get_gradients(self):\n",
|
||||
" grads = []\n",
|
||||
" for p in self.parameters():\n",
|
||||
" grad = None if p.grad is None else p.grad.data.cpu().numpy()\n",
|
||||
" grads.append(grad)\n",
|
||||
" return grads\n",
|
||||
"\n",
|
||||
" def set_gradients(self, gradients):\n",
|
||||
" for g, p in zip(gradients, self.parameters()):\n",
|
||||
" if g is not None:\n",
|
||||
" p.grad = torch.from_numpy(g)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "81e2bf73",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Defining the Parameter Server\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"The parameter server will hold a copy of the model.\n",
|
||||
"During training, it will:\n",
|
||||
"\n",
|
||||
"1. Receive gradients and apply them to its model.\n",
|
||||
"\n",
|
||||
"2. Send the updated model back to the workers.\n",
|
||||
"\n",
|
||||
"The ``@ray.remote`` decorator defines a remote process. It wraps the\n",
|
||||
"ParameterServer class and allows users to instantiate it as a\n",
|
||||
"remote actor."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "77b756a4",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"@ray.remote\n",
|
||||
"class ParameterServer(object):\n",
|
||||
" def __init__(self, lr):\n",
|
||||
" self.model = ConvNet()\n",
|
||||
" self.optimizer = torch.optim.SGD(self.model.parameters(), lr=lr)\n",
|
||||
"\n",
|
||||
" def apply_gradients(self, *gradients):\n",
|
||||
" summed_gradients = [\n",
|
||||
" np.stack(gradient_zip).sum(axis=0) for gradient_zip in zip(*gradients)\n",
|
||||
" ]\n",
|
||||
" self.optimizer.zero_grad()\n",
|
||||
" self.model.set_gradients(summed_gradients)\n",
|
||||
" self.optimizer.step()\n",
|
||||
" return self.model.get_weights()\n",
|
||||
"\n",
|
||||
" def get_weights(self):\n",
|
||||
" return self.model.get_weights()"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "ccf86e35",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Defining the Worker\n",
|
||||
"\n",
|
||||
"The worker will also hold a copy of the model. During training. it will\n",
|
||||
"continuously evaluate data and send gradients\n",
|
||||
"to the parameter server. The worker will synchronize its model with the\n",
|
||||
"Parameter Server model weights."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "18737d3a",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"@ray.remote\n",
|
||||
"class DataWorker(object):\n",
|
||||
" def __init__(self):\n",
|
||||
" self.model = ConvNet()\n",
|
||||
" self.data_iterator = iter(get_data_loader()[0])\n",
|
||||
"\n",
|
||||
" def compute_gradients(self, weights):\n",
|
||||
" self.model.set_weights(weights)\n",
|
||||
" try:\n",
|
||||
" data, target = next(self.data_iterator)\n",
|
||||
" except StopIteration: # When the epoch ends, start a new epoch.\n",
|
||||
" self.data_iterator = iter(get_data_loader()[0])\n",
|
||||
" data, target = next(self.data_iterator)\n",
|
||||
" self.model.zero_grad()\n",
|
||||
" output = self.model(data)\n",
|
||||
" loss = F.nll_loss(output, target)\n",
|
||||
" loss.backward()\n",
|
||||
" return self.model.get_gradients()"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "763fc48e",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Synchronous Parameter Server Training\n",
|
||||
"\n",
|
||||
"We'll now create a synchronous parameter server training scheme. We'll first\n",
|
||||
"instantiate a process for the parameter server, along with multiple\n",
|
||||
"workers."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "e3b47dc6",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"iterations = 200\n",
|
||||
"num_workers = 2\n",
|
||||
"\n",
|
||||
"ray.init(ignore_reinit_error=True)\n",
|
||||
"ps = ParameterServer.remote(1e-2)\n",
|
||||
"workers = [DataWorker.remote() for i in range(num_workers)]"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "3545f369",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"We'll also instantiate a model on the driver process to evaluate the test\n",
|
||||
"accuracy during training."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "f4db4f83",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"model = ConvNet()\n",
|
||||
"test_loader = get_data_loader()[1]"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "a7a721b7",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"Training alternates between:\n",
|
||||
"\n",
|
||||
"1. Computing the gradients given the current weights from the server\n",
|
||||
"2. Updating the parameter server's weights with the gradients."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "75da1bf4",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"print(\"Running synchronous parameter server training.\")\n",
|
||||
"current_weights = ps.get_weights.remote()\n",
|
||||
"for i in range(iterations):\n",
|
||||
" gradients = [worker.compute_gradients.remote(current_weights) for worker in workers]\n",
|
||||
" # Calculate update after all gradients are available.\n",
|
||||
" current_weights = ps.apply_gradients.remote(*gradients)\n",
|
||||
"\n",
|
||||
" if i % 10 == 0:\n",
|
||||
" # Evaluate the current model.\n",
|
||||
" model.set_weights(ray.get(current_weights))\n",
|
||||
" accuracy = evaluate(model, test_loader)\n",
|
||||
" print(\"Iter {}: \\taccuracy is {:.1f}\".format(i, accuracy))\n",
|
||||
"\n",
|
||||
"print(\"Final accuracy is {:.1f}.\".format(accuracy))\n",
|
||||
"# Clean up Ray resources and processes before the next example.\n",
|
||||
"ray.shutdown()"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "11c7f855",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Asynchronous Parameter Server Training\n",
|
||||
"\n",
|
||||
"We'll now create an asynchronous parameter server training scheme. We'll first\n",
|
||||
"instantiate a process for the parameter server, along with multiple\n",
|
||||
"workers."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "dcc5407d",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"print(\"Running Asynchronous Parameter Server Training.\")\n",
|
||||
"\n",
|
||||
"ray.init(ignore_reinit_error=True)\n",
|
||||
"ps = ParameterServer.remote(1e-2)\n",
|
||||
"workers = [DataWorker.remote() for i in range(num_workers)]"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "1dc42aa6",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"Here, workers will asynchronously compute the gradients given its\n",
|
||||
"current weights and send these gradients to the parameter server as\n",
|
||||
"soon as they are ready. When the Parameter server finishes applying the\n",
|
||||
"new gradient, the server will send back a copy of the current weights to the\n",
|
||||
"worker. The worker will then update the weights and repeat."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "4013539b",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"current_weights = ps.get_weights.remote()\n",
|
||||
"\n",
|
||||
"gradients = {}\n",
|
||||
"for worker in workers:\n",
|
||||
" gradients[worker.compute_gradients.remote(current_weights)] = worker\n",
|
||||
"\n",
|
||||
"for i in range(iterations * num_workers):\n",
|
||||
" ready_gradient_list, _ = ray.wait(list(gradients))\n",
|
||||
" ready_gradient_id = ready_gradient_list[0]\n",
|
||||
" worker = gradients.pop(ready_gradient_id)\n",
|
||||
"\n",
|
||||
" # Compute and apply gradients.\n",
|
||||
" current_weights = ps.apply_gradients.remote(*[ready_gradient_id])\n",
|
||||
" gradients[worker.compute_gradients.remote(current_weights)] = worker\n",
|
||||
"\n",
|
||||
" if i % 10 == 0:\n",
|
||||
" # Evaluate the current model after every 10 updates.\n",
|
||||
" model.set_weights(ray.get(current_weights))\n",
|
||||
" accuracy = evaluate(model, test_loader)\n",
|
||||
" print(\"Iter {}: \\taccuracy is {:.1f}\".format(i, accuracy))\n",
|
||||
"\n",
|
||||
"print(\"Final accuracy is {:.1f}.\".format(accuracy))"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "dfb8532b",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Final Thoughts\n",
|
||||
"\n",
|
||||
"This approach is powerful because it enables you to implement a parameter\n",
|
||||
"server with a few lines of code as part of a Python application.\n",
|
||||
"As a result, this simplifies the deployment of applications that use\n",
|
||||
"parameter servers and to modify the behavior of the parameter server.\n",
|
||||
"\n",
|
||||
"For example, sharding the parameter server, changing the update rule,\n",
|
||||
"switching between asynchronous and synchronous updates, ignoring\n",
|
||||
"straggler workers, or any number of other customizations,\n",
|
||||
"will only require a few extra lines of code."
|
||||
]
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"kernelspec": {
|
||||
"display_name": "Python 3",
|
||||
"language": "python",
|
||||
"name": "python3"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 5
|
||||
}
|
||||
@@ -0,0 +1,396 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "f867a908",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# Learning to Play Pong\n",
|
||||
"\n",
|
||||
"<a id=\"try-anyscale-quickstart-plot_pong_example\" href=\"https://console.anyscale.com/register/ha?render_flow=ray&utm_source=ray_docs&utm_medium=docs&utm_campaign=plot_pong_example\">\n",
|
||||
" <img src=\"../../_static/img/run-on-anyscale.svg\" alt=\"try-anyscale-quickstart\">\n",
|
||||
"</a>\n",
|
||||
"<br></br>\n",
|
||||
"\n",
|
||||
"```{tip}\n",
|
||||
"For a production-grade implementation of distributed\n",
|
||||
"reinforcement learning, use [Ray RLlib](https://docs.ray.io/en/master/rllib/index.html).\n",
|
||||
"```\n",
|
||||
"\n",
|
||||
"In this example, we'll train a **very simple** neural network to play Pong using\n",
|
||||
"Gymnasium.\n",
|
||||
"\n",
|
||||
"At a high level, we will use multiple Ray actors to obtain simulation rollouts\n",
|
||||
"and calculate gradient simultaneously. We will then centralize these\n",
|
||||
"gradients and update the neural network. The updated neural network will\n",
|
||||
"then be passed back to each Ray actor for more gradient calculation.\n",
|
||||
"\n",
|
||||
"This application is adapted, with minimal modifications, from\n",
|
||||
"Andrej Karpathy's [source code](https://gist.github.com/karpathy/a4166c7fe253700972fcbc77e4ea32c5)\n",
|
||||
"(see the accompanying [blog post](http://karpathy.github.io/2016/05/31/rl/)).\n",
|
||||
"\n",
|
||||
"```{image} /ray-core/images/pong-arch.svg\n",
|
||||
":align: center\n",
|
||||
"```\n",
|
||||
"\n",
|
||||
"To run the application, first install some dependencies.\n",
|
||||
"\n",
|
||||
"```bash\n",
|
||||
"pip install gymnasium[atari]==0.28.1\n",
|
||||
"```\n",
|
||||
"\n",
|
||||
"At the moment, on a large machine with 64 physical cores, computing an update\n",
|
||||
"with a batch of size 1 takes about 1 second, a batch of size 10 takes about 2.5\n",
|
||||
"seconds. A batch of size 60 takes about 3 seconds. On a cluster with 11 nodes,\n",
|
||||
"each with 18 physical cores, a batch of size 300 takes about 10 seconds. If the\n",
|
||||
"numbers you see differ from these by much, take a look at the\n",
|
||||
"**Troubleshooting** section at the bottom of this page and consider\n",
|
||||
"[submitting an issue](https://github.com/ray-project/ray/issues).\n",
|
||||
"\n",
|
||||
"**Note** that these times depend on how long the rollouts take, which in turn\n",
|
||||
"depends on how well the policy is doing. For example, a really bad policy will\n",
|
||||
"lose very quickly. As the policy learns, we should expect these numbers to\n",
|
||||
"increase."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "549e3475",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"import numpy as np\n",
|
||||
"import os\n",
|
||||
"import ray\n",
|
||||
"import time\n",
|
||||
"\n",
|
||||
"import gymnasium as gym"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "39e69bfd",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Hyperparameters\n",
|
||||
"\n",
|
||||
"Here we'll define a couple of the hyperparameters that are used."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "9cb838df",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"H = 200 # The number of hidden layer neurons.\n",
|
||||
"gamma = 0.99 # The discount factor for reward.\n",
|
||||
"decay_rate = 0.99 # The decay factor for RMSProp leaky sum of grad^2.\n",
|
||||
"D = 80 * 80 # The input dimensionality: 80x80 grid.\n",
|
||||
"learning_rate = 1e-4 # Magnitude of the update."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "a0efd0b2",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Helper Functions\n",
|
||||
"\n",
|
||||
"We first define a few helper functions:\n",
|
||||
"\n",
|
||||
"1. Preprocessing: The ``preprocess`` function will\n",
|
||||
"preprocess the original 210x160x3 uint8 frame into a one-dimensional 6400\n",
|
||||
"float vector.\n",
|
||||
"\n",
|
||||
"2. Reward Processing: The ``process_rewards`` function will calculate\n",
|
||||
"a discounted reward. This formula states that the \"value\" of a\n",
|
||||
"sampled action is the weighted sum of all rewards afterwards,\n",
|
||||
"but later rewards are exponentially less important.\n",
|
||||
"\n",
|
||||
"3. Rollout: The ``rollout`` function plays an entire game of Pong (until\n",
|
||||
"either the computer or the RL agent loses)."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "d20fd47c",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"def preprocess(img):\n",
|
||||
" # Crop the image.\n",
|
||||
" img = img[35:195]\n",
|
||||
" # Downsample by factor of 2.\n",
|
||||
" img = img[::2, ::2, 0]\n",
|
||||
" # Erase background (background type 1).\n",
|
||||
" img[img == 144] = 0\n",
|
||||
" # Erase background (background type 2).\n",
|
||||
" img[img == 109] = 0\n",
|
||||
" # Set everything else (paddles, ball) to 1.\n",
|
||||
" img[img != 0] = 1\n",
|
||||
" return img.astype(float).ravel()\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"def process_rewards(r):\n",
|
||||
" \"\"\"Compute discounted reward from a vector of rewards.\"\"\"\n",
|
||||
" discounted_r = np.zeros_like(r)\n",
|
||||
" running_add = 0\n",
|
||||
" for t in reversed(range(0, r.size)):\n",
|
||||
" # Reset the sum, since this was a game boundary (pong specific!).\n",
|
||||
" if r[t] != 0:\n",
|
||||
" running_add = 0\n",
|
||||
" running_add = running_add * gamma + r[t]\n",
|
||||
" discounted_r[t] = running_add\n",
|
||||
" return discounted_r\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"def rollout(model, env):\n",
|
||||
" \"\"\"Evaluates env and model until the env returns \"Terminated\" or \"Truncated\".\n",
|
||||
"\n",
|
||||
" Returns:\n",
|
||||
" xs: A list of observations\n",
|
||||
" hs: A list of model hidden states per observation\n",
|
||||
" dlogps: A list of gradients\n",
|
||||
" drs: A list of rewards.\n",
|
||||
"\n",
|
||||
" \"\"\"\n",
|
||||
" # Reset the game.\n",
|
||||
" observation, info = env.reset()\n",
|
||||
" # Note that prev_x is used in computing the difference frame.\n",
|
||||
" prev_x = None\n",
|
||||
" xs, hs, dlogps, drs = [], [], [], []\n",
|
||||
" terminated = truncated = False\n",
|
||||
" while not terminated and not truncated:\n",
|
||||
" cur_x = preprocess(observation)\n",
|
||||
" x = cur_x - prev_x if prev_x is not None else np.zeros(D)\n",
|
||||
" prev_x = cur_x\n",
|
||||
"\n",
|
||||
" aprob, h = model.policy_forward(x)\n",
|
||||
" # Sample an action.\n",
|
||||
" action = 2 if np.random.uniform() < aprob else 3\n",
|
||||
"\n",
|
||||
" # The observation.\n",
|
||||
" xs.append(x)\n",
|
||||
" # The hidden state.\n",
|
||||
" hs.append(h)\n",
|
||||
" y = 1 if action == 2 else 0 # A \"fake label\".\n",
|
||||
" # The gradient that encourages the action that was taken to be\n",
|
||||
" # taken (see http://cs231n.github.io/neural-networks-2/#losses if\n",
|
||||
" # confused).\n",
|
||||
" dlogps.append(y - aprob)\n",
|
||||
"\n",
|
||||
" observation, reward, terminated, truncated, info = env.step(action)\n",
|
||||
"\n",
|
||||
" # Record reward (has to be done after we call step() to get reward\n",
|
||||
" # for previous action).\n",
|
||||
" drs.append(reward)\n",
|
||||
" return xs, hs, dlogps, drs"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "7c00c00a",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Neural Network\n",
|
||||
"\n",
|
||||
"Here, a neural network is used to define a \"policy\"\n",
|
||||
"for playing Pong (that is, a function that chooses an action given a state).\n",
|
||||
"\n",
|
||||
"To implement a neural network in NumPy, we need to provide helper functions\n",
|
||||
"for calculating updates and computing the output of the neural network\n",
|
||||
"given an input, which in our case is an observation."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "8992067a",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"class Model(object):\n",
|
||||
" \"\"\"This class holds the neural network weights.\"\"\"\n",
|
||||
"\n",
|
||||
" def __init__(self):\n",
|
||||
" self.weights = {}\n",
|
||||
" self.weights[\"W1\"] = np.random.randn(H, D) / np.sqrt(D)\n",
|
||||
" self.weights[\"W2\"] = np.random.randn(H) / np.sqrt(H)\n",
|
||||
"\n",
|
||||
" def policy_forward(self, x):\n",
|
||||
" h = np.dot(self.weights[\"W1\"], x)\n",
|
||||
" h[h < 0] = 0 # ReLU nonlinearity.\n",
|
||||
" logp = np.dot(self.weights[\"W2\"], h)\n",
|
||||
" # Softmax\n",
|
||||
" p = 1.0 / (1.0 + np.exp(-logp))\n",
|
||||
" # Return probability of taking action 2, and hidden state.\n",
|
||||
" return p, h\n",
|
||||
"\n",
|
||||
" def policy_backward(self, eph, epx, epdlogp):\n",
|
||||
" \"\"\"Backward pass to calculate gradients.\n",
|
||||
"\n",
|
||||
" Arguments:\n",
|
||||
" eph: Array of intermediate hidden states.\n",
|
||||
" epx: Array of experiences (observations).\n",
|
||||
" epdlogp: Array of logps (output of last layer before softmax).\n",
|
||||
"\n",
|
||||
" \"\"\"\n",
|
||||
" dW2 = np.dot(eph.T, epdlogp).ravel()\n",
|
||||
" dh = np.outer(epdlogp, self.weights[\"W2\"])\n",
|
||||
" # Backprop relu.\n",
|
||||
" dh[eph <= 0] = 0\n",
|
||||
" dW1 = np.dot(dh.T, epx)\n",
|
||||
" return {\"W1\": dW1, \"W2\": dW2}\n",
|
||||
"\n",
|
||||
" def update(self, grad_buffer, rmsprop_cache, lr, decay):\n",
|
||||
" \"\"\"Applies the gradients to the model parameters with RMSProp.\"\"\"\n",
|
||||
" for k, v in self.weights.items():\n",
|
||||
" g = grad_buffer[k]\n",
|
||||
" rmsprop_cache[k] = decay * rmsprop_cache[k] + (1 - decay) * g ** 2\n",
|
||||
" self.weights[k] += lr * g / (np.sqrt(rmsprop_cache[k]) + 1e-5)\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"def zero_grads(grad_buffer):\n",
|
||||
" \"\"\"Reset the batch gradient buffer.\"\"\"\n",
|
||||
" for k, v in grad_buffer.items():\n",
|
||||
" grad_buffer[k] = np.zeros_like(v)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "c4a847bd",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Parallelizing Gradients\n",
|
||||
"\n",
|
||||
"We define an **actor**, which is responsible for taking a model and an env\n",
|
||||
"and performing a rollout + computing a gradient update."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "c95ee2f2",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"# This forces OpenMP to use 1 single thread, which is needed to \n",
|
||||
"# prevent contention between multiple actors. \n",
|
||||
"# See https://docs.ray.io/en/latest/ray-core/configure.html for \n",
|
||||
"# more details. \n",
|
||||
"os.environ[\"OMP_NUM_THREADS\"] = \"1\"\n",
|
||||
"# Tell numpy to only use one core. If we don't do this, each actor may\n",
|
||||
"# try to use all of the cores and the resulting contention may result\n",
|
||||
"# in no speedup over the serial version. Note that if numpy is using\n",
|
||||
"# OpenBLAS, then you need to set OPENBLAS_NUM_THREADS=1, and you\n",
|
||||
"# probably need to do it from the command line (so it happens before\n",
|
||||
"# numpy is imported).\n",
|
||||
"os.environ[\"MKL_NUM_THREADS\"] = \"1\"\n",
|
||||
"\n",
|
||||
"ray.init()\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"@ray.remote\n",
|
||||
"class RolloutWorker(object):\n",
|
||||
" def __init__(self):\n",
|
||||
" self.env = gym.make(\"ale_py:ALE/Pong-v5\")\n",
|
||||
"\n",
|
||||
" def compute_gradient(self, model):\n",
|
||||
" # Compute a simulation episode.\n",
|
||||
" xs, hs, dlogps, drs = rollout(model, self.env)\n",
|
||||
" reward_sum = sum(drs)\n",
|
||||
" # Vectorize the arrays.\n",
|
||||
" epx = np.vstack(xs)\n",
|
||||
" eph = np.vstack(hs)\n",
|
||||
" epdlogp = np.vstack(dlogps)\n",
|
||||
" epr = np.vstack(drs)\n",
|
||||
"\n",
|
||||
" # Compute the discounted reward backward through time.\n",
|
||||
" discounted_epr = process_rewards(epr)\n",
|
||||
" # Standardize the rewards to be unit normal (helps control the gradient\n",
|
||||
" # estimator variance).\n",
|
||||
" discounted_epr -= np.mean(discounted_epr)\n",
|
||||
" discounted_epr /= np.std(discounted_epr)\n",
|
||||
" # Modulate the gradient with advantage (the policy gradient magic\n",
|
||||
" # happens right here).\n",
|
||||
" epdlogp *= discounted_epr\n",
|
||||
" return model.policy_backward(eph, epx, epdlogp), reward_sum"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "1ce7f4da",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Running\n",
|
||||
"\n",
|
||||
"This example is easy to parallelize because the network can play ten games\n",
|
||||
"in parallel and no information needs to be shared between the games.\n",
|
||||
"\n",
|
||||
"In the loop, the network repeatedly plays games of Pong and\n",
|
||||
"records a gradient from each game. Every ten games, the gradients are\n",
|
||||
"combined together and used to update the network."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "e353bd1e",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"iterations = 20\n",
|
||||
"batch_size = 4\n",
|
||||
"model = Model()\n",
|
||||
"actors = [RolloutWorker.remote() for _ in range(batch_size)]\n",
|
||||
"\n",
|
||||
"running_reward = None\n",
|
||||
"# \"Xavier\" initialization.\n",
|
||||
"# Update buffers that add up gradients over a batch.\n",
|
||||
"grad_buffer = {k: np.zeros_like(v) for k, v in model.weights.items()}\n",
|
||||
"# Update the rmsprop memory.\n",
|
||||
"rmsprop_cache = {k: np.zeros_like(v) for k, v in model.weights.items()}\n",
|
||||
"\n",
|
||||
"for i in range(1, 1 + iterations):\n",
|
||||
" model_id = ray.put(model)\n",
|
||||
" gradient_ids = []\n",
|
||||
" # Launch tasks to compute gradients from multiple rollouts in parallel.\n",
|
||||
" start_time = time.time()\n",
|
||||
" gradient_ids = [actor.compute_gradient.remote(model_id) for actor in actors]\n",
|
||||
" for batch in range(batch_size):\n",
|
||||
" [grad_id], gradient_ids = ray.wait(gradient_ids)\n",
|
||||
" grad, reward_sum = ray.get(grad_id)\n",
|
||||
" # Accumulate the gradient over batch.\n",
|
||||
" for k in model.weights:\n",
|
||||
" grad_buffer[k] += grad[k]\n",
|
||||
" running_reward = (\n",
|
||||
" reward_sum\n",
|
||||
" if running_reward is None\n",
|
||||
" else running_reward * 0.99 + reward_sum * 0.01\n",
|
||||
" )\n",
|
||||
" end_time = time.time()\n",
|
||||
" print(\n",
|
||||
" \"Batch {} computed {} rollouts in {} seconds, \"\n",
|
||||
" \"running mean is {}\".format(\n",
|
||||
" i, batch_size, end_time - start_time, running_reward\n",
|
||||
" )\n",
|
||||
" )\n",
|
||||
" model.update(grad_buffer, rmsprop_cache, learning_rate, decay_rate)\n",
|
||||
" zero_grads(grad_buffer)"
|
||||
]
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"kernelspec": {
|
||||
"display_name": "Python 3",
|
||||
"language": "python",
|
||||
"name": "python3"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 5
|
||||
}
|
||||
@@ -0,0 +1,428 @@
|
||||
"""
|
||||
Reinforcement learning example using GPU-to-GPU Ray Direct Transport (RDT) and GRPO algorithm.
|
||||
|
||||
Based on: https://github.com/meta-pytorch/monarch/blob/0de4e6b4ad7da37e5dbb00a0e6fb61ef8105eac5/examples/presentation/demo.py
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import time
|
||||
from typing import Any
|
||||
|
||||
import ray
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.optim as optim
|
||||
from torch.distributions import Categorical
|
||||
|
||||
# -- TrajectorySlice --
|
||||
# TrajectorySlice holds one state's sampled actions and associated metadata:
|
||||
# - state: The 2D input vector fed to the generator model.
|
||||
# - actions: The generator model's predictions for this state.
|
||||
# - policy_version: Version of the generator model when these actions were generated.
|
||||
# - rewards: The per-action rewards computed by the scorer for this state.
|
||||
# - old_logps: The log-probabilities of the sampled actions under the policy that generated them.
|
||||
TrajectorySlice = dict[str, torch.Tensor | int]
|
||||
|
||||
|
||||
# -- Training --
|
||||
BATCH_SIZE = 32
|
||||
# Keep learning rate low so that the model does not jump outside
|
||||
# the trust policy region.
|
||||
LEARNING_RATE = 1e-6
|
||||
WEIGHT_DECAY = 1e-10
|
||||
# Adaptively reduces the learning rate to prevent large updates in a single step.
|
||||
GRAD_CLIP_NORM = 1.0
|
||||
|
||||
# -- GRPO algorithm --
|
||||
# Number of actions to sample for each state.
|
||||
GROUP_SIZE = 10
|
||||
# How far the new policy is allowed to stray from the older policies
|
||||
# AKA the "trust region".
|
||||
GRPO_CLIP_EPS = 0.1
|
||||
# Discard old experiences, so that the new model can gradually explore
|
||||
# further from the initial random policy.
|
||||
MAX_BUFFER_SIZE = BATCH_SIZE * GROUP_SIZE * 5
|
||||
|
||||
# -- Environment --
|
||||
STATE_DIM = 2 # The contextual bandit operates in 2D.
|
||||
ACTION_DIM = 8 # Eight compass directions: [W, NW, N, NE, E, SE, S, SW].
|
||||
# Unit direction vectors for the eight compass actions (W, NW, N, NE, E, SE, S, SW).
|
||||
DIAGONAL_MAGNITUDE = 2**0.5 / 2.0
|
||||
ACTION_DIRECTIONS = torch.tensor(
|
||||
[
|
||||
[-1.0, 0.0], # W
|
||||
[-DIAGONAL_MAGNITUDE, DIAGONAL_MAGNITUDE], # NW
|
||||
[0.0, 1.0], # N
|
||||
[DIAGONAL_MAGNITUDE, DIAGONAL_MAGNITUDE], # NE
|
||||
[1.0, 0.0], # E
|
||||
[DIAGONAL_MAGNITUDE, -DIAGONAL_MAGNITUDE], # SE
|
||||
[0.0, -1.0], # S
|
||||
[-DIAGONAL_MAGNITUDE, -DIAGONAL_MAGNITUDE], # SW
|
||||
],
|
||||
dtype=torch.float32,
|
||||
)
|
||||
|
||||
|
||||
# -- Model --
|
||||
# To demonstrate speed-ups from RDT, we use an oversized model.
|
||||
# Residual connections prevent vanishing gradients for deep models.
|
||||
class ResidualBlock(torch.nn.Module):
|
||||
def __init__(self, hidden_dim: int) -> None:
|
||||
super().__init__()
|
||||
self.norm = torch.nn.LayerNorm(hidden_dim)
|
||||
self.activation = torch.nn.ReLU()
|
||||
self.linear = torch.nn.Linear(hidden_dim, hidden_dim, bias=True)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
residual = x
|
||||
out = self.norm(x)
|
||||
out = self.activation(out)
|
||||
out = self.linear(out)
|
||||
return residual + out
|
||||
|
||||
|
||||
class ResidualMLP(torch.nn.Module): # Sized to ~50 MB of parameters.
|
||||
"""Model used for Generator and Learner.
|
||||
|
||||
It takes a 2D state vector as input and produces logits for each action.
|
||||
"""
|
||||
|
||||
def __init__(self, hidden_dim: int = 512, depth: int = 50):
|
||||
super().__init__()
|
||||
self.input = torch.nn.Linear(STATE_DIM, hidden_dim, bias=True)
|
||||
self.backbone = torch.nn.ModuleList(
|
||||
ResidualBlock(hidden_dim) for _ in range(depth - 1)
|
||||
)
|
||||
self.head = torch.nn.Linear(hidden_dim, ACTION_DIM, bias=True)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
x = self.input(x)
|
||||
for block in self.backbone:
|
||||
x = block(x)
|
||||
x = self.head(x)
|
||||
return x
|
||||
|
||||
|
||||
# -- Utilities --
|
||||
def sample_unit_vector(batch_size: int, dim: int = STATE_DIM) -> torch.Tensor:
|
||||
"""Sample unit vectors of shape [batch_size, dim] by normalizing Gaussian draws."""
|
||||
assert batch_size > 1, "Batch size must be greater than 1"
|
||||
v = torch.randn(batch_size, dim)
|
||||
norms = v.norm(dim=-1, keepdim=True) + 1e-8
|
||||
return v / norms
|
||||
|
||||
|
||||
# -- Actors --
|
||||
@ray.remote
|
||||
class ReplayBuffer:
|
||||
"""Storage for scored trajectory slices.
|
||||
|
||||
This class stores the past experiences (AKA trajectories, or slices) of the model.
|
||||
This allows the learner to sample and learn from the same experiences multiple times
|
||||
by comparing the latest model with previous models.
|
||||
|
||||
The sampler weights the trajectories by the policy version, such that trajectories produced
|
||||
by more recent versions of the model are more likely to be sampled.
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
# Each entry stores a TrajectorySlice with CPU tensors.
|
||||
self.storage: list[TrajectorySlice] = []
|
||||
|
||||
def put(self, slice: TrajectorySlice) -> None:
|
||||
"""Add a new slice to the buffer.
|
||||
|
||||
The buffer discards the oldest slices if the buffer gets too large to prevent memory leaks,
|
||||
and so that the latest model can gradually explore further from the initial random policy.
|
||||
"""
|
||||
self.storage.append(slice)
|
||||
if len(self.storage) > MAX_BUFFER_SIZE:
|
||||
self.storage = self.storage[-MAX_BUFFER_SIZE:]
|
||||
|
||||
def sample_from(self, n: int) -> list[TrajectorySlice]:
|
||||
"""Sample n scored trajectory slices.
|
||||
|
||||
Each slice is a 'group' of actions sampled from the same state.
|
||||
"""
|
||||
if self.size() < n:
|
||||
return []
|
||||
# The probability of sampling a slice is proportional to its policy version.
|
||||
total = sum(slice["policy_version"] for slice in self.storage)
|
||||
probs = [slice["policy_version"] / total for slice in self.storage]
|
||||
# Sample with replacement without exceeding the buffer's size.
|
||||
n = min(n, self.size())
|
||||
chosen = np.random.choice(self.size(), size=n, p=probs, replace=True)
|
||||
return [self.storage[i] for i in chosen]
|
||||
|
||||
def size(self) -> int:
|
||||
return len(self.storage)
|
||||
|
||||
|
||||
@ray.remote
|
||||
class Scorer:
|
||||
"""Evaluates actions and assigns rewards to trajectory slices.
|
||||
|
||||
This scorer implements an analytic contextual bandit reward: for a 2D unit
|
||||
context vector `s` and a discrete action `a` in {W, NW, N, NE, E, SE, S, SW},
|
||||
reward is cosine_similarity(s, direction[a]) == dot(s, unit_direction[a]).
|
||||
"""
|
||||
|
||||
def __init__(self, replay_buffer) -> None:
|
||||
self.replay_buffer = replay_buffer
|
||||
self.action_dirs = ACTION_DIRECTIONS # [ACTION_DIM, STATE_DIM]
|
||||
|
||||
@ray.method(tensor_transport="nixl") # CPU-CPU RDT
|
||||
def score_slices(self, batched_slices: dict) -> None:
|
||||
"""Score a batch of trajectory slices."""
|
||||
states = batched_slices["state"]
|
||||
actions = batched_slices["actions"]
|
||||
old_logps = batched_slices["old_logps"]
|
||||
policy_version = batched_slices["policy_version"]
|
||||
|
||||
# Unbatch the groups into separate slices so that they can be
|
||||
# sampled independently.
|
||||
for i in range(states.shape[0]):
|
||||
# Compute rewards on the CPU: rewards = dot(state, unit_dir).
|
||||
directions = self.action_dirs[actions[i]] # [GROUP_SIZE, STATE_DIM]
|
||||
rewards = torch.mv(directions, states[i])
|
||||
|
||||
scored = TrajectorySlice(
|
||||
policy_version=policy_version,
|
||||
state=states[i],
|
||||
actions=actions[i],
|
||||
old_logps=old_logps[i],
|
||||
rewards=rewards,
|
||||
)
|
||||
|
||||
self.replay_buffer.put.remote(scored)
|
||||
|
||||
|
||||
@ray.remote(num_gpus=1)
|
||||
class Learner:
|
||||
"""Updates policy based on collected experiences using GRPO algorithm."""
|
||||
|
||||
def __init__(self, replay_buffer) -> None:
|
||||
self.model = ResidualMLP().to("cuda")
|
||||
|
||||
# Use smaller betas to favor recent momentum history.
|
||||
self.optim = optim.AdamW(
|
||||
self.model.parameters(),
|
||||
lr=LEARNING_RATE,
|
||||
weight_decay=WEIGHT_DECAY,
|
||||
betas=(0.9, 0.9),
|
||||
)
|
||||
self.replay_buffer = replay_buffer
|
||||
|
||||
def _compute_advantages(self, rewards: torch.Tensor) -> torch.Tensor:
|
||||
"""Compute advantages from rewards.
|
||||
|
||||
In GRPO, advantages represent how much better a reward is compared to the mean reward for the group of actions
|
||||
Normalizing the advantages stabilizes training by maintaining a consistent scale of updates.
|
||||
"""
|
||||
# Unflatten rewards into [batch_size, GROUP_SIZE] in order to
|
||||
# compute per-state mean baselines.
|
||||
batch_size = rewards.shape[0] // GROUP_SIZE
|
||||
rewards_reshaped = rewards.view(batch_size, GROUP_SIZE)
|
||||
|
||||
# Compute the mean reward for each state's group of actions.
|
||||
baselines = rewards_reshaped.mean(dim=1, keepdim=True) # [batch_size, 1]
|
||||
|
||||
# Subtract the mean reward from each action's reward to get advantages.
|
||||
advantages = rewards_reshaped - baselines # [batch_size, GROUP_SIZE]
|
||||
|
||||
# Flatten the advantages back to the original shape.
|
||||
advantages = advantages.reshape(-1) # [batch_size * GROUP_SIZE]
|
||||
|
||||
# Normalize the advantages for training stability.
|
||||
advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)
|
||||
|
||||
return advantages
|
||||
|
||||
def _apply_policy_update(
|
||||
self,
|
||||
states: torch.Tensor,
|
||||
actions: torch.Tensor,
|
||||
old_logps: torch.Tensor,
|
||||
advantages: torch.Tensor,
|
||||
) -> dict[str, float]:
|
||||
"""Apply GRPO update to the model."""
|
||||
# Compute the new policy's action log-probabilities.
|
||||
dist_new = Categorical(logits=self.model(states))
|
||||
new_logps = dist_new.log_prob(actions)
|
||||
# Compare the new log-probabilities to the old log-probabilities to get the probability ratios.
|
||||
# This is a proxy for how different the new policy is from the old policy.
|
||||
ratio = (new_logps - old_logps).exp()
|
||||
unclipped = ratio * advantages
|
||||
# The 1 ± ε ratio defines the trust region. If the new policy's probability for an action is more than 1 ± ε times the old policy, clip the ratio
|
||||
# to prevent too-large updates.
|
||||
clipped = torch.clamp(ratio, 1 - GRPO_CLIP_EPS, 1 + GRPO_CLIP_EPS) * advantages
|
||||
loss = -torch.min(unclipped, clipped).mean()
|
||||
# Fraction of actions which did not contribute to the gradient update.
|
||||
clip_fraction = (
|
||||
((ratio < 1 - GRPO_CLIP_EPS) | (ratio > 1 + GRPO_CLIP_EPS)).float().mean()
|
||||
)
|
||||
|
||||
# Update the policy network.
|
||||
self.optim.zero_grad()
|
||||
loss.backward()
|
||||
# Clip the gradients to prevent exploding gradients and stabilize training.
|
||||
nn.utils.clip_grad_norm_(self.model.parameters(), GRAD_CLIP_NORM)
|
||||
self.optim.step()
|
||||
|
||||
return {
|
||||
"loss": loss.detach().item(),
|
||||
"clip_fraction": clip_fraction.detach().item(),
|
||||
}
|
||||
|
||||
def step(self) -> dict[str, Any]:
|
||||
"""Perform one training step and return metrics.
|
||||
|
||||
Each step samples a batch of trajectory slices from the replay buffer, computes the advantages, and updates the policy using the GRPO algorithm.
|
||||
"""
|
||||
slices: list[TrajectorySlice] = ray.get(
|
||||
self.replay_buffer.sample_from.remote(BATCH_SIZE)
|
||||
)
|
||||
while len(slices) < BATCH_SIZE:
|
||||
print(
|
||||
f"Not enough slices in the buffer to sample {BATCH_SIZE} slices. Waiting for more slices..."
|
||||
)
|
||||
time.sleep(0.05)
|
||||
slices = ray.get(self.replay_buffer.sample_from.remote(BATCH_SIZE))
|
||||
|
||||
# Prepare the tensors for the policy update.
|
||||
actions = torch.cat([s["actions"] for s in slices]).to("cuda")
|
||||
old_logps = torch.cat([s["old_logps"] for s in slices]).to("cuda")
|
||||
rewards = torch.cat([s["rewards"] for s in slices]).to("cuda")
|
||||
mean_rewards = torch.mean(rewards).item()
|
||||
states = torch.stack([s["state"] for s in slices])
|
||||
states = states.repeat_interleave(GROUP_SIZE, 0).to("cuda")
|
||||
|
||||
# Compute advantages and update the policy network using GRPO.
|
||||
advantages = self._compute_advantages(rewards)
|
||||
results = self._apply_policy_update(states, actions, old_logps, advantages)
|
||||
results["rewards"] = mean_rewards
|
||||
|
||||
return results
|
||||
|
||||
@ray.method(tensor_transport="nixl")
|
||||
def get_weights(self) -> dict[str, torch.Tensor]:
|
||||
"""Get the current model weights.
|
||||
|
||||
The tensor_transport="nixl" option enables NIXL via RDT to transfer model weight
|
||||
tensors. Without it, the weights will be transferred using the Ray object store.
|
||||
"""
|
||||
return self.model.state_dict()
|
||||
|
||||
|
||||
@ray.remote(num_gpus=1)
|
||||
class Generator:
|
||||
"""Holds the current policy network and generates unscored trajectory slices."""
|
||||
|
||||
def __init__(self, scorer) -> None:
|
||||
self.model = ResidualMLP().to("cuda").eval()
|
||||
self.scorer = scorer
|
||||
self.policy_version = 1
|
||||
|
||||
@ray.method(tensor_transport="nixl") # CPU-CPU RDT
|
||||
def generate(self, states: torch.Tensor):
|
||||
"""Generate actions using the current policy and send them and their metadata
|
||||
to the Scorer.
|
||||
|
||||
Note: GRPO requires *sampling* from the current policy (not just the most probable "greedy" action).
|
||||
"""
|
||||
with torch.no_grad():
|
||||
states = states.to("cuda")
|
||||
logits = self.model(states) # [batch_size, ACTION_DIM]
|
||||
|
||||
dist = Categorical(logits=logits)
|
||||
# Sample GROUP_SIZE actions for each state.
|
||||
actions = dist.sample((GROUP_SIZE,)) # [GROUP_SIZE, batch_size]
|
||||
logps = dist.log_prob(actions) # [GROUP_SIZE, batch_size]
|
||||
# Transpose actions and logprobs for compatibility with the states tensor.
|
||||
actions = actions.transpose(0, 1).contiguous() # [batch_size, GROUP_SIZE]
|
||||
logps = logps.transpose(0, 1).contiguous() # [batch_size, GROUP_SIZE]
|
||||
|
||||
# Create trajectory slices and enqueue them for scoring.
|
||||
slice_batch = {
|
||||
"policy_version": self.policy_version,
|
||||
"state": states.detach().cpu(),
|
||||
"actions": actions.detach().cpu(),
|
||||
"old_logps": logps.detach().cpu(),
|
||||
}
|
||||
self.scorer.score_slices.remote(slice_batch)
|
||||
|
||||
def update_weights(self, cuda_weights):
|
||||
"""Update the generator's weights from the learner's weights.
|
||||
|
||||
Note: the actor is single-threaded, so weight loads do not overlap with generation.
|
||||
"""
|
||||
first_tensor = next(iter(cuda_weights.values()))
|
||||
assert (
|
||||
first_tensor.device.type == "cuda"
|
||||
), "Expected CUDA tensors after GPU-to-GPU direct transfer"
|
||||
self.model.load_state_dict(cuda_weights)
|
||||
self.model.eval()
|
||||
self.policy_version += 1
|
||||
|
||||
|
||||
# -- Control loop --
|
||||
def train(total_steps: int) -> None:
|
||||
"""Run one end-to-end training session."""
|
||||
|
||||
# Instantiate one instance of each actor.
|
||||
replay_buf = ReplayBuffer.remote()
|
||||
learner = Learner.remote(replay_buf)
|
||||
scorer = Scorer.remote(replay_buf)
|
||||
generator = Generator.remote(scorer)
|
||||
|
||||
# Asynchronously initialize the generator with the current learner weights.
|
||||
weights_updated_ref = generator.update_weights.remote(learner.get_weights.remote())
|
||||
|
||||
# Pre-fill the ReplayBuffer before starting GRPO.
|
||||
# Generator is a single-threaded actor, so this generate call won't execute until after the
|
||||
# above update_weights call has completed.
|
||||
generator.generate.remote(sample_unit_vector(batch_size=BATCH_SIZE))
|
||||
step_results = []
|
||||
losses, rewards, clip_fractions = [], [], []
|
||||
for i in range(total_steps):
|
||||
states = sample_unit_vector(batch_size=BATCH_SIZE)
|
||||
generator.generate.remote(states)
|
||||
|
||||
# Wait until the generator has been updated before launching the next learner step.
|
||||
# Otherwise, the weights transfer could still be in progress during the next learner
|
||||
# update, and the generator may receive partially updated weights.
|
||||
ray.wait([weights_updated_ref])
|
||||
|
||||
# Asynchronously log every 20 steps.
|
||||
if len(step_results) >= 20:
|
||||
for step_result in ray.get(step_results):
|
||||
losses.append(step_result["loss"])
|
||||
rewards.append(step_result["rewards"])
|
||||
clip_fractions.append(step_result["clip_fraction"])
|
||||
print(
|
||||
f"Step {i}/{total_steps} | Loss: {sum(losses[-20:]) / 20} | Rewards: {sum(rewards[-20:]) / 20:.3f} | Fraction clipped: {sum(clip_fractions[-20:]) / 20:.3f}"
|
||||
)
|
||||
step_results.clear()
|
||||
|
||||
step_results.append(learner.step.remote())
|
||||
|
||||
# Update the generator with new weights.
|
||||
weights_updated_ref = generator.update_weights.remote(
|
||||
learner.get_weights.remote()
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument(
|
||||
"--steps",
|
||||
type=int,
|
||||
default=450,
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
ray.init(ignore_reinit_error=True)
|
||||
train(total_steps=args.steps)
|
||||
print("Done!")
|
||||
@@ -0,0 +1,251 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"collapsed": false
|
||||
},
|
||||
"source": [
|
||||
"# Speed up your web crawler by parallelizing it with Ray\n",
|
||||
"\n",
|
||||
"<a id=\"try-anyscale-quickstart-ray-core-web-crawler\" href=\"https://console.anyscale.com/register/ha?render_flow=ray&utm_source=ray_docs&utm_medium=docs&utm_campaign=ray-core-web-crawler\">\n",
|
||||
" <img src=\"../../_static/img/run-on-anyscale.svg\" alt=\"try-anyscale-quickstart\">\n",
|
||||
"</a>\n",
|
||||
"<br></br>"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"collapsed": false
|
||||
},
|
||||
"source": [
|
||||
"In this example we'll quickly demonstrate how to build a simple web scraper in Python and\n",
|
||||
"parallelize it with Ray Tasks with minimal code changes.\n",
|
||||
"\n",
|
||||
"To run this example locally on your machine, please first install `ray` and `beautifulsoup` with\n",
|
||||
"\n",
|
||||
"```\n",
|
||||
"pip install \"beautifulsoup4==4.11.1\" \"ray>=2.2.0\"\n",
|
||||
"```\n",
|
||||
"\n",
|
||||
"First, we'll define a function called `find_links` which takes a starting page (`start_url`) to crawl,\n",
|
||||
"and we'll take the Ray documentation as example of such a starting point.\n",
|
||||
"Our crawler simply extracts all available links from the starting URL that contain a given `base_url`\n",
|
||||
"(e.g. in our example we only want to follow links on `http://docs.ray.io`, not any external links).\n",
|
||||
"The `find_links` function is then called recursively with all the links we found this way, until a\n",
|
||||
"certain depth is reached.\n",
|
||||
"\n",
|
||||
"To extract the links from HTML elements on a site, we define a little helper function called\n",
|
||||
"`extract_links`, which takes care of handling relative URLs properly and sets a limit on the\n",
|
||||
"number of links returned from a site (`max_results`) to control the runtime of the crawler more easily.\n",
|
||||
"\n",
|
||||
"Here's the full implementation:"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 154,
|
||||
"metadata": {
|
||||
"collapsed": false
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"import requests\n",
|
||||
"from bs4 import BeautifulSoup\n",
|
||||
"\n",
|
||||
"def extract_links(elements, base_url, max_results=100):\n",
|
||||
" links = []\n",
|
||||
" for e in elements:\n",
|
||||
" url = e[\"href\"]\n",
|
||||
" if \"https://\" not in url:\n",
|
||||
" url = base_url + url\n",
|
||||
" if base_url in url:\n",
|
||||
" links.append(url)\n",
|
||||
" return set(links[:max_results])\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"def find_links(start_url, base_url, depth=2):\n",
|
||||
" if depth == 0:\n",
|
||||
" return set()\n",
|
||||
"\n",
|
||||
" page = requests.get(start_url)\n",
|
||||
" soup = BeautifulSoup(page.content, \"html.parser\")\n",
|
||||
" elements = soup.find_all(\"a\", href=True)\n",
|
||||
" links = extract_links(elements, base_url)\n",
|
||||
"\n",
|
||||
" for url in links:\n",
|
||||
" new_links = find_links(url, base_url, depth-1)\n",
|
||||
" links = links.union(new_links)\n",
|
||||
" return links"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"collapsed": false
|
||||
},
|
||||
"source": [
|
||||
"Let's define a starting and base URL and crawl the Ray docs to a `depth` of 2."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 162,
|
||||
"metadata": {
|
||||
"collapsed": false
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"base = \"https://docs.ray.io/en/latest/\"\n",
|
||||
"docs = base + \"index.html\""
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 163,
|
||||
"metadata": {
|
||||
"collapsed": false
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"CPU times: user 19.3 s, sys: 340 ms, total: 19.7 s\n",
|
||||
"Wall time: 25.8 s\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"591"
|
||||
]
|
||||
},
|
||||
"execution_count": 163,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"%time len(find_links(docs, base))"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"collapsed": false
|
||||
},
|
||||
"source": [
|
||||
"As you can see, crawling the documentation root recursively like this returns a\n",
|
||||
"total of `591` pages and the wall time comes in at around 25 seconds.\n",
|
||||
"\n",
|
||||
"Crawling pages can be parallelized in many ways.\n",
|
||||
"Probably the simplest way is to simple start with multiple starting URLs and call\n",
|
||||
"`find_links` in parallel for each of them.\n",
|
||||
"We can do this with [Ray Tasks](https://docs.ray.io/en/latest/ray-core/tasks.html) in a straightforward way.\n",
|
||||
"We simply use the `ray.remote` decorator to wrap the `find_links` function in a task called `find_links_task` like this:"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 157,
|
||||
"metadata": {
|
||||
"collapsed": false
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"import ray\n",
|
||||
"\n",
|
||||
"@ray.remote\n",
|
||||
"def find_links_task(start_url, base_url, depth=2):\n",
|
||||
" return find_links(start_url, base_url, depth)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"collapsed": false
|
||||
},
|
||||
"source": [
|
||||
"To use this task to kick off a parallel call, the only thing you have to do is use\n",
|
||||
"`find_links_tasks.remote(...)` instead of calling the underlying Python function directly.\n",
|
||||
"\n",
|
||||
"Here's how you run six crawlers in parallel, the first three (redundantly) crawl\n",
|
||||
"`docs.ray.io` again, the other three crawl the main entry points of the Ray RLlib,\n",
|
||||
"Tune, and Serve libraries, respectively:"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 160,
|
||||
"metadata": {
|
||||
"collapsed": false
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"links = [find_links_task.remote(f\"{base}{lib}/index.html\", base)\n",
|
||||
" for lib in [\"\", \"\", \"\", \"rllib\", \"tune\", \"serve\"]]"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 161,
|
||||
"metadata": {
|
||||
"collapsed": false
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"591\n",
|
||||
"591\n",
|
||||
"105\n",
|
||||
"204\n",
|
||||
"105\n",
|
||||
"CPU times: user 65.5 ms, sys: 47.8 ms, total: 113 ms\n",
|
||||
"Wall time: 27.2 s\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"%time for res in ray.get(links): print(len(res))"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"collapsed": false
|
||||
},
|
||||
"source": [
|
||||
"This parallel run crawls around four times the number of pages in roughly the same time as the initial, sequential run.\n",
|
||||
"Note the use of `ray.get` in the timed run to retrieve the results from Ray (the `remote` call promise gets resolved with `get`).\n",
|
||||
"\n",
|
||||
"Of course, there are much smarter ways to create a crawler and efficiently parallelize it, and this example\n",
|
||||
"gives you a starting point to work from."
|
||||
]
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"kernelspec": {
|
||||
"display_name": "Python 3",
|
||||
"language": "python",
|
||||
"name": "python3"
|
||||
},
|
||||
"language_info": {
|
||||
"codemirror_mode": {
|
||||
"name": "ipython",
|
||||
"version": 2
|
||||
},
|
||||
"file_extension": ".py",
|
||||
"mimetype": "text/x-python",
|
||||
"name": "python",
|
||||
"nbconvert_exporter": "python",
|
||||
"pygments_lexer": "ipython3",
|
||||
"version": "3.9.19"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 0
|
||||
}
|
||||
Reference in New Issue
Block a user