<commit_before>from collections import defaultdict
import os
import subprocess
import yaml
import multiprocessing
import gzip
import jobit
import pickle
import logging

from IMProv.IMProvDoc import IMProvDoc
from ProdCommon.FwkJobRep.ReportParser import readJobReport

from lobster import job, util
import jobit
import dash

def resolve_name(job, merged_job, name, name_format):
    base, ext = os.path.splitext(name)
    id = str(job) if merged_job == 0 else 'merged_{0}'.format(merged_job)

    return name_format.format(base=base, ext=ext[1:], id=id)

class MergeHandler(object):
    def __init__(self, id, dataset, chirp, jobdir, outname, num_outputs, outname_index, sdir, jobs):
        self.__id = id
        self.__dataset = dataset
        self.__chirp = chirp
        self.__jobdir = jobdir
        self.__outname = outname
        self.__sdir = sdir
        self.__jobs = jobs

        self.__reports = []
        self.__inputs = []

        if num_outputs == 1:
            self.__tag = str(id)
            self.__base = ''
        else:
            base, ext = os.path.splitext(outname)
            self.__tag = '{0}_{1}'.format(id, outname_index)
            self.__base = '{0}_'.format(base)

    @property
    def dataset(self):
        return self.__dataset

    @property
    def jobs(self):
        return self.__jobs

    @property
    def reports(self):
        return self.__reports

    @property
    def id(self):
        return self.__id

    @property
    def inputs(self):
        return self.__inputs

    @property
    def jobdir(self):
        return self.__jobdir

    @property
    def outname(self):
        return self.__outname

    @property
    def base(self):
        return self.__base

    @property
    def tag(self):
        return self.__tag

    def get_job_update(self, failed, outsize):
        if failed:
            status = jobit.FAILED
        else:
            status = jobit.SUCCESSFUL

        return [(job, self.__id, status, outsize) for job in self.__jobs]

    def get_job_info(self):
        args = ['output=' + self.__outname]
        if self.__chirp:
            args += ['chirp={0}'.format(self.__chirp)]
            args += [str('inputs=' + ','.join(self.__inputs))]

        files = ['file:' + os.path.basename(x) for x in self.__inputs]

        return args, files

    def merge_reports(self):
        reports = []
        for r in self.__reports:
            f = gzip.open(r)
            report = readJobReport(f)
            reports.extend(report)
            f.close()

        output = IMProvDoc("JobReports")
        for item in reports:
            output.addNode(item.save())

        outfile = gzip.open(os.path.join(self.__jobdir, 'report.xml.gz'), 'wb')
        outfile.write(output.makeDOMDocument().toprettyxml())
        outfile.close()

    def cleanup(self):
#         tdir = os.path.dirname(self.__jobdir) # FIXME Do we want to delete old report.xmls?
#         for r in self.__reports:
#             os.remove(r)
        for file in self.__inputs:
            fullpath = os.path.join(self.__sdir, os.path.basename(file))
            os.remove(fullpath)

