Merge pull request #6506 from grondo/mustache-env

mergify[bot] · web-flow · commit ce00cca7274b · 2025-01-06T19:33:27.000Z
allow mustache templates in job environment variables
diff --git a/doc/man1/common/job-env-rules.rst b/doc/man1/common/job-env-rules.rst
@@ -31,8 +31,13 @@ via a set of *RULE* expressions. The currently supported rules are
      * Advanced parameter substitution is not allowed, e.g. ``${var:-foo}``
        will raise an error.
 
+   ``VAL`` may also contain a mustache template, in which case the template
+   will be substituted in the job shell with the corresponding value before
+   launching job tasks. See `MUSTACHE TEMPLATES`_ for more information.
+
    Examples:
-       ``PATH=/bin``, ``PATH=$PATH:/bin``, ``FOO=${BAR}something``
+       ``PATH=/bin``, ``PATH=$PATH:/bin``, ``FOO=${BAR}something``,
+       ``PATH=${PATH}:/{{tmpdir}}/bin``
 
  * Otherwise, the rule is considered a pattern from which to match
    variables from the process environment if they do not exist in
diff --git a/src/bindings/python/flux/cli/base.py b/src/bindings/python/flux/cli/base.py
@@ -288,10 +288,11 @@ def get_filtered_environment(rules, environ=None):
     Filter environment dictionary 'environ' given a list of rules.
     Each rule can filter, set, or modify the existing environment.
     """
+    env_expand = {}
     if environ is None:
         environ = dict(os.environ)
     if rules is None:
-        return environ
+        return environ, env_expand
     for rule in rules:
         #
         #  If rule starts with '-' then the rest of the rule is a pattern
@@ -308,7 +309,8 @@ def get_filtered_environment(rules, environ=None):
             filename = os.path.expanduser(rule[1::])
             with open(filename) as envfile:
                 lines = [line.strip() for line in envfile]
-                environ = get_filtered_environment(lines, environ=environ)
+                environ, envx = get_filtered_environment(lines, environ=environ)
+                env_expand.update(envx)
         #
         #  Otherwise, the rule is an explicit variable assignment
         #   VAR=VAL. If =VAL is not provided then VAL refers to the
@@ -330,6 +332,11 @@ def get_filtered_environment(rules, environ=None):
                 for key, value in env.items():
                     if key not in environ:
                         environ[key] = value
+            elif "{{" in rest[0]:
+                #
+                #  Mustache template which should be expanded by job shell.
+                #  Place result in env_expand instead of environ:
+                env_expand[var] = rest[0]
             else:
                 #
                 #  Template lookup: use jobspec environment first, fallback
@@ -342,7 +349,7 @@ def get_filtered_environment(rules, environ=None):
                     raise
                 except KeyError as ex:
                     raise Exception(f"--env: Variable {ex} not found in {rule}")
-    return environ
+    return environ, env_expand
 
 
 class EnvFileAction(argparse.Action):
@@ -1001,7 +1008,13 @@ def jobspec_create(self, args):
         Create a jobspec from args and return it to caller
         """
         jobspec = self.init_jobspec(args)
-        jobspec.environment = get_filtered_environment(args.env)
+
+        jobspec.environment, env_expand = get_filtered_environment(args.env)
+        if env_expand:
+            # "expanded" environment variables are set in env-expand
+            # shell options dict and will be processed by the shell.
+            jobspec.setattr_shell_option("env-expand", env_expand)
+
         jobspec.cwd = args.cwd if args.cwd is not None else os.getcwd()
         rlimits = get_filtered_rlimits(args.rlimit)
         if rlimits:
diff --git a/src/shell/Makefile.am b/src/shell/Makefile.am
@@ -99,7 +99,8 @@ flux_shell_SOURCES = \
 	signal.c \
 	files.c \
 	hwloc.c \
-	rexec.c
+	rexec.c \
+	env-expand.c
 
 if HAVE_INOTIFY
 flux_shell_SOURCES += oom.c
diff --git a/src/shell/builtins.c b/src/shell/builtins.c
@@ -56,6 +56,7 @@ extern struct shell_builtin builtin_signal;
 extern struct shell_builtin builtin_oom;
 extern struct shell_builtin builtin_hwloc;
 extern struct shell_builtin builtin_rexec;
