<commit_before>import ROOT
import time
import os
import sys
import multiprocessing
from multiprocessing import Process
from operator import add, itemgetter
import uuid
from ..tree.filtering import *
from atlastools import datasets
from .. import routines
from .. import multilogging
import logging
import traceback
import shutil
try:
    import cPickle as pickle
except:
    import pickle

class Supervisor(Process):

    def __init__(self, name, outputname, fileset, nstudents, process, gridmode=False, args=None, **kwargs):
        
        Process.__init__(self) 
        self.name = name
        self.fileset = fileset
        self.outputname = outputname
        self.gridmode = gridmode
        if self.gridmode:
            self.nstudents = 1
        else:
            self.nstudents = min(nstudents, len(fileset.files))
        self.process = process
        self.students = []
        self.good_students = []
        self.student_outputs = []
        self.kwargs = kwargs
        self.logger = None
        self.args = args
        
    def run(self):

        ROOT.gROOT.SetBatch()
        
        # logging
        self.logging_queue = multiprocessing.Queue(-1)
        self.listener = multilogging.Listener("supervisor-%s-%s.log"% (self.name, self.outputname), self.logging_queue)
        self.listener.start()
        
        h = multilogging.QueueHandler(self.logging_queue)
        self.logger = logging.getLogger("Supervisor")
        self.logger.addHandler(h)
        self.logger.setLevel(logging.DEBUG)
        
        if not self.gridmode:
            sys.stdout = multilogging.stdout(self.logger)
            sys.stderr = multilogging.stderr(self.logger)
       
        try:
            self.__apply_for_grant()
            self.__supervise()
            self.__publish()
        except:
            print sys.exc_info()
            traceback.print_tb(sys.exc_info()[2])
        print "Done"
        self.logging_queue.put(None)
        self.listener.join()

    def __apply_for_grant(self):
        
        print "Will run on %i files:"% len(self.fileset.files)
        for filename in self.fileset.files:
            print "%s"% filename
        sys.stdout.flush()
        filesets = self.fileset.split(self.nstudents)
        self.output_queue = multiprocessing.Queue(-1)
        self.students = [
            self.process(
                name = self.name,
                fileset = fileset,
                output_queue = self.output_queue,
                logging_queue = self.logging_queue,
                gridmode = self.gridmode,
                args = self.args,
                **self.kwargs
            ) for fileset in filesets ]
        self.process_table = dict([(p.uuid, p) for p in self.students])
        self.good_students = []
            
    def __supervise(self):
        
        students = self.students[:]
        for p in students:
            p.start()
        nstudents = len(students)
        while students:
            while not self.output_queue.empty():
                id, output = self.output_queue.get()
                process = self.process_table[id]
                process.join()
                students.remove(process)
                if output is not None and process.exitcode == 0:
                    self.good_students.append(process)
                    self.student_outputs.append(output)
            time.sleep(1)

    def __publish(self, merge = True):
        
        if len(self.good_students) > 0:
            outputs = []
            event_filters = []
            object_filters = []
            for thing in self.student_outputs:
                event_filters.append(thing[0])
                object_filters.append(thing[1])
                outputs.append(thing[2])
            print "===== Cut-flow of event filters for dataset %s: ====\n"% self.outputname
            totalEvents = 0
            combinedEventFilterlist = reduce(FilterList.merge, event_filters)
            if len(combinedEventFilterlist) > 0:
                totalEvents = combinedEventFilterlist[0].total
            print ": Event Filters\n%s"% combinedEventFilterlist
            combinedObjectFilterlist = reduce(FilterList.merge, object_filters)
            print ": Object Filters\n%s"% combinedObjectFilterlist
            pickle.dump({"event": combinedEventFilterlist,
                         "object": combinedObjectFilterlist}, "cutflow.p")
            if merge:
                outputname = "%s.root" % self.outputname 
                if os.path.exists(outputname):
                    os.unlink(outputname)
                if len(outputs) == 1:
                    shutil.move(outputs[0], outputname)
                else:
                    os.system("hadd %s %s"%(outputname, " ".join(outputs)))
                    for output in outputs:
                        os.unlink(output)
            # set weights:
            if totalEvents != 0 and self.fileset.datatype != datasets.types['DATA']:
                outfile = ROOT.TFile.Open("%s.root"% self.outputname, "update")
                trees = routines.getTrees(outfile)
                for tree in trees:
                    tree.SetWeight(self.fileset.weight/totalEvents)
                    tree.Write("", ROOT.TObject.kOverwrite)
                outfile.Close()<commit_msg>updates<commit_after>import ROOT
import time
import os
import sys
import multiprocessing
from multiprocessing import Process
from operator import add, itemgetter
import uuid
from ..tree.filtering import *
from atlastools import datasets
from .. import routines
from .. import multilogging
import logging
import traceback
import shutil
import subprocess
try:
    import cPickle as pickle
