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
import opentracing


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
        with opentracing.tracer.start_active_span('parseAndValidate'):
            try:
                info = self.info_schema.loads(request.form.get('info'))
            except ValidationError as e:
                return abort(400, str(e))

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

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

            if model_file.content_length <= 0:
                return 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:
                return abort(411, f"Content length for weightsFile is not a positive integer or missing.")

        # create bucket if necessary
        with opentracing.tracer.start_active_span('ensureBuckets'):
            ensure_buckets()

        # Create the entry in the db
        with opentracing.tracer.start_active_span('sqlalchemy.create'):
            m = AIModel(id=info['id'], type=AIModelType.cnn, target_class_name=info['target_class_name'])
            db.session.add(m)

        # Put files into MinIO
        with opentracing.tracer.start_active_span('putObjectsInMinio'):
            with opentracing.tracer.start_active_span(
                    'minio.putObject',
                    tags={
                        "bucket": current_app.config['MINIO_CNN_BUCKET_NAME'],
                        "object_name": self.MODEL_DIRECTORY + str(m.id),
                        "length": model_file.content_length,
                        "component": "model"
                    }
            ):
                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
                )
            with opentracing.tracer.start_active_span(
                    'minio.putObject',
                    tags={
                        "bucket": current_app.config['MINIO_CNN_BUCKET_NAME'],
                        "object_name": self.WEIGHTS_DIRECTORY + str(m.id),
                        "length": model_file.content_type,
                        "component": "weights"
                    }
            ):
                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
                )

        with opentracing.tracer.start_active_span('sqlalchemy.commit'):
            db.session.commit()

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

    def delete(self, id_: str):

        with opentracing.tracer.start_active_span(
                'sqlalchemy.select',
                tags={"aimodel_type": AIModelType.cnn, "id": id_}
        ):
            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()

        with opentracing.tracer.start_active_span('removeFromMinio'):
            with opentracing.tracer.start_active_span(
                    'minio.removeObject',
                    tags={
                        "bucket": current_app.config['MINIO_CNN_BUCKET_NAME'],
                        "object_name": self.WEIGHTS_DIRECTORY + str(m.id),
                        "component": "weights"
                    }
            ):
                storage.connection.remove_object(current_app.config['MINIO_CNN_BUCKET_NAME'],
                                                 self.WEIGHTS_DIRECTORY + str(m.id))

            with opentracing.tracer.start_active_span(
                    'minio.removeObject',
                    tags={
                        "bucket": current_app.config['MINIO_CNN_BUCKET_NAME'],
                        "object_name": self.MODEL_DIRECTORY + str(m.id),
                        "component": "model"
                    }
            ):
                storage.connection.remove_object(current_app.config['MINIO_CNN_BUCKET_NAME'],
                                                 self.MODEL_DIRECTORY + str(m.id))

        with opentracing.tracer.start_active_span(
                'sqlalchemy.delete',
                tags={"aimodel_type": AIModelType.cnn, "id": id_}
        ):
            db.session.delete(m)

        with opentracing.tracer.start_active_span('sqlalchemy.commit'):
            db.session.commit()

        return '', 204

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

        with opentracing.tracer.start_active_span(
                'sqlalchemy.select',
                tags={"aimodel_type": AIModelType.cnn, "id": id_}
        ):
            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)
            component = "weights"
        else:
            path = self.MODEL_DIRECTORY + str(m.id)
            component = "model"

        with opentracing.tracer.start_active_span(
                'minio.getObject',
                tags={
                    "bucket": current_app.config['MINIO_CNN_BUCKET_NAME'],
                    "object_name": path,
                    "component": component
                }
        ):
            # Note: this only initiates the download, the file download itself is streamed
            try:
                data = storage.connection.get_object(current_app.config['MINIO_CNN_BUCKET_NAME'], path)
            except NoSuchKey:
                return abort(500, "The ID is stored in the database but not int the Object Store")

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