Merge pull request #6815 from grondo/issue#6810

remove restriction that taskmaps cannot assign differing counts of tasks to nodes in a job

Author: mergify[bot]

doc/man1/common/job-other-run.rst

@@ -13,17 +13,15 @@
    manual:TASKMAP
     An explicit RFC 34 taskmap is provided and used to manually map
     task ids to nodes. The provided *TASKMAP* must match the total number
-    of tasks in the job and the number of tasks per node assigned by
-    the job shell, so this option is not useful unless the total number
-    of nodes and tasks per node are known at job submission time.
+    of nodes and tasks in the job, so this option is not useful unless
+    the total number of nodes are known at job submission time.
 
    hostfile:FILE
     Assign tasks in order to hosts as they appear in FILE. FILE should
     have one or more lines each of which contains a host name or RFC
     29 Hostlist string. Each host assigned to the job must appear in
-    the hostfile and be assigned the same number of tasks as the default
-    taskmap from the shell.  If there are less hosts in the hostfile than
-    tasks in the job, then the list of hosts will be reused.
+    the hostfile.  If there are less hosts in the hostfile than tasks
+    in the job, then the list of hosts will be reused.
 
    However, shell plugins may provide other task mapping schemes, so
    check the current job shell configuration for a full list of supported

src/common/libtaskmap/taskmap.c

@@ -973,16 +973,6 @@ int taskmap_check (const struct taskmap *old,
                           "got %d total tasks, expected %d",
                           ntasks_new,
                           ntasks_old);
-    for (int i = 0; i < nnodes_old; i++) {
-        ntasks_old = taskmap_ntasks (old, i);
-        ntasks_new = taskmap_ntasks (new, i);
-        if (ntasks_old != ntasks_new)
-            return errprintf (errp,
-                              "node %d has %d tasks, expected %d",
-                              i,
-                              ntasks_new,
-                              ntasks_old);
-    }
     return 0;
 }
 

src/common/libtaskmap/taskmap.h

@@ -90,8 +90,8 @@ int taskmap_nnodes (const struct taskmap *map);
 int taskmap_total_ntasks (const struct taskmap *map);
 
 /*  Check if 'a' and 'b' taskmaps are compatible, i.e. they
- *  have equivalent numbers of total tasks, total nodes, and tasks assigned
- *  to each node.
+ *  have equivalent numbers of total tasks and total nodes.
+ *
  *  Returns 0 on success, -1 with error message in 'errp' on failure.
  */
 int taskmap_check (const struct taskmap *a,

src/common/libtaskmap/test/taskmap.c

@@ -522,12 +522,11 @@ struct check_test check_tests[] = {
     { "[[0,4,4,1]]", "[[0,4,1,4]]", 0, NULL },
     { "[[0,4,4,1]]", "[[0,4,2,2]]", 0, NULL },
     { "[[0,4,4,1]]", "[[0,4,3,1],[0,4,1,1]]", 0, NULL },
+    { "[[0,4,4,1]]", "[[0,2,4,1],[2,1,3,1],[3,1,5,1]]", 0, NULL },
     { "[[0,4,4,1]]", "[[0,4,3,1]]", -1,
       "got 12 total tasks, expected 16" },
     { "[[0,4,4,1]]", "[[0,2,8,1]]", -1,
       "got 2 nodes, expected 4" },
-    { "[[0,4,4,1]]", "[[0,2,4,1],[2,1,3,1],[3,1,5,1]]", -1,
-      "node 2 has 3 tasks, expected 4" },
     { NULL, NULL, 0, NULL },
 };
 

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 */

t/t2616-job-shell-taskmap.t

@@ -136,6 +136,15 @@ test_expect_success 'taskmap=manual works' '
 		| jq -e ".context.taskmap.map == [[1,3,4,1],[0,1,4,1]]" &&
 	test_check_taskmap $id
 '
+test_expect_success 'taskmap=manual can redistribute tasks amongst nodes' '
+	# put 2 extra tasks on rank 0
+	id=$(flux submit \
+	     --taskmap=manual:"[[0,1,3,1],[1,3,1,1]]" \
+	     -N4 -n6 ./map.sh) &&
+	flux job wait-event -p exec -f json $id shell.start \
+		| jq -e ".context.taskmap.map == [[0,1,3,1],[1,3,1,1]]" &&
+	test_check_taskmap $id
+'
 test_expect_success 'taskmap=manual with unknown taskmap fails' '
 	test_must_fail_or_be_terminated \
 		flux run --taskmap="manual:[]" true
@@ -149,8 +158,7 @@ test_expect_success 'invalid manual taskmaps fail predictably' '
 		-N4 --tasks-per-node=4 true ::: \
 		"{}" \
 		"[[1,3,4,1]]" \
-		"[[1,3,4,1],[0,1,3,1]]" \
-		"[[3,1,1,1],[0,3,5,1]]") &&
+		"[[1,3,4,1],[0,1,3,1]]") &&
 	for id in $ids; do
 		test_must_fail_or_be_terminated flux job attach $id
 	done