Files
sientia-dataops-laborious_t…/laborious/activities/activities.py
vitor-aignosi b68674fe64 SIENTIAPDE-1273
Enhance MLFlowRepository and Activities classes with new methods and metrics

- Added `check_artifact_exists` method to MLFlowRepository for verifying artifact presence in the MLflow Model Registry.
- Implemented `get_prediction_data` method in MLFlowRepository to retrieve prediction data from models.
- Updated Activities class to integrate ModelMetrics for improved metrics handling.
- Enhanced tests for artifact existence checks and prediction data retrieval, ensuring robust coverage for new functionalities.
- Updated various workflows to include `transform_table_name` in input data for better data handling.
2025-11-13 16:40:23 -03:00

134 lines
5.0 KiB
Python

from sientia_do.observability.metrics_controller import MetricsController
from temporalio import workflow
with workflow.unsafe.imports_passed_through():
from typing import Any
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
from sientia_do.observability.logger import Logger
from laborious.activities.gates import Gates
from laborious.activities.mlflow import MLFlow
from laborious.activities.opc import OPC
from laborious.activities.storage import Storage
from laborious.activities.model_metrics import ModelMetrics
class Activities(Storage, MLFlow, Gates, OPC, ModelMetrics):
"""
Main activities orchestrator for the Laborious system.
This class combines functionality from multiple activity classes to provide
a unified interface for all workflow operations. It manages database connections,
MLFlow model interactions, data quality validation, and OPC server communications.
The class implements multiple inheritance to combine specialized functionality:
- Postgres: Database operations and data persistence
- MLFlow: Model inference and transformation operations
- Gates: Data quality validation and filtering mechanisms
- OPC: Real-time data export to OPC servers
Attributes:
postgres_config (dict): PostgreSQL connection configuration
mlflow_config (dict): MLFlow server configuration
opc_config (dict): OPC server configuration
logger (Logger): Logging and observability instance
notification_handler (NotificationHandler): Notification management instance
"""
def __init__(
self,
postgres_config: dict[str, Any],
mlflow_config: dict[str, Any],
minio_config: dict[str, Any],
opc_config: dict[str, Any],
logger: Logger,
notification_handler: NotificationHandler,
):
"""
Initialize the Activities orchestrator with all required configurations.
This constructor initializes all parent classes with their respective
configurations and sets up the foundation for all activity operations.
Args:
postgres_config: PostgreSQL connection configuration dictionary
Required keys: host, port, user, password, dbname, min_connections, max_connections
mlflow_config: MLFlow server configuration dictionary
Required keys: host, port, username, password
opc_config: OPC server configuration dictionary
Can contain multiple server configurations
logger: Logger instance for observability and debugging
notification_handler: Notification handler for alerts and monitoring
Raises:
Exception: If any parent class initialization fails
"""
metrics_controller = MetricsController(logger=logger)
# Initialize parent classes
Storage.__init__(
self,
host=postgres_config['host'],
port=postgres_config['port'],
user=postgres_config['user'],
password=postgres_config['password'],
dbname=postgres_config['dbname'],
min_connections=postgres_config['min_connections'],
max_connections=postgres_config['max_connections'],
minio_config=minio_config,
logger=logger,
notification_handler=notification_handler,
metrics_controller=metrics_controller,
)
MLFlow.__init__(
self,
mlflow_host=mlflow_config['host'],
mlflow_port=mlflow_config['port'],
mlflow_username=mlflow_config['username'],
mlflow_password=mlflow_config['password'],
minio_config=minio_config,
logger=logger,
notification_handler=notification_handler,
metrics_controller=metrics_controller,
)
Gates.__init__(
self,
logger=logger,
notification_handler=notification_handler,
metrics_controller=metrics_controller,
)
OPC.__init__(
self,
opc_servers=opc_config,
logger=logger,
notification_handler=notification_handler,
metrics_controller=metrics_controller,
)
ModelMetrics.__init__(
self,
logger=logger,
notification_handler=notification_handler,
metrics_controller=metrics_controller,
)
async def shutdown(self):
"""
Gracefully shutdown all activities and clean up resources.
This method ensures proper cleanup of all resources including:
- PostgreSQL connection pools
- OPC server connections
- Any other resources that need explicit cleanup
The method should be called before the application terminates to ensure
proper resource cleanup and prevent resource leaks.
"""
Storage.close(self)
MLFlow.close(self)
Gates.close(self)
await OPC.close(self)
ModelMetrics.close(self)