#!/usr/bin/expect
############################################################################
# Purpose: Establish global state information for SLURM test suite
#
# To define site-specific state information, set the values in a file
# named 'globals.local'. Those values will override any specified here.
# for example:
#
# $ cat globals.local
# set slurm_dir "/usr/local"
# set mpicc     "/usr/local/bin/mpicc"
#
############################################################################
# Copyright (C) 2002-2006 The Regents of the University of California.
# Produced at Lawrence Livermore National Laboratory (cf, DISCLAIMER).
# Written by Morris Jette <jette1@llnl.gov>
# UCRL-CODE-217948.
# 
# This file is part of SLURM, a resource management program.
# For details, see <http://www.llnl.gov/linux/slurm/>.
#  
# SLURM 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.
# 
# SLURM 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 SLURM; if not, write to the Free Software Foundation, Inc.,
# 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA.
############################################################################

global sacct sbcast scancel scontrol sinfo smap squeue srun

# Conditional set.  Only set variable if variable does not yet exist.
proc cset {name value} {
    if {![uplevel 1 info exists $name]} {
	uplevel 1 "set $name \"$value\""
    }
}

if [file exists globals.local] {
    source globals.local
}

#
# Specify the slurm install directory.
# Used to locate binaries, libraries, and header files.
#
cset slurm_dir   "/usr"
cset sacct       "${slurm_dir}/bin/sacct"
cset sbcast      "${slurm_dir}/bin/sbcast"
cset scancel     "${slurm_dir}/bin/scancel"
cset scontrol    "${slurm_dir}/bin/scontrol"
cset sinfo       "${slurm_dir}/bin/sinfo"
cset smap        "${slurm_dir}/bin/smap"
cset squeue      "${slurm_dir}/bin/squeue"
cset srun        "${slurm_dir}/bin/srun"

cset partition "pdebug"

# If using MPICH-2 or other version of MPI requiring pmi libary, use this
#cset mpicc	"/home/jette/mpich2-install/bin/mpicc"
#cset use_pmi	1
# OR for other versions of MPICH, use this
cset mpicc       "/usr/local/bin/mpicc"
cset use_pmi	0

# If using XCPU job launch, specify directory location as needed
cset xcpu_dir	"/mnt/xcpu"

cset poe		"/usr/bin/poe"
cset mpirun	"mpirun"
cset totalviewcli	"/usr/local/bin/totalviewcli"

# Pattern to match your shell prompt
cset prompt "(%|#|\\\$) *$"

#
# Specify locations of other executable files used
# Only the shell names (e.g. bin_bash) must be full pathnames
#
cset bin_awk 	"awk"
cset bin_bash	"/bin/bash"
cset bin_cat	"cat"
cset bin_cc	"gcc"
cset bin_chmod	"chmod"
cset bin_cp	"cp"
cset bin_diff	"diff"
cset bin_echo	"echo"
cset bin_env	"env"
cset bin_file	"file"
cset bin_id	"id"
cset bin_grep    "grep"
cset bin_hostname "hostname"
cset bin_kill	"kill"
cset bin_make	"make"
cset bin_pkill	"pkill"
cset bin_ps	"ps"
cset bin_pwd	"pwd"
cset bin_rm	"rm"
cset bin_sleep   "sleep"
cset bin_sort    "sort"
cset bin_touch   "touch"
cset bin_uname	"uname"
cset bin_usleep  "usleep"
cset bin_wc	"wc"

#
# Let the commands complete without expect timing out waiting for a 
# response. Single node jobs submitted to the default partition should 
# be initiated within this number of seconds.
# for interactive slurm jobs: cset timeout $max_job_delay
#
cset max_job_delay 120

#
# Files must be propogated between nodes within this number of seconds.
# The delay may be due to NFS.
#
cset max_file_delay 60

#
# Specify the maximum number of tasks to use in the stress tests.
#
cset max_stress_tasks 4

