- Update pyinfra with absolute file path support (still supports dossierID fileID format) - Update CI, use new template
38 lines
1.4 KiB
Python
38 lines
1.4 KiB
Python
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.processor 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()
|