python.trfMPITools Namespace Reference

Classes
class	MPIType

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

Variables
	msg = logging.getLogger(__name__)
	mpiConfig = None

Function Documentation

getMPIRank()

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()

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()

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()

python.trfMPITools.mpiOutputs

(

)

Definition at line 118 of file trfMPITools.py.

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

mpiShouldValidate()

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()

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()

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 = None

Definition at line 22 of file trfMPITools.py.

msg

python.trfMPITools.msg = logging.getLogger(__name__)

Definition at line 20 of file trfMPITools.py.