+extern struct shell_builtin builtin_env_expand;
 
 static struct shell_builtin * builtins [] = {
     &builtin_tmpdir,
@@ -86,6 +87,7 @@ static struct shell_builtin * builtins [] = {
 #endif
     &builtin_hwloc,
     &builtin_rexec,
+    &builtin_env_expand,
     &builtin_list_end,
 };
 
diff --git a/src/shell/env-expand.c b/src/shell/env-expand.c
@@ -0,0 +1,113 @@
+/************************************************************\
+ * 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
+\************************************************************/
+
+#define FLUX_SHELL_PLUGIN_NAME "env-expand"
+
+#if HAVE_CONFIG_H
+#include "config.h"
+#endif
+
+#include <jansson.h>
+
+#include <flux/shell.h>
+
+#include "builtins.h"
+
+static int env_expand (flux_plugin_t *p,
+                       const char *topic,
+                       flux_plugin_arg_t *args,
+                       void *data)
+{
+    json_t *to_expand = NULL;
+    json_t *value;
+    void *tmp;
+    const char *key;
+    flux_shell_t *shell = flux_plugin_get_shell (p);
+
+    if (!(shell = flux_plugin_get_shell (p)))
+        return shell_log_errno ("unable to get shell handle");
+    if (flux_shell_getopt_unpack (shell, "env-expand", "o", &to_expand) != 1)
+        return 0;
+    json_object_foreach_safe (to_expand, tmp, key, value) {
+        const char *s = json_string_value (value);
+        char *result;
+        if (s == NULL) {
+            shell_log_error ("invalid value for env var %s", key);
+            continue;
+        }
+        result = flux_shell_mustache_render (shell, s);
+
+        /*  If mustache render was successful, then set it for the job and
+         *  remove the key from the env-expand object internally, so it isn't
+         *  expanded again in task_env_expand():
+         */
+        if (result && !strstr (result, "{{")) {
+            if (flux_shell_setenvf (shell, 1, key, "%s", result) < 0)
+                shell_log_errno ("failed to set %s=%s", key, result);
+            else
+                (void) json_object_del (to_expand, key);
+        }
+        free (result);
+    }
+    return 0;
+}
+
+/* Per-task environment variable mustache substitution.
+ * N.B.: Only templates that were not fully rendered by env_expand() above
+ * should remain in the `env-expand` shell options object.
+ */
+static int task_env_expand (flux_plugin_t *p,
+                            const char *topic,
+                            flux_plugin_arg_t *args,
+                            void *data)
+{
+    json_t *to_expand = NULL;
+    json_t *value;
+    const char *key;
+    flux_shell_t *shell;
+    flux_shell_task_t *task;
+    flux_cmd_t *cmd;
+
+    if (!(shell = flux_plugin_get_shell (p)))
+        return shell_log_errno ("unable to get shell handle");
+    if (flux_shell_getopt_unpack (shell, "env-expand", "o", &to_expand) != 1)
+        return 0;
+
+    if (!(task = flux_shell_current_task (shell))
+        || !(cmd = flux_shell_task_cmd (task)))
+        return -1;
+
+    json_object_foreach (to_expand, key, value) {
+        const char *s = json_string_value (value);
+        char *result;
+        if (s == NULL) {
+            shell_log_error ("invalid value for env var %s", key);
+            continue;
+        }
+        if (!(result = flux_shell_mustache_render (shell, s))) {
+            shell_log_errno ("failed to expand env var %s=%s", key, s);
+            continue;
+        }
+        if (flux_cmd_setenvf (cmd, 1, key, "%s", result) < 0)
+            shell_log_errno ("failed to set %s=%s", key, result);
+        free (result);
+    }
+    return 0;
+}
+
+struct shell_builtin builtin_env_expand = {
+    .name = FLUX_SHELL_PLUGIN_NAME,
+    .init = env_expand,
+    .task_init = task_env_expand,
+};
+
+/*
+ * vi:tabstop=4 shiftwidth=4 expandtab
+ */
diff --git a/t/t2620-job-shell-mustache.t b/t/t2620-job-shell-mustache.t
@@ -45,5 +45,18 @@ test_expect_success 'mustache: unsupported tag is left alone' '
 	test_debug "cat output.4" &&
 	test "$(cat output.4)" = "{{foo}} {{node.foo}} {{task.foo}}"
 '
-
+test_expect_success 'mustache: mustache templates can be rendered in env' '
+	flux run --env=TEST={{tmpdir}} -N2 \
+		sh -c "test \$TEST = \$FLUX_JOB_TMPDIR"
+'
+test_expect_success 'mustache: env variables can have per-task tags' '
+	flux run --env=TEST={{taskid}} -N2 -n4 \
+		sh -c "test \$TEST = \$FLUX_TASK_RANK" &&
+	flux run --env=T1={{size}} --env=T2={{taskid}} -N2 -n4 \
+		sh -c "test \$T1 -eq 4 -a \$T2 = \$FLUX_TASK_RANK"
+'
+test_expect_success 'mustache: invalid tags in env vars are left unexpanded' '
+	flux run --env=TEST={{task.foo}} \
+		sh -c "test \$TEST = {{task.foo}}"
+'
 test_done
