#!/usr/bin/env expect
############################################################################
# Purpose: Test of Slurm functionality
#          Test allocating sub-sets of GRES to job steps
############################################################################
# Copyright (C) 2018 SchedMD LLC
# Written by Morris Jette
#
# 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.,
# 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA.
############################################################################
source ./globals

set file_in1       "$test_dir/input1"
set file_in2       "$test_dir/input2"
set file_out       "$test_dir/output"
set job_id         0
set constrain_devices [expr {[get_config_param "ConstrainDevices"] eq "yes"}]

if {![check_config_select "cons_tres"]} {
	skip "This test is only compatible with select/cons_tres"
}
if {[get_config_param "FrontendName"] ne "MISSING"} {
	skip "This test is incompatible with front-end systems"
}

if {$constrain_devices} {
	log_info "Devices files are constrained by cgroups"
} else {
	log_info "Devices files are NOT constrained by cgroups"
}

proc cleanup {} {
	global job_id

	cancel_job $job_id
}

#
# Test --gpus-per-node option by job step
#
proc test_gpus_per_node_parallel_1_delayed {} {
	global file_in1 file_in2 file_out test_name job_id constrain_devices
	global sbatch srun squeue bin_cat number scontrol

	file delete $file_out
	make_bash_script $file_in1 "
		$scontrol -dd show job \${SLURM_JOBID}
		$srun --exact -n1 --gpus-per-node=1 --exact -n1 --mem=0 $file_in2 &
		$srun --exact -n1 --gpus-per-node=1 --exact -n1 --mem=0 $file_in2 &
		$srun --exact -n1 --gpus-per-node=1 --exact -n1 --mem=0 $file_in2 &
		wait
		exit 0"

	make_bash_script $file_in2 "
		echo 'STEP_ID:'\$SLURM_STEP_ID 'CUDA_VISIBLE_DEVICES:'\$CUDA_VISIBLE_DEVICES
		sleep 3
		if \[ \$SLURM_STEP_ID -eq 2 \]; then
			$squeue -s --name=$test_name
		fi
		exit 0"

	set job_id [submit_job -fail "--cpus-per-gpu=1 --gpus-per-node=2 -N1 -n3 -t1 -o $file_out -J $test_name $file_in1"]

	wait_for_job -fail $job_id "DONE"
	wait_for_file -fail $file_out

	set output [run_command_output -fail "$bin_cat $file_out"]
	subtest {![regexp "STEP_ID:($number) CUDA_VISIBLE_DEVICES:($number)," $output]} "Verify that all steps used only 1 GPU"
	subtest {[regexp -all "STEP_ID:($number) CUDA_VISIBLE_DEVICES:($number)" $output] == 3} "Verify that a GPU was used were used 3 times"
	subtest {[regexp -all "step creation temporarily disabled" $output] == 1} "Verify that 1 step was delayed"

	if {$constrain_devices} {
		subtest {[regexp -all "STEP_ID:($number) CUDA_VISIBLE_DEVICES:0" $output] == 3} "Verify that all GPUs are CUDA_VISIBLE_DEVICES:0 (with ConstrainDevices)"
	} else {
		subtest {[regexp -all "STEP_ID:($number) CUDA_VISIBLE_DEVICES:0" $output] == 2} "Verify that 2 GPUs are CUDA_VISIBLE_DEVICES:0 (without ConstrainDevices)"
		subtest {[regexp -all "STEP_ID:($number) CUDA_VISIBLE_DEVICES:1" $output] == 1} "Verify that 1 GPUs is CUDA_VISIBLE_DEVICES:1 (without ConstrainDevices)"
	}
}

