pTools.py

Go to the documentation of this file.

import pickle
import time
from multiprocessing import Event
 
from GaudiPython import FAILURE, SUCCESS, gbl
 
# Eoin Smith
# 3 Aug 2010
 
#
# This script contains the ancillary classes and functions used in the
#  GaudiPython Parallel model
#
# Classes :
# - HistoAgent : In charge of extracting Histograms from their Transient Store
#                 on a reader/worker, communicating them to the writer, and on
#                 the writer receiving them and rebuilding a single store
#
# - FileRecordsAgent : Similar to HistoAgent, but for FileRecords Data
#
# - LumiFSR : FSR data from different workers needs to be carefully merged to
#              replicate the serial version; this class aids in that task by
#              representing an LHCb LumiFSR object as a python class
#
# - PackedCaloHypos : Pythonization of an LHCb class, used for inspecting some
#                      differences in serial and parallel output
#
# - Syncer : This class is responsible for syncing processes for a specified
#             section of execution.  For example, one Syncer object might be
#             syncing Initialisation, one for Running, one for Finalisation.
#             Syncer uses multiprocessing.Event() objects as flags which are
#             visible across the N processes sharing them.
#             IMPORTANT : The Syncer objects in the GaudiPython Parallel model
#              ensure that there is no hanging; they in effect, allow a timeout
#              for Initialisation, Run, Finalise on all processes
#
# - SyncMini : A class wrapper for a multiprocessing.Event() object
#
# Methods :
# - getEventNumber(evt) : pass a valid instance of the GaudiPython TES
#                          ( AppMgr().evtsvc() ) to this to get the current
#                          Event Number as an integer (even from RawEvents!)
#
 
# used to convert stored histos (in AIDA format) to ROOT format
aida2root = gbl.Gaudi.Utils.Aida2ROOT.aida2root
 
