get_model Takes more too long to initialise so the endpoint times out about sagify HOT 6 CLOSED

from __future__ import absolute_import

import os
# Do not remove the following line
import sys;sys.path.append("..")  # NOQA

import logging
import concurrent.futures


_MODEL_PATH = os.path.join('/opt/ml/', 'model')  # Path where all your model(s) live in

log = logging.getLogger(__name__)

class ModelNotYetLoadedException(Exception):
    def __init__(self, message):
        super().__init__(message)

class ModelLoader:
    def __init__(self, load_method):
        log.info("setting up ModelLoader")
        # Invoke load method asynchronously so process does not block.
        self.model = None
        self.executor = concurrent.futures.ThreadPoolExecutor(max_workers=1)
        self.future = self.executor.submit(load_method)
        self.future.add_done_callback(self.done)

    def get_model(self):
        if self.future.done():
            return self.model

        log.error("get_model called before ready")
        raise ModelNotYetLoadedException("model not loaded")

    def done(self, future):
        """Callback method invoked when model load complete. Sets us to ready status."""
        log.info("model load done")
        self.model = future.result()
        log.info("shutting down executor")
        self.executor.shutdown(wait=False)

    def get_ready(self):
        return self.future.done()


class ModelService(object):
    model = None

    @staticmethod
    def load_model():
        # Load your model, your model that takes time, here
        from sklearn.externals import joblib

        return joblib.load(os.path.join(_MODEL_PATH, 'model.pkl'))

    @classmethod
    def init_model(cls):
        """Get the model object for this instance, loading it if it's not already loaded."""
        if cls.model is None:
            cls.model = ModelLoader(cls.load_model)
        return cls.model

    @classmethod
    def predict(cls, input):
        """For the input, do the predictions and return them."""
        try:
            return cls.model.get_model().predict(input)

        except ModelNotYetLoadedException as e:
            return 'model not yet loaded'

ModelService.init_model()


def predict(json_input):
    """
    Prediction given the request input
    :param json_input: [dict], request input
    :return: [dict], prediction
    """

    # TODO Transform json_input and assign the transformed value to model_input
    try:
        model_input = json_input['features']
        prediction = ModelService.predict(model_input)

        result = {'prediction': prediction.item()}

        return result
    except Exception as e:
        return {"error": str(e)}