<commit_before>'''
This software is released under an MIT/X11 open source license.

Copyright 2012-2014 Diffeo, Inc.
'''
from __future__ import absolute_import
import contextlib
import os
import copy
import time
import logging
import rejester
import multiprocessing
from rejester.workers import run_worker, MultiWorker
from rejester._logging import logger

from tests.rejester.test_task_master import task_master  ## a fixture that cleans up

import pytest


def test_task_register(task_master):
    worker = MultiWorker(task_master.registry.config)
    worker_id = worker.register()
    assert worker_id in task_master.workers()
    worker.unregister()
    assert worker_id not in task_master.workers()


def work_program(work_unit):
    ## just to show that this works, we get the config from the data
    ## and *reconnect* to the registry with a second instances instead
    ## of using work_unit.registry
    config = work_unit.data['config']
    sleeptime = float(work_unit.data.get('sleep', 9.0))
    task_master = rejester.TaskMaster(config)
    logger.info('executing work_unit %r ... %s', work_unit.key, sleeptime)
    time.sleep(sleeptime)  # pretend to work

def work_program_broken(work_unit):
    logger.info('executing "broken" work_unit %r', work_unit.key)

    ## just to show that this works, we get the config from the data
    ## and *reconnect* to the registry with a second instances instead
    ## of using work_unit.registry
    config = work_unit.data['config']
    task_master = rejester.TaskMaster(config)
    raise Exception('simulate broken work_unit')

## 1KB sized work_spec config
work_spec = dict(
    name = 'tbundle',
    desc = 'a test work bundle',
    min_gb = 8,
    config = dict(many=' ' * 2**10, params=''),
    module = 'tests.rejester.test_workers',
    run_function = 'work_program',
    terminate_function = 'work_program',
)

@contextlib.contextmanager
def run_multi_worker(task_master, duration):
    """Spawn a MultiWorker and clean up from it sanely.

    Use this as:

    >>> with run_multi_worker(task_master, 60) as deadline:
    ...   while time.time() < deadline:
    ...     # do some work
    ...     time.sleep(1)

    On exiting the context block, this will do a very clean shutdown
    of the workers; if anything goes wrong this will do an unclean shutdown
    of the workers.

    """
    p = multiprocessing.Process(target=run_worker, 
                                args=(MultiWorker, task_master.registry.config))
    num_workers = multiprocessing.cpu_count()
    logger.debug('expecting num_workers=%d', num_workers)

    start = time.time()
    end = start + duration
    finished_cleanly = False
    already_set_idle = False
    already_set_terminate = False
    p.start()

    try:
        yield end

        while time.time() < end:
            modes = task_master.mode_counts()
            logger.debug(modes)

            if modes[task_master.RUN] >= num_workers:
                logger.info('setting mode to IDLE')
                task_master.idle_all_workers()
                already_set_idle = True

            if modes[task_master.IDLE] >= 0 and already_set_idle:
                logger.info('setting mode to TERMINATE')
                task_master.set_mode(task_master.TERMINATE)
                already_set_terminate = True
        
            if p.exitcode is not None:
                assert already_set_idle
                assert already_set_terminate
                if p.exitcode == 0:
                    finished_cleanly = True
                break

            time.sleep(2)

        assert finished_cleanly, ('timed out after {0} seconds'
                                  .format(time.time() - start))

        p.join()
        logger.info('finished running %d worker processes', num_workers)
    finally:
        if p.is_alive():
            p.terminate()

def test_task_master_multi_worker(task_master):
    num_units = 10
    num_units_cursor = 0
    work_units = {'key' + str(x): dict(config=task_master.registry.config) 
                  for x in xrange(num_units_cursor, num_units_cursor + num_units)}<commit_msg>test_workers: make sure to p.join() if we kill the workers<commit_after>'''
This software is released under an MIT/X11 open source license.

Copyright 2012-2014 Diffeo, Inc.
'''
from __future__ import absolute_import
import contextlib
import os
import copy
import time
import logging
import multiprocessing
import sys

import pytest