#
# Test parallel step args with a job with --gpus-per-node
#
proc test_gpus_per_node_parallel {step_args} {
	global file_in1 file_in2 file_out test_name job_id constrain_devices
	global sbatch srun squeue bin_cat number scontrol

	make_bash_script $file_in1 "
		$scontrol -dd show job \${SLURM_JOBID}
		$srun --exact --gpus-per-node=0 --mem=0 $step_args $file_in2 &
		$srun --exact --gpus-per-node=0 --mem=0 $step_args $file_in2 &
		wait
		exit 0"

	make_bash_script $file_in2 "
		echo 'STEP_ID:'\$SLURM_STEP_ID 'CUDA_VISIBLE_DEVICES:'\$CUDA_VISIBLE_DEVICES
		sleep 3
		if \[ \$SLURM_STEP_ID -eq 1 \]; then
			$scontrol show step \$SLURM_JOB_ID.\$SLURM_STEP_ID
		fi
		exit 0"

	file delete $file_out
	set job_id [submit_job -fail "--cpus-per-gpu=2 --gpus-per-node=2 -N1 -n2 -t1 -o $file_out -J $test_name $file_in1"]

	wait_for_job -fail $job_id "DONE"
	wait_for_file -fail $file_out

	set output [run_command_output -fail "$bin_cat $file_out"]
	subtest {![regexp "STEP_ID:($number) CUDA_VISIBLE_DEVICES:($number)," $output]} "Verify that all steps used only 1 GPU"
	subtest {[regexp -all "STEP_ID:($number) CUDA_VISIBLE_DEVICES:($number)" $output] == 2} "Verify that a GPU was used were used 2 times"
	subtest {![regexp -all "step creation temporarily disabled" $output]} "Verify that all steps run in parallel"

	if {$constrain_devices} {
		subtest {[regexp -all "STEP_ID:($number) CUDA_VISIBLE_DEVICES:0" $output] == 2} "Verify that all GPUs are CUDA_VISIBLE_DEVICES:0 (with ConstrainDevices)"
	} else {
		subtest {[regexp -all "STEP_ID:($number) CUDA_VISIBLE_DEVICES:0" $output] == 1} "Verify that 1 GPUs is CUDA_VISIBLE_DEVICES:0 (without ConstrainDevices)"
		subtest {[regexp -all "STEP_ID:($number) CUDA_VISIBLE_DEVICES:1" $output] == 1} "Verify that 1 GPUs is CUDA_VISIBLE_DEVICES:1 (without ConstrainDevices)"
	}
}

#
# Test --gpus (per job or step) option by job step
#
proc test_gpus_per_node_different_gpus {} {
	global file_in1 file_in2 file_out test_name job_id constrain_devices
	global sbatch srun squeue bin_cat number scontrol

	make_bash_script $file_in1 "
		$scontrol -dd show job \${SLURM_JOBID}
		$srun --exact -n2 --gpus=2 --gpus-per-node=0 --mem=0 $file_in2 &
		$srun --exact -n1 --gpus=1 --gpus-per-node=0 --mem=0 $file_in2 &
		wait
		exit 0"

	make_bash_script $file_in2 "
		echo 'STEP_ID:'\$SLURM_STEP_ID 'CUDA_VISIBLE_DEVICES:'\$CUDA_VISIBLE_DEVICES
		sleep 3
		exit 0"

	file delete $file_out
	set job_id [submit_job -fail "--cpus-per-gpu=1 --gpus-per-node=3 -N1 -n3 -t2 -o $file_out -J $test_name $file_in1"]

	wait_for_job -fail $job_id "DONE"
        wait_for_file -fail $file_out

	set output [run_command "$bin_cat $file_out"]
	set step_2gpu [lreplace [regexp -inline "STEP_ID:$number CUDA_VISIBLE_DEVICES:($number),($number)" $output] 0 0]
	set step_1gpu [lreplace [regexp -inline -line "STEP_ID:$number CUDA_VISIBLE_DEVICES:($number$)" $output] 0 0]

	if {[llength $step_1gpu] != 1 || [llength $step_2gpu] != 2} {
		fail "Error parsing output fail to obtain all GPUs index ([llength $step_1gpu] != 1 || [llength $step_2gpu] != 2)"
	}

	if {$constrain_devices} {
		subtest {[lindex $step_2gpu 0] == 0 && $step_1gpu == 0} "Verify that if devices are constrained CUDA_VISIBLE_DEVICES start always with 0 in a step"
	} else {
		subtest {[lsearch $step_2gpu $step_1gpu] == -1} "Verify that if devices are NOT constrained, all CUDA_VISIBLE_DEVICES are unique"
	}
}

