PyCS/pycs/frontend/WebServer.py

import glob
import os
import threading

from logging.config import dictConfig

from flask import send_from_directory
from flask_socketio import SocketIO
from logging import config

from pycs import app

from pycs.database.Collection import Collection
from pycs.database.LabelProvider import LabelProvider
from pycs.database.Label import Label
from pycs.database.Model import Model
from pycs.database.Project import Project
from pycs.database.Result import Result

from pycs.frontend.endpoints.ListJobs import ListJobs
from pycs.frontend.endpoints.additional.FolderInformation import FolderInformation
from pycs.frontend.endpoints.base import ListView
from pycs.frontend.endpoints.data.GetFile import GetFile
from pycs.frontend.endpoints.data.GetPreviousAndNextFile import GetPreviousAndNextFile
from pycs.frontend.endpoints.data.GetResizedFile import GetResizedFile
from pycs.frontend.endpoints.data.RemoveFile import RemoveFile
from pycs.frontend.endpoints.data.UploadFile import UploadFile
from pycs.frontend.endpoints.jobs.RemoveJob import RemoveJob
from pycs.frontend.endpoints.labels.CreateLabel import CreateLabel
from pycs.frontend.endpoints.labels.EditLabelName import EditLabelName
from pycs.frontend.endpoints.labels.EditLabelParent import EditLabelParent
from pycs.frontend.endpoints.labels.RemoveLabel import RemoveLabel
from pycs.frontend.endpoints.pipelines.FitModel import FitModel
from pycs.frontend.endpoints.pipelines.PredictFile import PredictFile
from pycs.frontend.endpoints.pipelines.PredictModel import PredictModel
from pycs.frontend.endpoints.projects.CreateProject import CreateProject
from pycs.frontend.endpoints.projects.EditProjectDescription import EditProjectDescription
from pycs.frontend.endpoints.projects.EditProjectName import EditProjectName
from pycs.frontend.endpoints.projects.ExecuteExternalStorage import ExecuteExternalStorage
from pycs.frontend.endpoints.projects.ExecuteLabelProvider import ExecuteLabelProvider
from pycs.frontend.endpoints.projects.GetProjectModel import GetProjectModel
from pycs.frontend.endpoints.projects.ListFiles import ListFiles
from pycs.frontend.endpoints.projects.RemoveProject import RemoveProject
from pycs.frontend.endpoints.results.ConfirmResult import ConfirmResult
from pycs.frontend.endpoints.results.CreateResult import CreateResult
from pycs.frontend.endpoints.results.EditResultData import EditResultData
from pycs.frontend.endpoints.results.EditResultLabel import EditResultLabel
from pycs.frontend.endpoints.results.GetProjectResults import GetProjectResults
from pycs.frontend.endpoints.results.RemoveResult import RemoveResult
from pycs.frontend.endpoints.results.ResetResults import ResetResults
from pycs.frontend.notifications.NotificationManager import NotificationManager
from pycs.frontend.util.JSONEncoder import JSONEncoder
from pycs.jobs.JobRunner import JobRunner
from pycs.util.PipelineCache import PipelineCache


