diff --git a/orchestrator/worker/prepare_worker.py b/orchestrator/worker/prepare_worker.py new file mode 100644 index 0000000..c8b68f2 --- /dev/null +++ b/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'], + ), + ) diff --git a/orchestrator/worker/worker.py b/orchestrator/worker/worker.py index fe65d4f..8943b03 100644 --- a/orchestrator/worker/worker.py +++ b/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, ), ]