def _check_memory_threshold_for_model(device: str, model_size: Union[float, int], calling_func: str = None) -> bool:
    '''
    Note: this function should only be called by `_validate_model_into_device` for threading safeness.
    `_validate_model_into_device` is calle by `_update_available_models` which is already thread safe.

    Check the memory usage in the target device and decide whether we can add a new model
    Args:
        device: the target device to check
        model_size: the size of the model to load
    Returns:
        True if we have enough space
        False if we don't have enough space
    '''
    if calling_func not in ["unit_test", "_validate_model_into_device"]:
        raise RuntimeError(f"The function `{_check_memory_threshold_for_model.__name__}` should only be called by "
                           f"`unit_test` or `_validate_model_into_device` for threading safeness.")

    if device.startswith("cuda"):
        torch.cuda.synchronize(device)
        torch.cuda.empty_cache()
        used_memory = sum([_available_models[key].get("model_size", constants.DEFAULT_MODEL_SIZE) for key, values in
                           _available_models.items() if key.endswith(device)])
        threshold = float(read_env_vars_and_defaults(EnvVars.MARQO_MAX_CUDA_MODEL_MEMORY))
    elif device.startswith("cpu"):
        used_memory = sum([_available_models[key].get("model_size", constants.DEFAULT_MODEL_SIZE) for key, values in
                           _available_models.items() if key.endswith("cpu")])
        threshold = float(read_env_vars_and_defaults(EnvVars.MARQO_MAX_CPU_MODEL_MEMORY))
    else:
        raise ModelCacheManagementError(
            f"Unable to check the device cache for device=`{device}`. The model loading will proceed"
            f"without device cache check. This might break down Marqo if too many models are loaded.")
    if model_size > threshold:
        raise ModelCacheManagementError(
            f"You are trying to load a model with size = `{model_size}` into device = `{device}`, which is larger than the device threshold = `{threshold}`. "
            f"Marqo CANNOT find enough space for the model. Please change the threshold by adjusting the environment variables.\n"
            f"Please modify the threshold by setting the environment variable `MARQO_MAX_CUDA_MODEL_MEMORY` or `MARQO_MAX_CPU_MODEL_MEMORY`."
            f"You can find more detailed information at {marqo_docs.configuring_marqo()}.")
    return (used_memory + model_size) < threshold

    Note: this function should only be called by `_validate_model_into_device` for threading safeness.
    `_validate_model_into_device` is calle by `_update_available_models` which is already thread safe.

    Check the memory usage in the target device and decide whether we can add a new model
    Args:
        device: the target device to check
        model_size: the size of the model to load
    Returns:
        True if we have enough space
        False if we don't have enough space
    
src/marqo/s2_inference/s2_inference.py