import rejester
from rejester.workers import run_worker, MultiWorker
from rejester._logging import logger
from tests.rejester.test_task_master import task_master  ## a fixture that cleans up

def test_task_register(task_master):
    worker = MultiWorker(task_master.registry.config)
    worker_id = worker.register()
    assert worker_id in task_master.workers()
    worker.unregister()
    assert worker_id not in task_master.workers()


def work_program(work_unit):
    ## just to show that this works, we get the config from the data
    ## and *reconnect* to the registry with a second instances instead
    ## of using work_unit.registry
    config = work_unit.data['config']
    sleeptime = float(work_unit.data.get('sleep', 9.0))
    task_master = rejester.TaskMaster(config)
    logger.info('executing work_unit %r ... %s', work_unit.key, sleeptime)
    time.sleep(sleeptime)  # pretend to work

def work_program_broken(work_unit):
    logger.info('executing "broken" work_unit %r', work_unit.key)

    ## just to show that this works, we get the config from the data
    ## and *reconnect* to the registry with a second instances instead
    ## of using work_unit.registry
    config = work_unit.data['config']
    task_master = rejester.TaskMaster(config)
    raise Exception('simulate broken work_unit')

## 1KB sized work_spec config
work_spec = dict(
    name = 'tbundle',
    desc = 'a test work bundle',
    min_gb = 8,
    config = dict(many=' ' * 2**10, params=''),
    module = 'tests.rejester.test_workers',
    run_function = 'work_program',
    terminate_function = 'work_program',
)

@contextlib.contextmanager
def run_multi_worker(task_master, duration):
    """Spawn a MultiWorker and clean up from it sanely.

    Use this as:

    >>> with run_multi_worker(task_master, 60) as deadline:
    ...   while time.time() < deadline:
    ...     # do some work
    ...     time.sleep(1)

    On exiting the context block, this will do a very clean shutdown
    of the workers; if anything goes wrong this will do an unclean shutdown
    of the workers.

    """
    p = multiprocessing.Process(target=run_worker, 
                                args=(MultiWorker, task_master.registry.config))
    num_workers = multiprocessing.cpu_count()
    logger.debug('expecting num_workers=%d', num_workers)

    start = time.time()
    end = start + duration
    finished_cleanly = False
    already_set_idle = False
    already_set_terminate = False
    p.start()

    try:
        yield end

        while time.time() < end:
            modes = task_master.mode_counts()
            logger.debug(modes)

            if modes[task_master.RUN] >= num_workers:
                logger.info('setting mode to IDLE')
                task_master.idle_all_workers()
                already_set_idle = True

            if modes[task_master.IDLE] >= 0 and already_set_idle:
                logger.info('setting mode to TERMINATE')
                task_master.set_mode(task_master.TERMINATE)
                already_set_terminate = True
        
            if p.exitcode is not None:
                assert already_set_idle
                assert already_set_terminate
                if p.exitcode == 0:
                    finished_cleanly = True
                break

            time.sleep(2)

        assert finished_cleanly, ('timed out after {0} seconds'
                                  .format(time.time() - start))

        p.join()
        logger.info('finished running %d worker processes', num_workers)
    finally:
        if p.is_alive():
            logger.debug("killing worker processes")
            p.terminate()
            p.join(1.0)
        if p.is_alive():
            logger.warn("worker processes failed to die, hard killing")
            os.kill(p.pid, signal.SIGTERM)
            p.join(1.0)
        if p.is_alive():
            logger.critical("worker processes resisted SIGKILL, cannot recover")
            sys.exit()

def test_task_master_multi_worker(task_master):
    num_units = 10
    num_units_cursor = 0
    work_units = {'key' + str(x): dict(config=task_master.registry.config) 
                  for x in xrange(num_units_cursor, num_units_cursor + num_units)}
{'fraction_non_alphanumeric': 0.08850375505488157, 'fraction_numerical': 0.006701328711727325, 'mean_word_length': 2.6217573221757324, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 13, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1766915', 'n_tokens_mistral': 2604, 'n_tokens_neox': 2496, 'n_words': 716}