XGC-DiffusionOptimizer.py

#!/usr/bin/env python3

import argparse
import shutil
import os
import f90nml

#import xml.etree.ElementTree as ET
import lxml.etree as ET

from effis.composition import Workflow, EffisLogger, Input
from effis.composition.runner import slurm, perlmutter, frontier


class suchyta:

    perlmutter = {
        'src' : "/global/homes/e/esuchyta/software/src/XGC-Devel",
        'build' : "/global/homes/e/esuchyta/software/build/perlmutter/XGC-Devel",
        'examples': "/global/homes/e/esuchyta/software/src/XGC-Examples",
        'analysis': "/global/homes/e/esuchyta/software/src/XGC-Analysis",
        'adios': '/global/homes/e/esuchyta/software/install/perlmutter/adios2-gcc/lib/python3.11/site-packages',
        'gpus' : 4,
        'per': 32,  # 64 cores, but with 128 threads, which Slurm counts
    }

    frontier = {
        'src' : "/ccs/home/esuchyta/software/src/XGC-Devel",
        'build' : "/ccs/home/esuchyta/software/build/frontier/XGC-Devel",
        'examples': "/ccs/home/esuchyta/software/src/XGC-Examples",
        'analysis': "/ccs/home/esuchyta/software/src/XGC-Analysis",
        'adios': '/ccs/home/esuchyta/software/install/frontier/adios2-gcc12.2.0/lib/python3.11/site-packages',
        'gpus' : 8,
        'per': 14,  # Default gives one core to system stuff
    }


def SetupArgs():

    runner = Workflow.DetectRunnerInfo()
    if (runner is not None) and (not isinstance(runner, slurm)):
        EffisLogger.RaiseError(ValueError, "Current example batch setup is for Slurm")

    slurmparser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)

    # Run Configuration
    ##############################################
    slurmparser.add_argument(
        "-o", "--outdir",
        type=str,
        help="Path to run directory",
        required=True,
    )

    slurmparser.add_argument(
        "-x", "--xgc",
        type=str,
        help="Flavor of XGC",
        required=True,
        choices=(
            'xgca-cpp',
            'xgca-cpp-gpu',
            'xgc-es-cpp',
            'xgc-es-cpp-gpu',
            'xgc-eem-cpp',
            'xgc-eem-cpp-gpu',
        )
    )

    if isinstance(runner, (perlmutter, frontier)):
        slurmparser.add_argument(
            "-e", "--example",
            type=str,
            help="Which example to run",
            required=False,
            default=os.path.join(getattr(suchyta, runner.__class__.__name__)['examples'], "ecbc_low_beta_n6"),
        )
        slurmparser.add_argument(
            "-s", "--suchyta",
            action="store_true",
            help="Use Eric Suchyta's software",
        )
    else:
        slurmparser.add_argument(
            "-e", "--example",
            type=str,
            help="Which example to run",
            required=True,
        )

    ##############################################

    # Scheduler Properties
    ##############################################
    slurmparser.add_argument(
        "-c", "--charge",
        type=str,
        help="Account to charge",
        required=True,
        dest="Charge",
    )
    slurmparser.add_argument(
        "-n", "--nodes",
        type=int,
        help="Number of nodes to use",
        required=False,
        default=1,
        dest="Nodes",
    )
    slurmparser.add_argument(
        "-w", "--walltime",
        type=str, 
        help="Wall time to request",
        required=False,
        default="00:15:00",
        dest="Walltime",
    )
    if isinstance(runner, perlmutter):
        slurmparser.add_argument(
            "-q", "--qos",
            type=str,
            help="QOS",
            required=False,
            default="regular",
            dest="QOS"
        )
    elif isinstance(runner, frontier):
        slurmparser.add_argument(
            "-q", "--qos",
            type=str,
            help="QOS",
            required=False,
            dest="QOS"
        )

    ##############################################

    args = slurmparser.parse_args()
    if isinstance(runner, perlmutter) and args.suchyta:
        args.Constraint = 'gpu'

    return args, runner


