Ray on Discovery

Last updated September 09, 2026

Table of Contents

1 Build the environment
2 The job script
3 Example 1: Multiple CPUs on a single node
4 Example 2: Multi-node CPU jobs
5 Example 3: Multi-node, two GPUs per node
6 Example 4: vLLM on 2 nodes with 4 GPUs
7 Troubleshooting
8 Additional resources

Ray is an open-source framework for running Python code across multiple CPU cores, GPUs, or compute nodes. A function or a class is marked with a decorator (@ray.remote), and Ray then chooses which machine it runs on, moves the inputs there, and brings back the reference to the results. A “decorated function” is called a task and runs once. An instance of a “decorated class” is called an actor and stays alive between calls, so for example, a Machine Learning model is loaded into GPU memory once instead of on every call. The project documentation is available at docs.ray.io, and the Ray Core walkthrough is the quickest way to get started. Ray also ships higher level libraries built on that core, in particular Ray Train for distributed training and Ray Tune for hyperparameter search.

On an HPC cluster, Slurm and Ray have different roles. Slurm reserves the nodes, CPUs, GPUs, memory, and time for your job. Ray schedules work within those reserved resources. Basically, Ray is similar to MPI. Both run processes on one or more machines, and those processes can pass data directly to each other. However, they differ in two places. An MPI job starts a set number of copies of the same executable, while a Ray job starts one driver that creates workers as it runs. You then decide which rank does what in MPI, and Ray makes that decision for you, placing each piece of work on whatever core or GPU is free. Therefore, Ray may be more suited for work that is uneven or decided at run time (a parameter sweep, a queue of prompts), while MPI is the standard choice for tightly coupled numerical simulations or similar program whose processes communicate frequently and follow a regular pattern.

In a multi-node job, the script starts a Ray head process on the first node of the allocation, and a Ray worker process on each of the other nodes. The workers reach the head through an IP address and a port that the script picks. Once they have connected, the Python script running on the head node can place work on any core or GPU in the allocation, and Ray sends each piece of work to whichever resource is free.

In this guide, we wil give you 4 different Ray examples.

1 Build the environment

module purge
module load conda
conda create -n ray python=3.12 -y
conda activate ray
pip install "ray[default]==2.58.0"

Example 3 also needs pip install torch, and example 4 needs pip install vllm. A home directory holds 100 GB, which is enough for all of them. An environment that the rest of your group will use is built under /project2/<project_id> instead.

2 The job script

All four examples use this one script. It is saved as ray_job.slurm and submitted with sbatch ray_job.slurm. The #SBATCH lines that change from one example to the next are grouped together at the top, and the last line names the Python file to run.

#!/bin/bash
#SBATCH --account=<project_id>
#SBATCH --ntasks-per-node=1

# these lines change from one example to the next
#SBATCH --partition=main
#SBATCH --nodes=2
#SBATCH --cpus-per-task=16
#SBATCH --mem=32G
#SBATCH --time=00:30:00

module purge
eval "$(conda shell.bash hook)"
conda activate ray

nodes=($(scontrol show hostnames "$SLURM_JOB_NODELIST"))
head_ip=$(srun -N1 -n1 -w "${nodes[0]}" hostname --ip-address)
port=$((20000 + SLURM_JOB_ID % 10000))          # unique per job

# head process, on the first node
srun -N1 -n1 -w "${nodes[0]}" \
    ray start --head --node-ip-address="$head_ip" --port="$port" \
    --num-cpus="$SLURM_CPUS_PER_TASK" --block &
sleep 20

# one worker process on every other node
for node in "${nodes[@]:1}"; do
    srun -N1 -n1 -w "$node" \
        ray start --address="$head_ip:$port" \
        --num-cpus="$SLURM_CPUS_PER_TASK" --block &
done
sleep 20

export RAY_ADDRESS="$head_ip:$port"
ray status
python -u example1.py

ray status prints the cores and GPUs that Ray ended up with. It is the first thing to read when a job misbehaves.

With --nodes=1 the loop over the worker nodes has nothing to iterate over, so the same script also covers the single node examples.

Ray 2.49 and above provides ray symmetric-run, which replaces the head and worker block with a single srun. It cannot pin Ray’s internal port range, so it is only safe on nodes requested with --exclusive.

3 Example 1: Multiple CPUs on a single node

