Merge pull request #5186 from grondo/imp-exec-helper

mergify[bot] · web-flow · commit abd95c19d069 · 2023-05-22T14:36:41.000Z
job-exec: provide IMP exec helper when possible and use stdin/out for shell exec protocol
diff --git a/src/cmd/Makefile.am b/src/cmd/Makefile.am
@@ -112,7 +112,8 @@ dist_fluxcmd_SCRIPTS = \
 	flux-pstree.py \
 	flux-pgrep.py \
 	flux-queue.py \
-	flux-cancel.py
+	flux-cancel.py \
+	flux-imp-exec-helper
 
 fluxcmd_PROGRAMS = \
 	flux-terminus \
diff --git a/src/cmd/flux-imp-exec-helper b/src/cmd/flux-imp-exec-helper
@@ -0,0 +1,20 @@
+#!/bin/sh
+##############################################################
+# Copyright 2023 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
+##############################################################
+#
+# Helper for flux-imp exec functionality.
+# Emit input to IMP for jobid on stdout given jobid in $1
+#
+JOBID=${1:-$FLUX_JOB_ID}
+if test -z "$JOBID"; then
+    echo "flux-imp-exec-helper: Unable to determine jobid" >&2
+    exit 1
+fi
+printf '{"J": "%s"}' $(flux job info --orig $JOBID J)
diff --git a/src/common/libflux/Makefile.am b/src/common/libflux/Makefile.am
@@ -34,8 +34,7 @@ installed_conf_cppflags = \
 	-DINSTALLED_UPMI_PLUGINPATH=\"$(fluxlibdir)/upmi/plugins\" \
 	-DINSTALLED_CMDHELP_PATTERN=\"${datadir}/flux/help.d/*.json\" \
 	-DINSTALLED_NO_DOCS_PATH=\"${datadir}/flux/.nodocs\" \
-	-DINSTALLED_RUNDIR=\"${runstatedir}/flux\" \
-	-DINSTALLED_BINDIR=\"$(fluxcmddir)\"
+	-DINSTALLED_RUNDIR=\"${runstatedir}/flux\"
 
 intree_conf_cppflags = \
 	-DINTREE_MODULE_PATH=\"$(abs_top_builddir)/src/modules/.libs\" \
@@ -55,8 +54,7 @@ intree_conf_cppflags = \
 	-DINTREE_JOBTAP_PLUGINPATH=\"$(abs_top_builddir)/src/modules/job-manager/plugins/.libs\" \
 	-DINTREE_UPMI_PLUGINPATH=\"$(abs_top_builddir)/src/common/libpmi/plugins/.libs\" \
 	-DINTREE_CMDHELP_PATTERN=\"${abs_top_builddir}/etc/flux/help.d/*.json\" \
-	-DINTREE_NO_DOCS_PATH=\"${abs_top_builddir}/etc/flux/.nodocs\" \
-	-DINTREE_BINDIR=\"${abs_top_builddir}/src/cmd\"
+	-DINTREE_NO_DOCS_PATH=\"${abs_top_builddir}/etc/flux/.nodocs\"
 
 
 fluxcoreinclude_HEADERS = \
diff --git a/src/common/libflux/conf.c b/src/common/libflux/conf.c
@@ -66,7 +66,6 @@ static struct builtin builtin_tab[] = {
     { "upmi_pluginpath",INSTALLED_UPMI_PLUGINPATH,  INTREE_UPMI_PLUGINPATH },
     { "no_docs_path",   INSTALLED_NO_DOCS_PATH,     INTREE_NO_DOCS_PATH },
     { "rundir",         INSTALLED_RUNDIR,           NULL },
-    { "bindir",         INSTALLED_BINDIR,           INTREE_BINDIR },
     { NULL, NULL, NULL },
 };
 
diff --git a/src/modules/job-exec/exec.c b/src/modules/job-exec/exec.c
@@ -86,7 +86,7 @@ static void start_cb (struct bulk_exec *exec, void *arg)
     /*  This is going to be really slow. However, it should at least
      *   work for now. We wait for all imp's to start, then send input
      */
