117 lines
4.8 KiB
Python
117 lines
4.8 KiB
Python
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 model_manager.activities.experiment_tracking import ExperimentTracking
|
|
from model_manager.activities.gates import Gates
|
|
from model_manager.activities.minio import MinIO
|
|
from model_manager.activities.mlflow import MLFlow
|
|
|
|
|
|
class Activities(ExperimentTracking, MLFlow, MinIO, Gates):
|
|
"""
|
|
Main activities orchestrator for the Model Manager 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, MinIO storage operations, and data quality validation.
|
|
|
|
The class implements multiple inheritance to combine specialized functionality:
|
|
- ExperimentTracking: ML experiment lifecycle tracking and database operations (extends Postgres)
|
|
- MLFlow: Model inference and transformation operations
|
|
- MinIO: Object storage operations (file upload/download/delete)
|
|
- Gates: Data quality validation and filtering mechanisms
|
|
|
|
Attributes:
|
|
postgres_config (dict): PostgreSQL connection configuration
|
|
mlflow_config (dict): MLFlow server configuration
|
|
minio_config (dict): MinIO storage 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],
|
|
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
|
|
minio_config: MinIO storage configuration dictionary
|
|
Required keys: endpoint_url, access_key, secret_key, region, use_ssl
|
|
logger: Logger instance for observability and debugging
|
|
notification_handler: Notification handler for alerts and monitoring
|
|
|
|
Raises:
|
|
Exception: If any parent class initialization fails
|
|
"""
|
|
# Initialize parent classes
|
|
ExperimentTracking.__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'],
|
|
logger=logger,
|
|
notification_handler=notification_handler,
|
|
)
|
|
|
|
MLFlow.__init__(
|
|
self,
|
|
mlflow_host=mlflow_config['host'],
|
|
mlflow_port=mlflow_config['port'],
|
|
mlflow_username=mlflow_config['username'],
|
|
mlflow_password=mlflow_config['password'],
|
|
logger=logger,
|
|
notification_handler=notification_handler,
|
|
)
|
|
|
|
MinIO.__init__(
|
|
self,
|
|
endpoint_url=minio_config['endpoint_url'],
|
|
access_key=minio_config['access_key'],
|
|
secret_key=minio_config['secret_key'],
|
|
region=minio_config['region'],
|
|
use_ssl=minio_config['use_ssl'],
|
|
max_retry_attempts=minio_config['max_retry_attempts'],
|
|
retry_mode=minio_config['retry_mode'],
|
|
connect_timeout=minio_config['connect_timeout'],
|
|
read_timeout=minio_config['read_timeout'],
|
|
logger=logger,
|
|
notification_handler=notification_handler,
|
|
)
|
|
|
|
Gates.__init__(self, logger=logger, notification_handler=notification_handler)
|
|
|
|
async def shutdown(self):
|
|
"""
|
|
Gracefully shutdown all activities and clean up resources.
|
|
|
|
This method ensures proper cleanup of all resources including:
|
|
- PostgreSQL connection pools (via ExperimentTracking)
|
|
- 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.
|
|
"""
|
|
ExperimentTracking.close(self)
|