This script assumes that ssh is allowed between your nodes in cluster

EMR BootUp Script for Distributed TF

#!/bin/bash
## Author Sahil Badyal <[email protected]>
## This script setups the cluster env for tf
clusterid=$(aws emr list-clusters  --active | jq '.Clusters' | jq '.[]' -c | grep  <CLUSTER_NAME> | jq '.Id' | tr -d '"')
num=0
tot="$1"

if [ "2" != "" ]; then
    numParameterServers="$2"
else
    numParameterServers=1
fi

if [ $num -lt $tot ]; then
    while [ $num -lt $tot ]
    do
        ips=$(aws emr list-instances --cluster-id $clusterid | jq '.Instances' | jq '.[]' -c | jq '.PrivateIpAddress' | tr -d '"')
        num=$(aws emr list-instances --cluster-id $clusterid | grep 'ivateI' | wc -l)
        echo $num
        sleep 15
    done
else
    exit
fi
chief=$(ifconfig | grep 'inet addr:172' | awk '{print $2}' | sed 's/addr://g')
read -a arr <<< $ips
i=0
j=0
for x in "${arr[@]}"
do
    if [ "$x" != "$chief" ]; then
        if [ $i -lt $((tot-numParameterServers-1)) ]; then
            host[$i]=$x
        else
            ps[$j]=$x
            j=$((j+1))
        fi
    i=$((i+1))
    fi
done
workerList="\""
workerListE="\\\""
echo $chief
for x in "${host[@]}"
do
    workerListE="$workerListE$x:2222\\\",\\\""
    workerList=$workerList$x":2222\",\""
done
workerList=${workerList:0:-2}
workerListE=${workerListE:0:-3}
echo $workerListE
echo $workerList
psList="\""
psListE="\\\""
for x in "${ps[@]}"
do
    psListE="$psListE$x:2222\\\",\\\""
    psList=$psList$x":2222\",\""
done
psList=${psList:0:-2}
psListE=${psListE:0:-3}
echo $psListE
echo $psList
l=0
echo "export TF_CONFIG='{ \"cluster\": { \"chief\": [\"$chief:2222\"], \"worker\": [$workerList], \"ps\": [$psList] }, \"task\": {\"type\": \"chief\", \"index\": 0} }'" >> ~/.bashrc && source ~/.bashrc
for x in "${ps[@]}"
do
    config="'{ \\\"cluster\\\": { \\\"chief\\\": [\\\"$chief:2222\\\"], \\\"worker\\\": [$workerListE], \\\"ps\\\": [$psListE] }, \\\"task\\\": {\\\"type\\\": \\\"ps\\\", \\\"index\\\": $l}}'"
    com="echo \"export TF_CONFIG=$config\" >> ~/.bashrc && source ~/.bashrc"
    ssh -o "StrictHostKeyChecking no" -t hadoop@$x $com
    l=$((l+1))
done
k=0
for x in "${host[@]}"
do
    config="'{ \\\"cluster\\\": { \\\"chief\\\": [\\\"$chief:2222\\\"], \\\"worker\\\": [$workerListE], \\\"ps\\\": [$psListE] }, \\\"task\\\": {\\\"type\\\": \\\"worker\\\", \\\"index\\\": $k}}'"
    com="echo \"export TF_CONFIG=$config\" >> ~/.bashrc && source ~/.bashrc"
    ssh -o "StrictHostKeyChecking no" -t  hadoop@$x $com
    k=$((k+1))
done