Merge pull request #3150 from grondo/issue#3141

flux-mini: add environment manipulation options

doc/man1/flux-mini.rst

@@ -123,6 +123,119 @@ emitting the job's I/O to its stdout and stderr.
 **-l, --label-io**
    Add task rank prefixes to each line of output.
 
+ENVIRONMENT
+===========
+
+By default, ``flux-mini`` duplicates the current environment when
+submitting jobs. However, a set of environment manipulation options are
+provided to give fine control over the requested environment submitted
+with the job.
+
+**--env=RULE**
+   Control how environment variables are exported with *RULE*. See
+   *ENV RULE SYNTAX* section below for more information. Rules are
+   applied in the order in which they are used on the command line.
+   This option may be specified multiple times.
+
+**--env-remove=PATTERN**
+   Remove all environment variables matching *PATTERN* from the current
+   generated environment. If *PATTERN* starts with a ``/`` character,
+   then it is considered a regex(7), otherwise *PATTERN* is treated
+   as a shell glob(7). This option is equivalent to ``--env=-PATTERN``
+   and may be used multiple times.
+
+**--env-file=FILE**
+   Read a set of environment *RULES* from a *FILE*. This option is
+   equivalent to ``--env=^FILE`` and may be used multiple times.
+
+ENV RULES
+=========
+
+The ``--env*`` options of ``flux-mini`` allow control of the environment
+exported to jobs via a set of *RULE* expressions. The currently supported
+rules are
+
+ * If a rule begins with ``-``, then the rest of the rule is a pattern
+   which removes matching environment variables. If the pattern starts
+   with ``/``, it is a regex(7), optionally ending with ``/``, otherwise
+   the pattern is considered a shell glob(7) expression.
+
+   Examples:
+      ``-*`` or ``-/.*/`` filter all environment variables creating an
+      empty environment.
+
+ * If a rule begins with ``^`` then the rest of the rule is a filename
+   from which to read more rules, one per line. The ``~`` character is
+   expanded to the user's home directory.
+
+   Examples:
+      ``~/envfile`` reads rules from file ``$HOME/envfile``
+
+ * If a rule is of the form ``VAR=VAL``, the variable ``VAR`` is set
+   to ``VAL``. Before being set, however, ``VAL`` will undergo simple
+   variable substitution using the Python ``string.Template`` class. This
+   simple substitution supports the following syntax:
+
+     * ``$$`` is an escape; it is replaced with ``$``
+     * ``$var`` will substitute ``var`` from the current environment,
+       falling back to the process environment. An error will be thrown
+       if environment variable ``var`` is not set.
+     * ``${var}`` is equivalent to ``$var``
+     * Advanced parameter substitution is not allowed, e.g. ``${var:-foo}``
+       will raise an error.
+
+   Examples:
+       ``PATH=/bin``, ``PATH=$PATH:/bin``, ``FOO=${BAR}something``
+
+ * Otherwise, the rule is considered a pattern from which to match
+   variables from the process environment if they do not exist in
+   the generated environment. E.g. ``PATH`` will export ``PATH`` from the
+   current environment (if it has not already been set in the generated
+   environment), and ``OMP*`` would copy all environment variables that
+   start with ``OMP`` and are not already set in the generated environment.
+   It is important to note that if the pattern does not match any variables,
+   then the rule is a no-op, i.e. an error is *not* generated.
+
+   Examples:
+       ``PATH``, ``FLUX_*_PATH``, ``/^OMP.*/``
+
+Since ``flux-mini`` always starts with a copy of the current environment,
+the default implicit rule is ``*`` (or ``--env=*``). To start with an
+empty environment instead, the ``-*`` rule or ``--env-remove=*`` option
+should be used. For example, the following will only export the current
+``PATH`` to a job:
+
+::
+
+    flux mini run --env-remove=* --env=PATH ...
+
+
+Since variables can be expanded from the currently built environment, and
+``--env`` options are applied in the order they are used, variables can
+be composed on the command line by multiple invocations of ``--env``, e.g.:
+
+::
+
+    flux mini run --env-remove=* \
+                  --env=PATH=/bin --env='PATH=$PATH:/usr/bin' ...
+
+Note that care must be taken to quote arguments so that ``$PATH`` is not
+expanded by the shell.
+
+
+This works particularly well when specifying rules in a file:
+
+::
+
+    -*
+    OMP*
+    FOO=bar
+    BAR=${FOO}/baz
+
+The above file would first clear the environment, then copy all variables
+starting with ``OMP`` from the current environment, set ``FOO=bar``,
+and then set ``BAR=bar/baz``.
+
 
 EXIT STATUS
 ===========

doc/test/spell.en.pws

@@ -512,3 +512,6 @@ sysconfdir
 TRACEME
 WIFEXTED
 builtin