except:
    import pickle

class Supervisor(Process):

    def __init__(self, name, outputname, fileset, nstudents, process, gridmode=False, args=None, **kwargs):
        
        Process.__init__(self) 
        self.name = name
        self.fileset = fileset
        self.outputname = outputname
        self.gridmode = gridmode
        if self.gridmode:
            self.nstudents = 1
        else:
            self.nstudents = min(nstudents, len(fileset.files))
        self.process = process
        self.students = []
        self.good_students = []
        self.student_outputs = []
        self.kwargs = kwargs
        self.logger = None
        self.args = args
        
    def run(self):

        ROOT.gROOT.SetBatch()
        
        # logging
        self.logging_queue = multiprocessing.Queue(-1)
        self.listener = multilogging.Listener("supervisor-%s-%s.log"% (self.name, self.outputname), self.logging_queue)
        self.listener.start()
        
        h = multilogging.QueueHandler(self.logging_queue)
        self.logger = logging.getLogger("Supervisor")
        self.logger.addHandler(h)
        self.logger.setLevel(logging.DEBUG)
        
        if not self.gridmode:
            sys.stdout = multilogging.stdout(self.logger)
            sys.stderr = multilogging.stderr(self.logger)
       
        try:
            self.__apply_for_grant()
            self.__supervise()
            self.__publish()
        except:
            print sys.exc_info()
            traceback.print_tb(sys.exc_info()[2])
        print "Done"
        self.logging_queue.put(None)
        self.listener.join()

    def __apply_for_grant(self):
        
        print "Will run on %i files:"% len(self.fileset.files)
        for filename in self.fileset.files:
            print "%s"% filename
        sys.stdout.flush()
        filesets = self.fileset.split(self.nstudents)
        self.output_queue = multiprocessing.Queue(-1)
        self.students = [
            self.process(
                name = self.name,
                fileset = fileset,
                output_queue = self.output_queue,
                logging_queue = self.logging_queue,
                gridmode = self.gridmode,
                args = self.args,
                **self.kwargs
            ) for fileset in filesets ]
        self.process_table = dict([(p.uuid, p) for p in self.students])
        self.good_students = []
            
    def __supervise(self):
        
        students = self.students[:]
        for p in students:
            p.start()
        nstudents = len(students)
        while students:
            while not self.output_queue.empty():
                id, output = self.output_queue.get()
                process = self.process_table[id]
                process.join()
                students.remove(process)
                if output is not None and process.exitcode == 0:
                    self.good_students.append(process)
                    self.student_outputs.append(output)
            time.sleep(1)

    def __publish(self, merge = True):
        
        if len(self.good_students) > 0:
            outputs = []
            event_filters = []
            object_filters = []
            for thing in self.student_outputs:
                event_filters.append(thing[0])
                object_filters.append(thing[1])
                outputs.append(thing[2])
            print "===== Cut-flow of event filters for dataset %s: ====\n"% self.outputname
            totalEvents = 0
            combinedEventFilterlist = reduce(FilterList.merge, event_filters)
            if len(combinedEventFilterlist) > 0:
                totalEvents = combinedEventFilterlist[0].total
            print ": Event Filters\n%s"% combinedEventFilterlist
            combinedObjectFilterlist = reduce(FilterList.merge, object_filters)
            print ": Object Filters\n%s"% combinedObjectFilterlist
            pfile = open("cutflow.p",'w')
            pickle.dump({"event": combinedEventFilterlist,
                         "object": combinedObjectFilterlist}, pfile)
            pfile.close()
            if merge:
                outputname = "%s.root" % self.outputname 
                if os.path.exists(outputname):
                    os.unlink(outputname)
                if len(outputs) == 1:
                    shutil.move(outputs[0], outputname)
                else:
                    subprocess.call(["hadd", outputname] + outputs)
                    for output in outputs:
                        os.unlink(output)
            # set weights:
            if totalEvents != 0 and self.fileset.datatype != datasets.types['DATA']:
                outfile = ROOT.TFile.Open("%s.root"% self.outputname, "update")
                trees = routines.getTrees(outfile)
                for tree in trees:
                    tree.SetWeight(self.fileset.weight/totalEvents)
                    tree.Write("", ROOT.TObject.kOverwrite)
                outfile.Close()
{'fraction_non_alphanumeric': 0.09542557180352455, 'fraction_numerical': 0.0029996250468691415, 'mean_word_length': 1.7356410256410257, 'pattern_counts': {'":': 8, '<': 3, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3365513', 'n_tokens_mistral': 2951, 'n_tokens_neox': 2858, 'n_words': 613}