nopaque/web/app/tasks/job_utils.py

from datetime import datetime
from flask import current_app
from . import docker_client
from .. import db
from ..models import JobResult
import docker
import logging
import json
import os


def create_job_service(job):
    job_dir = os.path.join(current_app.config['DATA_DIR'],
                           str(job.user_id),
                           'jobs',
                           str(job.id))
    cmd = '{} -i /files -o /files/output'.format(job.service)
    if job.service == 'file-setup':
        cmd += ' -f {}'.format(job.secure_filename)
    cmd += ' --log-dir /files'
    cmd += ' --zip [{}]_{}'.format(job.service, job.secure_filename)
    cmd += ' ' + ' '.join(json.loads(job.service_args))
    service_args = {'command': cmd,
                    'constraints': ['node.role==worker'],
                    'labels': {'origin': 'nopaque',
                               'type': 'service.{}'.format(job.service),
                               'job_id': str(job.id)},
                    'mounts': [job_dir + ':/files:rw'],
                    'name': 'job_{}'.format(job.id),
                    'resources': docker.types.Resources(
                        cpu_reservation=job.n_cores * (10 ** 9),
                        mem_reservation=job.mem_mb * (10 ** 6)),
                    'restart_policy': docker.types.RestartPolicy()}
    service_image = ('gitlab.ub.uni-bielefeld.de:4567/sfb1288inf/'
                     + job.service + ':' + job.service_version)
    try:
        docker_client.services.create(service_image, **service_args)
    except docker.errors.APIError as e:
        logging.error('create_job_service({}): {} '.format(job.id, e)
                      + '(status: {} -> failed)'.format(job.status))
        job.status = 'failed'
    else:
        job.status = 'queued'
    finally:
        # TODO: send email
        pass


def checkout_job_service(job):
    service_name = 'job_{}'.format(job.id)
    try:
        service = docker_client.services.get(service_name)
    except docker.errors.NotFound as e:
        logging.error('checkout_job_service({}): {} '.format(job.id, e)
                      + '(status: {} -> submitted)'.format(job.status))
        job.status = 'submitted'
    # TODO: handle docker.errors.APIError and docker.errors.InvalidVersion
    else:
        service_tasks = service.tasks()
        if not service_tasks:
            return
        task_state = service_tasks[0].get('Status').get('State')
        if job.status == 'queued' and task_state != 'pending':
            job.status = 'running'
        elif job.status == 'queued' and task_state == 'complete':
            service.remove()
            job.end_date = datetime.utcnow()
            job.status = task_state
            if task_state == 'complete':
                results_dir = os.path.join(current_app.config['DATA_DIR'],
                                           str(job.user_id),
                                           'jobs',
                                           str(job.id),
                                           'output')
                results = filter(lambda x: x.endswith('.zip'),
                                 os.listdir(results_dir))
                for result in results:
                    job_result = JobResult(dir=results_dir,
                                           filename=result,
                                           job_id=job.id)
                    db.session.add(job_result)
        elif job.status == 'running' and task_state == 'failed':
            service.remove()
            job.end_date = datetime.utcnow()
            job.status = task_state
    finally:
        # TODO: send email
        pass


def remove_job_service(job):
    service_name = 'job_{}'.format(job.id)
    try:
        service = docker_client.services.get(service_name)
    except docker.errors.NotFound:
        # TODO: send email
        job.status = 'canceled'
    # TODO: handle docker.errors.APIError and docker.errors.InvalidVersion
    else:
        service.update(mounts=None)
        service.remove()