Parallel Python with Ray

A Python program runs in one process and, unless you do something about it, uses one CPU core. Training and evaluating models quickly needs more: several cores, a GPU, and sometimes several machines. Ray is a library for running ordinary Python functions and classes in parallel, on one machine or on a cluster, with the same code in both cases. Several sections of this book use it. The SmolVLM lab trains its model inside a Ray Train worker, and the figures and results of the book are computed by Ray jobs on a small GPU cluster. This section covers the three ideas Ray is built on (tasks, objects and actors), how Ray assigns CPUs and GPUs, and how Ray Train runs a PyTorch training loop. Everything up to that point runs in Google Colab or on your own computer. The last part shows how to build a Ray cluster from your own machines and submit work to it, and simulates a two-node cluster on one machine so that you can try it without a second computer.

import sys

if "google.colab" in sys.modules:
    %pip install -q "ray[default,train]"

import os
import time

import numpy as np
import ray

print(f"[INFO] Ray {ray.__version__}")

[INFO] Ray 2.60.0

Starting Ray

ray.init() starts a Ray instance on the current machine: a scheduler, an object store in shared memory, and a pool of worker processes. Your notebook or script is the driver, the process that submits work. The argument num_cpus=4 tells Ray to schedule at most four CPU tasks at a time. These are logical CPUs: Ray uses the number for scheduling and does not check it against the hardware, so the same code behaves the same in Colab, where the machine has two cores, and on a large workstation.

ray.init(num_cpus=4, include_dashboard=True, log_to_driver=False, ignore_reinit_error=True)
print(ray.cluster_resources())

{'CPU': 4.0, 'node:__internal_head__': 1.0, 'GPU': 1.0, 'accelerator_type:RTX-PRO-6000': 1.0, 'object_store_memory': 16320871833.0, 'memory': 50307916391.0, 'node:172.22.0.2': 1.0}

Ray reports the four logical CPUs it was given and the GPU it found. memory is the RAM that Ray may assign to tasks, and object_store_memory is the size of the object store, 10 GB here. The node: entries name the machine, and the accelerator type shows that this page was run on a machine with an RTX PRO 6000 GPU. In Colab without a GPU runtime, the GPU entries are missing.

Tasks

Decorating a function with @ray.remote turns it into a task. Calling f.remote(x) does not run f. It sends the call to the scheduler, which runs it in one of the worker processes, and returns at once with an ObjectRef, a handle to a result that may not exist yet. ray.get waits for the result and returns it. The function below sleeps for one second, standing in for a step that takes time, such as loading a file or evaluating a model on one batch. Eight calls in a loop take eight seconds. Eight remote calls run four at a time, so they finish in about two.

def slow_square(x):
    time.sleep(1.0)
    return x * x


@ray.remote
def slow_square_task(x):
    time.sleep(1.0)
    return x * x


start = time.perf_counter()
results = [slow_square(i) for i in range(8)]
print(f"plain Python: {results}  {time.perf_counter() - start:.1f} s")

start = time.perf_counter()
refs = [slow_square_task.remote(i) for i in range(8)]   # returns at once
print(f"submitted: {refs[0]}")
results = ray.get(refs)                                  # waits for all eight
print(f"Ray tasks:    {results}  {time.perf_counter() - start:.1f} s")

plain Python: [0, 1, 4, 9, 16, 25, 36, 49]  8.0 s
submitted: ObjectRef(c8ef45ccd0112571ffffffffffffffffffffffff0100000001000000)

Ray tasks:    [0, 1, 4, 9, 16, 25, 36, 49]  2.2 s

The eight calls took 8.0 s in a loop and 2.2 s as Ray tasks. Four tasks ran at a time, so the work took two rounds of one second, and the extra 0.2 s is the time to start the worker processes and move the results. ray.get on a list waits for all the results. When the tasks take different times, ray.wait returns as soon as some of them are done, so the driver can process results in the order they finish.

@ray.remote
def sleep_for(seconds):
    time.sleep(seconds)
    return seconds