#
# Test --gpus-per-task option by job step
#
proc test_gpus_per_node_with_gpus_per_task {} {
	global file_in1 file_in2 file_out test_name job_id constrain_devices
	global sbatch srun squeue bin_cat bin_grep number scontrol

	set job_gpus 3
	set step_gpus 2

	make_bash_script $file_in1 "
		$scontrol -dd show job \${SLURM_JOBID}
		$srun --exact -n1 --gpus-per-task=$step_gpus --gpus-per-node=0 --mem=0 $file_in2 &
		$srun --exact -n1 --gpus-per-task=$step_gpus --gpus-per-node=0 --mem=0 $file_in2 &
		$srun --exact -n1 --gpus-per-task=$step_gpus --gpus-per-node=0 --mem=0 $file_in2 &
		wait
		exit 0"

	make_bash_script $file_in2 "
		echo 'STEP_ID:'\$SLURM_STEP_ID 'CUDA_VISIBLE_DEVICES:'\$CUDA_VISIBLE_DEVICES
		sleep 3
		if \[ \$SLURM_STEP_ID -eq 2 \]; then
			$scontrol show step \$SLURM_JOB_ID.\$SLURM_STEP_ID
		fi
		exit 0"

	file delete $file_out
	set job_id [submit_job -fail "--cpus-per-gpu=1 --gpus-per-node=$job_gpus -N1 -n3 -t1 -o $file_out -J $test_name $file_in1"]

	wait_for_job -fail $job_id "DONE"
	wait_for_file -fail $file_out

	set match 0
	set bad_cuda_cnt 0
	set output [run_command_output -fail "$bin_cat $file_out"]
	subtest {![regexp "STEP_ID:($number) CUDA_VISIBLE_DEVICES:($number),($number)," $output]} "Verify that any step has more than 2 GPUs"
	subtest {[regexp -all "STEP_ID:($number) CUDA_VISIBLE_DEVICES:($number),($number)" $output] == 3} "Verify that all steps has 2 GPUs"
	subtest {[regexp -all "step creation temporarily disabled" $output] == 2} "Verify that 2 steps were delayed"
}

