#!/bin/tcsh -f
#
# DMC -- modified for an NxM MPI TRANSP job -- PPPL RedHat 4 PBS environment
#
# Assumption: ~/.cshrc sets USER_ENV_RESET_DISABLED to prevent system
# reset of PATH and module loads on shell startup.  Note use of "-f" switch.
#
# In this script, TRANSP runs as a serial process that from time to time
# launches MPI subprocesses, sharing data via files...
#
#============================================================================
# (comments from the PPPL system PBS script template...)
#
# --- this is a sample script for running a program on the beowulf
#     cluster. You modify this script for your needs, then use the
#     'qsub' command to submit it to the cluster's job scheduler.
#
#     The job scheduler is the Portable Batch System (PBS). It
#     uses directives in this file to determine where to run
#     the job, how many hosts to run it upon, where to put output
#     files, etc.
#
# remarks: a line beginning with # is a comment
#          a line beginning with #PBS is a pbs directive
#          assume (upper/lower) case to be sensitive
##PBS -q kruskal
#
# ------------------------------------------------------------
# Basic Information
#
# --- the name of your job (probably overridden in qsub command)
#PBS -N trmpi_nnxmm
 
# --- Mail options
#     mail on execution("b"), termination ("e"), or interruption ("a")
#PBS -m aeb
 
# ------------------------------------------------------------
# Job execution details
 
# --- run the job on 'n' nodes with 'm' processors per node (ppn)
#     (script to replace value, below, with nodes=<n>:ppn=<m>)
#PBS -l NxM

# --- specify the period of time the job will run in the format
#     hh:mm:ss (ex. 72 hours is 72:00:00)
#     run will gracefully exit after "walltime - 6" hours
#     and automatically restart. 
#     The 6 hours are grace time to be sure a restart file was written. 
#     3 days = 72+6 = 78
#PBS -l walltime=78:00:00

# --- submit this job to a special queue. If you do not specify a
#     special queue, the job will run in the default queue
#     There are three main queues at PPPL:
#     kestrel - the default queue used for multi node jobs
#     kite - a special queue of nodes connected with Infiniband
#     sque - nodes used for single CPU jobs (serial jobs)
#
# (for now, do not specify a queue here)
####     PBS -q sque

# --- do not rerun this job if it fails
#PBS -r n

# ------------------------------------------------------------
# When a batch job starts execution, a number of environment
# variables are predefined, which include:
#
#      Variables defined on the execution host.
#      Variables exported from the submission host with
#                -v (selected variables) and -V (all variables).
#      Variables defined by PBS.
#
# The following reflect the environment where the user ran qsub:
# PBS_O_HOST    The host where you ran the qsub command.
# PBS_O_LOGNAME Your user ID where you ran qsub.
# PBS_O_HOME    Your home directory where you ran qsub.
# PBS_O_WORKDIR The working directory where you ran qsub.
#
# These reflect the environment where the job is executing:
# PBS_ENVIRONMENT is set by PBS to "PBS_BATCH" if the job is a batch job, or
#                 to "PBS_INTERACTIVE" if the job is an interactive job.
# PBS_O_QUEUE   The original queue you submitted to.
# PBS_QUEUE     The queue the job is executing from.
# PBS_JOBID     The job's PBS identifier.
# PBS_JOBNAME   The job's name.

# export all my environment variables to the job
###NERSC -V
###NERSC -l mppwidth=N

# --- define standard output and error files for the job
#     default is <jobname>.o<jobid> and <jobname>.e<jobid>
##PBS -o test.out
##PBS -e test.err
# --- or combine both output and error streams into one file with the
#     job name and id as the default file name
#PBS -j oe

# ------------------------------------------------------------
# CHECK ENVIRONMENT

# 10/26/09 CLF
# If mounted filesystem, do not create/move directories
set NODE = `uname -n`
if ($NODE == stix.pppl.gov ) then
   set NO_MOVE = 1
endif

if ($?NERSC_HOST) then
   set NO_MOVE = 1
endif

###SWAPGCC source /u/pshare/tshare_gcc_setup.csh

set runtr_chk = `which runtr`

