#!/usr/bin/env python
"""
This script checks that the options in ``manager.yml``, ``scheduler.yml``,
and the environment on the local machine are properly configured.
"""
from __future__ import annotations
import sys
import os
import argparse
import abipy.flowtk as flowtk
import abipy.data as abidata
import abipy.tools.cli_parsers as cli
from monty import termcolor
from monty.termcolor import cprint
from monty.functools import prof_main
from abipy import abilab
[docs]
def show_managers(options):
"""
Print table with manager files provided by AbiPy.
"""
from tabulate import tabulate
table = []
root = os.path.join(abidata.dirpath, "managers")
yaml_paths = [os.path.join(root, f) for f in os.listdir(root) if f.endswith(".yml") and "_manager" in f]
for path in yaml_paths:
manager = flowtk.TaskManager.from_file(path)
hostname = os.path.basename(path).split("_")[0]
table.append([hostname, manager.qadapter.QTYPE, path])
if options.verbose > 1:
print(manager)
print(tabulate(table, headers=["hostname", "queue-type", "filepath"], tablefmt="rst"))
return 0
[docs]
def get_epilog() -> str:
return """\
Usage example:
abicheck.py ==> Test abipy installation and requirements.
abicheck.py -m ==> Print table with manager files provided by AbiPy.
abicheck.py -c ==> Install Yaml configuration files for manager and scheduler in ~/.abinit/abipy dir.
abicheck.py --with-flow ==> Consistency check followed by execution of AbiPy flow.
"""
[docs]
def get_parser(with_epilog=False):
parser = argparse.ArgumentParser(epilog=get_epilog() if with_epilog else "",
formatter_class=argparse.RawDescriptionHelpFormatter)
parser.add_argument('--loglevel', default="ERROR", type=str,
help="Set the loglevel. Possible values: CRITICAL, ERROR (default), WARNING, INFO, DEBUG")
parser.add_argument('-V', '--version', action='version', version=abilab.__version__)
parser.add_argument('-v', '--verbose', default=0, action='count', # -vv --> verbose=2
help='verbose, can be supplied multiple times to increase verbosity.')
parser.add_argument('--no-colors', default=False, action="store_true", help='Disable ASCII colors.')
parser.add_argument('--with-flow', default=False, action="store_true", help='Build and run small abipy flow for testing.')
parser.add_argument("-d", '--flow-dir', type=str, default=None,
help='Create AbiPy flow in this directory. If None, a default directory is used,')
parser.add_argument("-m", '--show-managers', default=False, action="store_true",
help="Print table with manager files provided by AbiPy.")
parser.add_argument("-c", '--create-config', default=False, action="store_true",
help="Create yaml configuration files in ~/abinit/.abipy with predefined settings.")
parser.add_argument("-f", '--force-reinstall', default=False, action="store_true",
help="Overwrite yaml configuration files if --create-config and files already exist.")
return parser
[docs]
@prof_main
def main():
def show_examples_and_exit(err_msg=None, error_code=1):
"""Display the usage of the script."""
sys.stderr.write(get_epilog())
if err_msg:
sys.stderr.write("Fatal Error\n" + err_msg + "\n")
sys.exit(error_code)
parser = get_parser(with_epilog=True)
# Parse the command line.
try:
options = parser.parse_args()
except Exception:
show_examples_and_exit(error_code=1)
cli.set_loglevel(options.loglevel)
if options.no_colors:
# Disable colors
termcolor.enable(False)
if options.show_managers:
# Show table with manager configuration files.
return show_managers(options)
if options.create_config:
# Install Yaml configuration files for manager and scheduler.
abilab.install_config_files(workdir=None, force_reinstall=options.force_reinstall)
errmsg = abilab.abicheck(verbose=options.verbose)
if errmsg:
cprint(errmsg, "red")
cprint("TIP: Use `--show-managers` to print the manager files provided by AbiPy.\n" +
"If abicheck.py is failing because it cannot find the manager.yml configuration file",
"yellow")
return 2
else:
cprint("\nAbipy requirements are properly configured\n", "green")
if not options.with_flow:
return 0
retcode = run_flow(options)
if retcode == 0:
cprint("\nTest flow completed successfully\n", "green")
return retcode
[docs]
def run_flow(options):
"""Run test flow, return exit code."""
import tempfile
workdir = tempfile.mkdtemp(dir=options.flow_dir)
cprint("Running small flow in workdir: %s" % workdir, "yellow")
print()
# Get the SCF and the NSCF input.
scf_input, nscf_input = make_scf_nscf_inputs()
# Build the flow.
flow = flowtk.bandstructure_flow(workdir, scf_input, nscf_input, manager=None)
flow.build_and_pickle_dump()
scheduler = flow.make_scheduler()
retcode = scheduler.start()
if retcode != 0:
cprint("Scheduler returned retcode %s" % retcode, "red")
return retcode
flow.show_status()
if not flow.all_ok:
cprint("Not all tasks in flow reached all_ok", "red")
retcode = 1
return retcode
if __name__ == "__main__":
sys.exit(main())