Skip to content

Commit ed07466

Browse files
mergify[bot]mergify[bot]
authored
Merge pull request #5844 from grondo/taskmap-hostfile
add `--taskmap=hostfile:FILE` support
2 parents 57ac53a + f31c26b commit ed07466

File tree

9 files changed

+332
-8
lines changed

9 files changed

+332
-8
lines changed

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

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,14 @@
1717
the job shell, so this option is not useful unless the total number
1818
of nodes and tasks per node are known at job submission time.
1919

20+
hostfile:FILE
21+
Assign tasks in order to hosts as they appear in FILE. FILE should
22+
have one or more lines each of which contains a host name or RFC
23+
29 Hostlist string. Each host assigned to the job must appear in
24+
the hostfile and be assigned the same number of tasks as the default
25+
taskmap from the shell. If there are less hosts in the hostfile than
26+
tasks in the job, then the list of hosts will be reused.
27+
2028
However, shell plugins may provide other task mapping schemes, so
2129
check the current job shell configuration for a full list of supported
2230
taskmap schemes.

‎doc/man1/flux-shell.rst‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -126,13 +126,13 @@ topics:
126126
**taskmap.SCHEME**
127127
Called when a taskmap scheme *SCHEME* is requested via the taskmap
128128
shell option or corresponding :option:`flux submit --taskmap` option.
129-
Plugins that want to offer a different taskmap scheme than the defaults of
130-
``block``, ``cyclic``, and ``manual`` can register a ``taskmap.*`` plugin
131-
callback and then users can request this mapping with the appropriate
132-
:option:`flux submit --taskmap=name`` option. The default block taskmap is
133-
passed to the plugin as "taskmap" in the plugin input arguments, and the
134-
plugin should return the new taskmap as a string in the output args. This
135-
callback is called before ``shell.init``.
129+
Plugins that want to offer a different taskmap scheme than the defaults
130+
of ``block``, ``cyclic``, ``hostfile``, and ``manual`` can register a
131+
``taskmap.*`` plugin callback and then users can request this mapping
132+
with the appropriate :option:`flux submit --taskmap=name`` option.
133+
The default block taskmap is passed to the plugin as "taskmap" in the
134+
plugin input arguments, and the plugin should return the new taskmap as a
135+
string in the output args. This callback is called before ``shell.init``.
136136

137137
**shell.connect**
138138
Called just after the shell connects to the local Flux broker. (Only

‎src/shell/Makefile.am‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -89,7 +89,8 @@ flux_shell_SOURCES = \
8989
doom.c \
9090
exception.c \
9191
rlimit.c \
92-
cyclic.c \
92+
taskmap/cyclic.c \
93+
taskmap/hostfile.c \
9394
signal.c \
9495
files.c \
9596
oom.c \

‎src/shell/builtins.c‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,7 @@ extern struct shell_builtin builtin_doom;
4949
extern struct shell_builtin builtin_exception;
5050
extern struct shell_builtin builtin_rlimit;
5151
extern struct shell_builtin builtin_cyclic;
52+
extern struct shell_builtin builtin_hostfile;
5253
extern struct shell_builtin builtin_signal;
5354
extern struct shell_builtin builtin_oom;
5455
extern struct shell_builtin builtin_hwloc;
@@ -74,6 +75,7 @@ static struct shell_builtin * builtins [] = {
7475
&builtin_exception,
7576
&builtin_rlimit,
7677
&builtin_cyclic,
78+
&builtin_hostfile,
7779
&builtin_signal,
7880
&builtin_oom,
7981
&builtin_hwloc,

‎src/shell/output.c‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -773,6 +773,7 @@ void shell_output_destroy (struct shell_output *out)
773773
zhash_destroy (&out->fds);
774774
}
775775
eventlogger_destroy (out->ev);
776+
idset_destroy (out->active_shells);
776777
free (out);
777778
errno = saved_errno;
778779
}

src/shell/cyclic.c renamed to src/shell/taskmap/cyclic.c

File renamed without changes.

‎src/shell/taskmap/hostfile.c‎

