Merge pull request #5721 from grondo/issue#5718

shell: write hwloc XML to a file with `HWLOC_XMLFILE` set with `-o hwloc.xmlfile`

Author: mergify[bot]

doc/man1/common/job-shell-options.rst

@@ -36,3 +36,6 @@
    * - :option:`exit-on-error`
      - Raise a fatal job exception immediately if first task exits with
        nonzero exit code.
+
+   * - :option:`hwloc.xmlfile`
+     - Write hwloc XML gathered by job to a file and set ``HWLOC_XMLFILE``

doc/man1/flux-shell.rst

@@ -438,6 +438,10 @@ plugins include:
 
   Send signal :option:`signal.signum` *TIME* seconds before job expiration.
 
+.. option:: hwloc.xmlfile
+
+  Write the job shell's copy of hwloc XML to a file and set ``HWLOC_XMLFILE``.
+
 .. warning::
   The directory referenced by :envvar:`FLUX_JOB_TMPDIR` is cleaned up when the
   job ends, is guaranteed to be unique, and is generally on fast local storage

doc/test/spell.en.pws

@@ -823,3 +823,4 @@ aspirational
 infeasible
 login
 workflow
+xmlfile

src/shell/Makefile.am

@@ -92,7 +92,8 @@ flux_shell_SOURCES = \
 	cyclic.c \
 	signal.c \
 	files.c \
-	oom.c
+	oom.c \
+	hwloc.c
 
 flux_shell_LDADD = \
 	$(builddir)/libshell.la \

src/shell/batch.c

@@ -97,11 +97,10 @@ batch_info_create (flux_shell_t *shell, json_t *batch)
         shell_debug ("Copying batch script size=%zu for job to %s",
                      len,
                      b->script);
-        if (write_all (fd, data, len) < 0) {
+        if (write_all (fd, data, len) < 0 || close (fd) < 0) {
             shell_log_error ("failed to write batch script");
             goto error;
         }
-        close (fd);
     }
     return b;
 error:

src/shell/builtins.c

@@ -51,6 +51,7 @@ extern struct shell_builtin builtin_rlimit;
 extern struct shell_builtin builtin_cyclic;
 extern struct shell_builtin builtin_signal;
 extern struct shell_builtin builtin_oom;
