3988894829

Build Type

github-actions

Committed by GitHub

Commit Message

Merge 89db93e8c into da1060a80

Pull Request Pull Request #190: Replace setup.py plugins with stand-alone scripts

Run Details

199 of 199 new or added lines in 7 files covered. (100.0%)

2010 of 2644 relevant lines covered (76.02%)

0.76 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

78.52

/py/desiutil/redirect.py

# 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


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)


@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

1	# Licensed under a 3-clause BSD style license - see LICENSE.rst
2	# -- coding: utf-8 --
3	"""	1✔
4	=================
5	desiutil.redirect
6	=================
7
8	Utilities for redirecting stdout / stderr to files.
9
10	"""
11
12	import os	1✔
13	import sys	1✔
14	import time	1✔
15	import io	1✔
16	import traceback	1✔
17	import logging	1✔
18	import ctypes	1✔
19
20	from contextlib import contextmanager	1✔
21
22	from .log import get_logger, _desiutil_log_root	1✔
23
24
25	# C file descriptors for stderr and stdout, used in redirection
26	# context manager.
27
28	_libc = None	1✔
29	_c_stdout = None	1✔
30	_c_stderr = None	1✔
31
32
33	def _get_libc():	1✔
34	"""Helper function to import libc once."""
35	global _libc
36	global _c_stdout
37	global _c_stderr
38	if _libc is None:	1✔
39	_libc = ctypes.CDLL(None)	1✔
40	try:	1✔
41	# Linux systems
42	_c_stdout = ctypes.c_void_p.in_dll(_libc, "stdout")	1✔
43	_c_stderr = ctypes.c_void_p.in_dll(_libc, "stderr")	1✔
44	except ValueError:	1✔
45	try:	1✔
46	# Darwin
47	_c_stdout = ctypes.c_void_p.in_dll(_libc, "__stdoutp")	1✔
48	_c_stderr = ctypes.c_void_p.in_dll(_libc, "__stderrp")	1✔
49	except ValueError:	1✔
50	# Neither!
51	pass	1✔
52	return (_libc, _c_stdout, _c_stderr)	1✔
53
54
55	@contextmanager	1✔
56	def stdouterr_redirected(to=None, comm=None):	1✔
57	"""Redirect stdout and stderr to a file.
58
59	The general technique is based on:
60
61	http://stackoverflow.com/questions/5081657
62	http://eli.thegreenplace.net/2015/redirecting-all-kinds-of-stdout-in-python/
63
64	If the optional communicator is specified, then each process redirects to
65	a different temporary file. Upon exit from the context the rank zero
66	process concatenates these in order to the final file result.
67
68	If the enclosing code raises an exception, the traceback is printed to the
69	log file.
70
71	Args:
72	to (str): The output file name.
73	comm (mpi4py.MPI.Comm): The optional MPI communicator.
74
75	"""
76	libc, c_stdout, c_stderr = _get_libc()	1✔
77
78	nproc = 1	1✔
79	rank = 0	1✔
80	MPI = None	1✔
81	if comm is not None:	1✔
82	# If we are already using MPI (comm is set), then we can safely
83	# import mpi4py.
84	from mpi4py import MPI	×
85
86	nproc = comm.size	×
87	rank = comm.rank	×
88
89	# The currently active POSIX file descriptors
90	fd_out = sys.stdout.fileno()	1✔
91	fd_err = sys.stderr.fileno()	1✔
92
93	# Save the original file descriptors so we can restore them later
94	saved_fd_out = os.dup(fd_out)	1✔
95	saved_fd_err = os.dup(fd_err)	1✔
96
97	# The DESI loggers.
98	desi_loggers = _desiutil_log_root	1✔
99
100	def _redirect(out_to, err_to):	1✔
101	# Flush the C-level buffers
102	if c_stdout is not None:	1✔
103	libc.fflush(c_stdout)	1✔
104	if c_stderr is not None:	1✔
105	libc.fflush(c_stderr)	1✔
106
107	# This closes the python file handles, and marks the POSIX
108	# file descriptors for garbage collection- UNLESS those
109	# are the special file descriptors for stderr/stdout.
110	sys.stdout.close()	1✔
111	sys.stderr.close()	1✔
112
113	# Close fd_out/fd_err if they are open, and copy the
114	# input file descriptors to these.
115	os.dup2(out_to, fd_out)	1✔
116	os.dup2(err_to, fd_err)	1✔
117
118	# Create a new sys.stdout / sys.stderr that points to the
119	# redirected POSIX file descriptors. In Python 3, these
120	# are actually higher level IO objects.
121	sys.stdout = io.TextIOWrapper(os.fdopen(fd_out, "wb"))	1✔
122	sys.stderr = io.TextIOWrapper(os.fdopen(fd_err, "wb"))	1✔
123
124	# update DESI logging to use new stdout
125	for name, logger in desi_loggers.items():	1✔
126	hformat = None	1✔
127	while len(logger.handlers) > 0:	1✔
128	h = logger.handlers[0]	1✔
129	if hformat is None:	1✔
130	hformat = h.formatter._fmt	1✔
131	logger.removeHandler(h)	1✔
132	# Add the current stdout.
133	ch = logging.StreamHandler(sys.stdout)	1✔
134	formatter = logging.Formatter(hformat, datefmt="%Y-%m-%dT%H:%M:%S")	1✔
135	ch.setFormatter(formatter)	1✔
136	logger.addHandler(ch)	1✔
137
138	def _open_redirect(filename):	1✔
139	# Open python file, which creates low-level POSIX file
140	# descriptor.
141	file_handle = open(filename, "wb")	1✔
142
143	# Redirect stdout/stderr to this new file descriptor.
144	_redirect(out_to=file_handle.fileno(), err_to=file_handle.fileno())	1✔
145	return file_handle	1✔
146
147	def _close_redirect(handle):	1✔
148	# Close python file handle, which will mark POSIX file descriptor for
149	# garbage collection. That is fine since we are about to overwrite those.
150	if handle is not None:	1✔
151	handle.close()	1✔
152
153	# Flush python handles for good measure
154	sys.stdout.flush()	1✔
155	sys.stderr.flush()	1✔
156
157	try:	1✔
158	# Restore old stdout and stderr
159	_redirect(out_to=saved_fd_out, err_to=saved_fd_err)	1✔
160	except Exception:	×
161	pass	×
162
163	# Redirect both stdout and stderr to the same file
164
165	if to is None:	1✔
166	to = "/dev/null"	×
167
168	if rank == 0:	1✔
169	log = get_logger(timestamp=True)	1✔
170	log.info("Begin log redirection to %s", to)	1✔
171
172	# Try to open the redirected file.
173
174	pto = to	1✔
175	if to != "/dev/null" and comm is not None:	1✔
176	pto = "{}_{}".format(to, rank)	×
177
178	fail_open = 0	1✔
179	file = None	1✔
180	try:	1✔
181	file = _open_redirect(pto)	1✔
182	except Exception:	×
183	log = get_logger()	×
184	log.error("Failed to open redirection file %s", pto)	×
185	fail_open = 1	×
186
187	if comm is not None:	1✔
188	fail_open = comm.allreduce(fail_open, op=MPI.SUM)	×
189
190	if fail_open > 0:	1✔
191	# Something went wrong on one or more processes, try to recover and exit
192	if rank == 0:	×
193	log = get_logger()	×
194	log.error("Failed to start redirect to %s", to)	×
195
196	_close_redirect(file)	×
197
198	# All processes raise an exception for the calling code to handle
199	msg = "Failed to start output redirect to {}".format(to)	×
200	raise RuntimeError(msg)	×
201
202	# Output should now be redirected. Run the code.
203
204	fail_run = 0	1✔
205	try:	1✔
206	yield # Allow code to be run with the redirected output	1✔
207	except Exception:	1✔
208	# We have an unhandled exception. Print a stack trace to the log.
209	exc_type, exc_value, exc_traceback = sys.exc_info()	1✔
210	lines = traceback.format_exception(exc_type, exc_value, exc_traceback)	1✔
211	print("".join(lines), flush=True)	1✔
212	fail_run = 1	1✔
213
214	# Check if any processes failed to run their code
215	if comm is not None:	1✔
216	fail_run = comm.allreduce(fail_run, op=MPI.SUM)	×
217
218	_close_redirect(file)	1✔
219
220	if comm is not None:	1✔
221	# Concatenate per-process files if we have multiple processes.
222	comm.barrier()	×
223	if rank == 0:	×
224	with open(to, "w") as outfile:	×
225	for p in range(nproc):	×
226	outfile.write(	×
227	"================= Process {} =================\n".format(p)
228	)
229	fname = "{}_{}".format(to, p)	×
230	with open(fname, "r") as infile:	×
231	outfile.write(infile.read())	×
232	os.remove(fname)	×
233	comm.barrier()	×
234
235	if rank == 0:	1✔
236	log = get_logger(timestamp=True)	1✔
237	log.info("End log redirection to %s", to)	1✔
238
239	# flush python handles for good measure
240	sys.stdout.flush()	1✔
241	sys.stderr.flush()	1✔
242
243	if fail_run > 0:	1✔
244	msg = "{} processes raised an exception while logs were redirected".format(	1✔
245	fail_run
246	)
247	if rank == 0:	1✔
248	log = get_logger()	1✔
249	log.error(msg)	1✔
250	raise RuntimeError(msg)	1✔
251
252	return	1✔

desihub / desiutil / 3988894829

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous