SLURM deployment#

If you have access to a SLURM cluster, you can refer to the following guide to run an Xorbits job. Other HPC job schedulers like Torque or LSF are similar. You are recommended to read the cluster deployment first to know some basic knowledge of a Xorbits cluster.

Installation#

HPC clusters have shared storage among all compute nodes. You can use module, conda or mamba to create a Python environment and use pip to install Xorbits on the shared storage, see installation document for reference.

Walkthrough using Xorbits with SLURM#

On a SLURM cluster, you are required to interact with the compute resources via sbatch command and a SLURM script file declaring specific compute resources.

In the SLURM script file, you will want to start a Xorbits cluster with multiple srun commands (tasks), and then execute your python script that connects to the Xorbits cluster. You need to first start a supervisor and then start the workers.

The below walkthrough will do the following:

  1. Set the proper headers where you ask for resources from the SLURM cluster.

  2. Load the proper environment/modules.

  3. Fetch a list of available compute nodes and their IP addresses.

  4. Launch a supervisor process on one of the nodes (called the head node).

  5. Launch worker processes on other worker nodes with the head node’s address.

  6. After the underlying Xorbits cluster is ready, submit the user-specified task.

SLURM script file#

In the SLURM script, you’ll need to tell SLURM to allocate nodes for your Xorbits job. In this example, we ask for 4 nodes, and on each node, we’ve set --cpus-per-task=24 and --ntasks-per-node=1 which means we need 24 CPUs per node. Modify this setting according to your workload. Similarly, you can also specify the number of GPUs per node via --gpus=1. You need to change --partition to select the partition in your site. You can also add other optional flags to your sbatch directives.

#!/bin/bash
#SBATCH --job-name=xorbits
#SBATCH --nodes=4
#SBATCH --cpus-per-task=24
#SBATCH --ntasks-per-node=1
#SBATCH --partition=cpu24c
#SBATCH --time=00:30:00

Load your environment#

You’ll need to install Xorbits into a specific environment using conda or module. In the SLURM script, you should load modules or your own conda environment. And on the compute nodes allocated, the environment will be switched to the one where Xorbits is installed. In this case, we install Xorbits in a conda environment called df.

# Example: module load xorbits
# Example: source activate my-env

source activate df

Obtain the nodes#

Next, we’ll want to obtain the head node and the IP address.

# Getting the node names
nodes=$(scontrol show hostnames "$SLURM_JOB_NODELIST")
nodes_array=($nodes)

Now the nodes_array is the list of all the nodes allocated for this job.

Choose the first node of the nodes_array as the head node. The hostname --ip-address command will return the IP address of the specific node. Get the IP address of the head node:

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

Start the supervisor#

After detecting the head node hostname and IP address, we’ll want to run the supervisor on the head node. We’ll do this by using srun to start the supervisor on the head node. xorbits-supervisor is the command line tool to start the supervisor. You should specify the IP, port, the web port. Note that you should sleep a few seconds as the supervisor need some time to start. Otherwise, worker nodes may not be able to connect to the supervisor.

port=16380
web_port=16379

echo "Starting HEAD at $head_node"
srun --nodes=1 --ntasks=1 -w "$head_node" \
    xorbits-supervisor -H "$head_node_ip" -p "$port" -w "$web_port" &
sleep 10

Start Workers#

The rest of the machines can be started as workers via command:

# 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]}
    port_i=$((port + i))

    echo "Starting WORKER $i at $node_i"
    srun --nodes=1 --ntasks=1 -w "$node_i" \
        xorbits-worker -H "$node_i"  -p "$port_i" -s "$head_node_ip":"$port" &
done
sleep 5

Connect to The Cluster#

Now, the Xorbits cluster is created, and address is the endpoint to connect. You can connect to the supervisor and submit your Xorbits job.

address=http://"$head_node_ip":"$web_port"

python -u test.py --endpoint "$address"

The test.py is like the following:

import argparse

import xorbits
import xorbits.numpy as np

parser = argparse.ArgumentParser(description="test")
parser.add_argument(
    "--endpoint",
    type=str,
    default="0.0.0.0",
    required=True,
)

args = parser.parse_args()

xorbits.init(args.endpoint)
print(np.random.rand(100, 100).mean())

Name this SLURM script file as xorbits_slurm.sh. Submit the job via:

sbatch xorbits_slurm.sh

Put all together#

The SLURM script looks like this:

#!/bin/bash

#SBATCH --job-name=xorbits
#SBATCH --nodes=4
#SBATCH --cpus-per-task=24
#SBATCH --ntasks-per-node=1
#SBATCH --partition=cpu24c
#SBATCH --time=00:30:00

source activate df

### Use the debug mode to see if the shell commands are correct.
### If you do not want the shell command logs, delete the following line.
set -x

# 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=16380
web_port=16379

echo "Starting HEAD at $head_node"
srun --nodes=1 --ntasks=1 -w "$head_node" \
    xorbits-supervisor -H "$head_node_ip" -p "$port" -w "$web_port" &
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]}
    port_i=$((port + i))

    echo "Starting WORKER $i at $node_i"
    srun --nodes=1 --ntasks=1 -w "$node_i" \
        xorbits-worker -H "$node_i"  -p "$port_i" -s "$head_node_ip":"$port" &
done
sleep 5

address=http://"$head_node_ip":"$web_port"

python -u test.py --endpoint "$address"