pending = [sleep_for.remote(s) for s in (2.0, 0.5, 1.5, 1.0)]
start = time.perf_counter()
while pending:
    done, pending = ray.wait(pending, num_returns=1)
    print(f"{time.perf_counter() - start:.1f} s: task that slept {ray.get(done[0])} s finished")

0.5 s: task that slept 0.5 s finished

1.0 s: task that slept 1.0 s finished

1.5 s: task that slept 1.5 s finished

2.0 s: task that slept 2.0 s finished

Objects

Arguments and results of tasks live in Ray’s object store, shared memory that every worker on a machine can read. ray.put places a value there and returns a reference to it. Passing the reference to many tasks, instead of the value, stores the data once: each worker reads the same copy, and a NumPy array is not even copied into the worker’s memory. Passing a large array directly to every call would serialize it again for each one. References can also be passed from one task to another. A task that receives a reference as an argument waits until that object exists, so a chain of calls becomes a small dependency graph that Ray runs as soon as each input is ready.

data = np.random.default_rng(0).normal(size=(4_000, 1_000))   # 32 MB
data_ref = ray.put(data)


@ray.remote
def column_means(x, cols):
    return x[:, cols].mean(axis=0)


@ray.remote
def combine(*parts):
    return np.concatenate(parts)


# Four tasks read the same stored array; a fifth combines their results
parts = [column_means.remote(data_ref, slice(i, i + 250)) for i in range(0, 1_000, 250)]
means = ray.get(combine.remote(*parts))
print(f"means of {means.shape[0]} columns, largest difference from NumPy: "
      f"{np.abs(means - data.mean(axis=0)).max():.2e}")

means of 1000 columns, largest difference from NumPy: 0.00e+00

Actors

A task has no memory: each call starts from its arguments. An actor is a class decorated with @ray.remote. Creating it starts a dedicated worker process that holds the object, and method calls run in that process one at a time, in the order they arrive. Actors hold state that many tasks update, such as a model’s parameters, a counter of finished work, or a connection to a database. The actor below keeps a running mean. Eight tasks each compute a partial sum over a slice of the data and send it to the same actor.

@ray.remote
class RunningMean:
    def __init__(self):
        self.total, self.count = 0.0, 0

    def add(self, total, count):
        self.total += total
        self.count += count
        return self.count

    def mean(self):
        return self.total / self.count


@ray.remote
def partial_sum(x, rows, accumulator):
    block = x[rows]
    return ray.get(accumulator.add.remote(float(block.sum()), block.size))


acc = RunningMean.remote()
counts = ray.get([partial_sum.remote(data_ref, slice(i, i + 500), acc) for i in range(0, 4_000, 500)])
print(f"values seen after each update: {sorted(counts)}")
print(f"mean from the actor: {ray.get(acc.mean.remote()):.6f}   NumPy: {data.mean():.6f}")

values seen after each update: [500000, 1000000, 1500000, 2000000, 2500000, 3000000, 3500000, 4000000]
mean from the actor: -0.000154   NumPy: -0.000154

Resources and GPUs

Every task and actor asks for resources, one CPU by default. @ray.remote(num_gpus=1) asks for a GPU as well. Ray starts the task only on a machine with a free GPU and sets CUDA_VISIBLE_DEVICES so that the task sees only the GPU assigned to it. A fraction such as num_gpus=0.5 lets two tasks share one GPU. Ray does not limit the memory each one uses, so the fractions must match what the tasks actually need. A task that asks for a resource the cluster does not have waits forever, which is why the cell checks first.

@ray.remote(num_gpus=1)
def which_gpu():
    import torch
    return torch.cuda.get_device_name(0), os.environ.get("CUDA_VISIBLE_DEVICES")


if ray.cluster_resources().get("GPU", 0) >= 1:
    name, visible = ray.get(which_gpu.remote())
    print(f"task ran on {name}, CUDA_VISIBLE_DEVICES={visible}")
else:
    print("this Ray instance has no GPU; in Colab, choose a GPU runtime to run the task")

task ran on NVIDIA RTX PRO 6000 Blackwell Max-Q Workstation Edition, CUDA_VISIBLE_DEVICES=0

