from image_prediction import logger
from image_prediction.config import Config
from image_prediction.locations import CONFIG_FILE
from image_prediction.pipeline import load_pipeline
from image_prediction.utils.banner import load_banner
from image_prediction.utils.process_wrapping import wrap_in_process
from pyinfra import config
from pyinfra.payload_processing import make_payload_processor
from pyinfra.queue.queue_manager import QueueManager

PYINFRA_CONFIG = config.get_config()
IMAGE_CONFIG = Config(CONFIG_FILE)

logger.setLevel(PYINFRA_CONFIG.logging_level_root)


# A component of the processing pipeline (probably tensorflow) does not release allocated memory (see RED-4206).
# See: https://stackoverflow.com/questions/39758094/clearing-tensorflow-gpu-memory-after-model-execution
# Workaround: Manage Memory with the operating system, by wrapping the processing in a sub-process.
# FIXME: Find more fine-grained solution or if the problem occurs persistently for python services,
@wrap_in_process
def process_data(data: bytes) -> list:
    pipeline = load_pipeline(verbose=IMAGE_CONFIG.service.verbose, batch_size=IMAGE_CONFIG.service.batch_size)
    return list(pipeline(data))


def main():
    logger.info(load_banner())

    process_payload = make_payload_processor(process_data, config=PYINFRA_CONFIG)

    queue_manager = QueueManager(PYINFRA_CONFIG)
    queue_manager.start_consuming(process_payload)


if __name__ == "__main__":
    main()