...
#SBATCH --partition=main
#SBATCH --nodes=1
#SBATCH --cpus-per-task=16
#SBATCH --mem=64G
#SBATCH --time=00:30:00
...
python -u example1.py

example1.py

import ray, time

@ray.remote
def run(id):
    time.sleep(5)                      # your real work goes here
    return id

ray.init()

start_time=time.perf_counter()

#64 processes are supposed to wait 5 seconds each
results = ray.get([run.remote(i) for i in range(64)])

print(f"{len(results)} runs finished")
print(f"elapsed time is = {time.perf_counter()-start_time}")

ray.init() reads RAY_ADDRESS from the environment and attaches to the cluster that the job script has already started. A call to run.remote(s) returns immediately with a reference to a result that does not exist yet, so all 64 calls are submitted in a single pass, and ray.get then waits for the 64 results. Each task asks for 1 CPU by default and this job holds 16 of them, so Ray runs 16 tasks at a time and the loop takes about 20 seconds instead of 320.

4 Example 2: Multi-node CPU jobs

...
#SBATCH --partition=main
#SBATCH --nodes=2
#SBATCH --cpus-per-task=16
#SBATCH --mem=64G
#SBATCH --time=00:30:00
...
python -u example2.py

example2.py

import ray, socket, time

@ray.remote
def run(id):
    time.sleep(5)
    return socket.gethostname()

ray.init()

start_time=time.perf_counter()

hosts = ray.get([run.remote(i) for i in range(64)])

print(f"64 runs on {len(set(hosts))} nodes: {set(hosts)}")
print(f"elapsed time is = {time.perf_counter()-start_time}")

The Python code is the same as in example 1 except for the hostname, which is returned here to show that the tasks did land on both machines. Two nodes of 16 cores give Ray 32 slots, so the same 64 tasks now take a little more than 10 seconds.

5 Example 3: Multi-node, two GPUs per node

...
#SBATCH --partition=gpu
#SBATCH --nodes=2
#SBATCH --cpus-per-task=8
#SBATCH --gpus-per-task=a40:2
#SBATCH --mem=64G
#SBATCH --time=00:30:00
...
python -u example3.py

example3.py

import socket
import ray
import torch
import os

@ray.remote(num_gpus=1)
def measure(n):
    device = torch.cuda.current_device()
    a = torch.randn(n, n, device=f"cuda:{device}")
    return {
        "n": n,
        "mean": (a @ a).mean().item(),
        "hostname": socket.gethostname(),
        "CUDA_VISIBLE_DEVICES": os.environ.get("CUDA_VISIBLE_DEVICES"),
        "gpu_name": torch.cuda.get_device_name(device),
    }

ray.init()

results =  ray.get([measure.remote(n) for n in [4000, 8000, 12000, 16000]])

for result in results: print(result)

Note: To run this example, you need to have installed PyTorch with CUDA in your Conda environment:

conda activate ray
python -m pip install "torch==2.13.0+cu129"  "torchvision==0.28.0+cu129" "torchaudio==2.11.0+cu129" "torchcodec==0.15.0+cu129" --index-url https://download.pytorch.org/whl/cu129

Two nodes with 2 A40 cards each give 4 GPUs, so the 4 calls all run at the same time. num_gpus=1 reserves one card for the task. Ray sets CUDA_VISIBLE_DEVICES before the function body starts, so cuda inside the function always refers to the card that Ray assigned, and a hard coded device index would be wrong. Two small tasks can share one card by asking for num_gpus=0.5.

The ray start lines need no GPU flag of their own, since Ray reads whichever devices Slurm exposed to each of them.

6 Example 4: vLLM on 2 nodes with 4 GPUs

Another example involves using Ray for distributed model inference.

Gemma 4 31B-it is a dense 31B model and ships as a 62.5 GB bf16 checkpoint. Two A40 cards hold 96 GB in total, and vLLM reserves 90 % of that by default, i.e. about 86 GB. This leaves roughly 24 GB for the KV cache and activations, which is why the model fits on a single node.

In this example, however, we would like to serve this model on 2 nodes, with 2 GPUs each, which allows for more room for KV cache and activations allowing more concurrent requests and longer context windows.

In this example, it is recommended to install HuggingFace and vLLM beforehand:

conda activate ray
pip install huggingface_hub
pip install "https://github.com/vllm-project/vllm/releases/download/v0.28.0/vllm-0.28.0+cu129-cp38-abi3-manylinux_2_28_x86_64.whl"