# Other common variables
set alpha                "\[a-zA-Z\]+"
set alpha_cap            "\[A-Z\]+"
set alpha_numeric        "\[a-zA-Z0-9\]+"
set alpha_numeric_under  "\[a-zA-Z0-9_\]+"
set alpha_under          "\[A-Z_\]+"
set end_of_line          "\[\r\n\]"
set number               "\[0-9\]+"
set float                "\[0-9\]+\\.?\[0-9\]+"

#
# Cache SlurmUser to check for SuperUser requests
#
cset super_user     0
cset super_user_set 0

################################################################
#
# Proc: cancel_job
#
# Purpose:  Cancel the specified job
#
# Returns: A non-zero return code indicates a failure.
#
# Input: job_id  -- The SLURM job id of a job we want to cancel.
#
################################################################

proc cancel_job { job_id } {
	global scancel bin_sleep

	send_user "cancelling $job_id\n"
	exec $scancel -q $job_id
	exec $bin_sleep 1
	return [wait_for_job $job_id "DONE"]
}


################################################################
#
# Proc: get_line_cnt
#
# Purpose:  Return size of the specified file
#
# Returns: Number of lines in the specified file.
#
# Input: file_name  -- Name of file to inspect.
#
################################################################
proc get_line_cnt { file_name } {
	global bin_wc number
	set lines 0
	spawn $bin_wc -l $file_name
	expect {
		-re "($number) " {
			set lines $expect_out(1,string)
			exp_continue
		}
		eof {
			wait
		}
	}
	return $lines
}

################################################################
#
# Proc: kill_srun
#
# Purpose:  Kill all srun commands associated with this user.
#	Issue two SIGINT, sleep 1 and a SIGKILL
#
# Returns: A non-zero return code indicates a failure.
#
# Note: usleep is used to avoid pkill race error message (the
# processes terminate after the process table is read and before
# the signal is sent to the detected processses)
#
################################################################

proc kill_srun {  } {
	global bin_id bin_pkill bin_sleep bin_usleep number

	spawn $bin_id -u
	expect {
		-re "($number)" {
			set uid $expect_out(1,string)
			exp_continue
		}
		eof {
			wait
		}
	}

	spawn $bin_pkill -INT -u $uid srun
	expect {
		eof {
			wait
		}
	}
	exec $bin_usleep 1000
	spawn $bin_pkill -INT -u $uid srun
	expect {
		eof {
			wait
		}
	}
	exec $bin_sleep  1
	spawn $bin_pkill -KILL -u $uid srun
	expect {
		eof {
			wait
		}
	}

	return 0
}


################################################################
#
# Proc: print_header
#
# Purpose:  Print header with test ID
#
# Input: job_id   -- The SLURM job id of a job we want to cancel.
#
################################################################

proc print_header { test_id } {

	send_user "============================================\n"
	send_user "TEST: $test_id\n"
}


################################################################
#
# Proc: wait_for_file
#
# Purpose:  Wait for the specified file to exist. This delay 
# provides time for NFS files to be propogated. The 
#
# Returns: A non-zero return code indicates a failure.
#
# Input: file_name   -- Name of the file to wait for.
#
################################################################

proc wait_for_file { file_name } {
	global bin_sleep max_file_delay

	for {set my_delay 0} {$my_delay <= $max_file_delay} {incr my_delay} {
		if {$my_delay > 0} {
			exec $bin_sleep 1
		}
		if [file exists $file_name] {
			return 0
		}
	}
	send_user "\nFAILURE: Timeout waiting for file $file_name\n"
	return 1
}


################################################################
#
# Proc: wait_for_job
#
# Purpose:  Wait for a previously submitted SLURM job to reach 
# the desired state, exponential back-off 1 to 10 seconds
#
# Returns: A non-zero return code indicates a failure.
#
# Input: job_id   -- The SLURM job id of a job we want to
#                    wait for.
#        desired_state -- The state you want the job to attain before
#                         returning.  Currently supports:
#                            DONE any terminated state
#                            RUNNING job is running
#
# NOTE: We sleep for two seconds before replying that a job is 
# done to give time for I/O completion (stdout/stderr files)
#
################################################################

