503d9aa485
Update dependencies and refactor input filter handling for consistency - Updated sientia-dataops-library dependency version from 1.10.3 to 1.10.4 in requirements.txt. - Refactored input filter handling in the Gates class to read policy and config keys in a case-insensitive manner. - Updated test cases to ensure consistency in filter key naming conventions across various scenarios.
505 lines
20 KiB
Python
505 lines
20 KiB
Python
from temporalio import activity, workflow
|
|
|
|
with workflow.unsafe.imports_passed_through():
|
|
import traceback
|
|
from typing import Any
|
|
|
|
import numpy as np
|
|
from pandas import to_datetime
|
|
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
|
|
from sientia_do.notifications.models import NotificationLevel
|
|
from sientia_do.observability.logger import Logger
|
|
from sientia_do.observability.metrics_controller import MetricsController
|
|
from sientia_do.repository.minio_repository import MinioRepository
|
|
from sientia_do.temporal.constants import (
|
|
DATETIME_FORMAT,
|
|
DATETIME_FORMAT_MS_WITH_TZ,
|
|
DATETIME_FORMAT_WITH_TZ,
|
|
now,
|
|
)
|
|
from sientia_do.utils.formatters import create_sample_dict
|
|
|
|
from laborious.utils.models.minio_dataframe_payload import MinioDataFramePayload
|
|
from laborious.utils.repository.minio_manager import MinioManager
|
|
from laborious.utils.repository.model_repository import MLFlowRepository
|
|
|
|
|
|
class MLFlow(MinioManager):
|
|
"""
|
|
MLFlow integration activities for model inference operations.
|
|
|
|
This class provides activities for interacting with MLFlow models, including
|
|
data transformation and prediction operations. It handles authentication,
|
|
data preprocessing, and model management with configurable retention policies.
|
|
|
|
The class implements comprehensive error handling and logging for all
|
|
MLFlow operations, ensuring reliable model inference in production environments.
|
|
|
|
Attributes:
|
|
mlflow_host (str): MLFlow server hostname
|
|
mlflow_port (int): MLFlow server port
|
|
mlflow_username (str): MLFlow authentication username
|
|
mlflow_password (str): MLFlow authentication password
|
|
model_monitoring_repository (MLFlowRepository): Repository for MLFlow operations
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
mlflow_host: str,
|
|
mlflow_port: int,
|
|
mlflow_username: str,
|
|
mlflow_password: str,
|
|
minio_repository: MinioRepository | None = None,
|
|
logger: Logger | None = None,
|
|
notification_handler: NotificationHandler | None = None,
|
|
metrics_controller: MetricsController | None = None,
|
|
):
|
|
"""
|
|
Initialize MLFlow activities with server configuration.
|
|
|
|
Args:
|
|
mlflow_host: MLFlow server hostname or IP address
|
|
mlflow_port: MLFlow server port number
|
|
mlflow_username: Username for MLFlow authentication
|
|
mlflow_password: Password for MLFlow authentication
|
|
logger: Logger instance for observability and debugging
|
|
notification_handler: Notification handler for alerts and monitoring
|
|
|
|
Raises:
|
|
Exception: If MLFlowRepository initialization fails
|
|
"""
|
|
MinioManager.__init__(
|
|
self, minio_repository, logger, notification_handler, metrics_controller
|
|
)
|
|
self.mlflow_host = mlflow_host
|
|
self.mlflow_port = mlflow_port
|
|
self.mlflow_username = mlflow_username
|
|
self.mlflow_password = mlflow_password
|
|
|
|
self.model_monitoring_repository = MLFlowRepository(
|
|
f'{mlflow_host}:{mlflow_port}',
|
|
mlflow_username,
|
|
mlflow_password,
|
|
logger,
|
|
notification_handler,
|
|
metrics_controller,
|
|
)
|
|
|
|
def close(self) -> None:
|
|
"""
|
|
Close the MLFlow activity and clean up resources.
|
|
"""
|
|
MinioManager.close(self)
|
|
|
|
def __del__(self):
|
|
self.close()
|
|
|
|
@activity.defn(name='request_transform')
|
|
async def request_transform(self, input_data: dict[str, Any]) -> MinioDataFramePayload:
|
|
"""
|
|
Transform input data using MLFlow models.
|
|
|
|
This activity processes input data through MLFlow model transformation,
|
|
including data preprocessing, format conversion, and validation. It handles
|
|
data deduplication, pivoting, and cleanup to ensure optimal model performance.
|
|
|
|
The transformation process includes:
|
|
1. Data deduplication based on variable and timestamp
|
|
2. Data pivoting for model input format
|
|
3. Null value handling and cleanup
|
|
4. MLFlow model transformation request
|
|
5. Response validation and logging
|
|
|
|
Args:
|
|
input_data: Configuration and data for transformation
|
|
Required keys:
|
|
- metadata (dict): Workflow execution metadata
|
|
- data (dict): Input data for transformation
|
|
- model_name (str): Name of the MLFlow model to use
|
|
- model_retention (int): Model retention period in minutes
|
|
|
|
Returns:
|
|
dict: Transformed data from MLFlow model
|
|
|
|
Raises:
|
|
Exception: If transformation fails or MLFlow model is unavailable
|
|
"""
|
|
metadata = input_data['metadata']
|
|
self.info('Transforming data...', metadata)
|
|
|
|
payload = MinioDataFramePayload.from_dict(input_data['data'])
|
|
data = await payload.retrieve(self.minio_repository, metadata)
|
|
|
|
model_name = input_data['model_name']
|
|
model_config = input_data.get('model_config', {})
|
|
|
|
self.debug('Raw input data:', metadata)
|
|
self.debug(data.head(5).to_string(), metadata)
|
|
|
|
# Sort by created_at in descending order and keep first occurrence of each variable/timestamp pair
|
|
data = data.sort_values('created_at', ascending=False).drop_duplicates(
|
|
subset=['variable', 'timestamp'], keep='first'
|
|
)
|
|
|
|
# Pivot data for model input format
|
|
data = data.pivot(index='timestamp', columns='variable', values='value')
|
|
data.fillna(np.nan, inplace=True)
|
|
|
|
data.columns.name = None
|
|
data.index.name = None
|
|
|
|
data['timestamp'] = data.index
|
|
|
|
self.debug(f'Processed input data: \n {data.to_csv()}', metadata)
|
|
|
|
# Request transformation from MLFlow model
|
|
response_data = await self.model_monitoring_repository.transform(
|
|
model_name, data, model_config, metadata
|
|
)
|
|
|
|
self.debug(
|
|
f'Transform raw response data: \n {create_sample_dict(response_data, max_items=5, max_depth=5)}',
|
|
metadata,
|
|
)
|
|
|
|
self.debug(
|
|
f'Transform response data: \n {create_sample_dict(response_data, max_items=5, max_depth=5)}',
|
|
metadata,
|
|
)
|
|
|
|
self.info('Data transformed successfully', metadata)
|
|
|
|
if not response_data.get('success', False):
|
|
return await MinioDataFramePayload.from_dataframe(
|
|
dataframe=None,
|
|
minio_repo=self.minio_repository,
|
|
model_name=model_name,
|
|
operation='transform',
|
|
status=response_data,
|
|
workflow_metadata=metadata,
|
|
last_timestamp=payload.last_timestamp,
|
|
)
|
|
|
|
return await MinioDataFramePayload.from_dataframe(
|
|
dataframe=response_data['content'],
|
|
minio_repo=self.minio_repository,
|
|
model_name=model_name,
|
|
operation='transform',
|
|
workflow_metadata=metadata,
|
|
status={
|
|
'success': True,
|
|
},
|
|
last_timestamp=payload.last_timestamp,
|
|
)
|
|
|
|
@activity.defn(name='request_predict')
|
|
async def request_predict(self, input_data: dict[str, Any]) -> MinioDataFramePayload:
|
|
"""
|
|
Execute predictions using MLFlow models.
|
|
|
|
This activity performs ML model inference using MLFlow models with the
|
|
transformed data. It handles data format conversion, null value processing,
|
|
and model prediction requests with comprehensive error handling.
|
|
|
|
The prediction process includes:
|
|
1. Data format validation and cleanup
|
|
2. Null value handling for model compatibility
|
|
3. MLFlow model prediction request
|
|
4. Response validation and logging
|
|
5. Performance monitoring and metrics
|
|
|
|
Args:
|
|
input_data: Configuration and data for prediction
|
|
Required keys:
|
|
- metadata (dict): Workflow execution metadata
|
|
- data (dict): Transformed data for prediction
|
|
- model_name (str): Name of the MLFlow model to use
|
|
- model_retention (int): Model retention period in minutes
|
|
|
|
Returns:
|
|
dict: Prediction results from MLFlow model
|
|
|
|
Raises:
|
|
Exception: If prediction fails or MLFlow model is unavailable
|
|
"""
|
|
metadata = input_data['metadata']
|
|
self.info('Predicting data...', metadata)
|
|
|
|
payload = MinioDataFramePayload.from_dict(input_data['data'])
|
|
data = await payload.retrieve(self.minio_repository, metadata)
|
|
|
|
model_name = input_data['model_name']
|
|
model_config = input_data.get('model_config', {})
|
|
|
|
self.debug(f'Input data for: \n {data.head(5).to_string()}', metadata)
|
|
|
|
# Convert numpy.nan to None for model compatibility
|
|
data.replace(np.nan, None, inplace=True)
|
|
|
|
data['timestamp'] = data.index
|
|
data['timestamp'] = to_datetime(
|
|
data['timestamp'], format=DATETIME_FORMAT_WITH_TZ
|
|
).dt.strftime(DATETIME_FORMAT)
|
|
|
|
# Request prediction from MLFlow model
|
|
response_data = await self.model_monitoring_repository.predict(
|
|
model_name, data, model_config, metadata
|
|
)
|
|
|
|
self.debug(
|
|
f'Prediction response data: \n {create_sample_dict(response_data, max_items=5, max_depth=5)}',
|
|
metadata,
|
|
)
|
|
|
|
self.info('Data predicted successfully', metadata)
|
|
|
|
if not response_data.get('success', False):
|
|
return await MinioDataFramePayload.from_dataframe(
|
|
dataframe=None,
|
|
minio_repo=self.minio_repository,
|
|
model_name=model_name,
|
|
operation='predict',
|
|
status=response_data,
|
|
workflow_metadata=metadata,
|
|
last_timestamp=payload.last_timestamp,
|
|
)
|
|
|
|
return await MinioDataFramePayload.from_dataframe(
|
|
dataframe=response_data['content'],
|
|
minio_repo=self.minio_repository,
|
|
model_name=model_name,
|
|
operation='predict',
|
|
workflow_metadata=metadata,
|
|
status={
|
|
'success': True,
|
|
},
|
|
last_timestamp=payload.last_timestamp,
|
|
)
|
|
|
|
@activity.defn(name='retrain_model')
|
|
async def retrain_model(self, input_data: dict[str, Any]) -> dict[str, Any]:
|
|
"""
|
|
Retrain MLFlow models with updated training data.
|
|
|
|
This activity orchestrates the complete model retraining process,
|
|
including data preparation, model retraining execution, and result
|
|
validation. It handles data preprocessing, column cleanup, and
|
|
comprehensive error handling for production model management.
|
|
|
|
The retraining process includes:
|
|
1. Data timestamp extraction and validation
|
|
2. Column cleanup and data preparation
|
|
3. Data pivoting for model input format
|
|
4. MLFlow model retraining execution
|
|
5. Result validation and error handling
|
|
|
|
Args:
|
|
input_data (dict): Input data containing:
|
|
- metadata (dict): Workflow execution metadata
|
|
- data (dict[str, Any]): Training data for model retraining
|
|
- model_name (str): Name of the MLFlow model to retrain
|
|
|
|
Returns:
|
|
dict: Retraining results containing:
|
|
- status (str): Retraining operation status
|
|
- timestamp (str): Timestamp of the retraining operation
|
|
- experiment (str): MLFlow experiment identifier
|
|
|
|
Raises:
|
|
Exception: If retraining fails or encounters critical errors
|
|
"""
|
|
|
|
if self.minio_repository is None:
|
|
raise ValueError('Minio repository not initialized')
|
|
|
|
metadata = input_data['metadata']
|
|
|
|
try:
|
|
# Payload-based retrain input (inline dict or MinIO offloaded).
|
|
payload = MinioDataFramePayload.from_dict(input_data['data'])
|
|
data = await payload.retrieve(self.minio_repository, metadata)
|
|
|
|
except Exception as e:
|
|
trace = traceback.format_exc()
|
|
await self.send_notification_async(
|
|
metadata=metadata,
|
|
notification_id='ERROR_LOADING_RETRAIN_DATA',
|
|
message=f'Error loading retrain data: {e}',
|
|
block='retrain_model',
|
|
level=NotificationLevel.ERROR,
|
|
attachment_content=trace,
|
|
)
|
|
self.error(trace, metadata)
|
|
return {
|
|
'success': False,
|
|
'message': f'Error loading retrain data: {e}',
|
|
'traceback': trace,
|
|
'timestamp': now().strftime(DATETIME_FORMAT_MS_WITH_TZ),
|
|
}
|
|
|
|
self.debug(f'Retrain data loaded successfully: shape {data.shape}', metadata)
|
|
|
|
model_name = input_data['model_name']
|
|
model_config = input_data.get('model_config', {})
|
|
|
|
self.info(f'Retraining model {model_name}...', metadata)
|
|
|
|
timestamp = data['timestamp'].max()
|
|
self.debug(f'Timestamp: {timestamp}', metadata)
|
|
|
|
# Sort by created_at in descending order and keep first occurrence of each variable/timestamp pair
|
|
if 'created_at' in data.columns:
|
|
data = data.sort_values('created_at', ascending=False).drop_duplicates(
|
|
subset=['variable', 'timestamp'], keep='first'
|
|
)
|
|
else:
|
|
data = data.drop_duplicates(subset=['variable', 'timestamp'], keep='first')
|
|
|
|
data.drop(columns=['model_id'], inplace=True, errors='ignore')
|
|
data.drop(columns=['created_at'], inplace=True, errors='ignore')
|
|
|
|
# Pivot data for model input format
|
|
data = data.pivot(index='timestamp', columns='variable', values='value')
|
|
data.fillna(np.nan, inplace=True)
|
|
# data.reset_index(inplace=True)
|
|
data.columns.name = None
|
|
|
|
data['timestamp'] = data.index
|
|
data['timestamp'] = to_datetime(
|
|
data['timestamp'], format=DATETIME_FORMAT_WITH_TZ
|
|
).dt.strftime(DATETIME_FORMAT)
|
|
data['timestamp'] = to_datetime(data['timestamp'], format=DATETIME_FORMAT)
|
|
|
|
data.columns.name = None
|
|
|
|
retrain_output = await self.model_monitoring_repository.retrain_model(
|
|
data=data, model_name=model_name, model_config=model_config, metadata=metadata
|
|
)
|
|
|
|
if not retrain_output['success']:
|
|
trace = retrain_output['traceback']
|
|
await self.send_notification_async(
|
|
metadata=metadata,
|
|
notification_id='RETRAIN_MODEL_ERROR',
|
|
message=f'Error retraining model {model_name}: {retrain_output["message"]}',
|
|
block='retrain_model',
|
|
level=NotificationLevel.ERROR,
|
|
attachment_content=trace,
|
|
)
|
|
self.error(trace, metadata=metadata)
|
|
|
|
return {**retrain_output, 'timestamp': timestamp}
|
|
|
|
@activity.defn(name='update_production_model')
|
|
async def update_production_model(self, input_data: dict[str, Any]) -> dict[Any, Any]:
|
|
"""
|
|
Update production model with newly trained model version.
|
|
|
|
This activity manages the critical process of updating production
|
|
models with newly trained versions. It handles model deployment,
|
|
status tracking, and comprehensive reporting for operational
|
|
visibility and audit trails.
|
|
|
|
The update process includes:
|
|
1. Production model update execution
|
|
2. Status and metadata tracking
|
|
3. Comprehensive reporting and logging
|
|
4. Error handling and notification
|
|
5. Audit trail maintenance
|
|
|
|
Args:
|
|
input_data (dict): Input data containing:
|
|
- metadata (dict): Workflow execution metadata
|
|
- model_name (str): Name of the MLFlow model to update
|
|
- experiment (str): MLFlow experiment identifier
|
|
- model_id (str): Unique identifier for the model version
|
|
- timestamp (str): Timestamp of the update operation
|
|
- status (str): Current status of the model update
|
|
|
|
Returns:
|
|
dict[Any, Any]: Comprehensive update report containing:
|
|
- model_id (str): Model version identifier
|
|
- model_name (str): Name of the updated model
|
|
- timestamp (str): Update operation timestamp
|
|
- status (str): Update operation status
|
|
- Additional MLFlow response metadata
|
|
|
|
Raises:
|
|
Exception: If production model update fails
|
|
"""
|
|
metadata = input_data['metadata']
|
|
model_name = input_data['model_name']
|
|
experiment = input_data['experiment']
|
|
self.info(
|
|
f'Updating production model {model_name} from experiment {experiment}...', metadata
|
|
)
|
|
|
|
try:
|
|
response = await self.model_monitoring_repository.update_production_model(
|
|
experiment=experiment, model_name=model_name, metadata=metadata
|
|
)
|
|
|
|
self.info(f'Production model {model_name} updated successfully', metadata)
|
|
return response
|
|
|
|
except Exception as e:
|
|
trace = traceback.format_exc()
|
|
await self.send_notification_async(
|
|
metadata=metadata,
|
|
notification_id='UPDATE_PRODUCTION_MODEL_ERROR',
|
|
message=f'Error updating production model {model_name}: {e}',
|
|
block='update_production_model',
|
|
level=NotificationLevel.ERROR,
|
|
attachment_content=trace,
|
|
)
|
|
self.error(trace, metadata=metadata)
|
|
raise e
|
|
|
|
@activity.defn(name='get_reference_data')
|
|
async def get_reference_data(self, input_data: dict[str, Any]) -> list[dict] | None:
|
|
"""
|
|
Get reference data from the MLflow Model Registry.
|
|
|
|
This method retrieves evaluation reference data stored as artifacts in the
|
|
MLflow Model Registry. The reference data is typically used for model
|
|
drift detection, performance comparison, and quality validation. The method
|
|
loads the data from a CSV artifact file and formats timestamps for
|
|
consistent processing.
|
|
|
|
The method handles:
|
|
1. Loading evaluation data artifact from MLflow Model Registry
|
|
2. Timestamp parsing and formatting for consistency
|
|
3. Data conversion to dictionary format for workflow consumption
|
|
4. Graceful handling of missing reference data
|
|
|
|
Args:
|
|
input_data (dict): Input data containing:
|
|
- metadata (dict): Workflow execution metadata
|
|
- model_name (str): Name of the MLFlow model to get reference data from
|
|
|
|
Returns:
|
|
list[dict[Hashable, Any]] | None: Reference data from the MLflow Model Registry
|
|
as a list of dictionaries. Returns None if reference data is not found
|
|
or if the artifact does not exist.
|
|
|
|
Raises:
|
|
Exception: If artifact loading fails or encounters errors during processing
|
|
"""
|
|
|
|
metadata = input_data['metadata']
|
|
model_name = input_data['model_name']
|
|
artifact = 'evaluation_data.csv'
|
|
|
|
reference_data = await self.model_monitoring_repository.load_artifact_dataframe(
|
|
model_name=model_name, artifact_path=artifact, metadata=metadata
|
|
)
|
|
|
|
if reference_data is None:
|
|
self.warning(f'Reference data not found for model {model_name}', metadata)
|
|
return None
|
|
|
|
reference_data['timestamp'] = to_datetime(reference_data['timestamp'])
|
|
reference_data['timestamp'] = reference_data['timestamp'].dt.strftime(DATETIME_FORMAT)
|
|
|
|
return reference_data.to_dict(orient='records')
|