from contextlib import redirect_stdout, redirect_stderr
from functools import wraps
from io import StringIO
from pathlib import Path
import shutil
import sys
import openmdao.core.component
[docs]
def get_storage_directory(
component,
storage_type: str = "logs",
get_iter: bool = False,
clean: bool = False,
):
"""
Get a storage directory for the component constructed here.
Take a component and create a storage directory (for, e.g. logs or init
files), mirroring the OpenMDAO model structure as subdirectories, returning
a pathlib.Path to the storage directory.
Parameters
----------
component : openmdao.core.Component
an OpenMDAO component for which we want to create a storage directory
storage_type : str, optional
the type of storage sub-directory to make, by default "logs"
get_iter : bool, optional
should the storage directory tree be given an iteration subdirectory, by
default False
clean : bool, optional
should the directory tree, if it already exists, be cleaned out, by
default False
Returns
-------
pathlib.Path
the path to the storage subdirectory created
"""
# the storage type we're doing (logs, discipline scripts, etc.)
storage_dir = [
storage_type,
]
# if there's an iteration number to grab, grab it and add it to the dir
iter = extract_iter(component) if get_iter else None
if iter:
storage_dir += [f"iter_{iter:04d}"]
# mirror the comp path for a log directory
subdir_logger = component.pathname.split(".")
# find the reports directory
dir_reports = Path(component._problem_meta["reports_dir"])
# put the storage directory next to it
path_storage = Path(dir_reports.parent, *storage_dir, *subdir_logger)
# make a clean log location for this component if permitted
try:
path_storage.mkdir(parents=True, exist_ok=False)
except FileExistsError: # handle a FileExists, but raise anything else
if clean:
shutil.rmtree(path_storage, ignore_errors=True)
path_storage.mkdir(parents=True, exist_ok=True)
else:
raise
return path_storage
[docs]
def name_create_log(component, iter: int = None):
"""
For a given component, clean and create component- and rank-unique logfiles.
Take a component and create logs, parallel to the reports file, mirroring
the OpenMDAO model structure with stdout and stderr files for each rank,
and finally return the file paths for the component to redirect stdout and
stderr to.
Parameters
----------
component : openmdao.core.component.Component
An OpenMDAO component that we want to capture stdout/stderr for
Returns
-------
pathlib.Path
a path to in the log system to dump stdout to
pathlib.Path
a path to in the log system to dump err to
"""
# make sure we are dealing with an OM component
if not isinstance(component, openmdao.core.component.Component):
raise TypeError(
f"Expected openmdao.core.component.Component, got {type(component)}"
)
path_logfile_template = (
get_storage_directory(component, "logs", True, clean=True)
/ f"%s_rank{component._comm.rank:03d}.txt"
)
path_logfile_stdout = Path(path_logfile_template.as_posix() % "stdout")
path_logfile_stderr = Path(path_logfile_template.as_posix() % "stderr")
# return stdout and stderr files
return path_logfile_stdout.absolute(), path_logfile_stderr.absolute()
[docs]
def component_log_capture(compute_func, iter: int = None):
"""
Decorator that redirects stdout and stderr to component-wise and rank-wise logfiles.
This decorator will redirect stdout and stderr to component-wise and
rank-wise logfiles, which are determined by the `name_create_log` function.
The decorator uses context managers to redirect output streams to these
files, ensuring that all print statements and errors within the function are
logged appropriately.
func : Callable
The function to be decorated. It should be a method of a class, as
`self` is expected as the first argument.
Callable
The wrapped function with stdout and stderr redirected to log files
during its execution.
"""
@wraps(compute_func)
def wrapper(self, *args, **kwargs):
# extract from modeling options the stdio_capture option iff it exists
stdio_capture = getattr(self, "modeling_options", {}).get("stdio_capture")
# bail out, returning the function w/ no changes if it doesn't
if not stdio_capture:
return compute_func(self, *args, **kwargs)
# if we get here, we want to capture stdio
# get log file paths
path_stdout_log, path_stderr_log = name_create_log(self)
try:
# use context manager to redirect stdout & stderr
with (
open(path_stdout_log, "a") as stdout_file,
open(path_stderr_log, "a") as stderr_file,
redirect_stdout(stdout_file),
redirect_stderr(stderr_file),
):
return compute_func(self, *args, **kwargs)
except Exception:
raise # make sure the exception is raised
return wrapper
[docs]
def prepend_tabs_to_stdio(func, tabs=1):
@wraps(func)
def wrapper(*args, **kwargs):
old_stdout = sys.stdout
sys.stdout = StringIO()
# run the function
returns = func(*args, **kwargs)
# get capture output and restore stdout
output = sys.stdout.getvalue()
sys.stdout = old_stdout
prefix = "\t"
tabset = "".join([prefix for t in range(tabs)])
if output:
for line in output.splitlines():
print(f"{tabset}{line}")
# pass through the returns of the function
return returns
return wrapper
