<commit_before>import sys, re, os, marshal, modutil, time, types, cPickle, cStringIO, random
from disco import func, util, comm
from disco.comm import json
from disco.error import DiscoError, JobException
from disco.eventmonitor import EventMonitor
from disco.netstring import encode_netstring_fd, decode_netstring_fd

class Params(object):
        def __init__(self, **kwargs):
                self._state = {}
                for k, v in kwargs.iteritems():
                        setattr(self, k, v)

        def __setattr__(self, k, v):
                if k[0] == '_':
                        self.__dict__[k] = v
                        return
                st_v = v
                st_k = "n_" + k
                try:
                        st_v = marshal.dumps(v.func_code)
                        st_k = "f_" + k
                except AttributeError:
                        pass
                self._state[st_k] = st_v
                self.__dict__[k] = v
        
        def __getstate__(self):
                return self._state

        def __setstate__(self, state):
                self._state = {}
                for k, v in state.iteritems():
                        if k.startswith('f_'):
                                t = lambda x: x
                                t.func_code = marshal.loads(v)
                                v = t
                        self.__dict__[k[2:]] = v

class Stats(object):
        def __init__(self, prof_data):
                self.stats = marshal.loads(prof_data)
        def create_stats(self):
                pass

class Disco(object):
        def __init__(self, host):
                self.host = "http://" + util.disco_host(host)[7:]

        def request(self, url, data = None, redir = False, offset = 0):
                try:
                        return comm.download(self.host + url, data = data, redir = redir, offset = offset)
                except Exception, e:
                        raise DiscoError('Got %s, make sure disco master is running at %s' % (e, self.host))

        def nodeinfo(self):
                return json.loads(self.request("/disco/ctrl/nodeinfo"))

        def joblist(self):
                return json.loads(self.request("/disco/ctrl/joblist"))
        
        def oob_get(self, name, key):
                try:
                        return util.load_oob(self.host, name, key)
                except comm.CommException, x:
                        if x.http_code == 404:
                                raise DiscoError("Unknown key or job name")
                        raise
        
        def oob_list(self, name):
                try:
                        r = self.request("/disco/ctrl/oob_list?name=%s" % name,
                                redir = True)
                except comm.CommException, x:
                        if x.http_code == 404:
                                raise DiscoError("Unknown key or job name")
                        raise
                return json.loads(r)

        def profile_stats(self, name, mode = ""):
                import pstats
                if mode:
                        prefix = "profile-%s-" % mode
                else:
                        prefix = "profile-"
                f = [s for s in self.oob_list(name) if s.startswith(prefix)]
                if not f:
                        raise JobException("No profile data", self.host, name)
                
                stats = pstats.Stats(Stats(self.oob_get(name, f[0])))
                for s in f[1:]:
                        stats.add(Stats(self.oob_get(name, s)))
                return stats
        
        def new_job(self, **kwargs):
                return Job(self, **kwargs)
        
        def kill(self, name):
                self.request("/disco/ctrl/kill_job", '"%s"' % name)
        
        def clean(self, name):
                self.request("/disco/ctrl/clean_job", '"%s"' % name)

        def purge(self, name):
                self.request("/disco/ctrl/purge_job", '"%s"' % name)

        def jobspec(self, name):
                r = self.request("/disco/ctrl/parameters?name=%s" % name,
                                 redir = True)
                return decode_netstring_fd(cStringIO.StringIO(r))

        def events(self, name, offset = 0):
                def event_iter(events):
                        offs = offset
                        lines = events.splitlines()
                        for i, l in enumerate(lines):
                                offs += len(l) + 1
                                if not len(l):
                                        continue
                                try:
                                        ent = tuple(json.loads(l))
                                except ValueError:
                                        break
                                # HTTP range request doesn't like empty ranges:
                                # Let's ensure that at least the last newline
                                # is always retrieved.
                                if i == len(lines) - 1:
                                        offs -= 1
                                yield offs, ent
                
                r = self.request("/disco/ctrl/rawevents?name=%s" % name,
                                 redir = True, offset = offset)
                
                if len(r) < 2:
                        return []
                else:
                        return event_iter(r)

        def results(self, jobspec, timeout = 2000):
                jobspec = JobSpecifier(jobspec)
                data    = json.dumps([timeout, list(jobspec.jobnames)])
                results = json.loads(self.request("/disco/ctrl/get_results", data))

                if len(jobspec) is 1:
                        return results[0][1]
                
                others, active = [], []
                for result in results:
                        if result[1][0] == 'active':
                                active.append(result)
                        else:
                                others.append(result)
                return others, active

        def jobinfo(self, name):
                jobinfo = self.request("/disco/ctrl/jobinfo?name=%s" % name)
                if jobinfo:
                        return json.loads(jobinfo)

        def wait(self, name, show = None, poll_interval = 2, timeout = None, clean = False):
                event_monitor = EventMonitor(show, disco=self, name=name)
                start_time    = time.time()
                while True:
                        status, results = self.results(name, timeout = poll_interval * 1000)
                        event_monitor.refresh()
                        if status == 'ready':
                                if clean:
                                        self.clean(name)
                                return results
                        if status != 'active':
                                raise JobException("Job status %s" % status, self.host, name)
                        if timeout and time.time() - start_time > timeout:
                                raise JobException("Timeout", self.host, name)