class MergeProvider(job.JobProvider):
    def __init__(self, config):
        super(MergeProvider, self).__init__(config)

        self.__chirp = self.config.get('stageout server', None)
        self.__sandbox = os.path.join(self.workdir, 'sandbox')
        self.__dash = dash.DummyMonitor(self.taskid)
        self.__mergehandlers = {}

        self.__store = jobit.JobitStore(self.config)
        self.__store.reset_merging()
        logging.info("registering unmerged jobs")
        self.__store.register_unmerged(config.get('max megabytes', 3500))

        self.__grid_files = [(os.path.join('/cvmfs/grid.cern.ch', x), os.path.join('grid', x)) for x in
                                 ['3.2.11-1/external/etc/profile.d/clean-grid-env-funcs.sh',
                                  '3.2.11-1/external/etc/profile.d/grid-env-funcs.sh',
                                  '3.2.11-1/external/etc/profile.d/grid-env.sh',
                                  '3.2.11-1/etc/profile.d/grid-env.sh',
                                  '3.2.11-1/glite/bin/voms-proxy-info',
                                  '3.2.11-1/glite/lib64/libvomsapi_nog.so.0.0.0',
                                  '3.2.11-1/glite/lib64/libvomsapi_nog.so.0',
                                  'etc/grid-security/certificates'
                                  ]
                             ]

        self.__common_inputs = [(self.__sandbox + ".tar.bz2", "sandbox.tar.bz2"),
                                (os.path.join(os.path.dirname(__file__), 'data', 'mtab'), 'mtab'),
                                (os.path.join(os.path.dirname(__file__), 'data', 'siteconfig'), 'siteconfig'),
                                (os.path.join(os.path.dirname(__file__), 'data', 'wrapper.sh'), 'wrapper.sh'),
                                (self.parrot_bin, 'bin'),
                                (self.parrot_lib, 'lib'),
                                (os.path.join(os.path.dirname(__file__), 'data', 'job.py'), 'job.py'),
                                (os.path.join(os.path.dirname(__file__), 'data', 'merge_cfg.py'), 'merge_cfg.py')
                                ] + self.__grid_files

        if not util.checkpoint(self.workdir, 'sandbox'):
            raise NotImplementedError

    def get_report(self, label, job, merge_job):
        if merge_job == 0:
            jobdir = self.get_jobdir(job, label, 'successful')
        else:
            jobdir = self.get_jobdir(merge_job, label, 'merged')
        return os.path.join(jobdir, 'report.xml.gz')

    def obtain(self, num=1):
        unmerged_jobs = self.retry(self.__store.pop_unmerged_jobs, (num,), {})
        if not unmerged_jobs or len(unmerged_jobs) == 0:
            return None

        tasks = []
        for merging_job, dset, jobs in unmerged_jobs:
            out_tag = 'merged_{0}'.format(merging_job)

            monitorid, syncid = self.__dash.register_job(merging_job)

            sdir = os.path.join(self.stageout, dset)
            jdir = self.create_jobdir(merging_job, dset, 'merging')

            for outname_index, local_outname in enumerate(self.outputs[dset]):
                base, ext = os.path.splitext(local_outname)<commit_msg>Changing naming to match jobit.py.<commit_after>from collections import defaultdict
import os
import subprocess
import yaml
import multiprocessing
import gzip
import jobit
import pickle
import logging

from IMProv.IMProvDoc import IMProvDoc
from ProdCommon.FwkJobRep.ReportParser import readJobReport

from lobster import job, util
import jobit
import dash

def resolve_name(job, merged_job, name, name_format):
    base, ext = os.path.splitext(name)
    id = str(job) if merged_job == 0 else 'merged_{0}'.format(merged_job)

    return name_format.format(base=base, ext=ext[1:], id=id)

class MergeHandler(object):
    def __init__(self, id, dataset, chirp, jobdir, outname, num_outputs, outname_index, sdir, jobs):
        self.__id = id
        self.__dataset = dataset
        self.__chirp = chirp
        self.__jobdir = jobdir
        self.__outname = outname
        self.__sdir = sdir
        self.__jobs = jobs

        self.__reports = []
        self.__inputs = []

        if num_outputs == 1:
            self.__tag = str(id)
            self.__base = ''
        else:
            base, ext = os.path.splitext(outname)
            self.__tag = '{0}_{1}'.format(id, outname_index)
            self.__base = '{0}_'.format(base)

    @property
    def dataset(self):
        return self.__dataset

    @property
    def jobs(self):
        return self.__jobs

    @property
    def reports(self):
        return self.__reports

    @property
    def id(self):
        return self.__id

    @property
    def inputs(self):
        return self.__inputs

    @property
    def jobdir(self):
        return self.__jobdir

    @property
    def outname(self):
        return self.__outname

    @property
    def base(self):
        return self.__base

    @property
    def tag(self):
        return self.__tag

    def get_job_update(self, failed, outsize):
        if failed:
            status = jobit.FAILED
        else:
            status = jobit.SUCCESSFUL

        return [(job, self.__id, status, outsize) for job in self.__jobs]

    def get_job_info(self):
        args = ['output=' + self.__outname]
        if self.__chirp:
            args += ['chirp={0}'.format(self.__chirp)]
            args += [str('inputs=' + ','.join(self.__inputs))]

        files = ['file:' + os.path.basename(x) for x in self.__inputs]

        return args, files

    def merge_reports(self):
        reports = []
        for r in self.__reports:
            f = gzip.open(r)
            report = readJobReport(f)
            reports.extend(report)
            f.close()

        output = IMProvDoc("JobReports")
        for item in reports:
            output.addNode(item.save())

        outfile = gzip.open(os.path.join(self.__jobdir, 'report.xml.gz'), 'wb')
        outfile.write(output.makeDOMDocument().toprettyxml())
        outfile.close()

    def cleanup(self):
