SIENTIAPDE-1243: Refactor and enhance model manager activities and workflows
This commit includes several changes: - Reorganized imports and class inheritance in activities.py, gates.py and mlflow.py for better readability and maintainability. - Improved error handling and logging in gates.py and mlflow.py. - Added input validation and filtering in gates.py to ensure data quality. - Enhanced prediction formatting and storage policy management in gates.py. - Updated metrics.py to use consistent naming conventions and labels. - Refactored connectors_config.py to use type hints and improve code clarity. - Updated conditional and MLFlow filters for better data quality checks. - Improved model repository logic for retraining and updating models. - Enhanced worker.py to include SDK metrics and improved error handling. - Refactored workflows for better modularity and error handling. - Updated tests to reflect the changes and improve test coverage.
This commit is contained in:
Bruno Domingues
@@ -25,32 +25,35 @@ Environment Variables:
|
||||
- PROJECT_NAME: Project name for notifications (default: model_manager)
|
||||
"""
|
||||
|
||||
from temporalio import workflow, client
|
||||
from temporalio.worker import Worker, PollerBehaviorAutoscaling
|
||||
from temporalio.runtime import Runtime, TelemetryConfig, PrometheusConfig
|
||||
from temporalio import client, workflow
|
||||
from temporalio.runtime import PrometheusConfig, Runtime, TelemetryConfig
|
||||
from temporalio.worker import PollerBehaviorAutoscaling, Worker
|
||||
|
||||
with workflow.unsafe.imports_passed_through():
|
||||
import asyncio
|
||||
import os
|
||||
import sys
|
||||
import asyncio
|
||||
from model_manager.workflows.minimal_retrain import MinimalRetrain
|
||||
from model_manager.workflows.predictions_batch import PredictionsBatch
|
||||
from model_manager.workflows.sub_workflows.prediction_process import PredictionProcess
|
||||
from model_manager.workflows.sub_workflows.format_and_export_prediction import \
|
||||
FormatAndExportPrediction
|
||||
from model_manager.activities.activities import Activities
|
||||
from model_manager.utils.connectors_config import (
|
||||
build_postgres_config,
|
||||
build_mlflow_config,
|
||||
build_mongodb_config
|
||||
)
|
||||
|
||||
from prometheus_client import start_http_server
|
||||
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
|
||||
from sientia_do.observability.logger import get_logger
|
||||
|
||||
from model_manager import metrics
|
||||
from prometheus_client import start_http_server
|
||||
from model_manager.activities.activities import Activities
|
||||
from model_manager.utils.connectors_config import (
|
||||
build_mlflow_config,
|
||||
build_mongodb_config,
|
||||
build_postgres_config,
|
||||
)
|
||||
from model_manager.workflows.minimal_retrain import MinimalRetrain
|
||||
from model_manager.workflows.predictions_batch import PredictionsBatch
|
||||
from model_manager.workflows.sub_workflows.format_and_export_prediction import (
|
||||
FormatAndExportPrediction,
|
||||
)
|
||||
from model_manager.workflows.sub_workflows.prediction_process import PredictionProcess
|
||||
|
||||
POD_ID = os.getenv('POD_ID')
|
||||
SDK_METRICS_PORT = int(os.getenv('HTTP_SDK_METRICS_PORT', "9091"))
|
||||
SDK_METRICS_PORT = int(os.getenv('HTTP_SDK_METRICS_PORT', '9091'))
|
||||
|
||||
|
||||
async def main():
|
||||
@@ -85,7 +88,7 @@ async def main():
|
||||
|
||||
logger.custom_info(f'Starting Worker with POD_ID: {POD_ID}', metadata)
|
||||
|
||||
logger.custom_info("Starting prometheus client...", metadata)
|
||||
logger.custom_info('Starting prometheus client...', metadata)
|
||||
start_prometheus_server()
|
||||
|
||||
logger.custom_info('Starting Notification Handler...', metadata)
|
||||
@@ -95,7 +98,7 @@ async def main():
|
||||
connection_string=mongo_config['connection_string'],
|
||||
database=mongo_config['database_name'],
|
||||
logger=logger,
|
||||
project_name=os.getenv('PROJECT_NAME', 'model-manager')
|
||||
project_name=os.getenv('PROJECT_NAME', 'model-manager'),
|
||||
)
|
||||
|
||||
logger.custom_info('Starting Activities...', metadata)
|
||||
@@ -104,16 +107,14 @@ async def main():
|
||||
postgres_config=build_postgres_config(),
|
||||
mlflow_config=build_mlflow_config(),
|
||||
logger=logger,
|
||||
notification_handler=notification_handler
|
||||
notification_handler=notification_handler,
|
||||
)
|
||||
|
||||
logger.custom_info(
|
||||
f'Starting SDK Metrics Server on port {SDK_METRICS_PORT}...', metadata)
|
||||
logger.custom_info(f'Starting SDK Metrics Server on port {SDK_METRICS_PORT}...', metadata)
|
||||
|
||||
new_runtime = Runtime(
|
||||
telemetry=TelemetryConfig(
|
||||
metrics=PrometheusConfig(
|
||||
bind_address=f"0.0.0.0:{SDK_METRICS_PORT}")
|
||||
metrics=PrometheusConfig(bind_address=f'0.0.0.0:{SDK_METRICS_PORT}')
|
||||
)
|
||||
)
|
||||
|
||||
@@ -122,7 +123,7 @@ async def main():
|
||||
temporal_client = await client.Client.connect(
|
||||
target_host=host,
|
||||
namespace=os.getenv('TEMPORAL_NAMESPACE', 'model-manager'),
|
||||
runtime=new_runtime
|
||||
runtime=new_runtime,
|
||||
)
|
||||
|
||||
logger.custom_info('Starting Workers...', metadata)
|
||||
@@ -136,20 +137,19 @@ async def main():
|
||||
activities.load_custom_query,
|
||||
activities.retrain_model,
|
||||
activities.update_production_model,
|
||||
activities.export_data_to_postgres
|
||||
activities.export_data_to_postgres,
|
||||
],
|
||||
max_concurrent_workflow_tasks=50,
|
||||
max_concurrent_activities=50,
|
||||
max_concurrent_local_activities=50,
|
||||
max_cached_workflows=200,
|
||||
workflow_task_poller_behavior=PollerBehaviorAutoscaling(),
|
||||
activity_task_poller_behavior=PollerBehaviorAutoscaling()
|
||||
activity_task_poller_behavior=PollerBehaviorAutoscaling(),
|
||||
),
|
||||
Worker(
|
||||
temporal_client,
|
||||
task_queue='predictions_batch-queue',
|
||||
workflows=[PredictionsBatch, PredictionProcess,
|
||||
FormatAndExportPrediction],
|
||||
workflows=[PredictionsBatch, PredictionProcess, FormatAndExportPrediction],
|
||||
activities=[
|
||||
# MLFlow
|
||||
activities.request_predict,
|
||||
@@ -165,15 +165,15 @@ async def main():
|
||||
activities.load_custom_query,
|
||||
activities.repeat_last_prediction,
|
||||
activities.export_data_to_postgres,
|
||||
activities.write_metrics
|
||||
activities.write_metrics,
|
||||
],
|
||||
max_concurrent_workflow_tasks=50,
|
||||
max_concurrent_activities=50,
|
||||
max_concurrent_local_activities=50,
|
||||
max_cached_workflows=200,
|
||||
workflow_task_poller_behavior=PollerBehaviorAutoscaling(),
|
||||
activity_task_poller_behavior=PollerBehaviorAutoscaling()
|
||||
)
|
||||
activity_task_poller_behavior=PollerBehaviorAutoscaling(),
|
||||
),
|
||||
]
|
||||
|
||||
handlers = []
|
||||
@@ -186,8 +186,8 @@ async def main():
|
||||
# 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.custom_error(f"An unhandled exception occurred: {e}", metadata)
|
||||
except BaseException as e: # noqa: BLE001
|
||||
logger.custom_error(f'An unhandled exception occurred: {e}', metadata)
|
||||
finally:
|
||||
if notification_handler:
|
||||
notification_handler.shutdown()
|
||||
@@ -216,12 +216,12 @@ def start_prometheus_server():
|
||||
SystemExit: If the metrics server fails to start
|
||||
"""
|
||||
try:
|
||||
port = int(os.getenv("HTTP_METRICS_PORT", 9090))
|
||||
port = int(os.getenv('HTTP_METRICS_PORT', 9090))
|
||||
start_http_server(port)
|
||||
print(f"Prometheus server started on port {port}.")
|
||||
print(f'Prometheus server started on port {port}.')
|
||||
metrics.APP_UP.labels(pod_id=POD_ID).set(1) # Mark app as UP
|
||||
except Exception as e:
|
||||
print(f"Failed to start Prometheus server: {e}")
|
||||
except Exception as e: # noqa: BLE001
|
||||
print(f'Failed to start Prometheus server: {e}')
|
||||
os._exit(1)
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user