Merge pull request #256 from avildema/int3_env_iodemo

TEST: Add env key to iodemo.sh
yosefe · Apr 5, 2022 · 4b273b0 · 4b273b0
2 parents c438ce4 + 148b3ce
commit 4b273b0
Show file tree

Hide file tree

Showing 3 changed files with 74 additions and 26 deletions.
diff --git a/buildlib/pr/io_demo/az-stage-io-demo.yaml b/buildlib/pr/io_demo/az-stage-io-demo.yaml
@@ -13,6 +13,8 @@ parameters:
   default: 20
 - name: analyzer_allow_list_args
   default: ''
+- name: extra_run_args
+  default: ''
 
 steps:
 - bash: |
@@ -62,6 +64,7 @@ steps:
         --map-by slot \
         --log-dir $(workspace)/${{ parameters.name }} \
         -i ${{ parameters.roce_iface }} \
+        ${{ parameters.extra_run_args }} \
         $(io_demo_exe) \
             -d 512:524288 \
             -P 2 \

diff --git a/buildlib/pr/io_demo/io-demo.yml b/buildlib/pr/io_demo/io-demo.yml
@@ -92,6 +92,7 @@ jobs:
             test_time: ${{ test.Value.duration }}
             test_intefrafe: ${{ test.Value.interface }}
             test_ucx_tls: ${{ test.Value.tls }}
+            test_extra_run_args: ${{ test.Value.extra_run_args }}
             initial_delay: ${{ coalesce(test.Value.initial_delay, parameters.initial_delay) }}
       maxParallel: 1
 
@@ -120,5 +121,6 @@ jobs:
           roce_iface: $(test_intefrafe)
           iodemo_tls: $(test_ucx_tls)
           initial_delay: $(initial_delay)
+          extra_run_args: $(test_extra_run_args)
           ${{ if eq(variables['Build.Reason'], 'PullRequest') }}:
             analyzer_allow_list_args: '--allow_list $(System.PullRequest.TargetBranch)'
diff --git a/test/apps/iodemo/run_io_demo.sh b/test/apps/iodemo/run_io_demo.sh
@@ -1,4 +1,4 @@
-#!/bin/sh
+#!/bin/bash
 #
 # Copyright (C) Mellanox Technologies Ltd. 2020.  ALL RIGHTS RESERVED.
 #
@@ -81,19 +81,23 @@ check_slurm_env()
 init_config()
 {
 	verbose=0
-	iodemo_exe=""
-	iodemo_client_args=""
+	net_if="bond0"
+	duration=30
+	bind_local=0
+	base_port_num=20000
 	num_clients=1
 	num_servers=1
-	map_by="node"
-	base_port_num=20000
 	tasks_per_node=1
-	net_if="bond0"
-	duration=30
+	map_by="node"
 	client_wait_time=2
 	launcher="pdsh -b -w"
 	dry_run=0
 	log_dir="$PWD"
+	iodemo_exe=""
+	iodemo_client_args=""
+	extra_env_client=""
+	extra_env_server=""
+	extra_env_all=""
 
 	# command line args will override slurm env vars
 	check_slurm_env
@@ -104,12 +108,10 @@ show_config()
 {
 	echo "Launch configuration:"
 	for key in \
-			host_list tasks_per_node map_by \
-			num_clients num_servers \
-			iodemo_exe iodemo_client_args \
-			net_if base_port_num \
-			duration client_wait_time \
-			launcher dry_run log_dir
+			host_list net_if duration bind_local base_port_num \
+			num_clients num_servers tasks_per_node map_by \
+			client_wait_time launcher dry_run log_dir \
+			iodemo_exe iodemo_client_args
 	do
 		show_var ${key}
 	done
@@ -132,6 +134,9 @@ usage()
 	echo "  -H|--hostlist <h1>,<h2>,..  List of host names to run on"$(show_default_value host_list)
 	echo "  -i|--netif <n>              Network interface to use"$(show_default_value net_if)
 	echo "  -d|--duration <seconds>     How much time to run the application"$(show_default_value duration)
+	echo "  --bind                      Bind to local IP address"
+	echo "  --env <role> <key>=<value>  Environment variable for <role> (client/server/all)"
+	echo "  --port-num <number>         TCP port number to start from"$(show_default_value base_port_num)
 	echo "  --num-clients <count>       Number of clients to run"$(show_default_value num_clients)
 	echo "  --num-servers <count>       Number of servers to run"$(show_default_value num_servers)
 	echo "  --tasks-per-node <count>    Maximal number of tasks per node"$(show_default_value tasks_per_node)
@@ -169,14 +174,32 @@ parse_args()
 			net_if="$2"
 			shift
 			;;
-		--log-dir)
-			log_dir="$2"
-			shift
-			;;
 		-d|--duration)
 			duration="$2"
 			shift
 			;;