# =========================== Classes =========================================
 
 
class HistoAgent:
    def __init__(self, gmpComponent):
        self._gmpc = gmpComponent
        self.hvt = self._gmpc.hvt
        self.histos = []
        self.qin = self._gmpc.hq
        self.log = self._gmpc.log
 
        # There are many methods for booking Histogram Objects to Histo store
        # here they are collected in a dictionary, with key = a relevant name
        self.bookingDict = {}
        self.bookingDict["DataObject"] = self.bookDataObject
        self.bookingDict["NTuple::Directory"] = self.bookDataObject
        self.bookingDict["NTuple::File"] = self.bookDataObject
        self.bookingDict["TH1D"] = self.bookTH1D
        self.bookingDict["TH2D"] = self.bookTH2D
        self.bookingDict["TH3D"] = self.bookTH3D
        self.bookingDict["TProfile"] = self.bookTProfile
        self.bookingDict["TProfile2D"] = self.bookTProfile2D
 
    def register(self, tup):
        # add a tuple of (worker-id, histoDict) to self.histos
        assert tup.__class__.__name__ == "tuple"
        self.histos.append(tup)
 
    def Receive(self):
        hstatus = self._gmpc.nWorkers + 1  # +1 for the Reader!
        while True:
            tup = self.qin.get()
            if tup == "HISTOS_SENT":
                self.log.debug("received HISTOS_SENT message")
                hstatus -= 1
                if not hstatus:
                    break
            else:
                self.register(tup)
        self._gmpc.sEvent.set()
        self.log.info("Writer received all histo bundles and set sync event")
        return SUCCESS
 
    def RebuildHistoStore(self):
        """
        Rebuild the Histogram Store from the histos received by Receive()
        If we have a histo which is not in the store,
        book and fill it according to self.bookingDict
        If we have a histo with a matching histo in the store,
        add the two histos, remembering that aida2root must be used on
        the Stored histo for compatibility.
        """
        errors = 0
        for tup in self.histos:
            workerID, histDict = tup
            added = 0
            booked = 0
 
            for n in histDict.keys():
                o = histDict[n]
                obj = self.hvt.retrieve(n)
 
                if obj:
                    try:
                        aida2root(obj).Add(o)
                    except Exception:
                        self.log.warning(f"FAILED TO ADD : {obj}")
                        errors += 1
                    added += 1
                else:
                    if o.__class__.__name__ in self.bookingDict.keys():
                        try:
                            self.bookingDict[o.__class__.__name__](n, o)
                        except Exception:
                            self.log.warning(
                                f"FAILED TO REGISTER : {o.__class__.__name__}\tto{n}"
                            )
                            errors += 1
                    else:
                        self.log.warning(
                            f"No booking method for: {n}\t{type(o)}\t{o.__class__.__name__}"
                        )
                        errors += 1
                    booked += 1
        hs = self.hvt.getHistoNames()
        self.log.info("Histo Store Rebuilt : ")
        self.log.info(f"  Contains {len(hs)} objects.")
        self.log.info(f"  Errors in Rebuilding : {errors}")
        return SUCCESS
 
    def bookDataObject(self, n, o):
        """
        Register a DataObject to the Histo Store
        """
        self._gmpc.hvt.registerObject(n, o)
 
    def bookTH1D(self, n, o):
        """
        Register a ROOT 1D THisto to the Histo Store
        """
        obj = self.hvt._ihs.book(
            n,
            o.GetTitle(),
            o.GetXaxis().GetNbins(),
            o.GetXaxis().GetXmin(),
            o.GetXaxis().GetXmax(),
        )
        aida2root(obj).Add(o)
 
    def bookTH2D(self, n, o):
        """
        Register a ROOT 2D THisto to the Histo Store
        """
        obj = self.hvt._ihs.book(
            n,
            o.GetTitle(),
            o.GetXaxis().GetNbins(),
            o.GetXaxis().GetXmin(),
            o.GetXaxis().GetXmax(),
            o.GetYaxis().GetNbins(),
            o.GetYaxis().GetXmin(),
            o.GetYaxis().GetXmax(),
        )
        aida2root(obj).Add(o)
 
    def bookTH3D(self, n, o):
        """
        Register a ROOT 3D THisto to the Histo Store
        """
        obj = self.hvt._ihs.book(
            n,
            o.GetTitle(),
            o.GetXaxis().GetXbins(),
            o.GetXaxis().GetXmin(),
            o.GetXaxis().GetXmax(),
            o.GetYaxis().GetXbins(),
            o.GetYaxis().GetXmin(),
            o.GetYaxis().GetXmax(),
            o.GetZaxis().GetXbins(),
            o.GetZaxis().GetXmin(),
            o.GetZaxis().GetXmax(),
        )
        aida2root(obj).Add(o)
 
    def bookTProfile(self, n, o):
        """
        Register a ROOT TProfile to the Histo Store
        """
        obj = self.hvt._ihs.bookProf(
            n,
            o.GetTitle(),
            o.GetXaxis().GetNbins(),
            o.GetXaxis().GetXmin(),
            o.GetXaxis().GetXmax(),
            o.GetOption(),
        )
        aida2root(obj).Add(o)
 
    def bookTProfile2D(self, n, o):
        """
        Register a ROOT TProfile2D to the Histo Store
        """
        obj = self.hvt._ihs.bookProf(
            n,
            o.GetTitle(),
            o.GetXaxis().GetNbins(),
            o.GetXaxis().GetXmin(),
            o.GetXaxis().GetXmax(),
            o.GetYaxis().GetNbins(),
            o.GetYaxis().GetXmin(),
            o.GetYaxis().GetXmax(),
        )
        aida2root(obj).Add(o)
 
 
