#!/bin/bash

# Torque prologue script to limit the processes in a job to the assigned 
# CPU list
# Note: this assumes that the number of CPUs in the system matches the number
#       of job slots that Torque considers on the node
#       This assumption is not verified!
#
# Author:   Ronald Starink   < ronalds AT nikhef.nl >
#
# This prologue script is invoked with arguments
# <id> <user> <group> <name> <resource-reqs> <queue> <account>

jobid=$(echo $1 | sed -re 's/\[|\]//g')

# set debug=1 to write debug output to $debugfile
debug=1
debugfile=/tmp/prologue.log

# Determine the job slot(s)
# The result is a single number or comma-separated list of numbers 
# that will be used as the CPU numbers which can be used by the job
slot=$(pbsnodes -a $(hostname) | \
       perl -ne 'if ( m/^\s*jobs/ ) {  s/\[|\]//g ; @a = m!(\d+)/'$jobid'!g ; print join(",",@a)."\n"; }')
[ $debug -gt 0 ] && echo jobid=$jobid, slot=$slot >> $debugfile

# Determine the ID of the parent process of this prologue script,
# which is the MOM for the job
ppid=$(ps -o ppid -p $$ | sed -n 's/^[ ]*\([0-9][0-9]*\)[ ]*$/\1/p' 2>> $debugfile )

# Assign the CPU(s) to the MOM of the current job
# This setting affects all child processes that will be started later
[ $debug -gt 0 ] && echo Setting taskset to CPU $slot for pid=$ppid >> $debugfile
res=$(taskset -cp $slot $ppid 2>&1)
[ $debug -gt 0 ] && echo $res >> $debugfile

exit 0