vitor-aignosi
5787d9f7a0
Refactor Gates activity to correctly sort data by timestamp and update exception handling in worker to include NOSONAR comment for clarity.
110 lines
3.5 KiB
Python
110 lines
3.5 KiB
Python
from temporalio import workflow, client
|
|
from temporalio.worker import Worker
|
|
|
|
|
|
with workflow.unsafe.imports_passed_through():
|
|
import os
|
|
import sys
|
|
import asyncio
|
|
from laborious.workflows.predictions_batch import PredictionsBatch
|
|
from laborious.workflows.sub_workflows.prediction_process import PredictionProcess
|
|
from laborious.workflows.sub_workflows.format_and_export_prediction import \
|
|
FormatAndExportPrediction
|
|
from laborious.activities.activities import Activities
|
|
from laborious.utils.connectors_config import (
|
|
build_postgres_config,
|
|
build_mlflow_config,
|
|
build_opc_config
|
|
)
|
|
from sientia_do.notifications.handlers import NotificationHandler
|
|
from sientia_do.temporal.utils.logger import get_logger
|
|
|
|
|
|
async def main():
|
|
host = os.getenv('TEMPORAL_HOST', 'localhost:7233')
|
|
logger = get_logger(__name__)
|
|
|
|
logger.info('Starting Worker...')
|
|
|
|
logger.info('Starting Notification Handler...')
|
|
|
|
notification_handler = NotificationHandler(
|
|
servers=os.getenv('KAFKA_BOOTSTRAP_SERVERS', 'http://localhost:9092'),
|
|
logger=logger,
|
|
project_name=os.getenv('PROJECT_NAME', 'laborious'),
|
|
)
|
|
|
|
logger.info('Starting Activities...')
|
|
|
|
activities = Activities(
|
|
postgres_config=build_postgres_config(),
|
|
mlflow_config=build_mlflow_config(),
|
|
opc_config=build_opc_config(),
|
|
logger=logger,
|
|
notification_handler=notification_handler
|
|
)
|
|
|
|
logger.info('Starting Temporal Client...')
|
|
|
|
temporal_client = await client.Client.connect(
|
|
target_host=host,
|
|
namespace=os.getenv('TEMPORAL_NAMESPACE', 'default')
|
|
)
|
|
|
|
logger.info('Starting Workers...')
|
|
|
|
workers = [
|
|
Worker(
|
|
temporal_client,
|
|
task_queue='predictions_batch-queue',
|
|
workflows=[PredictionsBatch, PredictionProcess,
|
|
FormatAndExportPrediction],
|
|
activities=[
|
|
# MLFlow
|
|
activities.request_predict,
|
|
activities.request_transform,
|
|
# Gates
|
|
activities.input_gate,
|
|
activities.mlflow_response_gate,
|
|
activities.mlflow_content_gate,
|
|
activities.format_prediction,
|
|
activities.format_default_prediction,
|
|
activities.get_last_timestamp,
|
|
# OPC
|
|
activities.write_opc_data,
|
|
# Postgres
|
|
activities.load_custom_query,
|
|
activities.repeat_last_prediction,
|
|
activities.export_data_to_postgres
|
|
],
|
|
max_concurrent_workflow_tasks=100,
|
|
max_concurrent_activities=100,
|
|
max_concurrent_local_activities=100,
|
|
max_concurrent_workflow_task_polls=100,
|
|
max_cached_workflows=50,
|
|
)
|
|
]
|
|
|
|
handlers = []
|
|
for w in workers:
|
|
handlers.append(w.run())
|
|
|
|
logger.info('Workers started successfully')
|
|
|
|
try:
|
|
# This will run the workers and wait for them to complete.
|
|
# If an exception occurs in any of the worker handlers, it will be propagated here.
|
|
await asyncio.gather(*handlers)
|
|
except BaseException as e: # NOSONAR
|
|
logger.error(f"An unhandled exception occurred: {e}")
|
|
finally:
|
|
if notification_handler:
|
|
notification_handler.shutdown()
|
|
if activities:
|
|
activities.shutdown()
|
|
# Exit with a non-zero status code to indicate failure to Kubernetes
|
|
sys.exit(1)
|
|
|
|
if __name__ == '__main__':
|
|
asyncio.run(main())
|