#         tdir = os.path.dirname(self.__jobdir) # FIXME Do we want to delete old report.xmls?
#         for r in self.__reports:
#             os.remove(r)
        for file in self.__inputs:
            fullpath = os.path.join(self.__sdir, os.path.basename(file))
            os.remove(fullpath)

class MergeProvider(job.JobProvider):
    def __init__(self, config):
        super(MergeProvider, self).__init__(config)

        self.__chirp = self.config.get('stageout server', None)
        self.__sandbox = os.path.join(self.workdir, 'sandbox')
        self.__dash = dash.DummyMonitor(self.taskid)
        self.__mergehandlers = {}

        self.__store = jobit.JobitStore(self.config)
        self.__store.reset_merging()
        logging.info("registering unmerged jobs")
        self.__store.register_unmerged(config.get('max megabytes', 3500))

        self.__grid_files = [(os.path.join('/cvmfs/grid.cern.ch', x), os.path.join('grid', x)) for x in
                                 ['3.2.11-1/external/etc/profile.d/clean-grid-env-funcs.sh',
                                  '3.2.11-1/external/etc/profile.d/grid-env-funcs.sh',
                                  '3.2.11-1/external/etc/profile.d/grid-env.sh',
                                  '3.2.11-1/etc/profile.d/grid-env.sh',
                                  '3.2.11-1/glite/bin/voms-proxy-info',
                                  '3.2.11-1/glite/lib64/libvomsapi_nog.so.0.0.0',
                                  '3.2.11-1/glite/lib64/libvomsapi_nog.so.0',
                                  'etc/grid-security/certificates'
                                  ]
                             ]

        self.__common_inputs = [(self.__sandbox + ".tar.bz2", "sandbox.tar.bz2"),
                                (os.path.join(os.path.dirname(__file__), 'data', 'mtab'), 'mtab'),
                                (os.path.join(os.path.dirname(__file__), 'data', 'siteconfig'), 'siteconfig'),
                                (os.path.join(os.path.dirname(__file__), 'data', 'wrapper.sh'), 'wrapper.sh'),
                                (self.parrot_bin, 'bin'),
                                (self.parrot_lib, 'lib'),
                                (os.path.join(os.path.dirname(__file__), 'data', 'job.py'), 'job.py'),
                                (os.path.join(os.path.dirname(__file__), 'data', 'merge_cfg.py'), 'merge_cfg.py')
                                ] + self.__grid_files

        if not util.checkpoint(self.workdir, 'sandbox'):
            raise NotImplementedError

    def get_report(self, label, job, merged_job):
        if merged_job == 0:
            jobdir = self.get_jobdir(job, label, 'successful')
        else:
            jobdir = self.get_jobdir(merged_job, label, 'merged')
        return os.path.join(jobdir, 'report.xml.gz')

    def obtain(self, num=1):
        unmerged_jobs = self.retry(self.__store.pop_unmerged_jobs, (num,), {})
        if not unmerged_jobs or len(unmerged_jobs) == 0:
            return None

        tasks = []
        for merging_job, dset, jobs in unmerged_jobs:
            out_tag = 'merged_{0}'.format(merging_job)

            monitorid, syncid = self.__dash.register_job(merging_job)

            sdir = os.path.join(self.stageout, dset)
            jdir = self.create_jobdir(merging_job, dset, 'merging')

            for outname_index, local_outname in enumerate(self.outputs[dset]):
                base, ext = os.path.splitext(local_outname)
{'fraction_non_alphanumeric': 0.1233201581027668, 'fraction_numerical': 0.009644268774703557, 'mean_word_length': 1.8771889924948828, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 8}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 18, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '476501', 'n_tokens_mistral': 4053, 'n_tokens_neox': 3849, 'n_words': 755}