add --keep option

mpienv/command/exec.py

@@ -22,6 +22,7 @@ def main():
     idx = 0
     dry_run = False
     verbose = False
+    keep = False
     while True:
         if args[idx] == '--dry-run':
             # --dry-run is mpienv's unique option and is not passed to mpiexec.
@@ -35,10 +36,13 @@ def main():
             # it is preserved and passed to mpiexec.
             idx += 1
             verbose = True
+        elif args[idx] == '--keep':
+            keep = True
+            args.pop(idx)
         else:
             break
 
-    mpienv.exec_(args, dry_run=dry_run, verbose=verbose)
+    mpienv.exec_(args, dry_run=dry_run, verbose=verbose, keep=keep)
 
 
 if __name__ == "__main__":

mpienv/mpibase.py

@@ -269,7 +269,7 @@ def lib_files(self):
     def libexec_files(self):
         assert False, "Must be overriden"
 
-    def _generate_exec_script(self, file_name, cmds):
+    def _generate_exec_script(self, file_name, cmds, keep: bool):
         with open(file_name, 'w') as f:
             for shell in ['/bin/bash', '/bin/ash', '/bin/sh']:
                 if os.path.exists(shell):
@@ -307,11 +307,12 @@ def _generate_exec_script(self, file_name, cmds):
             f.write(' '.join(cmds) + "\n")
 
             # remove the script itself
-            f.write('rm -f {}\n'.format(file_name))
+            if not keep:
+                f.write('rm -f {}\n'.format(file_name))
 
         os.chmod(file_name, 0o744)
 
-    def exec_(self, cmds, dry_run=False, verbose=False):
+    def exec_(self, cmds, keep, dry_run, verbose):
         # Determine the temporary shell script name
         # Determine the remote hosts
         # Transfer the shell script to remote hosts
@@ -326,7 +327,7 @@ def exec_(self, cmds, dry_run=False, verbose=False):
             print("hosts = {}".format(remote_hosts))
 
         # Generate a proxy shell script that runs user programs
-        self._generate_exec_script(tempfile, cmds)
+        self._generate_exec_script(tempfile, cmds, keep)
 
         # Copy script file
         for host in remote_hosts: