Using Ray on Savio

Ray enables parallel and distributed execution of Python functions and applications across multiple nodes, making it useful for machine learning, data processing and other such workloads.

Ray is provided as a module on Savio.

Loading Ray on Savio

To access Ray, load the Ray module:

module load ml/ray/2.54.1

This module includes Ray Core, Ray Train, Ray Tune, Ray Serve and Ray RLlib components, along with a Python environment that includes PyTorch 2.10 and torchvision 0.25.

Ray Cluster Open OnDemand App

We have a Ray Cluster OOD application that you can use to lauch a Ray cluster (one or more nodes) through Open OnDemand. It is available under Interactive Apps > Servers > Ray Cluster + Jupyter. This OOD application requests nodes exclusively which means that the nodes allocated to this application are not shared with other jobs/users.

Once the application is launched, you can open a Jupyter Lab session to run your Ray python scripts either through the notebook interface or through the terminal on the jupyter lab session.

You can also open the Ray Dashboard to monitor resource usage of the Ray cluster.

Example: Running a Ray Job with SLURM

The following sample SLURM script (submit-ray-savio.sh) launches a Ray cluster across two nodes. It initializes the head node with its IP and port, establishes worker nodes that connect to the head node, and sets the RAY_ADDRESS environment variable. Replace <account_name> with your allocation and <partition_name> with the partition you want to use (for example, savio3). Since Ray is designed to manage all resources on a node, use an exclusive partition when possible. Otherwise, request the full node in your SLURM script with --exclusive and --mem=0.

#!/bin/bash
#SBATCH --job-name=ray-pi
#SBATCH --account=<account_name>
#SBATCH --partition=<partition_name>
#SBATCH --nodes=2
#SBATCH --ntasks-per-node=1
#SBATCH --time=00:30:00
#SBATCH --qos=savio_normal
#SBATCH --output=ray-pi-%j.out
#SBATCH --error=ray-pi-%j.err

module load ml/ray/2.54.1

# Ray head node initialization
head_node=$(hostname)
head_node_ip=$(hostname --ip-address)

port=6379
echo "Starting Ray head node on $head_node with IP $head_node_ip"
srun -n 1 --nodes=1 -w ${head_node} \
    ray start --head \
              --port=${port} \
              --node-ip-address=${head_node_ip} \
              --block &

sleep 5

echo "Give Ray time to initialize"
# increase the sleep time if needed to ensure ray is properly initialized
sleep 15

export RAY_ADDRESS=${head_node_ip}:${port}

# Ray worker node(s) initialization
n_workers=$((SLURM_JOB_NUM_NODES - 1))

if [ "$n_workers" -gt 0 ]; then
    echo "Launching $n_workers worker nodes..."
    sleep 10
    srun -n $n_workers --nodes=$n_workers \
         --ntasks-per-node=1 \
         --exclude=$head_node \
         ray start --address=${head_node_ip}:${port} --block &
fi

# Run your ray python code here
python compute_pi.py

exit

The following Python script (compute_pi.py) demonstrates Ray functionality by calculating pi using distributed sampling across 100 billion experiments, using remote task decorators and ray.get() to collect results.

import ray
import random
import time
import math
from fractions import Fraction

ray.init(address='auto')

@ray.remote
def pi4_sample(sample_count):
    """pi4 sample runs sample_count experiments, and returns the
    fraction of time it was inside the circle.
    """
    in_count = 0
    for i in range(sample_count):
        x = random.random()
        y = random.random()
        if x*x + y*y <= 1:
            in_count += 1
    return Fraction(in_count, sample_count)

SAMPLE_COUNT = 1000 * 1000

FULL_SAMPLE_COUNT = 100 * 1000 * 1000 * 1000 # 100 billion samples!
BATCHES = int(FULL_SAMPLE_COUNT / SAMPLE_COUNT)
print(f'Doing {BATCHES} batches')
start = time.time()
results = []
for _ in range(BATCHES):
    results.append(pi4_sample.remote(sample_count = SAMPLE_COUNT))
output = ray.get(results)
end = time.time()
dur = end - start
print(f'Running {FULL_SAMPLE_COUNT} tests took {dur} seconds')

pi = sum(output)*4/len(output)
print(float(pi))
print (abs(pi-math.pi)/pi)

ray.shutdown()

Adding Worker Nodes to a Running Ray Cluster

Step 1. Retrieve the cluster details from a terminal on the node running the Ray head:

echo $RAY_ADDRESS
more $RAY_AUTH_TOKEN_PATH

You will need these values in the Slurm script to add a new worker node.

Step 2. Submit a worker node job using the add-ray-worker.sh script below, setting RAY_ADDRESS to the head node IP and port, RAY_AUTH_MODE to token, and RAY_AUTH_TOKEN to the value from Step 1.

#!/bin/bash

#SBATCH --job-name=ray-worker
#SBATCH --account=<account_name>
#SBATCH --partition=<partition_name>
#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --time=00:30:00
#SBATCH --qos=savio_normal
#SBATCH --output=ray-worker-%j.out
#SBATCH --error=ray-worker-%j.err

module load ml/ray/2.54.1

export RAY_ADDRESS=<head-node-ip>:<port>      # from echo $RAY_ADDRESS
export RAY_AUTH_MODE=token
export RAY_AUTH_TOKEN="<your-token>"          # from more $RAY_AUTH_TOKEN_PATH

srun -n 1 --nodes=1 \
          --ntasks-per-node=1 \
          ray start --address=${RAY_ADDRESS} --block

Note

The new worker node job's wall time should not exceed the remaining wall time of the Ray cluster session. The worker will disconnect when either job ends.