Merge pull request #5756 from grondo/shell-broker_rank

shell: fix incorrect values returned from `flux_shell_get_rank_info(3)`

Author: mergify[bot]

src/shell/shell.c

@@ -530,6 +530,7 @@ static json_t *flux_shell_get_rank_info_object (flux_shell_t *shell, int rank)
     char key [128];
     char *taskids = NULL;
     struct taskmap *map;
+    struct rcalc_rankinfo rankinfo;
 
     if (!shell->info)
         return NULL;
@@ -551,13 +552,16 @@ static json_t *flux_shell_get_rank_info_object (flux_shell_t *shell, int rank)
     if (!(taskids = get_rank_task_idset (map, rank)))
         return NULL;
 
+    if (rcalc_get_nth (shell->info->rcalc, rank, &rankinfo) < 0)
+        return NULL;
+
     o = json_pack_ex (&error, 0, "{ s:i s:i s:s s:{s:s s:s?}}",
-                   "broker_rank", rank,
+                   "broker_rank", rankinfo.rank,
                    "ntasks", taskmap_ntasks (map, rank),
                    "taskids", taskids,
                    "resources",
-                     "cores", shell->info->rankinfo.cores,
-                     "gpus",  shell->info->rankinfo.gpus);
+                     "cores", rankinfo.cores,
+                     "gpus",  rankinfo.gpus);
     free (taskids);
 
     if (o == NULL)

t/t2603-job-shell-initrc.t

@@ -4,7 +4,7 @@ test_description='Test flux-shell initrc.lua implementation'
 
 . `dirname $0`/sharness.sh
 
-test_under_flux 1
+test_under_flux 4
 
 FLUX_SHELL="${FLUX_BUILD_DIR}/src/shell/flux-shell"
 
@@ -280,4 +280,25 @@ test_expect_success 'flux-shell: initrc: shell.die function works' '
 		> ${name}.log 2>&1 &&
 	grep "FATAL: test: shell.die" ${name}.log
 '
+test_expect_success MULTICORE 'flux-shell: initrc: shell.rankinfo reports broker_rank' '
+	name=shell.rankinfo.broker_rank &&
+	cat >${name}.lua <<-EOF &&
+	ri0 = shell.get_rankinfo(0)
+	ri1 = shell.get_rankinfo(1)
+	if ri0.broker_rank ~= 1 or ri1.broker_rank ~= 3 then
+	  shell.log ("rankinfo(0).broker_rank = "
+	             ..shell.get_rankinfo(0).broker_rank)
+	  shell.log ("rankinfo(1).broker_rank = "
+	             ..shell.get_rankinfo(1).broker_rank)
+	  shell.die ("rankinfo.broker_rank incorrect!")
+	end
+	if ri0.ntasks ~= 2 or ri1.ntasks ~= 1 then
+	  shell.die ("got ri[0].ntasks = "..ri0.ntasks
+	             .." ri[1].ntasks = "..ri1.ntasks)
+	end
+	EOF
+	flux run -N2 -n3 --requires=rank:1,3 \
+	    -o verbose -o initrc=${name}.lua true
+'
+flux job info $(flux job last) R
 test_done