+		--bind)
+			bind_local=1
+			;;
+		--env)
+			role="$2"
+			if [ "$2" == "all" ]
+			then
+				extra_env_all="$extra_env_all $3"
+			elif [ "$2" == "client" ]
+			then
+				extra_env_client="$extra_env_client $3"
+			elif [ "$2" == "server" ]
+			then
+				extra_env_server="$extra_env_server $3"
+			fi
+			shift
+			shift
+			;;
+		--port-num)
+			base_port_num="$2"
+			shift
+			;;
 		--num-clients)
 			num_clients="$2"
 			shift
@@ -204,6 +227,10 @@ parse_args()
 		--dry-run)
 			dry_run=1
 			;;
+		--log-dir)
+			log_dir="$2"
+			shift
+			;;
 		[^-]*)
 			iodemo_exe="$key"
 			shift
@@ -294,7 +321,7 @@ build_server_args_list() {
 			iodemo_server_args+=" $key $value"
 			shift
 			;;
-		-q|-v|-H)
+		-q|-A|-v|-H)
 			iodemo_server_args+=" $key"
 			;;
 		*)
@@ -342,17 +369,17 @@ create_mapping_bynode()
 	# which means no node will have one process less.
 	# The expression "(x + N - 1) % N" yields a number in the range 0..N-1 and
 	# then adding 1 yields the equivalent of "x % N" in the range 1..N.
-	# 
+	#
 	remainder_client_index=$(((num_clients + num_hosts - 1) % num_hosts + 1))
 	remainder_server_index=$(((num_servers + num_hosts - 1) % num_hosts + 1))
 	show_var remainder_client_index
 	show_var remainder_client_index
-	
+
 	host_index=0
 	for host in $(split_list ${host_list})
 	do
 		# Add same amount of clients/servers on each host, except few last hosts
-		# which may have less (if mapping is not balanced) 
+		# which may have less (if mapping is not balanced)
 		num_clients_per_host[${host}]=$((max_clients_per_node - \
 		                                 (host_index >= remainder_client_index)))
 		num_servers_per_host[${host}]=$((max_servers_per_node - \
@@ -373,7 +400,7 @@ create_mapping_byslot()
 		                       ${remaining_clients})
 		num_servers_per_host[${host}]=${node_num_servers}
 		num_clients_per_host[${host}]=${node_num_clients}
-		
+
 		remaining_clients=$((remaining_clients - node_num_clients))
 		remaining_servers=$((remaining_servers - node_num_servers))
 	done
@@ -466,7 +493,7 @@ make_scripts()
 
 		# Add file header and startup
 		cat >${command_file} <<-EOF
-			#!/bin/sh
+			#!/bin/bash
 			#
 			# Launch script for io_demo on ${host} with ${num_servers_per_host[${host}]} servers and ${num_clients_per_host[${host}]} clients
 			#
@@ -601,6 +628,10 @@ make_scripts()
 			EOF
 		env | grep -P '^UCX_.*=|^PATH=|^LD_PRELOAD=|^LD_LIBRARY_PATH=' | \
 			xargs -L 1 echo "     export" >>${command_file}
+			for extra_env in ${extra_env_all}
+			do
+				echo "     export ${extra_env}" >>${command_file}
+			done
 		cat >>${command_file} <<-EOF
 			    cd $PWD
 			}
@@ -617,7 +648,7 @@ make_scripts()
 			cat >>${command_file} <<-EOF
 				function start_server_${i}() {
 				    mkdir -p ${log_dir}
-				    env IODEMO_ROLE=server_${i} ${cmd_prefix} \\
+				    env IODEMO_ROLE=server_${i} ${extra_env_server} ${cmd_prefix} \\
 				        ${iodemo_exe} \\
 				            ${iodemo_server_args} -p ${port_num} \\
 				            ${log_redirect} ${log_file} &
@@ -631,19 +662,31 @@ make_scripts()
 		for ((i=0;i<num_clients_per_host[${host}];++i))
 		do
 			log_file=${log_dir}/$(printf "iodemo_%s_client_%02d.log" ${host} $i)
+			if [ ${bind_local} -eq 1 ]
+			then
+				client_bind="-I ${ip_address_per_host[${host}]}"
+			else
+				client_bind=""
+			fi
 			is_verbose && echo ${log_file}
 			cat >>${command_file} <<-EOF
 				function start_client_${i}() {
 				    mkdir -p ${log_dir}
-				    env IODEMO_ROLE=client_${i} ${cmd_prefix} \\
+				    env IODEMO_ROLE=client_${i} ${extra_env_client} ${cmd_prefix} \\
 				        ${iodemo_exe} \\
-				            ${iodemo_client_args} ${client_connect_list} \\
+				            ${iodemo_client_args} \\
+				            ${client_connect_list} \\
+				            ${client_bind} \\
 				            ${log_redirect} ${log_file} &
 				}
 
 				EOF
 		done
 
+		show_var_verbose extra_env_all
+		show_var_verbose extra_env_client
+		show_var_verbose extra_env_server
+
 		# 'run_all' will start all servers, then clients, then wait for finish
 		cat >>${command_file} <<-EOF
 			start_all_servers() {