async def save_model_to_disk(training_id: str = '',
                             model: UploadFile = File(...)):
    """
    Receives a .zip file containing a tensorflow2 SavedModel object sent by a dataset microservice.
    Then, will store the .zip file in trained model docker volume, with the naming format of
    <training_id>.zip

    :param training_id: Training ID the model is associated with
    :param model: .zip file containing SavedModel object
    :return: {'status': 'success'} if saving is successful, else {'status': 'failure'}
    """
    logger.debug('Training ID: ' + training_id)

    if not get_training_result_by_training_id(training_id):
        return {
            'status': 'failure',
            'detail': 'Unable to find training result with specified ID',
            'training_id': training_id
        }

    upload_folder = open(os.path.join('/app/training_results', model.filename),
                         'wb+')
    shutil.copyfileobj(model.file, upload_folder)
    upload_folder.close()
    return {
        'status': 'success',
        'detail': 'Training results uploaded successfully',
        'training_id': training_id
    }
def ping_dataset(dataset_name):
    """
    Periodically ping a dataset's service to make sure that it is active and able to receive requests.
    If it's not, remove the dataset from the available_datasets map. This is a helper method that is
    not directly exposed via HTTP.

    :param dataset_name: Name of a registered dataset as a string
    """
    dataset_is_alive = True

    def kill_dataset():
        settings.available_datasets.pop(dataset_name)
        nonlocal dataset_is_alive
        dataset_is_alive = False
        logger.debug("Dataset " + dataset_name +
                     " is not responsive. Removing from available services...")

    while dataset_is_alive and not dependency.shutdown:
        try:
            r = requests.get('http://host.docker.internal:' +
                             str(settings.available_datasets[dataset_name]) +
                             '/status')
            r.raise_for_status()
            for increment in range(dependency.WAIT_TIME):
                if not dependency.shutdown:  # Check between increments to stop hanging on shutdown
                    time.sleep(1)
        except (requests.exceptions.ConnectionError,
                requests.exceptions.Timeout, requests.exceptions.HTTPError):
            kill_dataset()
            return

    if dependency.shutdown:
        logger.debug("Dataset [" + dataset_name +
                     "] Healthcheck Thread Terminated.")
 def kill_model():
     settings.available_models.pop(model_name)
     nonlocal model_is_alive
     model_is_alive = False
     logger.debug(
         "Model " + model_name +
         " is not responsive. Removing the model from available services..."
     )
def register_dataset(dataset: MicroserviceConnection):
    """
    Register a single dataset to the server by adding the name and port
    to available dataset settings. Also kick start a separate thread to keep track
    of the dataset service status. A valid dataset API key must be in the request
    header for this method to run.

    :param dataset: MicroserviceConnection object with name and port of dataset
    :return: {'status': 'success'} if saving is successful, else {'status': 'failure'}
    """
    # Do not accept calls if server is in process of shutting down
    if dependency.shutdown:
        return JSONResponse(
            status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
            content={
                'status':
                'failure',
                'detail':
                'Server is shutting down. Unable to complete new dataset registration.'
            })

    # Do not add duplicates of running datasets to server
    if dataset.name in settings.available_datasets:
        return {
            "status": "success",
            'dataset': dataset.name,
            'detail': 'Dataset has already been registered.'
        }

    # Ensure that we can connect back to dataset before adding it
    try:
        r = requests.get('http://host.docker.internal:' + str(dataset.port) +
                         '/status')
        r.raise_for_status()
    except (requests.exceptions.ConnectionError, requests.exceptions.Timeout,
            requests.exceptions.HTTPError):
        return {
            "status": "failure",
            'dataset': dataset.name,
            'detail': 'Unable to establish successful connection to dataset.'
        }

    # Register dataset to server and create thread to ensure dataset is responsive
    settings.available_datasets[dataset.name] = dataset.port
    pool.submit(ping_dataset, dataset.name)

    logger.debug("Dataset " + dataset.name +
                 " successfully registered to server.")

    return {
        "status": "success",
        'dataset': dataset.name,
        'detail': 'Dataset has been successfully registered to server.'
    }