Lines changed: 204 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,204 @@
1+
/************************************************************\
2+
* Copyright 2024 Lawrence Livermore National Security, LLC
3+
* (c.f. AUTHORS, NOTICE.LLNS, COPYING)
4+
*
5+
* This file is part of the Flux resource manager framework.
6+
* For details, see https://github.com/flux-framework.
7+
*
8+
* SPDX-License-Identifier: LGPL-3.0
9+
\************************************************************/
10+
11+
/* shell taskmap.hostfile plugin
12+
*
13+
* Read a list of hosts from a file, and assign tasks to hosts in order
14+
* they are listed.
15+
*/
16+
#define FLUX_SHELL_PLUGIN_NAME "taskmap.hostfile"
17+
18+
#if HAVE_CONFIG_H
19+
#include "config.h"
20+
#endif
21+
#include <stdio.h>
22+
#include <jansson.h>
23+
24+
#include <flux/core.h>
25+
#include <flux/taskmap.h>
26+
#include <flux/hostlist.h>
27+
28+
#include "src/common/libutil/errprintf.h"
29+
30+
#include "builtins.h"
31+
32+
33+
/* Create a taskmap that represents 'ntasks' tasks mapped across a set
34+
* of hosts in 'nodelist', ordered by hostlist 'hl'.
35+
*/
36+
char *taskmap_hostlist (int ntasks,
37+
struct hostlist *nodelist,
38+
struct hostlist *hl,
39+
flux_error_t *errp)
40+
{
41+
struct taskmap *map = NULL;
42+
char *result = NULL;
43+
const char *host = NULL;
44+
45+
if (!(map = taskmap_create ()))
46+
goto error;
47+
48+
/* Loop through hostlist hl until all tasks have been assigned to hosts
49+
*/
50+
while (ntasks > 0) {
51+
int rank;
52+
if (host == NULL)
53+
host = hostlist_first (hl);
54+
if ((rank = hostlist_find (nodelist, host)) < 0) {
55+
errprintf (errp, "host %s not found in job nodelist", host);
56+
goto error;
57+
}
58+
if (taskmap_append (map, rank, 1, 1) < 0) {
59+
errprintf (errp,
60+
"failed to append task to taskmap: %s",
61+
strerror (errno));
62+
goto error;
63+
}
64+
host = hostlist_next (hl);
65+
ntasks--;
66+
}
67+
result = taskmap_encode (map, TASKMAP_ENCODE_WRAPPED);
68+
error:
69+
taskmap_destroy (map);
70+
return result;
71+
}
72+
73+
static struct hostlist *hostlist_from_file (const char *path)
74+
{
75+
ssize_t n;
76+
size_t size;
77+
struct hostlist *hl = NULL;
78+
FILE *fp = NULL;
79+
char *line = NULL;
80+
81+
if (!(fp = fopen (path, "r"))) {
82+
shell_log_errno ("failed to open hostfile: %s", path);
83+
goto error;
84+
}
85+
if (!(hl = hostlist_create ())) {
86+
shell_log_errno ("failed to create hostlist");
87+
goto error;
88+
}
89+
while ((n = getline (&line, &size, fp)) != -1) {
90+
int len = strlen (line);
91+
if (line[len-1] == '\n')
92+
line[len-1] = '\0';
93+
if (strlen (line) > 0 && hostlist_append (hl, line) < 0) {
94+
shell_log_errno ("hostlist_append: %s", line);
95+
}
96+
}
97+
error:
98+
if (fp)
99+
fclose (fp);
100+
free (line);
101+
return hl;
102+
}
103+
104+
static struct hostlist *hostlist_from_R (flux_shell_t *shell)
105+
{
106+
size_t i;
107+
json_t *nodelist;
108+
json_t *val;
109+
struct hostlist *hl = NULL;
110+
111+
if (flux_shell_info_unpack (shell,
112+
"{s:{s:{s:o}}}",
113+
"R",
114+
"execution",
115+
"nodelist", &nodelist) < 0) {
116+
shell_log_errno ("unable to get job nodelist");
117+
return NULL;
118+
}
119+
if (!(hl = hostlist_create ())) {
120+
shell_log_errno ("hostlist_create");
121+
return NULL;
122+
}
123+
json_array_foreach (nodelist, i, val) {
124+
const char *host = json_string_value (val);
125+
if (!host)
126+
goto error;
127+
if (hostlist_append (hl, host) < 0) {
128+
shell_log_errno ("hostlist_append %s", host);
129+
goto error;
130+
}
131+
}
132+
return hl;
133+
error:
134+
hostlist_destroy (hl);
135+
return NULL;
136+
}
137+
138+
static int map_hostfile (flux_plugin_t *p,
139+
const char *topic,
140+
flux_plugin_arg_t *args,
141+
void *data)
142+
{
143+
flux_shell_t *shell;
144+
int rc = -1;
145+
const char *value = NULL;
146+
struct hostlist *hl = NULL;
147+
struct hostlist *nodelist = NULL;
148+
char *map = NULL;
149+
int ntasks;
150+
flux_error_t error;
151+
152+
if (!(shell = flux_plugin_get_shell (p)))
153+
return -1;
154+
155+
if (flux_plugin_arg_unpack (args,
156+
FLUX_PLUGIN_ARG_IN,
157+
"{s?s}",
158+
"value", &value) < 0) {
159+
shell_log_error ("unpack: %s", flux_plugin_arg_strerror (args));
160+
return -1;
161+
}
162+
163+
if (!(hl = hostlist_from_file (value))
164+
|| !(nodelist = hostlist_from_R (shell))) {
165+
shell_log_error ("failed to get hostlists from file and R");
166+
goto out;
167+
}
168+
if ((ntasks = taskmap_total_ntasks (flux_shell_get_taskmap (shell))) < 0)
169+
shell_log_error ("failed to get ntasks from current shell taskmap");
170+
171+
if (!(map = taskmap_hostlist (ntasks, nodelist, hl, &error))) {
172+
shell_log_error ("failed to map tasks with hostfile:%s: %s",
173+
value,
174+
error.text);
175+
goto out;
176+
}
177+
if (flux_plugin_arg_pack (args,
178+
FLUX_PLUGIN_ARG_OUT,
179+
"{s:s}",
180+
"taskmap", map) < 0) {
181+
shell_log_error ("failed to set new taskmap in plugin output args");
182+
goto out;
183+
}
184+
rc = 0;
185+
out:
186+
free (map);
187+
hostlist_destroy (hl);
188+
hostlist_destroy (nodelist);
189+
return rc;
190+
}
191+
192+
static int plugin_init (flux_plugin_t *p)
193+
{
194+
return flux_plugin_add_handler (p, "taskmap.hostfile", map_hostfile, NULL);
195+
}
196+
197+
struct shell_builtin builtin_hostfile = {
198+
.name = FLUX_SHELL_PLUGIN_NAME,
199+
.plugin_init = plugin_init,
200+
};
201+
202+
/*
203+
* vi:tabstop=4 shiftwidth=4 expandtab
204+
*/