@ ier = 0
if ( $#runtr_chk != 1 ) then
  @ ier++
else
  if ( ! -f $runtr_chk ) then
    @ ier++
  endif
endif

if ( $ier > 0 ) then
  echo " "
  if ( -f ~/.login ) then
     echo " runtr: not found -- so, source ~/.login "
     source ~/.login
  else
     echo " runtr: not found -- exit "
     exit 1
  endif
endif

echo " "
echo " which runtr: "
which runtr
echo " "

set runtr_chk = `which runtr`
if ( $#runtr_chk != 1 ) then
  echo " runtr: still not found after source ~/.login -- must stop "
  exit 1
endif

# ------------------------------------------------------------
# Log interesting information
#

echo "-------------------"
echo "This is a $PBS_ENVIRONMENT job"
echo "This job was submitted to the queue: $PBS_QUEUE"
echo "The job's id is: $PBS_JOBID"
echo "-------------------"
echo "The master node of this job is: $PBS_O_HOST"
echo " "
echo "Contents of PBS_NODEFILE = $PBS_NODEFILE"
echo " ----- PBS_NODEFILE ----- "
cat $PBS_NODEFILE
echo " --- End PBS_NODEFILE --- "
echo " "

# --- the nodes allocated to this job are listed in the
#     file PBS_NODEFILE

# --- count the number of processors allocated to this run 
# --- $NPROCS is required by mpirun.
set NPROCS=`wc -l < $PBS_NODEFILE`

set NNODES=`uniq $PBS_NODEFILE | wc -l`

@ npernode = $NPROCS / $NNODES

echo "This job is using $NPROCS CPU(s) on the following $NNODES node(s):"
echo "-----------------------"
uniq $PBS_NODEFILE | sort
echo "-----------------------"

set NODE_LIST = `uniq $PBS_NODEFILE | sort`

setenv TRANSP_NPROCS $NPROCS

###setenv NBI_NPROCS Nnbi 
###setenv TOR_NPROCS Ntor
###setenv PTR_NPROCS Nptr

# ------------------------------------------------------------
# Setup execution variables
#
# Setup working directory & data source directory
#
# --- PBS_O_WORKDIR is the working directory from
#     which this job was submitted using 'qsub'
#
echo "The job origination directory is $PBS_O_WORKDIR"

set tok = `tok_extract $PBS_O_WORKDIR`

if ( "$tok" == "none" ) then

  echo " ?pppl_trmpi_NxM.tcsh:  could not infer TOKAMAK from path:"
  echo "  $PBS_O_WORKDIR from which job was submitted."
  echo "  Execution must stop."
  exit 1

endif

set username = `whoami`

if (! $?NERSC_HOST) then
  if ( ! -d /local/$username ) then
    mkdir /local/$username
  endif

  if ( ! -d /local/$username/work ) then
    mkdir /local/$username/work
  endif

  setenv WORKDIR /local/$username/work



  if ( ! -d $WORKDIR/$tok ) then
    mkdir -p $WORKDIR/$tok
  endif

# work here...

  cd $WORKDIR/$tok

# relative path after /local:

  set suffix = $username/work/$tok

# get data from...

  set mroot = `echo $PBS_O_HOST | sed 's#\..*##'`
  echo " mroot: $mroot"

  set orig_dir = `echo $PBS_O_WORKDIR | sed "s#/local/#/l/$mroot/#"`
else
   set orig_dir = $PBS_O_WORKDIR
   cd  $orig_dir
   setenv CRAY_ROOTFS DSL
endif
echo " %pppl_trmpi_NxM.tcsh: looking for input data here: $orig_dir "



# ------------------------------------------------------------
# Run the program
#
# --- print out the current time, then use the 'time' command
#     to start the program and print useful stats
#
echo -n 'Started job at : ' ; date
echo " ====================================== "
echo " "

echo " hostname:"
hostname

echo " pwd:"
pwd

echo " "

#  if TRANSP environment has not been set up, $CONFIGDIR will be undefined.

#set tokid = $cwd:t
set tokid = $tok
echo " tokid: $tokid"
fgrep $tokid $CONFIGDIR/TOKAMAK.DAT
if ( $status != 0 ) then
  echo " ?pppl_trmpi_NxM.tcsh -- $tokid not found in $CONFIGDIR/TOKAMAK.DAT "
  exit 1 
endif

echo " "
echo " jobname: $PBS_JOBNAME"

#  jobname = q<runid>_NxM; strip off leading q and trailing _NxM jobsize info.
#    jobnames must start with alphabetic character

set runid = `echo $PBS_JOBNAME | sed 's#^q##' | sed 's#_.*##'`
echo " runid: $runid"

# in case of re-start, save previous pbs log
  set nxm = `echo $PBS_JOBNAME | sed 's/.*_//'`
  set log_nam=$orig_dir/q${runid}_${nxm}.joblog
  if ( -f $log_nam ) then
     echo "mv $log_nam  ${log_nam}_`date +'%d-%k'`"
     mv $log_nam  ${log_nam}_`date +'%d-%k'`
  endif

echo " ------------------------------------------------- "
echo " directory of origin files at start: "
ls $orig_dir

if ( ! -f $orig_dir/${runid}TR.DAT ) then
  echo " ?pppl_trmpi_NxM.tcsh -- ${runid}TR.DAT not found in $orig_dir ..."
  exit 1
endif
#-------------------------------------------------
# set runtr_arg 
echo " "
if ( $?RUNTR_ARG != 1 ) then
  echo " RUNTR_ARG environment variable not found; default applied:  all "
  set runtr_arg = all
else
  echo " RUNTR_ARG = $RUNTR_ARG"
  set runtr_arg = $RUNTR_ARG
endif
#-------------------------------------------------

# set log file log_level
if ( $?LOG_LEVEL != 1 ) then
  echo " "
  echo " LOG_LEVEL environment variable not found; default applied: 1 (warn) "
  set logfile_level = 1
else
  set logfile_level = $LOG_LEVEL
endif
echo " LOG_LEVEL environment variable = $logfile_level"
#-------------------------------------------------

#-------------------------------------------------
#  deal with distributed NUBEAM working files
#  ...and with other distributed working files...
#  as may exist on a restart...
#  to do this, need to form a list associating cpu# with nodename
#  (mpirun mi_proclist)

  setenv NUBEAM_WORKPATH ${runid}_fi
  if ( ! -d $NUBEAM_WORKPATH ) then
    mkdir $NUBEAM_WORKPATH
  else
   if($runtr_arg == "all") then
# if transp argument = "all" clean up $NUBEAM_WORKPATH 
    rm -f $NUBEAM_WORKPATH/* 
    rm -f *birth*
    endif
  endif

  setenv MYMPI_WORKPATH ${runid}_log
  if ( ! -d $MYMPI_WORKPATH ) then
    mkdir $MYMPI_WORKPATH
  endif

  if ( ! -d $orig_dir/$NUBEAM_WORKPATH ) then
    mkdir $orig_dir/$NUBEAM_WORKPATH
  else
   if($runtr_arg == "all") then
# if transp argument = "all" clean up $orig_dir/$NUBEAM_WORKPATH 
    rm -f $orig_dir/$NUBEAM_WORKPATH/* 
    rm -f $orig_dir/*birth*
    endif
  endif
  set orig_nubeam_wkpath = $orig_dir/$NUBEAM_WORKPATH

  if ( ! -d $orig_dir/$MYMPI_WORKPATH ) then
    mkdir $orig_dir/$MYMPI_WORKPATH
  endif
  set orig_mpi_wkpath = $orig_dir/$MYMPI_WORKPATH

  echo $NPROCS > $orig_nubeam_wkpath/Nnodes
  echo $npernode > $orig_nubeam_wkpath/Mpernode

  echo " "
  echo " get processor list... "
#  set plist = $NUBEAM_WORKPATH/proclist.dat
  set plist = proclist.dat
  rm -f $plist
    if ( $?MPI_CMD) then
       eval 'set mpx = "'"$MPI_CMD"'"'
    else
       set mpx = `which mpirun`
    endif  
  $mpx -np $NPROCS `which mpi_proclist` $plist
  if ( ! -f $plist ) then
    echo "mpi_proclist failed"
    exit 1
  endif

  set cpnums = `cat $plist | awk '{print $1}'`
  set cpnames = `cat $plist | awk '{print $2}' | sed 's#\..*##' | uniq`

if ( ! $?NO_MOVE) then
  echo " create work directories on each compute server node..."
  set targlist = ()
# 03/06/09 CLF: $NODE_LIST not usable on kruskal, use $cpnames
  set nubeam_dirs=$orig_dir/${runid}.nubeam_dirs
  rm -f $nubeam_dirs
  foreach node ( $cpnames )
    set target = "/l/$node/$suffix"
    run_gmkdir $target/$NUBEAM_WORKPATH
    run_gmkdir $target/$MYMPI_WORKPATH
    set targlist = ( $targlist $target )
    # for fi_entar
    if ( $node != $cpnames[1] ) then
      echo " echo $target >> $nubeam_dirs"
      echo $target >> $nubeam_dirs
    endif
  end

  # check (due to paranoia about NFS...)
  echo " check work directories on each compute server node..."
  @ ier = 0
  foreach node ( $cpnames )
    set target = "/l/$node/$suffix"
    echo " work directory: $target/$NUBEAM_WORKPATH "
    if ( ! -d $target/$NUBEAM_WORKPATH ) then
      @ ier++
      echo " ?? compute node work directory missing: $target/$NUBEAM_WORKPATH "
    else
      # delete any possibly left over files in the work directories
      set dlist = `ls $target/$NUBEAM_WORKPATH`
      @ j = 0
      while ( $j < $#dlist )
        @ j++
        rm -f $target/$NUBEAM_WORKPATH/$dlist[$j]
      end
    endif
    if ( ! -d $target/$MYMPI_WORKPATH ) then
      @ ier++
      echo " ?? compute node work directory missing: $target/$MYMPI_WORKPATH "
    else
      # delete any possibly left over files in the work directories
      set dlist = `ls $target/$MYMPI_WORKPATH`
      @ j = 0
      while ( $j < $#dlist )
        @ j++
        rm -f $target/$MYMPI_WORKPATH/$dlist[$j]
      end
    endif
  end
endif  # End not NO_MOVE

  if ( $ier > 0 ) then
    exit 1
  endif
 
 
#####
# moved part 2 to trmpi_mvnb
# distribute NUBEAM_WORKPATH work files

# 10/26/09 CLF: handle stix
if (! $?NO_MOVE) then
  $CODESYSDIR/pbs/trmpi_mvnb $NPROCS $NUBEAM_WORKPATH $orig_dir $suffix $plist
endif
#######

#-------------------------------------------------
# just copy the remaining files...
if ( $orig_dir != `pwd` ) then
  cp -r $orig_dir/${runid}* .
endif

#
# remove nx* files if not a tr_* run
#
if ( $runtr_arg == "all" ) then
  if ( -f nxlist.summary ) then
    set trname = `echo $username | sed -e "s#tr_##"`
    if ( "$username" == "$trname" ) then
      rm -f nxlist.summary nx_deleted.dat
    endif
  endif
endif

rm -f ${runid}.time_stop

#-------------------------------------------------
echo " ------------------------------------------------- "
echo " local files before start of TRANSP:"
ls


echo " ------------------------------------------------- "
echo " runtr $runid $runtr_arg wait "
echo "       logfile: ${runid}tr.log "

runtr $runid $runtr_arg wait -log $logfile_level

set rstat = $status

echo " runtr exit status (0 = normal): $rstat "
echo " ------------------------------------------------- "
set nocleanup = 0

if ( $rstat != 0 ) then
  echo " ------------------------------------------------- "
  if ( -f ${runid}.time_stop ) then
     set timeout=TRUE
     echo " *** Time Out *** "
  else
     echo " *** error exit *** "
  endif
  echo " copy back local files after TRANSP exit:"
  ls
  echo " (NUBEAM files first...) "
  echo " mv all results to: $orig_nubeam_wkpath"


  if (! $?NO_MOVE) then
    @ j = 0
    while ( $j < $#targlist )
      # copy MPI logs
      @ j++
      set targdir = $targlist[$j]/$MYMPI_WORKPATH
      echo " contents of: $targdir (w/changes)"
      set zlist = `ls $targdir`
      if ( $#zlist > 0 ) then
        cp -r $targdir/* $orig_mpi_wkpath
      endif

      # copy MPI NUBEAM sublist states
      set targdir = $targlist[$j]/$NUBEAM_WORKPATH
      echo " "
      echo " ------------------- "
      echo " contents of: $targdir (w/changes)"
      echo " "
      set zlist = `ls $targdir`
      if ( $#zlist == 0 ) then
        echo " (empty) "
      else
        foreach file ( $zlist )
          if ( ! -f $orig_nubeam_wkpath/$file ) then
            echo " moved: $file "
            mv $targdir/$file $orig_nubeam_wkpath/$file
	    set mvstatus = $?
	    if ( $mvstatus != 0 ) then
		set nocleanup = 1
		echo " Error occurs: do not cleanup  " 	
	    endif
          else
	    cmp -s $targdir/$file $orig_nubeam_wkpath/$file
	    set cstat = $status
	    if ( $cstat == 0 ) then
	      # if identical: remove compute server copy
              echo " unchanged: $file"
	      rm -f $targdir/$file
            else
              echo " changed: $file"
              mv $targdir/$file $orig_nubeam_wkpath/$file
	      set mvstatus = $?
	      if ( $mvstatus != 0 ) then
		set nocleanup = 1
		echo " Error occurs: do not cleanup  " 	
	      endif
            endif
          endif
        end
      endif
    end

# copy back remaining files...
    echo " " 
    echo " copy back remaining files... "
    echo " cp -r ${runid}* $orig_dir "
    echo " " 

    cp -r ${runid}* $orig_dir
    set mvstatus = $?
    if ( $mvstatus != 0 ) then
	set nocleanup = 1
	echo " Error occurs: do not cleanup  " 	
    endif
  endif  # End not NO_MOVE

  echo " ------------------------------------------------- "
  echo " adjust status files... "

  set zlist = `ls $orig_dir | grep "${runid}zz.runtrx"`
  if ( $#zlist > 0 ) then
    rm -f $orig_dir/${runid}zz.runtrx:*
  endif
  if ( $?timeout ) then
    date > $orig_dir/${runid}zz.runtrx:tim
  else
    date > $orig_dir/${runid}zz.runtrx:abo
  endif
  echo " ------------------------------------------------- "
  echo " directory of origin files at end: "
  ls $orig_dir
  set mvstatus = $?
  if ( $mvstatus != 0 ) then
    set nocleanup = 1
  endif


else
  echo " **NORMAL EXIT** "
  echo " clean up $runid files on master node directory of origin..."
  rm -r $orig_dir/${runid}*
endif

if (! $?NO_MOVE) then
    if ( $nocleanup != 0 ) then
	echo " EXITED and DID NOT clean up compute node local disk files "
	exit 1
    endif
    echo " clean up compute node local disk files "
    foreach node ( $NODE_LIST )
	set nnode = `echo $node | sed 's#\..*##'`
	set target = "/l/$nnode/$suffix"
	echo " ...on $target... "
    rm -r $target/${runid}*
    end
endif  # End not NO_MOVE

echo " "
echo " ====================================== "
if ( $?timeout ) then
  cd $orig_dir
  mv ${runid}.time_stop ${runid}.time_old
  rm ${runid}.timeleft
  echo -n 'Restart job at  : ' ; date
# get submit arguments from qstat
  set lines=`qstat -f $PBS_JOBID |grep -A4 submit_args`
  echo " "
  echo $lines
  echo " "
  set old_args=`echo $lines | sed -e 's/submit_args =//' -e 's/ pppl_trmpi.*//'`
  set new_args=`echo $old_args  |awk -F-  '{ for (i=2; i<=NF; i++) pairs[$i]++  } \
  END { for (i in pairs) print "-"i } ' |sed -e s'# ##'g -e s'#\(^-.\)#\1 #'`

  set new_args=`echo $new_args | sed -e 's/RUNTR_ARG=all/RUNTR_ARG=rs/'`
  set nxm = `echo $PBS_JOBNAME | sed 's/.*_//'`
  set pidi = `get_a_pid`
  set target=pppl_trmpi_${nxm}_${pidi}.tcsh
  set target0=$0
  set targetb=`echo $target0 | sed -e 's/_sub//'`
  cp $targetb $target
  echo "qsub $new_args $target" 
  qsub $new_args $target
  set rstat = $status
else
  echo -n 'Ended job at  : ' ; date
endif
echo " " 
echo " TRANSP exit status: $rstat "
echo " " 
exit $rstat