+OMP
+envfile
+regex

src/cmd/flux-mini.py

@@ -14,7 +14,11 @@
 import logging
 import argparse
 import json
+import fnmatch
+import re
 from itertools import chain
+from string import Template
+from collections import ChainMap
 
 import flux
 from flux import job
@@ -23,6 +27,122 @@
 from flux import debugged
 
 
+def filter_dict(env, pattern, reverseMatch=True):
+    """
+    Filter out all keys that match "pattern" from dict 'env'
+
+    Pattern is assumed to be a shell glob(7) pattern, unless it begins
+    with '/', in which case the pattern is a regex.
+    """
+    if pattern.startswith("/"):
+        pattern = pattern[1::].rstrip("/")
+    else:
+        pattern = fnmatch.translate(pattern)
+    regex = re.compile(pattern)
+    if reverseMatch:
+        return dict(filter(lambda x: not regex.match(x[0]), env.items()))
+    return dict(filter(lambda x: regex.match(x[0]), env.items()))
+
+
+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.
+    """
+    if environ is None:
+        environ = dict(os.environ)
+    if rules is None:
+        return environ
+    for rule in rules:
+        #
+        #  If rule starts with '-' then the rest of the rule is a pattern
+        #   which filters matching environment variables from the
+        #   generated environment.
+        #
+        if rule.startswith("-"):
+            environ = filter_dict(environ, rule[1::])
+        #
+        #  If rule starts with '^', then the result of the rule is a filename
+        #   from which to read more rules.
+        #
+        elif rule.startswith("^"):
+            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)
+        #
+        #  Otherwise, the rule is an explicit variable assignment
+        #   VAR=VAL. If =VAL is not provided then VAL refers to the
+        #   value for VAR in the current environment of this process.
+        #
+        #  Quoted shell variables are expanded using values from the
+        #   built environment, not the process environment. So
+        #   --env=PATH=/bin --env=PATH='$PATH:/foo' results in
+        #   PATH=/bin:/foo.
+        #
+        else:
+            var, *rest = rule.split("=", 1)
+            if not rest:
+                #
+                #  VAR alone with no set value pulls in all matching
+                #   variables from current environment that are not already
+                #   in the generated environment.
+                env = filter_dict(os.environ, var, reverseMatch=False)
+                for key, value in env.items():
+                    if key not in environ:
+                        environ[key] = value
+            else:
+                #
+                #  Template lookup: use jobspec environment first, fallback
+                #   to current process environment using ChainMap:
+                lookup = ChainMap(environ, os.environ)
+                try:
+                    environ[var] = Template(rest[0]).substitute(lookup)
+                except ValueError as ex:
+                    LOGGER.error("--env: Unable to substitute %s", rule)
+                    raise
+                except KeyError as ex:
+                    raise Exception(f"--env: Variable {ex} not found in {rule}")
+    return environ
+
+
+class EnvFileAction(argparse.Action):
+    """Convenience class to handle --env-file option
+
+    Append --env-file options to the "env" list in namespace, with "^"
+    prepended to the rule to indicate further rules are to be read
+    from the indicated file.
+
+    This is required to preserve ordering between the --env and --env-file
+    and --env-remove options.
+    """
+
+    def __call__(self, parser, namespace, values, option_string=None):
+        items = getattr(namespace, "env", [])
+        if items is None:
+            items = []
+        items.append("^" + values)
+        setattr(namespace, "env", items)
+
+
+class EnvFilterAction(argparse.Action):
+    """Convenience class to handle --env-remove option
+
+    Append --env-remove options to the "env" list in namespace, with "-"
+    prepended to the option argument.
+
+    This is required to preserve ordering between the --env and --env-remove
+    options.
+    """
+
+    def __call__(self, parser, namespace, values, option_string=None):
+        items = getattr(namespace, "env", [])
+        if items is None:
+            items = []
+        items.append("-" + values)
+        setattr(namespace, "env", items)
+
+
 class MiniCmd:
     """
     MiniCmd is the base class for all flux-mini subcommands
@@ -71,6 +191,33 @@ def create_parser(exclude_io=False):
             help="Set job attribute ATTR to VAL (multiple use OK)",
             metavar="ATTR=VAL",
         )
