run_adf_diag

#!/usr/bin/env python3

"""
script: run_adf_diag

goal:  To allow a user to easily run the CAM diagnostics package from
       the command line.

Inputs:  The ADF diagnostics config file (in YAML format).  If the file is
         located in a different directory then the system path must also be
         included.

Any problems or issues with this script should be posted on the
ADF Github Discussions page located online here:

https://github.com/NCAR/ADF/discussions

Please note that registration may be required before a message can
be posted.  However, feel free to search the forums for similar issues
(and possible solutions) without needing to register or sign in.

Good luck, and may all your plots be helpful!
"""

#++++++++++++++++++++++++++++++
#Import standard python modules
#++++++++++++++++++++++++++++++

import os
import os.path
import sys
import argparse

#+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
#Cap BLAS/OpenMP thread pools BEFORE NumPy (and thus OpenBLAS/MKL) is imported.
#
#ADF parallelizes across variables with a multiprocessing Pool.  Each worker
#imports NumPy, and the underlying BLAS/OpenMP libraries then start one thread
#per CPU core by default.  On many-core nodes (e.g. 128 cores) that becomes
#(workers x cores) threads, which quickly exceeds the per-user process limit
#(RLIMIT_NPROC) and fails with:
#    OpenBLAS blas_thread_init: pthread_create failed ...
#    Resource temporarily unavailable
#The process-level parallelism is already enough, so pin each worker's
#BLAS/OpenMP pool to a single thread.  These MUST be set before NumPy is
#imported to take effect.  setdefault() means an explicit user setting (e.g.
#`export OPENBLAS_NUM_THREADS=4`) still wins.
#+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
for _thread_var in ("OPENBLAS_NUM_THREADS", "OMP_NUM_THREADS",
                    "MKL_NUM_THREADS", "NUMEXPR_NUM_THREADS",
                    "VECLIB_MAXIMUM_THREADS"):
    os.environ.setdefault(_thread_var, "1")

#+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
#Force line-buffered stdout/stderr.
#
#When ADF output is redirected to a file (e.g. `run_adf_diag ... &> log`),
#Python block-buffers stdout, so progress messages accumulate in memory and
#are only written in large chunks -- and if the process is killed (e.g. an
#OOM kill on a memory-capped node), the un-flushed buffer is LOST, making the
#log stop well before the point where the run actually died.  Line-buffering
#flushes on every newline, so the log always reflects the true last step,
#which is essential for diagnosing where a run failed.  This makes the
#behaviour independent of the PYTHONUNBUFFERED environment variable.
#+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
try:
    sys.stdout.reconfigure(line_buffering=True)
    sys.stderr.reconfigure(line_buffering=True)
except (AttributeError, ValueError):
    #reconfigure requires Python 3.7+ and a regular text stream; if it is not
    #available, fall back silently (output is simply buffered as before).
    pass

#+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
#Collapse a noisy (harmless) dask/xarray performance advisory.
#
#Opening time-series/climo files with chunks={"time": N} whose boundaries do not
#align with the on-disk chunks triggers a "specified chunks separate the stored
#chunks ..." UserWarning -- once per opened file, so it floods the log.  It does
#NOT affect correctness, only I/O performance, so show it just once.
#+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
import warnings
warnings.filterwarnings("once",
                        message="The specified chunks separate the stored chunks")

#+++++++++++++++++++++++++++++
#Import ADF diagnostics module
#+++++++++++++++++++++++++++++

#Determine local directory path:
_LOCAL_PATH = os.path.dirname(os.path.abspath(__file__))

#Add lib directory to path:
_DIAG_LIB_PATH = os.path.join(_LOCAL_PATH,"lib")

#Check that path actually exists:
if os.path.isdir(_DIAG_LIB_PATH):
    #If true, then add to Python path:
    sys.path.append(_DIAG_LIB_PATH)
else:
    #If not, then raise error:
    raise FileNotFoundError("'./lib' directory not found. Has 'run_adf_diag' been moved?")

#Import ADF diagnostics object:
from adf_diag import AdfDiag

#Import ADF diagnostics error class:
from adf_base import AdfError

#################
#Helper functions
#################

#++++++++++++++++++++++++++++++
#Input argument parser function
#++++++++++++++++++++++++++++++

