Files
sientia-dataops-laborious-l…/laborious/worker/worker.py
vitor-aignosi 67fe4afaa6 SIENTIAPDE-994
Refactor and enhance the laborious workflow and utilities

- Removed outdated test file `test_predictions_batch.py` from workflows.
- Added `input_sample.json` for standardized input configuration.
- Introduced `connectors_config.py` to manage database and service configurations.
- Implemented a logging utility in `logger.py` for consistent logging across the application.
- Created `policies.py` to define retry policies for workflows.
- Developed comprehensive tests for `MLFlowRepository` in `test_model_repository.py`.
- Added extensive tests for `OpcRepository` in `test_opc_repository.py`.
- Updated `test_predictions_batch.py` to reflect new workflow structure and testing methodology.
2025-05-23 17:34:47 -03:00

97 lines
2.9 KiB
Python

from temporalio import workflow, client
from temporalio.worker import Worker
with workflow.unsafe.imports_passed_through():
import os
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.logger import get_logger
from laborious.utils.connectors_config import (
build_postgres_config,
build_mlflow_config,
build_opc_config
)
from sientia_do.notifications.handlers import NotificationHandler
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_SERVERS', 'http://localhost:9092'),
logger=logger,
project_name=os.getenv('PROJECT_NAME', 'laborious'),
pipeline_name='-',
trigger_name='-',
model_name='-',
model='-'
)
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-queue',
workflows=[PredictionsBatch, PredictionProcess,
FormatAndExportPrediction],
activities=[
# Base
activities.prepare_activity,
# 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
]
)
]
handlers = []
for w in workers:
handlers.append(w.run())
logger.info('Workers started successfully')
await asyncio.gather(*handlers)
if __name__ == '__main__':
asyncio.run(main())