#
# Test --gpus option by job step
#
proc test_gpus_per_node_with_gpus {} {
	global file_in1 file_in2 file_out test_name job_id constrain_devices
	global sbatch srun squeue bin_cat bin_grep number scontrol

	set job_gpus 2
	set step_gpus 2

	make_bash_script $file_in1 "
		$scontrol -dd show job \${SLURM_JOBID}
		$srun --exact -n2 --gpus=$step_gpus --gpus-per-node=0 --mem=0 $file_in2 &
		$srun --exact -n2 --gpus=$step_gpus --gpus-per-node=0 --mem=0 $file_in2 &
		$srun --exact -n2 --gpus=$step_gpus --gpus-per-node=0 --mem=0 $file_in2 &
		wait
		exit 0"

	make_bash_script $file_in2 "
		echo 'HOST:'\$SLURMD_NODENAME 'NODE_ID:'\$SLURM_NODEID 'STEP_ID:'\$SLURM_STEP_ID 'CUDA_VISIBLE_DEVICES:'\$CUDA_VISIBLE_DEVICES
		sleep 3
		if \[ \$SLURM_STEP_ID -eq 2 -a \$SLURM_NODEID -eq 0 \]; then
			$scontrol show step \$SLURM_JOB_ID.\$SLURM_STEP_ID
		fi
		exit 0"

	file delete $file_out
	set job_id [submit_job -fail "--cpus-per-gpu=2 --gpus-per-node=$job_gpus -N2 -n6 -t1 -o $file_out -J $test_name $file_in1"]

	wait_for_job -fail $job_id "DONE"
	wait_for_file -fail $file_out

	set output [run_command_output -fail "$bin_cat $file_out"]

	subtest {![regexp "STEP_ID:($number) CUDA_VISIBLE_DEVICES:($number)," $output]} "Verify that no more that 1 GPU is visible (per node)"
	subtest {[regexp -all "STEP_ID:0 CUDA_VISIBLE_DEVICES:($number)" $output] == 2} "Verify that step 0 had access to 2 GPUs"
	subtest {[regexp -all "STEP_ID:1 CUDA_VISIBLE_DEVICES:($number)" $output] == 2} "Verify that step 1 had access to 2 GPUs"
	subtest {[regexp -all "STEP_ID:2 CUDA_VISIBLE_DEVICES:($number)" $output] == 2} "Verify that step 2 had access to 2 GPUs"
	subtest {[regexp -all "step creation temporarily disabled" $output] == 1} "Verify that one step was delayed"

	if {$constrain_devices} {
		subtest {[regexp -all "STEP_ID:($number) CUDA_VISIBLE_DEVICES:0" $output] == 6} "Verify that all GPUs are CUDA_VISIBLE_DEVICES:0 due ConstrainDevices"
	} else {
		array set cuda_val {}
		regexp "STEP_ID:0 CUDA_VISIBLE_DEVICES:($number)" $output - cuda_val(0)
		regexp "STEP_ID:1 CUDA_VISIBLE_DEVICES:($number)" $output - cuda_val(1)
		regexp "STEP_ID:2 CUDA_VISIBLE_DEVICES:($number)" $output - cuda_val(2)
		subtest {$cuda_val(0) != $cuda_val(1)} "Verify that two first steps use different GPUs (without ConstrainDevices)"
		subtest {$cuda_val(2) == $cuda_val(0) || $cuda_val(2) == $cuda_val(1)} "Verify that the last step use one of the previous GPUs (without ConstrainDevices)"
	}
}

#
# Test --gpus option by job step
#
proc test_gpus_per_node_with_gpus_2_nodes {} {
	global file_in1 file_in2 file_out test_name job_id constrain_devices
	global sbatch srun squeue bin_cat bin_grep number scontrol

	set job_gpus 4

	make_bash_script $file_in1 "
		$scontrol -dd show job \${SLURM_JOBID}
		$srun --exact -n2 --gpus=6 --gpus-per-node=0 --mem=0 $file_in2 &
		$srun --exact -n2 --gpus=7 --gpus-per-node=0 --mem=0 $file_in2 &
		$srun --exact -n2 --gpus=8 --gpus-per-node=0 --mem=0 $file_in2 &
		wait
		exit 0"

	make_bash_script $file_in2 "
		echo 'HOST:'\$SLURMD_NODENAME 'NODE_ID:'\$SLURM_NODEID 'STEP_ID:'\$SLURM_STEP_ID 'CUDA_VISIBLE_DEVICES:'\$CUDA_VISIBLE_DEVICES
		sleep 3
		if \[ \$SLURM_STEP_ID -eq 2 -a \$SLURM_NODEID -eq 0 \]; then
			$scontrol show step \$SLURM_JOB_ID.\$SLURM_STEP_ID
		fi
		exit 0"

	file delete $file_out
	set job_id [submit_job -fail "--cpus-per-gpu=2 --gpus-per-node=$job_gpus -N2 -n6 -t1 -o $file_out -J $test_name $file_in1"]

	wait_for_job -fail $job_id "DONE"
	wait_for_file -fail $file_out

	set output [run_command -fail "$bin_cat $file_out"]

	subtest {![regexp "STEP_ID:($number) CUDA_VISIBLE_DEVICES:($number),($number),($number),($number)," $output]} "Verify that all steps has less than 5 GPUs per node"
	subtest {[regexp -all "STEP_ID:0 CUDA_VISIBLE_DEVICES:($number)" $output] == 2} "Verify that step 0 used 2 nodes"
	subtest {[regexp -all "STEP_ID:1 CUDA_VISIBLE_DEVICES:($number)" $output] == 2} "Verify that step 1 used 2 nodes"
	subtest {[regexp -all "STEP_ID:2 CUDA_VISIBLE_DEVICES:($number)" $output] == 2} "Verify that step 2 used 2 nodes"
	subtest {[regexp -all "step creation temporarily disabled" $output] == 2} "Verify that two steps were delayed"
	# TODO: Verify that each step got the right number of GPUs

}