def parse_arguments():

    """
    Parses command-line input arguments using the argparse
    python module and outputs the final argument object.
    """

    #Create parser object:
    parser = argparse.ArgumentParser(description='Command-line wrapper to run the ADF diagnostics package.')

    #Add input arguments to be parsed:
    #--------------------------------

    #Config file:
    parser.add_argument('configure_file', nargs='?', action='store', type=str,
                        help="YAML file used to configure CAM diagnostics.")

    #Flag method to point to config file.  Will be removed at some point in the future.
    parser.add_argument('--config-file', '--config_file', metavar='<DIAG_YAML_FILE>', action='store', type=str,
                        help="YAML file used to configure CAM diagnostics (deprecated).")

    #Debug setting:
    parser.add_argument('--debug', action='store_true', help="Turn on debug output.")

    #--------------------------------

    #Parse Argument inputs
    args = parser.parse_args()

    #If no config file argument is present, then throw an error:
    if (not args.configure_file) and (not args.config_file):
        emsg = "No Config file found, please run 'run_diag <DIAG_YAML_FILE>'.\n"
        emsg +="Where <DIAG_YAML_FILE> is the name of the yaml file used to configure"
        emsg +="the diagnostics package."
        parser.error(emsg)

    return args

############################
#Main CAM diagnostics script
############################

#Run code below if command is called
#directly from the command line:
if __name__ == "__main__":

    #+++++++++++++++++++++
    #Begin ADF diag script
    #+++++++++++++++++++++
    print('ADF diagnostics is starting...')

    #+++++++++++++++++++++++++++++++++++++++++++
    #Check that python is version 3.6 or greater
    #+++++++++++++++++++++++++++++++++++++++++++

    if sys.version_info[0] < 3:
        raise AdfError("Script only works with Python 3. Please switch python versions.")

    if sys.version_info[1] < 6:
        raise AdfError("Script only works with Python version 3.6 or greater.  Please update python.")

    #++++++++++++++++++++++++++++
    #Parse command-line arguments
    #++++++++++++++++++++++++++++

    args = parse_arguments()

    #Extract YAML config file name/path:
    if args.configure_file:
      config_yaml = args.configure_file
    else:
      config_yaml = args.config_file

    #Extract debug flag:
    config_debug = args.debug

    #+++++++++++++++++++++++++++++++++
    #Call main ADF diagnostics methods
    #+++++++++++++++++++++++++++++++++

    #Initalize CAM diagnostics object:
    diag = AdfDiag(config_yaml, debug=config_debug)

    #Create model time series.
    #Please note that this is an internal ADF function:
    diag.create_time_series()

    #Create model baseline time series (if needed):
    if not diag.compare_obs:
        diag.create_time_series(baseline=True)

    #Optionally generate FULL-RANGE time series (over the full available years,
    #2-D variables only) for the global-mean time series drift plot. This runs
    #only when the 'global_mean_ts_full_range' toggle is set AND
    #'global_mean_timeseries' is requested in 'plotting_scripts'. The full-range
    #pass skips any case whose full range already equals its climo range.
    _plot_scripts = diag.read_config_var("plotting_scripts") or []
    _gm_requested = any(
        ("global_mean_timeseries" in s) if isinstance(s, dict)
        else (s == "global_mean_timeseries")
        for s in _plot_scripts
    )
    if diag.get_basic_info("global_mean_ts_full_range") and _gm_requested:
        diag.create_time_series(full_range=True)
        if not diag.compare_obs:
            diag.create_time_series(baseline=True, full_range=True)

    #Call the CVDP:
    if diag.get_cvdp_info('cvdp_run'):
       diag.setup_run_cvdp()

    #Call the MDTF:
    if diag.get_mdtf_info('mdtf_run'):
        mdtf_proc = diag.setup_run_mdtf()  #returns mdtf_proc for subprocess control
    else:
        mdtf_proc = None

    #Create model climatology (climo) files.
    #This call uses the "time_averaging_scripts" specified
    #in the config file:
    diag.create_climo()

    #If a user is doing a model vs obs comparison, but
    #no observations were found, then stop here:
    if diag.compare_obs and not diag.var_obs_dict:
        print('\nADF diagnostics has completed successfully.')
        sys.exit(0)

    #Regrid model climatology files to match either
    #observations or CAM baseline climatologies.
    #This call uses the "regridding_scripts" specified
    #in the config file:
    diag.regrid_climo()

    #Perform analyses on the simulation(s).
    #This call uses the "analysis_scripts" specified in the
    #config file:
    diag.perform_analyses()

    #Create plots.
    #This call uses the "plotting_scripts" specified
    #in the config file:
    diag.create_plots()

    #Create website.
    #Please note that this is an internal ADF function:
    if diag.create_html:
        diag.create_website()

    # Check if sub-processes are still running (CVDP, MDTF)
    if mdtf_proc:
        mdtf_status = mdtf_proc.wait(timeout=None)
        if (mdtf_status != 0):
            print(f"ERROR: MDTF finished with code {mdtf_status}")
        else:
            print("MDTF finished successfully")
        
    #+++++++++++++++
    #End diag script
    #+++++++++++++++
    print('\nADF diagnostics has completed successfully.')
    sys.exit(0)