The weights are 62.5 GB, so they are downloaded once from a transfer node before the job is submitted.

export HF_HOME=/scratch1/$USER/huggingface
hf download google/gemma-4-31B-it

Edit the SLURM script:

...
#SBATCH --partition=gpu
#SBATCH --nodes=2
#SBATCH --cpus-per-task=16
#SBATCH --gpus-per-task=a40:2
#SBATCH --mem=192G
#SBATCH --time=02:00:00
...

export HF_HOME="/scratch1/$USER/huggingface"
export VLLM_WORKER_MULTIPROC_METHOD=spawn
export VLLM_USE_FLASHINFER_SAMPLER=0
export VLLM_PORT=$((30000 + SLURM_JOB_ID % 10000))

python -u example4.py

example4.py

import os
import ray
from vllm import LLM, SamplingParams

if __name__ == "__main__":
    ray.init(address="auto")

    llm = LLM(
        model="google/gemma-4-31B-it",
        tensor_parallel_size=2,
        pipeline_parallel_size=2,
        distributed_executor_backend="ray",
        max_model_len=128000,
    )

    outputs = llm.chat(
        messages=[
            {
                "role": "user",
                "content": (
                    "Explain gradient descent in plain English"
                ),
            }
        ],
        sampling_params=SamplingParams(
            temperature=0.0,
            max_tokens=128000,
        ),
        chat_template_kwargs={"enable_thinking": False},
    )

    print(outputs[0].outputs[0].text)

In this example, pipeline_parallel_size=2 which splits the 60 layers of the model into 2 groups of 30, one group per node, while tensor_parallel_size=2 splits each layer across the 2 cards inside a node. This is vLLM’s standard shape for a multi-node job, since matching the tensor parallel size to the GPUs in one node keeps the heavy per-layer traffic off the network. distributed_executor_backend="ray" then tells vLLM to place its workers on the Ray cluster rather than forking local processes, which is what allows them to sit on 2 different machines. Beyond ray.init() there is no Ray code here, since vLLM finds the running cluster and starts one worker per GPU on its own.

Please note that this is just a demonstration example. Here, using 2 nodes only allows for larger KV cache space and more concurrency. For this relatively small model, the multi-node inference pays off only at ~32 or more concurrent requests. For a single user, if your model fits in one node, it is generally recommended to use one node with TP=2 and PP=1. The gpu partition has a 200 Gbps NDR Infiniband network connection which may limit the inference speed.

7 Troubleshooting

If ray status shows fewer nodes than requested, first check the Slurm output for errors from the head and worker processes. Please note that ray start --block is supposed to keep running, and the 20-second waits in the script may not be enough for all workers to connect. Another possible cause is a port conflict when 2 Ray jobs share a node. The script above chooses the head port from SLURM_JOB_ID, which reduces the chance of a conflict but does not cover Ray’s other ports. For jobs that share nodes, those ports also need separate ranges, as described in the Ray Slurm documentation.

If Ray reports more CPUs than you requested, check that both ray start commands include --num-cpus="$SLURM_CPUS_PER_TASK". Without this setting, Ray may detect more CPUs than Slurm assigned to your job. Here, the setting tells Ray how many CPUs it can use on each node. ray status reports the total across the cluster, so in example 2 you should see 32 CPUs, or 16 per node.

With older versions of vLLM, Gemma 4 may fail to load because max_num_batched_tokens is smaller than the required multimodal token count. This can happen on GPUs with less than 70 GB of memory, where the default batch token budget is too small for the model’s video input budget. The issue was fixed in vLLM PR #43051. If you encounter it with an older version, try adding max_num_batched_tokens=4096 to LLM(...) in example 4.

If your job slows down when passing large results between tasks, check for object spilling with ray memory --stats-only. Ray normally stores these objects in /dev/shm and spills them to a filesystem when the object store fills up. On Discovery, /dev/shm uses RAM from your job’s memory allocation, with a filesystem limit of half the node’s memory. This does not mean that half of your requested memory is available to Ray’s object store. You need to request enough memory for both the Python processes and the objects they share. Also, Discovery’s /tmp is RAM-backed, so spilling there can still use your job’s memory. For large jobs that need to spill to disk, use a directory under /scratch1/$USER; see Ray’s object spilling instructions and CARC’s storage guide.

8 Additional resources

If you have questions about or need help with Ray, please submit a help ticket and we will assist you.