#
# Test --gpus-per-task option by job step
#
proc test_gpus_per_node_with_gpus_per_task_3 {} {
	global file_in1 file_in2 file_out test_name job_id constrain_devices
	global sbatch srun squeue bin_cat bin_grep number scontrol

	set job_gpus 4
	set step_gpus 2

	make_bash_script $file_in1 "
		$scontrol -dd show job \${SLURM_JOBID}
		$srun $file_in2
		$srun --exact -n3 --gpus-per-task=$step_gpus --gpus-per-node=0 --mem=0 $file_in2 &
		$srun --exact -n3 --gpus-per-task=$step_gpus --gpus-per-node=0 --mem=0 $file_in2 &
		wait
		exit 0"

	make_bash_script $file_in2 "
		echo 'STEP_ID:'\$SLURM_STEP_ID 'CUDA_VISIBLE_DEVICES:'\$CUDA_VISIBLE_DEVICES ' '
		sleep 3
		if \[ \$SLURM_STEP_ID -eq 1 -a \$SLURM_PROCID -eq 0 \]; then
			$scontrol show step \$SLURM_JOB_ID.\$SLURM_STEP_ID
		fi
		exit 0"

	file delete $file_out
	set job_id [submit_job -fail "--cpus-per-gpu=1 --gpus-per-node=$job_gpus -N2 -n4 -t1 -o $file_out -J $test_name $file_in1"]

	wait_for_job -fail $job_id "DONE"
	wait_for_file -fail $file_out

	set output [run_command_output -fail "$bin_cat $file_out"]
	subtest {![regexp "STEP_ID:($number) CUDA_VISIBLE_DEVICES:($number),($number),($number),($number)," $output]} "Verify that no more that 4 GPUs are visible in any step"
	subtest {[regexp -all "STEP_ID:(0) CUDA_VISIBLE_DEVICES:($number),($number),($number),($number)" $output] == 4} "Verify that job has access to 4 GPUs"
	subtest {[regexp -all "STEP_ID:(1) CUDA_VISIBLE_DEVICES:($number),($number)" $output] == 3} "Verify that step 1 has 3 tasks and 2 GPUs per task"
	subtest {[regexp -all "STEP_ID:(2) CUDA_VISIBLE_DEVICES:($number),($number)" $output] == 3} "Verify that step 2 has 3 tasks and 2 GPUs per task"
	subtest {[regexp -all "step creation temporarily disabled" $output] == 1} "Verify that one step was delayed"
}