class WebServer:
    """
    wrapper class for flask and socket.io which initializes most networking
    """

    # pylint: disable=line-too-long
    def __init__(self, app, settings: dict):

        dictConfig(settings["logging"])

        # initialize flask app instance
        self.app = app

        # run discovery modules
        Model.discover("models/")
        LabelProvider.discover("labels/")

        # start job runner
        self.logger.info('Starting job runner... ')
        self.jobs = JobRunner()

        # create pipeline cache
        self.logger.info('Creating pipeline cache')
        self.pipelines = PipelineCache(self.jobs)

        PRODUCTION = os.path.exists('webui/index.html')

        init_func = self.production_init if PRODUCTION else self.development_init
        kwargs, static_files = init_func()

        self.sio = SocketIO(self.app, **kwargs)#socketio.Server(**kwargs)
        # self.__app = socketio.WSGIApp(self.sio, self.app, static_files=static_files)
        self.host, self.port = settings['host'], settings['port']

        # set json encoder so database objects are serialized correctly
        self.app.json_encoder = JSONEncoder

        self.start_runner()
        self.init_notifications()
        self.define_routes()
        self.logger.info("Server initialized")



    def start_runner(self):
        app.logger.info(f"Main Thread ID: {threading.get_ident()}")
        self.jobs.start()
        self.pipelines.start()

    def stop_runner(self):
        self.jobs.stop()
        self.pipelines.stop()

    def wait_for_runner(self):
        self.jobs.wait_for_empty_queue()
        self.pipelines.wait_for_empty_queue()


    @property
    def logger(self):
        return self.app.logger

    def init_notifications(self):
        # create notification manager
        self.nm = NotificationManager(self.sio)

        self.jobs.on_create(self.nm.create_job)
        self.jobs.on_start(self.nm.edit_job)
        self.jobs.on_progress(self.nm.edit_job)
        self.jobs.on_finish(self.nm.edit_job)
        self.jobs.on_remove(self.nm.remove_job)

    def development_init(self):

        self.logger.info('Initializing development build')

        # set access control header to allow requests from Vue.js development server
        @self.app.after_request
        def after_request(response):
            # pylint: disable=unused-variable
            response.headers['Access-Control-Allow-Origin'] = '*'
            return response

        return dict(cors_allowed_origins='*', async_mode='eventlet'), None

    def production_init(self):

        self.logger.info('Initializing production build')

        kwargs = dict(async_mode='eventlet')
        if len(settings['allowedOrigins']) > 0:
            origins = settings['allowedOrigins']
            kwargs["cors_allowed_origins"] = origins

        # overwrite root path to serve index.html
        @self.app.route('/', methods=['GET'])
        def index():
            # pylint: disable=unused-variable
            return send_from_directory(os.path.join(os.getcwd(), 'webui'), 'index.html')

        return kwargs, self.static_files

    @property
    def static_files(self) -> dict:
        # find static files and folders
        static_files = {}

        for file_path in glob.glob('webui/*'):
            file_path = file_path.replace('\\', '/')
            static_files[file_path[5:]] = file_path

        # separately add svg files and set their correct mime type
        for svg_path in glob.glob('webui/img/*.svg'):
            svg_path = svg_path.replace('\\', '/')
            static_files[svg_path[5:]] = {'content_type': 'image/svg+xml', 'filename': svg_path}

        return static_files

    def define_routes(self):
        # additional
        self.app.add_url_rule(
            '/folder',
            view_func=FolderInformation.as_view('folder_information')
        )

        # jobs
        self.app.add_url_rule(
            '/jobs',
            view_func=ListJobs.as_view('list_jobs', self.jobs)
        )
        self.app.add_url_rule(
            '/jobs/<identifier>/remove',
            view_func=RemoveJob.as_view('remove_job', self.jobs)
        )

        # models
        self.app.add_url_rule(
            '/models',
            view_func=ListView.as_view('list_models', model_cls=Model)
        )
        self.app.add_url_rule(
            '/projects/<int:identifier>/model',
            view_func=GetProjectModel.as_view('get_project_model')
        )

        # labels
        self.app.add_url_rule(
            '/label_providers',
            view_func=ListView.as_view('label_providers', model_cls=LabelProvider)
        )
        self.app.add_url_rule(
            '/projects/<int:project_id>/labels',
            view_func=ListView.as_view('list_labels', model_cls=Label)
        )
        self.app.add_url_rule(
            '/projects/<int:identifier>/labels',
            view_func=CreateLabel.as_view('create_label', self.nm)
        )
        self.app.add_url_rule(
            '/projects/<int:project_id>/labels/<int:label_id>/remove',
            view_func=RemoveLabel.as_view('remove_label', self.nm)
        )
        self.app.add_url_rule(
            '/projects/<int:project_id>/labels/<int:label_id>/name',
            view_func=EditLabelName.as_view('edit_label_name', self.nm)
        )
        self.app.add_url_rule(
            '/projects/<int:project_id>/labels/<int:label_id>/parent',
            view_func=EditLabelParent.as_view('edit_label_parent', self.nm)
        )

        # collections
        self.app.add_url_rule(
            '/projects/<int:project_id>/collections',
            view_func=ListView.as_view('list_collections',
                model_cls=Collection, post_process=Collection.update_autoselect)
        )
        self.app.add_url_rule(
            '/projects/<int:project_id>/data/<int:collection_id>/<int:start>/<int:length>',
            view_func=ListFiles.as_view('list_collection_files')
        )

        # data
        self.app.add_url_rule(
            '/projects/<int:identifier>/data',
            view_func=UploadFile.as_view('upload_file', self.nm)
        )
        self.app.add_url_rule(
            '/projects/<int:project_id>/data/<int:start>/<int:length>',
            view_func=ListFiles.as_view('list_files')
        )
        self.app.add_url_rule(
            '/data/<int:identifier>/remove',
            view_func=RemoveFile.as_view('remove_file', self.nm)
        )
        self.app.add_url_rule(
            '/data/<int:file_id>',
            view_func=GetFile.as_view('get_file')
        )
        self.app.add_url_rule(
            '/data/<int:file_id>/<resolution>',
            view_func=GetResizedFile.as_view('get_resized_file')
        )
        self.app.add_url_rule(
            '/data/<int:file_id>/previous_next',
            view_func=GetPreviousAndNextFile.as_view('get_previous_and_next_file')
        )

        # results
        self.app.add_url_rule(
            '/projects/<int:project_id>/results',
            view_func=GetProjectResults.as_view('get_project_results')
        )
        self.app.add_url_rule(
            '/data/<int:file_id>/results',
            view_func=ListView.as_view('get_results', model_cls=Result)
        )
        self.app.add_url_rule(
            '/data/<int:file_id>/results',
            view_func=CreateResult.as_view('create_result', self.nm)
        )
        self.app.add_url_rule(
            '/data/<int:file_id>/reset',
            view_func=ResetResults.as_view('reset_results', self.nm)
        )
        self.app.add_url_rule(
            '/results/<int:result_id>/remove',
            view_func=RemoveResult.as_view('remove_result', self.nm)
        )
        self.app.add_url_rule(
            '/results/<int:result_id>/confirm',
            view_func=ConfirmResult.as_view('confirm_result', self.nm)
        )
        self.app.add_url_rule(
            '/results/<int:result_id>/label',
            view_func=EditResultLabel.as_view('edit_result_label', self.nm)
        )
        self.app.add_url_rule(
            '/results/<int:result_id>/data',
            view_func=EditResultData.as_view('edit_result_data', self.nm)
        )

        # projects
        self.app.add_url_rule(
            '/projects',
            view_func=ListView.as_view('list_projects', model_cls=Project)
        )
        self.app.add_url_rule(
            '/projects',
            view_func=CreateProject.as_view('create_project', self.nm, self.jobs)
        )
        self.app.add_url_rule(
            '/projects/<int:identifier>/label_provider',
            view_func=ExecuteLabelProvider.as_view('execute_label_provider',
                                                   self.nm, self.jobs)
        )
        self.app.add_url_rule(
            '/projects/<int:identifier>/external_storage',
            view_func=ExecuteExternalStorage.as_view('execute_external_storage',
                                                     self.nm, self.jobs)
        )
        self.app.add_url_rule(
            '/projects/<int:identifier>/remove',
            view_func=RemoveProject.as_view('remove_project', self.nm)
        )
        self.app.add_url_rule(
            '/projects/<int:identifier>/name',
            view_func=EditProjectName.as_view('edit_project_name', self.nm)
        )
        self.app.add_url_rule(
            '/projects/<int:identifier>/description',
            view_func=EditProjectDescription.as_view('edit_project_description', self.nm)
        )

        # pipelines
        self.app.add_url_rule(
            '/projects/<int:project_id>/pipelines/fit',
            view_func=FitModel.as_view('fit_model', self.jobs, self.pipelines)
        )
        self.app.add_url_rule(
            '/projects/<int:project_id>/pipelines/predict',
            view_func=PredictModel.as_view('predict_model', self.nm,
                                           self.jobs, self.pipelines)
        )
        self.app.add_url_rule(
            '/data/<int:file_id>/predict',
            view_func=PredictFile.as_view('predict_file', self.nm,
                                          self.jobs, self.pipelines)
        )

    def run(self):
        self.logger.info("Starting server...")
        return self.sio.run(self.app, host=self.host, port=self.port)
        # eventlet.wsgi.server(eventlet.listen((self.__host, self.__port)), self.__app)