SIENTIAPDE-1350: Configure separate task queues for train and cleanup workers, update sientia-dataops-library to 1.6.1 and refactor prometheus server startup to use logger.

model_manager/worker/worker.py

@@ -33,6 +33,7 @@ with workflow.unsafe.imports_passed_through():
 
     from prometheus_client import start_http_server
     from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
+    from sientia_do.observability.logger import Logger as SientiaLogger
 
     from model_manager import metrics
     from model_manager.activities.activities import Activities
@@ -48,6 +49,8 @@ with workflow.unsafe.imports_passed_through():
 
 POD_ID = os.getenv('POD_ID')
 SDK_METRICS_PORT = int(os.getenv('HTTP_SDK_METRICS_PORT', '9091'))
+TRAIN_TASK_QUEUE = os.getenv('TRAIN_TASK_QUEUE', 'train_model-queue')
+CLEANUP_TASK_QUEUE = os.getenv('CLEANUP_TASK_QUEUE', 'cleanup-queue')
 
 
 async def main():
@@ -77,13 +80,10 @@ async def main():
     }
 
     logger.custom_info(f'Starting Worker with POD_ID: {POD_ID}', metadata)
-
-    logger.custom_info('Starting prometheus client...', metadata)
-    start_prometheus_server()
-
+    start_prometheus_server(logger, metadata)
     logger.custom_info('Starting Notification Handler...', metadata)
-
     mongo_config = build_mongodb_config()
+
     notification_handler = NotificationHandler(
         connection_string=mongo_config['connection_string'],
         database=mongo_config['database_name'],
@@ -122,7 +122,7 @@ async def main():
     workers = [
         Worker(
             temporal_client,
-            task_queue='train_model-queue',
+            task_queue=TRAIN_TASK_QUEUE,
             workflows=[TrainModel],
             activities=[
                 activities.update_experiment_run,
@@ -139,7 +139,7 @@ async def main():
         ),
         Worker(
             temporal_client,
-            task_queue='cleanup-queue',
+            task_queue=CLEANUP_TASK_QUEUE,
             workflows=[CleanupFiles],
             activities=[
                 activities.cleanup_minio_files,
@@ -155,6 +155,7 @@ async def main():
     ]
 
     handlers = []
+
     for w in workers:
         handlers.append(w.run())
 
@@ -174,7 +175,7 @@ async def main():
         sys.exit(1)
 
 
-def start_prometheus_server():
+def start_prometheus_server(logger: SientiaLogger, metadata: dict[str, str | None]):
     """
     Starts the Prometheus metrics server for monitoring and observability.
 
@@ -194,10 +195,10 @@ def start_prometheus_server():
     try:
         port = int(os.getenv('HTTP_METRICS_PORT', 9090))
         start_http_server(port)
-        print(f'Prometheus server started on port {port}.')
+        logger.custom_info(f'Prometheus server started on port {port}.', metadata)
         metrics.APP_UP.labels(pod_id=POD_ID).set(1)  # Mark app as UP
     except Exception as e:  # noqa: BLE001
-        print(f'Failed to start Prometheus server: {e}')
+        logger.custom_critical(f'Failed to start Prometheus server: {e}', metadata)
         os._exit(1)