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 re
import select import select
import inspect import inspect
import multiprocessing
import llnl.util.tty as tty import llnl.util.tty as tty
import llnl.util.tty.color as color import llnl.util.tty.color as color
@ -103,155 +104,114 @@ class log_output(object):
"""Redirects output and error of enclosed block to a file. """Redirects output and error of enclosed block to a file.
Usage: Usage:
with log_output(open('logfile.txt', 'w')): with log_output('logfile.txt', 'w'):
# do things ... output will be logged. # do things ... output will be logged.
or: or:
with log_output(open('logfile.txt', 'w'), echo=True): with log_output('logfile.txt', echo=True):
# do things ... output will be logged # do things ... output will be logged
# and also printed to stdout. # 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. If echo is True, also prints the output to stdout.
""" """
def __init__(self, stream, echo=False, force_color=False, debug=False): def __init__(self, filename, echo=False, force_color=False, debug=False):
self.stream = stream self.filename = filename
# Various output options
# various output options
self.echo = echo self.echo = echo
self.force_color = force_color self.force_color = force_color
self.debug = debug self.debug = debug
# Default is to try file-descriptor reassignment unless the system # Default is to try file-descriptor reassignment unless the system
# out/err streams do not have an associated file descriptor # out/err streams do not have an associated file descriptor
self.directAssignment = False self.directAssignment = False
self.read, self.write = os.pipe()
def trace(self, frame, event, arg): # Spawn a daemon that writes what it reads from a pipe
"""Jumps to __exit__ on the child process.""" self.p = multiprocessing.Process(target=self._forward_redirected_pipe, args=(self.read,), name='logger_daemon')
raise _SkipWithBlock() self.p.daemon = True
self.p.start()
def __enter__(self): def __enter__(self):
"""Redirect output from the with block to a file. """Redirect output from the with block to a file.
This forks the with block as a separate process, with stdout Hijacks stdout / stderr and writes to the pipe
and stderr redirected back to the parent via a pipe. If connected to the logger daemon
echo is set, also writes to standard out.
""" """
# remember these values for later. # remember these values for later.
self._force_color = color._force_color self._force_color = color._force_color
self._debug = tty._debug self._debug = tty._debug
# Redirect this output to a pipe
self._redirect_to_pipe(self.write)
read, write = os.pipe() def _forward_redirected_pipe(self, read):
# Parent: read from child, skip the with block.
read_file = os.fdopen(read, 'r', 0)
with open(self.filename, 'w') as log_file:
with keyboard_input(sys.stdin):
while True:
rlist, _, _ = select.select([read_file, sys.stdin], [], [])
if not rlist:
break
self.pid = os.fork() # Allow user to toggle echo with 'v' key.
if self.pid: # Currently ignores other chars.
# Parent: read from child, skip the with block. if sys.stdin in rlist:
os.close(write) if sys.stdin.read(1) == 'v':
self.echo = not self.echo
read_file = os.fdopen(read, 'r', 0) # Handle output from the with block process.
with self.stream as log_file: if read_file in rlist:
with keyboard_input(sys.stdin): line = read_file.readline()
while True: if not line:
rlist, w, x = select.select([read_file, sys.stdin], [], [])
if not rlist:
break break
# Allow user to toggle echo with 'v' key. # Echo to stdout if requested.
# Currently ignores other chars. if self.echo:
if sys.stdin in rlist: sys.stdout.write(line)
if sys.stdin.read(1) == 'v':
self.echo = not self.echo
# handle output from the with block process. # Stripped output to log file.
if read_file in rlist: log_file.write(_strip(line))
line = read_file.readline() log_file.flush()
if not line:
break
# Echo to stdout if requested. def _redirect_to_pipe(self, write):
if self.echo: try:
sys.stdout.write(line) # Save old stdout and stderr
self._stdout = os.dup(sys.stdout.fileno())
# Stripped output to log file. self._stderr = os.dup(sys.stderr.fileno())
log_file.write(_strip(line))
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)
try:
# Save old stdout and stderr
self._stdout = os.dup(sys.stdout.fileno())
self._stderr = os.dup(sys.stderr.fileno())
# redirect to the pipe.
os.dup2(write, sys.stdout.fileno())
os.dup2(write, sys.stderr.fileno())
except AttributeError:
self.directAssignment = True
self._stdout = sys.stdout
self._stderr = sys.stderr
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
# redirect to the pipe.
os.dup2(write, sys.stdout.fileno())
os.dup2(write, sys.stderr.fileno())
except AttributeError:
self.directAssignment = True
self._stdout = sys.stdout
self._stderr = sys.stderr
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): def __exit__(self, exc_type, exception, traceback):
"""Exits on child, handles skipping the with block on parent.""" """Plugs back the original file descriptors
# Child should just exit here. for stdout and stderr
if self.pid == 0: """
# Flush the log to disk. # Flush the log to disk.
sys.stdout.flush() sys.stdout.flush()
sys.stderr.flush() sys.stderr.flush()
os.dup2(self._stdout, sys.stdout.fileno())
if exception: os.dup2(self._stderr, sys.stderr.fileno())
# 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. # restore output options.
color._force_color = self._force_color color._force_color = self._force_color
tty._debug = self._debug tty._debug = self._debug
# Suppresses exception if it's our own. def __del__(self):
return exc_type is _SkipWithBlock """Closes the pipes and joins the daemon"""
os.close(self.write)
self.p.join()
os.close(self.read)