cnn-classification-service/cnn_classification_service/main.py

91 lines
2.6 KiB
Python

#!/usr/bin/env python3
import jaeger_client
import opentracing
import logging
import os
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
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
results = MagicDoer.run_everything(msg) # <- This is where the magic happens
if 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 logging
logging.basicConfig(
stream=sys.stdout,
format="%(asctime)s - %(name)s [%(levelname)s]: %(message)s",
level=Config.LOG_LEVEL
)
# setup observability stuffs
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='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()