def register_model(model: MicroserviceConnection):
    """
    Register a single model to the server by adding the model's name and socket
    to available model settings. Also kick start a separate thread to keep track
    of the model service status. Models that are registered must use a valid API key.

    :param model: MicroserviceConnection object with the model name and model socket.
    :return: {'status': 'success'} if registration successful else {'status': 'failure'}
    """

    # Do not accept calls if server is in process of shutting down
    if dependency.shutdown:
        return JSONResponse(
            status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
            content={
                'status':
                'failure',
                'detail':
                'Server is shutting down. Unable to complete new model registration.'
            })

    # Do not add duplicates of running models to server
    if model.name in settings.available_models:
        return {
            "status": "success",
            'model': model.name,
            'detail': 'Model has already been registered.'
        }

    # Ensure that we can connect back to model before adding it
    try:
        r = requests.get(model.socket + '/status')
        r.raise_for_status()
    except (requests.exceptions.ConnectionError, requests.exceptions.Timeout,
            requests.exceptions.HTTPError):
        return {
            "status": "failure",
            'model': model.name,
            'detail': 'Unable to establish successful connection to model.'
        }

    # Register model to server and create thread to ensure model is responsive
    settings.available_models[model.name] = model.socket
    pool.submit(ping_model, model.name)

    logger.debug("Model " + model.name + " successfully registered to server.")

    return {
        "status": "success",
        'model': model.name,
        'detail': 'Model has been successfully registered to server.'
    }
def current_user_investigator(token: str = Depends(oauth2_scheme)):
    """
    Permission Checking Function to be used as a Dependency for API endpoints. This is used as a helper.
    This will either return a User object to the calling method if the user meets the authentication requirements,
    or it will raise a CredentialException and prevent the method that depends on this from continuing.

    :param token: User authentication token
    :return: User object if user has correct role, else raise dependency.CredentialException
    """
    user = get_current_user(token)
    if not any(role in [Roles.admin.name, Roles.investigator.name]
               for role in user.roles):
        logger.debug('User Roles')
        logger.debug(user.roles)

        raise CredentialException()

    return user
def ping_model(model_name):
    """
    Periodically ping a model's service to make sure that it is active. If it's not, remove the model from the
    available_models BaseSetting in dependency.py

    :param model_name: Name of model to ping. This is the name the model registered to the server with.
    """

    model_is_alive = True

    def kill_model():
        settings.available_models.pop(model_name)
        nonlocal model_is_alive
        model_is_alive = False
        logger.debug(
            "Model " + model_name +
            " is not responsive. Removing the model from available services..."
        )

    while model_is_alive and not dependency.shutdown:
        try:
            r = requests.get('http://host.docker.internal:' +
                             str(settings.available_models[model_name]) +
                             '/status')
            r.raise_for_status()
            for increment in range(dependency.WAIT_TIME):
                if not dependency.shutdown:  # Check between increments to stop hanging on shutdown
                    time.sleep(1)
        except (requests.exceptions.ConnectionError,
                requests.exceptions.Timeout, requests.exceptions.HTTPError):
            kill_model()
            return

    if dependency.shutdown:
        logger.debug("Model [" + model_name +
                     "] Healthcheck Thread Terminated.")
 def kill_dataset():
     settings.available_datasets.pop(dataset_name)
     nonlocal dataset_is_alive
     dataset_is_alive = False
     logger.debug("Dataset " + dataset_name +
                  " is not responsive. Removing from available services...")
def create_new_prediction_on_image(
        images: List[UploadFile] = File(...),
        models: List[str] = (),
        current_user: User = Depends(current_user_investigator)):
    """
    Create a new prediction request for any number of images on any number of models. This will enqueue the jobs
    and a worker will process them and get the results. Once this is complete, a user may later query the job
    status by the unique key that is returned from this method for each image uploaded.

    :param current_user: User object who is logged in
    :param images: List of file objects that will be used by the models for prediction
    :param models: List of models to run on images
    :return: Unique keys for each image uploaded in images.
    """

    # Start with error checking on the models list.
    # Ensure that all desired models are valid.
    if not models:
        return HTTPException(
            status_code=400,
            detail="You must specify models to process images with")

    invalid_models = []
    for model in models:
        if model not in settings.available_models:
            invalid_models.append(model)

    if invalid_models:
        error_message = "Invalid Models Specified: " + ''.join(invalid_models)
        return HTTPException(status_code=400, detail=error_message)

    # Now we must hash each uploaded image
    # After hashing, we will store the image file on the server.

    buffer_size = 65536  # Read image data in 64KB Chunks for hashlib
    hashes_md5 = {}

    # Process uploaded images
    for upload_file in images:
        file = upload_file.file
        md5 = hashlib.md5()
        while True:
            data = file.read(buffer_size)
            if not data:
                break
            md5.update(data)

        # Process image
        hash_md5 = md5.hexdigest()
        hashes_md5[upload_file.filename] = hash_md5

        file.seek(0)

        if get_image_by_md5_hash_db(hash_md5):
            image_object = get_image_by_md5_hash_db(hash_md5)
        else:  # If image does not already exist in db

            # Create a UniversalMLImage object to store data
            image_object = UniversalMLImage(
                **{
                    'file_names': [upload_file.filename],
                    'hash_md5': hash_md5,
                    'hash_sha1': 'TODO: Remove This Field',
                    'hash_perceptual': 'TODO: Remove This Field',
                    'users': [current_user.username],
                    'models': {}
                })

            # Add created image object to database
            add_image_db(image_object)

        # Associate the current user with the image that was uploaded
        add_user_to_image(image_object, current_user.username)

        # Associate the name the file was uploaded under to the object
        add_filename_to_image(image_object, upload_file.filename)

        # Copy image to the temporary storage volume for prediction
        new_filename = hash_md5 + os.path.splitext(upload_file.filename)[1]
        stored_image_path = "/app/prediction_images/" + new_filename
        stored_image = open(stored_image_path, 'wb+')
        shutil.copyfileobj(file, stored_image)

        for model in models:
            model_socket = settings.available_models[model]
            try:
                logger.debug('Creating Prediction Request. Hash: ' + hash_md5 +
                             ' Model: ' + model)
                request = requests.post(model_socket + '/predict',
                                        params={
                                            'image_md5_hash': hash_md5,
                                            'image_file_name': new_filename
                                        })
                request.raise_for_status(
                )  # Ensure prediction job hasn't errored.
            except (requests.exceptions.ConnectionError,
                    requests.exceptions.Timeout,
                    requests.exceptions.HTTPError):
                logger.error(
                    'Fatal error when creating prediction request. Hash: "' +
                    hash_md5 + '" Model: ' + model)

    return {"images": [hashes_md5[key] for key in hashes_md5]}