proc wait_for_job { job_id desired_state } {
    global scontrol

    # First verify that desired_state is supported
    switch $desired_state {
	"DONE" {}
	"RUNNING" {}
	default {
	    send_user "Unsupported desired state: $desired_state\n"
	    return 1
	}
    }

    set sleep_time  1
    while 1 {
	set fd [open "|$scontrol -o show job $job_id"]
	gets $fd line
	close $fd
	if {[regexp {JobState\s*=\s*(\w+)} $line foo state] != 1} {
	    set state "NOT_FOUND"
	}

	switch $state {
	    "NOT_FOUND" -
	    "CANCELLED" -
	    "FAILED" -
	    "TIMEOUT" -
	    "NODE_FAIL" -
	    "COMPLETED" {
		if {[string compare $desired_state "DONE"] == 0} {
		    send_user "Job $job_id is DONE\n"
		    sleep 2
		    return 0
		} else {
		    if {[string compare $desired_state "RUNNING"] == 0} {
			send_user "Job $job_id is $state, but we wanted RUNNING\n"
		    }
		    return 1
		}
	    }
	    "RUNNING" {
		if {[string compare $desired_state "RUNNING"] == 0} {
		    send_user "Job $job_id is RUNNING\n"
		    return 0
		}
		send_user "Job $job_id is in state $state, desire $desired_state\n"
	    }
	    default {
		send_user "Job $job_id is in state $state, desire $desired_state\n"
	    }
	}

	sleep $sleep_time
	set sleep_time  [expr $sleep_time * 2]
	if { $sleep_time > 10 } {
	    set sleep_time 10
	}
    }
}

################################################################
#
# Proc: test_front_end
#
# Purpose: Determine if the execution host is one in which the
# slurmd daemon executes on a front-end node rather than the 
# compute hosts (e.g. Blue Gene systems). This is based upon 
# the value of SelectType in the slurm.conf.
#
# Returns 1 if the system uses a front-end, 0 otherwise
#
################################################################

proc test_front_end { } {
	global scontrol bin_bash bin_grep xcpu_dir

	if [file exists $xcpu_dir] {
		return 1
	}

	log_user 0
	set bluegene 0
	spawn -noecho $bin_bash -c "exec $scontrol show config | $bin_grep SelectType"
	expect {
		"select/bluegene" {
			set bluegene 1
			exp_continue
		}
		eof {
			wait
		}
	}
	log_user 1
	
	return $bluegene
}

################################################################
#
# Proc: test_xcpu
#
# Purpose: Determine if the system xcpu for task launch
#
# Returns 1 if the system is xcpu, 0 otherwise
#
################################################################

proc test_xcpu { } {
	global xcpu_dir

	if [file exists $xcpu_dir] {
		return 1
	}

	return 0
}

################################################################
#
# Proc: test_bluegene
#
# Purpose: Determine if the system is a bluegene system
#
# Returns 1 if the system is a bluegene, 0 otherwise
#
################################################################

proc test_bluegene { } {
	global scontrol bin_bash bin_grep

	log_user 0
	set bluegene 0
	spawn -noecho $bin_bash -c "exec $scontrol show config | $bin_grep SelectType"
	expect {
		"select/bluegene" {
			set bluegene 1
			exp_continue
		}
		eof {
			wait
		}
	}
	log_user 1
	
	return $bluegene
}

################################################################
#
# Proc: test_aix
#
# Purpose: Determine if the system is AIX
#
# Returns 1 if the system is AIX, 0 otherwise
#
################################################################

proc test_aix {} {
    set fd [open "|uname"]
    gets $fd line
    close $fd
    if {[string compare $line AIX] == 0} {
	return 1
    } else {
	return 0
    }
}

################################################################
#
# Proc: test_super_user
#
# Purpose: Determine if user is a SLURM super user (i.e. user 
# root or configured SlurmUser)
#
################################################################

