Source code for yt.startup_tasks

# This handles the command line.

import argparse
import os
import signal
import sys

from yt.config import ytcfg
from yt.funcs import (
    mylog,
    paste_traceback,
    paste_traceback_detailed,
    signal_ipython,
    signal_print_traceback,
)
from yt.utilities import rpdb

exe_name = os.path.basename(sys.executable)


# At import time, we determined whether or not we're being run in parallel.
[docs] def turn_on_parallelism(): parallel_capable = False try: # we import this to check if mpi4py is installed from mpi4py import MPI # NOQA except ImportError as e: mylog.error( "Warning: Attempting to turn on parallelism, " "but mpi4py import failed. Try pip install mpi4py." ) raise e # Now we have to turn on the parallelism from the perspective of the # parallel_analysis_interface from yt.utilities.parallel_tools.parallel_analysis_interface import ( enable_parallelism, ) parallel_capable = enable_parallelism() return parallel_capable
# This fallback is for Paraview: # We use two signals, SIGUSR1 and SIGUSR2. In a non-threaded environment, # we set up handlers to process these by printing the current stack and to # raise a RuntimeError. The latter can be used, inside pdb, to catch an error # and then examine the current stack. try: signal.signal(signal.SIGUSR1, signal_print_traceback) mylog.debug("SIGUSR1 registered for traceback printing") signal.signal(signal.SIGUSR2, signal_ipython) mylog.debug("SIGUSR2 registered for IPython Insertion") except (ValueError, RuntimeError, AttributeError): # Not in main thread pass
[docs] class SetExceptionHandling(argparse.Action): def __call__(self, parser, namespace, values, option_string=None): # If we recognize one of the arguments on the command line as indicating a # different mechanism for handling tracebacks, we attach one of those handlers # and remove the argument from sys.argv. # if self.dest == "paste": sys.excepthook = paste_traceback mylog.debug("Enabling traceback pasting") elif self.dest == "paste-detailed": sys.excepthook = paste_traceback_detailed mylog.debug("Enabling detailed traceback pasting") elif self.dest == "detailed": import cgitb cgitb.enable(format="text") mylog.debug("Enabling detailed traceback reporting") elif self.dest == "rpdb": sys.excepthook = rpdb.rpdb_excepthook mylog.debug("Enabling remote debugging")
[docs] class SetConfigOption(argparse.Action): def __call__(self, parser, namespace, values, option_string=None): param, val = values.split("=") mylog.debug("Overriding config: %s = %s", param, val) ytcfg["yt", param] = val if param == "log_level": # special case mylog.setLevel(int(val))
[docs] class YTParser(argparse.ArgumentParser):
[docs] def error(self, message): """error(message: string) Prints a help message that is more detailed than the argparse default and then exits. """ self.print_help(sys.stderr) self.exit(2, f"{self.prog}: error: {message}\n")
parser = YTParser(description="yt command line arguments") parser.add_argument( "--config", action=SetConfigOption, help="Set configuration option, in the form param=value", ) parser.add_argument( "--paste", action=SetExceptionHandling, help="Paste traceback to paste.yt-project.org", nargs=0, ) parser.add_argument( "--paste-detailed", action=SetExceptionHandling, help="Paste a detailed traceback with local variables to " + "paste.yt-project.org", nargs=0, ) parser.add_argument( "--detailed", action=SetExceptionHandling, help="Display detailed traceback.", nargs=0, ) parser.add_argument( "--rpdb", action=SetExceptionHandling, help="Enable remote pdb interaction (for parallel debugging).", nargs=0, ) parser.add_argument( "--parallel", action="store_true", default=False, dest="parallel", help="Run in MPI-parallel mode (must be launched as an MPI task)", ) if not hasattr(sys, "argv") or sys.argv is None: sys.argv = [] unparsed_args: list[str] = [] parallel_capable = False if not ytcfg.get("yt", "internals", "command_line"): opts, unparsed_args = parser.parse_known_args() # THIS IS NOT SUCH A GOOD IDEA: # sys.argv = [a for a in unparsed_args] if opts.parallel: parallel_capable = turn_on_parallelism() subparsers = parser.add_subparsers( title="subcommands", dest="subcommands", description="Valid subcommands", ) else: subparsers = parser.add_subparsers( title="subcommands", dest="subcommands", description="Valid subcommands", ) def print_help(*args, **kwargs): parser.print_help() help_parser = subparsers.add_parser("help", help="Print help message") help_parser.set_defaults(func=print_help) if parallel_capable: pass elif ( exe_name in ["mpi4py", "embed_enzo", "python{}.{}-mpi".format(*sys.version_info)] or "_parallel" in dir(sys) or any("ipengine" in arg for arg in sys.argv) or any("cluster-id" in arg for arg in sys.argv) ): parallel_capable = turn_on_parallelism() else: parallel_capable = False