Added ability to store CNN models

model_service/views/cnn_view.py

@@ -0,0 +1,128 @@
+#!/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 json_required, storage, ensure_buckets
+
+
+class CNNView(FlaskView):
+    aimodel_schema = AIModelSchema(many=False)
+    aimodels_schema = AIModelSchema(many=True, exclude=['timestamp', 'details'])
+    default_schema = DefaultSchema(many=False)
+    info_schema = InfoSchema(many=False)
+
+    def index(self):
+        models = AIModel.query.filter_by(type=AIModelType.CNN).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.")
+
+        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()
+
+        # Put files into MinIO
+        storage.connection.put_object(current_app.config['MINIO_CNN_BUCKET_NAME'], "model/" + 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'], "weights/" + str(m.id), weights_file,
+                                      weights_file.content_length, content_type=weights_file.content_type)
+
+        m = AIModel(id=info['id'], type=AIModelType.CNN)
+
+        db.session.add(m)
+        db.session.commit()
+
+        return jsonify(self.aimodel_schema.dump(m)), 200
+
+    def get(self, _id: str):
+
+        if _id == "$default":
+            # TODO: Kitalálni, hogy inkább a latestestest-el térjen-e vissza
+            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 = "weights/" + str(m.id)
+        else:
+            path = "model/" + 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'])
+
+    @route('<_id>/details')
+    def get_details(self, _id: str):
+
+        if _id == "$default":
+            # TODO: Kitalálni, hogy inkább a latestestest-el térjen-e vissza
+            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()
+
+        return jsonify(self.aimodel_schema.dump(m))
+
+    def delete(self, _id: str):
+
+        if _id == "$default":
+            # TODO: Kitalálni, hogy inkább a latestestest-el térjen-e vissza
+            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'], "weights/" + str(m.id))
+        storage.connection.remove_object(current_app.config['MINIO_CNN_BUCKET_NAME'], "model/" + 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(400, str(e))
+
+        m = AIModel.query.filter_by(type=AIModelType.CNN, id=req['id']).first_or_404()
+
+        Default.query.filter_by(type=AIModelType.CNN).delete()
+        new_default = Default(type=AIModelType.CNN, aimodel=m)
+        db.session.add(new_default)
+        db.session.commit()
+
+        return '', 204