Add logging

worker.py

@@ -10,11 +10,19 @@ from typing import Tuple, Dict, Any
 
 from mineru_single import Processor
 
+import logging
+
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s [%(levelname)s] %(name)s - %(message)s"
+)
+logger = logging.getLogger(__name__)
 
 class RabbitMQWorker:
     def __init__(self, num_workers: int = 1):
         self.num_workers = num_workers
         self.rabbit_url = os.getenv("RABBITMQ_URL")
+        logger.info("Initializing RabbitMQWorker")
         self.processor = Processor()
 
         self.publisher_connection = None
@@ -23,6 +31,7 @@ class RabbitMQWorker:
 
     def setup_publisher(self):
         if not self.publisher_connection or self.publisher_connection.is_closed:
+            logger.info("Setting up publisher connection to RabbitMQ.")
             connection_params = pika.URLParameters(self.rabbit_url)
             connection_params.heartbeat = 1000 # Match the consumer heartbeat
             connection_params.blocked_connection_timeout = 500  # Increase timeout
@@ -30,6 +39,7 @@ class RabbitMQWorker:
             self.publisher_connection = pika.BlockingConnection(connection_params)
             self.publisher_channel = self.publisher_connection.channel()
             self.publisher_channel.queue_declare(queue="ml_server", durable=True)
+            logger.info("Publisher connection/channel established successfully.")
     
     def publish_message(self, body_dict: dict, headers: dict):
         """Use persistent connection for publishing"""
@@ -48,9 +58,10 @@ class RabbitMQWorker:
                         headers=headers
                     )
                 )
+                logger.info("Published message to ml_server queue (attempt=%d).", attempt + 1)
                 return True
             except Exception as e:
-                print(f"Publish attempt {attempt + 1} failed: {e}")
+                logger.error("Publish attempt %d failed: %s", attempt + 1, e)
                 # Close failed connection
                 if self.publisher_connection and not self.publisher_connection.is_closed:
                     try:
@@ -61,7 +72,7 @@ class RabbitMQWorker:
                 self.publisher_channel = None
                 
                 if attempt == max_retries - 1:
-                    print(f"Failed to publish after {max_retries} attempts")
+                    logger.error("Failed to publish after %d attempts", max_retries)
                     return False
                 time.sleep(2) 
 
@@ -70,7 +81,7 @@ class RabbitMQWorker:
         thread_id = threading.current_thread().name
         headers = properties.headers or {}
 
-        print(f"[Worker {thread_id}] Received message: {body}, headers: {headers}")
+        logger.info("[Worker %s] Received message: %s, headers: %s", thread_id, body, headers)
 
         try:
             if headers.get("request_type") == "process_files":
@@ -81,6 +92,7 @@ class RabbitMQWorker:
                 # Process files
                 for file in body_dict.get("input_files", []):
                     try:
+                        logger.info("[Worker %s] Processing %d input files.", thread_id, len(file))
                         context = {"key": file["key"], "body": self.processor.process(file["url"], file["key"])}
                         contexts.append(context)
                     except Exception as e:
@@ -94,18 +106,18 @@ class RabbitMQWorker:
                 
                 # Publish results
                 if self.publish_message(body_dict, headers):
-                    print(f"[Worker {thread_id}] Successfully published results")
+                    logger.info("[Worker %s] Successfully published results to ml_server.", thread_id)
                 else:
                     ch.basic_nack(delivery_tag=method.delivery_tag, requeue=True)
-                    print(f"[Worker {thread_id}] Failed to publish results")
+                    logger.error("[Worker %s] Failed to publish results.", thread_id)
                 
-                print(f"[Worker {thread_id}] Contexts: {contexts}")
+                logger.info("[Worker %s] Contexts: %s", thread_id, contexts)
             else:
                 ch.basic_ack(delivery_tag=method.delivery_tag)
-                print(f"[Worker {thread_id}] Unknown process")
+                logger.warning("[Worker %s] Unknown process type in headers: %s", thread_id, headers.get("request_type"))
                 
         except Exception as e:
-            print(f"Error in callback: {e}")
+            logger.error("Error in callback: %s", e)
             ch.basic_nack(delivery_tag=method.delivery_tag, requeue=True)
 
     def connect_to_rabbitmq(self):
@@ -113,6 +125,8 @@ class RabbitMQWorker:
         connection_params = pika.URLParameters(self.rabbit_url)
         connection_params.heartbeat = 1000
         connection_params.blocked_connection_timeout = 500
+        logger.info("Connecting to RabbitMQ for consumer with heartbeat=1000.")
+
         
         connection = pika.BlockingConnection(connection_params)
         channel = connection.channel()
@@ -123,21 +137,23 @@ class RabbitMQWorker:
             queue="gpu_server",
             on_message_callback=self.callback
         )
+        logger.info("Consumer connected. Listening on queue='gpu_server'...")
         return connection, channel
 
     def worker(self, channel):
         """Worker function"""
-        print(f"Worker started")
+        logger.info("Worker thread started. Beginning consuming...")
         try:
             channel.start_consuming()
         except Exception as e:
-            print(f"Worker stopped: {e}")
+            logger.error("Worker thread encountered an error: %s", e)
         finally:
+            logger.info("Worker thread shutting down. Closing channel.")
             channel.close()
 
     def start(self):
         """Start the worker threads"""
-        print(f"Starting {self.num_workers} workers")
+        logger.info("Starting %d workers in a ThreadPoolExecutor.", self.num_workers)
         while True:
             try:
                 with ThreadPoolExecutor(max_workers=self.num_workers) as executor:
@@ -145,7 +161,7 @@ class RabbitMQWorker:
                         connection, channel = self.connect_to_rabbitmq()
                         executor.submit(self.worker, channel)
             except Exception as e:
-                print(f"Connection lost, reconnecting... Error: {e}")
+                logger.error("Connection lost, reconnecting... Error: %s", e)
                 time.sleep(5)  # Wait before reconnecting
 
 def main():