> ## Documentation Index
> Fetch the complete documentation index at: https://aegean.ai/llms.txt
> Use this file to discover all available pages before exploring further.

# Parallel Python with Ray

> Tasks, objects, actors and Ray Train, on Colab and on your own cluster.

[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/pantelis/eng-ai-agents/blob/main/notebooks/python/ray/index.ipynb)

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](https://aegean.ai/aiml-common/lectures/vlm/finetuning/smolVLM/smolVLM.md) 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.

```python theme={null}
import sys

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

```python theme={null}
import os
import time

import numpy as np
import ray

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

```output theme={null}
[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.

```python theme={null}
ray.init(num_cpus=4, include_dashboard=True, log_to_driver=False, ignore_reinit_error=True)
print(ray.cluster_resources())
```

```output theme={null}
{'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.

```python theme={null}
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")
```

```output theme={null}
plain Python: [0, 1, 4, 9, 16, 25, 36, 49]  8.0 s
submitted: ObjectRef(c8ef45ccd0112571ffffffffffffffffffffffff0100000001000000)
```

```output theme={null}
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.

```python theme={null}
@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")
```

```output theme={null}
0.5 s: task that slept 0.5 s finished
```

```output theme={null}
1.0 s: task that slept 1.0 s finished
```

```output theme={null}
1.5 s: task that slept 1.5 s finished
```

```output theme={null}
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.

```python theme={null}
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}")
```

```output theme={null}
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.

```python theme={null}
@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}")
```

```output theme={null}
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.

```python theme={null}
@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")
```

```output theme={null}
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.

```python theme={null}
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}")
```

```output theme={null}
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.

```python theme={null}
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:

```bash theme={null}
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:

```bash theme={null}
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:

```bash theme={null}
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.

```python theme={null}
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))})
```

```output theme={null}
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}
```

```output theme={null}
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.

```python theme={null}
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])
```

```output theme={null}
status: SUCCEEDED
job sees 2 nodes and 4.0 CPUs
```

```python theme={null}
ray.shutdown()
cluster.shutdown()
```

## Key references

The [Ray documentation](https://docs.ray.io/en/latest/ray-core/walkthrough.html) covers tasks, actors and objects in more depth, and the [Ray Train guide](https://docs.ray.io/en/latest/train/getting-started-pytorch.html) covers distributed PyTorch training.

***

Found an error on this page? [Report an issue](https://github.com/aegean-ai/eaia-feedback/issues/new?template=page-issue.yml\&page=/aiml-common/lectures/python/ray).


This documentation is built and hosted on [Mintlify](https://mintlify.com), a developer documentation platform.