class JobSpecifier(list):
        def __init__(self, jobspec):
                super(JobSpecifier, self).__init__([jobspec] if type(jobspec) is str else jobspec)

        @property
        def jobnames(self):
                for job in self:
                        if type(job) is str:
                                yield job
                        elif type(job) is list:
                                yield job[0]
                        else:<commit_msg>bugfix: fixed a regression introduced by 46ccd2ce in disco.core.results() - results were not correct if jobspec was a list of only one element<commit_after>import sys, re, os, marshal, modutil, time, types, cPickle, cStringIO, random
from disco import func, util, comm
from disco.comm import json
from disco.error import DiscoError, JobException
from disco.eventmonitor import EventMonitor
from disco.netstring import encode_netstring_fd, decode_netstring_fd

class Params(object):
        def __init__(self, **kwargs):
                self._state = {}
                for k, v in kwargs.iteritems():
                        setattr(self, k, v)

        def __setattr__(self, k, v):
                if k[0] == '_':
                        self.__dict__[k] = v
                        return
                st_v = v
                st_k = "n_" + k
                try:
                        st_v = marshal.dumps(v.func_code)
                        st_k = "f_" + k
                except AttributeError:
                        pass
                self._state[st_k] = st_v
                self.__dict__[k] = v
        
        def __getstate__(self):
                return self._state

        def __setstate__(self, state):
                self._state = {}
                for k, v in state.iteritems():
                        if k.startswith('f_'):
                                t = lambda x: x
                                t.func_code = marshal.loads(v)
                                v = t
                        self.__dict__[k[2:]] = v

class Stats(object):
        def __init__(self, prof_data):
                self.stats = marshal.loads(prof_data)
        def create_stats(self):
                pass

