#!/usr/bin/expect
############################################################################
# Purpose: Test of BLUEGENE SLURM functionality
#          Stress test Dynamic mode block creation.
#
# Output:  "TEST: #.#" followed by "SUCCESS" if test was successful, OR
#          "FAILURE: ..." otherwise with an explanation of the failure, OR
#          anything else indicates a failure mode that must be investigated.
############################################################################
# Copyright (C) 2006-2007 The Regents of the University of California.
# Copyright (C) 2008 Lawrence Livermore National Security.
# Produced at Lawrence Livermore National Laboratory (cf, DISCLAIMER).
# Written by Danny Auble <da@llnl.gov>
# CODE-OCEC-09-009. All rights reserved.
#
# This file is part of SLURM, a resource management program.
# For details, see <https://computing.llnl.gov/linux/slurm/>.
# 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     "8.6"
set exit_code   0
set file_in     "test$test_id.input"
set job_id      0
set matches     0

# job parameters
set sleep_time  5

# make a bunch of blocks of the specified size (node_cnt)
proc run_batch_jobs { node_cnt job_cnt } {
	global sbatch number file_in
	set start_cnt 0
	for {set inx 0} {$inx < $job_cnt} {incr inx} {
		set sbatch_pid [spawn $sbatch --output=/dev/null -t5 -N$node_cnt-$node_cnt $file_in]
		expect {
			-re "Node count specification invalid" {
				send_user "This error was expected, no worries\n"
				return -1
			}
			-re "More processors requested than permitted" {
				send_user "This error was expected, no worries\n"
				return -1
			}
			-re "Requested node configuration is not available" {
				send_user "This error was expected, no worries\n"
				return -1
			}
			-re "Submitted batch job ($number)" {
				incr start_cnt
				exp_continue
			}
			-re "Unable to contact" {
				send_user "\nFAILURE: slurm appears to be down\n"
				exp_continue
			}
			timeout {
				send_user "\nFAILURE: sbatch not responding\n"
				slow_kill $sbatch_pid
				return 0
			}
			eof {
				wait
			}
		}
	}

	return $start_cnt
}

# Wait up to 900 seconds for all jobs to terminate
# Return 0 if all jobs done, remainin job count otherwise
proc wait_for_all_jobs { } {
	global scancel squeue bin_sleep file_in

	set last_matches 0
	send_user "Waiting for all jobs to terminate\n"
	for {set inx 0} {$inx < 600} {incr inx} {
		set matches 0
		log_user 0
		spawn $squeue -o %j
		expect {
			-re "$file_in" {
				incr matches
				exp_continue
			}
			-re "error" {
				set matches -1
			}
			timeout {
				send_user "No response from squeue\n"
				set matches -1
			}
			eof {
				wait
			}
		}
		log_user 1
		if {$matches == 0} {
			send_user "All jobs complete\n"
			break
		}
		if {$matches > 0} {
			send_user "  $matches jobs remaining\n"
#			Moab can slow throughput down a lot,
#			so don't return here
#			if {$matches == $last_matches} {
#				send_user "Running jobs hung\n"
#				break
#			}
#			set last_matches $matches
			exec $bin_sleep 15
		}
	}
	exec $scancel -n $file_in
	return $matches
}

proc run_bgl_test { } {
	global psets num_nodes

	set 32node_block_cnt 16
	set 128node_block_cnt 8
	set 512node_block_cnt 16
	set 1knode_block_cnt  8
	set 4knode_block_cnt  8
	set 8knode_block_cnt  4
	set 16knode_block_cnt 4
	set 32knode_block_cnt 4
	set started 0
	if {$psets >= 16} {
		incr started [run_batch_jobs 32 $32node_block_cnt]
	}

	incr started [run_batch_jobs 128 $128node_block_cnt]
	incr started [run_batch_jobs 512 $512node_block_cnt]

	if {$num_nodes >= 1024} {
		incr started [run_batch_jobs 1k $1knode_block_cnt]
		if {$num_nodes >= 4096} {
			incr started [run_batch_jobs 4k $4knode_block_cnt]
			if {$num_nodes >= 8192} {
				incr started [run_batch_jobs 8k $8knode_block_cnt]
				if {$num_nodes >= 16384} {
					incr started [run_batch_jobs 16k $16knode_block_cnt]
					if {$num_nodes >= 32768} {
						incr started [run_batch_jobs 32k $32knode_block_cnt]
					}
				}
			}
		}
		incr started [run_batch_jobs 1k $1knode_block_cnt]
	}

	incr started [run_batch_jobs 512 $512node_block_cnt]

	incr started [run_batch_jobs 128 $128node_block_cnt]

	if {$psets >= 16} {
		incr started [run_batch_jobs 32 $32node_block_cnt]
	}

	incr started [run_batch_jobs 512 $512node_block_cnt]

	return $started;
}


