log_output : moved from os.fork to multiprocessing.Process

This commit is contained in:
alalazo 2016-07-12 17:13:41 +02:00
parent 440e71ff13
commit 9af964a6d6

View File

@ -29,6 +29,7 @@
import re
import select
import inspect
import multiprocessing
import llnl.util.tty as tty
import llnl.util.tty.color as color
@ -103,21 +104,21 @@ class log_output(object):
"""Redirects output and error of enclosed block to a file.
Usage:
with log_output(open('logfile.txt', 'w')):
with log_output('logfile.txt', 'w'):
# do things ... output will be logged.
or:
with log_output(open('logfile.txt', 'w'), echo=True):
with log_output('logfile.txt', echo=True):
# do things ... output will be logged
# and also printed to stdout.
Closes the provided stream when done with the block.
Opens a stream in 'w' mode at instance creation and closes
it at instance deletion
If echo is True, also prints the output to stdout.
"""
def __init__(self, stream, echo=False, force_color=False, debug=False):
self.stream = stream
# various output options
def __init__(self, filename, echo=False, force_color=False, debug=False):
self.filename = filename
# Various output options
self.echo = echo
self.force_color = force_color
self.debug = debug
@ -125,36 +126,32 @@ def __init__(self, stream, echo=False, force_color=False, debug=False):
# Default is to try file-descriptor reassignment unless the system
# out/err streams do not have an associated file descriptor
self.directAssignment = False
self.read, self.write = os.pipe()
def trace(self, frame, event, arg):
"""Jumps to __exit__ on the child process."""
raise _SkipWithBlock()
# Spawn a daemon that writes what it reads from a pipe
self.p = multiprocessing.Process(target=self._forward_redirected_pipe, args=(self.read,), name='logger_daemon')
self.p.daemon = True
self.p.start()
def __enter__(self):
"""Redirect output from the with block to a file.
This forks the with block as a separate process, with stdout
and stderr redirected back to the parent via a pipe. If
echo is set, also writes to standard out.
Hijacks stdout / stderr and writes to the pipe
connected to the logger daemon
"""
# remember these values for later.
self._force_color = color._force_color
self._debug = tty._debug
# Redirect this output to a pipe
self._redirect_to_pipe(self.write)
read, write = os.pipe()
self.pid = os.fork()
if self.pid:
def _forward_redirected_pipe(self, read):
# Parent: read from child, skip the with block.
os.close(write)
read_file = os.fdopen(read, 'r', 0)
with self.stream as log_file:
with open(self.filename, 'w') as log_file:
with keyboard_input(sys.stdin):
while True:
rlist, w, x = select.select([read_file, sys.stdin], [], [])
rlist, _, _ = select.select([read_file, sys.stdin], [], [])
if not rlist:
break
@ -164,7 +161,7 @@ def __enter__(self):
if sys.stdin.read(1) == 'v':
self.echo = not self.echo
# handle output from the with block process.
# Handle output from the with block process.
if read_file in rlist:
line = read_file.readline()
if not line:
@ -176,19 +173,9 @@ def __enter__(self):
# Stripped output to log file.
log_file.write(_strip(line))
log_file.flush()
read_file.flush()
read_file.close()
# Set a trace function to skip the with block.
sys.settrace(lambda *args, **keys: None)
frame = inspect.currentframe(1)
frame.f_trace = self.trace
else:
# Child: redirect output, execute the with block.
os.close(read)
def _redirect_to_pipe(self, write):
try:
# Save old stdout and stderr
self._stdout = os.dup(sys.stdout.fileno())
@ -204,54 +191,27 @@ def __enter__(self):
output_redirect = os.fdopen(write, 'w')
sys.stdout = output_redirect
sys.stderr = output_redirect
if self.force_color:
color._force_color = True
if self.debug:
tty._debug = True
def __exit__(self, exc_type, exception, traceback):
"""Exits on child, handles skipping the with block on parent."""
# Child should just exit here.
if self.pid == 0:
"""Plugs back the original file descriptors
for stdout and stderr
"""
# Flush the log to disk.
sys.stdout.flush()
sys.stderr.flush()
if exception:
# Restore stdout on the child if there's an exception,
# and let it be raised normally.
#
# This assumes that even if the exception is caught,
# the child will exit with a nonzero return code. If
# it doesn't, the child process will continue running.
#
# TODO: think about how this works outside install.
# TODO: ideally would propagate exception to parent...
if self.directAssignment:
sys.stdout = self._stdout
sys.stderr = self._stderr
else:
os.dup2(self._stdout, sys.stdout.fileno())
os.dup2(self._stderr, sys.stderr.fileno())
return False
else:
# Die quietly if there was no exception.
os._exit(0)
else:
# If the child exited badly, parent also should exit.
pid, returncode = os.waitpid(self.pid, 0)
if returncode != 0:
os._exit(1)
# restore output options.
color._force_color = self._force_color
tty._debug = self._debug
# Suppresses exception if it's our own.
return exc_type is _SkipWithBlock
def __del__(self):
"""Closes the pipes and joins the daemon"""
os.close(self.write)
self.p.join()
os.close(self.read)