<commit_before>from celery import Celery
from fabric import state
from fabric.api import cd, run
from kombu import Queue
from logger import logger
from spring.wgen import WorkloadGen

from perfrunner.helpers.misc import uhex
from perfrunner.settings import BROKER_URL, REPO

CELERY_QUEUES = (Queue('Q1'), Queue('Q2'))
celery = Celery('workers', backend='amqp', broker=BROKER_URL)


@celery.task
def task_run_workload(settings, target, shutdown_event, ddocs):
    wg = WorkloadGen(settings, target, shutdown_event, ddocs)
    wg.run()


class WorkerManager(object):

    def __init__(self, cluster_spec):
        self.hosts = cluster_spec.get_workers()
        if self.hosts:
            self.is_remote = True

            state.env.user, state.env.password = \
                cluster_spec.get_ssh_credentials()
            state.output.running = False
            state.output.stdout = False

            self.temp_dir = '/tmp/{}'.format(uhex()[:12])
            self._initialize_project()
            self._start()
        else:
            self.is_remote = False

    def _initialize_project(self):
        for i, q in enumerate(CELERY_QUEUES):
            logger.info('Intializing remote worker environment')

            state.env.host_string = self.hosts[i]
            temp_dir = '{}-{}'.format(self.temp_dir, q.name)
            run('mkdir {}'.format(temp_dir))
            with cd(temp_dir):
                run('git clone {}'.format(REPO))
            with cd('{}/perfrunner'.format(temp_dir)):
                run('virtualenv -p python2.7 env')
                run('env/bin/pip install -r requirements.txt')

    def _start(self):
        for i, q in enumerate(CELERY_QUEUES):
            logger.info('Starting remote Celery worker')

            state.env.host_string = self.hosts[i]
            temp_dir = '{}-{}'.format(self.temp_dir, q.name)
            with cd('{}/perfrunner'.format(temp_dir)):
                run('dtach -n /tmp/perfrunner_{}.sock '
                    'env/bin/celery worker '
                    '-A perfrunner.helpers.worker -Q {} -c 1'.format(q.name))

    def run_workload(self, settings, target_iterator, shutdown_event=None,
                     ddocs=None):
        queues = (q.name for q in CELERY_QUEUES)
        curr_target = None
        curr_queue = None
        workers = []
        for target in target_iterator:
            if self.is_remote:
                logger.info('Starting workload generator remotely')
                if curr_target != target.node:
                    curr_target = target.node
                    curr_queue = queues.next()
                workers.append(task_run_workload.apply_async(
                    args=(settings, target, shutdown_event, ddocs),
                    queue=curr_queue))
            else:
                logger.info('Starting workload generator locally')
                task_run_workload.apply(
                    args=(settings, target, shutdown_event, ddocs)
                )
        for worker in workers:
            worker.wait()

    def terminate(self):
        if self.is_remote:
            for i, q in enumerate(CELERY_QUEUES):
                state.env.host_string = self.hosts[i]
                temp_dir = '{}-{}'.format(self.temp_dir, q.name)
                logger.info('Terminating remote Celery worker')
                run('killall -9 celery; exit 0')
                logger.info('Cleaning up remote worker environment')
                run('rm -fr {}'.format(temp_dir))
<commit_msg>fix queue name argument<commit_after>from celery import Celery
from fabric import state
from fabric.api import cd, run
from kombu import Queue
from logger import logger
from spring.wgen import WorkloadGen

from perfrunner.helpers.misc import uhex
from perfrunner.settings import BROKER_URL, REPO

CELERY_QUEUES = (Queue('Q1'), Queue('Q2'))
celery = Celery('workers', backend='amqp', broker=BROKER_URL)


@celery.task
def task_run_workload(settings, target, shutdown_event, ddocs):
    wg = WorkloadGen(settings, target, shutdown_event, ddocs)
    wg.run()


class WorkerManager(object):

    def __init__(self, cluster_spec):
        self.hosts = cluster_spec.get_workers()
        if self.hosts:
            self.is_remote = True

            state.env.user, state.env.password = \
                cluster_spec.get_ssh_credentials()
            state.output.running = False
            state.output.stdout = False

            self.temp_dir = '/tmp/{}'.format(uhex()[:12])
            self._initialize_project()
            self._start()
        else:
            self.is_remote = False

    def _initialize_project(self):
        for i, q in enumerate(CELERY_QUEUES):
            logger.info('Intializing remote worker environment')

            state.env.host_string = self.hosts[i]
            temp_dir = '{}-{}'.format(self.temp_dir, q.name)
            run('mkdir {}'.format(temp_dir))
            with cd(temp_dir):
                run('git clone {}'.format(REPO))
            with cd('{}/perfrunner'.format(temp_dir)):
                run('virtualenv -p python2.7 env')
                run('env/bin/pip install -r requirements.txt')

    def _start(self):
        for i, q in enumerate(CELERY_QUEUES):
            logger.info('Starting remote Celery worker')

            state.env.host_string = self.hosts[i]
            temp_dir = '{}-{}'.format(self.temp_dir, q.name)
            with cd('{}/perfrunner'.format(temp_dir)):
                run('dtach -n /tmp/perfrunner_{0}.sock '
                    'env/bin/celery worker '
                    '-A perfrunner.helpers.worker -Q {0} -c 1'.format(q.name))

    def run_workload(self, settings, target_iterator, shutdown_event=None,
                     ddocs=None):
        queues = (q.name for q in CELERY_QUEUES)
        curr_target = None
        curr_queue = None
        workers = []
        for target in target_iterator:
            if self.is_remote:
                logger.info('Starting workload generator remotely')
                if curr_target != target.node:
                    curr_target = target.node
                    curr_queue = queues.next()
                workers.append(task_run_workload.apply_async(
                    args=(settings, target, shutdown_event, ddocs),
                    queue=curr_queue))
            else:
                logger.info('Starting workload generator locally')
                task_run_workload.apply(
                    args=(settings, target, shutdown_event, ddocs)
                )
        for worker in workers:
            worker.wait()

    def terminate(self):
        if self.is_remote:
            for i, q in enumerate(CELERY_QUEUES):
                state.env.host_string = self.hosts[i]
                temp_dir = '{}-{}'.format(self.temp_dir, q.name)
                logger.info('Terminating remote Celery worker')
                run('killall -9 celery; exit 0')
                logger.info('Cleaning up remote worker environment')
                run('rm -fr {}'.format(temp_dir))

{'fraction_non_alphanumeric': 0.10794645400170891, 'fraction_numerical': 0.0028481913984619764, 'mean_word_length': 1.9632911392405064, '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': '6247139', 'n_tokens_mistral': 2105, 'n_tokens_neox': 1966, 'n_words': 454}