‎t/Makefile.am‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -208,6 +208,7 @@ TESTSCRIPTS = \
208208
t2614-job-shell-doom.t \
209209
t2615-job-shell-rlimit.t \
210210
t2616-job-shell-taskmap.t \
211+
t2616-job-shell-taskmap-hostfile.t \
211212
t2617-job-shell-stage-in.t \
212213
t2618-job-shell-signal.t \
213214
t2619-job-shell-hwloc.t \

‎t/t2616-job-shell-taskmap-hostfile.t‎

Lines changed: 107 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,107 @@
1+
#!/bin/sh
2+
#
3+
test_description='Test hostfile taskmap plugin support'
4+
5+
. `dirname $0`/sharness.sh
6+
7+
# Use "system" personality to get fake hostnames for hostfile use
8+
test_under_flux 4 system
9+
10+
# Test that actual task ranks match expected ranks.
11+
# Assumes job output is `echo $FLUX_TASK_RANK: $(flux getattr rank)`
12+
test_check_taskmap() {
13+
local id=$1
14+
flux job attach $id | sort -n >$id.output &&
15+
flux job taskmap --to=multiline $id >$id.expected &&
16+
test_cmp $id.expected $id.output
17+
}
18+
19+
test_expect_success 'create script for testing task mapping' '
20+
cat <<-EOF >map.sh &&
21+
#!/bin/sh
22+
echo \$FLUX_TASK_RANK: \$(flux getattr rank)
23+
EOF
24+
chmod +x map.sh
25+
'
26+
test_expect_success 'taskmap=hostfile works' '
27+
cat <<-EOF >h1 &&
28+
fake3
29+
fake2
30+
fake1
31+
fake0
32+
EOF
33+
expected="[[3,1,1,1],[2,1,1,1],[1,1,1,1],[0,1,1,1]]" &&
34+
id=$(flux submit --taskmap=hostfile:h1 -N4 -n4 ./map.sh) &&
35+
flux job attach -vEX $id &&
36+
flux job wait-event -p exec -f json $id shell.start &&
37+
flux job wait-event -p exec -f json $id shell.start \
38+
| jq -e ".context.taskmap.map == $expected" &&
39+
test_check_taskmap $id
40+
'
41+
test_expect_success 'taskmap=hostfile works with multiple tasks per node' '
42+
cat <<-EOF >h2 &&
43+
fake3
44+
fake3
45+
fake2
46+
fake2
47+
fake1
48+
fake1
49+
fake0
50+
fake0
51+
EOF
52+
expected="[[3,1,2,1],[2,1,2,1],[1,1,2,1],[0,1,2,1]]" &&
53+
id=$(flux submit --taskmap=hostfile:h2 -N4 --tasks-per-node=2 ./map.sh) &&
54+
flux job attach -vEX $id &&
55+
flux job wait-event -p exec -f json $id shell.start &&
56+
flux job wait-event -p exec -f json $id shell.start \
57+
| jq -e ".context.taskmap.map == $expected" &&
58+
test_check_taskmap $id
59+
'
60+
test_expect_success 'taskmap=hostfile reuses hosts in short hostlist' '
61+
cat <<-EOF >h3 &&
62+
fake3
63+
fake2
64+
fake1
65+
fake0
66+
EOF
67+
expected="[[3,1,1,1],[2,1,1,1],[1,1,1,1],[0,1,1,1],[3,1,1,1],[2,1,1,1],[1,1,1,1],[0,1,1,1]]" &&
68+
id=$(flux submit --taskmap=hostfile:h3 -N4 --tasks-per-node=2 ./map.sh) &&
69+
flux job attach -vEX $id &&
70+
flux job wait-event -p exec -f json $id shell.start &&
71+
flux job wait-event -p exec -f json $id shell.start \
72+
| jq -e ".context.taskmap.map == $expected" &&
73+
test_check_taskmap $id
74+
'
75+
test_expect_success 'taskmap=hostfile works with hostlists' '
76+
cat <<-EOF >h4 &&
77+
fake[1,2]
78+
fake[3,0]
79+
EOF
80+
expected="[[1,3,1,1],[0,4,1,1],[0,1,1,1]]" &&
81+
id=$(flux submit --taskmap=hostfile:h4 -N4 --tasks-per-node=2 ./map.sh) &&
82+
flux job attach -vEX $id &&
83+
flux job wait-event -p exec -f json $id shell.start &&
84+
flux job wait-event -p exec -f json $id shell.start \
85+
| jq -e ".context.taskmap.map == $expected" &&
86+
test_check_taskmap $id
87+
'
88+
test_expect_success 'taskmap=hostfile fails with invalid hostlist' '
89+
echo "fake[0-3">h5 &&
90+
test_must_fail_or_be_terminated \
91+
flux run --taskmap=hostfile:h5 -N4 hostname
92+
'
93+
test_expect_success 'taskmap=hostfile fails with incorrect hosts' '
94+
echo "foo[0-3]">h6 &&
95+
test_must_fail_or_be_terminated \
96+
flux run --taskmap=hostfile:h6 -N4 hostname
97+
'
98+
test_expect_success 'taskmap=hostfile fails when not all hosts present' '
99+
echo "foo[0,0,1,2]">h7 &&
100+
test_must_fail_or_be_terminated \
101+
flux run --taskmap=hostfile:h7 -N4 hostname
102+
'
103+
test_expect_success 'taskmap=hostfile fails with invalid filename' '
104+
test_must_fail_or_be_terminated \
105+
flux run --taskmap=hostfile:badfile -N4 hostname
106+
'
107+
test_done

0 commit comments

Comments
 (0)