-    if (job->multiuser) {
+    if (job->multiuser && !config_use_imp_helper ()) {
         char *input = NULL;
         json_t *o = json_pack ("{s:s}", "J", job->J);
         if (!o || !(input = json_dumps (o, JSON_COMPACT))) {
@@ -117,13 +117,13 @@ static void complete_cb (struct bulk_exec *exec, void *arg)
 static int exec_barrier_enter (struct bulk_exec *exec)
 {
     struct exec_ctx *ctx = bulk_exec_aux_get (exec, "ctx");
+    const char *stream = config_use_imp_helper () ?
+                         "stdin" :
+                         "FLUX_EXEC_PROTOCOL_FD";
     if (!ctx)
         return -1;
     if (++ctx->barrier_enter_count == bulk_exec_total (exec)) {
-        if (bulk_exec_write (exec,
-                             "FLUX_EXEC_PROTOCOL_FD",
-                             "exit=0\n",
-                             7) < 0)
+        if (bulk_exec_write (exec, stream, "exit=0\n", 7) < 0)
             return -1;
         ctx->barrier_enter_count = 0;
         ctx->barrier_completion_count++;
@@ -135,10 +135,7 @@ static int exec_barrier_enter (struct bulk_exec *exec)
          *   case where a shell exits while a barrier is already in progress
          *   is handled in exit_cb().
          */
-        if (bulk_exec_write (exec,
-                             "FLUX_EXEC_PROTOCOL_FD",
-                             "exit=1\n",
-                             7) < 0)
+        if (bulk_exec_write (exec, stream, "exit=1\n", 7) < 0)
             return -1;
     }
     return 0;
@@ -153,7 +150,8 @@ static void output_cb (struct bulk_exec *exec, flux_subprocess_t *p,
     struct jobinfo *job = arg;
     const char *cmd = flux_cmd_arg (flux_subprocess_get_cmd (p), 0);
 
-    if (streq (stream, "FLUX_EXEC_PROTOCOL_FD")) {
+    if (streq (stream, "FLUX_EXEC_PROTOCOL_FD")
+        || (config_use_imp_helper () && streq (stream, "stdout"))) {
         if (streq (data, "enter\n")
             && exec_barrier_enter (exec) < 0) {
             jobinfo_fatal_error (job,
@@ -293,10 +291,10 @@ static void exit_cb (struct bulk_exec *exec,
      */
     if (ctx->barrier_completion_count == 0
         || ctx->barrier_enter_count > 0) {
-        if (bulk_exec_write (exec,
-                             "FLUX_EXEC_PROTOCOL_FD",
-                             "exit=1\n",
-                             7) < 0)
+        const char *stream = config_use_imp_helper () ?
+                             "stdin" :
+                             "FLUX_EXEC_PROTOCOL_FD";
+        if (bulk_exec_write (exec, stream, "exit=1\n", 7) < 0)
             jobinfo_fatal_error (job, 0,
                                  "failed to terminate barrier: %s",
                                  strerror (errno));
@@ -351,6 +349,15 @@ static int exec_init (struct jobinfo *job)
         goto err;
     }
     if (job->multiuser) {
+            if (config_use_imp_helper ()
+                && flux_cmd_setenvf (cmd,
+                                     1,
+                                     "FLUX_IMP_EXEC_HELPER",
+                                     "flux imp-exec-helper %ju",
+                                     (uintmax_t) job->id) < 0) {
+            flux_log_error (job->h, "exec_init: flux_cmd_setenvf");
+            goto err;
+        }
         if (flux_cmd_argv_append (cmd, config_get_imp_path ()) < 0
             || flux_cmd_argv_append (cmd, "exec") < 0) {
             flux_log_error (job->h, "exec_init: flux_cmd_argv_append");
@@ -366,7 +373,7 @@ static int exec_init (struct jobinfo *job)
     /*  If more than one shell is involved in this job, set up a channel
      *   for exec system based barrier:
      */
-    if (idset_count (ranks) > 1) {
+    if (idset_count (ranks) > 1 && !config_use_imp_helper ()) {
         if (flux_cmd_add_channel (cmd, "FLUX_EXEC_PROTOCOL_FD") < 0
             || flux_cmd_setopt (cmd,
                                 "FLUX_EXEC_PROTOCOL_FD_LINE_BUFFER",
diff --git a/src/modules/job-exec/exec_config.c b/src/modules/job-exec/exec_config.c
@@ -17,12 +17,17 @@
 #include <jansson.h>
 #include <unistd.h>
 
+#if HAVE_FLUX_SECURITY
+#include <flux/security/version.h>
+#endif
+
 #include "exec_config.h"
 #include "ccan/str/str.h"
 
 static const char *default_cwd = "/tmp";
 static const char *default_job_shell = NULL;
 static const char *flux_imp_path = NULL;
+static bool use_imp_helper = false;
 
 static const char *jobspec_get_job_shell (json_t *jobspec)
 {
@@ -69,6 +74,11 @@ const char *config_get_imp_path (void)
     return flux_imp_path;
 }
 
+bool config_use_imp_helper (void)
+{
+    return use_imp_helper;
+}
+
 /*  Initialize common configurations for use by job-exec exec modules.
  */
 int config_init (flux_t *h, int argc, char **argv)
@@ -104,19 +114,34 @@ int config_init (flux_t *h, int argc, char **argv)
         return -1;
     }
 
+#if HAVE_FLUX_SECURITY
+    /* Use IMP helper by default for flux-security >= 0.9.0
+     */
+    if (FLUX_SECURITY_VERSION_MAJOR >= 0
+        && FLUX_SECURITY_VERSION_MINOR >= 9)
+        use_imp_helper = true;
+#endif /* HAVE_FLUX_SECURITY */
+
     if (argv && argc) {
         /* Finally, override values on cmdline */
         for (int i = 0; i < argc; i++) {
             if (strstarts (argv[i], "job-shell="))
                 default_job_shell = argv[i]+10;
             else if (strstarts (argv[i], "imp="))
                 flux_imp_path = argv[i]+4;
+            else if (streq (argv[i], "no-imp-helper"))
+                use_imp_helper = false;
         }
     }
 
     flux_log (h, LOG_DEBUG, "using default shell path %s", default_job_shell);
-    if (flux_imp_path)
-        flux_log (h, LOG_DEBUG, "using imp path %s", flux_imp_path);
+    if (flux_imp_path) {
+        flux_log (h,
+                  LOG_DEBUG,
+                  "using imp path %s (%s helper)",
+                  flux_imp_path,
+                  use_imp_helper ? "with" : "without");
+    }
     return 0;
 }
 
diff --git a/src/modules/job-exec/exec_config.h b/src/modules/job-exec/exec_config.h
@@ -21,6 +21,8 @@ const char *config_get_cwd (struct jobinfo *job);
 
 const char *config_get_imp_path (void);
 
+bool config_use_imp_helper (void);
+
 int config_init (flux_t *h, int argc, char **argv);
 
 #endif /* !HAVE_JOB_EXEC_CONFIG_EXEC_H */
diff --git a/src/modules/job-exec/job-exec.c b/src/modules/job-exec/job-exec.c
@@ -101,6 +101,7 @@
 
 #include "job-exec.h"
 #include "checkpoint.h"
+#include "exec_config.h"
 
 static double kill_timeout=5.0;
 
@@ -890,7 +891,7 @@ static void jobinfo_start_continue (flux_future_t *f, void *arg)
         jobinfo_fatal_error (job, errno, "initializing critical ranks");
         goto done;
     }
-    if (job->multiuser) {
+    if (job->multiuser && !config_use_imp_helper ()) {
         const char *J = jobinfo_kvs_lookup_get (f, "J");
         if (!J || !(job->J = strdup (J))) {
             jobinfo_fatal_error (job, errno, "reading J: %s", error.text);
@@ -1056,6 +1057,7 @@ static flux_future_t *jobinfo_start_init (struct jobinfo *job)
         || flux_future_push (f, "R", f_kvs) < 0)
         goto err;
     if (job->multiuser
+        && !config_use_imp_helper ()
         && (!(f_kvs = flux_jobid_kvs_lookup (h, job->id, 0, "J"))
         || flux_future_push (f, "J", f_kvs) < 0)) {
         goto err;
diff --git a/src/modules/job-exec/sdexec.c b/src/modules/job-exec/sdexec.c
@@ -86,11 +86,6 @@ struct sdexec
     bool jobinfo_tasks_complete_called;
 };
 
-static int sdexec_config (flux_t *h, int argc, char **argv)
-{
-    return config_init (h, argc, argv);
-}
-
 static void sdexec_destroy (void *data)
 {
     struct sdexec *se = data;
@@ -788,7 +783,6 @@ static int sdexec_cancel (struct jobinfo *job)
 
 struct exec_implementation sdexec = {
     .name =     "sdexec",
-    .config =   sdexec_config,
     .init =     sdexec_init,
     .exit =     sdexec_exit,
     .start =    sdexec_start,
diff --git a/src/shell/internal.h b/src/shell/internal.h
@@ -27,7 +27,7 @@ struct flux_shell {
     flux_jobid_t jobid;
     int broker_rank;
     char hostname [MAXHOSTNAMELEN + 1];
-    int protocol_fd;
+    int protocol_fd[2];
 
     optparse_t *p;
     flux_t *h;
diff --git a/src/shell/shell.c b/src/shell/shell.c
@@ -974,10 +974,12 @@ static int get_protocol_fd (int *pfd)
         }
         if (fd_set_cloexec (fd) < 0)
             return -1;
-        *pfd = fd;
+        pfd[0] = fd;
+        pfd[1] = fd;
         return 0;
     }
-    *pfd = -1;
+    pfd[0] = STDIN_FILENO;
+    pfd[1] = STDOUT_FILENO;
     return 0;
 }
 
@@ -1112,7 +1114,7 @@ static void shell_initialize (flux_shell_t *shell)
     if (gethostname (shell->hostname, sizeof (shell->hostname)) < 0)
         shell_die_errno (1, "gethostname");
 
-    if (get_protocol_fd (&shell->protocol_fd) < 0)
+    if (get_protocol_fd (shell->protocol_fd) < 0)
         shell_die_errno (1, "Failed to parse FLUX_EXEC_PROTOCOL_FD");
 
     if (!(shell->completion_refs = zhashx_new ()))
@@ -1240,10 +1242,10 @@ static int shell_barrier (flux_shell_t *shell, const char *name)
     if (shell->standalone || shell->info->shell_size == 1)
         return 0; // NO-OP
 
-    if (shell->protocol_fd < 0)
+    if (shell->protocol_fd[1] < 0)
         shell_die (1, "required FLUX_EXEC_PROTOCOL_FD not set");
 
-    if (dprintf (shell->protocol_fd, "enter\n") != 6)
+    if (dprintf (shell->protocol_fd[1], "enter\n") != 6)
         shell_die_errno (1, "shell_barrier: dprintf");
 
     /*  Note: The only expected values currently are "exit=0\n"
@@ -1254,7 +1256,7 @@ static int shell_barrier (flux_shell_t *shell, const char *name)
      *   elsewhere.
      */
     memset (buf, 0, sizeof (buf));
-    if (read (shell->protocol_fd, buf, 7) < 0)
+    if (read (shell->protocol_fd[0], buf, 7) < 0)
         shell_die_errno (1, "shell_barrier: read");
     if (!streq (buf, "exit=0\n"))
         exit (1);
diff --git a/src/test/docker/docker-run-checks.sh b/src/test/docker/docker-run-checks.sh
@@ -16,7 +16,7 @@ JOBS=2
 MOUNT_HOME_ARGS="--volume=$HOME:/home/$USER -e HOME"
 
 if test "$PROJECT" = "flux-core"; then
-  FLUX_SECURITY_VERSION=0.8.0
+  FLUX_SECURITY_VERSION=0.9.0
   POISON=t
 fi
 
diff --git a/t/job-exec/dummy.sh b/t/job-exec/dummy.sh
@@ -28,8 +28,8 @@
 #  drive tests of expected execution behavior.
 #
 
-die() { echo "dummy-shell: $@" >&2; exit 1; }
-log() { echo "dummy-shell: $@" >&2; }
+die() { echo "dummy-shell[$BROKER_RANK]: $@" >&2; exit 1; }
+log() { echo "dummy-shell[$BROKER_RANK]: $@" >&2; }
 
 #
 #  Declare globals:
@@ -104,10 +104,10 @@ get_command() {
 }
 
 get_traps() {
-    trap "/bin/echo got SIGTERM" 15
+    trap "log got SIGTERM" 15
     TRAP=$(json_get "$JOBSPEC" '.attributes.system.environment.TRAP')
     if test "x$TRAP" != "xnull" ; then
-        trap "/bin/echo got signal $TRAP" $TRAP
+        trap "log got signal $TRAP" $TRAP
     fi
 }
 
@@ -135,7 +135,7 @@ test_mock_failure() {
             #   entered the barrier before rank 1 unexpectedly exits.
             #   See caveats about best effort above.
             #
-            echo "after_barrier_entry: rank=$JOB_SHELL_RANK"
+            log "after_barrier_entry: rank=$JOB_SHELL_RANK"
             if test $JOB_SHELL_RANK -eq 1; then
                 log "after_barrier: exiting early on job shell rank 1"
                 sleep 1
@@ -149,9 +149,9 @@ test_mock_failure() {
 }
 
 barrier() {
-    if [[ ${NNODES} -gt 1 && -n ${FLUX_EXEC_PROTOCOL_FD} ]]; then
-        echo enter >&${FLUX_EXEC_PROTOCOL_FD}
-        read -u ${FLUX_EXEC_PROTOCOL_FD} line
+    if [[ ${NNODES} -gt 1 ]]; then
+        echo enter >&${FLUX_EXEC_PROTOCOL_FD:-1}
+        read -u ${FLUX_EXEC_PROTOCOL_FD:-0} line
         if [[ ${line##*=} -ne 0 ]]; then
             exit ${line##*=}
         fi
diff --git a/t/job-exec/imp.sh b/t/job-exec/imp.sh
@@ -9,7 +9,7 @@ case "$cmd" in
     exec)
         shift; 
         #  Copy IMP's input to a file so tests can check result:
-        cat - >imp-$(flux job id $2).input
+        ${FLUX_IMP_EXEC_HELPER:-cat} >imp-$(flux job id $2).input
         printf "test-imp: Running $* \n" >&2
         exec "$@" ;;
     kill)
diff --git a/t/t2402-job-exec-dummy.t b/t/t2402-job-exec-dummy.t
@@ -55,7 +55,7 @@ test_expect_success 'job-exec: job exception uses SIGKILL after kill-timeout' '
 	flux module reload job-exec kill-timeout=0.2 &&
 	cat <<-EOF >trap-sigterm.sh &&
 	#!/bin/sh
-	trap "echo trap-sigterm got SIGTERM" 15
+	trap "echo trap-sigterm got SIGTERM >&2" 15
 	flux kvs put trap-ready=1
 	sleep 60 &
 	pid=\$!
diff --git a/t/t2602-job-shell.t b/t/t2602-job-shell.t
diff --git a/t/valgrind/valgrind.supp b/t/valgrind/valgrind.supp
