Run Ray on SUNK - CoreWeave Docs

Using Ray on Slurm-on-Kubernetes (SUNK) provides Ray’s Python native parallelism with Slurm’s enterprise-level scheduler. By using Ray on SUNK, you let Slurm manage when and where jobs run, while Ray manages how distributed tasks are executed. This document covers integrating Ray’s distributed computing framework with SUNK’s Slurm-based job scheduler and covers container preparation, cluster management, and interactive debugging. Consider running Ray on SUNK if one of the following apply:

Most of your compute usage is through Slurm, and you need Ray to orchestrate workloads.
You want to support Ray and Slurm users from a single environment, using the Slurm scheduler for efficient resource allocation.

Prerequisites

Before completing the steps in the guide, be sure you have the following:

Access to a SUNK cluster with permissions to submit jobs.
Familiarity with Slurm commands, such srun, sbatch, and squeue.

Create a Ray cluster on Slurm

When working with containers on Slurm, you typically pull them as squash files to a shared directory in the Slurm login nodes. Within the shared directory, you can then execute the containers using multiple nodes without making simultaneous repository pulls. There are two methods for pulling containers:

Pulling a Ray container and saving it as a squash file using enroot. This is useful if you need to debug container errors.
Pulling a container as a sqsh file using srun. This is useful when your containers are well-tested and you only want to run jobs.

Pull a Ray container

To pull a Ray container, complete the following steps.

From the login node, create an interactive session on a GPU by running the following command: (Note that you can also create a VS Code tunnel to the GPU.)
```
srun -p NODE_TYPE --exclusive --pty bash -i
```
- NODE_TYPE: Change this option based on your available resources.
Here we are requesting an h100 node:
Example with h100 node
```
srun -p h100 --exclusive --pty bash -i
```
You should now be in a session on a GPU node. Verify you are in a shared directory, such as your home directory or /mnt/data.

Import and start the container. Note that this can be any container, but here we’ll use a rayproject nightly container:

# Import the container
enroot import docker://rayproject/ray:nightly-py39-gpu

# Create the container
enroot create rayproject+ray+nightly-py39-gpu.sqsh

# Start the container
enroot start --rw rayproject+ray+nightly-py39-gpu

You should see output similar to the following:

==========
== CUDA ==
==========

CUDA Version 12.1.1

Container image Copyright (c) 2016-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.

You now have an interactive session running on a single node within the container. Because you used the --rw flag to start the container, changes you make will be saved when you exit. If you are using a container from ghcr.io, prefix it with docker://, for example, docker://ghcr.io/. If your container requires permissions, you will need to configure your enroot credentials. For more information about enroot, see the SUNK Training guide.

Conda environment

Instead of using containers, you can create a conda environment that includes your dependencies. However, this is not recommended. Ray expects the versions of Ray and Python used in the head and worker nodes to be identical to those used in the Ray script being run. Conda environments with many users can be difficult to manage, but conda may be easier for testing and debugging when running Ray on Slurm. If using conda, you need to initialize your conda environment (only once).

To initialize a conda environment, determine its PATH and initialize it:
```
which conda
```
Use the conda PATH in the next command. For example, if the PATH is the following:
```
/home/ray/anaconda3/bin/conda
```
Enter the PATH below:
```
/home/ray/anaconda3/bin/conda init bash
source ~/.bashrc
```

Create an environment with the version of Ray that you want to use and any additional packages. Here the environment is called rayenv:

# Create a conda environment with a specific Python version
conda create --name rayenv python=3.12 pip

# Activate the ray environment
conda activate rayenv

# Install ray with selected packages
pip install -U "ray[data,train,tune,serve]"

You can now activate the same environment before you launch a Ray cluster or job.

Scripts for generating `sbatch` files

You can write a script to generate sbatch files for creating Ray clusters of the size and specifications you want. This lets you generate the right sbatch commands programmatically. For examples of these kinds of scripts, refer to the NERSC GitHub repository. The following example script starts up a Ray cluster and also runs a job. If you want to start a Ray cluster and use it interactively, make sure the sbatchscript does not exit by putting a sleep inf command in the script. This keeps the script running for a specified time as an interactive development cluster. You can exit it earlier if you don’t need it. Note that the instructions use the variable SLURM_GPUS_PER_TASK. This is not a default Slurm variable. You need to set it. In the following examples, it’s hardcoded to 8. Adjust the number of GPUs per task as you need. To see available capacity, use the sinfo command. The partitions of different skus will typically be separated by name, for example, h100, h200.

sinfo

The following is a sample sbatch script:

#!/bin/bash
#SBATCH --job-name=raytest
#SBATCH --nodes=2
#SBATCH --exclusive
#SBATCH --tasks-per-node=1 # we will launch one worker task per node
#SBATCH --cpus-per-task=8  # each worker task gets 8 CPUs. Adjust as needed.
#SBATCH --mem-per-cpu=1GB  # each cpu gets 1GB of memory. Adjust as needed.
#SBATCH --gpus-per-task=8  # each worker task will use 8 GPUs. Adjust as needed.
#SBATCH --time=1-00:00:00  # specify a time limit of one day

# Here we activate a conda environment named "rayenv" to load Ray
# and its dependencies. This assumes that you have already created a
# conda environment named "rayenv" with Ray installed.
eval "$(conda shell.bash hook)"
conda activate rayenv
# Getting the node names
nodes=$(scontrol show hostnames "$SLURM_JOB_NODELIST")
nodes_array=($nodes)

head_node=${nodes_array[0]}
head_node_ip=$(srun --nodes=1 --ntasks=1 -w "$head_node" hostname --ip-address)

# If we detect a space character in the head node IP, we'll
# convert it to an ipv4 address. This step is optional.
if [[ "$head_node_ip" == *" "* ]]; then
IFS=' ' read -ra ADDR <<<"$head_node_ip"
if [[ ${#ADDR[0]} -gt 16 ]]; then
 head_node_ip=${ADDR[1]}
else
 head_node_ip=${ADDR[0]}
fi
echo "IPV6 address detected. We split the IPV4 address as $head_node_ip"
fi

port=6379
ip_head=$head_node_ip:$port
export ip_head
echo "IP Head: $ip_head"

echo "Starting HEAD at $head_node"
echo srun --nodes=1 --ntasks=1 -w "$head_node" \
   ray start --head --node-ip-address="$head_node_ip" --port=$port \
   --num-cpus "${SLURM_CPUS_PER_TASK}" --num-gpus 1 --block

srun --nodes=1 --ntasks=1 -w "$head_node" \
   ray start --head --node-ip-address="$head_node_ip" --port=$port \
   --num-cpus "${SLURM_CPUS_PER_TASK}" --num-gpus 1 --block &

# Optional, though may be useful in certain versions of Ray < 1.0.
sleep 10

# Number of nodes other than the head node.
worker_num=$((SLURM_JOB_NUM_NODES - 1))

for ((i = 1; i <= worker_num; i++)); do
   node_i=${nodes_array[$i]}
   echo "Starting WORKER $i at $node_i"
   srun --nodes=1 --ntasks=1 -w "$node_i" \
       ray start --address "$ip_head" \
       --num-cpus "${SLURM_CPUS_PER_TASK}" --num-gpus 1 --block &
   sleep 5
done

# Update --container-image and the Python script you want to run.
# For an example, see the srun command after this example.
srun -J "cont" --overlap --container-image=/mnt/home/USER/RAY_CONTAINER_SQUASH_FILE --container-remap-root --container-mounts=/mnt/home:/mnt/home python3 -u /mnt/home/USER/PYTHON_SCRIPT "$SLURM_CPUS_PER_TASK"

Example srun command for the sbatch script above:

srun -J "cont" --overlap \
--container-image=/mnt/home/USER_NAME/rayproject+ray+nightly-py39-gpu.sqsh \
--container-remap-root \
--container-mounts=/mnt/home:/mnt/home \
python3 -u /mnt/home/USER_NAME/tune_basic_example.py "$SLURM_CPUS_PER_TASK"

USER_NAME: Replace with your username.
rayproject+ray+nightly-py39-gpu.sqsh: This is the Ray container squash file we downloaded in the Pull a Ray container section.
tune_basic_example.py: Create this file under your username with the contents found in the NERSC repository here.

Ray expects each person to be able to create their own Ray environments. Once you have created a Ray environment you can either interactively log in and launch jobs, or connect to a running container within that job to debug.

Pull a container as a `sqsh` file using `srun`

To pull a container as a sqsh file using srun, complete the following steps.

On a login node, execute a simple command—such as “echo hello”—on the development container. As an example, we use a nccl test to pull and save the container.

srun --container-image=ghcr.io#coreweave/nccl-tests:12.8.1-devel-ubuntu22.04-nccl2.26.2-1-0708d2e \
       --container-remap-root --no-container-mount-home \
       --container-save ${HOME}/nccl-test.sqsh echo hello

${HOME}/nccl-test.sqsh. Update ${HOME} and nccl-test.sqsh based on where you want to save the file and what name to give it.

You should see output similar to the following:

pyxis: imported docker image: ghcr.io#coreweave/nccl-tests:12.8.1-devel-ubuntu22.04-nccl2.26.2-1-0708d2e
hello
pyxis: exported container pyxis_228.0 to /mnt/home/username/nccl-test.sqsh

To pull the latest Ray container, run the following command:

srun --container-image=rayproject/ray:nightly-py39-gpu \
       --container-save ${HOME}/ray-nightly-py39-gpu echo hello

${HOME}/ray-nightly-py39-gpu. Update ${HOME} and ray-nightly-py39-gpu based on where you want to save the file and what name to give it.

You should see output similar to the following:

pyxis: importing docker image: rayproject/ray:nightly-py39-gpu
pyxis: imported docker image: rayproject/ray:nightly-py39-gpu
hello
pyxis: exported container pyxis_229.0 to /mnt/home/username/ray-nightly-py39-gpu

Interactively modify a container using `srun`

You can interactively modify a container using srun by mounting the home directory explicitly and saving the container. You can then install Ray inside the container:

srun --container-image=${HOME}/nccl-test.sqsh \
       --container-remap-root --container-mounts=/mnt/home:/mnt/home \
       --container-save ${HOME}/nccl-test-new.sqsh --pty bash -i

${HOME}/nccl-test.sqsh. Update ${HOME} and nccl-test.sqsh based on where you saved the container image from the previous steps.

When you exit the container, nccl-test-new.sqsh is saved.

Work with Ray on Slurm

Once your Ray environment is running, you can either connect interactively to the running Ray cluster for development or attach to a containerized Ray job to debug or inspect it. See instructions for each option below.

Interactively log in to a Ray environment

You can run the script by submitting the sbatch script. Note that the script includes the command to sleep an environment.

Use the example script above and save it as ray.devcluster.batch, and then run the following command:
```
sbatch ray.devcluster.batch
```

Your Ray environment will be a job. You can find out what job it is by running the following command:

squeue

You should see similar output like the following:

JOBID  PARTITION     NAME     USER       ST     TIME   NODES NODELIST(REASON)
1335   hpc-low       raytest  user_name  R      15:46  2     slurm-h100-231-[147,217]

Above we can see that the JOBID is 1335. This is a mini Ray environment with two h100 nodes. To connect to it and begin running things interactively, you would run the following command:
```
srun --jobid=1335 --overlap --pty bash -i
```
- jobid: Be sure to replace the jobid flag with the job you want to run.
- --overlap: The overlap flag specifies that you should attach to the head node of the specified job, here jobid=1335.

Attach to a Ray job using a container

If you launch a job using a container, you can connect to that container by providing the container name, as opposed to sleeping in the Ray environment for a job recently launched. To attach to a Ray job using a container, complete the following steps.

Run the squeue command to figure out what job it is:
```
squeue
```
Find the specific step that is running the container using sacct. sacct is a Slurm command that gives more information about the “steps” of the job. Steps are tasks launched within the allocation that use some or all of the resources.
```
sacct -j JOB_ID
```
- JOB_ID: Replace with your jobid.
This will start the container in the same allocation, so you can debug the run interactively. You want to be able to connect with the same container to a specific step.
```
srun --overlap --container-image=/mnt/data/coreweave/ray/nccl-test-new.sqsh --jobid 1336.3 --pty bash -i
```
- --container-image: Replace with the location of your container image.
- --jobid: Replace with your jobid.
In the example script, we gave it a name using the -J flag in the srun command. This makes it easier to find. Note this will take a minute because first it is launching all the Ray daemons. There will be multiple nodes in the cluster but you can only ssh into one. The node that it will pick by default will be the head node of the allocation, which will be where you can run ray status.
To connect to a specific node in your Ray environment, provide it on the command line with the --nodelist flag.
```
srun --overlap --nodelist=slurm-h100-227-211 \
     --container-image=/mnt/data/coreweave/ray/nccl-test-new.sqsh \
     --jobid 1336.3 --pty bash -i
```
- --nodelist: Replace with your node.
- --container-image: Replace with the location of your container image.
- --jobid: Replace with your jobid.

Stop a Ray environment

To delete a job, find the job id using squeue and cancel it using the following command:

scancel JOB_ID

JOB_ID: Replace with the job you want to delete.

SUNK

Documentation Index

​Prerequisites

​Create a Ray cluster on Slurm

​Pull a Ray container

​Conda environment

​Scripts for generating sbatch files

​Pull a container as a sqsh file using srun

​Interactively modify a container using srun

​Work with Ray on Slurm

​Interactively log in to a Ray environment

​Attach to a Ray job using a container

​Stop a Ray environment

Prerequisites

Create a Ray cluster on Slurm

Pull a Ray container

Conda environment

Scripts for generating `sbatch` files

Pull a container as a `sqsh` file using `srun`

Interactively modify a container using `srun`

Work with Ray on Slurm

Interactively log in to a Ray environment

Attach to a Ray job using a container

Stop a Ray environment