#!/usr/bin/env expect
############################################################################
# Purpose: Test heterogeneous job steps
############################################################################
# Copyright (C) 2021 SchedMD LLC.
# Written by Chad Vizino <chad@schedmd.com>
#
# This file is part of Slurm, a resource management program.
# For details, see <https://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 file_in1     "$test_name.in1"
set file_in2     "$test_name.in2"
set file_in3     "$test_name.in3"
set file_out     "$test_name.out"
set job_id_list  ""
set matches      0

proc cleanup {} {
	global job_id_list bin_rm file_in1 file_in2 file_in3 file_out
	cancel_job $job_id_list
	exec $bin_rm -f $file_in1 $file_in2 $file_in3 $file_out
}

proc get_first_node_of_job {job_id} {
	global scontrol

	set nodelist [get_job_param $job_id "NodeList"]
	set hostnames [run_command_output -fail "$scontrol show hostnames $nodelist"]
	set node [lindex $hostnames 0]

	return $node
}

#
# Requirements
#
set nodes [get_nodes_by_request "-N2"]
set nb_nodes [llength $nodes]
if { $nb_nodes != 2 } {
	skip "Insufficient nodes in default partition ($nodes)"
}
set gpu_cnt [get_highest_gres_count $nb_nodes "gpu"]
if {$gpu_cnt < 2} {
        skip "This test requires 2 or more GPUs on $nb_nodes nodes of the default partition"
}
if {[get_config_param "FrontendName"] ne "MISSING"} {
        skip "This test is incompatible with front-end systems"
}
if {[get_config_param "SchedulerType"] ne "sched/backfill"} {
	skip "This test requires SchedulerType = sched/backfill"
}

# het job
make_bash_script $file_in1 "$srun -l --mpi=none ./$file_in2 : ./$file_in3"
make_bash_script $file_in2 "$bin_echo FILE2"
make_bash_script $file_in3 "$bin_echo FILE3"

log_info "SUBTEST: Confirm het job srun runs within non-het job salloc"
set matches 0
set output [run_command_output -fail "$salloc -N2 ./$file_in1"]
if {[regexp "Granted job allocation ($number)" $output - job_id]} {
	lappend job_id_list $job_id
} else {
	fail "Allocation not granted ($output)"
}
set lines_matched [regexp -all -inline "($number): FILE($number)" $output]
foreach {re step file} $lines_matched {
	if {$step == 0 && $file == 2} {
		incr matches
	} elseif {$step == 1 && $file == 3} {
		incr matches
	} else {
		fail "Invalid output ($re from $lines_matched)"
	}
}
subtest {$matches == 2} "Job should be granted and output match the expected" "$matches != 2"

log_info "SUBTEST: Confirm het job srun runs within non-het job sbatch"
set matches 0
set job_id [submit_job -fail "-N2 -o $file_out ./$file_in1"]
lappend job_id_list $job_id
if {[wait_for_job $job_id "DONE"] != 0} {
	fail "Job $job_id did not complete"
}
if {[wait_for_file $file_out] != 0} {
	fail "No output file"
}
set output [run_command_output -fail "$bin_cat $file_out"]
set matches [regexp -all {FILE[23]} $output]
subtest {$matches == 2} "Job should be granted and output match the expected" "$matches != 2"

log_info "SUBTEST: Confirm het job srun runs on different nodes in non-het job sbatch"
set matches 0
set job_id [submit_job -fail "-N2 -o $file_out --wrap=\"$srun $bin_printenv SLURMD_NODENAME\""]
lappend job_id_list $job_id
if {[wait_for_job $job_id "DONE"] != 0} {
	fail "Job $job_id did not complete"
}
if {[wait_for_file $file_out] != 0} {
	fail "No output file"
}
set matches [run_command_output -fail "sort -u $file_out|wc -l"]
subtest {$matches == 2} "Job should be granted and output match the expected" "$matches != 2"

log_info "SUBTEST: Confirm het job srun with gres runs in non-het job sbatch"
set matches 0
set job_id [submit_job -fail "-N2 --gpus=2 -o $file_out --wrap=\"$srun --gpus=1 echo ok : --gpus=1 echo ok\""]
lappend job_id_list $job_id
if {[wait_for_job $job_id "DONE"] != 0} {
	fail "Job $job_id did not complete"
}
if {[wait_for_file $file_out] != 0} {
	fail "No output file"
}
set matches [run_command_output -fail "grep ok $file_out|wc -l"]
subtest {$matches == 2} "Job should be granted and output match the expected" "$matches != 2"

log_info "SUBTEST: Confirm het job srun fails in het job salloc"
set matches 0
set output [run_command_output -xfail -fail "$salloc $file_in1 : $file_in1"]
if {[regexp "Granted job allocation ($number)" $output - job_id]} {
	lappend job_id_list $job_id
} else {
	fail "Allocation not granted ($output)"
}
if {[regexp "srun: error: Allocation failure" $output]} {
	incr matches
}
subtest {$matches == 1} "Job should be granted but allocation should fail" "$matches != 1"

log_info "SUBTEST: Confirm het job srun fails when trying to overlap nodes"
set matches 0
set timeout $max_job_delay
set salloc_pid [spawn $salloc -N2 $bin_bash]
expect {
	-re "Granted job allocation ($number)" {
		set job_id $expect_out(1,string)
		lappend job_id_list $job_id
		send "$reset_bash_prompt\r"
		exp_continue
	}
	-re "\"$test_prompt\"" {
		# skip this, just echo of setting test_prompt
		exp_continue
	}
	-re "$test_prompt" {
		# job started
	}
	timeout {
		slow_kill [expr 0 - $salloc_pid]
		fail "salloc not responding"
	}
}

set node [get_first_node_of_job $job_id]
log_info "node=$node for job $job_id"

send "$srun -w $node : -w $node $bin_true\r"
expect {
	-re "$test_prompt" {
		#break
	}
	-re "overlaps with excluded $node" {
		incr matches
	}
	timeout {
		fail "not responding"
	}
}
send "exit\r"
subtest {$matches == 1} "srun should return overlaps with excluded" "$matches != 1"