+extern struct shell_builtin builtin_hwloc;
 
 static struct shell_builtin * builtins [] = {
     &builtin_tmpdir,
@@ -74,6 +75,7 @@ static struct shell_builtin * builtins [] = {
     &builtin_cyclic,
     &builtin_signal,
     &builtin_oom,
+    &builtin_hwloc,
     &builtin_list_end,
 };
 
@@ -91,6 +93,10 @@ static int shell_load_builtin (flux_shell_t *shell,
         || flux_plugin_add_handler (p, "shell.reconnect",
                                     sb->reconnect, NULL) < 0
         || flux_plugin_add_handler (p, "shell.init", sb->init, NULL) < 0
+        || flux_plugin_add_handler (p,
+                                    "shell.post-init",
+                                    sb->post_init,
+                                    NULL) < 0
         || flux_plugin_add_handler (p, "shell.exit", sb->exit, NULL) < 0
         || flux_plugin_add_handler (p, "task.init",  sb->task_init, NULL) < 0
         || flux_plugin_add_handler (p, "task.fork",  sb->task_fork, NULL) < 0

src/shell/builtins.h

@@ -21,6 +21,7 @@ struct shell_builtin {
     flux_plugin_f connect;
     flux_plugin_f reconnect;
     flux_plugin_f init;
+    flux_plugin_f post_init;
     flux_plugin_f task_init;
     flux_plugin_f task_exec;
     flux_plugin_f task_fork;

src/shell/hwloc.c

@@ -0,0 +1,99 @@
+/************************************************************\
+ * Copyright 2024 Lawrence Livermore National Security, LLC
+ * (c.f. AUTHORS, NOTICE.LLNS, COPYING)
+ *
+ * This file is part of the Flux resource manager framework.
+ * For details, see https://github.com/flux-framework.
+ *
+ * SPDX-License-Identifier: LGPL-3.0
+\************************************************************/
+
+/* hwloc options handler
+ */
+#define FLUX_SHELL_PLUGIN_NAME "hwloc"
+
+#if HAVE_CONFIG_H
+#include "config.h"
+#endif
+
+#include <sys/types.h>
+#include <sys/stat.h>
+#include <fcntl.h>
+#include <unistd.h>
+
+#include <flux/core.h>
+#include <flux/shell.h>
+
+#include "src/common/libutil/read_all.h"
+#include "ccan/str/str.h"
+
+#include "builtins.h"
+
+static int create_xmlfile (flux_shell_t *shell)
+{
+    int fd = -1;
+    char *xmlfile = NULL;
+    const char *hwloc_xml;
+    const char *tmpdir;
+
+    if (!(tmpdir = flux_shell_getenv (shell, "FLUX_JOB_TMPDIR")))
+        tmpdir = "/tmp";
+
+    if (flux_shell_get_hwloc_xml (shell, &hwloc_xml) < 0) {
+        shell_log_error ("failed to get shell hwloc xml");
+        goto error;
+    }
+    if (asprintf (&xmlfile, "%s/hwloc.xml", tmpdir) < 0) {
+        shell_log_error ("asprintf HWLOC_XMLFILE failed");
+        goto error;
+    }
+    if ((fd = open (xmlfile, O_CREAT|O_EXCL|O_WRONLY, 0640)) < 0) {
+        shell_log_errno ("%s", xmlfile);
+        goto error;
+    }
+    shell_debug ("Writing %ld bytes to HWLOC_XMLFILE=%s\n",
+                 (long int) strlen (hwloc_xml),
+                 xmlfile);
+    if (write_all (fd, hwloc_xml, strlen (hwloc_xml)) < 0 || close (fd) < 0) {
+        shell_log_errno ("failed to write HWLOC_XMLFILE");
+        goto error;
+    }
+    if (flux_shell_setenvf (shell, 0, "HWLOC_XMLFILE", "%s", xmlfile) < 0) {
+        shell_log_errno ("failed to set HWLOC_XMLFILE in job environment");
+        goto error;
+    }
+    return 0;
+error:
+    if (fd >= 0)
+        close (fd);
+    free (xmlfile);
+    return -1;
+}
+
+static int hwloc_post_init (flux_plugin_t *p,
+                            const char *topic,
+                            flux_plugin_arg_t *args,
+                            void *data)
+{
+    flux_shell_t *shell = flux_plugin_get_shell (p);
+    int xmlfile = 0;
+
+    if (flux_shell_getopt_unpack (shell,
+                                  "hwloc",
+                                  "{s?i}",
+                                    "xmlfile", &xmlfile) < 0)
+        return shell_log_errno ("failed to unpack hwloc options");
+
+    if (xmlfile && create_xmlfile (shell) < 0)
+        return shell_log_errno ("failed to write HWLOC_XMLFILE");
+    return 0;
+}
+
+struct shell_builtin builtin_hwloc = {
+    .name      = FLUX_SHELL_PLUGIN_NAME,
+    .post_init = hwloc_post_init,
+};
+
+/*
+ * vi:tabstop=4 shiftwidth=4 expandtab
+ */

t/Makefile.am

@@ -207,6 +207,7 @@ TESTSCRIPTS = \
 	t2616-job-shell-taskmap.t \
 	t2617-job-shell-stage-in.t \
 	t2618-job-shell-signal.t \
+	t2619-job-shell-hwloc.t \
 	t2710-python-cli-submit.t \
 	t2711-python-cli-run.t \
 	t2713-python-cli-bulksubmit.t \

t/t2619-job-shell-hwloc.t

@@ -0,0 +1,32 @@
+#!/bin/sh
+#
+test_description='Test flux-shell hwloc.xmlfile'
+
+. `dirname $0`/sharness.sh
+
+test_under_flux 2 job
+unset HWLOC_XMLFILE
+unset FLUX_HWLOC_XMLFILE
+
+test_expect_success 'shell: HWLOC_XMLFILE is not set for normal jobs' '
+	test_must_fail flux run printenv HWLOC_XMLFILE
+'
+test_expect_success 'shell: -o hwloc.xmlfile sets HWLOC_XMLFILE' '
+	flux run -o hwloc.xmlfile printenv HWLOC_XMLFILE
+'
+command -v hwloc-info >/dev/null && test_set_prereq HWLOC_INFO
+test_expect_success HWLOC_INFO 'shell: -o hwloc.xmlfile HWLOC_XMLFILE is valid' '
+	flux run -o hwloc.xmlfile sh -c "hwloc-info --input \$HWLOC_XMLFILE"
+'
+test_expect_success 'shell: -o hwloc.xmlfile sets HWLOC_XMLFILE per node' '
+	flux run -N2 --label-io -o hwloc.xmlfile printenv HWLOC_XMLFILE \
+		> xmlfile.out &&
+	test_debug "cat xmlfile.out" &&
+	sed -n "s/^0://p" xmlfile.out >path.0 &&
+	sed -n "s/^1://p" xmlfile.out >path.1 &&
+	test_must_fail test_cmp path.0 path.1
+'
+test_expect_success 'shell: bad hwloc.xmlfile value is an error' '
+	test_must_fail flux run -o hwloc.xmlfile=foo printenv HWLOC_XMLFILE
+'
+test_done