# =============================================================================
 
 
class FileRecordsAgent:
    def __init__(self, gmpComponent):
        self._gmpc = gmpComponent
        self.fsr = self._gmpc.fsr
        self.q = self._gmpc.fq
        self.log = self._gmpc.log
        self.objectsIn = []  # used for collecting FSR store objects
        self.objectsOut = []
 
    def localCmp(self, tupA, tupB):
        # sort tuples by a particular element
        # for the sort() method
        ind = 0
        valA = tupA[ind]
        valB = tupB[ind]
        if valA < valB:
            return -1
        elif valA > valB:
            return 1
        else:
            return 0
 
    def SendFileRecords(self):
        # send the FileRecords data as part of finalisation
 
        # Take Care of FileRecords!
        # There are two main things to consider here
        # 1) The DataObjects in the FileRecords Transient Store
        # 2) The fact that they are Keyed Containers, containing other objects
        #
        # The Lead Worker, nodeID=0, sends everything in the FSR store, as
        #   a completeness guarantee,
        #
        # send in form ( nodeID, path, object)
        self.log.info("Sending FileRecords...")
        lst = self.fsr.getHistoNames()
 
        # Check Validity
        if not lst:
            self.log.info("No FileRecords Data to send to Writer.")
            self.q.put("END_FSR")
            return SUCCESS
 
        # no need to send the root node
        if "/FileRecords" in lst:
            lst.remove("/FileRecords")
 
        for l in lst:
            o = self.fsr.retrieveObject(l)
            if hasattr(o, "configureDirectAccess"):
                o.configureDirectAccess()
            # lead worker sends everything, as completeness guarantee
            if self._gmpc.nodeID == 0:
                self.objectsOut.append((0, l, pickle.dumps(o)))
            else:
                # only add the Event Counter
                # and non-Empty Keyed Containers (ignore empty ones)
                if l == "/FileRecords/EventCountFSR":
                    tup = (self._gmpc.nodeID, l, pickle.dumps(o))
                    self.objectsOut.append(tup)
                elif "KeyedContainer" in o.__class__.__name__:
                    # It's a Keyed Container
                    nObjects = o.numberOfObjects()
                    if nObjects:
                        self.log.debug(f"Keyed Container {l} with {nObjects} objects")
                        tup = (self._gmpc.nodeID, l, pickle.dumps(o))
                        self.objectsOut.append(tup)
                else:
                    self.log.info(f"Ignoring {o.__class__.__name__} in FSR")
 
        self.log.debug("Done with FSR store, just to send to Writer.")
 
        if self.objectsOut:
            self.log.debug(f"{len(self.objectsOut)} FSR objects to Writer")
            for ob in self.objectsOut:
                self.log.debug(f"\t{ob[0]}")
            self.q.put(self.objectsOut)
        else:
            self.log.info("Valid FSR Store, but no data to send to Writer")
        self.log.info("SendFSR complete")
        self.q.put("END_FSR")
        return SUCCESS
 
    def Receive(self):
        # Receive contents of all Workers FileRecords Transient Stores
        self.log.info("Receiving FSR store data...")
        nc = self._gmpc.nWorkers
        while nc > 0:
            objects = self.q.get()
            if objects == "END_FSR":
                nc -= 1
                continue
                if nc == 0:
                    break
            # but if it's regular objects...
            for o in objects:
                self.objectsIn.append(o)
        # Now sort it by which worker it came from
        # an object is : (nodeID, path, pickledObject)
        self.objectsIn.sort(cmp=self.localCmp)
        self.log.info("All FSR data received")
        return SUCCESS
 
    def Rebuild(self):
        # objects is a list of (path, serializedObject) tuples
        for sourceNode, path, serialob in self.objectsIn:
            self.log.debug(f"Working with {path}")
            ob = pickle.loads(serialob)
            if hasattr(ob, "update"):
                ob.update()
            if hasattr(ob, "numberOfObjects"):
                nCont = ob.numberOfObjects()
                self.log.debug(f"\t {type(ob).__name__} has containedObjects : {nCont}")
            if sourceNode == 0:
                self.log.debug(f"Registering Object to : {path}")
                self.fsr.registerObject(path, ob)
            else:
                self.log.debug(f"Merging Object to : {path}")
                self.MergeFSRobject(sourceNode, path, ob)
        # As RecordStream has been split into Worker and Writer parts, the
        # count for output is wrong... fix that here, as every event received
        # by the writer is written (validation testing occurs on the worker)
 
        self.log.info("FSR Store Rebuilt.  Correcting EventCountFSR")
        if bool(self.fsr._idp):  # There might not be an FSR stream (Gauss)
            ecount = "/FileRecords/EventCountFSR"
            if self.fsr[ecount]:
                self.fsr[ecount].setOutput(self._gmpc.nIn)
                self.log.info(
                    f"Event Counter Output set : {ecount} : {self.fsr[ecount].output()}"
                )
            # Do some reporting
            self.log.debug("FSR store reconstructed!")
            lst = self.fsr.getHistoNames()
            if lst:
                for l in lst:
                    ob = self.fsr.retrieveObject(l)
                    if hasattr(ob, "configureDirectAccess"):
                        ob.configureDirectAccess()
                    if hasattr(ob, "containedObjects"):
                        # if ob.numberOfObjects() :
                        self.log.debug(
                            f"\t{l} (cont. objects : {ob.numberOfObjects()})"
                        )
                    else:
                        self.log.debug(f"\t{l}")
        self.log.info("FSR Store fully rebuilt.")
        return SUCCESS
 
    def MergeFSRobject(self, sourceNode, path, ob):
        # Merge Non-Empty Keyed Container from Worker>0
        if path == "/FileRecords/TimeSpanFSR":
            # TimeSpanFSR is a straightforward case
            self.ProcessTimeSpanFSR(path, ob)
        elif path == "/FileRecords/EventCountFSR":
            # Event Counter is also easy
            self.ProcessEventCountFSR(path, ob)
        # now other cases may not be so easy...
        elif (
            "KeyedContainer" in ob.__class__.__name__
            and "LumiFSR" in ob.__class__.__name__
        ):
            # Keyed Container of LumiFSRs : extract and re-register
            self.MergeLumiFSR(path, ob)
        else:
            self.log.info(f"Skipping Merge of {ob.__class__.__name__} at {path}")
 
    def ProcessTimeSpanFSR(self, path, ob):
        ob2 = self.fsr.retrieveObject(path)
        if ob.containedObjects().size():
            sz = ob.containedObjects().size()
            cob = ob2.containedObjects()[0]
            min = cob.earliest()
            max = cob.latest()
            for j in range(sz):
                cob = ob.containedObjects()[j]
                self.log.debug("Adding TimeSpanFSR")
                if cob.earliest() < min:
                    min = cob.earliest()
                if cob.latest() > max:
                    max = cob.latest()
                # this is annoying: it has to be rebuilt, without a key & added
                continue
            tsfsr = gbl.LHCb.TimeSpanFSR()
            tsfsr.setEarliest(min)
            tsfsr.setLatest(max)
            self.fsr[path].clear()
            self.fsr[path].add(tsfsr)
 
    def ProcessEventCountFSR(self, path, ob):
        self.log.debug("Event Count Input Addition")
        self.fsr[path].setInput(self.fsr[path].input() + ob.input())
 
    def MergeLumiFSR(self, path, keyedC):
        from ROOT import string
 
        # Fetch the first lumi
        keyedContainer = self.fsr.retrieveObject(path)
        # The LumiFSR KeyedContainers only have one object
        assert keyedContainer.numberOfObjects() == 1
        l = keyedContainer.containedObject(0)
        baseLumi = LumiFSR(l)
        # Now deal with the argument Non-empty Keyed Container of LumiFSRs
        nCont = keyedC.numberOfObjects()
        for i in range(nCont):
            obj = keyedC.containedObject(i)
            nextLumi = LumiFSR(obj)
            baseLumi.merge(nextLumi)
        # Now Rebuild and ReRegister
        newLumi = gbl.LHCb.LumiFSR()
        for r in baseLumi.runs:
            newLumi.addRunNumber(r)
        for f in baseLumi.files:
            newLumi.addFileID(string(f))
        for k in baseLumi.keys:
            increment, integral = baseLumi.info[k]
            newLumi.addInfo(k, increment, integral)
        # clear existing Keyed Container
        self.fsr[path].clear()
        # Add newly merged lumiFSR
        self.fsr[path].add(newLumi)
        return SUCCESS
 
 
