ml/registry.py

#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership. The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License. You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
#

"""
ML registry

Registry object that will keep information about available algorithms and corresponding endpoints.
"""

from ml.classifiers import Classifier
from typing import Dict
from api.models import Algorithm, Dataset


def has_empty_values(data: dict):

    for key, value in data.items():
        if value == None:
            raise ValueError(f'{key} cannot be null')

    return False


class MLRegistry:
    def __init__(self):
        self.classifiers: Dict[int, Classifier] = {}

    def add_algorithms(self,
                       attrs=[{
                           "classifier": None,
                           "description": None,
                           "status": None,
                           "version": None,
                           "dataset": None,
                           "region": None,
                           "created_by": None
                       }]):

        for attr in attrs:

            if not has_empty_values(attr):
                #get dataset
                dataset, _ = Dataset.objects.get_or_create(
                    name=attr['dataset'], region=attr['region'])

                # get algorithm
                algorithm, _ = Algorithm.objects.get_or_create(
                    classifier=attr['classifier'].__class__.__name__,
                    description=attr['description'],
                    version=attr['version'],
                    status=attr['status'],
                    dataset=dataset,
                    created_by=attr['created_by'])
                self.classifiers[algorithm.id] = attr['classifier']

        return self.classifiers