nopaque/app/scheduler.py

from apscheduler.schedulers.background import BackgroundScheduler
from datetime import datetime
from sqlalchemy import create_engine
from sqlalchemy.orm import scoped_session, sessionmaker
import os


class Scheduler(BackgroundScheduler):
    def __init__(self, app=None):
        super().__init__()
        self.app = app
        if app is not None:
            self.init_app(app)

    def init_app(self, app):
        engine = create_engine(app.config['SQLALCHEMY_DATABASE_URI'])
        self.Session = scoped_session(sessionmaker(bind=engine))
        if not app.debug or os.environ.get('WERKZEUG_RUN_MAIN') == 'true':
            self.add_job(self.check_jobs, id='jobs', seconds=10, trigger='interval')
            self.start()

    def check_jobs(self):
        from .models import Job
        from . import swarm

        session = self.Session()
        jobs = session.query(Job)
        submitted_jobs = jobs.filter_by(status='submitted').all()
        foo_jobs = jobs.filter(Job.status != 'complete',
                               Job.status != 'failed',
                               Job.status != 'submitted').all()
        for job in submitted_jobs:
            swarm.run(job)
            job.status = 'scheduled'
        for job in foo_jobs:
            '''
            ' TODO: Handle service not found error.
            '''
            service = swarm.docker.services.get(str(job.id))
            job.status = service.tasks()[0].get('Status').get('State')
            if job.status == 'complete' or job.status == 'failed':
                job.end_date = datetime.utcnow()
                service.remove()
        session.commit()
        self.Session.remove()