<commit_before>"""Core dispatcher logic"""
from aiodocker.exceptions import DockerError
import asyncio
from enum import Enum
import logging
from aioredis import RedisError
import os
import time

from .download import download
from .mail import send_job_mail, send_error_mail
from .models import Control, Job


class JobOutcome(Enum):
    SUCCESS = 'success'
    FAILURE = 'failure'
    TIMEOUT = 'timeout'


async def dispatch(app):
    """Run the dispatcher main process."""
    pool = app['engine']
    db = await pool.acquire()
    run_conf = app['run_conf']
    run_conf.up()
    while True:
        try:
            if run_conf.want_less_jobs():
                app.logger.debug("%s shutting down a task", run_conf.name)
                run_conf.down()
                break

            uid = await db.brpoplpush(run_conf.queue, '{}:queued'.format(run_conf.name), timeout=5)
            if uid is None:
                await asyncio.sleep(5)
                continue

            job_db_conn = await pool.acquire()
            job = Job(job_db_conn, uid)
            await job.fetch()

            job.dispatcher = run_conf.name
            job.state = 'queued'
            job.status = 'queued: {}'.format(run_conf.name)
            await job.commit()

            await run_container(job, db, app)
            pool.release(job_db_conn)

        except RedisError as exc:
            app.logger.error("Got redis error: %r", exc)
            raise SystemExit()
        except asyncio.CancelledError:
            break
        except Exception as exc:
            app.logger.error("Got unhandled exception %s: '%s'", type(exc), str(exc))
            raise SystemExit()
    pool.release(db)


async def run_container(job, db, app):
    """Run a docker container for the given job

    :param job: A Job object representing the job to run
    :param db: A Redis database connection from the pool
    :param app: The app object with all the central config values
    :return:
    """
    docker = app['docker']
    timeout_tasks = app['timeout_tasks']
    containers = app['containers']
    run_conf = app['run_conf']

    # TODO: put download jobs in a separate queue, handle them in a separate task?
    if job.download != '':
        try:
            ret = await download(job, app)
        except asyncio.TimeoutError:
            ret = None
        if not ret:
            job.state = 'failed'
            job.status = 'failed: Downloading from NCBI failed.'
            await job.commit()
            return

    app.logger.debug("Dispatching job %s", job)
    job.state = 'running'
    job.status = 'running'
    await job.commit()

    await db.lrem('{}:queued'.format(run_conf.name), 1, job.job_id)
    await db.lpush('jobs:running', job.job_id)

    # start a container that will run forever
    container = await docker.containers.create(
        config={
            'Cmd': create_commandline(job, run_conf),
            'Image': run_conf.image,
            'HostConfig': create_host_config(job, run_conf),
            'User': run_conf.uid_string,
        }
    )
    await container.start()
    app.logger.debug("Started %s", container._id[:8])
    containers[container._id] = (container, job)

    event = asyncio.Future(loop=app.loop)

    task = asyncio.ensure_future(follow(container, job, event))

    def timeout_handler():
        asyncio.ensure_future(cancel(container, event), loop=app.loop)

    timeout = app.loop.call_later(run_conf.timeout, timeout_handler)
    timeout_tasks[container._id] = timeout

    res, warnings, errors = await event
    if res == JobOutcome.SUCCESS:
        timeout.cancel()
        job.state = 'done'
        job.status = 'done'
    elif res == JobOutcome.FAILURE:
        timeout.cancel()
        job.state = 'failed'

        backtrace = [l.strip() for l in await container.log(stderr=True, tail=50)]
        await send_error_mail(app, job, warnings, errors, backtrace)

        if not errors:
            errors = backtrace

        error_lines = '\n'.join(errors)
        job.status = 'failed: Job returned errors: \n{}'.format(error_lines)

    else:
        task.cancel()
        job.state = 'failed'
        job.status = 'failed: Runtime exceeded'
        await send_error_mail(app, job, warnings, errors, [])

    del timeout_tasks[container._id]

    await job.commit()

    await db.lrem('jobs:running', 1, job.job_id)
    await db.lpush('jobs:completed', job.job_id)

    app.logger.debug('Done with %s', container._id[:8])

    await container.delete(force=True)
    del containers[container._id]

    await send_job_mail(app, job, warnings, errors)

    app.logger.debug('Finished job %s', job)


async def follow(container, job, event):
    """Follow a container log

    :param container: a DockerContainer to follow
    :param job: the Job object running on the container
    :param event: the Future the parent task uses to track this run
    """
    timestamp = 0
    warnings = set()
    errors = set()

    def order_logs(logs):
        return sorted(list(logs))

    while True:
        try:
            log = await container.log(stderr=True, stdout=True, since=timestamp)
            for line in log:
                line = line.strip()

                if line.startswith('INFO'):
                    job.status = 'running: {}'.format(line[25:])
                    job.changed()
                    await job.commit()
                elif line.startswith('WARNING'):
                    warnings.add(line)
                elif line.startswith('ERROR'):
                    errors.add(line)

                if line.endswith('SUCCESS'):
                    event.set_result((JobOutcome.SUCCESS, order_logs(warnings), order_logs(errors)))
                    return
                elif line.endswith('FAILED'):
                    event.set_result((JobOutcome.FAILURE, order_logs(warnings), order_logs(errors)))
                    return
                timestamp = int(time.time())
            asyncio.sleep(5)
        except asyncio.TimeoutError:
            # Docker is dumb and times out after 5 minutes, just retry
            pass<commit_msg>core: Don't busy-loop reading the container logs<commit_after>"""Core dispatcher logic"""
from aiodocker.exceptions import DockerError
import asyncio
from enum import Enum
import logging
from aioredis import RedisError
import os
import time

