cnn-classification-service/cnn_classification_service/main.py

#!/usr/bin/env python3
import time

import jaeger_client
import opentracing
from opentracing.ext import tags
from opentracing.propagation import Format

import logging
import sys
import pika
import json

from sentry_sdk.integrations.logging import LoggingIntegration
import sentry_sdk

from config import Config

from magic_doer import MagicDoer
from classifier_cache import ClassifierCache


def message_callback(channel, method, properties, body):
    try:
        msg = json.loads(body.decode('utf-8'))
    except (UnicodeDecodeError, json.JSONDecodeError) as e:
        logging.warning(f"Invalid message recieved: {e}")
        channel.basic_ack(delivery_tag=method.delivery_tag)  # We don't want this to be requeue
        return

    logging.debug(f"Handling message: {msg}")

    span_ctx = opentracing.tracer.extract(Format.TEXT_MAP, msg)
    span_tags = {tags.SPAN_KIND: tags.SPAN_KIND_CONSUMER}

    with opentracing.tracer.start_active_span(
            'main.handleMessage', finish_on_close=True, child_of=span_ctx, tags=span_tags
    ) as scope:

        with opentracing.tracer.start_active_span('magicDoer.runEverything'):
            try:
                results = MagicDoer.run_everything(msg)  # <- This is where the magic happens
            except Exception as e:
                logging.error(f"Something went wrong during handling sample {msg['tag']} run: {e}; msg: {msg}")
                logging.exception(e)
                channel.basic_nack(delivery_tag=method.delivery_tag, requeue=True)
                return

        if results:
            opentracing.tracer.inject(scope.span.context, Format.TEXT_MAP, results)
            logging.debug(f"Publishing message: {results}")
            channel.basic_publish(
                exchange=Config.PIKA_OUTPUT_EXCHANGE,
                routing_key='classification-result',
                body=json.dumps(results).encode("utf-8")
            )

        channel.basic_ack(delivery_tag=method.delivery_tag)


def main():
    # setup observability stuffs
    if (not Config.PIKA_URL) or (not Config.PIKA_OUTPUT_EXCHANGE):
        logging.error("Mandatory config parameters unset: PIKA_URL or PIKA_OUTPUT_EXCHANGE")
        raise KeyError

    if Config.SENTRY_DSN:
        sentry_logging = LoggingIntegration(
            level=logging.DEBUG,  # Capture info and above as breadcrumbs
            event_level=logging.ERROR  # Send errors as events
        )
        sentry_sdk.init(
            dsn=Config.SENTRY_DSN,
            integrations=[sentry_logging],
            traces_sample_rate=1.0,
            send_default_pii=True,
            release=Config.RELEASE_ID,
            environment=Config.RELEASEMODE,
            _experiments={"auto_enabling_integrations": True}
        )

    jaeger_client.Config(config={}, service_name='cnn-classification-service', validate=True).initialize_tracer()

    # Start the memes

    logging.info("Connecting to MQ service...")
    connection = pika.BlockingConnection(pika.connection.URLParameters(Config.PIKA_URL))
    channel = connection.channel()
    channel.exchange_declare(exchange=Config.PIKA_INPUT_EXCHANGE, exchange_type='direct')

    queue_declare_result = channel.queue_declare(queue='cnnqueue', exclusive=False)
    queue_name = queue_declare_result.method.queue

    channel.queue_bind(exchange=Config.PIKA_INPUT_EXCHANGE, routing_key='sample', queue=queue_name)

    channel.basic_qos(prefetch_count=1)

    channel.basic_consume(queue=queue_name, on_message_callback=message_callback, auto_ack=False)

    logging.info("Connection complete! Listening to messages...")
    try:
        channel.start_consuming()
    except KeyboardInterrupt:
        logging.info("SIGINT Received! Stopping stuff...")
        channel.stop_consuming()

    time.sleep(2)  # lol
    opentracing.tracer.close()


def test_loader():
    logging.info("Testing if model loading works...")
    cc = ClassifierCache(Config.MODEL_INFO_URL)
    details, classifier = cc.get_default_classifier()
    logging.info(f"Loaded classifier: {classifier}")
    logging.info(f"Details: {details}")


if __name__ == '__main__':
    # setup logging

    logging.basicConfig(
        stream=sys.stdout,
        format="%(asctime)s - %(name)s [%(levelname)s]: %(message)s",
        level=Config.LOG_LEVEL
    )

    if '--test-loader' in sys.argv:
        test_loader()
    else:
        main()