mirror of
https://github.com/wassname/ray.git
synced 2026-08-12 12:20:11 +08:00
[docs] Add more guideline on using ray in slurm cluster (#12819)
Co-authored-by: Sumanth Ratna <sumanthratna@gmail.com> Co-authored-by: PENG Zhenghao <pengzh@ie.cuhk.edu.hk> Co-authored-by: Richard Liaw <rliaw@berkeley.edu>
This commit is contained in:
PENG Zhenghao
co-authored by
Sumanth Ratna
Richard Liaw
Sumanth Ratna
Richard Liaw
parent
d98235cc84
commit
e63da54931
@@ -0,0 +1,29 @@
|
||||
# trainer.py
|
||||
from collections import Counter
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
import ray
|
||||
|
||||
num_cpus = int(sys.argv[1])
|
||||
|
||||
ray.init(address=os.environ["ip_head"])
|
||||
|
||||
print("Nodes in the Ray cluster:")
|
||||
print(ray.nodes())
|
||||
|
||||
|
||||
@ray.remote
|
||||
def f():
|
||||
time.sleep(1)
|
||||
return ray.services.get_node_ip_address()
|
||||
|
||||
|
||||
# The following takes one second (assuming that
|
||||
# ray was able to access all of the allocated nodes).
|
||||
for i in range(60):
|
||||
start = time.time()
|
||||
ip_addresses = ray.get([f.remote() for _ in range(num_cpus)])
|
||||
print(Counter(ip_addresses))
|
||||
end = time.time()
|
||||
print(end - start)
|
||||
@@ -0,0 +1,8 @@
|
||||
:orphan:
|
||||
|
||||
.. _slurm-basic:
|
||||
|
||||
slurm-basic.sh
|
||||
~~~~~~~~~~~~~~
|
||||
|
||||
.. literalinclude:: /cluster/examples/slurm-basic.sh
|
||||
@@ -0,0 +1,65 @@
|
||||
#!/bin/bash
|
||||
# shellcheck disable=SC2206
|
||||
#SBATCH --job-name=test
|
||||
#SBATCH --cpus-per-task=5
|
||||
#SBATCH --mem-per-cpu=1GB
|
||||
#SBATCH --nodes=4
|
||||
#SBATCH --tasks-per-node=1
|
||||
#SBATCH --time=00:30:00
|
||||
|
||||
set -x
|
||||
|
||||
# __doc_head_address_start__
|
||||
|
||||
# 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
|
||||
# __doc_head_address_end__
|
||||
|
||||
# __doc_head_ray_start__
|
||||
port=6379
|
||||
ip_head=$head_node_ip:$port
|
||||
export ip_head
|
||||
echo "IP Head: $ip_head"
|
||||
|
||||
echo "Starting HEAD at $head_node"
|
||||
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 "${SLURM_GPUS_PER_TASK}" --block &
|
||||
# __doc_head_ray_end__
|
||||
|
||||
# __doc_worker_ray_start__
|
||||
# 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 "${SLURM_GPUS_PER_TASK}" --block &
|
||||
sleep 5
|
||||
done
|
||||
# __doc_worker_ray_end__
|
||||
|
||||
# __doc_script_start__
|
||||
# ray/doc/source/cluster/examples/simple-trainer.py
|
||||
python -u simple-trainer.py
|
||||
@@ -0,0 +1,103 @@
|
||||
# slurm-launch.py
|
||||
# Usage:
|
||||
# python slurm-launch.py --exp-name test \
|
||||
# --command "rllib train --run PPO --env CartPole-v0"
|
||||
|
||||
import argparse
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
template_file = Path(__file__) / "slurm-template.sh"
|
||||
JOB_NAME = "${JOB_NAME}"
|
||||
NUM_NODES = "${NUM_NODES}"
|
||||
NUM_GPUS_PER_NODE = "${NUM_GPUS_PER_NODE}"
|
||||
PARTITION_OPTION = "${PARTITION_OPTION}"
|
||||
COMMAND_PLACEHOLDER = "${COMMAND_PLACEHOLDER}"
|
||||
GIVEN_NODE = "${GIVEN_NODE}"
|
||||
LOAD_ENV = "${LOAD_ENV}"
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument(
|
||||
"--exp-name",
|
||||
type=str,
|
||||
required=True,
|
||||
help="The job name and path to logging file (exp_name.log).")
|
||||
parser.add_argument(
|
||||
"--num-nodes",
|
||||
"-n",
|
||||
type=int,
|
||||
default=1,
|
||||
help="Number of nodes to use.")
|
||||
parser.add_argument(
|
||||
"--node",
|
||||
"-w",
|
||||
type=str,
|
||||
help="The specified nodes to use. Same format as the "
|
||||
"return of 'sinfo'. Default: ''.")
|
||||
parser.add_argument(
|
||||
"--num-gpus",
|
||||
type=int,
|
||||
default=0,
|
||||
help="Number of GPUs to use in each node. (Default: 0)")
|
||||
parser.add_argument(
|
||||
"--partition",
|
||||
"-p",
|
||||
type=str,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--load-env",
|
||||
type=str,
|
||||
help="The script to load your environment ('module load cuda/10.1')")
|
||||
parser.add_argument(
|
||||
"--command",
|
||||
type=str,
|
||||
required=True,
|
||||
help="The command you wish to execute. For example: "
|
||||
" --command 'python test.py'. "
|
||||
"Note that the command must be a string.")
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.node:
|
||||
# assert args.num_nodes == 1
|
||||
node_info = "#SBATCH -w {}".format(args.node)
|
||||
else:
|
||||
node_info = ""
|
||||
|
||||
job_name = "{}_{}".format(args.exp_name,
|
||||
time.strftime("%m%d-%H%M", time.localtime()))
|
||||
|
||||
partition_option = "#SBATCH --partition={}".format(
|
||||
args.partition) if args.partition else ""
|
||||
|
||||
# ===== Modified the template script =====
|
||||
with open(template_file, "r") as f:
|
||||
text = f.read()
|
||||
text = text.replace(JOB_NAME, job_name)
|
||||
text = text.replace(NUM_NODES, str(args.num_nodes))
|
||||
text = text.replace(NUM_GPUS_PER_NODE, str(args.num_gpus))
|
||||
text = text.replace(PARTITION_OPTION, partition_option)
|
||||
text = text.replace(COMMAND_PLACEHOLDER, str(args.command))
|
||||
text = text.replace(LOAD_ENV, str(args.load_env))
|
||||
text = text.replace(GIVEN_NODE, node_info)
|
||||
text = text.replace(
|
||||
"# THIS FILE IS A TEMPLATE AND IT SHOULD NOT BE DEPLOYED TO "
|
||||
"PRODUCTION!",
|
||||
"# THIS FILE IS MODIFIED AUTOMATICALLY FROM TEMPLATE AND SHOULD BE "
|
||||
"RUNNABLE!")
|
||||
|
||||
# ===== Save the script =====
|
||||
script_file = "{}.sh".format(job_name)
|
||||
with open(script_file, "w") as f:
|
||||
f.write(text)
|
||||
|
||||
# ===== Submit the job =====
|
||||
print("Starting to submit job!")
|
||||
subprocess.Popen(["sbatch", script_file])
|
||||
print(
|
||||
"Job submitted! Script file is at: <{}>. Log file is at: <{}>".format(
|
||||
script_file, "{}.log".format(job_name)))
|
||||
sys.exit(0)
|
||||
@@ -0,0 +1,8 @@
|
||||
:orphan:
|
||||
|
||||
.. _slurm-launch:
|
||||
|
||||
slurm-launch.py
|
||||
~~~~~~~~~~~~~~~
|
||||
|
||||
.. literalinclude:: /cluster/examples/slurm-launch.py
|
||||
@@ -0,0 +1,9 @@
|
||||
:orphan:
|
||||
|
||||
.. _slurm-template:
|
||||
|
||||
slurm-template.sh
|
||||
~~~~~~~~~~~~~~~~~
|
||||
|
||||
.. literalinclude:: /cluster/examples/slurm-template.sh
|
||||
:language: bash
|
||||
@@ -0,0 +1,64 @@
|
||||
#!/bin/bash
|
||||
# shellcheck disable=SC2206
|
||||
# THIS FILE IS GENERATED BY AUTOMATION SCRIPT! PLEASE REFER TO ORIGINAL SCRIPT!
|
||||
# THIS FILE IS A TEMPLATE AND IT SHOULD NOT BE DEPLOYED TO PRODUCTION!
|
||||
${PARTITION_OPTION}
|
||||
#SBATCH --job-name=${JOB_NAME}
|
||||
#SBATCH --output=${JOB_NAME}.log
|
||||
${GIVEN_NODE}
|
||||
### This script works for any number of nodes, Ray will find and manage all resources
|
||||
#SBATCH --nodes=${NUM_NODES}
|
||||
#SBATCH --exclusive
|
||||
### Give all resources to a single Ray task, ray can manage the resources internally
|
||||
#SBATCH --ntasks-per-node=1
|
||||
#SBATCH --gpus-per-task=${NUM_GPUS_PER_NODE}
|
||||
|
||||
# Load modules or your own conda environment here
|
||||
# module load pytorch/v1.4.0-gpu
|
||||
# conda activate ${CONDA_ENV}
|
||||
${LOAD_ENV}
|
||||
|
||||
# ===== DO NOT CHANGE THINGS HERE UNLESS YOU KNOW WHAT YOU ARE DOING =====
|
||||
# This script is a modification to the implementation suggest by gregSchwartz18 here:
|
||||
# https://github.com/ray-project/ray/issues/826#issuecomment-522116599
|
||||
redis_password=$(uuidgen)
|
||||
export redis_password
|
||||
|
||||
nodes=$(scontrol show hostnames "$SLURM_JOB_NODELIST") # Getting the node names
|
||||
nodes_array=($nodes)
|
||||
|
||||
node_1=${nodes_array[0]}
|
||||
ip=$(srun --nodes=1 --ntasks=1 -w "$node_1" hostname --ip-address) # making redis-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 [[ "$ip" == *" "* ]]; then
|
||||
IFS=' ' read -ra ADDR <<< "$ip"
|
||||
if [[ ${#ADDR[0]} -gt 16 ]]; then
|
||||
ip=${ADDR[1]}
|
||||
else
|
||||
ip=${ADDR[0]}
|
||||
fi
|
||||
echo "IPV6 address detected. We split the IPV4 address as $ip"
|
||||
fi
|
||||
|
||||
port=6379
|
||||
ip_head=$ip:$port
|
||||
export ip_head
|
||||
echo "IP Head: $ip_head"
|
||||
|
||||
echo "STARTING HEAD at $node_1"
|
||||
srun --nodes=1 --ntasks=1 -w "$node_1" \
|
||||
ray start --head --node-ip-address="$ip" --port=$port --redis-password="$redis_password" --block &
|
||||
sleep 30
|
||||
|
||||
worker_num=$((SLURM_JOB_NUM_NODES - 1)) #number of nodes other than the head node
|
||||
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" --redis-password="$redis_password" --block &
|
||||
sleep 5
|
||||
done
|
||||
|
||||
# ===== Call your code below =====
|
||||
${COMMAND_PLACEHOLDER}
|
||||
Reference in New Issue
Block a user