#!/usr/bin/env expect
############################################################################
# Purpose: Test of SLURM functionality
#          Test of --spread-job option
#
# Output:  "TEST: #.#" followed by "SUCCESS" if test was successful, OR
#          "WARNING: ..." with an explanation of why the test can't be made, OR
#          "FAILURE: ..." otherwise with an explanation of the failure, OR
#          anything else indicates a failure mode that must be investigated.
############################################################################
# Copyright (C) 2016 SchedMD LLC
#
# This file is part of SLURM, a resource management program.
# For details, see <http://slurm.schedmd.com/>.
# Please also read the included file: DISCLAIMER.
#
# 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.,
# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301  USA.
############################################################################
source ./globals

set test_id     "1.98"

print_header $test_id

if {[test_alps]} {
	send_user "\nWARNING: This test is incompatible with ALPS systems\n"
	exit $exit_code
} elseif {[test_serial]} {
	send_user "\nWARNING: This test is incompatible with serial system\n"
	exit 0
} elseif {[test_bluegene]} {
	send_user "\nWARNING: This test is incompatible with bluegene system\n"
	exit 0
} elseif { [test_xcpu] } {
	send_user "\nWARNING: This test is incompatible with XCPU systems\n"
	exit 0
}

if {[test_launch_poe]} {
	set node_name_env MP_I_UPMD_HOSTNAME
} else {
	set node_name_env SLURMD_NODENAME
}

#
# Run two tasks per idle node in default partition, up to 64 tasks
#
set node_cnt 2
set partition [default_partition]
set task_cnt [available_nodes $partition "idle"]
if {$task_cnt < $node_cnt} {
    send_user "\nWARNING: not enough nodes currently available ($task_cnt avail, $node_cnt needed)\n"
    exit 0
}
send_user "TEST 1\n"
incr task_cnt $task_cnt
if {$task_cnt > 64} {
	set task_cnt 64
}

for {set inx 0} {$inx < $task_cnt} {incr inx} {
	set host($inx) ""
}
set srun_pid [spawn $srun -n $task_cnt -l -t1 --spread-job $bin_printenv $node_name_env]
expect {
	-re "($number): *($alpha_numeric_under)" {
		set host($expect_out(1,string)) $expect_out(2,string)
		exp_continue
	}
	timeout {
		send_user "\nFAILURE: srun not responding\n"
		slow_kill $srun_pid
		exit 1
	}
	eof {
		wait
	}
}

# Determine how the tasks were distributed across nodes
set host_cnt 0
for {set inx 0} {$inx < $task_cnt} {incr inx} {
	for {set jnx 0} {$jnx < $host_cnt} {incr jnx} {
		if {[string compare $host($inx) $unique($jnx)] == 0} {
			incr uniq_cnt($jnx)
			break
		}
	}
	if {$jnx >= $host_cnt} {
		set unique($jnx) $host($inx)
		set uniq_cnt($jnx) 1
		incr host_cnt
	}
}

set min_ntasks $uniq_cnt(0)
set max_ntasks $uniq_cnt(0)
for {set jnx 0} {$jnx < $host_cnt} {incr jnx} {
	if {$min_ntasks > $uniq_cnt($jnx)} {
		set min_ntasks $uniq_cnt($jnx)
	}
	if {$max_ntasks < $uniq_cnt($jnx)} {
		set max_ntasks $uniq_cnt($jnx)
	}
}
send_user "Spread $task_cnt tasks over $host_cnt hosts\n"
send_user "Ntasks per node range from $min_ntasks to $max_ntasks\n"

# Make sure tasks were spead over a reasonable number of nodes
set min_nodes $node_cnt
if {$min_nodes > $task_cnt} {
	set min_nodes $task_cnt
}
incr min_nodes -1
if {$min_nodes < 1} {
	set $min_nodes 1
}
if {$host_cnt < $min_nodes} {
	send_user "FAILURE: Did not spread job across nodes ($host_cnt < $min_nodes)\n"
	exit 1
}

#
# Run one tasks per idle node in default partition, up to 64 tasks
#
send_user "\n\nTEST 2\n"
set task_cnt $host_cnt
for {set inx 0} {$inx < $task_cnt} {incr inx} {
	set host($inx) ""
}
set srun_pid [spawn $srun -n $task_cnt -l -t1 --spread-job --ntasks-per-node=1 $bin_printenv $node_name_env]
expect {
	-re "($number): *($alpha_numeric_under)" {
		set host($expect_out(1,string)) $expect_out(2,string)
		exp_continue
	}
	timeout {
		send_user "\nFAILURE: srun not responding\n"
		slow_kill $srun_pid
		exit 1
	}
	eof {
		wait
	}
}

# Determine how the tasks were distributed across nodes
set host_cnt 0
for {set inx 0} {$inx < $task_cnt} {incr inx} {
	for {set jnx 0} {$jnx < $host_cnt} {incr jnx} {
		if {[string compare $host($inx) $unique($jnx)] == 0} {
			incr uniq_cnt($jnx)
			break
		}
	}
	if {$jnx >= $host_cnt} {
		set unique($jnx) $host($inx)
		set uniq_cnt($jnx) 1
		incr host_cnt
	}
}

set min_ntasks $uniq_cnt(0)
set max_ntasks $uniq_cnt(0)
for {set jnx 0} {$jnx < $host_cnt} {incr jnx} {
	if {$min_ntasks > $uniq_cnt($jnx)} {
		set min_ntasks $uniq_cnt($jnx)
	}
	if {$max_ntasks < $uniq_cnt($jnx)} {
		set max_ntasks $uniq_cnt($jnx)
	}
}
send_user "Spread $task_cnt tasks over $host_cnt hosts\n"
send_user "Ntasks per node range from $min_ntasks to $max_ntasks\n"

# Make sure tasks were spead over a reasonable number of nodes
if {$host_cnt < $task_cnt} {
	send_user "FAILURE: Did not spread job across nodes ($host_cnt < $task_cnt)\n"
	exit 1
}
send_user "SUCCESS\n" 
exit 0
