wmakeScheduler

#!/bin/bash
#------------------------------------------------------------------------------
# =========                 |
# \\      /  F ield         | OpenFOAM: The Open Source CFD Toolbox
#  \\    /   O peration     |
#   \\  /    A nd           | Copyright (C) 1991-2008 OpenCFD Ltd.
#    \\/     M anipulation  |
#-------------------------------------------------------------------------------
# License
#     This file is part of OpenFOAM.
#
#     OpenFOAM is free software; you can redistribute it and/or modify it
#     under the terms of the GNU General Public License as published by the
#     Free Software Foundation; either version 2 of the License, or (at your
#     option) any later version.
#
#     OpenFOAM is distributed in the hope that it will be useful, but WITHOUT
#     ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
#     FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License
#     for more details.
#
#     You should have received a copy of the GNU General Public License
#     along with OpenFOAM; if not, write to the Free Software Foundation,
#     Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA
#
# Script
#     wmakeScheduler
#
# Description
#     Scheduler for network distributed compilations using wmake.
#     - WM_HOSTS contains a list of hosts and number of concurrent processes
#     eg,
#        export WM_HOSTS="hostA:1 hostB:2 hostC:1"
#     - WM_COLOURS contains a list of colours to cycle through
#        export WM_COLOURS="black blue green cyan red magenta yellow"
#
#     Sources the relevant cshrc/bashrc if not set.
#
#     WM_PROJECT_DIR, WM_PROJECT and WM_PROJECT_VERSION will have been set
#     before calling this routine.
#     FOAM_INST_DIR may possibly have been set (to find installation)
#
# Usage
#     wmakeScheduler COMMAND
#         run 'COMMAND' on one of the slots listed in WM_HOSTS
#
#     wmakeScheduler -count
#         count the total number of slots available in WM_HOSTS
#         eg,  export WM_NCOMPPROCS=$(wmakeScheduler -count)
#
#-------------------------------------------------------------------------------
lockDir=$HOME/.wmakeScheduler

# fallback - 1 core on current host
: ${WM_HOSTS:=$HOST:1}

# count the total number of slots available and exit
if [ "$1" = "-count" ]
then
   expr $(
      for slotGroup in $WM_HOSTS
      do
         n=${slotGroup##*:}
         [ "$n" = "${slotGroup%%:*}" ] && n=1  # missing ':'
         echo "+ ${n:-1}"
      done
   )
   exit 0
fi

# where to source WM_PROJECT settings in a remote shell
# This code tries to figure out which cshrc or bashrc to execute.
# !! Assumes remote computer running same shell and startup files
# in same location

sourceFoam=false    # fallback command
case $SHELL in
*/csh | */tcsh )    # [t]csh vs bash|ksh|sh
   shellRc=cshrc
   ;;
*)
   shellRc=bashrc
   ;;
esac

# check ~/.$WM_PROJECT/$WM_PROJECT_VERSION/
# check ~/.$WM_PROJECT/
# check <installedProject>/etc/
if [ "$WM_PROJECT" ]
then
   for i in \
      $HOME/.$WM_PROJECT/$WM_PROJECT_VERSION \
      $HOME/.$WM_PROJECT \
      $WM_PROJECT_DIR/etc \
      ;
   do
      if [ -f "$i/$shellRc" ]
      then
         sourceFoam="$i/$shellRc"
         break
      fi
   done
fi

# Construct test string for remote execution.
# Source WM_PROJECT settings if WM_PROJECT environment not set.
# attempt to preserve the installation directory 'FOAM_INST_DIR'
case $sourceFoam in
*/bashrc)
   if [ "$FOAM_INST_DIR" ]
   then
      sourceFoam='[ "$WM_PROJECT" ] || '"FOAM_INST_DIR=$FOAM_INST_DIR . $sourceFoam"
   else
      sourceFoam='[ "$WM_PROJECT" ] || '". $sourceFoam"
   fi
   ;;

*/cshrc)
   # TODO: csh equivalent to bash code (preserving FOAM_INST_DIR)
   sourceFoam='if ( ! $?WM_PROJECT ) source '"$sourceFoam"
   ;;
esac

# quote double-quotes for remote command line
rcmd=$(echo $* | sed -e s/\"/\'\"\'/g)
## the same, without forking (not ksh, maybe not /bin/sh either)
# rcmd=$(while [ "$#" -gt 0 ]; do echo "${1//\"/'\"'}"; shift; done)


# Convert WM_COLOURS into an array
declare colourList
nColours=0
for col in $WM_COLOURS
do
   colourList[$nColours]=$col
   ((nColours = $nColours + 1))
done

# Bashism: make pipe fail early.
# This ensures the return value of the command is returned and not of the
# colouring pipe etc.
set -o pipefail


#
# colour output by argument 1
#
colourPipe()
{
   if [ "$1" ]
   then
   (
      while read line
      do
         setterm -foreground $1
         echo "$line"
      done
      setterm -foreground default
   )
   else
      cat
   fi
}


#
# prefix message with [HOSTNAME]
#
prefixPipe()
{
   while read line
   do
      echo "[$@] $line"
   done
}


colourIndex=0

while :
do
   for slotGroup in $WM_HOSTS
   do
      # split 'host:N', but catch 'host:' and 'host' too
      host=${slotGroup%%:*}
      n=${slotGroup##*:}
      [ "$n" = "$host" ] && n=1  # missing ':'
      : ${n:=1}

      i=0
      while [ "$i" -lt "$n" ]
      do
         lockFile="$lockDir/$host:$i"
         if lockfile -r0 "$lockFile" 2>/dev/null
         then
            if [ "$nColours" -gt 0 ]
            then
               # Set colour and index to next colour
               colour="${colourList[$colourIndex]}"
               colourIndex=$(expr $colourIndex + 1)
               [ "$colourIndex" -lt "$nColours" ] || colourIndex=0

               if [ "$host" = "$HOST" ]; then
                  eval $* 2>&1 | colourPipe "$colour"
               elif [ -n "$JOB_ID" ]; then
                  qrsh -inherit -v PWD $host "$rcmd"
               else
                  ssh $host "$sourceFoam 2>/dev/null; cd $PWD && $rcmd" 2>&1 | colourPipe "$colour"
               fi
               retval=$?
            else
               if [ "$host" = "$HOST" ]; then
                  eval $* 2>&1 | prefixPipe "$host"
               elif [ -n "$JOB_ID" ]; then
                  qrsh -inherit -v PWD $host "$rcmd" | prefixPipe "$host"
               else
                  ssh $host "$sourceFoam 2>/dev/null; cd $PWD && $rcmd" 2>&1 | prefixPipe "$host"
               fi
               retval=$?
            fi

            # Release lock
            rm -f "$lockFile" 2>/dev/null
            exit $retval
         fi
         i=$(expr $i + 1)
      done
   done
   # Did not find any free slots. Rest a bit.
   sleep 1
done

if [ "$nColours" -gt 0 ]
then
   setterm -foreground default
fi

#------------------------------------------------------------------------------