run_pyscreener_distributed_example.batch

#!/bin/bash
#SBATCH -J pyscreener_distributed
#SBATCH -o %x_%j.out
#SBATCH -e %x_%j.err

#SBATCH -N 1
#SBATCH --ntasks-per-node 4
#SBATCH -c 4
#SBATCH --mem-per-cpu 1000

source activate pyscreener

######################## DO NOT CHANGE THINGS HERE ############################
redis_password=$( uuidgen 2> /dev/null )
export redis_password

nodes=$( scontrol show hostnames $SLURM_JOB_NODELIST )
nodes_array=( $nodes )

node_1=${nodes_array[0]} 
ip=$( srun -N 1 -n 1 -w $node_1 hostname --ip-address )
port=$( python -c 'import socket; s=socket.socket(); s.bind(("", 0)); print(s.getsockname()[1]); s.close()' )
ip_head=$ip:$port

export ip_head
echo "IP Head: $ip_head"

echo "STARTING HEAD at $node_1"
srun -N 1 -n 1 -w $node_1 \
    ray start --head --node-ip-address=$ip --port=$port --redis-password=$redis_password \
    --num-cpus $SLURM_CPUS_ON_NODE --temp-dir $TMPDIR \
    --block > /dev/null 2>& 1 &
sleep 30

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 -N 1 -n 1 -w $node_i \
        ray start --address $ip_head --redis-password=$redis_password \
        --num-cpus $SLURM_CPUS_ON_NODE --temp-dir $TMPDIR \
        --block > /dev/null 2>& 1 &
    sleep 1
done
sleep 30
###############################################################################

pyscreener --config integration-tests/configs/test_vina.ini --ncpu $SLURM_CPUS_PER_TASK