<commit_before>import copy
import time

from rq.job import Status
from rq import Queue, cancel_job
from rq import Connection

from .objs import KitchenSinkJob, KitchenSinkRedisQueue, pull_intermediate_results

class TaskQueue(object):
    def __init__(self, redis_conn):
        self.conn = redis_conn
        self.queues = {}

    def get_queue(self, name):
        if not name in self.queues:
            with Connection(self.conn):
                queue = KitchenSinkRedisQueue(name)
                self.queues[name] = queue
        return self.queues[name]

    def enqueue(self, queue_names, func, args, kwargs, metadata={}, interval=0.01):

        job = KitchenSinkJob.create(
            func, args, kwargs, connection=self.conn,
            result_ttl=None, status=Status.QUEUED,
            description=None, depends_on=None, timeout=None)
        for k,v in metadata.items():
            job.meta[k] = v
        job.save()
        for queue_name in queue_names:
            queue = self.get_queue(queue_name)
            queue.enqueue_job(job)
            #time.sleep(interval)
        return job._id, job.get_status()

    # def enqueue(self, queue_name, func, args, kwargs, metadata={}):
    #     queue = self.get_queue(queue_name)
    #     job = self.get_queue(queue_name).enqueue_call(func, args=args, kwargs=kwargs)
    #     for k,v in metadata.items():
    #         job.meta[k] = v
    #     job.save()
    #     return job._id, job.get_status()

    def bulkstatus(self, job_ids, timeout=5.0):
        with Connection(self.conn):
            jobs = [KitchenSinkJob(job_id) for job_id in job_ids]
        statuses = [job.get_status() for job in jobs]
        print ('****STATUSES', statuses)
        if any([x in {Status.FINISHED, Status.FAILED} for x in statuses]):
            messages = pull_intermediate_results(self.conn, job_ids, timeout=0)
        else:
            messages = pull_intermediate_results(self.conn, job_ids, timeout=timeout)
        statuses = [job.get_status() for job in jobs]
        metadata = {}
        results = {}
        messages_by_job = {}
        for job_id, msg in messages:
            messages_by_job.setdefault(job_id, []).append(msg)
        for job_id, job, status in zip(job_ids, jobs, statuses):
            job.refresh()
            mdata = copy.copy(job.meta)
            mdata['status'] = status
            mdata['msgs'] = messages_by_job.get(job_id, [])
            metadata[job_id] = mdata
            if status == Status.FINISHED:
                if job.return_value is None:
                    pass
                    #import pdb; pdb.set_trace()
                results[job_id] = job.return_value
            elif status == Status.FAILED:
                results[job_id] = job.exc_info
            else:
                results[job_id] = None
        return [(metadata[job_id], results[job_id]) for job_id in job_ids]

    def status(self, job_id, timeout=None):
        """timeout is None - default to non-blocking
        """
        with Connection(self.conn):
            job = KitchenSinkJob(job_id)
        job.refresh()
        status = job.get_status()
        # if the job is finished, grab all messages and get out!
        if status == Status.FINISHED:
            messages = job.pull_intermediate_results(timeout=0)
        else:
            messages = job.pull_intermediate_results()
            #then refresh the job again
            job.refresh()
            status = job.get_status()
        metadata = copy.copy(job.meta)
        metadata['status'] = status
        metadata['msgs'] = messages

        if status == Status.FINISHED:
            return metadata, job.return_value
        elif status == Status.FAILED:
            return metadata, job.exc_info
        else:
            return metadata, None
    def cancel(self, job_id):
        with Connection(self.conn):
            cancel_job(job_id)
<commit_msg>removing debug info<commit_after>import copy
import time

from rq.job import Status
from rq import Queue, cancel_job
from rq import Connection

from .objs import KitchenSinkJob, KitchenSinkRedisQueue, pull_intermediate_results

class TaskQueue(object):
    def __init__(self, redis_conn):
        self.conn = redis_conn
        self.queues = {}

    def get_queue(self, name):
        if not name in self.queues:
            with Connection(self.conn):
                queue = KitchenSinkRedisQueue(name)
                self.queues[name] = queue
        return self.queues[name]

    def enqueue(self, queue_names, func, args, kwargs, metadata={}, interval=0.01):

        job = KitchenSinkJob.create(
            func, args, kwargs, connection=self.conn,
            result_ttl=None, status=Status.QUEUED,
            description=None, depends_on=None, timeout=None)
        for k,v in metadata.items():
            job.meta[k] = v
        job.save()
        for queue_name in queue_names:
            queue = self.get_queue(queue_name)
            queue.enqueue_job(job)
            #time.sleep(interval)
        return job._id, job.get_status()

    # def enqueue(self, queue_name, func, args, kwargs, metadata={}):
    #     queue = self.get_queue(queue_name)
    #     job = self.get_queue(queue_name).enqueue_call(func, args=args, kwargs=kwargs)
    #     for k,v in metadata.items():
    #         job.meta[k] = v
    #     job.save()
    #     return job._id, job.get_status()

    def bulkstatus(self, job_ids, timeout=5.0):
        with Connection(self.conn):
            jobs = [KitchenSinkJob(job_id) for job_id in job_ids]
        statuses = [job.get_status() for job in jobs]
        if any([x in {Status.FINISHED, Status.FAILED} for x in statuses]):
            messages = pull_intermediate_results(self.conn, job_ids, timeout=0)
        else:
            messages = pull_intermediate_results(self.conn, job_ids, timeout=timeout)
        statuses = [job.get_status() for job in jobs]
        metadata = {}
        results = {}
        messages_by_job = {}
        for job_id, msg in messages:
            messages_by_job.setdefault(job_id, []).append(msg)
        for job_id, job, status in zip(job_ids, jobs, statuses):
            job.refresh()
            mdata = copy.copy(job.meta)
            mdata['status'] = status
            mdata['msgs'] = messages_by_job.get(job_id, [])
            metadata[job_id] = mdata
            if status == Status.FINISHED:
                results[job_id] = job.return_value
            elif status == Status.FAILED:
                results[job_id] = job.exc_info
            else:
                results[job_id] = None
        return [(metadata[job_id], results[job_id]) for job_id in job_ids]

    def status(self, job_id, timeout=None):
        """timeout is None - default to non-blocking
        """
        with Connection(self.conn):
            job = KitchenSinkJob(job_id)
        job.refresh()
        status = job.get_status()
        # if the job is finished, grab all messages and get out!
        if status == Status.FINISHED:
            messages = job.pull_intermediate_results(timeout=0)
        else:
            messages = job.pull_intermediate_results()
            #then refresh the job again
            job.refresh()
            status = job.get_status()
        metadata = copy.copy(job.meta)
        metadata['status'] = status
        metadata['msgs'] = messages

        if status == Status.FINISHED:
            return metadata, job.return_value
        elif status == Status.FAILED:
            return metadata, job.exc_info
        else:
            return metadata, None
    def cancel(self, job_id):
        with Connection(self.conn):
            cancel_job(job_id)

{'fraction_non_alphanumeric': 0.10177453027139875, 'fraction_numerical': 0.0018267223382045928, 'mean_word_length': 2.1057536466774716, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, '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': '2210751', 'n_tokens_mistral': 2267, 'n_tokens_neox': 2248, 'n_words': 549}