Spaces:
Running
Running
File size: 3,611 Bytes
e8a6607 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 | import logging
import os
import signal
import time
import sys
import asyncio
# from kafka_python.topic_setup import create_topics
# from ModerationService.image_moderation_consumer import ImageModerationConsumer
# from TextGeneration.generate_ai_desc_consumer import GenAiDescConsumer
from classInfra import worker_service
from QueueAndWorker.model_manager import load_model
# Check if the environment is production. Default to 'development' if not set.
ENVIRONMENT = os.getenv("APP_ENV", "development")
# Automatically pick the level based on the environment
if ENVIRONMENT == "production":
LOG_LEVEL = logging.INFO # Hides DEBUG noise in production
else:
LOG_LEVEL = logging.DEBUG # Shows EVERYTHING in development
logging.basicConfig(
level=LOG_LEVEL,
format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
)
# logging.getLogger("kafka").setLevel(logging.WARNING)
# Works all of these if set logging.DEBUG
# logger.debug()
# logger.info()
# logger.warning()
# logger.error()
# logger.critical()
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
# def kafka_init():
# create_topics()
# # START PRODUCERS
# producer_manager.create_producer(
# producer_name="image-moderation-producer",
# max_retries=6
# )
# producer_manager.create_producer(
# producer_name="text-gen-producer",
# max_retries=6
# )
# # START CONSUMERS
# image_moderation_consumer = ImageModerationConsumer()
# image_moderation_consumer.start()
# text_generation_consumer = GenAiDescConsumer()
# text_generation_consumer.start()
shutdown_event = None
def signal_handler(sig, frame):
logger.info(
"🛑 Shutdown signal received"
)
shutdown_event.set()
# MAIN APPLICATION
async def main():
global shutdown_event
exit_code = 0
# logger.info(
# "🚀 Starting moderation service"
# )
# logger.info(
# "Starting topic creation"
# )
try:
# Simulate adding jobs via our Producer class
# await queue_service.connect()
# Start up the background worker service
await worker_service.start()
load_model()
# CREATE TOPICS
# 💡 FIX: Force the blocking Kafka call to execute on a separate thread
# await asyncio.to_thread(kafka_init)
shutdown_event = asyncio.Event()
# Wait until Ctrl+C or SIGTERM
await shutdown_event.wait()
# KEEP SERVICE RUNNING - no need this because we already has await shutdown_event.wait()
# while True:
# await asyncio.sleep(60)
except Exception as e:
logger.exception("❌ Error starting service")
exit_code = 1
finally:
logging.info("Cleaning up...")
try:
# consumer_manager.stop_all_consumers()
# producer_manager.stop_all_producers()
# Clean queue and worker
await worker_service.stop()
# await queue_service.disconnect()
logger.info(
"✅ Graceful shutdown completed"
)
except Exception as e:
logger.exception("❌ Graceful shutdown failed")
finally:
sys.exit(exit_code)
# SERVICE ENTRY POINT
if __name__ == "__main__":
signal.signal(signal.SIGINT, signal_handler)
signal.signal(signal.SIGTERM, signal_handler)
asyncio.run(main()) # no need of sys.exit(0) at last, asyncio will do that |