tf_multiworker.sh

#!/bin/bash
# To run: bash tf_multiworker.sh <num_workers>
# Where <num_workers> is an integer factor of the number of physical cores on the machine
# OpenMPI must be installed for this to run

if [ "$1" == "-h" ]; then
  echo "Usage: `basename $0` [<num_workers>]"
  echo "   where "
  echo "   [<num_workers>] = Number of TensorFlow workers per node"
  echo " "
  exit 0
fi

# Environment variables
export KMP_AFFINITY=granularity=fine,noverbose,compact,1,0
export KMP_BLOCKTIME=0

# Grep hardware parameters
export num_sockets=`lscpu | grep "Socket(s)" | cut -d':' -f2 | sed "s/ //g"`   # Number of sockets per node
export physical_cores=`lscpu | grep "Core(s) per socket" | cut -d':' -f2 | sed "s/ //g"` # Total number of physical cores per socket

# Set worker values
export num_workers=$1  # Total number of workers across machine
export cpu_per_proc=$(( $physical_cores * $num_sockets / $num_workers ))  # Calculate # of cores/worker

# Run tf_cnn_benchmarks.py, assigning cpu_per_proc cores per process
echo "Running $num_workers workers, assigning $cpu_per_proc cores per worker..."

if [ "$1" == "1" ]; then
  bash intel_tf_cnn_benchmarks.sh $physical_cores
fi

mpirun -np $num_workers -H localhost --map-by socket -cpus-per-proc $cpu_per_proc --oversubscribe bash intel_tf_cnn_benchmarks.sh $cpu_per_proc