def WriteXML(outdir):
    filename = os.path.join(outdir, "adios2cfg.xml")
    tree = ET.parse(filename, parser=ET.XMLParser(remove_blank_text=True))

    adios_config = tree.getroot()

    #kind = "BP5"
    ios = adios_config.findall('io')
    for io in ios:
        engines = io.findall('engine')
        for engine in engines:
            #engine.set('type', kind)
            engine.text = None

    kind = "BP5"
    dname = 'diagnosis.diffusion_profiles'
    ios = adios_config.findall(f"./io[@name='{dname}']")
    if len(ios) == 0:
        ios = [ET.SubElement(adios_config, "io")]
    for io in ios:
        io.set("name", dname)
        engines = io.findall('engine')
        if len(engines) == 0:
            engines = [ET.SubElement(io, "engine")]
        for engine in engines:
            engine.set("type", kind)

    dname = "input.diffusion_coefficients"

    ET.indent(tree)
    tree.write(filename)



def EditNML(Directory):

    n = f90nml.Namelist()
    n["sml_param"] = {}
    n["sml_param"]["sml_mstep"] = 6

    n["diag_param"] = {}
    n["diag_param"]['diag_1d_on'] = True
    n["diag_param"]['diag_1d_period'] = 1
    n["diag_param"]['diag_3d_on'] = True
    n["diag_param"]['diag_3d_period'] = 1

    n["diag_param"]['diag_diff_profiles_on'] = True
    n["diag_param"]['diag_diff_profiles_period'] = 1

    n["diff_param"] = {}
    n["diff_param"]['diff_update_period'] = 2
    n["diff_param"]["diff_update_on"] = True

    filename = os.path.join(Directory, "input")
    tmp = "{0}.tmp".format(filename)
    shutil.copy(filename, tmp)
    os.remove(filename)

    f90nml.patch(tmp, n, out_path=filename)
    os.remove(tmp)




if __name__ == "__main__":

    args, runner = SetupArgs()

    if ('suchyta' in dir(args)) and args.suchyta:
        setup = getattr(suchyta, runner.__class__.__name__)
    else:
        EffisLogger.RaiseError(RuntimeError, "Edit the example for your setup")

    extra = {}
    for key in ('Nodes', 'Walltime', 'Charge', 'QOS', 'Constraint'):
        if key in dir(args):
            extra[key] = getattr(args, key)

    if isinstance(runner, frontier):
        extra['SchedulerDirectives'] = ["--threads-per-core", "2"]

    MyWorkflow = Workflow(
        Runner=runner,
        Directory=args.outdir,
        Subdirs=False,
        **extra,
    )

    # ITG_core_deltaf - xgc-es-cpp-gpu
    # ecbc_low_beta_n6 - xgc-eem-cpp-gpu
    # neoclassical_core - xgca-cpp-gpu

    AppInfo = {
        'cmd': os.path.join(setup['build'], 'bin', args.xgc),
        'Nodes': args.Nodes,
        'Ranks':  args.Nodes * setup['gpus'],
        'RanksPerNode': setup['gpus'],
        'GPUsPerRank': 1,
        'CoresPerRank': setup['per'],
        'MPIRunnerArguments': ['--gpu-bind', 'closest'],
        'SetupFile': os.path.join(setup['build'], "setup.sh"),
    }
    if isinstance(runner, frontier):
        AppInfo['MPIRunnerArguments'] += ["--threads-per-core", "2"]

    Simulation = MyWorkflow.Application(
        Name="XGC",
        LogFile="XGC.log",
        **AppInfo,
    )
    Simulation.Input += os.path.join(setup['src'], "quickstart", "inputs", "petsc.rc")
    Simulation.Input += os.path.join(setup['src'], "quickstart", "inputs", "adios2cfg.xml")

    usedir = args.example
    if os.path.isfile(args.example):
        usedir = os.path.dirname(args.example)
        Simulation.Input += Input(args.example, rename="input")

    for filename in os.listdir(usedir):
        Simulation.Input += os.path.join(usedir, filename)


    env = {}
    pythonpath = [
        setup['analysis'],
        setup['adios']
    ]
    if 'PYTHONPATH' in os.environ:
        pythonpath += [os.environ['PYTHONPATH']]
    env['PYTHONPATH'] = ':'.join(pythonpath)

    Analysis = MyWorkflow.Application(
        Runner=None,
        Environment=env,
        Name="DiffusionOptmizer",
        LogFile="DiffusionOptimizer.log",
        cmd="python3",
        CommandLineArguments=os.path.join(setup['analysis'], "Diffusion_Optimizer", "diffusion_optimizer.py"),
    )

    MyWorkflow.Create()
    WriteXML(Simulation.Directory)
    EditNML(Simulation.Directory)

    MyWorkflow.Submit()