chore: import upstream snapshot with attribution

doc/source/cluster/doc_code/simple-trainer.py

@@ -0,0 +1,32 @@
+import os
+import socket
+import sys
+import time
+
+# trainer.py
+from collections import Counter
+
+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 socket.gethostbyname("localhost")
+
+
+# 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)

doc/source/cluster/doc_code/slurm-basic.sh

@@ -0,0 +1,45 @@
+#!/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]}
+
+port=6379
+ip_head=$head_node:$port
+export ip_head
+echo "IP Head: $ip_head"
+# __doc_head_address_end__
+
+# __doc_symmetric_run_start__
+# Start Ray cluster using symmetric_run.py on all nodes.
+# Symmetric run will automatically start Ray on all nodes and run the script ONLY the head node.
+# Use the '--' separator to separate Ray arguments and the entrypoint command.
+# The --min-nodes argument ensures all nodes join before running the script.
+
+# All nodes (including head and workers) will execute this block.
+# The entrypoint (simple-trainer.py) will only run on the head node.
+srun --nodes="$SLURM_JOB_NUM_NODES" --ntasks="$SLURM_JOB_NUM_NODES" \
+    ray symmetric-run \
+    --address "$ip_head" \
+    --min-nodes "$SLURM_JOB_NUM_NODES" \
+    --num-cpus="${SLURM_CPUS_PER_TASK}" \
+    --num-gpus="${SLURM_GPUS_PER_TASK}" \
+    -- \
+    python -u simple-trainer.py "$SLURM_CPUS_PER_TASK"
+# __doc_symmetric_run_end__
+
+# __doc_script_start__
+# The entrypoint script (simple-trainer.py) will be run on the head node by symmetric_run.

doc/source/cluster/doc_code/slurm-launch.py

@@ -0,0 +1,109 @@
+# 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')",
+        default="",
+    )
+    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)

doc/source/cluster/doc_code/slurm-template.sh

@@ -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}

doc/source/cluster/doc_code/xgboost_submit.py

@@ -0,0 +1,19 @@
+from ray.job_submission import JobSubmissionClient
+
+client = JobSubmissionClient("http://127.0.0.1:8265")
+
+kick_off_xgboost_benchmark = (
+    # Clone ray. If ray is already present, don't clone again.
+    "git clone https://github.com/ray-project/ray || true; "
+    # Run the benchmark.
+    "python ray/release/train_tests/xgboost_lightgbm/train_batch_inference_benchmark.py"
+    " xgboost --size=100G --disable-check"
+)
+
+
+submission_id = client.submit_job(
+    entrypoint=kick_off_xgboost_benchmark,
+)
+
+print("Use the following command to follow this Job's logs:")
+print(f"ray job logs '{submission_id}' --follow")

doc/source/cluster/doc_code/yarn/dashboard.py

@@ -0,0 +1,18 @@
+import skein
+import sys
+from urllib.parse import urlparse
+
+if __name__ == "__main__":
+    if len(sys.argv) < 2:
+        print("Usage: python dashboard.py <dashboard-address>")
+        sys.exit(1)
+    address = sys.argv[1]
+    # Check if the address is a valid URL
+    result = urlparse(address)
+    if not all([result.scheme, result.netloc]):
+        print("Error: Invalid dashboard address. Please provide a valid URL.")
+        sys.exit(1)
+
+    print("Registering dashboard " + address + " on skein.")
+    app = skein.ApplicationClient.from_current()
+    app.ui.add_page("ray-dashboard", address, "Ray Dashboard")

doc/source/cluster/doc_code/yarn/example.py

@@ -0,0 +1,50 @@
+import sys
+import time
+from collections import Counter
+
+import ray
+
+
+@ray.remote
+def get_host_name(x):
+    import platform
+    import time
+
+    time.sleep(0.01)
+    return x + (platform.node(),)
+
+
+def wait_for_nodes(expected):
+    # Wait for all nodes to join the cluster.
+    while True:
+        num_nodes = len(ray.nodes())
+        if num_nodes < expected:
+            print(
+                "{} nodes have joined so far, waiting for {} more.".format(
+                    num_nodes, expected - num_nodes
+                )
+            )
+            sys.stdout.flush()
+            time.sleep(1)
+        else:
+            break
+
+
+def main():
+    wait_for_nodes(4)
+
+    # Check that objects can be transferred from each node to each other node.
+    for i in range(10):
+        print("Iteration {}".format(i))
+        results = [get_host_name.remote(get_host_name.remote(())) for _ in range(100)]
+        print(Counter(ray.get(results)))
+        sys.stdout.flush()
+
+    print("Success!")
+    sys.stdout.flush()
+    time.sleep(20)
+
+
+if __name__ == "__main__":
+    ray.init(address="localhost:6379")
+    main()

doc/source/cluster/doc_code/yarn/ray-skein.yaml

@@ -0,0 +1,71 @@
+name: ray
+
+services:
+    # Head service.
+    ray-head:
+        # There should only be one instance of the head node per cluster.
+        instances: 1
+        resources:
+            # The resources for the head node.
+            vcores: 1
+            memory: 2048
+        files:
+            # ray/doc/source/cluster/doc_code/yarn/example.py
+            example.py: example.py
+            # ray/doc/source/cluster/doc_code/yarn/dashboard.py
+            dashboard.py: dashboard.py
+        #     # A packaged python environment using `conda-pack`. Note that Skein
+        #     # doesn't require any specific way of distributing files, but this
+        #     # is a good one for python projects. This is optional.
+        #     # See https://jcrist.github.io/skein/distributing-files.html
+        #     environment: environment.tar.gz
+        script: |
+            # Activate the packaged conda environment
+            #  - source environment/bin/activate
+
+            # This gets the IP address of the head node.
+            RAY_HEAD_ADDRESS=$(hostname -i)
+
+            # This stores the Ray head address in the Skein key-value store so that the workers can retrieve it later.
+            skein kv put current --key=RAY_HEAD_ADDRESS --value=$RAY_HEAD_ADDRESS
+
+            # This command starts all the processes needed on the ray head node.
+            # By default, we set object store memory and heap memory to roughly 200 MB. This is conservative
+            # and should be set according to application needs.
+            #
+            ray start --head --port=6379 --object-store-memory=200000000 --memory 200000000 --num-cpus=1 --dashboard-host=$RAY_HEAD_ADDRESS
+
+            # This registers the Ray dashboard on Skein, which can be accessed on Skein's web UI.
+            python dashboard.py "http://$RAY_HEAD_ADDRESS:8265"
+
+            # This executes the user script.
+            python example.py
+
+            # After the user script has executed, all started processes should also die.
+            ray stop
+            skein application shutdown current
+    # Worker service.
+    ray-worker:
+        # The number of instances to start initially. This can be scaled
+        # dynamically later.
+        instances: 4
+        resources:
+            # The resources for the worker node
+            vcores: 1
+            memory: 2048
+        # files:
+        #     environment: environment.tar.gz
+        depends:
+            # Don't start any worker nodes until the head node is started
+            - ray-head
+        script: |
+            # Activate the packaged conda environment
+            #  - source environment/bin/activate
+
+            # This command gets any addresses it needs (e.g. the head node) from
+            # the skein key-value store.
+            RAY_HEAD_ADDRESS=$(skein kv get --key=RAY_HEAD_ADDRESS current)
+
+            # The below command starts all the processes needed on a ray worker node, blocking until killed with sigterm.
+            # After sigterm, all started processes should also die (ray stop).
+            ray start --object-store-memory=200000000 --memory 200000000 --num-cpus=1 --address=$RAY_HEAD_ADDRESS:6379 --block; ray stop