+        parser.add_argument(
+            "--env",
+            action="append",
+            help="Control how environment variables are exported. If RULE "
+            + "starts with '-' apply rest of RULE as a remove filter (see "
+            + "--env-remove), if '^' then read rules from a file "
+            + "(see --env-file). Otherwise, set matching environment variables "
+            + "from the current environment (--env=PATTERN) or set a value "
+            + "explicitly (--env=VAR=VALUE). Rules are applied in the order "
+            + "they are used on the command line. (multiple use OK)",
+            metavar="RULE",
+        )
+        parser.add_argument(
+            "--env-remove",
+            action=EnvFilterAction,
+            help="Remove environment variables matching PATTERN. "
+            + "If PATTERN starts with a '/', then it is matched "
+            + "as a regular expression, otherwise PATTERN is a shell "
+            + "glob expression. (multiple use OK)",
+            metavar="PATTERN",
+        )
+        parser.add_argument(
+            "--env-file",
+            action=EnvFileAction,
+            help="Read a set of environment rules from FILE. (multiple use OK)",
+            metavar="FILE",
+        )
         parser.add_argument(
             "--input",
             type=str,
@@ -135,7 +282,7 @@ def submit(self, args):
         """
         jobspec = self.init_jobspec(args)
         jobspec.cwd = os.getcwd()
-        jobspec.environment = dict(os.environ)
+        jobspec.environment = get_filtered_environment(args.env)
         if args.time_limit is not None:
             jobspec.duration = args.time_limit
 

t/t2700-mini-cmd.t

@@ -182,6 +182,59 @@ test_expect_success HAVE_JQ 'flux mini --job-name works' '
 	flux mini submit --dry-run --job-name=foobar hostname >name.out &&
 	test $(jq ".attributes.system.job.name" name.out) = "\"foobar\""
 '
+test_expect_success HAVE_JQ 'flux-mini --env=-*/--env-remove=* works' '
+	flux mini submit --dry-run --env=-* hostname > no-env.out &&
+	jq -e ".attributes.system.environment == {}" < no-env.out &&
+	flux mini submit --dry-run --env-remove=* hostname > no-env2.out &&
+	jq -e ".attributes.system.environment == {}" < no-env2.out
+'
+test_expect_success HAVE_JQ 'flux-mini --env=VAR works' '
+	FOO=bar flux mini submit --dry-run \
+	    --env=-* --env FOO hostname >FOO-env.out &&
+	jq -e ".attributes.system.environment == {\"FOO\": \"bar\"}" FOO-env.out
+'
+test_expect_success HAVE_JQ 'flux-mini --env=PATTERN works' '
+	FOO_ONE=bar FOO_TWO=baz flux mini submit --dry-run \
+	    --env=-* --env="FOO_*" hostname >FOO-pattern-env.out &&
+	jq -e ".attributes.system.environment == \
+           {\"FOO_ONE\": \"bar\", \"FOO_TWO\": \"baz\"}" FOO-pattern-env.out &&
+	FOO_ONE=bar FOO_TWO=baz flux mini submit --dry-run \
+	    --env=-* --env="/^FOO_.*/" hostname >FOO-pattern2-env.out &&
+	jq -e ".attributes.system.environment == \
+           {\"FOO_ONE\": \"bar\", \"FOO_TWO\": \"baz\"}" FOO-pattern2-env.out
 
-
+'
+test_expect_success HAVE_JQ 'flux-mini --env=VAR=VAL works' '
+	flux mini submit --dry-run \
+	    --env=-* --env PATH=/bin hostname >PATH-env.out &&
+	jq -e ".attributes.system.environment == {\"PATH\": \"/bin\"}" PATH-env.out &&
+	FOO=bar flux mini submit --dry-run \
+	    --env=-* --env FOO=\$FOO:baz hostname >FOO-append.out &&
+	jq -e ".attributes.system.environment == {\"FOO\": \"bar:baz\"}" FOO-append.out
+'
+test_expect_success 'flux-mini --env=VAR=${VAL:-default} fails' '
+    test_expect_code 1 flux mini run --dry-run \
+        --env=* --env=VAR=\${VAL:-default} hostname >env-fail.err 2>&1 &&
+    test_debug "cat env-fail.err" &&
+    grep "Unable to substitute" env-fail.err
+'
+test_expect_success 'flux-mini --env=VAR=$VAL fails when VAL not in env' '
+    unset VAL &&
+    test_expect_code 1 flux mini run --dry-run \
+        --env=* --env=VAR=\$VAL hostname >env-notset.err 2>&1 &&
+    test_debug "cat env-notset.err" &&
+    grep "env: Variable .* not found" env-notset.err
+'
+test_expect_success HAVE_JQ 'flux-mini --env-file works' '
+	cat <<-EOF >envfile &&
+	-*
+	FOO=bar
+	BAR=\${FOO}/baz
+	EOF
+    for arg in "--env=^envfile" "--env-file=envfile"; do
+	  flux mini submit --dry-run ${arg} hostname >envfile.out &&
+	  jq -e ".attributes.system.environment == \
+	       {\"FOO\":\"bar\", \"BAR\":\"bar/baz\"}" envfile.out
+    done
+'
 test_done