model-service/model_service/views/cnn_view.py

#!/usr/bin/env python3
from flask import request, jsonify, current_app, abort, Response
from flask_classful import FlaskView, route
from model import db, Default, AIModel, AIModelType
from minio.error import NoSuchKey
from schemas import AIModelSchema, DefaultSchema, InfoSchema
from marshmallow.exceptions import ValidationError
from utils import multipart_required, storage, ensure_buckets


class CNNView(FlaskView):
    route_base = 'cnn'

    aimodel_schema = AIModelSchema(many=False)
    info_schema = InfoSchema(many=False)

    MODEL_DIRECTORY = "model/"
    WEIGHTS_DIRECTORY = "weights/"

    @multipart_required
    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.")

        weights_file = request.files['weightsFile']

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

        # create bucket if necessary
        ensure_buckets()

        # Create the entry in the db
        m = AIModel(id=info['id'], type=AIModelType.cnn, target_class_name=info['target_class_name'])

        # Put files into MinIO
        storage.connection.put_object(
            current_app.config['MINIO_CNN_BUCKET_NAME'],
            self.MODEL_DIRECTORY + str(m.id),
            model_file,
            model_file.content_length,
            content_type=model_file.content_type
        )

        storage.connection.put_object(
            current_app.config['MINIO_CNN_BUCKET_NAME'],
            self.WEIGHTS_DIRECTORY + str(m.id),
            weights_file,
            weights_file.content_length,
            content_type=weights_file.content_type
        )

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

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

    def delete(self, id_: str):

        if id_ == "$default":
            default = Default.query.filter_by(type=AIModelType.cnn).first_or_404()
            m = default.aimodel
        else:
            m = AIModel.query.filter_by(type=AIModelType.cnn, id=id_).first_or_404()

        storage.connection.remove_object(current_app.config['MINIO_CNN_BUCKET_NAME'], self.WEIGHTS_DIRECTORY + str(m.id))
        storage.connection.remove_object(current_app.config['MINIO_CNN_BUCKET_NAME'], self.MODEL_DIRECTORY + str(m.id))

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

        return '', 204

    @route('<id_>/file')
    def get_file(self, id_: str):

        if id_ == "$default":
            default = Default.query.filter_by(type=AIModelType.cnn).first_or_404()
            m = default.aimodel
        else:
            m = AIModel.query.filter_by(type=AIModelType.cnn, id=id_).first_or_404()

        if "weights" in request.args:
            path = self.WEIGHTS_DIRECTORY + str(m.id)
        else:
            path = self.MODEL_DIRECTORY + str(m.id)

        try:
            data = storage.connection.get_object(current_app.config['MINIO_CNN_BUCKET_NAME'], path)
        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'])