model-service/model_service/views/model_view.py

#!/usr/bin/env python3
import tempfile
import os
from flask import request, jsonify, current_app, abort, Response
from flask_classful import FlaskView, route
from model import db, Default, AIModel
from minio.error import BucketAlreadyExists, BucketAlreadyOwnedByYou, ResponseError, NoSuchKey
from schemas import AIModelSchema, DefaultSchema, InfoSchema
from marshmallow.exceptions import ValidationError
from utils import json_required, storage
from pyAudioAnalysis.audioTrainTest import load_model, load_model_knn


class ModelView(FlaskView):
    aimodel_schema = AIModelSchema(many=False)
    aimodels_schema = AIModelSchema(many=True,
                                    exclude=['timestamp', 'mid_window', 'mid_step', 'short_window', 'short_step',
                                             'compute_beat'])
    default_schema = DefaultSchema(many=False)
    info_schema = InfoSchema(many=False)

    def _ensure_buckets(self):
        for bucket_name in [current_app.config['MINIO_MEANS_BUCKET_NAME'],
                            current_app.config['MINIO_MODEL_BUCKET_NAME']]:
            try:
                storage.connection.make_bucket(bucket_name)

            except BucketAlreadyOwnedByYou as err:
                pass
            except BucketAlreadyExists as err:
                pass
                # Everything else should be raised

    def index(self):
        models = AIModel.query.all()
        return jsonify(self.aimodels_schema.dump(models)), 200

    def post(self):

        # get important data from the request
        try:
            info = self.info_schema.loads(request.form.get('info'))
        except ValidationError as e:
            abort(400, str(e))

        # check for conflict
        m = AIModel.query.filter_by(id=info['id']).first()
        if m:
            abort(409)

        # get and validate file
        model_file = request.files['modelFile']

        if model_file.content_length <= 0:
            abort(411, f"Content length for modelFile is not a positive integer or missing.")

        means_file = request.files['meansFile']

        if means_file.content_length <= 0:
            abort(411, f"Content length for meansFile is not a positive integer or missing.")

        # create bucket if necessary
        self._ensure_buckets()

        # Temporarily save the file, because pyAudioAnalysis can only read files
        _, temp_model_filename = tempfile.mkstemp()
        temp_means_filename = temp_model_filename + "MEANS"

        model_file.save(temp_model_filename)
        means_file.save(temp_means_filename)

        try:

            if info['type'] == 'knn':
                _, _, _, _, mid_window, mid_step, short_window, short_step, compute_beat \
                    = load_model_knn(temp_model_filename)
            else:
                _, _, _, _, mid_window, mid_step, short_window, short_step, compute_beat \
                    = load_model(temp_model_filename)

            # Because of pyAudiomeme the files already saved, so we just use the file uploader functions
            storage.connection.fput_object(
                current_app.config['MINIO_MODEL_BUCKET_NAME'],
                str(info['id']),
                temp_model_filename
            )

            storage.connection.fput_object(
                current_app.config['MINIO_MEANS_BUCKET_NAME'],
                str(info['id']),
                temp_means_filename
            )

        finally:
            os.remove(temp_model_filename)
            os.remove(temp_means_filename)

        m = AIModel(id=info['id'], mid_window=mid_window, mid_step=mid_step, short_window=short_window,
                    short_step=short_step, compute_beat=compute_beat, type=info['type'])

        db.session.add(m)
        db.session.commit()

        return jsonify(self.aimodel_schema.dump(m)), 200

    def get(self, _id: str):

        if _id == "$default":
            default = Default.query.first_or_404()  # TODO: Kitalálni, hogy inkább a latestestest-el térjen-e vissza
            m = default.aimodel
        else:
            m = AIModel.query.filter_by(id=_id).first_or_404()

        if "means" in request.args:
            bucket = current_app.config['MINIO_MEANS_BUCKET_NAME']
        else:
            bucket = current_app.config['MINIO_MODEL_BUCKET_NAME']

        try:
            data = storage.connection.get_object(bucket, str(m.id))
        except NoSuchKey:
            abort(500, "The ID is stored in the database but not int the Object Store")

        return Response(data.stream(), mimetype=data.headers['Content-type'])

    @route('<_id>/details')
    def get_details(self, _id: str):

        if _id == "$default":
            default = Default.query.first_or_404()  # TODO: Kitalálni, hogy inkább a latestestest-el térjen-e vissza
            m = default.aimodel
        else:
            m = AIModel.query.filter_by(id=_id).first_or_404()

        return jsonify(self.aimodel_schema.dump(m))

    def delete(self, _id: str):

        if _id == '$default':
            default = Default.query.first_or_404()
            m = default.aimodel
        else:
            m = AIModel.query.filter_by(id=_id).first_or_404()

        storage.connection.remove_object(current_app.config['MINIO_MODEL_BUCKET_NAME'], str(m.id))
        storage.connection.remove_object(current_app.config['MINIO_MEANS_BUCKET_NAME'], str(m.id))

        db.session.delete(m)
        db.session.commit()

        return '', 204

    @json_required
    @route('$default', methods=['PUT'])
    def put_default(self):

        try:
            req = self.default_schema.load(request.json)
        except ValidationError as e:
            abort(404, str(e))

        m = AIModel.query.filter_by(id=req['id']).first_or_404()

        Default.query.delete()
        new_default = Default(aimodel=m)
        db.session.add(new_default)
        db.session.commit()

        return '', 204