Skip to content

Commit

Permalink
python: cli: allow mustache templates in --env* options
Browse files Browse the repository at this point in the history
Problem: It would be useful to expand mustache templates in environment
variable values, but this is currently not supported in the submission
cli commands.

When processing the `--env*` set of options, detect any provided
environment variable values that appear to contain a mustache
template. Return these in a separate dict from the main environ
dictionary and add them to `attributes.system.shel.optionsenv-expand`
in jobspec.  The job shell will then expand these variables before
running the job.
  • Loading branch information
grondo committed Dec 13, 2024
1 parent 1957539 commit 3eb221f
Showing 1 changed file with 17 additions and 4 deletions.
21 changes: 17 additions & 4 deletions src/bindings/python/flux/cli/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand All @@ -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]

Check warning on line 339 in src/bindings/python/flux/cli/base.py

View check run for this annotation

Codecov / codecov/patch

src/bindings/python/flux/cli/base.py#L339

Added line #L339 was not covered by tests
else:
#
# Template lookup: use jobspec environment first, fallback
Expand All @@ -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):
Expand Down Expand Up @@ -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)

Check warning on line 1016 in src/bindings/python/flux/cli/base.py

View check run for this annotation

Codecov / codecov/patch

src/bindings/python/flux/cli/base.py#L1016

Added line #L1016 was not covered by tests

jobspec.cwd = args.cwd if args.cwd is not None else os.getcwd()
rlimits = get_filtered_rlimits(args.rlimit)
if rlimits:
Expand Down

0 comments on commit 3eb221f

Please sign in to comment.