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 41 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 55 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 123 of file trfMPITools.py.

def mergeOutputs():
    """Merge outputs into rank 0"""
    if mpiConfig is None:
        msg.warning("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
    }
    num_ranks = len(rank_dirs) + 1
    # First wait for all ranks to reach this point so we don't start merging before some outputs are fully closed
    open("athena_done", "a").close()
    files_to_check = [
        (rank, f"../rank-{rank}/athena_done") for rank in range(0, num_ranks)
    ]
    count = 0
    files_to_check = list(
        it.filterfalse(lambda f: os.path.exists(f[1]), files_to_check)
    )
    while files_to_check:
        if count % 10 == 0 and getMPIRank() == 0:
            msg.info(
                f"{count // 10 + 1}: Waiting for all ranks to finish athena: {list(map(lambda x: x[0], files_to_check))}"
            )
        count += 1
        sleep(6)
        files_to_check = list(
            it.filterfalse(lambda f: os.path.exists(f[1]), files_to_check)
        )
    # Now continue: First the logs
    if getMPIRank() == 0:
        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()
        # Then everything else
        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:
            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")
            continue
        for idx in range(getMPIRank(), len(merge_lists), num_ranks):
            my_merge = merge_lists[idx]
            if len(my_merge[1]) < 1:
                msg.info(
                    f"In rank {getMPIRank()}, no inputs for ../rank-0/{my_merge[0]}"
                )
                continue
            msg.info(
                f"In rank {getMPIRank()}, merging into ../rank-0/{my_merge[0]}. Inputs are \n{pprint.pformat(my_merge[1])}"
            )
            try:
                merge_helper.selfMerge(f"../rank-0/{my_merge[0]}", my_merge[1])
            except Exception as e:
                msg.error(
                    f"Merge failure in rank {getMPIRank()} merging into {my_merge[0]}: {e}"
                )
                with open("../rank-0/merge_failure", "a") as f:
                    f.write(
                        f"Merge failure in rank {getMPIRank()} merging into {my_merge[0]}: {e}\n"
                    )
    # Create a file to indicate we are done
    open("done_merging", "a").close()
    if getMPIRank() == 0:
        # In rank 0, wait until all other ranks have finished merging
        files_to_check = [
            (rank, f"../rank-{rank}/done_merging") for rank in range(0, num_ranks)
        ]
        count = 0
        files_to_check = list(
            it.filterfalse(lambda f: os.path.exists(f[1]), files_to_check)
        )
        while files_to_check:
            if count % 10 == 0:
                msg.info(
                    f"Waiting for all ranks to finish merging: {list(map(lambda x: x[0], files_to_check))}"
                )
            count += 1
            sleep(6)
            files_to_check = list(
                it.filterfalse(lambda f: os.path.exists(f[1]), files_to_check)
            )
        if not os.path.exists("merge_failure"):
            msg.info("All ranks done merging")
        else:
            msg.error("ERRORS WHILE MERGING")
            raise RuntimeError("Output merging error")

mpiOutputs()

python.trfMPITools.mpiOutputs

(

)

Definition at line 119 of file trfMPITools.py.

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

mpiShouldValidate()

python.trfMPITools.mpiShouldValidate

(

)

Definition at line 111 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 67 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 34 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 23 of file trfMPITools.py.

msg

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

Definition at line 21 of file trfMPITools.py.