SIENTIAPDE-1478

Refactor worker initialization by introducing prepare_worker function to streamline worker setup for orchestrator, alerts, and reports workflows, enhancing code clarity and maintainability.

orchestrator/worker/prepare_worker.py

@@ -0,0 +1,72 @@
+import os
+import re
+from collections.abc import Sequence
+from typing import Any
+
+from sientia_do.observability.logger import Logger
+from temporalio.client import Client
+from temporalio.worker import PollerBehaviorAutoscaling, Worker
+
+parameters = [
+    ('MAX_CONCURRENT_WORKFLOW_TASKS', '200'),
+    ('MAX_CONCURRENT_ACTIVITIES', '200'),
+    ('MAX_CONCURRENT_LOCAL_ACTIVITIES', '200'),
+    ('MAX_CACHED_WORKFLOWS', '200'),
+    ('WORKFLOW_POLLER_BEHAVIUR_MINIMUM', '10'),
+    ('WORKFLOW_POLLER_BEHAVIUR_INITIAL', '100'),
+    ('WORKFLOW_POLLER_BEHAVIUR_MAXIMUM', '200'),
+    ('ACTIVITY_POLLER_BEHAVIUR_MINIMUM', '10'),
+    ('ACTIVITY_POLLER_BEHAVIUR_INITIAL', '100'),
+    ('ACTIVITY_POLLER_BEHAVIUR_MAXIMUM', '200'),
+]
+
+
+def camel_to_snake(text: str) -> str:
+    """Convert camelCase or PascalCase to snake_case."""
+    text = re.sub('(.)([A-Z][a-z]+)', r'\1_\2', text)
+    text = re.sub('([a-z0-9])([A-Z])', r'\1_\2', text)
+    return text.lower()
+
+
+def prepare_worker(
+    main_workflow: type,
+    other_workflows: Sequence[type],
+    activities: Sequence[Any],
+    temporal_client: Client,
+    logger: Logger,
+) -> Worker:
+    main_workflow_name = main_workflow.__name__.upper()
+
+    queue_name = f'{camel_to_snake(main_workflow.__name__)}-queue'
+
+    local_workflow_parameters = {}
+
+    for parameter in parameters:
+        local_workflow_parameters[parameter[0]] = int(
+            os.getenv(main_workflow_name + '_' + parameter[0], parameter[1])
+        )
+
+    logger.info(f'Preparing worker for {main_workflow_name} with queue {queue_name}')
+
+    return Worker(
+        temporal_client,
+        task_queue=queue_name,
+        workflows=[main_workflow, *other_workflows],
+        activities=[*activities],
+        max_concurrent_workflow_tasks=local_workflow_parameters['MAX_CONCURRENT_WORKFLOW_TASKS'],
+        max_concurrent_activities=local_workflow_parameters['MAX_CONCURRENT_ACTIVITIES'],
+        max_concurrent_local_activities=local_workflow_parameters[
+            'MAX_CONCURRENT_LOCAL_ACTIVITIES'
+        ],
+        max_cached_workflows=local_workflow_parameters['MAX_CACHED_WORKFLOWS'],
+        workflow_task_poller_behavior=PollerBehaviorAutoscaling(
+            minimum=local_workflow_parameters['WORKFLOW_POLLER_BEHAVIUR_MINIMUM'],
+            initial=local_workflow_parameters['WORKFLOW_POLLER_BEHAVIUR_INITIAL'],
+            maximum=local_workflow_parameters['WORKFLOW_POLLER_BEHAVIUR_MAXIMUM'],
+        ),
+        activity_task_poller_behavior=PollerBehaviorAutoscaling(
+            minimum=local_workflow_parameters['ACTIVITY_POLLER_BEHAVIUR_MINIMUM'],
+            initial=local_workflow_parameters['ACTIVITY_POLLER_BEHAVIUR_INITIAL'],
+            maximum=local_workflow_parameters['ACTIVITY_POLLER_BEHAVIUR_MAXIMUM'],
+        ),
+    )

orchestrator/worker/worker.py

@@ -7,6 +7,8 @@ with workflow.unsafe.imports_passed_through():
     import os
     import sys
 
+    from orchestrator.worker.prepare_worker import prepare_worker
+
     from prometheus_client import start_http_server
     from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
     from sientia_do.observability.logger import get_logger
@@ -17,7 +19,6 @@ with workflow.unsafe.imports_passed_through():
         build_email_config,
         build_mongodb_config,
         build_postgres_config,
-        # build_couchbase_config,
         build_redis_config,
         build_temporal_config,
     )
@@ -104,10 +105,10 @@ async def main():
     logger.custom_info('Starting Workers...', metadata=metadata)
 
     workers = [
-        Worker(
-            temporal_client,
-            task_queue='orchestrator-queue',
-            workflows=[Orchestrator],
+        prepare_worker(
+            temporal_client=temporal_client,
+            main_workflow=Orchestrator,
+            other_workflows=[],
             activities=[
                 # Redis
                 activities.load_active_ingestors,
@@ -137,11 +138,12 @@ async def main():
                 activities.report_slot_orchestration,
                 activities.format_schedule_config,
             ],
+            logger=logger,
         ),
-        Worker(
-            temporal_client,
-            task_queue='alerts-queue',
-            workflows=[Alerts, LoadNotificationPackage, ProcessNotifications],
+        prepare_worker(
+            temporal_client=temporal_client,
+            main_workflow=Alerts,
+            other_workflows=[LoadNotificationPackage, ProcessNotifications],
             activities=[
                 # Load notifications
                 activities.get_last_data_timestamp,
@@ -158,11 +160,12 @@ async def main():
                 # Store notification cache
                 activities.store_notification_cache,
             ],
+            logger=logger,
         ),
-        Worker(
-            temporal_client,
-            task_queue='reports-queue',
-            workflows=[Reports, LoadNotificationPackage, ProcessNotifications],
+        prepare_worker(
+            temporal_client=temporal_client,
+            main_workflow=Reports,
+            other_workflows=[LoadNotificationPackage, ProcessNotifications],
             activities=[
                 # Load notifications
                 activities.get_last_data_timestamp,
@@ -177,6 +180,7 @@ async def main():
                 activities.format_log_report,
                 activities.export_data_to_postgres,
             ],
+            logger=logger,
         ),
     ]