PyCS/pycs/jobs/JobRunner.py

from concurrent.futures import ThreadPoolExecutor
from time import time
from types import GeneratorType
from typing import Callable, List, Generator, Optional, Any

import eventlet
# from eventlet import spawn, spawn_n, tpool
from eventlet.event import Event


from pycs.database.Project import Project
from pycs.jobs.Job import Job
from pycs.jobs.JobGroupBusyException import JobGroupBusyException

from pycs.util.green_worker import GreenWorker

class JobRunner(GreenWorker):
    """
    run jobs in a thread pool, but track progress and process results in eventlet queue
    """

    # pylint: disable=too-many-arguments
    def __init__(self):
        super().__init__()
        self.__jobs = []
        self.__groups = {}

        self.__executor = ThreadPoolExecutor(1)

        self.__create_listeners = []
        self.__start_listeners = []
        self.__progress_listeners = []
        self.__finish_listeners = []
        self.__remove_listeners = []

    def list(self) -> List[Job]:
        """
        get a list of all jobs including finished ones

        :return: list of job objects
        """
        return self.__jobs

    def on_create(self, callback: Callable[[Job], None]) -> None:
        """
        register a callback that is executed each time a job is created

        :param callback: callback function
        :return:
        """
        self.__create_listeners.append(callback)

    def on_start(self, callback: Callable[[Job], None]) -> None:
        """
        register a callback that is executed each time a job is started

        :param callback: callback function
        :return:
        """
        self.__start_listeners.append(callback)

    def on_progress(self, callback: Callable[[Job], None]) -> None:
        """
        register a callback that is executed each time a job changes it's progress

        :param callback: callback function
        :return:
        """
        self.__progress_listeners.append(callback)

    def on_finish(self, callback: Callable[[Job], None]) -> None:
        """
        register a callback that is executed each time a job is finished

        :param callback: callback function
        :return:
        """
        self.__finish_listeners.append(callback)

    def on_remove(self, callback: Callable[[Job], None]) -> None:
        """
        register a callback that is executed each time a job is removed

        :param callback: callback function
        :return:
        """
        self.__remove_listeners.append(callback)

    def remove(self, id):
        """
        remove a job using its unique id

        :param id: job id
        :return:
        """
        for i in range(len(self.__jobs)):
            if self.__jobs[i].id == id:
                if self.__jobs[i].finished is not None:
                    job = self.__jobs[i]
                    del self.__jobs[i]

                    for callback in self.__remove_listeners:
                        callback(job)

                return

    def run(self,
            project: Project,
            job_type: str,
            name: str,
            group: str,
            executable: Callable[[Any], Optional[Generator[float, None, None]]],
            *args,
            progress: Callable[[Any], float] = None,
            result: Callable[[Any], None] = None,
            result_event: Event = None,
            **kwargs) -> Job:
        """
        add a job to run it in a thread pool

        :param project: project the job is associated with
        :param job_type: job type
        :param name: job name
        :param group: job group (raises JobGroupBusyException if there is already a job running
                      with the same group id)
        :param executable: function to execute
        :param args: arguments for executable
        :param progress: is called everytime executable yields a value
        :param result: is called as soon as executable returns a value
        :param result_event: eventlet event to be called as soon as executable returns a value
        :param kwargs: named arguments for executable
        :return: job object
        """
        # create job object
        job = Job(project, job_type, name)

        # abort if group is busy
        # otherwise add to groups dict
        if group is not None:
            if group in self.__groups:
                raise JobGroupBusyException

            self.__groups[group] = job

        # add to job list
        self.__jobs.append(job)

        # execute create listeners
        for callback in self.__create_listeners:
            callback(job)

        # add to execution queue
        # self.__queue.put((group, executable, job, progress, result, result_event, args, kwargs))
        self.queue.put((group, executable, job, progress, result, result_event, args, kwargs))

        # return job object
        return job


    def _job_started(self, job):
        job.start()
        for callback in self.__start_listeners:
            callback(job)

    def _job_progress(self, job, progress):
        job.update(progress=progress)
        for callback in self.__progress_listeners:
            callback(job)

    def _job_finished(self, job):
        job.finish()
        for callback in self.__finish_listeners:
            callback(job)


    def process_iterator(self, generator, job, progress_fun):
        try:
            iterator = iter(generator)

            while True:
                # run until next progress event
                future = self.__executor.submit(next, iterator)
                progress = eventlet.tpool.execute(future.result)

                # execute progress function
                if progress_fun is not None:
                    if isinstance(progress, tuple):
                        progress = progress_fun(*progress)
                    else:
                        progress = progress_fun(progress)

                self._job_progress(job, progress=progress)

        except StopIteration as stop_iteration_exception:
            return stop_iteration_exception.value


    def start_work(self, group, executable, job, progress_fun, result_fun, result_event, args, kwargs):
        """ started as coroutine. Starts other tasks in separate green threads. """

        # execute start listeners
        self._job_started(job)
        try:
            # result = generator = executable(*args, **kwargs)
            future = self.__executor.submit(executable, *args, **kwargs)
            result = generator = eventlet.tpool.execute(future.result)

            if isinstance(generator, GeneratorType):
                result = self.process_iterator(generator, job, progress_fun)

            self._job_progress(job, progress=1)

            # execute result function
            if result_fun is not None:
                if isinstance(result, tuple):
                    result_fun(*result)
                else:
                    result_fun(result)

            # execute event
            if result_event is not None:
                result_event.send(result)

        # save exceptions to show in ui
        except Exception as e:
            import traceback
            traceback.print_exc()
            job.exception = f'{type(e).__name__} ({str(e)})'

        # remove from group dict
        if group is not None:
            del self.__groups[group]

        self._job_finished(job)

    def __run(self):

        while True:

            # get execution function and job from queue
            group, executable, job, progress_fun, result_fun, result_event, args, kwargs \
                = self.__queue.get(block=True)

            # execute start listeners
            job.started = int(time())
            job.updated = int(time())

            for callback in self.__start_listeners:
                callback(job)

            # run function and track progress
            try:
                # result = generator = executable(*args, **kwargs)
                future = self.__executor.submit(executable, *args, **kwargs)
                result = generator = tpool.execute(future.result)

                if isinstance(generator, GeneratorType):
                    iterator = iter(generator)

                    try:
                        while True:
                            # run until next progress event
                            future = self.__executor.submit(next, iterator)
                            progress = tpool.execute(future.result)
                            # progress = next(iterator)


                            # execute progress function
                            if progress_fun is not None:
                                if isinstance(progress, tuple):
                                    progress = progress_fun(*progress)
                                else:
                                    progress = progress_fun(progress)

                            # execute progress listeners
                            job.progress = progress
                            job.updated = int(time())

                            for callback in self.__progress_listeners:
                                callback(job)
                    except StopIteration as stop_iteration_exception:
                        result = stop_iteration_exception.value

                # update progress
                job.progress = 1
                job.updated = int(time())

                for callback in self.__progress_listeners:
                    callback(job)

                # execute result function
                if result_fun is not None:
                    if isinstance(result, tuple):
                        result_fun(*result)
                    else:
                        result_fun(result)

                # execute event
                if result_event is not None:
                    result_event.send(result)

            # save exceptions to show in ui
            except Exception as exception:
                import traceback
                traceback.print_exc()
                job.exception = f'{type(exception).__name__} ({str(exception)})'

            # remove from group dict
            if group is not None:
                del self.__groups[group]

            # finish job
            job.finished = int(time())
            job.updated = int(time())

            for callback in self.__finish_listeners:
                callback(job)