shell: allow taskmaps to redistribute tasks amongst nodes

Problem: taskmap_check() now allows an updated taskmap to change
the counts of tasks per node, if total task and node count remains
unchanged, but the affected rcalc data in the shell is not updated,
leading to a possible mismatch between various shell components.

Add a rcalc_update_map() function which can be passed a new taskmap
and updates ntasks on each shell rank internally.

Also update the rankinfo object for the current shell after
rcalc_update_map() returns.

This should now allow a custom taskmap to redistribute task counts
amongst the nodes of a job.

Author: grondo

src/shell/info.c

@@ -263,6 +263,18 @@ int shell_info_set_taskmap (struct shell_info *info,
     if (!(taskids = taskmap_taskids (map, info->shell_rank))
         || !(copy = idset_copy (taskids)))
         return -1;
+
+    /* Update rcalc using the new taskmap.
+     * Be sure to update the local rankinfo structure, which may have changed.
+     */
+    if (rcalc_update_map (info->rcalc, map) < 0
+        || rcalc_get_nth (info->rcalc,
+                          info->shell_rank,
+                          &info->rankinfo) < 0) {
+        shell_log_error ("unable to update task counts from new taskmap");
+        return -1;
+    }
+
     idset_destroy (info->taskids);
     info->taskids = copy;
     taskmap_destroy (info->taskmap);

src/shell/rcalc.c

@@ -16,6 +16,7 @@
 #include <unistd.h>
 #include <errno.h>
 #include <jansson.h>
+#include <flux/taskmap.h>
 
 #include "src/common/libczmqcontainers/czmq_containers.h"
 #include "src/common/libidset/idset.h"
@@ -52,6 +53,22 @@ struct rcalc {
     struct allocinfo *alloc;
 };
 
+int rcalc_update_map (rcalc_t *r, struct taskmap *map)
+{
+    if (r == NULL || map == NULL) {
+        errno = EINVAL;
+        return -1;
+    }
+    /* Update task counts for each rank in rcalc based on new taskmap.
+     * N.B.: ignores alloc[i]->ncores_avail since this variable is known
+     * to not be used after the initial distribution of tasks.
+     */
+    for (int i = 0; i < r->nranks; i++)
+        r->alloc[i].ntasks = taskmap_ntasks (map, i);
+
+    return 0;
+}
+
 void rcalc_destroy (rcalc_t *r)
 {
     if (r == NULL)

src/shell/rcalc.h

@@ -14,6 +14,7 @@
 #include <sched.h>
 #include <jansson.h>
 #include <flux/core.h>
+#include <flux/taskmap.h>
 
 typedef struct rcalc rcalc_t;
 
@@ -61,4 +62,8 @@ int rcalc_get_rankinfo (rcalc_t *r, int rank, struct rcalc_rankinfo *ri);
 /*  Fill in rcalc_rankinfo for the nth rank in the rcalc_t list */
 int rcalc_get_nth (rcalc_t *r, int id, struct rcalc_rankinfo *ri);
 
+/* Update rcalc with a new taskmap
+ */
+int rcalc_update_map (rcalc_t *r, struct taskmap *map);
+
 #endif /* !SHELL_RCALC_H */