add logging

Author: AlpinDale

pyproject.toml

@@ -10,6 +10,7 @@ requires-python = ">=3.8"
 dependencies = [
     "aiohttp>=3.8.0",
     "pyyaml>=6.0",
+    "loguru>=0.7.0",
 ]
 
 [project.optional-dependencies]

src/aphrodite_loadbalancer/__main__.py

@@ -1,12 +1,21 @@
 import asyncio
 import sys
 
+from loguru import logger
+
 from .loadbalancer import LoadBalancer
 
 
+logger.remove()
+logger.add(
+    sys.stderr,
+    format="<green>{time:HH:mm:ss}</green> | <level>{message}</level>",
+)
+
+
 def main():
     if len(sys.argv) != 2:
-        print('Usage: aphrodite-loadbalancer <config.yaml>')
+        logger.error('Usage: aphrodite-loadbalancer <config.yaml>')
         sys.exit(1)
 
     config_path = sys.argv[1]
@@ -21,7 +30,7 @@ async def async_main(config_path: str):
         while True:
             await asyncio.sleep(3600)
     except KeyboardInterrupt:
-        print('\nShutting down...')
+        logger.info('\nShutting down...')
     finally:
         await balancer.cleanup()
 

src/aphrodite_loadbalancer/loadbalancer.py

@@ -6,6 +6,7 @@
 import aiohttp
 import yaml
 from aiohttp import web
+from loguru import logger
 
 
 class LoadBalancer:
@@ -56,11 +57,11 @@ async def monitor_health(self):
                 is_healthy = await self.health_check(endpoint)
 
                 if not is_healthy and was_healthy:
-                    print(f'[Health] Endpoint {endpoint} is down')
+                    logger.warning(f'Endpoint {endpoint} is down')
                     self.unhealthy_endpoints.add(i)
                     self._create_weighted_cycles()
                 elif is_healthy and not was_healthy:
-                    print(f'[Health] Endpoint {endpoint} is back up')
+                    logger.info(f'Endpoint {endpoint} is back up')
                     self.unhealthy_endpoints.remove(i)
                     self._create_weighted_cycles()
 
@@ -75,7 +76,7 @@ def _create_weighted_cycles(self):
                 weighted_indices.extend([i] * weight)
 
         if not weighted_indices:
-            print('WARNING: All endpoints are unhealthy!')
+            logger.critical('All endpoints are unhealthy!')
             weighted_indices = list(range(len(self.endpoints)))
 
         random.shuffle(weighted_indices)
@@ -94,12 +95,12 @@ async def start(self, port: int):
         await runner.setup()
         site = web.TCPSite(runner, '0.0.0.0', port)
         await site.start()
-        print(f'Load balancer running on http://0.0.0.0:{port}')
+        logger.info(f'Load balancer running on http://0.0.0.0:{port}')
 
         for i, (endpoint, weight) in enumerate(
             zip(self.endpoints, self.weights)
         ):
-            print(f'Endpoint {i}: {endpoint} (weight: {weight})')
+            logger.info(f'Endpoint {i}: {endpoint} (weight: {weight})')
 
     async def handle_request(self, request: web.Request) -> web.StreamResponse:
         cors_headers = {
@@ -126,6 +127,10 @@ async def handle_request(self, request: web.Request) -> web.StreamResponse:
         if request.query_string:
             path += f'?{request.query_string}'
         target_url = f"{target_url.rstrip('/')}/{path.lstrip('/')}"
+        
+        logger.info(
+            f"Routing {request.method} {path} to endpoint "
+            f"{endpoint_index}: {target_url}")
 
         try:
             assert self.client_session is not None
@@ -149,7 +154,7 @@ async def handle_request(self, request: web.Request) -> web.StreamResponse:
                 return response
 
         except Exception as e:
-            print(f'Request failed: {str(e)}')
+            logger.error(f'Request failed: {str(e)}')
             raise
 
     async def cleanup(self):