vitor-aignosi
89b9892a5b
Update dependencies and enhance ML model retraining functionality - Updated sientia-dataops-library version in requirements.txt from 1.3.3 to 1.3.4. - Incremented image tag in values.yaml from 0.2.4 to 0.2.5 and added a new environment variable MONGODB_TTL_INDEX_HOURS. - Introduced new methods in MLFlowRepository for model retraining and production model updates, including error handling and logging. - Added retrain_model and update_production_model activities in mlflow.py to support model management workflows. - Modified MongoDB connection settings in connectors_config.py for improved security and configuration flexibility.
130 lines
4.3 KiB
Python
130 lines
4.3 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.minimal_retrain import MinimalRetrain
|
|
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,
|
|
build_mongodb_config
|
|
)
|
|
from sientia_do.notifications.handlers import CoreNotificationHandler as 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...')
|
|
|
|
mongo_config = build_mongodb_config()
|
|
notification_handler = NotificationHandler(
|
|
connection_string=mongo_config['connection_string'],
|
|
database=mongo_config['database_name'],
|
|
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', 'laborious')
|
|
)
|
|
|
|
logger.info('Starting Workers...')
|
|
|
|
workers = [
|
|
Worker(
|
|
temporal_client,
|
|
task_queue='minimal_retrain-queue',
|
|
workflows=[MinimalRetrain],
|
|
activities=[
|
|
activities.load_custom_query,
|
|
activities.retrain_model,
|
|
activities.update_production_model,
|
|
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,
|
|
),
|
|
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())
|