class Disco(object):
        def __init__(self, host):
                self.host = "http://" + util.disco_host(host)[7:]

        def request(self, url, data = None, redir = False, offset = 0):
                try:
                        return comm.download(self.host + url, data = data, redir = redir, offset = offset)
                except Exception, e:
                        raise DiscoError('Got %s, make sure disco master is running at %s' % (e, self.host))

        def nodeinfo(self):
                return json.loads(self.request("/disco/ctrl/nodeinfo"))

        def joblist(self):
                return json.loads(self.request("/disco/ctrl/joblist"))
        
        def oob_get(self, name, key):
                try:
                        return util.load_oob(self.host, name, key)
                except comm.CommException, x:
                        if x.http_code == 404:
                                raise DiscoError("Unknown key or job name")
                        raise
        
        def oob_list(self, name):
                try:
                        r = self.request("/disco/ctrl/oob_list?name=%s" % name,
                                redir = True)
                except comm.CommException, x:
                        if x.http_code == 404:
                                raise DiscoError("Unknown key or job name")
                        raise
                return json.loads(r)

        def profile_stats(self, name, mode = ""):
                import pstats
                if mode:
                        prefix = "profile-%s-" % mode
                else:
                        prefix = "profile-"
                f = [s for s in self.oob_list(name) if s.startswith(prefix)]
                if not f:
                        raise JobException("No profile data", self.host, name)
                
                stats = pstats.Stats(Stats(self.oob_get(name, f[0])))
                for s in f[1:]:
                        stats.add(Stats(self.oob_get(name, s)))
                return stats
        
        def new_job(self, **kwargs):
                return Job(self, **kwargs)
        
        def kill(self, name):
                self.request("/disco/ctrl/kill_job", '"%s"' % name)
        
        def clean(self, name):
                self.request("/disco/ctrl/clean_job", '"%s"' % name)

        def purge(self, name):
                self.request("/disco/ctrl/purge_job", '"%s"' % name)

        def jobspec(self, name):
                r = self.request("/disco/ctrl/parameters?name=%s" % name,
                                 redir = True)
                return decode_netstring_fd(cStringIO.StringIO(r))

        def events(self, name, offset = 0):
                def event_iter(events):
                        offs = offset
                        lines = events.splitlines()
                        for i, l in enumerate(lines):
                                offs += len(l) + 1
                                if not len(l):
                                        continue
                                try:
                                        ent = tuple(json.loads(l))
                                except ValueError:
                                        break
                                # HTTP range request doesn't like empty ranges:
                                # Let's ensure that at least the last newline
                                # is always retrieved.
                                if i == len(lines) - 1:
                                        offs -= 1
                                yield offs, ent
                
                r = self.request("/disco/ctrl/rawevents?name=%s" % name,
                                 redir = True, offset = offset)
                
                if len(r) < 2:
                        return []
                else:
                        return event_iter(r)

        def results(self, jobspec, timeout = 2000):
                jobspec = JobSpecifier(jobspec)
                data    = json.dumps([timeout, list(jobspec.jobnames)])
                results = json.loads(self.request("/disco/ctrl/get_results", data))

                if type(jobspec) == str:
                        return results[0][1]
                
                others, active = [], []
                for result in results:
                        if result[1][0] == 'active':
                                active.append(result)
                        else:
                                others.append(result)
                return others, active

        def jobinfo(self, name):
                jobinfo = self.request("/disco/ctrl/jobinfo?name=%s" % name)
                if jobinfo:
                        return json.loads(jobinfo)

        def wait(self, name, show = None, poll_interval = 2, timeout = None, clean = False):
                event_monitor = EventMonitor(show, disco=self, name=name)
                start_time    = time.time()
                while True:
                        status, results = self.results(name, timeout = poll_interval * 1000)
                        event_monitor.refresh()
                        if status == 'ready':
                                if clean:
                                        self.clean(name)
                                return results
                        if status != 'active':
                                raise JobException("Job status %s" % status, self.host, name)
                        if timeout and time.time() - start_time > timeout:
                                raise JobException("Timeout", self.host, name)

class JobSpecifier(list):
        def __init__(self, jobspec):
                super(JobSpecifier, self).__init__([jobspec]
                        if type(jobspec) is str else jobspec)

        @property
        def jobnames(self):
                for job in self:
                        if type(job) is str:
                                yield job
                        elif type(job) is list:
                                yield job[0]
                        else:
{'fraction_non_alphanumeric': 0.08996138996138996, 'fraction_numerical': 0.0042471042471042475, 'mean_word_length': 1.0246221990620115, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 36, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '76672', 'n_tokens_mistral': 3994, 'n_tokens_neox': 3683, 'n_words': 989}