#
# Test --gpus-per-task option by job step
#
proc test_gpus_per_node_with_gpus_per_task_5 {} {
	global file_in1 file_in2 file_out test_name job_id constrain_devices
	global sbatch srun squeue bin_cat bin_grep number scontrol

	set job_gpus 4
	set step_gpus 2

	make_bash_script $file_in1 "
		$scontrol -dd show job \${SLURM_JOBID}
		$srun --exact -N1 -n1 --gpus-per-task=$step_gpus --gpus-per-node=0 --mem=0 $file_in2 &
		$srun --exact -N1 -n1 --gpus-per-task=$step_gpus --gpus-per-node=0 --mem=0 $file_in2 &
		$srun --exact -N1 -n1 --gpus-per-task=$step_gpus --gpus-per-node=0 --mem=0 $file_in2 &
		$srun --exact -N1 -n1 --gpus-per-task=$step_gpus --gpus-per-node=0 --mem=0 $file_in2 &
		$srun --exact -N1 -n1 --gpus-per-task=$step_gpus --gpus-per-node=0 --mem=0 $file_in2 &
		wait
		exit 0"

	make_bash_script $file_in2 "
		echo 'STEP_ID:'\$SLURM_STEP_ID 'CUDA_VISIBLE_DEVICES:'\$CUDA_VISIBLE_DEVICES ' '
		sleep 3
		if \[ \$SLURM_STEP_ID -eq 1 -a \$SLURM_PROCID -eq 0 \]; then
			$scontrol show step \$SLURM_JOB_ID.\$SLURM_STEP_ID
		fi
		exit 0"

	file delete $file_out
	set job_id [submit_job -fail "--cpus-per-gpu=1 --gpus-per-node=$job_gpus -N2 -n5 -t1 -o $file_out -J $test_name $file_in1"]

	wait_for_job -fail $job_id "DONE"
	wait_for_file -fail $file_out

	set output [run_command_output -fail "$bin_cat $file_out"]
	subtest {![regexp "STEP_ID:($number) CUDA_VISIBLE_DEVICES:($number),($number)," $output]} "Verify that no more that 2 GPUs are visible in any step"
	subtest {[regexp -all "STEP_ID:($number) CUDA_VISIBLE_DEVICES:($number),($number)" $output] == 5} "Verify that all 5 steps has access to 2 GPUs"
	subtest {[regexp -all "step creation temporarily disabled" $output] == 1} "Verify that one step was delayed"
}

# Test calls
if {[llength [get_nodes_by_request "--cpus-per-gpu=1 --gpus-per-node=2 -N1 -n3"]] != 1} {
	skip_following_testprocs "Testproc needs at least --cpus-per-gpu=1 --gpus-per-node=2 -N1 -n3"
}
testproc_alias "test_gpus_per_node_parallel_1_delayed" test_gpus_per_node_parallel_1_delayed
testproc_alias "test_gpus_per_node_parallel_per task"  test_gpus_per_node_parallel "-n1 --gpus-per-task=1"
testproc_alias "test_gpus_per_node_parallel_gpus"      test_gpus_per_node_parallel "-n1 --gpus=1"
run_following_testprocs

if {[llength [get_nodes_by_request "--cpus-per-gpu=1 --gpus-per-node=3 -N1 -n3"]] != 1} {
	skip_following_testprocs "Testproc needs at least --cpus-per-gpu=1 --gpus-per-node=3 -N1 -n3"
}
testproc_alias "test_gpus_per_node_different_gpus"     test_gpus_per_node_different_gpus
testproc_alias "test_gpus_per_node_with_gpus_per_task" test_gpus_per_node_with_gpus_per_task
run_following_testprocs

if {[llength [get_nodes_by_request "--cpus-per-gpu=2 --gpus-per-node=4 -N2 -n6"]] != 2} {
	skip_following_testprocs "Testproc needs at least --cpus-per-gpu=2 --gpus-per-node=4 -N2 -n6"
}
testproc_alias "test_gpus_per_node_with_gpus"         test_gpus_per_node_with_gpus
testproc_alias "test_gpus_per_node_with_gpus_2_nodes" test_gpus_per_node_with_gpus_2_nodes
run_following_testprocs

if {[llength [get_nodes_by_request "--cpus-per-gpu=1 --gpus-per-node=4 -N2 -n5"]] != 2} {
	skip_following_testprocs "Testproc needs at least --cpus-per-gpu=1 --gpus-per-node=4 -N2 -n5"
}
testproc_alias "test_gpus_per_node_with_gpus_per_task_3" test_gpus_per_node_with_gpus_per_task_3
testproc_alias "test_gpus_per_node_with_gpus_per_task_5" test_gpus_per_node_with_gpus_per_task_5
run_following_testprocs
