# Licensed under a 3-clause BSD style license - see LICENSE.rst
# -*- coding: utf-8 -*-
"""
=================
desiutil.redirect
=================
Utilities for redirecting stdout / stderr to files.
"""
import os
import sys
import time
import io
import traceback
import logging
import ctypes
from contextlib import contextmanager
from .log import get_logger, _desiutil_log_root
# C file descriptors for stderr and stdout, used in redirection
# context manager.
_libc = None
_c_stdout = None
_c_stderr = None
[docs]def _get_libc():
"""Helper function to import libc once."""
global _libc
global _c_stdout
global _c_stderr
if _libc is None:
_libc = ctypes.CDLL(None)
try:
# Linux systems
_c_stdout = ctypes.c_void_p.in_dll(_libc, "stdout")
_c_stderr = ctypes.c_void_p.in_dll(_libc, "stderr")
except ValueError:
try:
# Darwin
_c_stdout = ctypes.c_void_p.in_dll(_libc, "__stdoutp")
_c_stderr = ctypes.c_void_p.in_dll(_libc, "__stderrp")
except ValueError:
# Neither!
pass
return (_libc, _c_stdout, _c_stderr)
[docs]@contextmanager
def stdouterr_redirected(to=None, comm=None):
"""Redirect stdout and stderr to a file.
The general technique is based on:
http://stackoverflow.com/questions/5081657
http://eli.thegreenplace.net/2015/redirecting-all-kinds-of-stdout-in-python/
If the optional communicator is specified, then each process redirects to
a different temporary file. Upon exit from the context the rank zero
process concatenates these in order to the final file result.
If the enclosing code raises an exception, the traceback is printed to the
log file.
Args:
to (str): The output file name.
comm (mpi4py.MPI.Comm): The optional MPI communicator.
"""
libc, c_stdout, c_stderr = _get_libc()
nproc = 1
rank = 0
MPI = None
if comm is not None:
# If we are already using MPI (comm is set), then we can safely
# import mpi4py.
from mpi4py import MPI
nproc = comm.size
rank = comm.rank
# The currently active POSIX file descriptors
fd_out = sys.stdout.fileno()
fd_err = sys.stderr.fileno()
# Save the original file descriptors so we can restore them later
saved_fd_out = os.dup(fd_out)
saved_fd_err = os.dup(fd_err)
# The DESI loggers.
desi_loggers = _desiutil_log_root
def _redirect(out_to, err_to):
# Flush the C-level buffers
if c_stdout is not None:
libc.fflush(c_stdout)
if c_stderr is not None:
libc.fflush(c_stderr)
# This closes the python file handles, and marks the POSIX
# file descriptors for garbage collection- UNLESS those
# are the special file descriptors for stderr/stdout.
sys.stdout.close()
sys.stderr.close()
# Close fd_out/fd_err if they are open, and copy the
# input file descriptors to these.
os.dup2(out_to, fd_out)
os.dup2(err_to, fd_err)
# Create a new sys.stdout / sys.stderr that points to the
# redirected POSIX file descriptors. In Python 3, these
# are actually higher level IO objects.
sys.stdout = io.TextIOWrapper(os.fdopen(fd_out, "wb"))
sys.stderr = io.TextIOWrapper(os.fdopen(fd_err, "wb"))
# update DESI logging to use new stdout
for name, logger in desi_loggers.items():
hformat = None
while len(logger.handlers) > 0:
h = logger.handlers[0]
if hformat is None:
hformat = h.formatter._fmt
logger.removeHandler(h)
# Add the current stdout.
ch = logging.StreamHandler(sys.stdout)
formatter = logging.Formatter(hformat, datefmt="%Y-%m-%dT%H:%M:%S")
ch.setFormatter(formatter)
logger.addHandler(ch)
def _open_redirect(filename):
# Open python file, which creates low-level POSIX file
# descriptor.
file_handle = open(filename, "wb")
# Redirect stdout/stderr to this new file descriptor.
_redirect(out_to=file_handle.fileno(), err_to=file_handle.fileno())
return file_handle
def _close_redirect(handle):
# Close python file handle, which will mark POSIX file descriptor for
# garbage collection. That is fine since we are about to overwrite those.
if handle is not None:
handle.close()
# Flush python handles for good measure
sys.stdout.flush()
sys.stderr.flush()
try:
# Restore old stdout and stderr
_redirect(out_to=saved_fd_out, err_to=saved_fd_err)
except Exception:
pass
# Redirect both stdout and stderr to the same file
if to is None:
to = "/dev/null"
if rank == 0:
log = get_logger(timestamp=True)
log.info("Begin log redirection to %s", to)
# Try to open the redirected file.
pto = to
if to != "/dev/null" and comm is not None:
pto = "{}_{}".format(to, rank)
fail_open = 0
file = None
try:
file = _open_redirect(pto)
except Exception:
log = get_logger()
log.error("Failed to open redirection file %s", pto)
fail_open = 1
if comm is not None:
fail_open = comm.allreduce(fail_open, op=MPI.SUM)
if fail_open > 0:
# Something went wrong on one or more processes, try to recover and exit
if rank == 0:
log = get_logger()
log.error("Failed to start redirect to %s", to)
_close_redirect(file)
# All processes raise an exception for the calling code to handle
msg = "Failed to start output redirect to {}".format(to)
raise RuntimeError(msg)
# Output should now be redirected. Run the code.
fail_run = 0
try:
yield # Allow code to be run with the redirected output
except Exception:
# We have an unhandled exception. Print a stack trace to the log.
exc_type, exc_value, exc_traceback = sys.exc_info()
lines = traceback.format_exception(exc_type, exc_value, exc_traceback)
print("".join(lines), flush=True)
fail_run = 1
# Check if any processes failed to run their code
if comm is not None:
fail_run = comm.allreduce(fail_run, op=MPI.SUM)
_close_redirect(file)
if comm is not None:
# Concatenate per-process files if we have multiple processes.
comm.barrier()
if rank == 0:
with open(to, "w") as outfile:
for p in range(nproc):
outfile.write(
"================= Process {} =================\n".format(p)
)
fname = "{}_{}".format(to, p)
with open(fname, "r") as infile:
outfile.write(infile.read())
os.remove(fname)
comm.barrier()
if rank == 0:
log = get_logger(timestamp=True)
log.info("End log redirection to %s", to)
# flush python handles for good measure
sys.stdout.flush()
sys.stderr.flush()
if fail_run > 0:
msg = "{} processes raised an exception while logs were redirected".format(
fail_run
)
if rank == 0:
log = get_logger()
log.error(msg)
raise RuntimeError(msg)
return