persist environment

Signed-off-by: vsoch <[email protected]>

jobspec/transformer/flux/steps.py

@@ -85,10 +85,10 @@ def write_job_script(self, command):
         command = command.strip()
         match = re.match("#!/bin/(?P<executable>bash|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 []:

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]}