def create_new_prediction_on_image(
        images: List[UploadFile] = File(...),
        models: List[str] = (),
        current_user: User = Depends(current_user_investigator)):
    """
    Create a new prediction request for any number of images on any number of models. This will enqueue the jobs
    and a worker will process them and get the results. Once this is complete, a user may later query the job
    status by the unique key that is returned from this method for each image uploaded.

    :param current_user: User object who is logged in
    :param images: List of file objects that will be used by the models for prediction
    :param models: List of models to run on images
    :return: Unique keys for each image uploaded in images.
    """

    # Start with error checking on the models list.
    # Ensure that all desired models are valid.
    if not models:
        return HTTPException(
            status_code=400,
            detail="You must specify models to process images with")

    invalid_models = []
    for model in models:
        if model not in settings.available_models:
            invalid_models.append(model)

    if invalid_models:
        error_message = "Invalid Models Specified: " + ''.join(invalid_models)
        return HTTPException(status_code=400, detail=error_message)

    # Now we must hash each uploaded image
    # After hashing, we will store the image file on the server.

    buffer_size = 65536  # Read image data in 64KB Chunks for hashlib
    hashes_md5 = {}

    # Process uploaded images
    for upload_file in images:
        file = upload_file.file
        md5 = hashlib.md5()
        sha1 = hashlib.sha1()
        while True:
            data = file.read(buffer_size)
            if not data:
                break
            md5.update(data)
            sha1.update(data)

        # Process image
        hash_md5 = md5.hexdigest()
        hash_sha1 = sha1.hexdigest()
        hashes_md5[upload_file.filename] = hash_md5

        if get_image_by_md5_hash_db(hash_md5):
            image_object = get_image_by_md5_hash_db(hash_md5)
        else:  # If image does not already exist in db

            # Generate perceptual hash
            hash_perceptual = str(
                imagehash.phash(Image.open('/app/images/' + file_name)))

            # Create a UniversalMLImage object to store data
            image_object = UniversalMLImage(
                **{
                    'file_names': [upload_file.filename],
                    'hash_md5': hash_md5,
                    'hash_sha1': hash_sha1,
                    'hash_perceptual': hash_perceptual,
                    'users': [current_user.username],
                    'models': {}
                })

            # Add created image object to database
            add_image_db(image_object)

        # Associate the current user with the image that was uploaded
        add_user_to_image(image_object, current_user.username)

        # Associate the name the file was uploaded under to the object
        add_filename_to_image(image_object, upload_file.filename)

        for model in models:
            random_tail = ''.join(
                random.choices(string.ascii_uppercase + string.digits, k=10))

            job_id = hash_md5 + '---' + model + '---' + random_tail

            model_socket = settings.available_models[model]
            logger.debug('Adding Job For For Image ' + hash_md5 +
                         ' With Model ' + model + ' With ID ' + job_id)
            # Submit a job to use scene detection model
            prediction_queue.enqueue(get_model_prediction,
                                     model_socket,
                                     hash_md5,
                                     model,
                                     upload_file,
                                     job_id=job_id)

    return {"images": [hashes_md5[key] for key in hashes_md5]}