From 6dce1b9e802f7a3d332ce5d7e50e601af4b4a21a Mon Sep 17 00:00:00 2001 From: vsoch Date: Wed, 12 Jun 2024 06:12:13 -0600 Subject: [PATCH] persist environment Signed-off-by: vsoch --- jobspec/transformer/flux/steps.py | 7 +++---- jobspec/utils.py | 6 +++--- 2 files changed, 6 insertions(+), 7 deletions(-) diff --git a/jobspec/transformer/flux/steps.py b/jobspec/transformer/flux/steps.py index cf8349d..7aa74b0 100644 --- a/jobspec/transformer/flux/steps.py +++ b/jobspec/transformer/flux/steps.py @@ -85,10 +85,10 @@ def write_job_script(self, command): command = command.strip() match = re.match("#!/bin/(?Pbash|sh|python)", command) terms = match.groupdict() - tmpfile = utils.get_tmpfile(prefix="jobscript-") + tmpfile = utils.get_tmpfile(prefix="jobscript-", suffix=".sh") - # Clean up (activate later) - command += f"\n#rm -rf {tmpfile}" + # Clean self up + command += f"\nrm -rf {tmpfile}" utils.write_file(command, tmpfile) return [terms["executable"], tmpfile] @@ -119,7 +119,6 @@ def prepare(self, command=None, waitable=False): for key, value in attributes.get("environment", {}).items(): cmd += [f"--env={key}={value}"] - print(cmd) # Note that you need to install our frobnicator plugin # for this to work. See the examples/depends_on directory for depends_on in task.get("depends_on") or []: diff --git a/jobspec/utils.py b/jobspec/utils.py index 674ff34..3162411 100644 --- a/jobspec/utils.py +++ b/jobspec/utils.py @@ -40,7 +40,7 @@ def recursive_find(base, pattern="[.]py"): yield filepath -def get_tmpfile(tmpdir=None, prefix=""): +def get_tmpfile(tmpdir=None, prefix="", suffix=None): """ Get a temporary file with an optional prefix. """ @@ -51,7 +51,7 @@ def get_tmpfile(tmpdir=None, prefix=""): if tmpdir: prefix = os.path.join(tmpdir, os.path.basename(prefix)) - fd, tmp_file = tempfile.mkstemp(prefix=prefix) + fd, tmp_file = tempfile.mkstemp(prefix=prefix, suffix=suffix) os.close(fd) return tmp_file @@ -120,7 +120,7 @@ def run_command(cmd, stream=False, check_output=False, return_code=0): If check_output is True, check against an expected return code. """ stdout = subprocess.PIPE if not stream else None - output = subprocess.Popen(cmd, stderr=subprocess.STDOUT, stdout=stdout) + output = subprocess.Popen(cmd, stderr=subprocess.STDOUT, stdout=stdout, env=os.environ.copy()) t = output.communicate()[0], output.returncode output = {"message": t[0], "return_code": t[1]}