proc run_bgp_test { } {
	global psets num_nodes

	set 16node_block_cnt 32
	set 32node_block_cnt 16
	set 64node_block_cnt 8
	set 128node_block_cnt 8
	set 256node_block_cnt 8
	set 512node_block_cnt 16
	set 1knode_block_cnt  8
	set 4knode_block_cnt  8
	set 8knode_block_cnt  4
	set 16knode_block_cnt 4
	set 32knode_block_cnt 4
	set started 0

	if {$psets >= 32} {
		incr started [run_batch_jobs 16 $16node_block_cnt]
	}

	if {$psets >= 16} {
		incr started [run_batch_jobs 32 $32node_block_cnt]
	}
	if {$psets >= 8} {
		incr started [run_batch_jobs 64 $64node_block_cnt]
	}

	incr started [run_batch_jobs 128 $128node_block_cnt]
	incr started [run_batch_jobs 256 $256node_block_cnt]
	incr started [run_batch_jobs 512 $512node_block_cnt]

	if {$num_nodes >= 1024} {
		incr started [run_batch_jobs 1k $1knode_block_cnt]
		if {$num_nodes >= 4096} {
			incr started [run_batch_jobs 4k $4knode_block_cnt]
			if {$num_nodes >= 8192} {
				incr started [run_batch_jobs 8k $8knode_block_cnt]
				if {$num_nodes >= 16384} {
					incr started [run_batch_jobs 16k $16knode_block_cnt]
					if {$num_nodes >= 32768} {
						incr started [run_batch_jobs 32k $32knode_block_cnt]
					}
				}
			}
		}
		incr started [run_batch_jobs 1k $1knode_block_cnt]
	}

	incr started [run_batch_jobs 512 $512node_block_cnt]
	incr started [run_batch_jobs 256 $256node_block_cnt]
	incr started [run_batch_jobs 128 $128node_block_cnt]
	if {$psets >= 8} {
		incr started [run_batch_jobs 64 $64node_block_cnt]
	}

	if {$psets >= 16} {
		incr started [run_batch_jobs 32 $32node_block_cnt]
	}

	if {$psets >= 32} {
		incr started [run_batch_jobs 16 $16node_block_cnt]
	}

	incr started [run_batch_jobs 512 $512node_block_cnt]

	return $started;
}


# TEST STARTS HERE

print_header $test_id

if {[test_bluegene] == 0} {
	send_user "\nWARNING: This test is only compatable with bluegene systems\n"
	exit $exit_code
}

if {[string compare [get_bluegene_layout] Dynamic]} {
	send_user "\nWARNING: This test is only compatable with dynamic bluegene systems\n"
	exit $exit_code
}

set psets [get_bluegene_psets]

if {$psets == 0} {
	send_user "\nFAILURE: No psets are set on this system\n"
	exit 1
}

set num_nodes [expr [get_node_cnt] * [get_bluegene_cnodes_per_mp]]
if {$num_nodes == 0} {
	send_user "\nFAILURE: No nodes are found on this system\n"
	exit 1
}

set procs_per_cnode [get_bluegene_procs_per_cnode]

if {$procs_per_cnode == 0} {
	send_user "\nFAILURE: Couldn't determine procs per cnode\n"
	exit 1
}

set type [get_bluegene_type]

if {$type == 0} {
	send_user "\nFAILURE: No bluegene type found \n"
	exit 1
}


make_bash_script $file_in "$bin_sleep $sleep_time"
set timeout [expr $max_job_delay + $sleep_time]


if {![string compare $type "P"] || ![string compare $type "Q"]} {
	set started [run_bgp_test]
} elseif {![string compare $type "L"]} {
	set started [run_bgl_test]
} else {
	send_user "\nFAILURE: unknown bluegene system type '$type'\n";
	exit 1
}

if {!$started} {
	send_user "\nFAILURE: No jobs were started\n";
	exit 1
}

send_user "Started $started jobs\n"

if {[wait_for_all_jobs] != 0} {
	send_user "\nFAILURE: some submitted jobs failed to terminate\n"
	set exit_code 1
}

if {$exit_code == 0} {
	exec rm -f $file_in
	send_user "\nSUCCESS\n"
}
exit $exit_code
