diff --git a/src/watts/fileutils.py b/src/watts/fileutils.py index 97b4392d..4d52dee5 100644 --- a/src/watts/fileutils.py +++ b/src/watts/fileutils.py @@ -84,7 +84,7 @@ class tee_stderr(_TeeStream): _stream = "stderr" -def run(args): +def run(args, wait=None): """Function that mimics subprocess.run but actually writes to sys.stdout and sys.stderr (not the same as the underlying file descriptors) @@ -94,15 +94,16 @@ def run(args): p = subprocess.Popen(args, stdout=subprocess.PIPE, stderr=subprocess.PIPE, universal_newlines=True) - while True: - select.select([p.stdout, p.stderr], [], []) + select.select([p.stdout, p.stderr], [], []) - stdout_data = p.stdout.read() - stderr_data = p.stderr.read() - if stdout_data: - sys.stdout.write(stdout_data) - if stderr_data: - sys.stderr.write(stderr_data) + try: + stdout_data, stderr_data = p.communicate(timeout=wait) + except subprocess.TimeoutExpired: + p.kill() + stdout_data, stderr_data = p.communicate() - if p.poll() is not None: - break + if stdout_data: + sys.stdout.write(stdout_data) + + if stderr_data: + sys.stderr.write(stderr_data) diff --git a/src/watts/plugin.py b/src/watts/plugin.py index f3881755..a00f41af 100644 --- a/src/watts/plugin.py +++ b/src/watts/plugin.py @@ -223,7 +223,8 @@ def postrun(self, params: Parameters, name: str, **kwargs) -> Results: return results_cls(params, name, time, inputs, outputs, **kwargs) def run(self, mpi_args: Optional[List[str]] = None, - extra_args: Optional[List[str]] = None): + extra_args: Optional[List[str]] = None, + wait: Optional[int] = 60): """Run plugin Parameters @@ -234,9 +235,13 @@ def run(self, mpi_args: Optional[List[str]] = None, extra_args Additional command-line arguments to append after the main command + wait + The time in seconds that `watts` will wait before terminating the process. + Default is 60 seconds. + """ if mpi_args is None: mpi_args = [] if extra_args is None: extra_args = [] - run_proc(mpi_args + self.execute_command + extra_args) + run_proc(mpi_args + self.execute_command + extra_args, wait)