from .download import download
from .mail import send_job_mail, send_error_mail
from .models import Control, Job


class JobOutcome(Enum):
    SUCCESS = 'success'
    FAILURE = 'failure'
    TIMEOUT = 'timeout'


async def dispatch(app):
    """Run the dispatcher main process."""
    pool = app['engine']
    db = await pool.acquire()
    run_conf = app['run_conf']
    run_conf.up()
    while True:
        try:
            if run_conf.want_less_jobs():
                app.logger.debug("%s shutting down a task", run_conf.name)
                run_conf.down()
                break

            uid = await db.brpoplpush(run_conf.queue, '{}:queued'.format(run_conf.name), timeout=5)
            if uid is None:
                await asyncio.sleep(5)
                continue

            job_db_conn = await pool.acquire()
            job = Job(job_db_conn, uid)
            await job.fetch()

            job.dispatcher = run_conf.name
            job.state = 'queued'
            job.status = 'queued: {}'.format(run_conf.name)
            await job.commit()

            await run_container(job, db, app)
            pool.release(job_db_conn)

        except RedisError as exc:
            app.logger.error("Got redis error: %r", exc)
            raise SystemExit()
        except asyncio.CancelledError:
            break
        except Exception as exc:
            app.logger.error("Got unhandled exception %s: '%s'", type(exc), str(exc))
            raise SystemExit()
    pool.release(db)


async def run_container(job, db, app):
    """Run a docker container for the given job

    :param job: A Job object representing the job to run
    :param db: A Redis database connection from the pool
    :param app: The app object with all the central config values
    :return:
    """
    docker = app['docker']
    timeout_tasks = app['timeout_tasks']
    containers = app['containers']
    run_conf = app['run_conf']

    # TODO: put download jobs in a separate queue, handle them in a separate task?
    if job.download != '':
        try:
            ret = await download(job, app)
        except asyncio.TimeoutError:
            ret = None
        if not ret:
            job.state = 'failed'
            job.status = 'failed: Downloading from NCBI failed.'
            await job.commit()
            return

    app.logger.debug("Dispatching job %s", job)
    job.state = 'running'
    job.status = 'running'
    await job.commit()

    await db.lrem('{}:queued'.format(run_conf.name), 1, job.job_id)
    await db.lpush('jobs:running', job.job_id)

    # start a container that will run forever
    container = await docker.containers.create(
        config={
            'Cmd': create_commandline(job, run_conf),
            'Image': run_conf.image,
            'HostConfig': create_host_config(job, run_conf),
            'User': run_conf.uid_string,
        }
    )
    await container.start()
    app.logger.debug("Started %s", container._id[:8])
    containers[container._id] = (container, job)

    event = asyncio.Future(loop=app.loop)

    task = asyncio.ensure_future(follow(container, job, event))

    def timeout_handler():
        asyncio.ensure_future(cancel(container, event), loop=app.loop)

    timeout = app.loop.call_later(run_conf.timeout, timeout_handler)
    timeout_tasks[container._id] = timeout

    res, warnings, errors = await event
    if res == JobOutcome.SUCCESS:
        timeout.cancel()
        job.state = 'done'
        job.status = 'done'
    elif res == JobOutcome.FAILURE:
        timeout.cancel()
        job.state = 'failed'

        backtrace = [l.strip() for l in await container.log(stderr=True, tail=50)]
        await send_error_mail(app, job, warnings, errors, backtrace)

        if not errors:
            errors = backtrace

        error_lines = '\n'.join(errors)
        job.status = 'failed: Job returned errors: \n{}'.format(error_lines)

    else:
        task.cancel()
        job.state = 'failed'
        job.status = 'failed: Runtime exceeded'
        await send_error_mail(app, job, warnings, errors, [])

    del timeout_tasks[container._id]

    await job.commit()

    await db.lrem('jobs:running', 1, job.job_id)
    await db.lpush('jobs:completed', job.job_id)

    app.logger.debug('Done with %s', container._id[:8])

    await container.delete(force=True)
    del containers[container._id]

    await send_job_mail(app, job, warnings, errors)

    app.logger.debug('Finished job %s', job)


async def follow(container, job, event):
    """Follow a container log

    :param container: a DockerContainer to follow
    :param job: the Job object running on the container
    :param event: the Future the parent task uses to track this run
    """
    timestamp = 0
    warnings = set()
    errors = set()

    def order_logs(logs):
        return sorted(list(logs))

    while True:
        try:
            log = await container.log(stderr=True, stdout=True, since=timestamp)
            for line in log:
                line = line.strip()

                if line.startswith('INFO'):
                    job.status = 'running: {}'.format(line[25:])
                    job.changed()
                    await job.commit()
                elif line.startswith('WARNING'):
                    warnings.add(line)
                elif line.startswith('ERROR'):
                    errors.add(line)

                if line.endswith('SUCCESS'):
                    event.set_result((JobOutcome.SUCCESS, order_logs(warnings), order_logs(errors)))
                    return
                elif line.endswith('FAILED'):
                    event.set_result((JobOutcome.FAILURE, order_logs(warnings), order_logs(errors)))
                    return
                timestamp = int(time.time())
            await asyncio.sleep(5)
        except asyncio.TimeoutError:
            # Docker is dumb and times out after 5 minutes, just retry
            await asyncio.sleep(5)
            pass
{'fraction_non_alphanumeric': 0.10707332900713822, 'fraction_numerical': 0.002190136275146009, 'mean_word_length': 2.182498709344347, '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': 12, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2051080', 'n_tokens_mistral': 3601, 'n_tokens_neox': 3485, 'n_words': 966}