cnn-classification-service/cnn_classification_service/main.py

91 lines
2.6 KiB
Python
Raw Normal View History

2020-07-27 17:58:48 +02:00
#!/usr/bin/env python3
2021-07-30 12:17:18 +02:00
import jaeger_client
import opentracing
2020-07-27 17:58:48 +02:00
import logging
import os
import sys
import pika
import json
from sentry_sdk.integrations.logging import LoggingIntegration
import sentry_sdk
2021-07-30 11:43:38 +02:00
from config import Config
2020-10-06 00:41:54 +02:00
from magic_doer import MagicDoer
2020-07-27 17:58:48 +02:00
2020-10-05 22:19:25 +02:00
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
2020-07-28 18:10:19 +02:00
2021-07-30 11:47:28 +02:00
results = MagicDoer.run_everything(msg) # <- This is where the magic happens
2020-10-05 22:19:25 +02:00
2021-06-14 03:12:44 +02:00
if results:
channel.basic_publish(
2021-07-30 11:43:38 +02:00
exchange=Config.PIKA_OUTPUT_EXCHANGE,
2021-06-14 03:12:44 +02:00
routing_key='classification-result',
body=json.dumps(results).encode("utf-8")
)
2021-06-14 03:16:56 +02:00
channel.basic_ack(delivery_tag=method.delivery_tag)
2020-07-27 17:58:48 +02:00
def main():
2021-07-30 12:17:18 +02:00
# setup logging
2020-10-05 22:19:25 +02:00
logging.basicConfig(
stream=sys.stdout,
format="%(asctime)s - %(name)s [%(levelname)s]: %(message)s",
2021-07-30 11:46:06 +02:00
level=Config.LOG_LEVEL
2020-10-05 22:19:25 +02:00
)
2020-07-27 17:58:48 +02:00
2021-07-30 12:17:18 +02:00
# setup observability stuffs
2021-07-30 11:43:38 +02:00
if Config.SENTRY_DSN:
2020-07-27 17:58:48 +02:00
sentry_logging = LoggingIntegration(
level=logging.DEBUG, # Capture info and above as breadcrumbs
event_level=logging.ERROR # Send errors as events
)
sentry_sdk.init(
2021-07-30 11:43:38 +02:00
dsn=Config.SENTRY_DSN,
2020-07-27 17:58:48 +02:00
integrations=[sentry_logging],
2020-10-19 22:30:00 +02:00
traces_sample_rate=1.0,
2020-07-27 17:58:48 +02:00
send_default_pii=True,
2021-07-30 11:43:38 +02:00
release=Config.RELEASE_ID,
environment=Config.RELEASEMODE,
2020-10-19 22:30:00 +02:00
_experiments={"auto_enabling_integrations": True}
2020-07-27 17:58:48 +02:00
)
2021-07-30 12:17:18 +02:00
jaeger_client.Config(config={}, service_name='cnn-classification-service', validate=True).initialize_tracer()
# Start the memes
2020-07-27 17:58:48 +02:00
logging.info("Connecting to MQ service...")
2021-07-30 11:43:38 +02:00
connection = pika.BlockingConnection(pika.connection.URLParameters(Config.PIKA_URL))
2020-07-27 17:58:48 +02:00
channel = connection.channel()
2021-07-30 11:43:38 +02:00
channel.exchange_declare(exchange=Config.PIKA_INPUT_EXCHANGE, exchange_type='direct')
2020-07-27 17:58:48 +02:00
2020-10-24 00:58:28 +02:00
queue_declare_result = channel.queue_declare(queue='cnnqueue', exclusive=False)
2020-07-27 17:58:48 +02:00
queue_name = queue_declare_result.method.queue
2021-07-30 11:43:38 +02:00
channel.queue_bind(exchange=Config.PIKA_INPUT_EXCHANGE, routing_key='feature', queue=queue_name)
2020-10-06 00:41:54 +02:00
2020-10-24 00:50:55 +02:00
channel.basic_qos(prefetch_count=1)
channel.basic_consume(queue=queue_name, on_message_callback=message_callback, auto_ack=False)
2020-07-27 17:58:48 +02:00
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__':
2020-07-28 18:10:19 +02:00
main()