parse

#!/usr/bin/env python3
import os
import subprocess
import threading
import sys
import shlex
import tempfile
import re
import datetime
import random
import time

if not os.path.exists('tools/parse'):
    print('Running from incorrect directory')
    exit(1)

if 'EPNSYNCMODE' in os.environ and int(os.environ['EPNSYNCMODE']):
    sys.path.insert(0, '/usr/share/Modules/init')
    import python as mod

if len(sys.argv) != 4:
    print('Incorrect number of arguments provided, syntax is parse [description library file] [topology name] [output file name]')
    exit(1)

if not 'FILEWORKDIR' in os.environ:
    print('$FILEWORKDIR env variable missing')
    exit(1)

if not 'DDWORKFLOW' in os.environ and not 'DDMODE' in os.environ:
    print('Need either $DDWORKFLOW or $DDMODE env variable')
    exit(1)

if not 'RECOSHMSIZE' in os.environ or not 'DDSHMSIZE' in os.environ:
    print('RECOSHMSIZE or DDSHMSIZE setting missing')
    exit(1)

NO_PROCESSING_MODE=0
if not 'DDWORKFLOW' in os.environ:
    os.environ['DDWORKFLOW'] = 'tools/datadistribution_workflows/dd-' + os.environ['DDMODE'] + '.xml'
    if os.environ['DDMODE'] == 'discard' or os.environ['DDMODE'] == 'disk':
        NO_PROCESSING_MODE=1

print('Using topology', sys.argv[2], 'of library', sys.argv[1])

if 'WORKFLOWMODE' in os.environ:
    if not os.environ['WORKFLOWMODE'] in ['dds', 'print']:
        print('Invalid WORKFLOWMODE provided')
        raise
else:
    os.environ['WORKFLOWMODE'] = 'dds'

if 'O2_NO_CATCHALL_EXCEPTIONS' in os.environ:
    os.unsetenv("O2_NO_CATCHALL_EXCEPTIONS")

if 'RECO_NUM_NODES_OVERRIDE' in os.environ and os.environ['RECO_NUM_NODES_OVERRIDE'] != '' and int(os.environ['RECO_NUM_NODES_OVERRIDE']) > 0:
    reco_num_nodes_override = int(os.environ['RECO_NUM_NODES_OVERRIDE'])
    os.environ['RECO_NUM_NODES_WORKFLOW'] = str(reco_num_nodes_override)
else:
    reco_num_nodes_override = 0

modulecmd = " module purge ; "
if 'GEN_TOPO_OVERRIDE_MODULE_DIR' in os.environ:
    mod.module('unuse', '/opt/alisw/el9/modulefiles')
    mod.module('use', os.environ['GEN_TOPO_OVERRIDE_MODULE_DIR'])
    modulecmd += ' module unuse /opt/alisw/el9/modulefiles 2>&1 ; module use ' + os.environ['GEN_TOPO_OVERRIDE_MODULE_DIR'] + ' 2>&1 ; '

