Start ray processes in tmux

ray.tmux

#!/usr/bin/env bash

set -x

# NOTES
# - Requires tmux 2.3 or greater
# - Run in the appropriate Python environment
# - Check out the keybindings for helpful tips
# - We could also do this in tmuxinator

# Arguments
# $1: head-node-ip-address
# $2: redis-port

HEAD_NODE_IP_ADDRESS=$1
REDIS_PORT=$2

if [ -z "$HEAD_NODE_IP_ADDRESS" ]; then
  HEAD_NODE_IP_ADDRESS="127.0.0.1"
fi

if [ -z "$REDIS_PORT" ]; then
  REDIS_PORT=6379
fi

OBJECT_STORE_NAME=/tmp/plasma_store
RAYLET_NAME=/tmp/raylet
RESOURCES="CPU,1"

# KEYBINDINGS
# toggle pane title visibility
tmux bind t run 'zsh -c "arr=( off top ) && tmux setw pane-border-status \${arr[\$(( \${arr[(I)#{pane-border-status}]} % 2 + 1 ))]}"'

# rename pane
tmux bind T command-prompt -p "(rename-pane)" -I "#T" "select-pane -T '%%'"

# Start new tmux session with redis
tmux new-session -d -s ray-debug "ipython"
tmux select-pane -t 0 -T "ipython"

tmux split-window -h "python -c \"import ray; ray.services.start_redis('$HEAD_NODE_IP_ADDRESS', port=$REDIS_PORT, num_redis_shards=1, redirect_output=False, cleanup=False)\"; read"
tmux select-pane -t 1 -T "redis"

# Start monitor
# NOTE: Assume redis-address is same as head-node-ip-address
tmux split-window "python -u ./python/ray/monitor.py --redis-address=$HEAD_NODE_IP_ADDRESS:$REDIS_PORT; read"
tmux select-pane -t 2 -T "monitor"

# Start raylet monitor
tmux split-window -h "./python/ray/core/src/ray/raylet/raylet_monitor $HEAD_NODE_IP_ADDRESS $REDIS_PORT; read"
tmux select-pane -t 3 -T "raylet-monitor"

# Start log monitor
tmux split-window "python -u ./python/ray/log_monitor.py --redis-address=$HEAD_NODE_IP_ADDRESS:$REDIS_PORT --node-ip-address=$HEAD_NODE_IP_ADDRESS; read"
tmux select-pane -t 4 -T "log-monitor"

# Start object store
tmux split-window -h "./python/ray/core/src/plasma/plasma_store -s $OBJECT_STORE_NAME -m 1000000000; read"
tmux select-pane -t 5 -T "objstore"

# Start raylet
START_WORKER_COMMAND="python ./python/ray/workers/default_worker.py --node-ip-address=$HEAD_NODE_IP_ADDRESS --object-store-name=$OBJECT_STORE_NAME --raylet-name=$RAYLET_NAME --redis-address=$HEAD_NODE_IP_ADDRESS:$REDIS_PORT"
echo $START_WORKER_COMMAND
START_RAYLET="./python/ray/core/src/ray/raylet/raylet $RAYLET_NAME $OBJECT_STORE_NAME $HEAD_NODE_IP_ADDRESS $HEAD_NODE_IP_ADDRESS $REDIS_PORT 1 \"$START_WORKER_COMMAND\" \"$RESOURCES\""
echo $START_RAYLET
tmux split-window "$START_RAYLET; read"
tmux select-pane -t 6 -T "raylet"

# Start gateway
tmux split-window -h  "python ./python/ray/gateway.py -s $OBJECT_STORE_NAME -m $RAYLET_NAME -p 5002"
tmux select-pane -t 7 -T "gateway"

# Set layout (so tmux doesn't yell at us)
tmux select-layout tiled

# Start socat
tmux split-window "socat TCP-LISTEN:5001,reuseaddr,fork UNIX-CONNECT:$RAYLET_NAME; read"
tmux select-pane -t 8 -T "socat"

# Set layout (again)
tmux select-layout tiled

# Attach to tmux session
tmux -2 attach-session -d

echo "Stopping all ray processes"
ray stop

echo "Cleaning up all files"
rm -rf /tmp/ray* /tmp/plasma* /tmp/raylet*