proc test_super_user { } {
	global alpha_numeric bin_id number scontrol super_user super_user_set

	if {$super_user_set != 0} {
		return $super_user
	}

#
#	Check if user root
#
	log_user 0
	spawn $bin_id -u
	set uid -1
	expect {
		-re "($number)" {
			set uid $expect_out(1,string)
			exp_continue
		}
		eof {
			wait
		}
	}
	if {$uid == 0} {
		log_user 1
		set super_user 1
		set super_user_set 1
		return $super_user
	}

#
#	Check if SlurmUser
#
	spawn $bin_id -un
	set user ""
	expect {
		-re "($alpha_numeric)" {
			set user $expect_out(1,string)
			exp_continue
		}
		eof {
			wait
		}
	}
	spawn $scontrol show control
	set slurm_user ""
	expect {
		-re "SlurmUser *= ($alpha_numeric)" {
			set slurm_user $expect_out(1,string)
			exp_continue
		}
		eof {
			wait
		}
	}
	if {[string compare $user $slurm_user] == 0} {
		set super_user 1
	}
	set super_user_set 1
	log_user 1
	return $super_user
}

# Create a 16 bit hex number from a signed decimal number
# Replace all non-decimal characters
# Courtesy of Chris Cornish
# http://aspn.activestate.com/ASPN/Cookbook/Tcl/Recipe/415982
proc dec2hex16 {value} {
	regsub -all {[^0-x\.-]} $value {} newtemp
	set value [string trim $newtemp]
	if {$value < 32767 && $value > -32768} {
		set tempvalue [format "%#010X" [expr $value]]
		return [string range $tempvalue 6 9]
	} elseif {$value < 32768} {
		return "8000"
	} else {
		return "7FFF"
	}
}


################################################################
#
# Proc: available_nodes
#
# Purpose: Check to see if a given partition has a at least
#          "num_nodes" number of nodes in the alloc, idle, or comp
#          state.  This can be used to avoid launching a job that
#          will never run because nodes are in the "drained" state
#          or otherwise unavailable.
#
# Returns: Returns the number of available nodes in the partition, or
#          -1 on failure.
#
# Input: partition - name of a partition
#
################################################################

proc available_nodes { partition } {
    global sinfo

    set available -1
    set fd [open "|$sinfo --noheader --partition pdebug --state idle,alloc,comp --format %D"]
    gets $fd line
    close $fd
    regexp {\d+} $line available
    return $available
}


################################################################
#
# Proc: default_partition
#
# Purpose: Use scontrol to determine the name of the default partition
#
# Returns: Name of the current default partition
#
################################################################

proc default_partition {} {
    global scontrol

    set name ""
    set fd [open "|$scontrol --all --oneliner show partition"]
    while {[gets $fd line] != -1} {
	if {[regexp {^PartitionName=(\w+).*Default=YES} $line name]
	    == 1} {
	    break
	}
    }
    close $fd

    if {[string length $name] == 0} {
	send_user "ERROR: could not identify the default partition"
    }

    return $name
}


################################################################
#
# Proc: make_bash_script
#
# Purpose: Create a bash script of name "script_name", and
#          make the body of the script "script_contents".
#          make_bash_script removes the file if it already exists,
#          then generates the #! line, and then dumps "script_contents"
#          to the file.  Finally, it makes certain that the script
#          is executable.
#
# Returns: Nothing.
#
# Input: script_name - file name for the bash script
#        script_contents - body of the script, not including the
#                          initial #! line.
#
################################################################

proc make_bash_script { script_name script_contents } {
    global bin_bash bin_chmod

    file delete $script_name
    set fd [open $script_name "w"]
    puts $fd "#!$bin_bash"
    puts $fd $script_contents
    close $fd
    exec $bin_chmod 700 $script_name
}

################################################################
# Check if we are user root or SlurmUser
#
# Return 1 if true, 0 if false
################################################################

proc is_super_user { } {
	global alpha_numeric bin_id scontrol

	log_user 0
	set user_name "nobody"
	spawn $bin_id -u -n
	expect {
		-re "($alpha_numeric)" {
			set user_name $expect_out(1,string)
			exp_continue
		}
		eof {
			wait
		}
	}
	if {[string compare $user_name "root"] == 0} {
		log_user 1
		return 1
	}

	set found_user 0
	spawn $scontrol show config
	expect {
		-re "SlurmUser *= $user_name" {
			set found_user 1
			exp_continue
		}
		eof {
			wait
		}
	}
	log_user 1
	return $found_user
}
