python.trfMPITools Namespace Reference

Classes
class	MPIType

Functions
def	signalError (message)
def	getMPIRank ()
def	getMPIType ()
def	setupMPIConfig (output, dataDict)
def	mpiShouldValidate ()
def	mpiOutputs ()
def	mergeOutputs ()

Variables
	msg
	mpiConfig

Function Documentation

getMPIRank()

def python.trfMPITools.getMPIRank

(

)

Return MPI rank

Definition at line 40 of file trfMPITools.py.

def getMPIRank():
    """Return MPI rank"""
    if mpiConfig is not None:
        return int(mpiConfig["rank"])
    if "RANK" not in os.environ:
        return -1
    else:
        try:
            return int(os.environ["RANK"])
        except ValueError:
            signalError("$RANK environment variable is not an integer")
            return -2  # Only here to placate PyRight
 
 

getMPIType()

def python.trfMPITools.getMPIType

(

)

Return MPI type

Definition at line 54 of file trfMPITools.py.

def getMPIType():
    """Return MPI type"""
    if mpiConfig is not None:
        return mpiConfig["type"]
    if "RANK" not in os.environ:
        return MPIType.NOMPI
    if getMPIRank() == 0:
        return MPIType.MPIMASTER
    else:
        return MPIType.MPIWORKER
 
 

mergeOutputs()

def python.trfMPITools.mergeOutputs

(

)

Merge outputs into rank 0

Definition at line 122 of file trfMPITools.py.

def mergeOutputs():
    """Merge outputs into rank 0"""
    if mpiConfig is None:
        msg.warn("trfMPITools.mergeOutputs called when we are not in MPI mode")
        return
    rank_dir_regex = re.compile("rank-([0-9]+)$")
    rank_dirs = {
        int(m.group(1)): m.string
        for m in (rank_dir_regex.search(d.path) for d in os.scandir("..") if d.is_dir())
        if m and int(m.group(1)) > 0
    }
    if getMPIRank() == 0:
        msg.info("Rank output directories are:\n{}".format(pprint.pformat(rank_dirs)))
    all_merge_inputs = list(
        map(
            lambda f: f.path,
            filter(
                lambda f: f.is_file(),
                it.chain.from_iterable(map(os.scandir, rank_dirs.values())),
            ),
        )
    )
    # Remove PoolFileCatalog
    try:
        os.remove("PoolFileCatalog.xml")
    except FileNotFoundError:
        pass
    for dtype, defn in mpiConfig["outputs"].items():
        if getMPIRank() == 0:
            msg.info(f"Output type is {dtype}")
        merge_helper = deepcopy(defn)
        merge_helper.multipleOK = True
        if getMPIRank() == 0:
            for fn in defn.list_to_remove:
                # remove empty files from rank 0
                try:
                    os.remove(fn)
                except FileNotFoundError:
                    pass
        merge_lists = []
        for fn in defn.value:
            msg.info(f"Generating merge list by filtering for {fn} in {all_merge_inputs}")
            merge_inputs = sorted(filter(lambda s: s.endswith(fn), all_merge_inputs))
            # Add to list
            merge_helper.value.extend(merge_inputs)
            merge_lists.append((fn, merge_inputs))
        # Remove non-existent output files from mpiOutputs
        defn.value = [x[0] for x in merge_lists if len(x[1]) >= 1]
        # Merge each final output in a different rank
        if getMPIRank() >= len(merge_lists):
            msg.info(f"In rank {getMPIRank()}, not merging")
            break
        my_merge = merge_lists[getMPIRank()]
        if len(my_merge[1]) < 1:
            msg.info(f"In rank {getMPIRank()}, no inputs for ../rank-0/{my_merge[0]}")
            open("done_merging", "a").close()
            break
        msg.info(
            f"In rank {getMPIRank()}, merging into ../rank-0/{my_merge[0]}. Inputs are \n{pprint.pformat(my_merge[1])}"
        )
        merge_helper.selfMerge(f"../rank-0/{my_merge[0]}", my_merge[1])
    # Create a file to indicate we are done
    open("done_merging", "a").close()
    if getMPIRank() == 0:
        from functools import reduce
        from operator import and_
        from time import sleep
        import sqlite3 as sq3
        from glob import glob
 
        # Merge log databases
        conn = sq3.connect("mpilog.db")
        cur = conn.cursor()
        tables = ["ranks", "files", "event_log"]
        for db in glob("../rank-[1-9]*/mpilog.db"):
            cur.execute("ATTACH DATABASE ? as db", (db,))
            for table in tables:
                upsert = "INSERT OR IGNORE" if table == "files" else "INSERT"
                cur.execute(f"{upsert} INTO {table} SELECT * from db.{table}")
            conn.commit()
            cur.execute("DETACH DATABASE db")
        conn.close()
 
        # In rank 0, wait until all other ranks have finished merging
        files_to_check = [
            f"../rank-{rank}/done_merging" for rank in range(0, len(merge_lists))
        ]
        count = 0
        check = [os.path.exists(f) for f in files_to_check]
        while not reduce(and_, check):
            if count % 10 == 0:
                msg.info("Waiting for other ranks to finish merging")
                msg.debug(f"Looking for {files_to_check}")
                msg.debug(f"Result: {check}")
            count = count + 1
            sleep(6)
            check = [os.path.exists(f) for f in files_to_check]
        msg.info("All ranks done merging")