Inside the task, the assigned GPU is device 0 and the only one visible. Ray set CUDA_VISIBLE_DEVICES to it, so code in the task can use cuda:0 without knowing which physical GPU it was given.

Training with Ray Train

Ray Train runs a PyTorch training loop on several workers. You write the loop as a function, and TorchTrainer starts the workers, sets up distributed data parallel training between them, and collects the metrics they report. prepare_model wraps the model for data parallelism and moves it to the worker’s device, and ray.train.report reports metrics to Ray Train. To print the loss of every epoch here, the first worker also sends it to an actor, the same pattern as RunningMean above. ScalingConfig sets the number of workers and whether each one gets a GPU. The same function runs with one worker in Colab and with eight on a cluster. The example fits a linear model

y = \mathbf{w}^\top \mathbf{x} + b

to synthetic data with two CPU workers. Each worker generates the same data set and trains on its own half of the rows, and the gradients are averaged across the workers at every step.

import torch
from torch import nn

import ray.train
from ray.train import ScalingConfig
from ray.train.torch import TorchTrainer, get_device, prepare_model


@ray.remote
class LossLog:
    def __init__(self):
        self.rows = []

    def log(self, epoch, loss):
        self.rows.append((epoch, loss))

    def get(self):
        return self.rows


def train_loop(config):
    rank = ray.train.get_context().get_world_rank()
    world = ray.train.get_context().get_world_size()

    g = torch.Generator().manual_seed(0)
    X = torch.randn(4_096, 8, generator=g)
    w_true = torch.arange(1.0, 9.0)
    y = X @ w_true + 0.5 + 0.1 * torch.randn(4_096, generator=g)
    X, y = X[rank::world], y[rank::world]                 # each worker trains on its own share
    X, y = X.to(get_device()), y.to(get_device())         # the worker's GPU, or the CPU

    model = prepare_model(nn.Linear(8, 1))                # data-parallel wrapper
    opt = torch.optim.SGD(model.parameters(), lr=config["lr"])
    for epoch in range(config["epochs"]):
        for i in range(0, len(X), 128):
            loss = nn.functional.mse_loss(model(X[i:i + 128]).squeeze(-1), y[i:i + 128])
            opt.zero_grad()
            loss.backward()
            opt.step()
        ray.train.report({"epoch": epoch + 1, "loss": loss.item()})
        if rank == 0:
            ray.get(config["log"].log.remote(epoch + 1, loss.item()))


loss_log = LossLog.remote()
trainer = TorchTrainer(
    train_loop,
    train_loop_config={"lr": 0.05, "epochs": 5, "log": loss_log},
    scaling_config=ScalingConfig(num_workers=2, use_gpu=False),
)
trainer.fit()
for epoch, loss in ray.get(loss_log.get.remote()):
    print(f"epoch {epoch}: loss on the last batch of worker 0 = {loss:.4f}")

epoch 1: loss on the last batch of worker 0 = 6.3473
epoch 2: loss on the last batch of worker 0 = 0.2082
epoch 3: loss on the last batch of worker 0 = 0.0149
epoch 4: loss on the last batch of worker 0 = 0.0100
epoch 5: loss on the last batch of worker 0 = 0.0101

The loss falls from 6.35 after the first epoch to about 0.01 after the third and then stays there. That is the variance of the noise added to

y

(

0.1^2 = 0.01

): the model has found

\mathbf{w}

and

b

, and what remains is noise that no model can predict. Each worker trained on 2,048 of the 4,096 rows, and because the gradients were averaged at every step, the two workers kept identical copies of the model throughout. To train on GPUs, set use_gpu=True. Ray then gives each worker one GPU. prepare_model moves the model to that GPU, and get_device returns the same device, so the loop already moves the data there. That is the only change.

ray.shutdown()

Your own Ray cluster

A Ray cluster is a head node and any number of worker nodes. The head runs the scheduler state (the global control store, GCS) and the dashboard. Every node, the head included, runs a raylet that starts worker processes and holds that node’s part of the object store. The machines must reach each other over the network, and all of them must have the same versions of Python and Ray. On the machine that will be the head, start Ray from a terminal:

ray start --head --port=6379 --dashboard-host=0.0.0.0

The command prints the address of the head, such as 192.168.1.20:6379. On each other machine, join the cluster:

ray start --address=192.168.1.20:6379

ray status on any node lists the nodes and their resources. The dashboard, at http://192.168.1.20:8265, shows the nodes, the running tasks and actors, and their logs. Port 6379 must be open between the machines, and port 8265 from wherever you open the dashboard. A program connects to the cluster in one of two ways. A script started on one of the cluster’s nodes calls ray.init(address="auto"), which finds the running cluster instead of starting a new one, and every .remote call is then scheduled across all the nodes. A program on any other machine, such as your laptop, submits a job through the dashboard address:

ray job submit --address http://192.168.1.20:8265 --working-dir . -- python train.py

Ray uploads the current directory, runs python train.py on the cluster, and streams its output back. Inside train.py, ray.init() connects to the cluster the job runs on. ray stop on a node shuts Ray down on that node.

A two-node cluster on one machine

ray.cluster_utils.Cluster starts a head and extra worker nodes as separate processes on one machine. Ray uses it for its own tests, and it is a convenient way to see multi-node scheduling without a second computer. The cell starts a head with two CPUs and a worker node with two more, and runs eight tasks that report the node they ran on. The SPREAD scheduling strategy asks Ray to distribute the tasks over the nodes instead of filling one node first.

from ray.cluster_utils import Cluster

cluster = Cluster(initialize_head=True, head_node_args={"num_cpus": 2, "include_dashboard": True})
cluster.add_node(num_cpus=2)
context = ray.init(address=cluster.address, log_to_driver=False)
print(f"nodes: {len(ray.nodes())}, resources: {ray.cluster_resources()}")


@ray.remote(scheduling_strategy="SPREAD")
def where():
    time.sleep(0.5)
    return ray.get_runtime_context().get_node_id()[:8]


node_ids = ray.get([where.remote() for _ in range(8)])
print("tasks per node:", {n: node_ids.count(n) for n in sorted(set(node_ids))})

nodes: 2, resources: {'memory': 131896811520.0, 'node:172.22.0.2': 2.0, 'CPU': 4.0, 'object_store_memory': 314572800.0, 'node:__internal_head__': 1.0}

tasks per node: {'616bdd92': 6, '7892b178': 2}

The cluster has two nodes with two CPUs each. Both run on the same machine, so they share its IP address but have different node IDs. SPREAD asks Ray to distribute the tasks but does not guarantee an even split: in this run six of the eight tasks ran on one node and two on the other. The same cluster accepts jobs. The cell below does what ray job submit does from a terminal, with the Python client: it submits a short script, waits for it to finish and prints its output.

from ray.job_submission import JobStatus, JobSubmissionClient

client = JobSubmissionClient(f"http://{context.dashboard_url}")  # the head's dashboard
job_id = client.submit_job(
    entrypoint="python -c \"import ray; ray.init(); print('job sees', len(ray.nodes()), 'nodes and', "
               "ray.cluster_resources()['CPU'], 'CPUs')\""
)
while client.get_job_status(job_id) not in {JobStatus.SUCCEEDED, JobStatus.FAILED, JobStatus.STOPPED}:
    time.sleep(1)
print(f"status: {client.get_job_status(job_id)}")
print(client.get_job_logs(job_id).strip().splitlines()[-1])

status: SUCCEEDED
job sees 2 nodes and 4.0 CPUs

ray.shutdown()
cluster.shutdown()

Key references

The Ray documentation covers tasks, actors and objects in more depth, and the Ray Train guide covers distributed PyTorch training.

Found an error on this page? Report an issue.

Prerequisites

Probability

Linear Algebra

Calculus

Complex Numbers

Programming

Parallel Python with Ray

Starting Ray

Tasks

Objects

Actors

Resources and GPUs

Training with Ray Train

Your own Ray cluster

A two-node cluster on one machine

Key references