cnn-classification-service/cnn_classification_service/main.py

#!/usr/bin/env python3
import logging
import os
import sys
import pika
import json

from sentry_sdk.integrations.logging import LoggingIntegration
from sentry_sdk import start_transaction
import sentry_sdk

from magic_doer import MagicDoer


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}")
        return

    with start_transaction(op="cnn-classification-service", name="classify-soundfile"):
        results = MagicDoer.run_everything(msg)  # <- This is where the magic happens

    if results:
        channel.basic_publish(
            exchange=os.environ['PIKA_OUTPUT_EXCHANGE'],
            routing_key='classification-result',
            body=json.dumps(results).encode("utf-8")
        )
        channel.basic_ack(delivery_tag=method.delivery_tag)


def main():
    logging.basicConfig(
        stream=sys.stdout,
        format="%(asctime)s - %(name)s [%(levelname)s]: %(message)s",
        level=logging.DEBUG if '--debug' in sys.argv else logging.INFO
    )

    SENTRY_DSN = os.environ.get("SENTRY_DSN")
    if 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=SENTRY_DSN,
            integrations=[sentry_logging],
            traces_sample_rate=1.0,
            send_default_pii=True,
            release=os.environ.get('RELEASE_ID', 'test'),
            environment=os.environ.get('RELEASEMODE', 'dev'),
            _experiments={"auto_enabling_integrations": True}
        )

    logging.info("Connecting to MQ service...")
    connection = pika.BlockingConnection(pika.connection.URLParameters(os.environ['PIKA_URL']))
    channel = connection.channel()
    channel.exchange_declare(exchange=os.environ['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=os.environ['PIKA_INPUT_EXCHANGE'], routing_key='feature', 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()


if __name__ == '__main__':
    main()