f = open(sys.argv[1], 'r')
for line in f:
    line = line.strip()
    if len(line) == 0:
        continue
    if line[0] == '#':
        continue
    args = shlex.split(line)
    if len(args) <= 1:
        print('Toplogy must have at least name and O2 version')
        raise
    if len(args[0]) == 0:
        print('Empty topology name forbitten')
        raise
    if not args[0].endswith(':'):
        print('Topology name ', args[0], "not followed by ':'")
        raise
    if args[0] == sys.argv[2] + ':':
        reconodes = 0
        reconodesmin = 0
        recoworkflows = []
        recoworkflows_mi100 = []
        calibworkflows = []
        calibworkflowsdds = []
        print('Found topology', sys.argv[2], '-', args)
        if 'EPNSYNCMODE' in os.environ and int(os.environ['EPNSYNCMODE']) and (not 'GEN_TOPO_RUN_HOME' in os.environ or not int(os.environ['GEN_TOPO_RUN_HOME'])):
            restore_O2DPG_ROOT = 'O2DPG_ROOT' in os.environ
            if restore_O2DPG_ROOT:
                restore_O2DPG_ROOT_val = os.environ['O2DPG_ROOT']
            if 'OVERRIDE_PDPSUITE_VERSION' in os.environ and os.environ['OVERRIDE_PDPSUITE_VERSION'] != '':
                args[1] = os.environ['OVERRIDE_PDPSUITE_VERSION']
            for i in args[1].split():
                if 'GEN_TOPO_CACHEABLE' in os.environ and os.environ['GEN_TOPO_CACHEABLE'] == '1':
                    if i.find('/') == -1 or i.find('/latest') != -1:
                        print('Must not use non-versioned module', i, 'in cacheable workflow (i.e. with repository hash)')
                        raise
                print('Loading module', i)
                mod.module('load', i)
            if restore_O2DPG_ROOT:
                os.environ['O2DPG_ROOT'] = restore_O2DPG_ROOT_val
        if len(args) > 2 and not 'O2_ROOT' in os.environ:
            print('O2 not loaded')
            raise
        with tempfile.TemporaryDirectory(prefix='o2_workflow_') as tmpdir:
            if 'GEN_TOPO_OVERRIDE_TEMPDIR' in os.environ:
                tmpdir = os.environ['GEN_TOPO_OVERRIDE_TEMPDIR']
            os.environ['GEN_TOPO_QC_JSON_FILE'] = os.path.abspath(tmpdir + '/qc.json')
            if os.path.exists(os.environ['GEN_TOPO_QC_JSON_FILE']):
                os.remove(os.environ['GEN_TOPO_QC_JSON_FILE'])
            if NO_PROCESSING_MODE and len(args) > 2:
                print('Cannot use DPL workflow together with DD mode', os.environ['DDMODE'])
                raise

            def processSubtopology(i, wfCount, args, tmpdir, reconodes, reconodesmin, recoworkflows, recoworkflows_mi100, calibworkflows, calibworkflowsdds, threadsOK, mi100node):
                filename = tmpdir + '/wf' + str(wfCount) + '.dds'
                calibcores = '1'
                if args[i].startswith('reco'):
                    wf = args[i].split(',', 3)
                    is_calib_workflow = 0
                elif args[i].startswith('calib'):
                    wf = args[i].split(',', 2)
                    wf.append(wf[2])
                    calibcores = wf[1]
                    wf[1] = '1';
                    wf[2] = wf[1]
                    is_calib_workflow = 1
                else:
                    print('Invalid workflow type', args[i])
                    raise
                print('Adding', wf[0], 'workflow (', wf[2], '-', wf[1], 'nodes):', wf[3])
                reconodes = max(reconodes, int(wf[1]))
                reconodesmin = max(reconodesmin, int(wf[2]))
                command_preopt = 'GEN_TOPO_CALIB_WORKFLOW=' + str(is_calib_workflow)
                if not is_calib_workflow:
                    command_preopt += ' SHMSIZE=' + (str(int(int(os.environ['RECOSHMSIZE']) * 3 / 2)) if mi100node else os.environ['RECOSHMSIZE'])
                    if mi100node:
                        command_preopt += ' EPN_NODE_MI100=1 EPN_GLOBAL_SCALING="3/2"'
                        if os.getenv('DISABLE_MI100_SERIALIZATION', '0') == '1':
                            command_preopt += ' DISABLE_MI100_SERIALIZATION=1'
                        if os.getenv('FULL_MI100_SERIALIZATION', '0') == '1':
                            command_preopt += ' FULL_MI100_SERIALIZATION=1'
                if reco_num_nodes_override == 0:
                    command_preopt += ' RECO_NUM_NODES_WORKFLOW=' + wf[1]
                if is_calib_workflow:
                    command_preopt += ' GEN_TOPO_CALIB_NCORES=' + calibcores
                if os.environ['WORKFLOWMODE'] == 'dds':
                    command_preopt += ' WORKFLOWMODE_FILE=' + filename
                command = command_preopt + ' GLOBALDPLOPT+=" -b --dds-workflow-suffix _' + wf[0] + str(wfCount) + '" ' + wf[3]
                if os.environ['WORKFLOWMODE'] == 'print':
                    command += ' > ' + filename
                print('Running DPL command', command)
                starttime = time.time()
                retVal = subprocess.run(command, shell=True, stderr=subprocess.PIPE, stdout=subprocess.PIPE, input="")
                print('Execution time: ', time.time() - starttime)
                tmpchk = retVal.stderr.decode() + retVal.stdout.decode()
                haserror = 0
                if tmpchk.find('[FATAL]') != -1 or tmpchk.find('[ERROR]') != -1 or tmpchk.find(': command not found') != -1 or tmpchk.find('No such file or directory') != -1 or tmpchk.find('bash: syntax error') != -1:
                    print("Found error messages during topology generation")
                    haserror = 1
                if retVal.returncode != 0:
                    print("Non-zero return value " + str(retVal.returncode) + " during topology generation")
                    haserror = 1
                if haserror != 0:
                    print('Error running command', command)
                    print('stderr:', retVal.stderr.decode())
                    print('stdout:', retVal.stdout.decode())
                    raise
                threadsOK[wfCount] = 1
                if args[i].startswith('reco'):
                    if mi100node:
                        recoworkflows_mi100.append(filename)
                    else:
                        recoworkflows.append(filename)
                elif args[i].startswith('calib'):
                    if os.environ['WORKFLOWMODE'] == 'dds':
                        isempty = 1
                        ftmp = open(filename, 'r')
                        rg = re.compile('^.*<decltask')
                        for line in ftmp:
                            if re.match(rg, line):
                                isempty = 0
                                break
                        if isempty:
                            return
                    calibworkflows.append(filename)
                    calibworkflowsdds.append(filename + ':' + calibcores)

            subtopologyThreads = []
            threadsOK = []
            wfCount = 0
            for i in range(2, len(args)):
                for nodetype in range(0, 2):
                    if nodetype == 1 and (not args[i].startswith('reco') or not 'GEN_TOPO_MI100_NODES' in os.environ or os.environ['GEN_TOPO_MI100_NODES'] == 0):
                        continue
                    threadsOK.append(0)
                    t = threading.Thread(target = processSubtopology, args = (i, wfCount, args, tmpdir, reconodes, reconodesmin, recoworkflows, recoworkflows_mi100, calibworkflows, calibworkflowsdds, threadsOK, nodetype))
                    t.start()
                    if 'GEN_TOPO_SINGLE_THREAD' in os.environ and int(os.environ['GEN_TOPO_SINGLE_THREAD']):
                        t.join()
                        if not threadsOK[wfCount]:
                            raise
                    subtopologyThreads.append(t)
                    wfCount += 1
            for t in subtopologyThreads:
                t.join()
            for i in range(0, wfCount):
                if not threadsOK[i]:
                    raise
            if reco_num_nodes_override > 0:
                reconodes = reco_num_nodes_override
                reconodesmin = min(reconodes, reconodesmin)
            if 'RECO_MAX_FAIL_NODES_OVERRIDE' in os.environ and os.environ['RECO_MAX_FAIL_NODES_OVERRIDE'] != '':
                reconodesmin = max(1, reconodes - int(os.environ['RECO_MAX_FAIL_NODES_OVERRIDE']))
            if os.environ['WORKFLOWMODE'] == 'dds':
                odccommand = os.environ['GEN_TOPO_ODC_EPN_TOPO_CMD']
                if 'GEN_TOPO_ODC_EPN_TOPO_ARGS' in os.environ:
                  odccommand += ' ' + os.environ['GEN_TOPO_ODC_EPN_TOPO_ARGS']
                if True:
                    replacestring = ''
                    dd_env_variables = ['DD_DISK_FRACTION', 'TF_SINK_NUM_TFS', 'DD_STF_MIN_ID', 'DD_STF_MAX_ID', 'EPN_DD_TEST', 'EPN_DD_TEST_2', 'EPN_DD_TEST_3', 'SHM_MANAGER_SHMID']
                    for i in dd_env_variables:
                        if i in os.environ:
                            replacestring += ' ' + i + '=' + os.environ[i]
                    fddin = open(os.environ['DDWORKFLOW'], 'rt')
                    filename = tmpdir + '/wf_dd.dds';
                    fddout = open(filename, 'wt')
                    for line in fddin:
                        fddout.write(line.replace('GEN_TOPO_TFBUILDER_ENV_VARIABLES', replacestring))
                    fddin.close()
                    fddout.close()
                    odccommand += ' --dd ' + filename
                if len(recoworkflows):
                    odccommand += ' --reco ' + ' '.join(recoworkflows)
                if len(recoworkflows_mi100):
                    odccommand += ' --reco100 ' + ' '.join(recoworkflows_mi100)
                if len(calibworkflows):
                    calibworkflowsdds.sort(key=lambda x:int(x.split(':')[1])*-1)
                    odccommand += ' --calib ' + ' '.join(calibworkflowsdds)
                if 'GEN_TOPO_STDERR_LOGGING' in os.environ and int(os.environ['GEN_TOPO_STDERR_LOGGING']):
                    odccommand += ' --mon tools/monitoring_workflows/epnstderrlog.xml'
                if args[1] != '':
                    odccommand += ' --prependexe "' + modulecmd + ' module load ' + args[1] + ' 2>&1 ; export O2_NO_CATCHALL_EXCEPTIONS=1 ; '
                    if 'GEN_TOPO_ENV_AT_RUNTIME' in os.environ:
                        env_at_runtime = os.environ['GEN_TOPO_ENV_AT_RUNTIME'].split(',')
                        for i in env_at_runtime:
                            odccommand += 'export ' + i + '=' + os.environ[i] + ' ; '
                    odccommand += '"'
                if os.path.exists(os.environ['GEN_TOPO_QC_JSON_FILE']):
                    odccommand += f" --assets {os.environ['GEN_TOPO_QC_JSON_FILE']}"
                odccommand += ' -o ' + sys.argv[3]
                print('Running topology merger command', odccommand)
                # Add a custom PYTHONPATH in order to import the python packages needed for the topology merger
                starttime = time.time()
                try:
                    out = subprocess.check_output(odccommand, stderr=subprocess.STDOUT, shell=True)
                    print(out.decode('utf-8'))
                    print('Topology merger finished successfully')
                except subprocess.CalledProcessError as e:
                    print(e.output.decode('utf-8'))
                    print('Error running topology merger!')
                    raise
                print('Execution time: ', time.time() - starttime)
                if os.path.exists(os.environ['GEN_TOPO_QC_JSON_FILE']):
                    command = 'sed -i \'s/&/\&amp;/g; s/</\&lt;/g; s/>/\&gt;/g; s/"/\&quot;/g; s/\'"\'"\'/\&#39;/g\' ' + os.environ['GEN_TOPO_QC_JSON_FILE']
                    print('Running SED command for DPL JSON XML Escaping', command)
                    if os.system(command) != 0:
                        print('Error running sed on JSON file')
            else:
                outf = open(sys.argv[3], 'w+')
                for i in recoworkflows:
                    outf.write('# RECO workflow\n\n' + open(i, 'r').read() + '\n\n')
                for i in recoworkflows_mi100:
                    outf.write('# RECO MI100 workflow\n\n' + open(i, 'r').read() + '\n\n')
                for i in calibworkflows:
                    outf.write('# CALIB workflow\n\n' + open(i, 'r').read() + '\n\n')
        print('Done')
        exit(0)

print('Could not find workflow', sys.argv[2])
exit(1)