mpiOutputs()

def python.trfMPITools.mpiOutputs

(

)

Definition at line 118 of file trfMPITools.py.

def mpiOutputs():
    return mpiConfig["outputs"].values()
 
 

mpiShouldValidate()

def python.trfMPITools.mpiShouldValidate

(

)

Definition at line 110 of file trfMPITools.py.

def mpiShouldValidate():
    if getMPIType() == MPIType.NOMPI:
        return True  # validate if we're not in MPI mode
    if getMPIRank() == 0:
        return True  # validate in rank 0
    return False  # don't validate in other ranks
 
 

setupMPIConfig()

def python.trfMPITools.setupMPIConfig	(		output,
			dataDict
	)

Check environment is correct if we are in MPI mode, and setup dictionaries

Definition at line 66 of file trfMPITools.py.

def setupMPIConfig(output, dataDict):
    """Check environment is correct if we are in MPI mode, and setup dictionaries"""
    global mpiConfig
    if "RANK" not in os.environ:
        signalError(
            "Running in MPI mode but the $RANK environment variable is not set!"
        )
    rank = getMPIRank()
    if not os.getcwd().endswith("rank-{}".format(getMPIRank())):
        signalError(
            "Running in MPI mode with rank {0} but working directory is not called rank-{0}".format(
                getMPIRank()
            )
        )
    mpiType = getMPIType()
    mpiConfig = {}
    mpiConfig["rank"] = rank
    mpiConfig["type"] = mpiType
    mpiConfig["outputs"] = {
        dataType: deepcopy(dataDict[dataType]) for dataType in output
    }
    # expand any [  ] lists in output filenames
    output_proc_regex = re.compile(r"(.+)\[(.*)](.*)")
    for v in mpiConfig["outputs"].values():
        v.multipleOK = True
        new_list = []
        list_to_remove = []
        for fn in v.value:
            if ("[" in fn) and ("]" in fn):
                match = output_proc_regex.match(fn)
                new_list.extend(
                    [
                        f"{match.group(1)}{it}{match.group(3)}"
                        for it in match.group(2).split(",")
                    ]
                )
                list_to_remove.append(match.group(1))
            else:
                new_list.append(fn)
                list_to_remove.append(fn)
        v.value = new_list
        v.list_to_remove = list(set(list_to_remove))
 
 

signalError()

def python.trfMPITools.signalError

(

message

)

Definition at line 33 of file trfMPITools.py.

def signalError(message):
    msg.error(message)
    raise trfExceptions.TransformSetupException(
        trfExit.nameToCode("TRF_SETUP"), message
    )
 
 

Variable Documentation

mpiConfig

python.trfMPITools.mpiConfig

Definition at line 22 of file trfMPITools.py.

msg

python.trfMPITools.msg

Definition at line 20 of file trfMPITools.py.