# =============================================================================
 
 
class LumiFSR:
    def __init__(self, lumi):
        # lumi looks like :
        # {  runs : 69857 69858
        #    files : root:/castor/cer.../069857_0000000006.raw
        #    info (key/incr/integral) : 0 8 0 / 1 8 259 / 2 8 76 ... }
 
        # class variables
        self.runs = []
        self.files = []
        self.info = {}
        self.keys = []
 
        # get run numbers
        for r in lumi.runNumbers():
            self.runs.append(r)
        # get file ids
        for f in lumi.fileIDs():
            self.files.append(f)
        # Now the tricky bit, the info is not accessible via Python
        # except as a string
        s = str(lumi)
        sa = s.split("info (key/incr/integral) : ")[-1]
        sa = sa.split("/")[:-1]
        for rec in sa:
            k, i, t = rec.split()
            k = int(k)
            i = int(i)
            t = int(t)
            self.info[k] = (i, t)
        self.keys = self.info.keys()
 
    def merge(self, otherLumi):
        assert otherLumi.__class__.__name__ == "LumiFSR"
        # add any extra runs
        for r in otherLumi.runs:
            if r in self.runs:
                pass
            else:
                self.runs.append(r)
        self.runs.sort()
        # add any extra fileIDs
        for f in otherLumi.files:
            if f in self.files:
                pass
            else:
                self.files.append(f)
        self.files.sort()
        # Now add any extra records
        for k in otherLumi.keys:
            increment, integral = otherLumi.info[k]
            if k in self.keys:
                myIncrement, myIntegral = self.info[k]
                self.info[k] = (myIncrement + increment, myIntegral + integral)
            else:
                self.info[k] = (increment, integral)
        # don't forget to update keys
        self.keys = self.info.keys()
 
    def __repr__(self):
        s = "LumiFSR Python class\n"
        s += "\tRuns : \n"
        for r in self.runs:
            s += f"\t\t{r}\n"
        s += "\tFiles : \n"
        for f in self.files:
            s += f"\t\t{f}\n"
        s += "\tInfo : \n"
        for k in self.keys:
            increment, integral = self.info[k]
            s += f"\t\t{k}\t{increment}\t{integral}\n"
        return s
 
 
# =============================================================================
 
 
class PackedCaloHypo:
    def __init__(self, o):
        cl = "LHCb::PackedCaloHypo"
        assert o.__class__.__name__ == cl
        self.centX = o.centX
        self.centY = o.centY
        self.cerr = (o.cerr00, o.cerr10, o.cerr11)
        self.cov = (o.cov00, o.cov10, o.cov11, o.cov20, o.cov21, o.cov22)
        self.firstCluster = o.firstCluster
        self.firstDigit = o.firstDigit
        self.firstHypo = o.firstHypo
        self.hypothesis = o.hypothesis
        self.key = o.key
        self.lastCluster = o.lastCluster
        self.lastDigit = o.lastDigit
        self.lastHypo = o.lastHypo
        self.lh = o.lh
        self.pos = (o.posE, o.posX, o.posY)
        self.z = o.z
 
    def __repr__(self):
        s = "PackedCaloHypo : \n"
        s += f"\tcentX        : {self.centX}\n"
        s += f"\tcentY        : {self.centY}\n"
        s += f"\tcerr         : {self.cerr}\n"
        s += f"\tcov          : {self.cov}\n"
        s += f"\tfirstCluster : {self.firstCluster}\n"
        s += f"\tfirstDigit   : {self.firstDigit}\n"
        s += f"\tfirstHypo    : {self.firstHypo}\n"
        s += f"\thypothesis   : {self.hypothesis}\n"
        s += f"\tkey          : {self.key}\n"
        s += f"\tlastCluster  : {self.lastCluster}\n"
        s += f"\tlastDigit    : {self.lastDigit}\n"
        s += f"\tlastHypo     : {self.lastHypo}\n"
        s += f"\tlh           : {self.lh}\n"
        s += f"\tpos          : {self.pos}\n"
        s += f"\tz            : {self.z}\n"
        s += "---------------------------------------\n"
        return s
 
 
# =============================================================================
 
 
class SyncMini:
    def __init__(self, event, lastEvent=None):
        self.event = event
        self.t = 0.0
        self.lastEvent = None
        if lastEvent:
            self.lastEvent = lastEvent
 
    def check(self):
        return self.event.is_set()
 
    def checkLast(self):
        return self.lastEvent.is_set()
 
    def reset(self):
        self.event.clear()
        self.t = time.time()
 
    def getTime(self):
        return self.t
 
    def set(self):
        self.event.set()
 
    def __repr__(self):
        s = "---------- SyncMini --------------\n"
        s += f"    Status : {self.event.is_set()}\n"
        s += f"         t : {self.t:5.2f}\n"
        if self.lastEvent:
            s += f"Last Event : {self.lastEvent.is_set()}\n"
        s += "----------------------------------\n"
        return s
 
 
# =============================================================================
 
 
class Syncer:
    def __init__(
        self, nWorkers, log, manyEvents=False, limit=None, step=None, firstEvent=None
    ):
        # Class to help synchronise the sub-processes
        self.limit = limit
        self.step = step
        self.d = {}
        self.manyEvents = manyEvents
 
        for i in range(-2, nWorkers):
            self.d[i] = SyncMini(Event(), lastEvent=Event())
            if self.manyEvents:
                self.limitFirst = firstEvent
 
        self.keys = list(self.d.keys())
        self.nWorkers = nWorkers
        self.log = log
 
    def syncAll(self, step="Not specified"):
        # is it this method, or is it the rolling version needed?
        # if so, drop through...
 
        if self.manyEvents:
            sc = self.syncAllRolling()
            return sc
 
        # Regular version ----------------------------
        for i in range(0, self.limit, self.step):
            if self.checkAll():
                self.log.info(f"{step} : All procs done @ {i} s")
                break
            else:
                time.sleep(self.step)
 
        # Now the time limit is up... check the status one final time
        if self.checkAll():
            self.log.info(f"All processes : {step} ok.")
            return SUCCESS
        else:
            self.log.critical(f"Some process is hanging on : {step}")
            for k in self.keys:
                self.log.critical(f"{step} : Proc/Stat : {k}/{self.d[k].check()}")
            return FAILURE
 
    def syncAllRolling(self):
        # Keep track of the progress of Event processing
        # Each process syncs after each event, so keep clearing
        #  the sync Event, and re-checking
        # Note the time between True checks too, if the time
        #  between events exceeds singleEvent, this is considered a hang
 
        # set the initial time
        begin = time.time()
        firstEv = {}
        timers = {}
        for k in self.keys:
            self.d[k].reset()
            firstEv[k] = False
            timers[k] = 0.0
 
        active = self.keys
        while True:
            # check the status of each sync object
            for k in active:
                sMini = self.d[k]
 
                if sMini.check() or sMini.checkLast():
                    if sMini.checkLast() and sMini.check():
                        # if last Event set,then event loop finished
                        active.remove(k)
                        alive = time.time() - begin
                        self.log.info(f"Audit : Node {k} alive for {alive:5.2}")
                    else:
                        sMini.reset()
                else:
                    # the event still has not been checked, how long is that?
                    # is it the first Event?
                    wait = time.time() - sMini.getTime()
                    cond = wait > self.limit
                    if not firstEv[k]:
                        cond = wait > self.limitFirst
                        firstEv[k] = True
                    if cond:
                        # It is hanging!
                        self.log.critical(f"Single event wait : {wait:5.2f}")
                        self.processHang()
                        return FAILURE
 
            # Termination Criteria : if all procs have been removed, we're done
            if self.checkLastEvents():
                self.log.info("TC met for event loop")
                break
            else:
                # sleep, loop again
                time.sleep(self.step)
 
        self.log.info("All processes Completed all Events ok")
        return SUCCESS
 
    def processHang(self):
        self.log.critical("Some proc is hanging during Event processing!")
        for k in self.keys:
            self.log.critical(f"Proc/Stat : {k} / {self.d[k].check()}")
        return
 
    def checkAll(self):
        # Check the status of each Sync object
        # return True or False
        currentStatus = [mini.check() for mini in self.d.values()]
        return all(currentStatus)
 
    def checkLastEvents(self):
        # check if all of the lastEvents are set to true in self.d[k][1]
        stat = [sMini.checkLast() for sMini in self.d.values()]
        return all(stat)
 
 
# =========================== Methods =========================================
 
 
def getEventNumber(evt):
    # The class-independent version of the Event Number Retrieval method
    #
    n = None
    # First Attempt : Unpacked Event Data
    lst = ["/Event/Gen/Header", "/Event/Rec/Header"]
    for l in lst:
        try:
            n = evt[l].evtNumber()
            return n
        except Exception:
            # No evt number at this path
            continue
 
    # second attepmt : try DAQ/RawEvent data
    # The Evt Number is in bank type 16, bank 0, data pt 4
    try:
        n = evt["/Event/DAQ/RawEvent"].banks(16)[0].data()[4]
        return n
    except Exception:
        pass
 
    # Default Action
    return n
 
 
# ================================= EOF =======================================